Skip to main content

memra_engine/
lib.rs

1//! memra engine: Stage-1 correctness-first forward-pass kernels + ops, on sm_120 via cudarc.
2
3use cudarc::driver::sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES;
4use cudarc::driver::{
5    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut,
6    DeviceSlice, LaunchConfig, PushKernelArg,
7};
8use cudarc::nvrtc::Ptx;
9use std::sync::{Arc, Mutex};
10
11const GDN_K2_DYNAMIC_SHARED_BYTES: u32 = 67_072;
12
13/// The default dynamic-shared-memory launch bound the naive SDPA family lives under: past
14/// `T_kv * 4 > 48KB` (T_kv > 12288) the smem kernel cannot launch — the measured
15/// dspark/full-attn long-ctx crash class. `sdpa_naive` dispatches to the byte-identical
16/// gmem-scores twin above this line.
17const SDPA_NAIVE_SMEM_MAX: usize = 48 * 1024;
18
19/// Guard on the gmem twin's `n_head * T * T_kv * 4`-byte scores workspace. The shapes that
20/// legitimately hit the smem bound are tall-KV blocks (T <= draft block size), which land in
21/// the tens of MB; 1 GiB refuses a square T==T_kv misuse before it silently eats the card.
22const SDPA_NAIVE_GMEM_WS_MAX: usize = 1 << 30;
23
24#[cfg(debug_assertions)]
25pub(crate) fn debug_assert_tensor_stream_device<T>(
26    tensor: &CudaSlice<T>,
27    stream: &CudaStream,
28    site: &str,
29) {
30    let tensor_dev = tensor.ordinal();
31    let stream_dev = stream.context().ordinal();
32    assert_eq!(
33        tensor_dev, stream_dev,
34        "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
35    );
36}
37
38fn ensure_tensor_stream_device<T>(
39    tensor: &impl DeviceSlice<T>,
40    stream: &CudaStream,
41    site: &str,
42) -> Result<(), Box<dyn std::error::Error>> {
43    let tensor_dev = tensor.stream().context().ordinal();
44    let stream_dev = stream.context().ordinal();
45    if tensor_dev != stream_dev {
46        return Err(format!(
47            "PP cross-device tensor access at {site}: tensor on dev{tensor_dev}, \
48             stream on dev{stream_dev}"
49        )
50        .into());
51    }
52    Ok(())
53}
54
55pub use memra_gguf;
56pub use memra_runtime;
57
58pub mod forward;
59pub mod hybrid;
60pub mod hybrid_forward;
61pub mod hyper;
62pub mod model;
63pub mod sigrouter_contract;
64pub mod vision;
65pub mod vision_gemma;
66pub mod vision_glm5;
67pub mod vision_pre;
68pub mod vision_step;
69/// The dual cache lives in the shared `memra-kv` crate (Phase D extraction); this
70/// re-export keeps every `crate::cache::` / `memra_engine::cache::` path unchanged.
71pub mod cache {
72    pub use memra_kv::*;
73}
74pub mod decode;
75pub mod decode_batch;
76pub mod dflash;
77pub mod eagle;
78/// Measured expert-placement map (`MEMRA_EP_MAP`; glm5 alias honored) — the fail-closed
79/// `memra-ep-map-v1` reader every family's EP shard builders consume (fleet-shared by
80/// design; glm5 is the first consumer). (LAW:coactivation-expert-placement; maps are
81/// minted by the shared fleet tool from `MEMRA_MOE_WEIGHT_TRACE` traces). No CUDA deps.
82pub mod ep_map;
83pub mod gemma_spec;
84pub mod glm5_decode_graph;
85pub mod glm5_sel_ledger;
86pub mod glm5_tp;
87/// glm5_next T-parallel speculative verify: the rows-walk verify, per-step KDA state-column
88/// rollback, latent/kpool truncation, and the MEMRA_GLM5_SPEC-gated draft->verify->rollback
89/// loop over the native MTP head (lane/glm5-tparallel-verify).
90pub mod glm_spec;
91pub mod graph_update;
92pub mod kda;
93/// MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1.
94/// Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent
95/// oracle for the MLA kernel family (`research/mla-bringup-20260801/DESIGN.md`). No CUDA deps.
96pub mod mla;
97pub mod mla_ffi;
98pub mod moe_sel_dump;
99pub mod moesd;
100pub mod parallel;
101pub mod plan_backend;
102pub mod pp;
103pub mod progress;
104/// qwen4_exp (Qwen3.8-Flash-Next) GPU eager forward — onboarding phase 7, correctness arm
105/// gated against memra-reference (research/qwen4exp-bringup-20260829/GPU-EAGER.md).
106pub mod qwen4exp_gpu;
107pub mod round_stream;
108pub mod spec;
109/// Per-burst spec-round phase attribution (`MEMRA_SPEC_TRACE`; glm5 alias honored) —
110/// the draft/verify/accept/rollback/maintenance split every spec family owns, with
111/// caller-tagged emit lines so banked receipts keep their grep shape. No CUDA deps
112/// beyond the stream drains at phase boundaries.
113pub mod spec_phase;
114pub mod tp;
115pub mod tp_transport;
116pub use memra_sampling as sampler;
117
118/// In-house MoE router GEMV on the spec-verify small-t path (DEFAULT ON since 2026-07-10:
119/// battery green on 35B p2/p3 K=1..8, acceptance bit-identical, +2-4% spec e2e — replaces
120/// ~240 per-column cuBLAS gemv launches/round). MEMRA_ROUTER_KERNEL=0 is the rollback seam.
121/// MoE grouped f16 GEMM door (experimental until gated), f16-mirror numeric class:
122/// per-layer expert dequant to f16 + one grouped f16 GEMM over the CSR groups.
123///   MEMRA_MOE_F16G=1  cublasGemmGroupedBatchedEx (round 46 arc 2). The grouped API issues
124///                     through cublas-internal streams NOT ordered with ours — v1 pays a full
125///                     stream sync per projection (round-47 ledgered defect).
126///   MEMRA_MOE_F16G=2  single-kernel grouped GEMM on the engine stream (round 49): ordered by
127///                     construction, zero syncs, f32 C with the act row-scale folded in.
128/// DEFAULT (2026-08-01, round 49 promotion): mode 1 on the Hopper lane — with the 41/41
129/// dequant coverage fix the q35 board-2048 prime measured 5490 (MMQ) / 8380 (mode 1,
130/// +53%) / 7990 (mode 2) x3 interleaved on the H100, argmax MATCH — the last board loss
131/// flips. The 5090 measured FLAT (858GB/s makes the dequant-workspace traffic cancel the
132/// GEMM win) — but that verdict is for expert banks the int8-MMA MMQ arm can take
133/// (IQ3_S/IQ4_XS/Q4_0). MEMRA_MOE_F16G=0 kills anywhere.
134///
135/// HOPPER RE-VERDICT (2026-08-02, lane/h100-flip-full): mode 2 with full direct coverage
136/// (Q4_K/Q6_K/IQ4_XS/IQ3_S tile loaders, lane/iq-direct-loaders) + the deep tail
137/// (lane/sk-tail-form) FLIPS past cublas mode 1 on the H100 — q35 board-2048 prime
138/// 13163.6 (mode 2, cross=32) vs 8626.5 (mode 1) vs 8073.4 (round-51 sk form), +52.6%,
139/// interleaved x5 zero overlap, argmax MATCH 30/30. The round-54 NO-FLIP (8547 vs 8112)
140/// was coverage-priced at 5.2% direct; ~100% coverage kills the workspace pass and the
141/// verdict inverts. Hopper naked default -> mode 2 (this arm); the gemma (gelu) site
142/// stays env-explicit-only via moe_f16g_gemma_on (Err => closed, unaffected by this arm).
143///
144/// MODE-2 DEFAULT (sm_120a naked, 2026-08-02, lane/f16g-default-rearb): with the direct
145/// tile loaders covering Q4_K/Q6_K/IQ4_XS/IQ3_S, the sk visitor beats the int8-MMA MMQ
146/// tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7% / pp2048
147/// +30.6% — research/iq-direct-loaders-20260802 §3-5, confirmed + full battery in
148/// research/f16g-default-rearb-20260802/), so every f16g-admitted expert layer rides
149/// mode 2 naked. Decode/verify stay on dp4a (t >= 16 floor). f16-mirror numeric class
150/// for naked q35/KAT prefill+prime — new token-sha anchors stamped in the rearb lane.
151///
152/// AUTO-KQUANT (mode 3, 2026-08-02, lane/q4k-expert-prefill): the previous sm_120a
153/// default, kept reachable via MEMRA_MOE_F16G=3. The mode-2 sk form is admitted ONLY for
154/// layers the MMA MMQ arm rejects (k-quant expert projections — Q3_K/Q4_K/Q6_K), i.e.
155/// exactly where the baseline is the per-pair moe_pairs_matvec_q8_em fallback with zero
156/// token reuse (Ornith-35B Q4_K_M board-2048 1098.2 -> 3453.7, 3.14x,
157/// research/q4k-expert-prefill-20260802/). Its "IQ banks keep their measured-faster MMQ
158/// tiles" ruling was priced BEFORE the IQ direct loaders and is refuted on the 5090 —
159/// the k-quant-only admission survives as the rollback seam, not the default.
160/// The gemma (gelu) site stays env-explicit-only (moe_f16g_gemma_on).
161pub fn moe_f16g_mode() -> u8 {
162    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
163    *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
164        Ok("0") => 0,
165        Ok("2") => 2,
166        Ok("3") => 3,
167        Ok(_) => 1,
168        // Both arches independently re-arbitrated to mode 2 on 2026-08-02
169        // (5090: lane/f16g-default-rearb; H100: lane/h100-flip-full) — unset = 2 everywhere.
170        Err(_) => 2,
171    })
172}
173/// Mode-2 sk kernel form policy (round 51, lane/sk-bm128): the single-kernel grouped GEMM runs
174/// as a persistent problem-visitor over the real CSR tiles with two tile forms. Returns
175/// (shape_sel, cross) for the FFI:
176///   MEMRA_F16G_SK=0    -> (-1, _): the round-49 grid-scan kernel (rollback seam).
177///   MEMRA_F16G_SK=32   -> all groups on the 32x64x32 2-stage form (cross = i32::MAX).
178///   MEMRA_F16G_SK=128  -> all groups on the 128x64x64 3-stage form (cross = 1; groups fall
179///                         back to 32x64 in-launcher when the device/in_f can't take it).
180///   unset              -> hybrid split: groups with m_e >= MEMRA_F16G_SK_CROSS ride the 128
181///                         form. Default cross = 64 (5090 sweep 2026-08-01, receipts
182///                         research/sk-bm128-20260801/; H100 re-swept on the direct+tail
183///                         form 2026-08-02, lane/h100-flip-full: {16,32,64} ->
184///                         12868/13192/13225 — 64 wins there too, the pre-direct 32
185///                         verdict was stale).
186pub fn moe_f16g_sk_params() -> (i32, i32) {
187    static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
188    *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
189        Ok("0") => (-1, 0),
190        Ok("32") => (0, i32::MAX),
191        Ok("128") => (0, 1),
192        _ => {
193            let cross = std::env::var("MEMRA_F16G_SK_CROSS")
194                .ok()
195                .and_then(|v| v.parse().ok())
196                .unwrap_or(64);
197            (0, cross)
198        }
199    })
200}
201/// DIRECT-FROM-QUANT sk tile loaders (lane/kquant-tile-loaders, 2026-08-02; IQ classes added
202/// by lane/iq-direct-loaders): Q4_K/Q6_K/IQ4_XS/IQ3_S expert projections on the mode-2/3 sk
203/// visitor forms dequant their weight tiles in-register from the quant superblocks instead of
204/// running the per-(layer,projection) dequant pass into an f16 workspace (41.8% of Ornith-35B
205/// t=512 kernel time — the pp512 wall, research/q4k-expert-prefill-20260802 §5; the IQ classes
206/// are 94.8% of q35's bank bytes — the h100-sk-direct coverage pricing). Bit-identical to the
207/// workspace path by construction (kernel-check "f16g-kq-direct" gates it bitwise) — a
208/// data-movement change, not a numeric-class change. Default ON; MEMRA_F16G_DIRECT=0 reverts
209/// to the workspace path everywhere; MEMRA_F16G_DIRECT=kq keeps the k-quant loaders and
210/// reverts only the IQ classes (the iq-direct-loaders A/B seam — the pre-lane shipped config).
211pub fn moe_f16g_direct_on(qtype: i32) -> bool {
212    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
213    let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
214        Ok("0") => 0,
215        Ok("kq") => 1,
216        _ => 2,
217    });
218    match m {
219        0 => false,
220        1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
221        _ => true,
222    }
223}
224/// DEEP-TAIL sk form (lane/sk-tail-form, 2026-08-02): groups below the visitor crossover ride
225/// a 32x64x64 3-STAGE cp.async tile instead of the round-51 32x64x32 2-stage — the same 32-row
226/// tile (zero extra padding), 2 k-blocks in flight instead of 1 and half the syncs per k. The
227/// H100 ncu pricing (research/sk-bm128-20260801) put the 2-stage tail at 31% of the sk GEMM
228/// stage under q35's routing skew. Bit-identical to every other sk form by construction
229/// (kernel-check "f16g-sk" gates all tail arms maxdiff==0); exists in both the workspace-f16
230/// and direct-from-quant variants. Default ON; MEMRA_F16G_TAIL=0 = rollback to the 2-stage
231/// tail. in_f % 64 != 0 falls back in-launcher.
232pub fn moe_f16g_tail_on() -> bool {
233    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
234    *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
235}
236
237/// Per-model door for the gemma-MoE (gelu) grouped path: round 49's Hopper default
238/// REGRESSED g26 board-2048 prefill -8.3% interleaved x5 on-box (def median 10380,
239/// wild 8.9k-11.7k spread; off 11317, ±0.13%) — the +6-15% probe verdict didn't
240/// survive the board workload (stale-verdict law, round 50). The silu/qwen class
241/// keeps the round-49 default (q35 +53% board-2048). Explicit MEMRA_MOE_F16G=1/2
242/// still opens this door for A/B.
243pub fn moe_f16g_gemma_on() -> bool {
244    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
245    *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
246}
247
248/// Fused act-epilogue (silu/gelu-mul + q8_1_mmq quantize in one launch) for the MoE prefill
249/// MMA arms. Byte-identical to the two-pass path (kernel-check gated) — default ON.
250/// MEMRA_MOE_FUSE_ACTQ=0 is the rollback/A-B seam.
251pub fn moe_fuse_actq_on() -> bool {
252    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
253    *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
254}
255
256/// Door `MEMRA_GLM5_Q8_FUSE` (lane/b200-q8-fuse-20260902, DEFAULT OFF pending the box A/B):
257/// on the glm5_next mHC decode trunk (`hyper_range_decode` / `hyper_range_decode_ws_body`),
258/// fold the FFN-input rms_norm and its consumer's standalone `quantize_q8_1` launch into
259/// ONE `rms_norm_zq8_f32` launch. Byte-identical to the unfused chain (see that kernel's
260/// header in cu/kernels.cu); this door only changes launch count. See docs/FLAGS.md and
261/// research/b200-q8-fuse-20260902/LANE.md.
262pub fn glm5_q8_fuse_on() -> bool {
263    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
264    *ON.get_or_init(|| std::env::var("MEMRA_GLM5_Q8_FUSE").as_deref() == Ok("1"))
265}
266
267/// `MEMRA_GLM5_Q8_FUSE_ATTN=1` (lane/glm5-attn-norm-zq8-20260904, default OFF): the
268/// ATTENTION-input norm of a plain KDA layer in the glm5_next T=1 walk runs `rms_norm_zq8_f32`
269/// and hands its q8_1 view to the fused six-projection launcher, which then skips its own
270/// `quantize_q8_1_into`. The FFN-input twin is `MEMRA_GLM5_Q8_FUSE`. Read PER CALL. Why and
271/// receipts: docs/FLAGS.md.
272pub fn glm5_q8_fuse_attn_on() -> bool {
273    std::env::var("MEMRA_GLM5_Q8_FUSE_ATTN").as_deref() == Ok("1")
274}
275
276/// Engagement counter for `MEMRA_GLM5_Q8_FUSE_ATTN`; gates take a delta.
277pub static GLM5_Q8_FUSE_ATTN_DISPATCHES: std::sync::atomic::AtomicU64 =
278    std::sync::atomic::AtomicU64::new(0);
279
280/// Door `MEMRA_GLM5_DECODE_GRAPH` (lane/b200-glm5-graph-20260902, DEFAULT ON since 2026-09-04):
281/// capture the glm5_next T=1 decode walk as replayable per-stage CUDA graphs instead of
282/// issuing every kernel per token. Unset/`1` arms it, `=0` is the eager walk. Read PER CALL so
283/// `=0` is a live rollback seam, never a process-lifetime latch.
284///
285/// WHY ON (receipts in docs/FLAGS.md and darklanes `research/glm5-b200-20260902/LANE.md`,
286/// cells graphab + graphgates 2026-09-04, 2x B200 SXM PP-2, assembled serving posture): the
287/// corrected door (memra#131 root-caused and fixed in #168) is +9.89% at c1 plain
288/// (66.40 -> 72.97 tok/s, interleaved x3, greedy tape identical to the eager walk on six
289/// boots), passes the 1M-context gate with every self-check green, is inert on the DFlash2
290/// spec route (engaged=0, wall unchanged), and is neutral on the 8-turn vendor-sampled
291/// twin with short turns (the per-session capture + warm + check cost eats the gain there:
292/// that cost is the next lever, not a reason to stay eager). Every refusal shape falls
293/// through to the eager walk byte-identically, and the first replay of every captured run
294/// is compared bitwise against an eager step (`MEMRA_GLM5_GRAPH_SELFCHECK_N`), latching the
295/// stage eager on any mismatch.
296///
297/// WHAT IS CAPTURED: the maximal CONTIGUOUS runs of KDA-mixer layers inside each pipeline
298/// stage's `[lo, hi)` range, hc glue and routed MoE included. WHAT STAYS EAGER: every
299/// MLA/DSA layer (its launch geometry is derived on the host from `layer.len`, see
300/// `HybridModel::mla_attn_cached_pre_wo`), the decode tail, prefill, and the spec verify
301/// walk (which keeps `MEMRA_SPEC_VERIFY_GRAPH`). See docs/FLAGS.md and
302/// research/b200-glm5-graph-20260902/LANE.md.
303pub fn glm5_decode_graph_on() -> bool {
304    glm5_decode_graph_on_from(std::env::var("MEMRA_GLM5_DECODE_GRAPH").ok().as_deref())
305}
306
307/// The pure parse behind [`glm5_decode_graph_on`]: only an explicit `0` disarms the door;
308/// unset, `1`, and any other value arm it. Kept separate so the default can be unit-tested
309/// without mutating the process environment (the OFF arm of every gate sets `=0` and this is
310/// the contract that makes that arm non-vacuous).
311pub fn glm5_decode_graph_on_from(v: Option<&str>) -> bool {
312    !matches!(v.map(str::trim), Some("0"))
313}
314
315/// `MEMRA_GLM5_GRAPH_HOST_MOE=1` — BISECT knob for `MEMRA_GLM5_DECODE_GRAPH`, gate harness only.
316/// The door has TWO enablers and box run 5 showed they fail independently: (1) the T=1
317/// device-table MoE arm that removes the per-layer router readback, and (2) the capture/replay
318/// itself. With this set the door stays ON but the MoE arm stands down to the host oracle, and
319/// the capture then refuses BY NAME (a host readback inside a capture region is illegal), so a
320/// run isolates enabler 2's absence from enabler 1's behaviour instead of confounding them.
321pub fn glm5_graph_host_moe() -> bool {
322    std::env::var("MEMRA_GLM5_GRAPH_HOST_MOE").as_deref() == Ok("1")
323}
324
325/// `MEMRA_GLM5_GRAPH_NO_CAPTURE=1` — THE OTHER HALF OF THE BISECT, and the half that was missing.
326///
327/// `MEMRA_GLM5_GRAPH_HOST_MOE=1` turns OFF the device-table MoE arm AND makes the capture refuse
328/// by name, so box run 6 compared "neither enabler" against "both enablers". That is not a
329/// bisect, and calling it one was wrong: it could never attribute the defect to one of the two.
330/// This knob supplies the missing cell — the device-table MoE arm ENGAGES exactly as it does in
331/// serving, and the capture never happens, so the whole walk runs eagerly.
332///
333/// The rig has since cleared the MoE arm end to end, including at serving scale (288 experts,
334/// `in_f` 4096, `expert_stride` 4718592) driven by the box's own routing dump, so the expected
335/// result is a CORRECT tape — which would pin the defect on the capture and exonerate the arm.
336/// A wrong tape here would instead mean the arm behaves differently in situ than in the fixture,
337/// and would say so on the first run rather than after another round of guessing.
338pub fn glm5_graph_no_capture() -> bool {
339    std::env::var("MEMRA_GLM5_GRAPH_NO_CAPTURE").as_deref() == Ok("1")
340}
341
342/// `MEMRA_GLM5_VROWS_T1_DEV=1` — the T=1 device-table MoE arm, forced ON with no capture, no
343/// graph, and no `MEMRA_GLM5_DECODE_GRAPH` anywhere in the run. Default OFF, gate harness only.
344///
345/// From 2026-09-03 the arm is keyed on an OPEN CAPTURE REGION rather than on the decode-graph
346/// door, because that is the only place it is required (a host sel/w readback cannot live inside
347/// a capture). That keying is what makes the door's eager fall-through byte-identical again, and
348/// it also means the arm can no longer be observed on a plain decode run — so this knob puts it
349/// back within reach of a bisect. It is the cell box takes 4 through 11 never had: they set one
350/// env that turned on BOTH the arm and the capture, so a wrong tape could not be attributed to
351/// either. Run this alone and the answer is unambiguous.
352/// `MEMRA_GLM5_GRAPH_RECAPTURE=1` (default OFF): when a captured stage is invalidated (its
353/// `cache.pos` moved, or a recurrent-state buffer was re-seated rather than overwritten), REBUILD
354/// it instead of latching that stage to the eager walk for the rest of the session.
355///
356/// Default OFF is a decision, not an omission. Box run 3 (2026-09-02) died in the teardown with
357/// `CUDA_ERROR_INVALID_VALUE`: it destroyed a stage's execs, and freed every buffer they baked,
358/// with a replay of those same execs still outstanding. That is a destroy-in-use, and the fix is
359/// ordering (drain the stream, THEN drop, and refuse to drop at all if the drain fails), which is
360/// what the armed path now does. But the latch is not a workaround, it is a correct product
361/// behaviour on its own: an invalidated stage falls through to the byte-identical eager walk, so
362/// the only cost of NOT rebuilding is that one stage of one session stops being graphed. Nothing
363/// is wrong, only slower.
364///
365/// So the door ships with the latch and this knob exists to take the rebuild's receipt. It is
366/// what the gate's forced-re-seat arm needs in order to exercise the invalidation path at all:
367/// with the knob off that arm was asserting on a path the engine deliberately does not take, and
368/// box take 13 duly reported `VACUOUS RE-CAPTURE ARM` on a run whose tokens were all correct.
369/// Counter: `GLM5_DECODE_GRAPH_RECAPTURES`.
370pub fn glm5_graph_recapture_on() -> bool {
371    std::env::var("MEMRA_GLM5_GRAPH_RECAPTURE").as_deref() == Ok("1")
372}
373
374pub fn glm5_vrows_t1_dev_forced() -> bool {
375    glm5_vrows_t1_dev_forced_from(
376        std::env::var("MEMRA_GLM5_VROWS_T1_DEV").ok().as_deref(),
377        env!("MEMRA_BUILT_CUDA_ARCH"),
378        glm5_graph_no_capture() && glm5_decode_graph_on(),
379    )
380}
381
382/// The pure parse behind [`glm5_vrows_t1_dev_forced`] (arch-keyed since 2026-09-04): `1` forces
383/// the T=1 device-table MoE arm on every eager layer, `0` leaves the eager layers on the host
384/// readback (the arm still engages inside a capture, which is keyed on the open region, not on
385/// this), and UNSET forces it on `100a` builds. Receipt (darklanes
386/// research/glm5-b200-20260902/LANE.md, t1devab 2026-09-04, 2x B200 pair, composed defaults):
387/// host 79.60/79.81/79.31 -> 79.60 vs forced 84.11/84.19/83.29 -> 84.11, +5.67%, tape
388/// 9437b599f6b9d2a9 on all six boots; the 11 eager MLA-layer MoE calls each lose a pinned
389/// readback plus a device drain. `no_capture_bisect` is the `MEMRA_GLM5_GRAPH_NO_CAPTURE` knob's
390/// own forcing, unchanged.
391pub fn glm5_vrows_t1_dev_forced_from(
392    v: Option<&str>,
393    built_arch: &str,
394    no_capture_bisect: bool,
395) -> bool {
396    match v.map(str::trim) {
397        Some("1") => true,
398        Some("0") => no_capture_bisect,
399        _ => built_arch == "100a" || no_capture_bisect,
400    }
401}
402
403#[cfg(test)]
404mod glm5_vrows_t1_dev_default_tests {
405    use super::glm5_vrows_t1_dev_forced_from;
406
407    #[test]
408    fn arch_keyed_default_with_override_and_bisect_knob() {
409        assert!(glm5_vrows_t1_dev_forced_from(None, "100a", false));
410        assert!(!glm5_vrows_t1_dev_forced_from(None, "120a", false));
411        assert!(glm5_vrows_t1_dev_forced_from(None, "120a", true));
412        assert!(glm5_vrows_t1_dev_forced_from(Some("1"), "120a", false));
413        assert!(!glm5_vrows_t1_dev_forced_from(Some("0"), "100a", false));
414        assert!(glm5_vrows_t1_dev_forced_from(Some("0"), "100a", true));
415    }
416}
417
418/// `MEMRA_GLM5_GRAPH_TRACE=1` — GATE-HARNESS trace for `MEMRA_GLM5_DECODE_GRAPH`, never a
419/// serving flag. Prints one line per captured-run boundary per token, on BOTH arms and at the
420/// SAME layer boundaries, with a checksum of the stream state leaving that segment. Box run 4
421/// produced token 0 at every step with the door running cleanly and no error anywhere: the only
422/// way to tell "the captured range wrote nothing the remainder reads" from "the state is wrong
423/// from layer N onward" is to compare the two arms segment by segment, and `nz=` in the line
424/// separates an all-zero hidden from a wrong-but-live one on sight.
425pub fn glm5_graph_trace_on() -> bool {
426    std::env::var("MEMRA_GLM5_GRAPH_TRACE").as_deref() == Ok("1")
427}
428
429/// How many times the T=1 device-table MoE arm has dumped its input shape under
430/// `MEMRA_GLM5_GRAPH_TRACE`. Capped at two layers: the question is what the arm is HANDED on the
431/// real artifact, and two routed layers answer it without turning a 64-step run into a log flood.
432/// Trace-dump budget for the decode-graph door's MoE seam (`MEMRA_GLM5_GRAPH_TRACE`), keyed by
433/// `(kind, arm, layer)`.
434///
435/// It was a pair of process-global counters capped at 4, and take 10 showed exactly why that is
436/// the wrong shape: the gate runs its EAGER arm first, that arm spent the whole budget on one
437/// layer, and the run printed four identical `arm=host il=3` lines and NOT ONE `arm=device`
438/// line — the arm the run existed to observe. A budget for a two-arm comparison has to be per
439/// arm, and a per-layer dump has to be keyed by layer or it reprints the first one.
440///
441/// [`glm5_trace_reset`] clears it at every arm switch so the second arm starts with a full
442/// budget rather than inheriting the first arm's exhaustion.
443/// `(kind, arm, layer)` — one dump slot. Named so the map type stays readable at both use sites.
444type Glm5TraceKey = (&'static str, String, u16);
445
446fn glm5_trace_slots() -> &'static Mutex<std::collections::BTreeSet<Glm5TraceKey>> {
447    static S: std::sync::OnceLock<Mutex<std::collections::BTreeSet<Glm5TraceKey>>> =
448        std::sync::OnceLock::new();
449    S.get_or_init(Default::default)
450}
451
452/// Claim the one dump slot for `(kind, arm, il)`. Returns false once that exact line has printed,
453/// and false past `GLM5_TRACE_MAX_LAYERS` distinct layers for this `(kind, arm)` so a 64-step run
454/// cannot become a log flood.
455pub(crate) fn glm5_trace_take_slot(kind: &'static str, arm: &str, il: u16) -> bool {
456    const GLM5_TRACE_MAX_LAYERS: usize = 8;
457    let mut s = glm5_trace_slots().lock().unwrap();
458    if s.iter().filter(|(k, a, _)| *k == kind && a == arm).count() >= GLM5_TRACE_MAX_LAYERS {
459        return false;
460    }
461    s.insert((kind, arm.to_string(), il))
462}
463
464/// Clear the trace budget. The gate calls this at every arm switch; without it the first arm's
465/// exhaustion silences the second (take 10).
466pub fn glm5_trace_reset() {
467    glm5_trace_slots().lock().unwrap().clear();
468}
469
470/// Captured-run replays (one per graph launch), captures, and the layer count currently
471/// covered by captured runs — the door's engagement receipt, read by the gate bin.
472pub static GLM5_DECODE_GRAPH_REPLAYS: std::sync::atomic::AtomicU64 =
473    std::sync::atomic::AtomicU64::new(0);
474pub static GLM5_DECODE_GRAPH_CAPTURES: std::sync::atomic::AtomicU64 =
475    std::sync::atomic::AtomicU64::new(0);
476pub static GLM5_DECODE_GRAPH_LAYERS: std::sync::atomic::AtomicU64 =
477    std::sync::atomic::AtomicU64::new(0);
478/// Stages torn down and rebuilt by the armed re-capture path (`MEMRA_GLM5_GRAPH_RECAPTURE`).
479pub static GLM5_DECODE_GRAPH_RECAPTURES: std::sync::atomic::AtomicU64 =
480    std::sync::atomic::AtomicU64::new(0);
481
482/// True while a glm5 decode-graph CAPTURE is open on this process. Two engine pools must not
483/// hand a captured graph a buffer they will later re-issue to eager work: a replay would
484/// scribble whatever landed there (the draft-graph root cause, see `capture_graph_retained`).
485/// While this is set, `vws_recycle*` drops instead of returning the buffer to the verify
486/// workspace, so every transient the captured body took stays owned by the graph.
487pub(crate) static GLM5_GRAPH_CAPTURE_OPEN: std::sync::atomic::AtomicBool =
488    std::sync::atomic::AtomicBool::new(false);
489
490pub(crate) fn glm5_graph_capture_open() -> bool {
491    GLM5_GRAPH_CAPTURE_OPEN.load(std::sync::atomic::Ordering::Relaxed)
492}
493
494/// PREFILL router m-invariance (lane/concat-prime-exact, 2026-08-02). The batched cuBLASLt
495/// router GEMM changes a row's logits when OTHER rows join the call (probed: first change at
496/// m=65 on the Ornith-35B router, 3.9e-3 — while the MMQ/f16 trunk GEMMs are bit-identical
497/// across m). Feeding a top-k discontinuity, that made a served request's expert selection a
498/// function of its CO-ARRIVALS under cross-request prime batching. The in-house router GEMV
499/// is m-invariant, so prefill uses it too and routing depends on a session's own tokens only.
500/// DEFAULT ON: it is the serving isolation contract, and it is the same kernel decode and spec
501/// verify already use (dispatch parity, one router kernel for every t).
502/// MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched GEMM.
503pub fn router_prefill_exact_on() -> bool {
504    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
505    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
506}
507
508pub fn router_kernel_on() -> bool {
509    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
510    *ON.get_or_init(|| {
511        let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
512        if !on {
513            eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
514        }
515        on
516    })
517}
518
519/// FAST-ROUTER batch twin (lane/fast-router, 2026-08-02). The concat-prime exactness fix
520/// (router_prefill_exact_on) routes prefill through router_gemv — m-invariant, but a
521/// per-(expert,token) GEMV program with zero operand reuse, so q35 board-2048 prefill paid
522/// -10% on the 5090. router_gemv_f32_w8_batch register-tiles (8x8 expert-x-token) the same
523/// per-row FP chains (BIT-IDENTICAL per row — kernel-check sweeps m=1..2048 on real router
524/// weights), so the t crossover below is pure perf, not a numeric config. Swept on-box
525/// (research/fast-router-20260802/crossover-router*.jsonl): plain wins t<=4, batch +7-9%
526/// at t=8, 1.9x at t=16 rising to 3.45x at t=2048 — MIN_T=8. Decode t=1 and spec verify
527/// t<8 keep the plain w8 form. MEMRA_ROUTER_BATCH=0 forces plain at every t (rollback
528/// seam, perf-only: bits are equal by the kernel-check gate).
529/// Killed arms (same sweep, JSONL is the record): the 8x16 tile lost to 8x8 at every t
530/// (128-accumulator register pressure beats the halved w-traffic), and the same-shape
531/// sigmoid_dot_rows twin (out_f=1) measured 0.62-0.89x at every prefill t
532/// (launch-latency-bound, ~7us/layer at m=2048) — both bit-identity-PASSED before dying.
533pub const ROUTER_BATCH_MIN_T: usize = 8;
534pub fn router_batch_on() -> bool {
535    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
536    *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
537}
538mod cpu_experts;
539#[cfg(memra_cutlass)]
540pub mod cutlass_ffi;
541pub mod dsv4_ffi;
542pub mod dsv4_gpu;
543pub mod f16_ffi;
544pub mod fp8_ffi;
545pub mod mmq_ffi;
546pub mod moe_cache;
547pub mod prime_graph;
548pub mod spill;
549mod spill_pread;
550
551// Fatbins are EMBEDDED (crates-release lane, 2026-08-04): build.rs still writes them to
552// OUT_DIR, but the bytes ship inside the binary via include_bytes! and load through
553// cuModuleLoadData. Distribution contract: a prebuilt or cargo-installed binary must be
554// self-contained — the old baked OUT_DIR *paths* pointed at the builder's temp dir and
555// broke every machine that wasn't the build machine. Same bytes, same module image;
556// the runtime MEMRA_GEMM_FATBIN tune-seam override below is preserved.
557const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
558const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
559/// kda.cu: the glm5_next Kimi Delta Attention mixer (per-channel-decay delta rule).
560const KDA_FATBIN: &[u8] = include_bytes!(env!("MEMRA_KDA_FATBIN"));
561const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
562const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
563const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
564const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
565/// spec_sample.cu: sampled-spec primitives (Philox Gumbel-max / softmax gather / residual sampler).
566const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
567
568/// TUNE SEAM (tools/sweep): a RUNTIME `MEMRA_GEMM_FATBIN=<path>` overrides the baked-in
569/// qmatvec_gemm.cu fatbin path (build.rs bakes the same name at COMPILE time via
570/// cargo:rustc-env — that constant is the default). Lets the sweep harness swap in a
571/// `-D`-tuned fatbin per process with NO rust rebuild. Unset at runtime => the
572/// compile-time default (zero behavior change).
573fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
574    assert!(
575        !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
576        "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
577    );
578    match std::env::var("MEMRA_GEMM_FATBIN") {
579        Ok(path) => std::borrow::Cow::Owned(
580            std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
581        ),
582        Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
583    }
584}
585
586/// Phase A (ARCHITECTURE-H100.md): sm_90a re-enables the portable-PTX tensor-core paths
587/// (int8 mma.m16n8k32/k16.s8, bf16 m16n8k16, ldmatrix, cp.async — all sm_80-class, native
588/// on Hopper) that the portable boot lane gates off. Dispatch guards that used to test
589/// `cfg!(memra_portable_cuda)` test this instead; sm_89 keeps the pure-portable behavior.
590/// The sm_120a/sm_100a-only MMA kinds (mxf4nvf4, kind::f8f6f4) are NOT covered — their
591/// launchers stay fail-closed stubs on 90a and their dispatch arms stay arch-gated.
592pub(crate) const fn portable_mma_gated() -> bool {
593    cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
594}
595
596/// Refuse an env force that would reach a kernel THIS BUILD DOES NOT CONTAIN.
597///
598/// Doors of the shape `MEMRA_X=1 => true` are arch-blind: they were written so an operator could
599/// force a promoted path on, and the default arm (`cfg!(memra_hopper_mma)` or similar) is the only
600/// thing that consulted the arch. On a portable build the forced path then reaches
601/// `Engine::func`, which resolves lazily and ends in `panic!("kernel {name} not in any fatbin")` —
602/// a confusing crash naming a kernel the operator never heard of, several frames from the switch
603/// they actually flipped.
604///
605/// Found 2026-08-23 by tools/fatbin-lookup-census.py, which listed 20 looked-up kernels absent
606/// from the sm_89 fatbins. 18 of those turned out to be correctly unreachable (the GDN varlen
607/// chain is gated through `gdn_mma_enabled`, which starts with `!portable_mma_gated()`); these
608/// env doors were the two that were genuinely reachable, and only by explicit operator action.
609///
610/// Same shape and same message style as `gemm_fatbin_bytes`'s refusal above — one idiom for
611/// "this switch cannot work on this build", so it fails at the switch instead of at the lookup.
612#[track_caller]
613pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
614    assert!(
615        !portable_mma_gated(),
616        "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
617         portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
618    );
619}
620
621/// The GDN K4/K5 mma pair's UNSET-env default — ONE definition for the three read sites
622/// (gdn_mma_enabled, the k123 pre-work, gdn_scan_chunked's dispatch). They read the env
623/// per call ON PURPOSE (kernel-check toggles it to pin both configs), so the shared part
624/// is this compile-time constant: ON for Hopper-MMA builds (the original 90a promotion)
625/// and for sm_120a builds (lane/moeprime-nvfp4-direct, 2026-08-21 — measured on one RTX
626/// PRO 6000 ornith15 pp14715 +6-8% and the local 5090 q38-27b +1-2%, both orders both
627/// rigs). A site defaulting differently from its peers arms the mma pre-work while the
628/// scan takes the scalar route — measured as a 0.8% LOSS, the drift this helper kills.
629pub(crate) const fn gdn_mma_default_on() -> bool {
630    cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
631}
632
633/// const str-eq (std `==` on &str is not const-stable on this toolchain floor).
634const fn konst_eq(a: &str, b: &str) -> bool {
635    let (a, b) = (a.as_bytes(), b.as_bytes());
636    if a.len() != b.len() {
637        return false;
638    }
639    let mut i = 0;
640    while i < a.len() {
641        if a[i] != b[i] {
642            return false;
643        }
644        i += 1;
645    }
646    true
647}
648
649/// The legacy quantized prefill GEMMs are tuned and validated for sm_120a; sm_90a re-admits
650/// them through the Hopper-MMA lane (int8 m16n8k32.s8 is sm_80-class PTX).  Keep the policy
651/// in a pure helper so the dispatch guard can be regression-tested without constructing an
652/// Engine or allocating a GPU tensor.
653const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
654    (!portable_cuda || hopper_mma) && !no_gemm
655}
656
657// ---- KV-cache format selection (kvbytes lane, 2026-07-08; default OFF = daily config) ----
658// `MEMRA_KV_K` = q8_0 (default, 34 B/32elem) | fp8 (raw e4m3, 32 B — the -6% K-bytes arm)
659// `MEMRA_KV_V` = q5_1 (default, 24 B/32elem) | q4_0 (18 B, -25% V bytes) | fp8 (32 B, +33%)
660// A non-default format is a NEW NUMERIC CONFIG: its own run-gen argmax baseline is legal,
661// but the gate battery (kernel-check, run-spec self-consistency) must pass WITHIN it and
662// the choice is explicit env, never silent. flash_attn.cu is compiled once per format pair
663// (build.rs); the kernels keep their names — Engine::new just loads the matching fatbin.
664const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
665const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
666const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
667const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
668const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
669
670/// KV format policy moved to the shared `memra-kv` crate (Phase D); re-exported so the
671/// fatbin router below and every existing `crate::kv_blk_bytes()` call site is unchanged.
672pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
673
674/// The flash_attn fatbin matching the selected KV formats.
675fn flash_fatbin_bytes() -> &'static [u8] {
676    match kv_cache_formats() {
677        ("q8_0", "q5_1") => FLASH_FATBIN,
678        ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
679        ("q8_0", "fp8") => FLASH_FATBIN_VF8,
680        ("fp8", "q5_1") => FLASH_FATBIN_KF8,
681        ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
682        ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
683        other => unreachable!("kv_cache_formats returned {other:?}"),
684    }
685}
686
687/// TUNE SEAM (tools/sweep): kernel1 (Q8_0/Q4_K/Q5_K) launch-tile override,
688/// `MEMRA_GEMM_K1_LAUNCH="BM,BN,NWARP"`. MUST match the `-D K1_BM/K1_BN/NWARP` the swept
689/// fatbin was compiled with (the .cu tile and the host launch grid/block have to agree —
690/// the hardcoded (128,128,8) in qmatvec_gemm/qmatvec_gemm_raw is the shipped default).
691/// Kernel2 (Q6_K/NVFP4) launch is untouched. Unset or malformed => None => shipped
692/// defaults (zero behavior change).
693fn k1_launch_override() -> Option<(u32, u32, u32)> {
694    static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
695    *K1.get_or_init(|| {
696        let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
697        let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
698        match p.as_slice() {
699            [bm, bn, w] => Some((*bm, *bn, *w)),
700            _ => None,
701        }
702    })
703}
704
705/// H100 wgmma prefill-GEMM seam (task 8, ARCHITECTURE-H100.md): OPT-IN (MEMRA_WGMMA=1).
706/// v0 verdict (2026-07-26, N=5 pp512 9B-Q8_0): wgmma 3845 tok/s vs MMQ 8692 — the
707/// standalone harness's "688us MMQ ref" was a pp2048-shape figure, so v0 (unpipelined,
708/// 64x64 tile, wait_group<0> every 32-K step) is ~3x SLOWER per launch at m=512 model
709/// shapes. Default stays MMQ until the pipelined version beats it N=5 (repo law).
710/// Correctness stays pinned regardless: kernel-check's wgmma case is cfg-gated, not env-gated.
711pub(crate) fn wgmma_gemm_enabled() -> bool {
712    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
713    *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
714}
715
716/// TUNE SEAM: keys per FA-decode split (`MEMRA_FA_SPLIT` forces a fixed size; default 64). Smaller
717/// splits raise grid.y so grid = n_head_kv * n_splits fills the 82 SMs at short/mid ctx (vec path
718/// launches only n_head_kv=8 CTAs per split). Swept clock-locked 2026-07-03 (graph tg128): 32 beat
719/// 64 at ctx 128/512 (+0.5/+1.2%) and lost at 2048 (-3%) — BUT the adaptive 32/64 default BROKE the
720/// MTP spec-decode exact-match gate (run-spec K=1/2 self-consistency FAIL with 32; PASS with 64):
721/// the split count changes the combine's FP summation order, and the spec verify's batched forward
722/// only argmax-matches single-step decode under the 64-split order on real prompts. Spec exactness
723/// (the bigger lever) outranks a <=1.2% decode win -> default stays FIXED 64; sweeps use the env.
724/// Takes t_kv so eager, _dc capture, and fa_geom_eager stay signature-compatible for future
725/// adaptive retries (any retry MUST pass run-spec self-consistency first).
726/// Minimum t_kv for the warp-per-token vec FA path (below it the scalar path's 4x-more-blocks
727/// hides latency better — measured crossover, see `fa_decode`). Shared by fa_decode / fa_decode_dc /
728/// fa_geom_eager / fa_decode_rows-eligibility (spec verify) so the kernel pick NEVER diverges
729/// between eager decode and the verify (the spec-exactness law).
730pub const FA_VEC_MIN_TKV: usize = 96;
731/// Env-overridable crossover (MEMRA_FA_VEC_MIN, default FA_VEC_MIN_TKV). The 96 floor was
732/// measured on the qwen geometry (nkv=2); gemma4 SWA layers run nkv=8 = 4x the vec grid,
733/// which moves the crossover — sweep per model, adopt per the battery.
734pub fn fa_vec_min_tkv() -> usize {
735    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
736    *V.get_or_init(|| {
737        std::env::var("MEMRA_FA_VEC_MIN")
738            .ok()
739            .and_then(|v| v.parse().ok())
740            .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
741    })
742}
743
744/// f16-P/V class (DEFAULT since 2026-07-23 stamp v4; MEMRA_FA_F16PV=0 = f32-class rollback):
745/// llama-fa=1-style f16 P + f16 P@V accumulation on the hd512/SWA prefill stamps
746/// (KQ/softmax/normalize stay f32). Laptop stamp: 12B 1.045x, 31B 0.979x vs llama.
747///
748/// SPEC-SERVING FLIP (2026-07-26, the wkv acceptance-law pattern): with MEMRA_DRAFT set the
749/// default is OFF. f16 P/V shifts the PRIME's hidden states/KV in the sub-argmax logit
750/// space the drafter feeds on — argmax gates stay MATCH while depth acceptance falls off a
751/// cliff (26B d1736 0.883 -> 0.405, -40% e2e; f16pv-off alone restores 0.846/314 tok/s —
752/// the perf-ci acceptance battery is the only gate that sees this class). Explicit
753/// MEMRA_FA_F16PV always wins; plain serving keeps the f16 prefill win.
754pub fn fa_f16pv_on() -> bool {
755    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
756    *ON.get_or_init(|| {
757        std::env::var("MEMRA_FA_F16PV")
758            .map(|v| v != "0")
759            .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
760    })
761}
762
763/// hd512 head-pair arm (DEFAULT since stamp v4; MEMRA_FA512_HP=0 reverts to sp16): GQA
764/// ncols2=2 — 2 heads per CTA share each staged K/V tile, Q register-resident. Engages
765/// when n_head is even and the GQA group (n_head/n_head_kv) is even.
766pub fn fa512_hp_on() -> bool {
767    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
768    *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
769}
770
771/// SWA head-pair arm (DEFAULT since stamp v4; MEMRA_FAW_HP=0 reverts to p1): llama-class
772/// windowed geometry — 32 q-rows x 2 heads per CTA sharing staged K/V, f16 P@V
773/// accumulation. Even n_head and even GQA group required (guarded per call).
774pub fn faw_hp_on() -> bool {
775    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
776    *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
777}
778
779/// 4-warp sp16 experiment arm (MEMRA_FA512_W4=1, requires the f16pv door): GEMM0 split-K
780/// 4-way + GEMM1 4x128 O-dims. Own partial-sum order — oracle-band gated. Returns warp
781/// count (2 = base sp16). 8-warp arm measured NEGATIVE 2026-07-23 (jsonl) and removed.
782pub fn fa512_wide_warps() -> usize {
783    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
784    *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
785        Ok("1") => 4,
786        _ => 2,
787    })
788}
789
790/// hd-512 vec crossover floor (MEMRA_FA512_MIN, default 512) — shared by fa_decode dispatch
791/// and the gemma global-layer rows/parity call sites.
792pub fn fa512_min_tkv() -> usize {
793    static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
794    *FA512_MIN.get_or_init(|| {
795        std::env::var("MEMRA_FA512_MIN")
796            .ok()
797            .and_then(|v| v.parse().ok())
798            .unwrap_or(512)
799    })
800}
801/// Per-model crossover default, set at model load BEFORE the first decode (per-model
802/// numeric-config adoption law). qwen keeps the measured 96; gemma4 (nkv=8 SWA) measured
803/// vec-always fastest: 119.9 (96) / 130.0 (48) / 133.2 (1) tok/s tg128-regime, 2026-07-10.
804pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
805    std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
806/// Per-model windowed-split default (MEMRA_FA_SPW overrides): gemma MoE (26B, nkv=8) measured
807/// 32 (grid-limited t=1 under the raw-e4m3 sV ceiling, 2026-07-12); dense gemma (31B)
808/// measured 64 (37.13/37.12 vs 36.87/36.86 at 1.7k, N=2 — different attention geometry).
809pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
810/// Per-model hd512 (gemma globals) split default (MEMRA_FA_SP512 overrides): 26B measured 16
811/// (2026-07-11 N=2), dense 31B measured 32 (36.86/36.93 vs 36.73/36.73 at 1.7k, 2026-07-12).
812/// fused t=1 q4_0 pair/triple row mapping: true = mr1 (one row/warp). Per-model default
813/// (dense gemma wins +1.1% short / +0.6% depth on the 31B; MoE 26B REGRESSES −1.2% —
814/// its shared-expert fused2 shapes lose to the finer grid). MEMRA_Q40_MR env still wins.
815pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
816    std::sync::atomic::AtomicBool::new(false);
817/// Per-model router-GEMV form (2026-07-31): the 8-warp twin is +8.8% on the H100 q35
818/// decode step (router was 14.8% of it) with argmax + spec self-consistency green on
819/// qwen-class MoE both rigs. The gemma-4 26B knife-edge block (2026-07-31, single
820/// synthetic prompt) was RE-ARBITRATED 2026-08-01 on 6 real prompts — gate outcomes
821/// identical to the lone-warp arm, +13% g26 decode — so gemma4 rides the default too
822/// (research/g26-decode-20260801/). MEMRA_ROUTER_V2 env overrides either way.
823pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
824    std::sync::atomic::AtomicBool::new(true);
825pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
826    std::sync::atomic::AtomicUsize::new(16);
827/// Per-model rms_norm block size (per-model numeric-config law: the per-thread partial-sum
828/// split changes with blockDim -> different FP order -> battery-arbitrated per model).
829/// qwen keeps the shipped 256; gemma4 adopts 1024 (single-row 2816-col norms are one-block
830/// latency-bound at 256 threads — 7us/launch measured).
831pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
832/// gemma4 fa split ladder switch (set at model load; see fa_split_keys).
833pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
834/// Per-model stream-k entry override for SPEC serving (-1 = unset → env/default;
835/// 0 = force tiling; 1 = admit the deterministic form selector). The former timing
836/// selector made identical boots choose different fold orders; `MEMRA_MMQ_SK_FORM` is the
837/// explicit numerical-form seam. mmq_ffi reads this before the env.
838pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
839/// Per-model FP8-KV door — lives in memra-kv next to the format policy it drives
840/// (re-export keeps `crate::KV_FP8_FORCE` setters in model.rs/hybrid.rs working).
841pub use memra_kv::KV_FP8_FORCE;
842/// bf16 matvec family block size (MEMRA_MMV_BLOCK, default 128, clamped to [64, 256] and a
843/// multiple of 32 — the f32acc twin's shared reduce caps at 256). NUMERIC-CLASS knob: the
844/// per-thread stride and reduction order change with the block, same acceptance class as
845/// MEMRA_RMS_BLOCK (fresh-tape identity + battery at the pinned value).
846pub(crate) fn mmv_block() -> u32 {
847    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
848    *V.get_or_init(|| {
849        std::env::var("MEMRA_MMV_BLOCK")
850            .ok()
851            .and_then(|v| v.parse().ok())
852            .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
853            .unwrap_or(128)
854    })
855}
856
857/// MEMRA_B200_MATVEC_ARM=1: the sm_100a occupancy arms for the plain-decode MoE/matvec family
858/// (lane/b200-matvec-occupancy-20260902, docs/FLAGS.md). The B200 census (2026-09-02,
859/// GLM-5.3-Flash NVFP4, PP2 decode) found `moe_gate_up_preclamp8_q8` / `moe_down8_fma_q8` /
860/// `matvec_bf16_f32acc_x4_rows` running ~3-9x their roofline byte estimate on 2x B200 — an
861/// occupancy/latency signature from kernels tuned for the RTX PRO 6000's 188-SM/1.8-TB/s shape,
862/// not the B200's 148-SM/8-TB/s one. Restricted to `sm_100a` BUILDS (`MEMRA_BUILT_CUDA_ARCH`,
863/// baked in at compile time): setting the var on an `sm_120a` build is a documented no-op, so
864/// the naked sm_120a defaults stay byte-identical (per-hardware arm selection law, CLAUDE.md).
865/// Default OFF everywhere; the arms are BIT-IDENTICAL per-output twins pending their B200 A/B —
866/// see docs/FLAGS.md and research/b200-matvec-occupancy-20260902/LANE.md.
867pub(crate) fn b200_matvec_arm_on() -> bool {
868    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
869    *V.get_or_init(|| {
870        env!("MEMRA_BUILT_CUDA_ARCH") == "100a"
871            && std::env::var("MEMRA_B200_MATVEC_ARM").as_deref() == Ok("1")
872    })
873}
874
875/// MEMRA_B200_GEMV_V2=1: the sm_100a HBM-speed rewrite of the t=1 decode matvec class
876/// (lane/b200-gemv-hbm-20260902, docs/FLAGS.md, research/b200-gemv-hbm-20260902/LANE.md).
877///
878/// WHY A REWRITE AND NOT ANOTHER OCCUPANCY ARM. The B200 census has these kernels at 11-34% of
879/// the 8 TB/s HBM3e wall with every existing door ON, and the previous lane's warp-packing and
880/// prefetch arms bought ~5% — so the residual is not block-slot occupancy. It is BYTES IN
881/// FLIGHT PER SM: Little's law at 8 TB/s and a ~700 ns HBM3e round trip wants ~5.6 MB of reads
882/// outstanding across the die (~38 KB per SM) at all times, and the shipped kernels hold one
883/// weight load per thread per K step behind a serially dependent fma chain.
884///
885/// The v2 family is the SAME arithmetic, rescheduled: 8 rows per block accumulated
886/// CONCURRENTLY with the activation loaded once and reused across them, a two-stage software
887/// pipeline that issues 10 independent 16 B `ld.global.nc` loads before the first fma consumes
888/// one, the 8 rows' reductions run in lockstep so a block pays ONE barrier chain instead of
889/// four, a warp-shuffle tail, `__launch_bounds__`, and grids that cover the die (the down
890/// projection goes from `out_f` warps wide to `out_f * n_used`).
891///
892/// EVERY DISPATCHED ARM IS BIT-IDENTICAL to its shipped twin per output element. The one
893/// exception is the split-K arm (`matvec_bf16_v2_sk` + its fixed-order combine), a NAMED
894/// numeric class `bf16_gemv_v2_splitk` that only engages when a shape's row grid cannot cover
895/// two waves of CTAs on this die; the shipped GLM-5.3 decode shapes never reach it.
896///
897/// Restricted to `sm_100a` BUILDS (`MEMRA_BUILT_CUDA_ARCH`, baked in at compile time), like
898/// `MEMRA_B200_MATVEC_ARM`: on an sm_120a build the var is a documented no-op and the naked
899/// sm_120a defaults stay byte-identical, per the per-hardware arm selection law. Default OFF
900/// everywhere pending its B200 A/B.
901/// `MEMRA_B200_GEMV_V2` as a LEVEL, not a boolean: `0`/unset off, `1` = the v2 family,
902/// `2` = v2 plus the cp.async-staged v3 bf16 walk wherever it fits (v3 falls back to v2 per
903/// call when the shape's dynamic shared memory would exceed the 48 KB default cap or when the
904/// shape wants split-K). Any other value is off, deliberately: a typo must not silently arm a
905/// kernel arm. Same `sm_100a`-BUILD restriction as before.
906pub(crate) fn b200_gemv_v2_level() -> u8 {
907    static V: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
908    *V.get_or_init(|| {
909        if env!("MEMRA_BUILT_CUDA_ARCH") != "100a" {
910            return 0;
911        }
912        match std::env::var("MEMRA_B200_GEMV_V2").as_deref() {
913            Ok("1") => 1,
914            Ok("2") => 2,
915            _ => 0,
916        }
917    })
918}
919
920pub(crate) fn b200_gemv_v2_on() -> bool {
921    b200_gemv_v2_level() >= 1
922}
923
924/// `MEMRA_E4M3_ROW_ILP` (lane/glm5-b200-mint-consume-20260904) — loads-in-flight twin of the
925/// per-tensor e4m3 row walk.
926///
927/// DEFAULT **OFF**, deliberately and with the reason stated (new-flags law). The change is pure
928/// scheduling: `e4m3_row_dot_ilp<4>` issues four blocks' weight and activation loads before any of
929/// the four dependent fma chains consumes one, and folds the four block sums into `acc` in the
930/// SAME ascending order the serial walk uses — bit-identical per row, not a new numeric class.
931/// It ships OFF because it has no ON-part receipt yet: the B200 pair that would price it is still
932/// staging the mint, and this lane does not default a door it has not measured. The rig 5090 can
933/// prove the identity and never the speed (rig exactness-only law).
934///
935/// WHY IT IS EXPECTED TO PAY. This lane's ncu census found every big matvec long-scoreboard-bound
936/// at 71-76%, and the same lever already paid on both sibling walks (`matvec_bf16_v2`/`v3` and
937/// `q8_0_mmvq_row1_rp_v2_ilp`). The e4m3 family had no twin, and the B200 hybrid mint makes it the
938/// KDA hot path by quantizing all six KDA projections to per-tensor e4m3 on 34 of 45 layers.
939/// `MEMRA_E4M3_ROW_ILP`: 0 = shipped serial walk, 1 = ILP loads-in-flight. Nothing else.
940///
941/// A ladder of wider blocks (8, 16, 32 rows) and CTA-staged activations was built and priced on
942/// 2x B200 (3 reps x 15 interleaved iters, every arm bit-identical). Every one of them LOST:
943/// r8 0.988-0.993, r8+staged 0.932-0.934, r16 0.970-0.972, r16+staged 0.919-0.926,
944/// r32+staged 0.739-0.746, against ILP-at-4-rows' 1.003-1.015. The arms are gone; the reasons are
945/// the durable part and live in `qmatvec.cu` beside this kernel and in
946/// VERDICT:e4m3-fused6-wider-blocks-and-staged-activation-KILLED.
947pub fn e4m3_row_ilp_level() -> u32 {
948    static LVL: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
949    *LVL.get_or_init(|| u32::from(std::env::var("MEMRA_E4M3_ROW_ILP").as_deref() == Ok("1")))
950}
951
952pub fn e4m3_row_ilp_on() -> bool {
953    e4m3_row_ilp_level() > 0
954}
955
956/// `MEMRA_Q8_ROW_ILP` (lane/glm5-q8-row-ilp-20260904; default ON on sm_100a builds since
957/// 2026-09-04, OFF elsewhere, `=0`/`=1` override): the W8-posture q8_0 row
958/// walk (`qmatvec_q8_0_mmvq_rp_v2` at t=1 and the fused `qmatvec_kda6_q8f32_rp_v2`) takes its
959/// `_ilp` twin, the same per-row program with four blocks' loads per lane issued ahead of the
960/// dp4a chains. Read PER CALL (a live rollback seam). Why and receipts: the kernel header in
961/// cu/qmatvec.cu and docs/FLAGS.md.
962pub(crate) fn q8_row_ilp_on() -> bool {
963    q8_row_ilp_on_from(
964        std::env::var("MEMRA_Q8_ROW_ILP").ok().as_deref(),
965        env!("MEMRA_BUILT_CUDA_ARCH"),
966    )
967}
968
969/// The pure parse behind [`q8_row_ilp_on`]: `1` arms, `0` disarms, unset follows the BUILD ARCH
970/// (ON for `100a`, OFF otherwise): the twins carry a 2x B200 receipt (+2.16% at c1, darklanes
971/// research/glm5-b200-20260902/LANE.md, q8ab) and no SM120 one.
972pub fn q8_row_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
973    match v.map(str::trim) {
974        Some("1") => true,
975        Some("0") => false,
976        _ => built_arch == "100a",
977    }
978}
979
980/// Engagement counter for `MEMRA_Q8_ROW_ILP` (both launch sites); gates take a delta.
981pub static Q8_ROW_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
982    std::sync::atomic::AtomicU64::new(0);
983
984/// Snapshot of [`Q8_ROW_ILP_DISPATCHES`].
985pub fn q8_row_ilp_dispatches() -> u64 {
986    Q8_ROW_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
987}
988
989fn q8_row_ilp_note(site: &str) {
990    if Q8_ROW_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
991        eprintln!(
992            "[q8-row-ilp] engaged at {site}: q8_0 row walk with four blocks' loads per lane \
993             ahead of the dp4a chains (MEMRA_Q8_ROW_ILP=1)"
994        );
995    }
996}
997
998/// `MEMRA_NVFP4_ROW_ILP` (lane/glm5-nvfp4-row-ilp-20260904; default ON on sm_100a builds since
999/// 2026-09-04, OFF elsewhere, `=0`/`=1` override): the NVFP4
1000/// split-plane trunk matvec (`qmatvec_nvfp4_mmvq_mr2_rp`, and `qmatvec_nvfp4_mmvq_rp` when the
1001/// B200 grid-fill arm picks mr1) takes its `_ilp` twin: four groups' loads per lane issued
1002/// ahead of the table lookups and dp4a chains, same per-row accumulation order. Read PER CALL.
1003/// Why and receipts: the kernel header in cu/qmatvec.cu and docs/FLAGS.md.
1004pub(crate) fn nvfp4_row_ilp_on() -> bool {
1005    nvfp4_row_ilp_on_from(
1006        std::env::var("MEMRA_NVFP4_ROW_ILP").ok().as_deref(),
1007        env!("MEMRA_BUILT_CUDA_ARCH"),
1008    )
1009}
1010
1011/// The pure parse behind [`nvfp4_row_ilp_on`]: `1` arms, `0` disarms, unset follows the BUILD
1012/// ARCH (ON for `100a`, OFF otherwise): the twins carry a 2x B200 receipt (+1.98% alone, +2.55%
1013/// with the grid fill, darklanes research/glm5-b200-20260902/LANE.md, nvab) and no SM120 one.
1014pub fn nvfp4_row_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
1015    match v.map(str::trim) {
1016        Some("1") => true,
1017        Some("0") => false,
1018        _ => built_arch == "100a",
1019    }
1020}
1021
1022/// Engagement counter for `MEMRA_NVFP4_ROW_ILP`; gates take a delta.
1023pub static NVFP4_ROW_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
1024    std::sync::atomic::AtomicU64::new(0);
1025
1026/// Snapshot of [`NVFP4_ROW_ILP_DISPATCHES`].
1027pub fn nvfp4_row_ilp_dispatches() -> u64 {
1028    NVFP4_ROW_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1029}
1030
1031/// `MEMRA_B200_MR1_FILL=<blocks per SM>` (lane/glm5-nvfp4-row-ilp-20260904; default 16 since
1032/// 2026-09-04, receipt +2.03% alone / +2.55% with the ILP twin on the pair; 2 was the
1033/// lane/b200-matvec-occupancy-20260902 threshold): the B200 grid-fill arm of the NVFP4 m=1
1034/// decode matvec (under `MEMRA_B200_MATVEC_ARM=1`) forces mr1 (two warps per row pair -> one
1035/// warp per row, the shipped `qmatvec_nvfp4_mmvq_rp`) when the mr2 grid would be fewer than
1036/// this many four-warp blocks per SM. At 2 the 4096-row shapes (512 blocks, 2,048 warps on a
1037/// 148-SM part) stay mr2; 16 is one full wave of four-warp blocks (64 warp slots).
1038pub(crate) fn b200_mr1_fill() -> u32 {
1039    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1040    *V.get_or_init(|| {
1041        std::env::var("MEMRA_B200_MR1_FILL")
1042            .ok()
1043            .and_then(|v| v.trim().parse::<u32>().ok())
1044            .filter(|&n| n >= 1)
1045            .unwrap_or(16)
1046    })
1047}
1048
1049/// Dynamic shared memory one v3 CTA needs at this block size: `STAGES * R * (nb*8) * 2` bytes of
1050/// stage buffers plus `R * nb * 4` bytes of reduction window. Mirrors `MEMRA_GEMV_V3_STAGES` and
1051/// `MEMRA_GEMV_V3_REDOFF` in cu/qmatvec.cu; the two MUST move together.
1052pub(crate) const GEMV_V3_STAGES: usize = 2;
1053
1054/// Warps per block for the v2 q8_0 W8-posture twins. Mirrors `MEMRA_Q8_V2_ROWS` in
1055/// cu/qmatvec.cu (the shipped kernels use `MEMRA_MMVQ_ROWS` = 4); the two MUST move together.
1056pub(crate) const Q8_V2_ROWS: u32 = 8;
1057pub(crate) fn gemv_v3_smem_bytes(nb: usize) -> usize {
1058    GEMV_V3_STAGES * GEMV_V2_ROWS * (nb * 8) * 2 + GEMV_V2_ROWS * nb * 4
1059}
1060
1061/// True if a v3 launch fits the 48 KB default dynamic-shared-memory cap at `mmv_block()`.
1062/// 36 KB at the default 128; 72 KB at 256, which does NOT fit and falls back to v2 rather than
1063/// opting into `cudaFuncAttributeMaxDynamicSharedMemorySize` for a door that is still pending
1064/// its receipt.
1065pub(crate) fn gemv_v3_fits() -> bool {
1066    gemv_v3_smem_bytes(mmv_block() as usize) <= 48 * 1024
1067}
1068
1069/// Rows per block for the v2 bf16 GEMV family. Mirrors `MEMRA_GEMV_V2_ROWS` in cu/qmatvec.cu;
1070/// the two MUST move together (the launcher's grid and dynamic-smem size are derived from it).
1071pub(crate) const GEMV_V2_ROWS: usize = 8;
1072
1073/// Engagement counters for `MEMRA_B200_GEMV_V2` (lane/b200-gemv-hbm-20260902), one per arm,
1074/// the `KDA_FUSED6_*` precedent in kda.rs. Counted at each arm's own call site
1075/// (`moe_grouped_prefill_dispatches` precedent): a door that never actually took its path must
1076/// not be indistinguishable from one that did, and that has to hold PER ARM, not per door. In
1077/// the W8 posture the decode t=1 trunk and the t=2..=8 verify walk both engage in the same
1078/// process; with one shared print-once gate whichever fired second never announced, and a box
1079/// A/B could not attribute engagement to the arm that actually ran. Each counter gates its own
1080/// print-once line.
1081///
1082/// W8 verify width arm, `qmatvec_q8_0_rows_tw_v2` (t in 2..=8).
1083pub(crate) static GEMV_V2_Q8_ROWS_TW_DISPATCHES: std::sync::atomic::AtomicU64 =
1084    std::sync::atomic::AtomicU64::new(0);
1085/// W8 decode t=1 trunk arm, `qmatvec_q8_0_rp_v2`.
1086pub(crate) static GEMV_V2_Q8_RP_DISPATCHES: std::sync::atomic::AtomicU64 =
1087    std::sync::atomic::AtomicU64::new(0);
1088/// bf16 row arm, `matvec_bf16_v2` (level 1) or `matvec_bf16_v3` (level 2); the line names which.
1089pub(crate) static GEMV_V2_BF16_DISPATCHES: std::sync::atomic::AtomicU64 =
1090    std::sync::atomic::AtomicU64::new(0);
1091
1092/// MEMRA_B200_BF16_GEMV_LT=1: cuBLASLt REFERENCE door for the t=1 bf16 decode row matvec
1093/// (lane/b200-gemv-hbm-20260902, docs/FLAGS.md).
1094///
1095/// WHAT IT IS FOR. The B200 census puts `matvec_bf16_f32acc_x4_rows` at 23.6us for 64 MB of
1096/// bf16 weight reads = 2.7 TB/s, 34% of the 8 TB/s HBM3e wall, and `qmatvec_kda6_bf16f32` at
1097/// 93.8us for ~200 MB = 2.1 TB/s (26%). Before writing a faster memra kernel it is worth
1098/// knowing what a TUNED VENDOR LIBRARY reaches on the same bytes on this part, because that
1099/// number bounds what "a well-scheduled GEMV" looks like on sm_100a. This door routes those
1100/// rows through `cublasLtMatmul` (m=1, bf16 x bf16 -> f32, the `memra_bf16_pp_gemm` TN plan
1101/// with the per-device handle from cu/f16_prefill.cu) so the box can measure it directly.
1102///
1103/// IT IS A NAMED NUMERIC CLASS, NOT A BIT-IDENTICAL TWIN. Two things change: the ACTIVATION is
1104/// cast f32 -> bf16 before the multiply (the shipped kernel keeps the f32 activation and only
1105/// widens the bf16 weight), and the summation order over K is cuBLASLt's, not the shipped
1106/// per-thread chain + red[] tree. Class name: `bf16_gemv_lt` (the same class
1107/// `MEMRA_PP_BF16`'s prefill GEMM already ships under, at m=1). Because of that it is a
1108/// REFERENCE door: default OFF, never a serving default, and it is not a candidate for
1109/// promotion without its own argmax/serving acceptance.
1110///
1111/// Restricted to `sm_100a` BUILDS (`MEMRA_BUILT_CUDA_ARCH`, baked in at compile time), like
1112/// `MEMRA_B200_MATVEC_ARM`: setting it on an sm_120a build is a documented no-op, so the naked
1113/// sm_120a defaults stay byte-identical.
1114pub(crate) fn b200_bf16_gemv_lt_on() -> bool {
1115    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1116    *V.get_or_init(|| {
1117        env!("MEMRA_BUILT_CUDA_ARCH") == "100a"
1118            && std::env::var("MEMRA_B200_BF16_GEMV_LT").as_deref() == Ok("1")
1119    })
1120}
1121
1122/// MEMRA_STEP_TP_W8=1: q8_0 mirror of the step TP attention projections for DECODE.
1123///
1124/// NUMERIC-CLASS door, same class and acceptance as `MEMRA_STEP_TP_QKV_FUSED` /
1125/// `MEMRA_BF16_MMV`: the per-row arithmetic becomes an int8 dp4a dot
1126/// with per-32 scales instead of a bf16xf32 fma chain, so a bit-tape cannot apply and the
1127/// acceptance is the argmax gate plus the boot battery. Motivation is measured, not assumed
1128/// (`decode-kernel-census`, 2026-08-25): the fused qkv shape runs 23.0 us in bf16 at
1129/// 1.83 TB/s and 14.0 us in q8_0 at 1.60, and o_proj 24.2 -> 11.7 us — together
1130/// ~-1.0 ms of a 13.16 ms token. Default OFF.
1131/// MEMRA_W8_HYBRID=1 opts the door's HYBRID half in (LM head, shared expert, dense FFN).
1132/// Default OFF on measurement AND on residency: it moved decode +0.1% (the W8 trace showed it
1133/// only ever mirrored the shexp down rows, which SHEXP_OVERLAP already hides), while costing
1134/// ~1.7 GB per card on top of the attention mirrors' ~0.9 GB — and at the model's NATURAL
1135/// 262144-token context the full set does not fit: `MEMRA_STEP_TP_W8=1` there dies in
1136/// CUDA_ERROR_OUT_OF_MEMORY while plain decode runs 76.03 tok/s.
1137/// STEP37 SERVING DEFAULTS (owner flip, 2026-08-27). The step37 serving shape — the t-row walk,
1138/// the q8 W8 doors, the SWA ring, the NVFP4 draft heads, and this lane's three verify fixes —
1139/// was gated door by door (byte tape == plain, acceptance unchanged, run-spec K=1..8 PASS,
1140/// interleaved x5 wall, vendor-default sampled cell with engagement receipts: greedy 93.18 vs
1141/// 81.95 plain, sampled 81.79 vs 78.50) and the owner ordered the defaults ON. The doors' call
1142/// sites are not all family-scoped (the W8 mirror routing sits inside generic matmul paths), so
1143/// the default arms AT MODEL LOAD when the plan compiles to the SlidingGatedMoe program, never
1144/// globally. Every door keeps a per-flag env override: `=1` forces ON for any family, `=0` is
1145/// the kill switch — the rollback seam the FLAGS rows name. Per-process: a process that loads a
1146/// step37-class model arms the defaults for its lifetime.
1147static STEP37_SERVING_DEFAULTS: std::sync::atomic::AtomicBool =
1148    std::sync::atomic::AtomicBool::new(false);
1149
1150pub fn arm_step37_serving_defaults() {
1151    STEP37_SERVING_DEFAULTS.store(true, std::sync::atomic::Ordering::Relaxed);
1152    crate::cache::set_swa_ring_default(true);
1153    eprintln!(
1154        "[step37-defaults] serving doors armed ON for the SlidingGatedMoe program \
1155         (per-flag =0 kills, =1 forces; owner flip 2026-08-27)"
1156    );
1157}
1158
1159pub(crate) fn step37_defaults_armed() -> bool {
1160    STEP37_SERVING_DEFAULTS.load(std::sync::atomic::Ordering::Relaxed)
1161}
1162
1163/// Tri-state door: `=1` ON, `=0` OFF, unset = the family default (ON once a step37-class model
1164/// armed it, OFF otherwise). The env parse is cached; the family default is read live because
1165/// arming happens at model load, possibly after another door's first read.
1166pub(crate) fn step37_door(cell: &'static std::sync::OnceLock<Option<bool>>, name: &str) -> bool {
1167    match *cell.get_or_init(|| match std::env::var(name).ok().as_deref() {
1168        Some("1") => Some(true),
1169        Some("0") => Some(false),
1170        _ => None,
1171    }) {
1172        Some(forced) => forced,
1173        None => step37_defaults_armed(),
1174    }
1175}
1176
1177pub(crate) fn w8_hybrid_on() -> bool {
1178    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1179    step37_door(&ENV, "MEMRA_W8_HYBRID")
1180}
1181
1182pub(crate) fn step_tp_w8_on() -> bool {
1183    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1184    step37_door(&ENV, "MEMRA_STEP_TP_W8")
1185}
1186
1187/// MEMRA_GLM5_W8=1 (default OFF, unset/0 = bf16): q8_0 mirror of the bf16-resident glm5_next
1188/// KDA and MLA decode projections, modeled on `MEMRA_STEP_TP_W8`'s hybrid half but its OWN
1189/// independent door — a strict boolean, NOT step37-family-armed and NOT gated behind
1190/// `MEMRA_W8_HYBRID`. Reuses the SAME building block: `matvec_bf16_via_q8_mirror` /
1191/// `matvec_bf16_via_q8_mirror_t` (pointer-keyed, built on first decode use, `w8_mirrors`/
1192/// `w8_act` caches shared with the step37 door). NUMERIC-CLASS door, same class and
1193/// acceptance shape as `MEMRA_STEP_TP_W8`: the per-row arithmetic becomes an int8 dp4a dot
1194/// with per-32 scales instead of a bf16xf32 fma chain, so the acceptance is the argmax gate
1195/// (`glm5_w8_gate`), not a bit tape. Motivation (nsys, 2x B200, GLM-5.3-Flash NVFP4 mint,
1196/// resident PP2, plain decode t=1): ~15 GB/token weight reads per token, of which the
1197/// BF16-resident KDA/MLA projections (`matvec_bf16_f32acc_x4_rows`, 211 launches/token,
1198/// ~13.5 GB/token) dominate; the mirror halves that class's per-weight bytes (2 B bf16 -> ~
1199/// 1.0625 B q8_0). See docs/FLAGS.md for the bytes/token arithmetic and rollback seam.
1200pub(crate) fn glm5_w8_on() -> bool {
1201    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1202    *ON.get_or_init(|| std::env::var("MEMRA_GLM5_W8").as_deref() == Ok("1"))
1203}
1204
1205/// Dispatch counter for `MEMRA_GLM5_W8`'s decode-tier engagement, announced once (the
1206/// no-announce-cannot-be-read-both-ways lesson from `MEMRA_W8_VIEW`).
1207pub(crate) static GLM5_W8_DISPATCHES: std::sync::atomic::AtomicU64 =
1208    std::sync::atomic::AtomicU64::new(0);
1209
1210/// MEMRA_W8_VIEW=1: extend the W8 hybrid half to the ROW-RANGE-VIEW GEMVs, i.e. the lo halves
1211/// that `MEMRA_HEAD_SPLIT` and `MEMRA_SHEXP_OVERLAP` keep on rank 0. NOT a step37 family door
1212/// and NOT armed by `arm_step37_serving_defaults`: it stays off until it carries its own
1213/// interleaved speed rows and its own argmax gate. Unset or `=0` is the rollback seam.
1214pub(crate) fn w8_view_on() -> bool {
1215    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1216    *ON.get_or_init(|| std::env::var("MEMRA_W8_VIEW").as_deref() == Ok("1"))
1217}
1218
1219/// MEMRA_Q8T_WONCE=1: the q8 t-column verify kernels take their weight-once `_tw` twins — one
1220/// row grid, each weight int4 loaded once and dotted against all t columns — instead of the `_t`
1221/// forms, whose column grid axis plus __ldcs (streaming, evict-first) re-reads the fully-shared
1222/// weights from DRAM once per column (nsys 2026-08-27: qkv_rp_t 1.67x, b4_rp_t 1.43x a
1223/// single-column call for 2 columns, where weight-bound scaling says ~1.1x). Per-column float
1224/// program unchanged (same lane-strided blk order, own accumulator chain, same reduce); default
1225/// off until the byte tape says so.
1226/// MEMRA_STEP_GEMM_PRIME: prime chunks (t>=16) route the routed MoE through the grouped f16 GEMM
1227/// over the resident NVFP4 banks instead of the per-token device routes. FAMILY-DEFAULT ON since
1228/// 2026-08-28 because on the server route it is the only prime that WORKS: measured there, walk
1229/// = ERR (tail chunk missing from the distributed kv), fallback chunked prime = 29 s on a
1230/// ~450-token prompt and a 90 s TIMEOUT at 4k, grouped GEMM = 3.5-4.9 s with coherent output.
1231/// `=0` is the kill switch back to the fallback prime.
1232pub(crate) fn step_gemm_prime_on() -> bool {
1233    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1234    step37_door(&ENV, "MEMRA_STEP_GEMM_PRIME")
1235}
1236
1237pub(crate) fn q8t_wonce_on() -> bool {
1238    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1239    step37_door(&ENV, "MEMRA_Q8T_WONCE")
1240}
1241
1242/// MEMRA_TOPK_FAST=1: barrier-lean sigmoid top-k twin (warp-local top-k + one merge).
1243/// Selection and weight arithmetic identical to the round-robin kernel — a latency twin.
1244/// MEMRA_SIG_EXPF_DEV=1: device-libm expf sigmoid router (numeric-class door — the
1245/// host-glibc transcription is FP64-rate-bound on consumer Blackwell). New tape + battery.
1246pub(crate) fn sig_expf_dev_on() -> bool {
1247    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1248    *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
1249}
1250
1251pub(crate) fn topk_fast_on() -> bool {
1252    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1253    *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
1254}
1255
1256/// Select the sigmoid-router kernel without ever sending a shape wider than the fast
1257/// kernels' fixed eight-pick scratch. The generic and dexp kernels support the full host
1258/// contract; both `_fast` twins index `[warp][8]` storage and would write out of bounds for
1259/// `n_used > 8` (Hermes `0d220d8c9a3eb634`).
1260fn sigmoid_topk_kernel(sig_expf: bool, fast: bool, n_used: usize) -> &'static str {
1261    match (sig_expf, fast && n_used <= 8) {
1262        (true, true) => "moe_router_sigmoid_topk_f32_dexp_fast",
1263        (true, false) => "moe_router_sigmoid_topk_f32_dexp",
1264        (false, true) => "moe_router_sigmoid_topk_f32_fast",
1265        (false, false) => "moe_router_sigmoid_topk_f32",
1266    }
1267}
1268
1269#[cfg(test)]
1270mod sigmoid_topk_dispatch_tests {
1271    #[test]
1272    fn fast_kernel_refuses_wide_topk_and_composes_with_dexp() {
1273        use super::sigmoid_topk_kernel;
1274
1275        assert_eq!(
1276            sigmoid_topk_kernel(false, true, 8),
1277            "moe_router_sigmoid_topk_f32_fast"
1278        );
1279        assert_eq!(
1280            sigmoid_topk_kernel(true, true, 8),
1281            "moe_router_sigmoid_topk_f32_dexp_fast"
1282        );
1283        assert_eq!(
1284            sigmoid_topk_kernel(false, true, 9),
1285            "moe_router_sigmoid_topk_f32"
1286        );
1287        assert_eq!(
1288            sigmoid_topk_kernel(true, true, 9),
1289            "moe_router_sigmoid_topk_f32_dexp"
1290        );
1291    }
1292}
1293
1294pub(crate) fn rms_block() -> u32 {
1295    static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1296    *V.get_or_init(|| {
1297        std::env::var("MEMRA_RMS_BLOCK")
1298            .ok()
1299            .and_then(|v| v.parse().ok())
1300            .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1301    })
1302}
1303
1304pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
1305    static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1306    if let Some(forced) = *S.get_or_init(|| {
1307        std::env::var("MEMRA_FA_SPLIT")
1308            .ok()
1309            .and_then(|v| v.parse().ok())
1310            .filter(|&s: &usize| s >= 8 && s % 8 == 0)
1311    }) {
1312        return forced;
1313    }
1314    // CTX-ADAPTIVE default (2026-07-05 40k sweep: sp32 24.5 vs sp128 26.0 tok/s = +5.8% — at
1315    // deep ctx the n_splits count explodes (40k/32 = 1265 splits x 8 kv-heads) and the combine
1316    // + partial-buffer cost dominates; at short ctx small splits fill the SMs). Exactness: split
1317    // size only changes the PARTITION of keys; the rows/combine order per split is fixed and the
1318    // gate battery (kernel-check + run-spec K=1..8) arbitrates every default change.
1319    //
1320    // SM-AWARE SHORT-CTX RUNG (2026-07-06 rtx6000): the 32-key rung was tuned on the 82-SM 5090.
1321    // On 188 SMs the vec grid (n_head_kv x n_splits CTAs) starves at short ctx — the 35B has
1322    // n_head_kv=2, so ctx128/split32 = 8 CTAs on 188 SMs. Measured on rtx6000 (N=1 sweep + N=3
1323    // interleaved confirm): 35B ctx128 sp16 179 vs sp32 161 (+11%), ctx512 178 vs 158, ctx2048
1324    // flat, ctx>=4096 sp64 edges sp16 by ~3%; 27B ctx128 70.9 vs 66.3 (+7%); 9B 177 vs 163
1325    // (+9%). Rigs <=100 SMs keep the validated 5090 ladder EXACTLY (default unchanged there —
1326    // rig-divergence law: this branch is measured on 188 SMs only).
1327    // gemma4 all-16 ladder probe REVERTED (2026-07-10): +1.3 plain at d1736 (157.5 vs 156.2)
1328    // but depth VERIFY collapsed (spec 203.5 -> 169 — the windowed rows' per-row combine over
1329    // 64 splits). The mixed default (swa nkv=8 -> 32, globals nkv=2 -> 8-ladder) stays; a
1330    // caller-split policy would break row-vs-decode split parity. FA_SP_GEMMA kept as a seam.
1331    if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
1332        && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
1333    {
1334        return if t_kv <= 8192 {
1335            16
1336        } else if t_kv <= 16384 {
1337            64
1338        } else {
1339            128
1340        };
1341    }
1342    let big_rig = fa_sm_count() >= 128;
1343    if big_rig {
1344        let _ = n_head_kv;
1345        if t_kv <= 2048 {
1346            // MEMRA_FA_SP_SHORT=N: the SHORT rung only (the SWA layers' capped t_kv lands
1347            // here on step37: 33 of 45 layers at t_kv=512). At 16 the tile loop runs
1348            // HALF-EMPTY (FA_DEC_TILE=32 -> nt=16 per split), so the V staging pass moves a
1349            // half tile per iteration and the combine carries 2x the partials; 32 makes each
1350            // split exactly one full tile. A global MEMRA_FA_SPLIT cannot isolate this — it
1351            // moves the deep-ctx rung too, where more splits measured worse.
1352            // NUMERIC-CLASS door (key partition -> different per-split partials/combine):
1353            // new tape + battery, exactly like every other split-ladder change.
1354            static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1355            if let Some(sp) = *SHORT.get_or_init(|| {
1356                std::env::var("MEMRA_FA_SP_SHORT")
1357                    .ok()
1358                    .and_then(|v| v.parse().ok())
1359                    .filter(|&s: &usize| s >= 8 && s % 8 == 0)
1360            }) {
1361                return sp;
1362            }
1363            16
1364        } else if t_kv <= 16384 {
1365            64
1366        } else {
1367            128
1368        }
1369    } else if n_head_kv <= 4 {
1370        // KV-HEAD-AWARE RUNG (2026-07-08, 5090): the 8192->32 rung was validated on kv=8 models
1371        // (27B/9B: 8 heads x n_splits fills 82 SMs). The 35B has n_head_kv=2 — at ctx512/sp32
1372        // the vec grid is 2 x 20 = 40 CTAs on 82 SMs (half idle). Measured (35B, run-gen 128tok
1373        // N=1 sweep + N=3 confirm): sp8 162.1 / sp16 161.3 / sp32 159.4 at short ctx.
1374        // DEPTH TAPER (same day, the deep-ctx lesson re-learned on this rung): sp8 at d6257 =
1375        // 782 splits -> combine + partial-buffer cost dominates (141.2 tok/s); the d6257 sweep
1376        // says sp64 = 153.0 (sp16/32 147, sp96 147.6, sp128 141). Few-kv-head models need the
1377        // taper EARLIER than kv=8 (per-split grid 4x thinner, same per-split combine cost).
1378        // Crossover hunt: sp8 vs sp64 = 156.7/155.9 at d3072, 151.7/155.6 at d4096 -> boundary 3072.
1379        // RUNG RE-SWEPT UNDER THE DEEP KERNEL (2026-08-02, lane/ladder-3072 — the stale-verdict
1380        // law: the 3072 boundary was calibrated on the conflicted v4 core; the deep rewrite cut
1381        // vec cost ~1.2-1.4x while combine scales with n_splits, so sp8's combine bill
1382        // dominates far earlier). Kernel receipts (quiet-rig nsys, deep vec + combine us):
1383        // d1024 sp8 17.1 vs sp64 10.6; d2048 31.0 vs 12.2; d3072 44.0 vs 18.3. e2e run-gen
1384        // tg128 N=3 interleaved (KAT + q35, research/ladder-3072-20260802/): sp8 loses at
1385        // EVERY depth >= 1024 (KAT d2048 182.6 vs 188.0 = -2.9%, d3072 175.9 vs 186.4 =
1386        // -5.6%; q35 d4096 169.2 vs 182.6 = -7.4%); d512 flat (+-0.2%, inside noise). sp32
1387        // ties sp64 within noise in the mid band and loses at d4096 -> no extra rung.
1388        // Boundary 3072 -> 512: sp8 keeps only the short-ctx band it was validated on
1389        // (ctx128-512); sp64 takes over where the deep kernel made combine the bill.
1390        if t_kv <= 512 {
1391            8
1392        } else if t_kv <= 16384 {
1393            64
1394        } else {
1395            128
1396        }
1397    } else {
1398        if t_kv <= 8192 {
1399            32
1400        } else if t_kv <= 16384 {
1401            64
1402        } else {
1403            128
1404        }
1405    }
1406}
1407
1408/// SM count of device 0, cached (used by fa_split_keys' rig-size rung; primary-context query,
1409/// same attribute Engine::batched_variant reads).
1410pub(crate) fn fa_sm_count() -> i32 {
1411    static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
1412    *N.get_or_init(|| {
1413        cudarc::driver::result::init().ok();
1414        cudarc::driver::result::device::get(0)
1415            .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
1416                d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
1417            .unwrap_or(82)
1418    })
1419}
1420
1421/// FA-prefill kernel-name suffix for a head_dim (the template-stamped twins in flash_attn.cu):
1422/// 256 = the original names (qwen35 class, dispatch unchanged), 128 = `_hd128` (MiniMax-M3).
1423/// Any other dim errors — callers gate to sdpa_naive before dispatching FA.
1424#[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1425fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
1426    match head_dim {
1427        256 => Ok(""),
1428        128 => Ok("_hd128"),
1429        d => Err(format!(
1430            "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
1431                          callers must gate to sdpa_naive"
1432        )
1433        .into()),
1434    }
1435}
1436
1437/// Quant type codes matching qmatvec.cu QType enum.
1438pub const QT_Q8_0: i32 = 0;
1439pub const QT_Q4_K: i32 = 1;
1440pub const QT_Q6_K: i32 = 2;
1441pub const QT_Q5_K: i32 = 3;
1442pub const QT_Q3_K: i32 = 4;
1443pub const QT_IQ4_XS: i32 = 5;
1444pub const QT_IQ3_S: i32 = 6;
1445pub const QT_NVFP4: i32 = 7;
1446/// Slot-major v2 bank permutation of `QT_NVFP4` (see tp.rs `nvfp4_matrix_v2_permute`) — only the
1447/// grouped-prefill dequant consumes this tag; every direct/dp4a lane must keep refusing it.
1448pub const QT_NVFP4_V2: i32 = 107;
1449/// Checkpoint-native FP8-E4M3 (MEMRA_ST_E4M3, lane e4m3dec): raw safetensors e4m3 weight bytes
1450/// [out_f, in_f] row-major (row_bytes == in_f), per-tensor f32 weight_scale in GpuTensor `scale`
1451/// (fused at the mmvq write / post-matmul scale_inplace). Decode = qmatvec_e4m3_mmvq (+ _b2/_b4/_b8
1452/// batched twins); prefill (m>=16) = the cuBLASLt FP8 GEMM on the SAME resident bytes (fp8_ffi.rs)
1453/// — ONE weight copy total, no Q8_0 re-encode duplicate.
1454pub const QT_F8_E4M3: i32 = 10;
1455/// Device-side tag for the A6 SPLIT-PLANE repacked NVFP4 layout (Stage-A generic kernel only;
1456/// GpuTensor keeps qtype=QT_NVFP4 + an `rp` flag — this tag never lives in a GpuTensor).
1457pub const QT_NVFP4_RP: i32 = 9;
1458/// Unquantized f32 weight (safetensors MoE Path A: experts dequantized to f32 host-resident).
1459pub const QT_F32: i32 = 8;
1460pub const QT_BF16: i32 = 11;
1461pub const QT_Q4_0: i32 = 12; // gemma-4 QAT GGUF weight format (18B/32: fp16 d + nibbles)
1462/// GGUF Q2_K. Appended after the existing Q4_0 code so kernel ABI values do not move.
1463/// Mixed-expert artifacts use the generic f32-dequant staged kernel until a target-rig-gated
1464/// dp4a/MMQ implementation exists.
1465pub const QT_Q2_K: i32 = 13;
1466/// Checkpoint-native FP8-E4M3 with a BLOCK-128 weight-scale GRID (lane/fp8-blk128-decode,
1467/// 2026-08-05) — the Qwen-official FP8 / DeepSeek-V3 scale class. Same raw e4m3 bytes as
1468/// `QT_F8_E4M3` ([out_f, in_f] row-major, row_bytes == in_f), but the dequant scale is
1469/// `GpuTensor::Quant.blk` (`Fp8BlockScales`, [ceil(out_f/128), ceil(in_f/128)] f32) and the
1470/// scalar `scale` field is 1.0 by the layout contract.
1471///
1472/// WHY A DISTINCT CODE rather than `QT_F8_E4M3` + a `blk` flag: every existing QT_F8_E4M3
1473/// consumer (qmatvec_e4m3_mmvq and its batched/fused twins, e4m3_fused_params,
1474/// matmul_pre_dual_noscale's F8 arm, try_fp8_gemm) threads exactly ONE scalar weight scale. Under
1475/// a shared code, any consumer that was not taught the grid would still MATCH and would dequant
1476/// every tile at scale 1.0 — a silent numeric corruption. Under a distinct code every untaught
1477/// consumer refuses loudly instead (`mmvq_supports`/`gemm_supports`/`mmq_supports` return false;
1478/// the mmvq name match panics), so a missed dispatch site is a crash or a refusal receipt, never
1479/// wrong numbers. Decode = `qmatvec_e4m3_blk_mmvq`; prefill (m>=16) = the per-block FP8 MMQ tile
1480/// on the SAME resident bytes+grid (fp8_ffi::try_fp8_blk_mmq) — ONE weight copy total.
1481pub const QT_F8_E4M3_BLK: i32 = 14;
1482
1483/// Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).
1484pub struct Engine {
1485    pub gpu: memra_runtime::Gpu,
1486    module: Arc<CudaModule>,
1487    hybrid: Arc<CudaModule>,
1488    /// Kimi Delta Attention kernels (cu/kda.cu) — separate fatbin, resolved through `func`.
1489    kda: Arc<CudaModule>,
1490    qmatvec: Arc<CudaModule>,
1491    flash: Arc<CudaModule>,
1492    /// FP8-GLOBALS module (2026-07-11): the kf8vf8 fatbin loaded ALONGSIDE the default —
1493    /// gemma GLOBAL layers (hd512) append + attend in e4m3 (dequant-latency arc, HANDOVER).
1494    /// Lazy: loaded on first global-format use; None until then.
1495    flash_g: std::sync::OnceLock<Arc<CudaModule>>,
1496    gemm: Arc<CudaModule>,
1497    router: Arc<CudaModule>,
1498    /// Sampled-spec kernels (research/sampled-spec-impl-map.md piece A).
1499    sample: Arc<CudaModule>,
1500    /// EDGE-1 §B: one shared SLRU expert-residency cache, lazily built on first MoE dispatch under
1501    /// MEMRA_MOE_CACHE. `Mutex` makes it multi-agent safe (§E.2); the lock covers only lookup/admit/
1502    /// memcpy-issue (µs), NOT the GEMM, so streams still overlap. `None` => cache disabled.
1503    moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
1504    /// MEMRA_STEP_TP_W8, hybrid half: q8_0 mirrors of bf16 GEMV weights that do NOT live in a
1505    /// TP resident bank (the LM head, the shared expert, the dense-FFN layers), keyed by the
1506    /// bf16 slab's device pointer and built on first decode use. The mirror is 1.0625 B/w
1507    /// against bf16's 2, and the raw slab stays resident, so prefill keeps its arithmetic.
1508    /// KEYED ON (pointer, in_f, out_f), not on the pointer alone: a row-range VIEW of a slab
1509    /// carries the PARENT's base pointer when the range starts at row 0, so a pointer-only key
1510    /// would hand the head-split lo half (4096 x 64448) the full head's mirror (4096 x 128896)
1511    /// and read 2x past the rows it owns. The shape is part of the identity of a mirror.
1512    /// SECOND CONSUMER (2026-09-02): `MEMRA_GLM5_W8` reuses this SAME cache for the glm5_next
1513    /// KDA/MLA decode trunk — independent door, same building block, same key shape.
1514    w8_mirrors: Mutex<std::collections::HashMap<(u64, u32, u32), CudaSlice<u8>>>,
1515    /// Per-`in_f` q8_1 activation scratch for those mirrors (allocating per call would cost
1516    /// more than the door saves).
1517    #[allow(clippy::type_complexity)]
1518    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1519    w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
1520    /// Exact retained expert-block lengths collected after model load. Mixed-layout models use
1521    /// this inventory to preallocate fixed-address size classes instead of sizing every slot to
1522    /// the single largest block. The cache still owns every address for its full lifetime.
1523    moe_cache_layout: Mutex<Option<Vec<usize>>>,
1524    /// CAPTURE-RETAIN mode (graph arc, 2026-07-12): while a graph capture (and its allocator
1525    /// warmups) runs, every Engine allocation is ALSO kept alive here — a captured graph's
1526    /// transient buffers must never return to the pool, or later allocations (e.g. the spec
1527    /// verify between replays) reuse their addresses and the replay reads/writes live memory
1528    /// (the draft-graph corruption root cause). Fast-path cost when off: one relaxed atomic.
1529    capture_keep_on: std::sync::atomic::AtomicBool,
1530    /// VERIFY-EXACT scope (dflash lane, 2026-07-13): when set, matmul/matmul_pre skip the
1531    /// m>=16 prefill-GEMM branches so a t>=16 batched VERIFY rides the decode-exact b-tier
1532    /// class (the parity law). The t=16 dflash verify tripped the GEMM threshold — 770us/
1533    /// matmul (54% of the round) AND a different FP order than decode (issue-10 landmine).
1534    verify_exact: std::sync::atomic::AtomicBool,
1535    capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
1536    /// EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).
1537    pub copy_stream: Arc<CudaStream>,
1538    /// Resident CUTLASS NVFP4 prefill scratch (workspace + a_packed + sfa_linear + sfa_sw + y + alpha),
1539    /// allocated ONCE and grown to the largest prefill GEMM shape, then reused per-call. Removes the
1540    /// 6 fresh allocations + alpha htod that `cutlass_fp4_gemm` did every prefill matmul (~200/prefill).
1541    /// Safe as a single shared buffer because all GPU compute serializes on the one `gpu.stream` worker
1542    /// thread (the server runs one GPU worker; no concurrent CUTLASS GEMMs share this scratch). `None`
1543    /// until the first CUTLASS FP4 GEMM. Mutex guards lazy build/grow only (matches `moe_cache`).
1544    #[cfg(memra_cutlass)]
1545    cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
1546    /// FP8-ACT PREFILL scratch (MEMRA_PP_FP8): quantized-activation buffer + scale block + cuBLASLt
1547    /// workspace, allocated once and grown to the largest prefill m*k (see fp8_ffi.rs). `None`
1548    /// until the first FP8 prefill GEMM; Mutex guards lazy build/grow only (matches cutlass_scratch).
1549    fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
1550    /// f16-P/V door: pooled V re-encode buffer (bf16->f16) for the hd512 _pre path. Lazy-grow;
1551    /// per-call cudaMalloc was a laptop-regression suspect (VRAM pressure, 31B nkv=4 = 4x bytes).
1552    fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
1553    /// Pooled fa-decode split partials (part_o, part_m, part_l): per-call zeros() was 3
1554    /// alloc+memset pairs per fa launch (~144 mem nodes per decode token — the graph door's
1555    /// residual launch tax) — lazy-grow, memset-prefix per use, stream-ordered reuse.
1556    #[allow(clippy::type_complexity)]
1557    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1558    fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
1559    /// Retired fa-part pool generations (#68): old buffers whose addresses captured graphs may
1560    /// have baked — kept alive for the Engine's lifetime instead of returning to the async pool
1561    /// (see the RETIRE-ON-GROW comment at the realloc sites). Doubling growth bounds the total.
1562    #[allow(clippy::type_complexity)]
1563    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1564    fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
1565    /// name -> resolved CudaFunction (capture-safe lookups; see `func`).
1566    fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
1567    f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
1568    /// RANK1 LEVER (parallel argmax): resident pass-1 partials scratch (part_v[NB] f32, part_i[NB] i32),
1569    /// allocated ONCE on first parallel-argmax call and reused. Stable pointers so the 2-pass argmax
1570    /// is CUDA-graph-capturable (the buffer is referenced by both captured passes; lazy-allocated
1571    /// before capture under the generate_graph tracking-off window so it carries no events).
1572    argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
1573    /// ARC B (chunk-prime dequant-once): resident bf16 K/V workspace for `fa_prefill_view_ws`
1574    /// ((K bytes, V bytes) u8 buffers holding [t_kv, kv_dim] bf16). Grown lazily to the largest
1575    /// (t_kv, kv_dim) seen, REUSED across layers/chunks/calls (contents rewritten per launch —
1576    /// safe because all compute serializes on the one gpu.stream). ~82MB at 40k ctx on the 27B.
1577    prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
1578    /// LAUNCH-STRUCTURE STAGE 1: persistent PINNED (cacheable, flags=0) host staging buffer for the
1579    /// fused-router sel/w readback — one async DtoH pair + ONE sync instead of two synced dtohs.
1580    /// Grown lazily; reused every MoE layer (single-threaded decode serializes on the sync).
1581    router_stage: Mutex<Option<PinnedStage>>,
1582    /// Persistent hc-glue decode workspace (MEMRA_HC_DECODE_WS, lane/glm5-decode-diet lever 2).
1583    /// Pooled per engine like `fa_part_pool`: the buffers are pure per-step scratch (every
1584    /// element fully overwritten before read each step), so one slot per engine is correct
1585    /// even across sessions; the walk TAKES it for the step and puts it back, and a second
1586    /// concurrent walk on the same engine simply falls back to fresh allocations.
1587    hyper_decode_ws: Mutex<Option<crate::hyper::HyperDecodeWs>>,
1588    /// Per-session MLA PRE/POST handoff buffers (`MEMRA_MLA_SEG_WS`,
1589    /// lane/glm5-mla-capture-20260904); see [`crate::hybrid_forward::MlaSegWs`].
1590    mla_seg_ws: Mutex<Option<crate::hybrid_forward::MlaSegWs>>,
1591    /// Verify-walk allocation workspace (MEMRA_VERIFY_WS — glm5-alias
1592    /// MEMRA_GLM5_VERIFY_WS honored, OFF-wins; lane/glm5-matvec door W, generalized
1593    /// lane/glm5-extract-general — the pool is family-agnostic by content): the
1594    /// `MEMRA_HC_DECODE_WS` pattern extended to the spec verify walk, whose ~1380
1595    /// `cuMemAllocAsync`+Free pairs/token the t=1 workspace door structurally never reaches
1596    /// (spec decodes through the t=K+1 walk — diet-battery WINDOW.md). Size-keyed free-lists;
1597    /// verify-only call sites (the rows-exact matmul class, the KDA rows arm, the MoE vrows
1598    /// staging) draw from and recycle into it. Reuse is byte-identical by the same contract
1599    /// that makes `uninit` legal at those sites: every element is fully overwritten before
1600    /// any read, by the SAME unchanged kernels. Per-engine = per-stream, so stream ordering
1601    /// makes recycle-then-reuse safe exactly like free-then-alloc on the async pool.
1602    verify_ws: Mutex<VerifyWs>,
1603    /// Resident device mirrors of the per-expert NVFP4 `weight_scale_2` macro planes, keyed by
1604    /// `(layer, plane)` with plane 0/1/2 = gate/up/down (MEMRA_MOE_VROWS_DEV_TABLES, door D).
1605    /// The device table build needs `macro_scale(ex)` where the selection lives; the host plane
1606    /// is an immutable `Vec<f32>` of n_expert entries for the process lifetime, so ONE upload
1607    /// per (layer, plane) serves every subsequent layer-call — 3 x n_expert x 4 B (3.5 KB at
1608    /// 288 experts), 126 buffers = ~145 KB for a 42-MoE-layer model. Uploading per call instead
1609    /// would ADD three HtoD to a door whose whole purpose is removing two.
1610    vrows_macro_dev: Mutex<std::collections::HashMap<(u16, u8), CudaSlice<f32>>>,
1611    /// Resident all-ones f32 vector for the UNGATED shared-expert add (`MEMRA_HTOD_DIET`,
1612    /// door H). A family whose plan carries no `ffn_gate_inp_shexp` (GLM-5.3-Flash is the
1613    /// first) makes `moe_shexp_add` take the `g = 1.0` arm, which re-uploaded a freshly
1614    /// allocated `vec![1.0f32; t]` on EVERY MoE layer-call — 42 pageable HtoD per ship round
1615    /// to move a constant on the glm5 serving geometry. Grown to the largest t
1616    /// seen; the buffer may be LONGER than t because `add_scaled_rows_f32` reads only
1617    /// `scale[0..nrows]`.
1618    shexp_ones: Mutex<Option<CudaSlice<f32>>>,
1619}
1620
1621/// Size-keyed device-buffer free-lists for the verify walk (door W — see the field doc on
1622/// [`Engine::verify_ws`]). Exact-length keying: the walk's shapes quantize to a few
1623/// classes per round (t in 2..=8 times fixed widths), so hit rates are structural, and an
1624/// exact-size buffer keeps every `debug_assert_eq!(len, ...)` at the launchers intact.
1625#[derive(Default)]
1626pub struct VerifyWs {
1627    f32_pool: std::collections::HashMap<usize, Vec<CudaSlice<f32>>>,
1628    i8_pool: std::collections::HashMap<usize, Vec<CudaSlice<i8>>>,
1629    u64_pool: std::collections::HashMap<usize, Vec<CudaSlice<u64>>>,
1630    held_bytes: usize,
1631}
1632
1633/// Per-size-class retention cap: enough for every live shape class of one round plus the
1634/// stash generation, small enough that a shape drift cannot hoard VRAM.
1635const VWS_PER_CLASS_CAP: usize = 16;
1636/// Total retention cap (bytes). The round's recurring buffers are t*8192-f32-class and MoE
1637/// staging (<= ~1 MiB each); 256 MiB holds every class with an order of magnitude of slack.
1638const VWS_HELD_BYTES_CAP: usize = 256 << 20;
1639
1640impl VerifyWs {
1641    fn take<T>(
1642        pool: &mut std::collections::HashMap<usize, Vec<CudaSlice<T>>>,
1643        held: &mut usize,
1644        n: usize,
1645    ) -> Option<CudaSlice<T>> {
1646        let s = pool.get_mut(&n)?.pop()?;
1647        *held -= n * std::mem::size_of::<T>();
1648        Some(s)
1649    }
1650    fn put<T>(
1651        pool: &mut std::collections::HashMap<usize, Vec<CudaSlice<T>>>,
1652        held: &mut usize,
1653        s: CudaSlice<T>,
1654    ) {
1655        let n = s.len();
1656        let bytes = n * std::mem::size_of::<T>();
1657        if *held + bytes > VWS_HELD_BYTES_CAP {
1658            return; // drop: falls to the ordinary async free
1659        }
1660        let v = pool.entry(n).or_default();
1661        if v.len() >= VWS_PER_CLASS_CAP {
1662            return;
1663        }
1664        v.push(s);
1665        *held += bytes;
1666    }
1667}
1668
1669/// Device-scratch allocation census (lane/glm5-decode-diet): bumped by every `alloc_uninit`
1670/// and `zeros` call — the class the launch-diet census measured at 2,358
1671/// `cuMemAllocAsync+Free` calls/token. The decode-workspace gate reads deltas per step; the
1672/// cost axis is the CALL COUNT (the box's measured ~1.06 us/driver call), which is exactly
1673/// what this counts. Relaxed atomic: one increment per allocation, noise-level.
1674pub static SCRATCH_ALLOC_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1675
1676/// FAVENDOR lane env gate (2026-07-08): MEMRA_FA_V2=1 dispatches the llama-fattn-vec-mechanism
1677/// decode kernels (fa_decode_vec_q_v2 / fa_decode_vec_q_rows_v2 / fa_decode_vec_q_v2_dc):
1678/// tile-batched online softmax (one alpha rescale per 32-key tile instead of per key) + wide-load
1679/// block dequant in the staging phase. NOTE rev2: llama's register streaming (no smem) was ALSO
1680/// tried and measured 2x WORSE at depth in our gqa-warps frame — the smem KV-tile broadcast stays
1681/// (see the kernel comment). NEW NUMERIC CONFIG (tile-level softmax regrouping changes FP order vs
1682/// the per-key twins) — own argmax baseline; eager decode, the spec-verify rows path AND the
1683/// graph _dc path switch TOGETHER (the spec-exactness law). Default OFF. Read per call (not
1684/// OnceLock) so the gate battery can A/B within one process, matching the MEMRA_NO_FA_VEC pattern.
1685fn fa_v2_on() -> bool {
1686    // DEFAULT ON since 2026-07-08 (MEMRA_FA_V2=0 reverts): tile-batched online softmax, e2e
1687    // measured across every model x depth — 35B 168.7->173.4 (d512) / 153.1->158.5 (d6257),
1688    // 9B 131.2->132.7 / 108.4->124.5 (+15% — the engine-wide depth-slope fix), 27B 47.2->47.7 /
1689    // 42.2->44.9. One-time numeric-config change; kernel-check + argmax + spec self-consistency
1690    // + graph bit-identity green on all three models.
1691    std::env::var("MEMRA_FA_V2")
1692        .map(|v| v != "0")
1693        .unwrap_or(true)
1694}
1695
1696/// FA v3 gate (default ON since 2026-07-09; MEMRA_FA_V3=0 reverts to v2 — research/fa/fa_v3_design.md):
1697/// HYBRID decode twins (fa_decode_vec_q_v3 / _rows_v3 / _v3_dc): llama's int8-dp4a K.Q with
1698/// register-quantized Q (no K dequant, no K smem) + OUR CTA-shared staged bf16 V tile + OUR
1699/// split partition/combine. NEW NUMERIC CONFIG (int8 Q quantization changes the K.Q accumulation
1700/// vs the bf16-roundtrip FMA chain) — own argmax baseline; eager decode, the spec-verify rows
1701/// path AND the graph _dc path switch TOGETHER (the spec-exactness law). Read per call so the
1702/// gate battery can A/B within one process (the MEMRA_FA_V2 pattern).
1703/// `MEMRA_FA_PART_ZERO=1`: zero every freshly grown fa partial bank. DEFAULT OFF,
1704/// diagnostic only. See `fa_part_alloc` for what it discriminates and why it is not a fix.
1705pub(crate) fn fa_part_zero_on() -> bool {
1706    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1707    *ON.get_or_init(|| std::env::var("MEMRA_FA_PART_ZERO").as_deref() == Ok("1"))
1708}
1709
1710pub(crate) fn fa_v3_on() -> bool {
1711    // DEFAULT ON since 2026-07-09 (MEMRA_FA_V3=0 reverts to v2): dp4a-K hybrid FA decode —
1712    // fa kernel -21-23% at depth (micro), 35B spec p3 +5% (190->200, the last spec cell),
1713    // d6257 +1.7%. Own numeric config; full battery green on 35B+9B incl graph bit-identity.
1714    std::env::var("MEMRA_FA_V3")
1715        .map(|v| v != "0")
1716        .unwrap_or(true)
1717}
1718
1719/// The v3 dp4a K path reads RAW q8_0 bytes (34B blocks) and stages q5_1 V verbatim — it is only
1720/// correct on the DEFAULT KV formats — and needs dpl % 4 == 0 consecutive quants per lane
1721/// (head_dim % 128 == 0; both daily models are hd256). All three dispatch sites share this
1722/// predicate so the twins can never diverge.
1723fn fa_v4_mode() -> &'static str {
1724    static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
1725    M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
1726}
1727fn fa_v4_on() -> bool {
1728    fa_v4_mode() != "0"
1729} // DEFAULT ON 2026-07-10 (MEMRA_FA_V4=0 rollback)
1730/// t_kv-conditional v4 pick (gemma depth lesson 2026-07-10: v4's key-per-lane pipeline starves
1731/// at the 1024-window with short splits — MEMRA_FA_V4=0 measured depth plain 158.0 vs 156.7).
1732/// Threshold MEMRA_FA_V4_MAX (default usize::MAX = unchanged behavior; gemma sets 1024 at load
1733/// via FA_V4_MAX_DEFAULT). Applied at EVERY dispatch site (eager, rows, rows_w, dc) so verify
1734/// stays kernel-family-identical to decode at the same t_kv.
1735/// Per-model deep-ctx smem floor default (MEMRA_FA_SMEM_TKV env overrides): gemma pushes it
1736/// above the 1024 window so the windowed decode + verify rows share the REGISTER family.
1737pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
1738    std::sync::atomic::AtomicUsize::new(1024);
1739pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
1740    std::sync::atomic::AtomicUsize::new(usize::MAX);
1741pub fn fa_v4_at_pub(t_kv: usize) -> bool {
1742    fa_v4_at(t_kv)
1743}
1744fn fa_v4_at(t_kv: usize) -> bool {
1745    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
1746    let mx = *M.get_or_init(|| {
1747        std::env::var("MEMRA_FA_V4_MAX")
1748            .ok()
1749            .and_then(|v| v.parse().ok())
1750            .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1751    });
1752    fa_v4_on() && t_kv < mx
1753}
1754/// FA-DEEP gate (2026-08-02, lane fa-decode-deep): deep-ctx v4 twins
1755/// (fa_decode_vec_q_v4_deep / _deep_dc) — the depth-decode lane's priced fix. Unlike
1756/// v2/v3/v4 this is NOT a numeric config: the deep twins run the v4 program VERBATIM
1757/// (same split partition, same softmax/accumulation order, same partials/combine) and only
1758/// move the smem physical layout (bank de-conflict row pads) + the load schedule (next-tile
1759/// L2 prefetch) — kernel-check pins bitdiff==0 vs the v4 twins across depths, so eager /
1760/// rows-verify / graph / seqs stay mutually bit-identical wherever the threshold falls.
1761/// Engages at t_kv >= MEMRA_FA_DEEP_MIN. The swept floor is 0 = ALWAYS ON where v4 ran
1762/// (fa-deep-bench fine grid 96..6144, 2026-08-02: deep flat-or-better at EVERY depth,
1763/// 1.01-1.26x, no losing cell — so there is no engagement boundary and no new
1764/// capture-recapture edge; the env stays as a sweep/diagnostic seam only).
1765/// MEMRA_FA_DEEP=0 is the rollback seam. Read per call so the battery + bench can A/B
1766/// within one process (the v2/v3 pattern).
1767pub const FA_DEEP_MIN_DEFAULT: usize = 0;
1768fn fa_deep_at(t_kv: usize) -> bool {
1769    if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
1770        return false;
1771    }
1772    let min = std::env::var("MEMRA_FA_DEEP_MIN")
1773        .ok()
1774        .and_then(|v| v.parse().ok())
1775        .unwrap_or(FA_DEEP_MIN_DEFAULT);
1776    t_kv >= min
1777}
1778/// Public twin (kernel-check builds the deep-vs-v4 bit pin; bench sweeps the floor).
1779pub fn fa_deep_at_pub(t_kv: usize) -> bool {
1780    fa_deep_at(t_kv)
1781}
1782
1783fn fa_v3_active(head_dim: usize) -> bool {
1784    // v3's dp4a-K walk reads raw q8_0 K bytes — no e4m3 arm; the fp8-KV arm (MEMRA_KV_FP8)
1785    // must fall back like any non-default KV format (the rows_dc stream path asserts on it).
1786    fa_v3_on()
1787        && head_dim.is_multiple_of(128)
1788        && kv_cache_formats() == ("q8_0", "q5_1")
1789        && !Engine::kv_fp8_on()
1790}
1791
1792/// BATCHED-TICK increment 2 (2026-08-01): true iff a row at this t_kv would take the v4
1793/// eager arm in `fa_decode_kvmod`'s dispatch — the exact precondition for the z-batched
1794/// `fa_decode_vec_q_seqs_v4` twin to reproduce its per-seq program bit-identically.
1795/// Mirrors the kvmod predicates: vec on + above the vec floor + hd256 + inside the v4
1796/// window + the PRODUCTION v4 body (the noB3/stage phase probes are wrong-output) + the
1797/// default flash module (no fp8-KV g-module). Callers must ALSO group rows on one
1798/// `fa_split_keys` rung (the rows-twins' straddle law) before batching.
1799pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1800    std::env::var("MEMRA_NO_FA_VEC").is_err()
1801        && t_kv >= fa_vec_min_tkv()
1802        && head_dim == 256
1803        && fa_v4_at(t_kv)
1804        && !matches!(fa_v4_mode(), "noB3" | "stage")
1805        && !Engine::kv_fp8_on()
1806}
1807/// Public twin of the crate-private split ladder (kernel-check builds the seqs-vs-loop pin).
1808pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1809    fa_split_keys(t_kv, n_head_kv)
1810}
1811
1812/// A raw pinned (page-locked, CACHEABLE — flags=0, not write-combined) host allocation for
1813/// DtoH staging. cudarc's `alloc_pinned` uses CU_MEMHOSTALLOC_WRITECOMBINED, which is right for
1814/// HtoD streams but pathologically slow for host READS — the router readback is host-read-heavy,
1815/// so we allocate through `result::malloc_host` with flags=0 directly.
1816struct PinnedStage {
1817    ptr: *mut u8,
1818    cap: usize,
1819}
1820unsafe impl Send for PinnedStage {}
1821impl PinnedStage {
1822    fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1823        let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1824        Ok(PinnedStage { ptr, cap })
1825    }
1826}
1827impl Drop for PinnedStage {
1828    fn drop(&mut self) {
1829        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1830    }
1831}
1832
1833/// Owned page-locked CACHEABLE host buffer (flags=0, deliberately NOT write-combined) for the
1834/// prefix-cache host tier (lane/kv-host-spill-20260830). Same allocation class as `PinnedStage`
1835/// above and for the same reason: `ctx().alloc_pinned` is CU_MEMHOSTALLOC_WRITECOMBINED, which
1836/// is right for H2D-only staging but pathologically slow for host READS (see the HostBuf CAVEAT
1837/// in model.rs), and these bytes are CPU-read by the MEMRA_KV_HOST_VERIFY digest arm. Public
1838/// because the server's host-tier cache owns these buffers across requests.
1839pub struct PinnedHostBuf {
1840    ptr: *mut u8,
1841    len: usize,
1842}
1843// Safety: the allocation is process-wide page-locked host memory; the raw pointer is owned by
1844// this struct alone and freed exactly once in Drop (identical justification to PinnedStage).
1845unsafe impl Send for PinnedHostBuf {}
1846impl PinnedHostBuf {
1847    /// Allocate `len` pinned cacheable bytes (a zero-length request still pins one byte so the
1848    /// pointer stays valid, mirroring the device planes' `alloc_u8(kb.max(1))` convention).
1849    pub fn new(len: usize) -> Result<Self, Box<dyn std::error::Error>> {
1850        let ptr = unsafe { cudarc::driver::result::malloc_host(len.max(1), 0)? } as *mut u8;
1851        Ok(PinnedHostBuf { ptr, len })
1852    }
1853    pub fn len(&self) -> usize {
1854        self.len
1855    }
1856    pub fn is_empty(&self) -> bool {
1857        self.len == 0
1858    }
1859    pub fn as_slice(&self) -> &[u8] {
1860        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
1861    }
1862    pub fn as_mut_slice(&mut self) -> &mut [u8] {
1863        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
1864    }
1865}
1866impl Drop for PinnedHostBuf {
1867    fn drop(&mut self) {
1868        let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1869    }
1870}
1871
1872/// Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks
1873/// x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.
1874pub const ARGMAX_NB: usize = 256;
1875
1876/// crate-visible alias for the batched FA3 shim entry (hybrid_forward's batch arm).
1877pub(crate) use memra_fa3_vl as fa3_vl_raw;
1878
1879unsafe extern "C" {
1880    /// FA3 v10 shim (cu/fa3_prefill.cu): TMA-swizzled wgmma FA, fresh causal hd256.
1881    fn memra_fa3_prefill(
1882        q16: *const core::ffi::c_void,
1883        k16: *const core::ffi::c_void,
1884        v16: *const core::ffi::c_void,
1885        o: *mut f32,
1886        t: i32,
1887        h: i32,
1888        hkv: i32,
1889        d: i32,
1890        scale: f32,
1891        stream: *mut core::ffi::c_void,
1892    ) -> i32;
1893    /// batched varlen twin: host arrays of device pointers per seq (B <= 8).
1894    pub(crate) fn memra_fa3_vl(
1895        q16s: *const *const core::ffi::c_void,
1896        k16s: *const *const core::ffi::c_void,
1897        v16s: *const *const core::ffi::c_void,
1898        os: *const *mut f32,
1899        ts: *const i32,
1900        b: i32,
1901        h: i32,
1902        hkv: i32,
1903        d: i32,
1904        scale: f32,
1905        stream: *mut core::ffi::c_void,
1906    ) -> i32;
1907}
1908
1909/// STAGE-2 GROUPED DECODE: 8 expert weight-block device pointers passed BY VALUE as one kernel
1910/// param (matches the CUDA `wptr8_t` struct: 8x 64-bit pointers, `#[repr(C)]` => identical
1911/// layout). The pointers are SLRU cache-slot base addresses — fixed for the engine's lifetime
1912/// (slots are never re-allocated), so passing raw values is stable across the launch.
1913#[repr(C)]
1914#[derive(Clone, Copy)]
1915pub struct WPtr8(pub [u64; 8]);
1916unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1917
1918/// task #18 varlen GDN: per-seq args for gdn_chunk_{state,output}_mma_vl — one launch
1919/// runs all B<=8 sequences' K4/K5 (CUDA `gdnseq_t`/`gdnvl_t`, layout-identical repr(C)).
1920/// Raw addresses are valid for the launch: every referenced buffer outlives the call and
1921/// all work is on the single compute stream (same discipline as the f16 GEMM FFI).
1922#[repr(C)]
1923#[derive(Clone, Copy, Default)]
1924pub struct GdnSeqVl {
1925    pub kb16: u64,
1926    pub gcum: u64,
1927    pub beta: u64,
1928    pub u: u64,
1929    pub wb16: u64,
1930    pub y: u64,
1931    pub ssnap: u64,
1932    pub state_in: u64,
1933    pub state_out: u64,
1934    pub q: u64,
1935    pub p: u64,
1936    pub o: u64,
1937    pub k: u64,
1938    pub v: u64,
1939    pub g: u64,
1940    pub a: u64,
1941    pub w: u64,
1942    pub t: i32,
1943    pub nc: i32,
1944}
1945unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1946#[repr(C)]
1947#[derive(Clone, Copy)]
1948pub struct GdnVl8(pub [GdnSeqVl; 8]);
1949unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1950
1951/// task #22: per-seq wgmma-fused extras (CUDA `gdnw_t`/`gdnwvl_t`) — qb16 mirror +
1952/// pre-masked Pb16, riding NEXT TO GdnSeqVl so the base struct stays untouched.
1953#[repr(C)]
1954#[derive(Clone, Copy, Default)]
1955pub struct GdnWVl {
1956    pub qb16: u64,
1957    pub pb16: u64,
1958}
1959unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1960#[repr(C)]
1961#[derive(Clone, Copy)]
1962pub struct GdnWVl8(pub [GdnWVl; 8]);
1963unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1964
1965/// task #18 increment 3: per-seq PREP/TAIL args (CUDA `gdnprep_t`/`gdnprepvl_t`).
1966#[repr(C)]
1967#[derive(Clone, Copy, Default)]
1968pub struct GdnPrepVl {
1969    pub qkv: u64,
1970    pub conv_state: u64,
1971    pub conv_out: u64,
1972    pub q_g: u64,
1973    pub k_g: u64,
1974    pub v_g: u64,
1975    pub q_l2: u64,
1976    pub k_l2: u64,
1977    pub beta_raw: u64,
1978    pub alpha: u64,
1979    pub beta: u64,
1980    pub g_log: u64,
1981    pub o: u64,
1982    pub z: u64,
1983    pub gn: u64,
1984    pub gn16: u64,
1985    pub kb16: u64,
1986    pub qb16: u64,
1987    pub t: i32,
1988    pub pad: i32,
1989}
1990unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1991#[repr(C)]
1992#[derive(Clone, Copy)]
1993pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1994unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1995
1996/// task #18 (attn side): per-seq varlen FA args (CUDA `faseq_t`/`favl_t`).
1997#[repr(C)]
1998#[derive(Clone, Copy, Default)]
1999pub struct FaSeqVl {
2000    pub q: u64,
2001    pub k16: u64,
2002    pub v16: u64,
2003    pub o: u64,
2004    pub kf: u64,
2005    pub vf: u64,
2006    pub t: i32,
2007    pub pad: i32,
2008}
2009unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
2010#[repr(C)]
2011#[derive(Clone, Copy)]
2012pub struct FaVl8(pub [FaSeqVl; 8]);
2013unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
2014
2015/// task #18 (attn pre-FA): per-seq split/norm/rope/append args (CUDA `attnpre_t`).
2016#[repr(C)]
2017#[derive(Clone, Copy, Default)]
2018pub struct AttnPreVl {
2019    pub qf: u64,
2020    pub kf: u64,
2021    pub vf: u64,
2022    pub q: u64,
2023    pub gate: u64,
2024    pub qn: u64,
2025    pub kn: u64,
2026    pub kc: u64,
2027    pub vc: u64,
2028    pub t: i32,
2029    pub pad: i32,
2030}
2031unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
2032#[repr(C)]
2033#[derive(Clone, Copy)]
2034pub struct AttnPreVl8(pub [AttnPreVl; 8]);
2035unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
2036
2037/// task #18 increment 2: one sequence's FULL chunk-buffer set (alloc-only; the
2038/// varlen K1-K5 chain fills them).
2039pub struct GdnChunkBufs {
2040    pub gcum: CudaSlice<f32>,
2041    pub a: CudaSlice<f32>,
2042    pub p: CudaSlice<f32>,
2043    pub u: CudaSlice<f32>,
2044    pub w: CudaSlice<f32>,
2045    pub kb16: CudaSlice<u8>,
2046    pub wb16: CudaSlice<u8>,
2047    pub y16: CudaSlice<u8>,
2048    pub ssnap16: CudaSlice<u8>,
2049    pub qb16: CudaSlice<u8>,
2050    pub pb16: CudaSlice<u8>,
2051    pub o: CudaSlice<f32>,
2052    pub t: usize,
2053    pub nc: usize,
2054}
2055
2056/// STAGE-2 GROUPED DECODE: the 8 routed-expert weights by value (CUDA `f32x8_t`).
2057#[repr(C)]
2058#[derive(Clone, Copy)]
2059pub struct F32x8(pub [f32; 8]);
2060unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
2061
2062/// Harness timing contract: wall nanos of the LAST generate/generate_spec prompt prime on this
2063/// process. Bench binaries read it right after the call to print gen-only throughput without the
2064/// prime-subtraction hack (which amplifies prime jitter into the gen number at long prompts).
2065pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2066
2067/// Fused MoE-epilogue dispatches taken since process start (`MEMRA_MOE_FUSED_EPI`), incremented
2068/// once per (token, layer) that actually runs `moe_fused_epi_token_q8`.
2069///
2070/// This exists because the arm cannot be observed any other way: setting `MEMRA_MOE_STATS` /
2071/// `MEMRA_MOE_TRACE` / `MEMRA_MOE_WEIGHT_TRACE` / `MEMRA_MOE_INPUT_TRACE_DIR` sets
2072/// `observe_routes` in `moe_ffn_inner`, which DIVERTS dispatch to the host-routed path — so a
2073/// gate that tried to prove the fused arm ran by tracing would prove it about a different
2074/// program. Read it via [`moe_fused_epilogue_dispatches`] around a workload.
2075pub static MOE_FUSED_EPI_DISPATCHES: std::sync::atomic::AtomicU64 =
2076    std::sync::atomic::AtomicU64::new(0);
2077
2078/// Snapshot of [`MOE_FUSED_EPI_DISPATCHES`]. Gates take a before/after pair around a workload and
2079/// assert on the delta, anchoring on the arm's own invocation rather than on a flag being set
2080/// (LAW:wiring-assertions-match-prose).
2081pub fn moe_fused_epilogue_dispatches() -> u64 {
2082    MOE_FUSED_EPI_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2083}
2084
2085/// Verify-rows batched MoE dispatches taken since process start (lane/glm5-vrest): incremented
2086/// once per (layer, verify-call) that runs the pairs-shaped routed-expert program
2087/// (`moe_gate_up_preclamp8_q8_rows` + `moe_down8_fma_q8_rows`) instead of the per-(token,expert)
2088/// sequential loop. Rides `MEMRA_GLM5_VERIFY_BATCH`'s arm — no flag of its own. Same rationale
2089/// as [`MOE_FUSED_EPI_DISPATCHES`]: the observation envs divert dispatch, so gates anchor on the
2090/// arm's own invocation (LAW:wiring-assertions-match-prose).
2091pub static MOE_VROWS_DISPATCHES: std::sync::atomic::AtomicU64 =
2092    std::sync::atomic::AtomicU64::new(0);
2093
2094/// Snapshot of [`MOE_VROWS_DISPATCHES`] — gates take a before/after delta around a workload.
2095pub fn moe_vrows_dispatches() -> u64 {
2096    MOE_VROWS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2097}
2098
2099/// `MEMRA_BF16_TCOLS_WIDE` (lane/glm5-matvec door T, default ON since the 2026-08-31 mv-battery
2100/// flip; `=0` is the rollback seam): FloatBf16 rows calls at
2101/// t=2..=16 ride the weight-once t-column twins (`matvec_bf16_f32acc_x4_tcols` for t<=8, the
2102/// NEW `..._tcols16` for 9..=16) instead of the grid.y=t weight-rereading `_rows` kernel. The
2103/// motivating call is the DFlash2 drafter's block head: `eh.matmul(head, rows, 15)` re-read
2104/// the 1.269 GB lm head 15x per spec round (diet-battery c8-ship census, 5.31 ms/round —
2105/// 13% of decode GPU). Bit-identical per (row, token) by the tcols class's standing
2106/// construction; gated by `glm5_matvec_doors_gpu`. Read per call — the rollback seam.
2107fn bf16_tcols_wide_on() -> bool {
2108    std::env::var("MEMRA_BF16_TCOLS_WIDE").as_deref() != Ok("0")
2109}
2110
2111/// Engagement counter for the wide-t tcols door (`MEMRA_BF16_TCOLS_WIDE`), incremented at the
2112/// door's own dispatch (LAW:wiring-assertions-match-prose). Read via
2113/// [`bf16_tcols_wide_dispatches`].
2114pub static BF16_TCOLS_WIDE_DISPATCHES: std::sync::atomic::AtomicU64 =
2115    std::sync::atomic::AtomicU64::new(0);
2116
2117/// Snapshot of [`BF16_TCOLS_WIDE_DISPATCHES`] — gates take a before/after delta.
2118pub fn bf16_tcols_wide_dispatches() -> u64 {
2119    BF16_TCOLS_WIDE_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2120}
2121
2122/// `MEMRA_BF16_TCOLS_X1` (lane/glm5-matvec door X, default ON since the 2026-08-31 mv-battery
2123/// flip; `=0` is the rollback seam): the tcols dispatch takes
2124/// the one-row-per-block grid twin (`matvec_bf16_f32acc_x1_tcols`, grid.x = out_f) instead of
2125/// the 4-rows-per-block form. WHY: the trunk kda shapes (out_f 4096/8192) launch 1024/2048
2126/// blocks — ~one resident wave, and the census pins them at 1.05 TB/s (59% of peak) while the
2127/// SAME kernel at the lm head's 38720-block grid runs 1.43 TB/s (80%). Per-row program and
2128/// tree verbatim — bit-identical. Gated by `glm5_matvec_doors_gpu`. Read per call.
2129fn bf16_tcols_x1_on() -> bool {
2130    std::env::var("MEMRA_BF16_TCOLS_X1").as_deref() != Ok("0")
2131}
2132
2133/// Engagement counter for the x1-grid tcols door (`MEMRA_BF16_TCOLS_X1`).
2134pub static BF16_TCOLS_X1_DISPATCHES: std::sync::atomic::AtomicU64 =
2135    std::sync::atomic::AtomicU64::new(0);
2136
2137/// Snapshot of [`BF16_TCOLS_X1_DISPATCHES`] — gates take a before/after delta.
2138pub fn bf16_tcols_x1_dispatches() -> u64 {
2139    BF16_TCOLS_X1_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2140}
2141
2142/// `MEMRA_BF16_TCOLS_RED_FUSED=1` (lane/glm5-door-r door R, default OFF): the tcols
2143/// dispatches take the `_rf` fused-reduce-tail twins (`matvec_bf16_f32acc_x1_tcols_rf` /
2144/// `..._x4_tcols_rf` / `..._x4_tcols16_rf`). WHY (moe-loc LANE.md §2.2): after door X the
2145/// kda trunk's tcols calls sit at 67.0% of peak because the reduce tail runs t SEPARATE
2146/// strided trees — ~30 block-wide barriers at t=3.34 (135 at the drafter head's t=15)
2147/// against a 4-iteration main loop; the kernel is barrier/tail-bound. The twins share ONE
2148/// barrier sequence across the t columns (`red[t*blockDim]`, dynamic shared) and run levels
2149/// s<=16 as a `__shfl_down_sync` chain at the IDENTICAL pairing and operand order — 9t -> 3
2150/// barriers per block, bit-identical by pairing preservation (gated with a shifted-pairing
2151/// red in `glm5_matvec_doors_gpu`). Engages only when `MEMRA_MMV_BLOCK` is a power of two
2152/// (the fused tail's block-wide loop must pass exactly through s=32; the default 128 is).
2153/// Read per call — unset or `=0` is byte-for-byte the standing tcols program.
2154fn bf16_tcols_red_fused_on() -> bool {
2155    std::env::var("MEMRA_BF16_TCOLS_RED_FUSED").as_deref() == Ok("1")
2156}
2157
2158/// Engagement counter for the fused-reduce-tail tcols door (`MEMRA_BF16_TCOLS_RED_FUSED`),
2159/// incremented at the door's own dispatch (LAW:wiring-assertions-match-prose).
2160pub static BF16_TCOLS_RED_FUSED_DISPATCHES: std::sync::atomic::AtomicU64 =
2161    std::sync::atomic::AtomicU64::new(0);
2162
2163/// Snapshot of [`BF16_TCOLS_RED_FUSED_DISPATCHES`] — gates take a before/after delta.
2164pub fn bf16_tcols_red_fused_dispatches() -> u64 {
2165    BF16_TCOLS_RED_FUSED_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2166}
2167
2168/// `MEMRA_MOE_VROWS_PACK=1` (lane/glm5-matvec door M, default OFF): the verify-rows MoE pair
2169/// launches its `_w4` warp-packed twins — MEMRA_MMVQ_ROWS = 4 warps per block on threadIdx.y
2170/// (the qmatvec mmvq family's standing shape) instead of one warp per block. The unpacked
2171/// launch caps residency at the blocks/SM limit (<=67% of warp slots) and schedules ~65k
2172/// one-warp blocks per launch; per-warp body verbatim, bit-identical per (row, pair). Gated
2173/// by `glm5_matvec_doors_gpu`. Read per call.
2174pub(crate) fn moe_vrows_pack_on() -> bool {
2175    std::env::var("MEMRA_MOE_VROWS_PACK").as_deref() == Ok("1")
2176}
2177
2178/// Engagement counter for the warp-packed verify-rows MoE door (`MEMRA_MOE_VROWS_PACK`).
2179pub static MOE_VROWS_PACK_DISPATCHES: std::sync::atomic::AtomicU64 =
2180    std::sync::atomic::AtomicU64::new(0);
2181
2182/// Snapshot of [`MOE_VROWS_PACK_DISPATCHES`] — gates take a before/after delta.
2183pub fn moe_vrows_pack_dispatches() -> u64 {
2184    MOE_VROWS_PACK_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2185}
2186
2187/// `MEMRA_MOE_VROWS_ILP` (lane/glm5-moe-rows-ilp-20260904; default ON on sm_100a builds since
2188/// 2026-09-04, OFF elsewhere, `=0`/`=1` override): the verify-rows MoE pair launches its `_ilp` twins, the same per-warp program with the loads of four (then two)
2189/// groups per lane issued ahead of their math. Composes with door M (`_w4_ilp`). ONLY the
2190/// interleaved NVFP4 expert layout (`QT_NVFP4`): the launchers refuse by name for any other
2191/// qtype and keep the shipped kernel. Read PER CALL. Why and receipts: the kernel header in
2192/// cu/qmatvec.cu and docs/FLAGS.md.
2193pub(crate) fn moe_vrows_ilp_on() -> bool {
2194    moe_vrows_ilp_on_from(
2195        std::env::var("MEMRA_MOE_VROWS_ILP").ok().as_deref(),
2196        env!("MEMRA_BUILT_CUDA_ARCH"),
2197    )
2198}
2199
2200/// The pure parse behind [`moe_vrows_ilp_on`]: `1` arms, `0` disarms, unset follows the BUILD
2201/// ARCH (ON for `100a`, OFF otherwise), the per-hardware arm selection law: the twins carry a
2202/// 2x B200 receipt (+6.0% at c1, darklanes research/glm5-b200-20260902/LANE.md, ilpab) and no
2203/// SM120 one, so an sm_120a build keeps its measured default until it has its own.
2204pub fn moe_vrows_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
2205    match v.map(str::trim) {
2206        Some("1") => true,
2207        Some("0") => false,
2208        _ => built_arch == "100a",
2209    }
2210}
2211
2212/// Engagement counter for the ILP verify-rows MoE door (`MEMRA_MOE_VROWS_ILP`), both launches.
2213pub static MOE_VROWS_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
2214    std::sync::atomic::AtomicU64::new(0);
2215
2216/// Snapshot of [`MOE_VROWS_ILP_DISPATCHES`] — gates take a before/after delta.
2217pub fn moe_vrows_ilp_dispatches() -> u64 {
2218    MOE_VROWS_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2219}
2220
2221/// One line, once, when the ILP door is armed but the expert layout is not the interleaved
2222/// NVFP4 it was written for: a silent fall-through would let a box A/B read the door's absence.
2223fn moe_vrows_ilp_refuse(site: &str, qt: i32) {
2224    static SAID: std::sync::Once = std::sync::Once::new();
2225    SAID.call_once(|| {
2226        eprintln!(
2227            "[moe-vrows-ilp] REFUSED at {site}: MEMRA_MOE_VROWS_ILP is on but the expert qtype is \
2228             {qt}, neither QT_NVFP4 ({QT_NVFP4}) nor QT_NVFP4_V2 ({QT_NVFP4_V2}); the shipped \
2229             kernel runs (the ILP twins hoist the NVFP4 group loads, interleaved or slot-major, \
2230             and no other layout)"
2231        );
2232    });
2233}
2234
2235/// `MEMRA_MOE_VROWS_DEV_TABLES=1` (lane/glm5-moe-loc door D, default OFF): the verify-rows MoE
2236/// pair builds its `ptrs`/`scl` tables ON DEVICE from the router's own `sel`/`w` device output
2237/// (`moe_vrows_tables_from_sel`) instead of on the host, and the layer routes through the
2238/// readback-free `moe_router_sigmoid_topk` rather than `..._host`. WHY: the host table build is
2239/// the ONLY consumer of the selection on the serving shape, and it costs a full
2240/// `cuStreamSynchronize` + 2 DtoH + 2 pageable HtoD + 2 host Vec allocations per MoE layer-call
2241/// = 42 device-wide drains + 84 DtoH + 84 HtoD per ship round. Bit-identical: same integer
2242/// `base + ex*stride`, same macro-plane lookups, same single `w * macro_down` product. Read per
2243/// call; fails closed to the host path whenever any host-visible route consumer is armed.
2244pub(crate) fn moe_vrows_dev_tables_on() -> bool {
2245    std::env::var("MEMRA_MOE_VROWS_DEV_TABLES").as_deref() == Ok("1")
2246}
2247
2248/// Engagement counter for the device-side vrows table build (`MEMRA_MOE_VROWS_DEV_TABLES`).
2249pub static MOE_VROWS_DEV_TABLES_DISPATCHES: std::sync::atomic::AtomicU64 =
2250    std::sync::atomic::AtomicU64::new(0);
2251
2252/// Snapshot of [`MOE_VROWS_DEV_TABLES_DISPATCHES`] — gates take a before/after delta.
2253pub fn moe_vrows_dev_tables_dispatches() -> u64 {
2254    MOE_VROWS_DEV_TABLES_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2255}
2256
2257/// Router readbacks (one full `cuStreamSynchronize` + 2 DtoH each) that door D skipped. The
2258/// count receipt for the host seam: a gate asserts it moves 1:1 with
2259/// [`MOE_VROWS_DEV_TABLES_DISPATCHES`] on the ON arm and stays flat on the OFF arm.
2260pub static MOE_VROWS_ROUTER_SYNCS_AVOIDED: std::sync::atomic::AtomicU64 =
2261    std::sync::atomic::AtomicU64::new(0);
2262
2263/// Snapshot of [`MOE_VROWS_ROUTER_SYNCS_AVOIDED`].
2264pub fn moe_vrows_router_syncs_avoided() -> u64 {
2265    MOE_VROWS_ROUTER_SYNCS_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2266}
2267
2268/// `MEMRA_MOE_VROWS_DEDUP_STAT=1` (lane/glm5-moe-loc, default OFF — a MEASUREMENT instrument,
2269/// not a serving door): on the host table-build arm, count the pair union's expert VISITS and
2270/// DISTINCT experts per layer-call into [`MOE_VROWS_PAIR_VISITS`] /
2271/// [`MOE_VROWS_PAIR_DISTINCT`]. WHY IT EXISTS: the pair runs at ~90% of this card class's
2272/// theoretical DRAM peak (moe-loc LANE.md §1), so cross-row expert-slab dedup is the ONLY
2273/// remaining byte lever, and its size is exactly `1 - distinct/visits` — an unmeasured routing
2274/// property whose independent-routing bound is 3.2% but whose structural ceiling is 70%. This
2275/// counter turns a speculative kernel campaign into a priced decision for the cost of a host
2276/// bitset. Requires `MEMRA_MOE_VROWS_DEV_TABLES=0` (door D removes the host selection).
2277fn moe_vrows_dedup_stat_on() -> bool {
2278    std::env::var("MEMRA_MOE_VROWS_DEDUP_STAT").as_deref() == Ok("1")
2279}
2280
2281/// Expert VISITS (t x n_used) summed over vrows layer-calls under `MEMRA_MOE_VROWS_DEDUP_STAT`.
2282pub static MOE_VROWS_PAIR_VISITS: std::sync::atomic::AtomicU64 =
2283    std::sync::atomic::AtomicU64::new(0);
2284
2285/// DISTINCT experts in the pair union, summed over the same layer-calls. The dedup lever is
2286/// `1 - distinct/visits`; equal counters mean routing is disjoint across the verify rows and
2287/// there is no byte to save.
2288pub static MOE_VROWS_PAIR_DISTINCT: std::sync::atomic::AtomicU64 =
2289    std::sync::atomic::AtomicU64::new(0);
2290
2291/// `(visits, distinct)` for one layer-call's pair union — the dedup lever's whole arithmetic.
2292/// `visits` is `t * n_used`, the slab reads the pair performs today; `distinct` is how many of
2293/// them are to a DIFFERENT expert. `1 - distinct/visits` is the share of the pair's 9.86 ms/round
2294/// that a dedup kernel could remove, and nothing else about the pair is removable (it already
2295/// runs at ~90% of theoretical DRAM peak). Split out from the call site so the counting itself is
2296/// unit-testable on planted overlaps rather than inferred from a live routing tape.
2297pub(crate) fn vrows_overlap_counts(sel_all: &[u32]) -> (u64, u64) {
2298    let mut seen = std::collections::HashSet::with_capacity(sel_all.len());
2299    for &ex in sel_all {
2300        seen.insert(ex);
2301    }
2302    (sel_all.len() as u64, seen.len() as u64)
2303}
2304
2305/// vrows layer-calls the dedup instrument has observed — the reporting cadence's clock.
2306static MOE_VROWS_DEDUP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2307
2308/// AN INSTRUMENT HAS TO SPEAK. A box window greps a server log; it cannot read a Rust atomic, so
2309/// the dedup counters emit their cumulative ratio on the first vrows layer-call and every 42
2310/// after (42 = the MoE layer count, i.e. about one line per decode round). The reported
2311/// `repeat` IS the dedup lever's ceiling: the share of the pair's 9.86 ms/round that reading a
2312/// shared expert slab once could remove, and the only removable share that exists (LANE.md §1 —
2313/// the pair already runs at ~90% of theoretical DRAM peak).
2314fn moe_vrows_dedup_report() {
2315    let n = MOE_VROWS_DEDUP_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2316    if n != 0 && !n.is_multiple_of(42) {
2317        return;
2318    }
2319    let (visits, distinct) = moe_vrows_pair_overlap();
2320    if visits == 0 {
2321        return;
2322    }
2323    let repeat = 100.0 * (1.0 - distinct as f64 / visits as f64);
2324    eprintln!(
2325        "[moe-vrows-dedup] layer-calls={} visits={visits} distinct={distinct} \
2326         repeat={repeat:.2}% = the cross-row expert-slab dedup ceiling on the vrows pair \
2327         (MEMRA_MOE_VROWS_DEDUP_STAT=1)",
2328        n + 1
2329    );
2330}
2331
2332/// Gate hook for [`vrows_overlap_counts`] — the counting is the whole instrument, so it is gated
2333/// on planted overlaps (disjoint / partial / identical) rather than inferred from a live tape.
2334pub fn vrows_overlap_counts_for_test(sel_all: &[u32]) -> (u64, u64) {
2335    vrows_overlap_counts(sel_all)
2336}
2337
2338/// Snapshot of the dedup instrument as `(visits, distinct)`.
2339pub fn moe_vrows_pair_overlap() -> (u64, u64) {
2340    (
2341        MOE_VROWS_PAIR_VISITS.load(std::sync::atomic::Ordering::Relaxed),
2342        MOE_VROWS_PAIR_DISTINCT.load(std::sync::atomic::Ordering::Relaxed),
2343    )
2344}
2345
2346/// `MEMRA_MOE_VROWS_DEDUP_ORDER=1` (lane/glm5-dedup door E, default OFF): the verify-rows
2347/// gate/up launch takes the `_ord` twin — grid TRANSPOSED so the pair index is the fastest
2348/// dimension, walking an EXPERT-MAJOR order plane appended to the pointer table. WHY: the
2349/// struct-battery instrument measured a **21.96% repeat fraction** across the pair's expert
2350/// visits (2.55M visits, 6.9x the 3.21% independent-routing bound), and the pair is already at
2351/// 90.2% of theoretical DRAM peak, so the only lever left is not re-reading a slab a sibling
2352/// verify row already read — which requires the repeat visit to be SCHEDULED inside the reuse
2353/// window. Bit-identical by construction: every output is a pure function of its `(o, pr)`
2354/// coordinate and no block communicates, so re-indexing which block computes which output moves
2355/// no bits (`glm5_dedup_sched_gpu`). The WIN is a scheduling property, unpriceable on an
2356/// exactness-only rig — hence default OFF with the box pricing the flip.
2357///
2358/// Refused by name, falling closed to the shipped schedule: door M (`MEMRA_MOE_VROWS_PACK`, the
2359/// refuted 4-warp pack) takes precedence in the launcher, and the door engages only when the
2360/// order plane is actually present (`ptrs.len() >= 4*n_pairs`), so a direct launcher call with a
2361/// 3-plane table keeps the shipped program.
2362pub(crate) fn moe_vrows_dedup_order_on() -> bool {
2363    std::env::var("MEMRA_MOE_VROWS_DEDUP_ORDER").as_deref() == Ok("1")
2364}
2365
2366/// `MEMRA_MOE_VROWS_DOWN_TMAJ=1` (lane/glm5-dedup door E-down, default OFF): the verify-rows down
2367/// launch takes the `_tmaj` twin — grid transposed to `(t, out_f)` so the t verify rows at one
2368/// output row are adjacent blocks and a repeated expert's down row is read once for every token
2369/// sharing it. The down chain's slot-ordered `__fmaf_rn` accumulation is INSIDE the block and is
2370/// untouched (it keeps its original slot order — the vrest gate-4 bit bar); only the grid moves.
2371/// Split from [`moe_vrows_dedup_order_on`] as its own flag so the box can attribute the two
2372/// halves of the lever separately (gate/up is 2/3 of the pair's bytes, down 1/3). Same refusals:
2373/// door M wins, and `out_f > 65535` falls closed (a grid.y bound, not a serving shape).
2374fn moe_vrows_down_tmaj_on() -> bool {
2375    std::env::var("MEMRA_MOE_VROWS_DOWN_TMAJ").as_deref() == Ok("1")
2376}
2377
2378/// Engagement counter for the expert-major gate/up schedule (`MEMRA_MOE_VROWS_DEDUP_ORDER`).
2379pub static MOE_VROWS_DEDUP_ORDER_DISPATCHES: std::sync::atomic::AtomicU64 =
2380    std::sync::atomic::AtomicU64::new(0);
2381
2382/// Snapshot of [`MOE_VROWS_DEDUP_ORDER_DISPATCHES`] — gates take a before/after delta.
2383pub fn moe_vrows_dedup_order_dispatches() -> u64 {
2384    MOE_VROWS_DEDUP_ORDER_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2385}
2386
2387/// Engagement counter for the token-major down schedule (`MEMRA_MOE_VROWS_DOWN_TMAJ`).
2388pub static MOE_VROWS_DOWN_TMAJ_DISPATCHES: std::sync::atomic::AtomicU64 =
2389    std::sync::atomic::AtomicU64::new(0);
2390
2391/// Snapshot of [`MOE_VROWS_DOWN_TMAJ_DISPATCHES`] — gates take a before/after delta.
2392pub fn moe_vrows_down_tmaj_dispatches() -> u64 {
2393    MOE_VROWS_DOWN_TMAJ_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2394}
2395
2396/// AVOIDED SLAB READS — the box receipt for door E. Every layer-call adds `visits - distinct`,
2397/// i.e. the expert-slab reads whose repeat visit the expert-major schedule places inside the
2398/// reuse window. Multiply by the per-visit slab bytes (gate+up 9.4372 MB, down 4.7186 MB at the
2399/// serving geometry) for the bytes the schedule makes avoidable; that product is the CEILING of
2400/// the win, not the win (the realized share is a cache/scheduling property the box prices).
2401///
2402/// HOST-ARM ONLY, by construction: with door D on there is no host-side selection to count and a
2403/// 4-byte readback would reintroduce the very `cuStreamSynchronize` door D removed. The counting
2404/// boot is therefore `MEMRA_MOE_VROWS_DEV_TABLES=0`, exactly like the dedup instrument — while
2405/// [`MOE_VROWS_DEDUP_ORDER_DISPATCHES`] moves in BOTH table arms.
2406pub static MOE_VROWS_SLAB_READS_AVOIDED: std::sync::atomic::AtomicU64 =
2407    std::sync::atomic::AtomicU64::new(0);
2408
2409/// Snapshot of [`MOE_VROWS_SLAB_READS_AVOIDED`].
2410pub fn moe_vrows_slab_reads_avoided() -> u64 {
2411    MOE_VROWS_SLAB_READS_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2412}
2413
2414/// The EXPERT-MAJOR order plane, host build — the stable sort by `(expert id, pair index)` whose
2415/// bit-for-bit twin is the `moe_vrows_order_from_sel` counting rank. Returned as the `[n_pairs]`
2416/// tail plane the pointer table carries at `[3*n_pairs ..)`, and split out from the call site so
2417/// the device kernel can be gated against it directly.
2418pub(crate) fn vrows_expert_major_order(sel_all: &[u32]) -> Vec<u64> {
2419    let mut ord: Vec<u64> = (0..sel_all.len() as u64).collect();
2420    // Stable by construction: `sort_by_key` on the expert id keeps ascending pair order inside
2421    // each expert's run, so per-token slot order survives within a shared expert.
2422    ord.sort_by_key(|&p| sel_all[p as usize]);
2423    ord
2424}
2425
2426/// Gate hook for [`vrows_expert_major_order`] — the permutation is the whole door, so it is gated
2427/// against the device build and on planted selections rather than inferred from a live tape.
2428pub fn vrows_expert_major_order_for_test(sel_all: &[u32]) -> Vec<u64> {
2429    vrows_expert_major_order(sel_all)
2430}
2431
2432// ---- THE FLAG-ALIAS LAW for boolean doors (lane/glm5-extract2, phase 2) ------------------
2433//
2434// A door extracted from a family name to its general name keeps the FAMILY NAME HONORED:
2435// every banked gate script, box battery and in-flight lane sets the old name today, so
2436// refusing it would break receipts mid-bank for the price of one extra env read. Phase 1
2437// established the pattern for the two default-ON doors it moved (`MEMRA_VERIFY_WS`
2438// OFF-wins; `MEMRA_SPEC_TRACE` general-wins-loudly) and for the one VALUED door
2439// (`MEMRA_EP_MAP`, [`ep_map::resolve_ep_map_env`], which refuses a disagreeing pair at load).
2440// [`alias_door_from`] is the same law for a DEFAULT-OFF BOOLEAN door read PER CALL.
2441
2442/// Pure two-name resolution for a default-OFF boolean door (unit-tested without env
2443/// mutation — the phase-1 co-refusal-test pattern). Returns `(armed, the name the operator
2444/// actually set)` so every downstream refusal names the flag they typed, exactly as
2445/// [`ep_map::resolve_ep_map_env`] does for the valued seam.
2446///
2447/// * either name `=1` arms the door; anything else (including `=0`) is a deliberate pin;
2448/// * both set to the SAME value resolves to the general name;
2449/// * both set to DISAGREEING values is an operator error and is refused — `Err` carries the
2450///   message naming BOTH flags. The CALLER falls closed to the shipped program rather than
2451///   picking a precedence winner.
2452pub(crate) fn alias_door_from(
2453    general: (&'static str, Option<&str>),
2454    alias: (&'static str, Option<&str>),
2455) -> Result<(bool, &'static str), String> {
2456    match (general.1, alias.1) {
2457        (Some(g), Some(a)) if g != a => Err(format!(
2458            "{}={g:?} and {}={a:?} disagree — the alias and the general flag name ONE door \
2459             (unset one); refused rather than silently picking a precedence winner, and the \
2460             door falls closed to the shipped program",
2461            general.0, alias.0
2462        )),
2463        (Some(g), _) => Ok((g == "1", general.0)),
2464        (None, Some(a)) => Ok((a == "1", alias.0)),
2465        (None, None) => Ok((false, general.0)),
2466    }
2467}
2468
2469/// Env-reading wrapper over [`alias_door_from`]. A disagreeing pair FALLS CLOSED (door not
2470/// armed = the shipped program) and prints the refusal ONCE PER PROCESS through `latch`.
2471///
2472/// COST, stated because "read-site only" is true of the ARITHMETIC and not of the lookups:
2473/// honoring two names doubles the `env::var` calls on a per-call door (door H goes from ~64 to
2474/// ~128 lookups per ship round across `i32_mirror_store` and the shexp add), and `env::var`
2475/// takes the process environ lock. That is the price of not breaking every banked script, it
2476/// is paid only on doors whose call sites are already per-layer rather than per-token, and it
2477/// is unmeasured on a rig that cannot time host effects (LAW:rig-exactness-only). If a door
2478/// ever moves to a per-token site, resolve it once behind a `OnceLock` and give up the
2479/// in-process arm flipping the gates use today — that is the trade, named in advance.
2480///
2481/// It does not panic and it does not return `Result`: this is read per call inside the round,
2482/// and an abort in the GPU worker thread exits the process and kills every live session
2483/// (engine panics are fleet-fatal). A per-call door refuses by NOT ARMING; the loud line is
2484/// the operator's receipt that neither value won.
2485fn alias_door(
2486    general: &'static str,
2487    alias: &'static str,
2488    latch: &'static std::sync::atomic::AtomicBool,
2489) -> (bool, &'static str) {
2490    let g = std::env::var(general).ok();
2491    let a = std::env::var(alias).ok();
2492    match alias_door_from((general, g.as_deref()), (alias, a.as_deref())) {
2493        Ok(resolved) => resolved,
2494        Err(msg) => {
2495            if !latch.swap(true, std::sync::atomic::Ordering::Relaxed) {
2496                eprintln!("[flag-alias] {msg}");
2497            }
2498            (false, general)
2499        }
2500    }
2501}
2502
2503/// `MEMRA_HTOD_DIET=1` (default OFF; generalized from `MEMRA_GLM5_HTOD_DIET`, which stays
2504/// honored per the flag-alias law above — door H, lane/glm5-moe-loc): ENGINE-GENERIC HtoD
2505/// hygiene. Nothing in either class is family knowledge; both are "the host uploaded bytes
2506/// the device already had".
2507///
2508/// 1. The UNGATED shared-expert add re-uploaded a fresh `vec![1.0f32; t]` per MoE layer-call
2509///    (42 pageable HtoD/round to move a CONSTANT) — it now reads a resident ones buffer
2510///    ([`Engine::shexp_ones`]). Applies to every MoE family whose plan carries no
2511///    `ffn_gate_inp_shexp`.
2512/// 2. The latent-plane `len_d` i32 mirror took `memcpy_htod(&[v], ..)`, a SYNCHRONIZING
2513///    pageable copy, at 11 walk sites + 11 rollback sites per round. It now takes
2514///    [`Engine::i32_set_k`], the existing async twin whose value rides the kernel argument —
2515///    whose own doc already says the copy form is "fine at stream-idle boundaries, poison
2516///    mid-round". Applies to every latent-KV consumer ([`Engine::i32_mirror_store`] is an
2517///    Engine method, not a family method).
2518///
2519/// Both write identical values to identical buffers and both are stream-ordered, so the arms are
2520/// bit-identical by construction. Default OFF because no box timing receipt exists (rig is
2521/// exactness-only): 64 driver calls/round of measured count, UNPRICED wall. Read per call.
2522pub fn htod_diet_on() -> bool {
2523    htod_diet_armed().0
2524}
2525
2526/// Once-per-process latch for door H's disagreeing-pair line.
2527static HTOD_DIET_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2528    std::sync::atomic::AtomicBool::new(false);
2529
2530/// Resolve door H, returning the armed flag name for refusals/announces.
2531pub(crate) fn htod_diet_armed() -> (bool, &'static str) {
2532    alias_door(
2533        "MEMRA_HTOD_DIET",
2534        "MEMRA_GLM5_HTOD_DIET",
2535        &HTOD_DIET_ALIAS_WARNED,
2536    )
2537}
2538
2539/// HtoD calls avoided by door H (`MEMRA_HTOD_DIET`): the count receipt. A gate asserts it
2540/// tracks the layer-call count on the ON arm and stays flat on the OFF arm.
2541pub static HTOD_DIET_AVOIDED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2542
2543/// Snapshot of [`HTOD_DIET_AVOIDED`] — gates take a before/after delta.
2544pub fn htod_diet_avoided() -> u64 {
2545    HTOD_DIET_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2546}
2547
2548/// `MEMRA_EP_DIET=1` (default OFF; generalized from `MEMRA_GLM5_EP_DIET`, which stays honored
2549/// per the flag-alias law — lane/glm5-ep-diet): the EP DISPATCH DIET door, general to any
2550/// expert-parallel MoE walk. What the door names is a movement CLASS, not a family: one bulk
2551/// peer activation fan-out per layer-call instead of per-token uploads, compact peer staging
2552/// with one bulk return instead of a per-slot round-trip dribble, and one scatter launch
2553/// instead of the `t*n_used` sequential axpy chain. The glm5 TP-2 walk is today's CONSUMER
2554/// (its kernels, its combine order, its counters in `glm5_tp.rs`); hy3/step EP walks arm the
2555/// same door for their own walks.
2556///
2557/// The glm5 consumer's contract, unchanged: same per-slot expert kernels, same slot-ordered combine chain, restructured
2558/// data movement: ONE bulk peer z fan-out per layer-call (skipped entirely when no peer-owned
2559/// expert routed), zero per-slot host round-trips (peer rows stage compact on the peer and
2560/// return in ONE bulk DtoH+HtoD), and the t*n_used sequential `axpy_f32` combine launches
2561/// collapse into ONE `moe_pairs_scatter` launch — whose kernel header carries the
2562/// byte-identity contract vs the zeros+sequential-axpy chain. Decode stays BYTE-identical to
2563/// the v1 walk (and therefore to plain) by construction; `glm5-tp-gate` re-proves it with the
2564/// door pinned ON. Default OFF: the rig is exactness-only and the door changes the round's
2565/// SYNC STRUCTURE (the class the diet window warned does not always transfer from counts to
2566/// wall) — it ships with count receipts and the box window prices the wall. Read per call;
2567/// `=0`/unset restores the v1 per-slot walk byte-for-byte.
2568pub fn ep_diet_on() -> bool {
2569    ep_diet_armed().0
2570}
2571
2572/// Once-per-process latch for the EP-diet door's disagreeing-pair line.
2573static EP_DIET_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2574    std::sync::atomic::AtomicBool::new(false);
2575
2576/// Resolve the EP-diet door, returning the armed flag name — the co-refusal in `hybrid.rs`
2577/// names the flag the operator actually set.
2578pub(crate) fn ep_diet_armed() -> (bool, &'static str) {
2579    alias_door("MEMRA_EP_DIET", "MEMRA_GLM5_EP_DIET", &EP_DIET_ALIAS_WARNED)
2580}
2581
2582/// `MEMRA_EP_GROUPED_PRIME=1` (default OFF; generalized from `MEMRA_GLM5_EP_GROUPED_PRIME`,
2583/// which stays honored per the flag-alias law — lane/glm5-ep-diet): the EP GROUPED-PRIME door,
2584/// general to any expert-parallel MoE walk — "run the family's own chunked grouped MoE prefill
2585/// program per rank over each rank's resident expert slab, then add the peer's bulk-returned
2586/// partial". The glm5 TP-2 walk is today's consumer.
2587///
2588/// The glm5 consumer's contract, unchanged: port the chunked
2589/// grouped MoE prefill (`MEMRA_MOE_GROUPED_PREFILL`, the plain walk's default-ON 85->616-639
2590/// tok/s prefill program) through the glm5 TP-2 EP walk: the SAME sigmoid host-oracle
2591/// routing, per-rank expert-major CSR restricted to each rank's owned experts, one grouped
2592/// f16 GEMM per projection PER RANK over the rank's resident EP slab (pointer tables minted
2593/// at arm time), per-rank slot-ordered scatter, then root adds the peer's bulk-returned
2594/// partial. Fires only where the plain grouped arm would (f16g-eligible qtypes, PRE-clamp,
2595/// n_used<=8); everything else — including the rig fixture's Q8_0 bank — falls closed to the
2596/// (dieted) sequential EP walk. Numeric class: per-expert GEMMs are the plain grouped arm's;
2597/// the ONE reassociation is the per-token root+peer partial add (band-gated, never claimed
2598/// byte). Read per call.
2599pub fn ep_grouped_prime_on() -> bool {
2600    ep_grouped_prime_armed().0
2601}
2602
2603/// Once-per-process latch for the EP grouped-prime door's disagreeing-pair line.
2604static EP_GROUPED_PRIME_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2605    std::sync::atomic::AtomicBool::new(false);
2606
2607/// Resolve the EP grouped-prime door, returning the armed flag name for the co-refusal.
2608pub(crate) fn ep_grouped_prime_armed() -> (bool, &'static str) {
2609    alias_door(
2610        "MEMRA_EP_GROUPED_PRIME",
2611        "MEMRA_GLM5_EP_GROUPED_PRIME",
2612        &EP_GROUPED_PRIME_ALIAS_WARNED,
2613    )
2614}
2615
2616/// `MEMRA_TOPK_SHARDS` (lane/glm5-matvec door K, default ON since the 2026-08-31 mv-battery
2617/// flip; `=0` is the rollback seam): `topk_rows` runs the exact
2618/// two-launch shard split (per-(row,shard) partial top-k + per-row shard merge) instead of the
2619/// one-block-per-row kernel. The standing kernel puts n_rows blocks on the card (the DFlash2
2620/// selector: 15 blocks on 188 SMs, 9.3 MB read in 1.31 ms = 7 GB/s). Top-k under the total
2621/// order (value desc, column asc) is a discrete selection, so the shard split is
2622/// OUTPUT-IDENTICAL by construction (same insertion comparisons, same tie rules in both
2623/// stages); gated by `glm5_matvec_doors_gpu` incl. planted-tie fixtures. Read per call.
2624fn topk_shards_on() -> bool {
2625    std::env::var("MEMRA_TOPK_SHARDS").as_deref() != Ok("0")
2626}
2627
2628/// Engagement counter for the sharded top-k door (`MEMRA_TOPK_SHARDS`).
2629pub static TOPK_SHARDS_DISPATCHES: std::sync::atomic::AtomicU64 =
2630    std::sync::atomic::AtomicU64::new(0);
2631
2632/// Snapshot of [`TOPK_SHARDS_DISPATCHES`] — gates take a before/after delta.
2633pub fn topk_shards_dispatches() -> u64 {
2634    TOPK_SHARDS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2635}
2636
2637/// `MEMRA_ALLOC_TRACE=1` (gate-harness instrument, default OFF, never a serving flag): every
2638/// Engine device allocation funnel (`alloc_uninit` and the zeroed / typed / host-upload
2639/// wrappers) prints `[alloc-trace] <bytes> bytes from <file>:<line>` naming the CALLER
2640/// (`#[track_caller]`). Why: the nsys trace of the B200 GLM-5.3-Flash decode (2026-09-03)
2641/// counted ~1,545 cuMemAllocAsync / cuMemFreeAsync pairs per token (1.8 ms of host API time per
2642/// token) with no host stacks; this names the Rust lines that churn the pool.
2643pub fn alloc_trace_on() -> bool {
2644    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2645    *V.get_or_init(|| std::env::var("MEMRA_ALLOC_TRACE").as_deref() == Ok("1"))
2646}
2647
2648#[track_caller]
2649pub(crate) fn alloc_trace_hit(bytes: usize) {
2650    if alloc_trace_on() {
2651        let loc = std::panic::Location::caller();
2652        eprintln!(
2653            "[alloc-trace] {bytes} bytes from {}:{}",
2654            loc.file(),
2655            loc.line()
2656        );
2657    }
2658}
2659
2660/// `MEMRA_DTOH_TRACE=1` (gate-harness instrument, default OFF, never a serving flag): every
2661/// Engine device-to-host copy prints one line, `[dtoh-trace] <bytes> bytes from <file>:<line>`,
2662/// naming the CALLER of the wrapper (`#[track_caller]`). Why: an nsys trace of the B200 decode
2663/// (2026-09-03) showed two blocking DtoH calls per token (4 B after `argmax_final_f32`, 2112 B
2664/// after `moe_router_sigmoid_topk_f32`) each costing ~1.3 ms of queue drain, and the trace has
2665/// no host stacks; this names the Rust line that owns each drain.
2666pub fn dtoh_trace_on() -> bool {
2667    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2668    *V.get_or_init(|| std::env::var("MEMRA_DTOH_TRACE").as_deref() == Ok("1"))
2669}
2670
2671#[track_caller]
2672pub(crate) fn dtoh_trace_hit(bytes: usize) {
2673    if dtoh_trace_on() {
2674        let loc = std::panic::Location::caller();
2675        eprintln!(
2676            "[dtoh-trace] {bytes} bytes from {}:{}",
2677            loc.file(),
2678            loc.line()
2679        );
2680    }
2681}
2682
2683/// `MEMRA_MOE_EXPERT_RP=1` (default OFF, memra#147): the device-RESIDENT NVFP4 expert slabs are
2684/// repacked at upload into the slot-major per-row layout the engine already names
2685/// `QT_NVFP4_V2` (per row: slot g's 16 quant bytes at g*16, its two UE4M3 scale bytes at
2686/// nsb*16 + g*2; `nvfp4_expert_split_repack`, the same bytes as tp.rs
2687/// `nvfp4_matrix_v2_permute`) and `DevExps::rp` is set. Readers are told `QT_NVFP4_V2`
2688/// (`rp_qt`): every expert dot goes through `expert_dot_g`'s V2 case on the shared pinned core
2689/// (one 16B window per lane-group at a 16B lane stride instead of five scattered 4B loads at
2690/// a 36B stride: root ncu measured 24.97 sectors per warp request on
2691/// `moe_gate_up_preclamp8_q8_w4`, 4 is coalesced), and the grouped prefill takes its existing
2692/// V2 dequant / `kq_fetch<V2>` arms. Host bytes, the SLRU cache and the TP upload paths stay
2693/// interleaved and untouched. A resident-slab reader not yet handed the V2 qtype refuses with
2694/// a named error (`moe_rp_refuse`) rather than reading repacked bytes interleaved.
2695pub fn moe_expert_rp_on() -> bool {
2696    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2697    *V.get_or_init(|| std::env::var("MEMRA_MOE_EXPERT_RP").as_deref() == Ok("1"))
2698}
2699
2700/// The qtype a kernel is told for an expert slab: `QT_NVFP4_V2` (the slot-major per-row layout,
2701/// tp.rs `nvfp4_matrix_v2_permute`) when the slab it will read is a repacked resident slab, the
2702/// tensor's own qtype otherwise.
2703pub fn rp_qt(rp: bool, qt: i32) -> i32 {
2704    if rp && qt == QT_NVFP4 {
2705        QT_NVFP4_V2
2706    } else {
2707        qt
2708    }
2709}
2710
2711/// A resident-slab reader that has no split-plane arm refuses, by name, instead of reading the
2712/// repacked bytes with the interleaved walk (which would be a plausible-looking wrong answer).
2713pub fn moe_rp_refuse(rp: bool, path: &str) -> Result<(), Box<dyn std::error::Error>> {
2714    if rp {
2715        return Err(format!(
2716            "{path}: the resident expert slab is split-plane (MEMRA_MOE_EXPERT_RP=1) and this \
2717             path reads experts interleaved; it is not wired for the door (memra#147). Boot \
2718             without MEMRA_MOE_EXPERT_RP for this model or wire the path."
2719        )
2720        .into());
2721    }
2722    Ok(())
2723}
2724
2725/// `MEMRA_VERIFY_WS` (lane/glm5-matvec door W, default ON since the 2026-08-31 mv-battery
2726/// flip; `=0` is the rollback seam; generalized from `MEMRA_GLM5_VERIFY_WS`, which stays
2727/// honored as the family alias — OFF-WINS composition: either name `=0` disables, so every
2728/// banked gate arm and box script pinning the old name keeps its exact semantics, and the
2729/// old name is never silently dead): the verify walk's
2730/// recurring buffers draw from the engine's size-keyed free-lists and recycle back instead
2731/// of one `cuMemAllocAsync`+Free pair per buffer (~1380+1370 driver calls/token on the ship
2732/// shape — diet-battery apisum; `MEMRA_HC_DECODE_WS` owns only the t=1 walk and never
2733/// reaches the spec serving shape). Byte-identical by the sites' own full-overwrite uninit
2734/// contract; gated by `glm5_matvec_doors_gpu` (multi-call byte identity + the
2735/// `SCRATCH_ALLOC_CALLS` delta receipt). Read per call — the rollback seam.
2736/// `MEMRA_HYPER_BATCH_SOLO=1` (default OFF, read per call — the rollback seam is its absence):
2737/// at B=1 the batched hc decode walk delegates to the solo walk `hyper_range_decode`.
2738///
2739/// WHY IT EXISTS. glm5 PP-N serving decodes through `hyper_batch_range_decode`, which is the only
2740/// hc decode walk it reaches and the only one with neither the allocation-workspace door
2741/// (`MEMRA_HC_DECODE_WS`) nor the decode-graph door (`MEMRA_GLM5_DECODE_GRAPH`) — both of those
2742/// guard `hyper_range_decode_eager`, reachable only via `hyper_range_decode`. Measured on the
2743/// 2x B200 pair 2026-09-03: forcing `MEMRA_HC_DECODE_WS=1` on serving moved 55.85 -> 55.96 tok/s
2744/// (noise) and printed its engagement line ZERO times, because the walk was never entered. At B=1
2745/// the batch walk also pays a per-layer `h_row` allocation plus a `dtod_copy_view` of the single
2746/// row, which the solo walk does not.
2747///
2748/// BYTE IDENTITY is the batch walk's own gated contract ("row b of a B-row step must be
2749/// BIT-IDENTICAL to session b decoding alone through `decode_step_hyper`", `glm5-hyper-batch-gate`,
2750/// red-armed), so at B=1 the delegation is that contract's own right-hand side.
2751pub(crate) fn hyper_batch_solo_on() -> bool {
2752    std::env::var("MEMRA_HYPER_BATCH_SOLO").as_deref() == Ok("1")
2753}
2754
2755/// `MEMRA_HC_PRE_BLOCK=<n>` (default 128, lane/b200-hcpre-wide-20260903): the CUDA block
2756/// width of the fused hyper-connection pre-chain, when `MEMRA_HC_FUSED_PRE=2` selected the
2757/// v2 arm. 128 is v2 verbatim.
2758///
2759/// WHY THIS EXISTS. `memra_dsv4_hc_pre_fused_v2` launches one block of 128 threads PER ROW.
2760/// At t=1 decode there is one row, so the whole call occupies ONE SM of a B200's 148. nsys
2761/// on the 2x B200 pair in the current best posture (2026-09-03) makes it the largest kernel
2762/// in the decode profile: 17.5% of kernel time, 31.1 us average, 23,220 launches over 256
2763/// profiled tokens = 90.7 per token, which is exactly the 2 sites (attn, mlp) on each of 45
2764/// layers. Those 31 us move about 128 KB, i.e. 4.1 GB/s, so the kernel is latency-bound on
2765/// four warps rather than limited by its arithmetic.
2766///
2767/// EXACTNESS, stated rather than assumed. Stage 3 (the collapse) is bit-identical at any
2768/// width: each output sums the same hc terms in the same order and only the owning thread
2769/// moves. Stage 2 (Sinkhorn) is warp-0-only at every width. Stage 1 (rowsq) is NOT: a wider
2770/// block gives `dsv4_block_sum` a different partition of the row, so the double accumulation
2771/// order changes. The f32 narrowing of `1/sqrt(tot/w + eps)` is expected to absorb a
2772/// last-ulp double difference, but expected is not constructed, so any width other than 128
2773/// is the named numeric class `hc_pre_rowsq_blockwide` and carries an argmax gate plus a
2774/// greedy tape before it can be a default.
2775///
2776/// Refuses a value that is not a power of two in [32, 1024], by name, at read time — a bad
2777/// width would otherwise reach the launcher and return an opaque 40023.
2778/// `MEMRA_HC_PRE_SINK_REG=1` (default OFF, read per call — the rollback seam is its absence):
2779/// run the fused hc pre-chain's Sinkhorn stage in registers with `__shfl_sync` instead of
2780/// shared memory. Only meaningful alongside `MEMRA_HC_PRE_BLOCK` (it rides the v3 kernel).
2781///
2782/// WHY, from two nsys measurements rather than an argument. The same kernel at two block
2783/// widths on 2x B200 (2026-09-03): 128 threads -> 31.194 us, 1024 threads -> 26.609 us. Stages
2784/// 1 and 3 scale with the block and stage 2 does not (warp-0-only at every width), so
2785/// S + P = 31.194 and S + P/8 = 26.609 give P = 5.24 us and **S = 25.95 us**. The Sinkhorn is
2786/// 83% of the kernel: 90 launches x 25.95 us = 2.34 ms of an 18.44 ms token, 12.7% of the
2787/// token, to normalise an hc x hc matrix (16 floats at hc=4) for `hc_sinkhorn_iters` = 20
2788/// rounds. It is not arithmetic — per round the shared path does ~2*hc dependent shared loads
2789/// per lane plus six `__syncwarp` and a shared `atomicOr`, on ONE warp with nothing resident to
2790/// cover the latency.
2791///
2792/// BIT-IDENTICAL BY CONSTRUCTION, and that is the point of the design. `comb` lives one element
2793/// per lane and every row/column sum is gathered with `__shfl_sync` IN THE SAME ORDER the
2794/// shared loop used, so the same addends land in the same sequence in the same running float.
2795/// This is NOT a numeric class and needs no argmax gate — unlike `hc_pre_rowsq_blockwide`,
2796/// which the same door family does carry. A tree reduction would have been fewer instructions
2797/// and a different association; it is deliberately not used.
2798pub(crate) fn hc_pre_sink_reg() -> bool {
2799    hc_pre_sink_reg_from(
2800        std::env::var("MEMRA_HC_PRE_SINK_REG").ok().as_deref(),
2801        env!("MEMRA_BUILT_CUDA_ARCH"),
2802    )
2803}
2804
2805/// The pure parse behind [`hc_pre_sink_reg`] (arch-keyed since 2026-09-04): `1` arms, `0`
2806/// disarms, unset = ON on `100a` builds (receipt +7.86% alone, +10.25% with the 512-wide block,
2807/// tape 9437b599f6b9d2a9, darklanes research/glm5-b200-20260902/LANE.md hcpreab), OFF elsewhere.
2808pub fn hc_pre_sink_reg_from(v: Option<&str>, built_arch: &str) -> bool {
2809    match v.map(str::trim) {
2810        Some("1") => true,
2811        Some("0") => false,
2812        _ => built_arch == "100a",
2813    }
2814}
2815
2816/// `MEMRA_HC_PRE_SPLIT_COLLAPSE`: run the hc pre-chain's stage 3 as a SEPARATE multi-block
2817/// kernel (`memra_dsv4_hc_collapse`, which already exists as the unfused chain's third kernel)
2818/// instead of inline in the fused kernel's single block.
2819///
2820/// DEFAULT OFF pending its model-scale row (new-flags law). WHY IT SHOULD PAY: at decode the
2821/// fused kernel's grid is the SEQUENCE LENGTH, so s = 1 puts all three stages on ONE block --
2822/// measured 8.77 us for 146 KB, about 16.6 GB/s, roughly a single SM's reach. Widening the block
2823/// saturates inside that SM (1024 measured worse than 512: 9.02 vs 8.77 us) because the limit is
2824/// outstanding loads per SM, not threads; BLOCKS are the axis that multiplies memory-level
2825/// parallelism, and the standalone collapse at grid(d/256, s) = 16 blocks measures 1.8 us for
2826/// stage 3's 81 KB. Costs one extra launch per site (90 per token).
2827///
2828/// EXACTNESS: bit-identical by construction. Each output is `sum_c pre[c] * x[c*d+i]` with the
2829/// c-sum inside ONE thread, so partitioning i across blocks moves no arithmetic, and the split
2830/// kernel reads back the exact `pre` bits the fused kernel wrote. Stage 1's reduction is NOT
2831/// split for exactly the reason stage 3 can be: repartitioning a reduction changes its order.
2832pub fn hc_pre_split_collapse() -> bool {
2833    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2834    *ON.get_or_init(|| std::env::var("MEMRA_HC_PRE_SPLIT_COLLAPSE").as_deref() == Ok("1"))
2835}
2836
2837/// The build-arch default of `MEMRA_HC_PRE_BLOCK`: 512 on `100a`, 128 elsewhere.
2838pub fn hc_pre_block_default(built_arch: &str) -> usize {
2839    if built_arch == "100a" { 512 } else { 128 }
2840}
2841
2842pub(crate) fn hc_pre_block() -> usize {
2843    // Arch-keyed default since 2026-09-04: 512 on `100a` builds (receipt +10.25% with the
2844    // register Sinkhorn, tape 9437b599f6b9d2a9 unchanged, darklanes
2845    // research/glm5-b200-20260902/LANE.md hcpreab), 128 (= v2 verbatim) everywhere else.
2846    let default = hc_pre_block_default(env!("MEMRA_BUILT_CUDA_ARCH"));
2847    match std::env::var("MEMRA_HC_PRE_BLOCK") {
2848        Err(_) => default,
2849        Ok(v) if v.is_empty() => default,
2850        Ok(v) => match v.parse::<usize>() {
2851            Ok(n) if (32..=1024).contains(&n) && n.is_power_of_two() => n,
2852            _ => {
2853                static SAID: std::sync::Once = std::sync::Once::new();
2854                SAID.call_once(|| {
2855                    eprintln!(
2856                        "[hc-pre-block] MEMRA_HC_PRE_BLOCK={v:?} is not a power of two in \
2857                         [32, 1024]; using the build default {default}"
2858                    );
2859                });
2860                default
2861            }
2862        },
2863    }
2864}
2865
2866fn verify_ws_on() -> bool {
2867    verify_ws_on_from(
2868        std::env::var("MEMRA_VERIFY_WS").ok().as_deref(),
2869        std::env::var("MEMRA_GLM5_VERIFY_WS").ok().as_deref(),
2870    )
2871}
2872
2873/// The pure OFF-wins composition over the general name and the glm5 alias (unit-tested
2874/// without env mutation; default ON, either name `=0` disables).
2875fn verify_ws_on_from(general: Option<&str>, glm5_alias: Option<&str>) -> bool {
2876    general != Some("0") && glm5_alias != Some("0")
2877}
2878
2879/// Engagement counter for the verify-walk workspace (`MEMRA_VERIFY_WS`): incremented
2880/// once per POOL HIT (a reused buffer = one avoided alloc + one avoided free). Gates anchor
2881/// on the delta; `SCRATCH_ALLOC_CALLS` carries the complementary real-alloc count.
2882pub static VERIFY_WS_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2883
2884/// Snapshot of [`VERIFY_WS_HITS`] — gates take a before/after delta.
2885pub fn verify_ws_hits() -> u64 {
2886    VERIFY_WS_HITS.load(std::sync::atomic::Ordering::Relaxed)
2887}
2888
2889/// Engagement counter for the glm5_next tensor-core MLA prefill chain
2890/// (`MEMRA_MLA_TC_PREFILL`), incremented once per (layer, chunk) dispatch at the chain's own
2891/// invocation, AFTER the strided-batched GEMM decline check — a declined shape does not count.
2892/// A gate that must prove "the TC arm ran N times for this workload" reads this delta; the
2893/// once-per-boot announce line dedups and cannot carry a count
2894/// (LAW:wiring-assertions-match-prose).
2895pub static MLA_TC_PREFILL_DISPATCHES: std::sync::atomic::AtomicU64 =
2896    std::sync::atomic::AtomicU64::new(0);
2897
2898/// Snapshot of [`MLA_TC_PREFILL_DISPATCHES`]. Gates take a before/after pair around a workload
2899/// and assert on the delta — including the DECODE byte-identity gate, whose assertion is that
2900/// this stays FLAT across t=1 steps with the flag on.
2901pub fn mla_tc_prefill_dispatches() -> u64 {
2902    MLA_TC_PREFILL_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2903}
2904
2905/// Engagement counter for the glm5_next expert-grouped MoE PREFILL arm
2906/// (`MEMRA_MOE_GROUPED_PREFILL`), incremented once per (layer, chunk) dispatch at the arm's own
2907/// call site. Same reason the fused-epilogue counter exists: the observation env vars divert
2908/// dispatch, so a counter at the invocation is the only honest engagement receipt
2909/// (LAW:wiring-assertions-match-prose). Read via [`moe_grouped_prefill_dispatches`].
2910pub static MOE_GROUPED_PREFILL_DISPATCHES: std::sync::atomic::AtomicU64 =
2911    std::sync::atomic::AtomicU64::new(0);
2912
2913/// Snapshot of [`MOE_GROUPED_PREFILL_DISPATCHES`]. Gates take a before/after pair around a
2914/// workload and assert on the delta.
2915pub fn moe_grouped_prefill_dispatches() -> u64 {
2916    MOE_GROUPED_PREFILL_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2917}
2918
2919/// RAII guard from `Engine::exact_scope`: restores the pre-scope `verify_exact` value on
2920/// drop, so error propagation (`?`) can never leave the engine latched in the
2921/// decode-exact matmul program (hermes finding, fixed 2026-08-23). Holds the flag, not
2922/// the Engine, so the restoration contract is unit-testable without a GPU.
2923#[must_use = "dropping immediately ends the exact scope"]
2924pub struct ExactScope<'a> {
2925    flag: &'a std::sync::atomic::AtomicBool,
2926    prev: bool,
2927}
2928
2929impl<'a> ExactScope<'a> {
2930    pub(crate) fn set(flag: &'a std::sync::atomic::AtomicBool, on: bool) -> Self {
2931        let prev = flag.load(std::sync::atomic::Ordering::Relaxed);
2932        flag.store(on, std::sync::atomic::Ordering::Relaxed);
2933        ExactScope { flag, prev }
2934    }
2935}
2936
2937impl Drop for ExactScope<'_> {
2938    fn drop(&mut self) {
2939        self.flag
2940            .store(self.prev, std::sync::atomic::Ordering::Relaxed);
2941    }
2942}
2943
2944#[cfg(test)]
2945mod verify_ws_flag_tests {
2946    use super::verify_ws_on_from;
2947
2948    #[test]
2949    fn off_wins_across_general_and_alias() {
2950        // default ON
2951        assert!(verify_ws_on_from(None, None));
2952        // either name =0 disables (the banked gate arms pin the ALIAS =0; the general
2953        // name must be exactly as loud)
2954        assert!(!verify_ws_on_from(Some("0"), None));
2955        assert!(!verify_ws_on_from(None, Some("0")));
2956        assert!(!verify_ws_on_from(Some("1"), Some("0")));
2957        assert!(!verify_ws_on_from(Some("0"), Some("1")));
2958        // explicit ON on either name keeps the default
2959        assert!(verify_ws_on_from(Some("1"), None));
2960        assert!(verify_ws_on_from(None, Some("1")));
2961    }
2962}
2963
2964#[cfg(test)]
2965mod moe_vrows_ilp_default_tests {
2966    use super::moe_vrows_ilp_on_from;
2967
2968    #[test]
2969    fn nvfp4_row_ilp_arch_keyed_default() {
2970        use super::nvfp4_row_ilp_on_from;
2971        assert!(nvfp4_row_ilp_on_from(None, "100a"));
2972        assert!(!nvfp4_row_ilp_on_from(None, "120a"));
2973        assert!(nvfp4_row_ilp_on_from(Some("1"), "120a"));
2974        assert!(!nvfp4_row_ilp_on_from(Some("0"), "100a"));
2975    }
2976
2977    #[test]
2978    fn q8_row_ilp_arch_keyed_default() {
2979        use super::q8_row_ilp_on_from;
2980        assert!(q8_row_ilp_on_from(None, "100a"));
2981        assert!(!q8_row_ilp_on_from(None, "120a"));
2982        assert!(q8_row_ilp_on_from(Some("1"), "120a"));
2983        assert!(!q8_row_ilp_on_from(Some("0"), "100a"));
2984    }
2985
2986    #[test]
2987    fn arch_keyed_default_with_explicit_override() {
2988        assert!(moe_vrows_ilp_on_from(None, "100a"));
2989        assert!(!moe_vrows_ilp_on_from(None, "120a"));
2990        assert!(!moe_vrows_ilp_on_from(None, "90a"));
2991        assert!(moe_vrows_ilp_on_from(Some("1"), "120a"));
2992        assert!(!moe_vrows_ilp_on_from(Some("0"), "100a"));
2993        assert!(!moe_vrows_ilp_on_from(Some(" 0 "), "100a"));
2994    }
2995}
2996
2997#[cfg(test)]
2998mod hc_pre_default_tests {
2999    use super::{hc_pre_block_default, hc_pre_sink_reg_from};
3000
3001    #[test]
3002    fn arch_keyed_defaults_with_explicit_override() {
3003        assert_eq!(hc_pre_block_default("100a"), 512);
3004        assert_eq!(hc_pre_block_default("120a"), 128);
3005        assert!(hc_pre_sink_reg_from(None, "100a"));
3006        assert!(!hc_pre_sink_reg_from(None, "120a"));
3007        assert!(hc_pre_sink_reg_from(Some("1"), "120a"));
3008        assert!(!hc_pre_sink_reg_from(Some("0"), "100a"));
3009        assert_eq!(
3010            crate::hyper::hc_fused_pre_arm_from(None, "100a"),
3011            crate::hyper::HcFusedPreArm::V2
3012        );
3013        assert_eq!(
3014            crate::hyper::hc_fused_pre_arm_from(None, "120a"),
3015            crate::hyper::HcFusedPreArm::Off
3016        );
3017        assert_eq!(
3018            crate::hyper::hc_fused_pre_arm_from(Some("0"), "100a"),
3019            crate::hyper::HcFusedPreArm::Off
3020        );
3021        assert_eq!(
3022            crate::hyper::hc_fused_pre_arm_from(Some("1"), "100a"),
3023            crate::hyper::HcFusedPreArm::V1
3024        );
3025    }
3026}
3027
3028#[cfg(test)]
3029mod glm5_decode_graph_default_tests {
3030    use super::glm5_decode_graph_on_from;
3031
3032    /// Default ON since 2026-09-04; only an explicit `0` disarms. The OFF arm of every gate
3033    /// and test sets `=0`, and this is the contract that keeps that arm non-vacuous.
3034    #[test]
3035    fn unset_and_one_arm_only_zero_disarms() {
3036        assert!(glm5_decode_graph_on_from(None));
3037        assert!(glm5_decode_graph_on_from(Some("1")));
3038        assert!(glm5_decode_graph_on_from(Some("")));
3039        assert!(!glm5_decode_graph_on_from(Some("0")));
3040        assert!(!glm5_decode_graph_on_from(Some(" 0 ")));
3041    }
3042}
3043
3044#[cfg(test)]
3045mod alias_door_tests {
3046    use super::alias_door_from;
3047
3048    const G: &str = "MEMRA_EP_DIET";
3049    const A: &str = "MEMRA_GLM5_EP_DIET";
3050
3051    fn r(g: Option<&str>, a: Option<&str>) -> Result<(bool, &'static str), String> {
3052        alias_door_from((G, g), (A, a))
3053    }
3054
3055    #[test]
3056    fn default_off_and_either_name_arms() {
3057        // unset/unset: the door is OFF and the general name is what a refusal would cite
3058        assert_eq!(r(None, None).unwrap(), (false, G));
3059        // either name =1 arms it, and the ARMED NAME is the one the operator set
3060        assert_eq!(r(Some("1"), None).unwrap(), (true, G));
3061        assert_eq!(r(None, Some("1")).unwrap(), (true, A));
3062        // =0 is a deliberate pin on either name, never an arming
3063        assert_eq!(r(Some("0"), None).unwrap(), (false, G));
3064        assert_eq!(r(None, Some("0")).unwrap(), (false, A));
3065        // anything that is not "1" is not an arming (no truthiness guessing)
3066        assert_eq!(r(None, Some("on")).unwrap(), (false, A));
3067        assert_eq!(r(Some(""), None).unwrap(), (false, G));
3068    }
3069
3070    #[test]
3071    fn agreeing_pair_resolves_to_the_general_name() {
3072        assert_eq!(r(Some("1"), Some("1")).unwrap(), (true, G));
3073        assert_eq!(r(Some("0"), Some("0")).unwrap(), (false, G));
3074    }
3075
3076    #[test]
3077    fn disagreeing_pair_refuses_and_names_both() {
3078        for (g, a) in [("1", "0"), ("0", "1")] {
3079            let err = r(Some(g), Some(a)).expect_err("a disagreeing pair must refuse");
3080            assert!(
3081                err.contains(G),
3082                "the refusal must name the general flag: {err}"
3083            );
3084            assert!(err.contains(A), "the refusal must name the alias: {err}");
3085            // and it must say which way it falls, so an operator reading the line knows the
3086            // door is CLOSED rather than guessing a precedence winner
3087            assert!(err.contains("falls closed"), "{err}");
3088        }
3089    }
3090}
3091
3092#[cfg(test)]
3093mod exact_scope_tests {
3094    use std::sync::atomic::{AtomicBool, Ordering};
3095
3096    #[test]
3097    fn error_path_restores_verify_exact() {
3098        // TOOTH (hermes finding, fixed 2026-08-23): dspark_spec_session_burst called
3099        // set_verify_exact(true)/(false) manually with `?`s in between — any error left
3100        // the engine latched in the decode-exact matmul program for every later request.
3101        // The RAII scope must restore across an error propagation.
3102        let flag = AtomicBool::new(false);
3103        let failing = |flag: &AtomicBool| -> Result<(), &'static str> {
3104            let _scope = super::ExactScope::set(flag, true);
3105            assert!(flag.load(Ordering::Relaxed), "scope arms the flag");
3106            Err("draft forward failed")? // the `?` exit the manual pair leaked on
3107        };
3108        assert!(failing(&flag).is_err());
3109        assert!(
3110            !flag.load(Ordering::Relaxed),
3111            "error propagation must restore the pre-scope value"
3112        );
3113        // Nested/previous-value contract: a scope entered while already ON restores ON.
3114        let flag = AtomicBool::new(true);
3115        {
3116            let _scope = super::ExactScope::set(&flag, true);
3117        }
3118        assert!(flag.load(Ordering::Relaxed));
3119        // Early drop ends the scope exactly where the manual `false` used to sit.
3120        let flag = AtomicBool::new(false);
3121        let scope = super::ExactScope::set(&flag, true);
3122        drop(scope);
3123        assert!(!flag.load(Ordering::Relaxed));
3124    }
3125}
3126
3127impl Engine {
3128    pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
3129        let gpu = memra_runtime::Gpu::new(ordinal)?;
3130        // ARCH GUARD (unified dual-arch engine): the fatbins carry single-arch SASS, so a
3131        // binary/device mismatch otherwise dies at first module load with an opaque CUDA
3132        // error. Fail early with the rebuild hint instead. MEMRA_ARCH_CHECK=0 skips.
3133        if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
3134            use cudarc::driver::sys::CUdevice_attribute_enum as A;
3135            let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
3136                .and_then(|d| unsafe {
3137                    Ok((
3138                        cudarc::driver::result::device::get_attribute(
3139                            d,
3140                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
3141                        )?,
3142                        cudarc::driver::result::device::get_attribute(
3143                            d,
3144                            A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
3145                        )?,
3146                    ))
3147                })
3148                .unwrap_or((0, 0));
3149            let built = env!("MEMRA_BUILT_CUDA_ARCH");
3150            let ok = matches!(
3151                (built, maj, min),
3152                ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
3153            );
3154            if !ok {
3155                return Err(format!(
3156                    "memra was built for sm_{built} but device {ordinal} reports compute \
3157                     capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
3158                     auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
3159                )
3160                .into());
3161            }
3162        }
3163        // Default async-pool RELEASE_THRESHOLD is 0: freed blocks return to the OS at every
3164        // sync, so cuMemAllocAsync NODES inside captured graphs re-map memory on EVERY
3165        // cuGraphLaunch (measured 226us/launch on the gemma graph door, 2026-07-23 osrt).
3166        // Pinning the threshold keeps the pool cached -> alloc nodes become pointer bumps.
3167        unsafe {
3168            use cudarc::driver::sys;
3169            let dev: sys::CUdevice = ordinal as sys::CUdevice;
3170            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3171            if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
3172                let mut thresh: u64 = u64::MAX;
3173                let _ = sys::cuMemPoolSetAttribute(
3174                    pool,
3175                    sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
3176                    &mut thresh as *mut u64 as *mut core::ffi::c_void,
3177                );
3178            }
3179        }
3180        let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
3181        let hybrid = gpu
3182            .ctx
3183            .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
3184        let kda = gpu.ctx.load_module(Ptx::from_binary(KDA_FATBIN.to_vec()))?;
3185        let qmatvec = gpu
3186            .ctx
3187            .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
3188        let flash = gpu
3189            .ctx
3190            .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
3191        let gemm = gpu
3192            .ctx
3193            .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
3194        let router = gpu
3195            .ctx
3196            .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
3197        let sample = gpu
3198            .ctx
3199            .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
3200        let copy_stream = gpu.ctx.new_stream()?;
3201        // DECODE EVENT-TRACKING ELISION — DEFAULT ON (2026-07-05; MEMRA_EVT=1 = escape hatch).
3202        // cudarc is in multi-stream mode (main stream +
3203        // copy_stream are both created streams), so with tracking on EVERY launch arg records a
3204        // read/write CudaEvent and inserts cuStreamWaitEvent on prior events. On the 35B MoE decode
3205        // that is ~19k cuStreamWaitEvent + ~9k cuEventRecord + ~6k event create/destroy per token
3206        // (~7 ms/tok host time, measured nsys 2026-07-04 rtx6000), and +4.6% measured on 27B decode —
3207        // protecting NOTHING: every hot-path kernel/memcpy runs on the ONE gpu.stream.
3208        // CROSS-STREAM HAZARD AUDIT: MoeSlotCache in-memory prefetch uses copy_stream. Every
3209        // overwrite explicitly records the prior compute point and makes copy_stream wait; every
3210        // consumer explicitly waits for the copy completion event. The opt-in positioned-read
3211        // proof stays on gpu.stream and retains an explicit event solely to guard pinned-source
3212        // reuse. Graph-capture sites use only gpu.stream, so these handoffs never rely on cudarc's
3213        // implicit event tracking.
3214        // SAFETY: single-stream ordering is total; the runtime mem-pool is configured with
3215        // internal-dependency reuse (memra-runtime), so alloc reuse is stream-ordered too.
3216        if std::env::var("MEMRA_EVT")
3217            .map(|v| v == "1")
3218            .unwrap_or(false)
3219        {
3220            // escape hatch: keep cudarc's implicit cross-stream event tracking.
3221        } else {
3222            unsafe {
3223                gpu.ctx.disable_event_tracking();
3224            }
3225        }
3226        Ok(Self {
3227            gpu,
3228            module,
3229            hybrid,
3230            kda,
3231            qmatvec,
3232            flash,
3233            flash_g: std::sync::OnceLock::new(),
3234            gemm,
3235            router,
3236            sample,
3237            moe_cache: Mutex::new(None),
3238            w8_mirrors: Mutex::new(std::collections::HashMap::new()),
3239            w8_act: Mutex::new(std::collections::HashMap::new()),
3240            moe_cache_layout: Mutex::new(None),
3241            copy_stream,
3242            capture_keep_on: std::sync::atomic::AtomicBool::new(false),
3243            verify_exact: std::sync::atomic::AtomicBool::new(false),
3244            capture_keep: Mutex::new(Vec::new()),
3245            argmax_partials: Mutex::new(None),
3246            prime_deqw_ws: Mutex::new(None),
3247            router_stage: Mutex::new(None),
3248            hyper_decode_ws: Mutex::new(None),
3249            mla_seg_ws: Mutex::new(None),
3250            verify_ws: Mutex::new(VerifyWs::default()),
3251            vrows_macro_dev: Mutex::new(std::collections::HashMap::new()),
3252            shexp_ones: Mutex::new(None),
3253            fp8_scratch: Mutex::new(None),
3254            fa_vf16_scratch: Mutex::new(None),
3255            fa_part_pool: Mutex::new(None),
3256            fa_part_retired: Mutex::new(Vec::new()),
3257            fn_cache: Mutex::new(Default::default()),
3258            f16_scratch: Mutex::new(None),
3259            #[cfg(memra_cutlass)]
3260            cutlass_scratch: Mutex::new(None),
3261        })
3262    }
3263
3264    pub fn ctx(&self) -> &Arc<CudaContext> {
3265        &self.gpu.ctx
3266    }
3267
3268    /// Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks
3269    /// parked in the pool because `Engine::new` pins RELEASE_THRESHOLD to u64::MAX above.
3270    ///
3271    /// Why this is a public engine surface: `mem_get_info`'s `free` DOES NOT SEE these bytes —
3272    /// they are mapped to this process, so `free` counts them as gone, yet the very next
3273    /// `alloc_u8` is satisfied from them without touching `free` at all. Any admission or
3274    /// budget decision that reads `free` alone therefore under-counts real headroom by exactly
3275    /// this amount. Effective allocatable headroom is `free + pool_cached_bytes()`.
3276    ///
3277    /// MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB
3278    /// during the burst — SMALL. The admission gate adds it because a term that can only ever
3279    /// under-count headroom does not belong in a gate that queues real work, but the honest
3280    /// reading of this number is that pool caching is NOT where a long-running server's VRAM
3281    /// hides on this path: reserved ~= used throughout, so the memory the driver reports as
3282    /// gone is genuinely LIVE (see `pool_reserved_used` for the diagnostic pair).
3283    ///
3284    /// Returns 0 if the pool cannot be queried (never a false-positive headroom claim).
3285    pub fn pool_cached_bytes(&self) -> usize {
3286        let (reserved, used) = self.pool_reserved_used();
3287        reserved.saturating_sub(used)
3288    }
3289
3290    /// Bytes the driver's per-device CUDA GRAPH memory pool currently holds RESERVED
3291    /// (`cuDeviceGetGraphMemAttribute` RESERVED_MEM_CURRENT) — the backing store of every
3292    /// captured alloc node, which on this engine means the dspark verify-graph pool
3293    /// (decode/step graphs bake pre-allocated buffers and own no alloc nodes). This memory
3294    /// is DISTINCT from the async pool above: `mem_get_info`'s `free` already excludes it,
3295    /// it is never released back (the vgraph pool has no eviction by design), and it GROWS
3296    /// as new (segment, vt)/(vt, rung, hi) keys capture — the growth is what
3297    /// `dspark_vg_admission_debt` charges at admission. Returns 0 if the attribute cannot
3298    /// be queried (never a false headroom claim, matching `pool_cached_bytes`).
3299    pub fn device_graph_mem_reserved(&self) -> usize {
3300        use cudarc::driver::sys as cus;
3301        let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
3302            return 0;
3303        };
3304        let mut bytes: u64 = 0;
3305        let rc = unsafe {
3306            cus::cuDeviceGetGraphMemAttribute(
3307                dev,
3308                cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
3309                &mut bytes as *mut u64 as *mut std::ffi::c_void,
3310            )
3311        };
3312        if rc == cus::cudaError_enum::CUDA_SUCCESS {
3313            bytes as usize
3314        } else {
3315            0
3316        }
3317    }
3318
3319    /// Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is
3320    /// what the pool has mapped from the driver; used is what is live inside it. Exposed for
3321    /// admission/VRAM diagnostics — the pair distinguishes "memory is parked in the pool and
3322    /// `free` cannot see it" (reserved >> used) from "memory is genuinely held live by some
3323    /// owner" (reserved ~= used), which are opposite bugs with opposite fixes.
3324    /// (0, 0) if the pool cannot be queried.
3325    /// Release every CACHED (freed-but-retained) block of the default async mempool
3326    /// back to the driver (deploy-headroom lane, 2026-08-27). The boot-time
3327    /// RELEASE_THRESHOLD=u64::MAX pin keeps freed blocks cached for graph-launch speed,
3328    /// which is right for steady serving and wrong at a blue/green overlap: a green
3329    /// PROCESS cannot use blue's cached pool. cuMemPoolTrimTo(0) frees only unused
3330    /// blocks — live allocations are untouched; later allocs re-map once. Returns the
3331    /// bytes released (reserved delta), 0 if the pool cannot be queried.
3332    pub fn pool_trim_to_zero(&self) -> usize {
3333        use cudarc::driver::sys;
3334        let (before, _) = self.pool_reserved_used();
3335        unsafe {
3336            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3337            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3338                != sys::CUresult::CUDA_SUCCESS
3339            {
3340                return 0;
3341            }
3342            let _ = sys::cuMemPoolTrimTo(pool, 0);
3343        }
3344        let (after, _) = self.pool_reserved_used();
3345        before.saturating_sub(after)
3346    }
3347
3348    pub fn pool_reserved_used(&self) -> (usize, usize) {
3349        use cudarc::driver::sys;
3350        unsafe {
3351            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3352            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3353                != sys::CUresult::CUDA_SUCCESS
3354            {
3355                return (0, 0);
3356            }
3357            let (mut reserved, mut used) = (0u64, 0u64);
3358            if sys::cuMemPoolGetAttribute(
3359                pool,
3360                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
3361                &mut reserved as *mut u64 as *mut core::ffi::c_void,
3362            ) != sys::CUresult::CUDA_SUCCESS
3363            {
3364                return (0, 0);
3365            }
3366            if sys::cuMemPoolGetAttribute(
3367                pool,
3368                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
3369                &mut used as *mut u64 as *mut core::ffi::c_void,
3370            ) != sys::CUresult::CUDA_SUCCESS
3371            {
3372                return (0, 0);
3373            }
3374            (reserved as usize, used as usize)
3375        }
3376    }
3377
3378    /// Async-pool HIGH-WATER pair since the last reset: (RESERVED_MEM_HIGH, USED_MEM_HIGH)
3379    /// in bytes, then reset both watermarks to their CURRENT values
3380    /// (lane/step37-vram-admission-20260830). This is the instrument the boot admission
3381    /// calibration reads: engine transients are allocated and freed INSIDE one step, so any
3382    /// tick-boundary sampling of `mem_get_info`/pool-current sees nothing of the peak — the
3383    /// driver-kept watermark is the only honest record of how deep a burst actually dipped.
3384    /// (0, 0) if the pool cannot be queried (never a false claim, matching
3385    /// `pool_cached_bytes`).
3386    pub fn pool_high_water_reset(&self) -> (usize, usize) {
3387        use cudarc::driver::sys;
3388        unsafe {
3389            let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3390            if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3391                != sys::CUresult::CUDA_SUCCESS
3392            {
3393                return (0, 0);
3394            }
3395            let (mut reserved, mut used) = (0u64, 0u64);
3396            if sys::cuMemPoolGetAttribute(
3397                pool,
3398                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
3399                &mut reserved as *mut u64 as *mut core::ffi::c_void,
3400            ) != sys::CUresult::CUDA_SUCCESS
3401            {
3402                return (0, 0);
3403            }
3404            if sys::cuMemPoolGetAttribute(
3405                pool,
3406                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
3407                &mut used as *mut u64 as *mut core::ffi::c_void,
3408            ) != sys::CUresult::CUDA_SUCCESS
3409            {
3410                return (0, 0);
3411            }
3412            // Setting a *_HIGH attribute resets the watermark to the pool's current value
3413            // (the value argument must be 0 per the driver contract).
3414            let mut zero: u64 = 0;
3415            let _ = sys::cuMemPoolSetAttribute(
3416                pool,
3417                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
3418                &mut zero as *mut u64 as *mut core::ffi::c_void,
3419            );
3420            let mut zero2: u64 = 0;
3421            let _ = sys::cuMemPoolSetAttribute(
3422                pool,
3423                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
3424                &mut zero2 as *mut u64 as *mut core::ffi::c_void,
3425            );
3426            (reserved as usize, used as usize)
3427        }
3428    }
3429
3430    /// Ambient stream (by value since M1-PP2 increment 2): the thread's pp2 stage stream
3431    /// when a stage scope is active, else the main compute stream — see `Gpu::stream`.
3432    pub fn stream(&self) -> Arc<CudaStream> {
3433        self.gpu.stream()
3434    }
3435    /// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
3436    /// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
3437    pub fn gkv_on() -> bool {
3438        memra_kv::gkv_on()
3439    }
3440
3441    /// FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated
3442    /// window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama
3443    /// same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA)
3444    /// layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module.
3445    /// SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the
3446    /// MTP drafter's acceptance — its single swa attention reads the windowed cache and
3447    /// e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88
3448    /// -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default
3449    /// keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its
3450    /// depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays
3451    /// ON for both — no acceptance cost measured.
3452    pub fn wkv_on() -> bool {
3453        memra_kv::wkv_on()
3454    }
3455
3456    /// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
3457    /// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
3458    /// module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat,
3459    /// 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON
3460    /// (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.
3461    pub fn kv_fp8_on() -> bool {
3462        memra_kv::kv_fp8_on()
3463    }
3464
3465    /// fa kernel routed by head_dim: hd512 (gemma globals) resolves from the kf8vf8 module
3466    /// when the fp8-globals arm is on; everything else from the default flash module.
3467    fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
3468        if head_dim == 512 && Self::gkv_on() {
3469            self.func_g(name)
3470        } else {
3471            self.func(name)
3472        }
3473    }
3474
3475    /// Kernel from the FP8-GLOBALS (kf8vf8) flash module — gemma global-layer arm only.
3476    /// Format-AGNOSTIC kernels (e.g. fa_decode_combine_f32) are not compiled into the
3477    /// per-format fatbins; fall back to the base modules for those.
3478    fn func_g(&self, name: &str) -> CudaFunction {
3479        let m = self.flash_g.get_or_init(|| {
3480            self.gpu
3481                .ctx
3482                .load_module(cudarc::nvrtc::Ptx::from_binary(
3483                    FLASH_FATBIN_KF8VF8.to_vec(),
3484                ))
3485                .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
3486        });
3487        let key = format!("g:{name}");
3488        if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
3489            return f.clone();
3490        }
3491        let f = match m.load_function(name) {
3492            Ok(f) => f,
3493            Err(_) => self.func(name),
3494        };
3495        self.fn_cache.lock().unwrap().insert(key, f.clone());
3496        f
3497    }
3498
3499    fn func(&self, name: &str) -> CudaFunction {
3500        // Resolution cache: cuModuleGetFunction fails inside a CUDA-graph capture region,
3501        // so capture-time lookups MUST be host-memory hits (warmups populate the cache).
3502        if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
3503            return f.clone();
3504        }
3505        let f = self
3506            .module
3507            .load_function(name)
3508            .or_else(|_| self.hybrid.load_function(name))
3509            .or_else(|_| self.kda.load_function(name))
3510            .or_else(|_| self.qmatvec.load_function(name))
3511            .or_else(|_| self.flash.load_function(name))
3512            .or_else(|_| self.gemm.load_function(name))
3513            .or_else(|_| self.router.load_function(name))
3514            .or_else(|_| self.sample.load_function(name))
3515            .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
3516        self.fn_cache
3517            .lock()
3518            .unwrap()
3519            .insert(name.to_string(), f.clone());
3520        f
3521    }
3522
3523    /// Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then
3524    /// dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.
3525    pub fn scatter_trim_logits(
3526        &self,
3527        src: &CudaSlice<f32>,
3528        d2t: &CudaSlice<u32>,
3529        dst: &mut CudaSlice<f32>,
3530        d_vocab: usize,
3531        n_vocab: usize,
3532    ) -> Result<(), Box<dyn std::error::Error>> {
3533        let f1 = self.func("scatter_trim_logits_f32");
3534        let f2 = self.func("scatter_trim_logits_pass2_f32");
3535        let (dv, nv) = (d_vocab as i32, n_vocab as i32);
3536        let cfg1 = LaunchConfig {
3537            grid_dim: (256, 1, 1),
3538            block_dim: (256, 1, 1),
3539            shared_mem_bytes: 0,
3540        };
3541        let __s_b1 = self.gpu.stream();
3542        let mut b1 = __s_b1.launch_builder(&f1);
3543        b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
3544        unsafe {
3545            b1.launch(cfg1)?;
3546        }
3547        let cfg2 = LaunchConfig {
3548            grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
3549            block_dim: (256, 1, 1),
3550            shared_mem_bytes: 0,
3551        };
3552        let __s_b2 = self.gpu.stream();
3553        let mut b2 = __s_b2.launch_builder(&f2);
3554        b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
3555        unsafe {
3556            b2.launch(cfg2)?;
3557        }
3558        Ok(())
3559    }
3560
3561    // ---- FILTERED-SPEC (feat/filtered-spec): top-k/p/min-p transforms applied symmetrically
3562    // to p and q — rejection sampling stays distribution-exact for the filtered target. ----
3563
3564    /// Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the
3565    /// filter (top_k, top_p, min_p) at `temp`. Rows index into x with row_stride f32s.
3566    #[allow(clippy::too_many_arguments)]
3567    pub fn filter_stats(
3568        &self,
3569        x: &CudaSlice<f32>,
3570        row_stride: usize,
3571        rows: &CudaSlice<i32>,
3572        out_th: &mut CudaSlice<f32>,
3573        out_z: &mut CudaSlice<f32>,
3574        out_max: &mut CudaSlice<f32>,
3575        n: usize,
3576        nrow: usize,
3577        temp: f32,
3578        top_k: i32,
3579        top_p: f32,
3580        min_p: f32,
3581    ) -> Result<(), Box<dyn std::error::Error>> {
3582        // A top-K-selection form of this kernel (3 vocab passes vs the search's ~51) was
3583        // implemented and REFUTED on 2026-08-21 (lane/moebatch-q35moe): the 248k-vocab row is
3584        // L2-resident, so the extra passes are near-free while the per-thread selection list
3585        // spills to local memory — B=8 tick 12.8/11.2 ms (cap 64/32) vs 10.4 ms for this
3586        // kernel, and serve c8 agg ~648 vs ~666. The receipts row is the record; the real
3587        // filtered-sampling win was batching the per-row launches (decode_batch.rs).
3588        //
3589        // COOPERATIVE MULTI-BLOCK FORM (lane/samplat, 2026-08-21): the surviving inefficiency
3590        // was WIDTH, not passes — one block per row left ~94% of the device idle for ~620us
3591        // per B=8 serve tick (5.9% of the tick, box4 nsys receipt). filter_stats_coop_f32
3592        // splits each row across 16 blocks with grid-synced bisection totals — same algorithm,
3593        // slice-partial f32 sums (accepted device-sampling class; sample-check arbitrates).
3594        // MEMRA_FILTER_COOP=0 is the rollback seam to the single-block form.
3595        //
3596        // DETERMINISTIC KEYING (hermes finding, fixed 2026-08-23): the old admission
3597        // `16*nrow <= sm_count` fell back to the single-block program PER CALL when a tick
3598        // carried too many rows — and the two programs are NOT bit-identical (measured
3599        // ~1e-7 rel on the renorm mass: different f32 partial-sum order), so a request's
3600        // sampling threshold arithmetic depended on how many rows shared its serve tick.
3601        // Coop is now THE program on every coop-capable device: rows are CHUNKED to the
3602        // co-residency cap (sm_count/16 rows per cooperative launch) and each row's
3603        // arithmetic uses only its own 16 slices + its own ws region, so the per-row bits
3604        // are independent of batch width by construction — the kernel-check
3605        // FILTER-COOP-CHUNK cell pins exactly that. The single-block program remains only
3606        // behind the deployment-keyed seams: MEMRA_FILTER_COOP=0, or a device with
3607        // sm_count < 16 (fixed per device class, never per call).
3608        static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3609        let coop_on =
3610            *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
3611        if coop_on && self.sm_count() >= 16 {
3612            let cap = self.sm_count() as usize / 16;
3613            let mut done = 0usize;
3614            while done < nrow {
3615                let chunk = cap.min(nrow - done);
3616                self.filter_stats_coop_chunk(
3617                    x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
3618                    top_p, min_p,
3619                )?;
3620                done += chunk;
3621            }
3622            return Ok(());
3623        }
3624        self.filter_stats_plain_program(
3625            x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
3626        )
3627    }
3628
3629    /// One cooperative `filter_stats` launch over rows `row0..row0+chunk` (pub so the
3630    /// kernel-check FILTER-COOP-CHUNK cell can pin batch-width independence directly).
3631    /// The kernel indexes `rows`/outputs by blockIdx.y, so the chunk is expressed as
3632    /// sub-views at `row0` — per-row arithmetic is untouched by the offset.
3633    #[allow(clippy::too_many_arguments)]
3634    pub fn filter_stats_coop_chunk(
3635        &self,
3636        x: &CudaSlice<f32>,
3637        row_stride: usize,
3638        rows: &CudaSlice<i32>,
3639        row0: usize,
3640        out_th: &mut CudaSlice<f32>,
3641        out_z: &mut CudaSlice<f32>,
3642        out_max: &mut CudaSlice<f32>,
3643        n: usize,
3644        chunk: usize,
3645        temp: f32,
3646        top_k: i32,
3647        top_p: f32,
3648        min_p: f32,
3649    ) -> Result<(), Box<dyn std::error::Error>> {
3650        let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
3651        let f = self.func("filter_stats_coop_f32");
3652        let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
3653        let cfg = LaunchConfig {
3654            grid_dim: (16, chunk as u32, 1),
3655            block_dim: (512, 1, 1),
3656            shared_mem_bytes: 0,
3657        };
3658        let rows_v = rows.slice(row0..row0 + chunk);
3659        let mut th_v = out_th.slice_mut(row0..row0 + chunk);
3660        let mut z_v = out_z.slice_mut(row0..row0 + chunk);
3661        let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
3662        let __s_b = self.gpu.stream();
3663        let mut b = __s_b.launch_builder(&f);
3664        b.arg(x)
3665            .arg(&rs)
3666            .arg(&rows_v)
3667            .arg(&mut th_v)
3668            .arg(&mut z_v)
3669            .arg(&mut mx_v)
3670            .arg(&mut ws)
3671            .arg(&ni)
3672            .arg(&nr)
3673            .arg(&temp)
3674            .arg(&top_k)
3675            .arg(&top_p)
3676            .arg(&min_p);
3677        unsafe {
3678            b.launch_cooperative(cfg)?;
3679        }
3680        Ok(())
3681    }
3682
3683    /// The single-block-per-row `filter_stats` program (the pre-coop form; the
3684    /// MEMRA_FILTER_COOP=0 rollback and the occupancy fallback). Gate-callable twin of
3685    /// `filter_stats_coop_program`.
3686    #[allow(clippy::too_many_arguments)]
3687    pub fn filter_stats_plain_program(
3688        &self,
3689        x: &CudaSlice<f32>,
3690        row_stride: usize,
3691        rows: &CudaSlice<i32>,
3692        out_th: &mut CudaSlice<f32>,
3693        out_z: &mut CudaSlice<f32>,
3694        out_max: &mut CudaSlice<f32>,
3695        n: usize,
3696        nrow: usize,
3697        temp: f32,
3698        top_k: i32,
3699        top_p: f32,
3700        min_p: f32,
3701    ) -> Result<(), Box<dyn std::error::Error>> {
3702        let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
3703        let f = self.func("filter_stats_f32");
3704        let cfg = LaunchConfig {
3705            grid_dim: (nrow as u32, 1, 1),
3706            block_dim: (1024, 1, 1),
3707            shared_mem_bytes: 0,
3708        };
3709        let __s_b = self.gpu.stream();
3710        let mut b = __s_b.launch_builder(&f);
3711        b.arg(x)
3712            .arg(&rs)
3713            .arg(rows)
3714            .arg(&mut *out_th)
3715            .arg(&mut *out_z)
3716            .arg(&mut *out_max)
3717            .arg(&ni)
3718            .arg(&nr)
3719            .arg(&temp)
3720            .arg(&top_k)
3721            .arg(&top_p)
3722            .arg(&min_p);
3723        unsafe {
3724            b.launch(cfg)?;
3725        }
3726        Ok(())
3727    }
3728
3729    /// out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).
3730    #[allow(clippy::too_many_arguments)]
3731    pub fn softmax_gather_filtered(
3732        &self,
3733        x: &CudaSlice<f32>,
3734        row_stride: usize,
3735        ids: &CudaSlice<u32>,
3736        rows: &CudaSlice<i32>,
3737        th: &CudaSlice<f32>,
3738        z: &CudaSlice<f32>,
3739        out: &mut CudaSlice<f32>,
3740        n: usize,
3741        npair: usize,
3742        temp: f32,
3743    ) -> Result<(), Box<dyn std::error::Error>> {
3744        let f = self.func("softmax_gather_filtered_f32");
3745        let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
3746        let cfg = LaunchConfig {
3747            grid_dim: (npair as u32, 1, 1),
3748            block_dim: (256, 1, 1),
3749            shared_mem_bytes: 0,
3750        };
3751        let __s_b = self.gpu.stream();
3752        let mut b = __s_b.launch_builder(&f);
3753        b.arg(x)
3754            .arg(&rs)
3755            .arg(ids)
3756            .arg(rows)
3757            .arg(th)
3758            .arg(z)
3759            .arg(&mut *out)
3760            .arg(&ni)
3761            .arg(&np)
3762            .arg(&temp);
3763        unsafe {
3764            b.launch(cfg)?;
3765        }
3766        Ok(())
3767    }
3768
3769    /// Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.
3770    #[allow(clippy::too_many_arguments)]
3771    pub fn residual_sample_filtered(
3772        &self,
3773        p: &CudaSlice<f32>,
3774        q: Option<&CudaSlice<f32>>,
3775        n: usize,
3776        temp: f32,
3777        seed: u64,
3778        stream_pos: u32,
3779        p_stats: (f32, f32, f32),
3780        q_stats: (f32, f32, f32),
3781        out_tok: &mut CudaSlice<u32>,
3782    ) -> Result<(), Box<dyn std::error::Error>> {
3783        let f = self.func("residual_sample_filtered_f32");
3784        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3785        let has_q: i32 = q.is_some() as i32;
3786        let qbuf = q.unwrap_or(p);
3787        let (pm, pth, pz) = p_stats;
3788        let (qm, qth, qz) = q_stats;
3789        let cfg = LaunchConfig {
3790            grid_dim: (1, 1, 1),
3791            block_dim: (1024, 1, 1),
3792            shared_mem_bytes: 0,
3793        };
3794        let __s_b = self.gpu.stream();
3795        let mut b = __s_b.launch_builder(&f);
3796        b.arg(p)
3797            .arg(qbuf)
3798            .arg(&has_q)
3799            .arg(&ni)
3800            .arg(&temp)
3801            .arg(&slo)
3802            .arg(&shi)
3803            .arg(&stream_pos)
3804            .arg(&pm)
3805            .arg(&pth)
3806            .arg(&pz)
3807            .arg(&qm)
3808            .arg(&qth)
3809            .arg(&qz)
3810            .arg(&mut *out_tok);
3811        unsafe {
3812            b.launch(cfg)?;
3813        }
3814        Ok(())
3815    }
3816
3817    /// Sparse-q filtered residual sample: token ~ norm(max(0, fp - q)) where fp is the
3818    /// FILTERED softmax of `p` (stats from `filter_stats`) and q is a PROBABILITY vector
3819    /// supported on `cand_ids` (<=32 ids — the DFlash2 selector's candidate-set proposal;
3820    /// lane/dspark-sampled-admission-20260820). Same event semantics/Philox tag as
3821    /// `residual_sample_filtered` — one uniform per (seed, stream_pos).
3822    #[allow(clippy::too_many_arguments)]
3823    pub fn residual_sample_sparse_q(
3824        &self,
3825        p: &CudaSlice<f32>,
3826        cand_ids: &CudaSlice<u32>,
3827        q_probs: &CudaSlice<f32>,
3828        n_cand: usize,
3829        n: usize,
3830        temp: f32,
3831        seed: u64,
3832        stream_pos: u32,
3833        p_stats: (f32, f32, f32),
3834        out_tok: &mut CudaSlice<u32>,
3835    ) -> Result<(), Box<dyn std::error::Error>> {
3836        assert!(
3837            (1..=32).contains(&n_cand),
3838            "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
3839        );
3840        let f = self.func("residual_sample_sparse_q_f32");
3841        let (ni, nc) = (n as i32, n_cand as i32);
3842        let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3843        let (pm, pth, pz) = p_stats;
3844        let cfg = LaunchConfig {
3845            grid_dim: (1, 1, 1),
3846            block_dim: (1024, 1, 1),
3847            shared_mem_bytes: 0,
3848        };
3849        let __s_b = self.gpu.stream();
3850        let mut b = __s_b.launch_builder(&f);
3851        b.arg(p)
3852            .arg(cand_ids)
3853            .arg(q_probs)
3854            .arg(&nc)
3855            .arg(&ni)
3856            .arg(&temp)
3857            .arg(&slo)
3858            .arg(&shi)
3859            .arg(&stream_pos)
3860            .arg(&pm)
3861            .arg(&pth)
3862            .arg(&pz)
3863            .arg(&mut *out_tok);
3864        unsafe {
3865            b.launch(cfg)?;
3866        }
3867        Ok(())
3868    }
3869
3870    /// Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).
3871    #[allow(clippy::too_many_arguments)]
3872    pub fn gumbel_perturb_filtered(
3873        &self,
3874        x: &CudaSlice<f32>,
3875        y: &mut CudaSlice<f32>,
3876        n: usize,
3877        seed: u64,
3878        stream_pos: u32,
3879        temp: f32,
3880        row_max: f32,
3881        th: f32,
3882    ) -> Result<(), Box<dyn std::error::Error>> {
3883        let f = self.func("gumbel_perturb_filtered_f32");
3884        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3885        let cfg = LaunchConfig {
3886            grid_dim: (n.div_ceil(256) as u32, 1, 1),
3887            block_dim: (256, 1, 1),
3888            shared_mem_bytes: 0,
3889        };
3890        let __s_b = self.gpu.stream();
3891        let mut b = __s_b.launch_builder(&f);
3892        b.arg(x)
3893            .arg(&mut *y)
3894            .arg(&ni)
3895            .arg(&slo)
3896            .arg(&shi)
3897            .arg(&stream_pos)
3898            .arg(&temp)
3899            .arg(&row_max)
3900            .arg(&th);
3901        unsafe {
3902            b.launch(cfg)?;
3903        }
3904        Ok(())
3905    }
3906
3907    /// Keskar penalties applied IN PLACE to a logits buffer: history token ids get
3908    /// rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps
3909    /// filtered rejection sampling exact for the penalized target.
3910    #[allow(clippy::too_many_arguments)]
3911    pub fn penalize_logits(
3912        &self,
3913        x: &mut CudaSlice<f32>,
3914        hist: &CudaSlice<u32>,
3915        n_hist: usize,
3916        rep: f32,
3917        freq: f32,
3918        present: f32,
3919        n: usize,
3920    ) -> Result<(), Box<dyn std::error::Error>> {
3921        if n_hist == 0 {
3922            return Ok(());
3923        }
3924        let f = self.func("penalize_logits_f32");
3925        let (nh, ni) = (n_hist as i32, n as i32);
3926        let cfg = LaunchConfig {
3927            grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
3928            block_dim: (128, 1, 1),
3929            shared_mem_bytes: 0,
3930        };
3931        let __s_b = self.gpu.stream();
3932        let mut b = __s_b.launch_builder(&f);
3933        b.arg(&mut *x)
3934            .arg(hist)
3935            .arg(&nh)
3936            .arg(&rep)
3937            .arg(&freq)
3938            .arg(&present)
3939            .arg(&ni);
3940        unsafe {
3941            b.launch(cfg)?;
3942        }
3943        Ok(())
3944    }
3945
3946    /// Rows variant: penalize `nrow` contiguous rows of length n in one launch.
3947    #[allow(clippy::too_many_arguments)]
3948    pub fn penalize_logits_rows(
3949        &self,
3950        x: &mut CudaSlice<f32>,
3951        hist: &CudaSlice<u32>,
3952        n_hist: usize,
3953        rep: f32,
3954        freq: f32,
3955        present: f32,
3956        n: usize,
3957        nrow: usize,
3958    ) -> Result<(), Box<dyn std::error::Error>> {
3959        if n_hist == 0 || nrow == 0 {
3960            return Ok(());
3961        }
3962        let f = self.func("penalize_logits_rows_f32");
3963        let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
3964        let cfg = LaunchConfig {
3965            grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
3966            block_dim: (128, 1, 1),
3967            shared_mem_bytes: 0,
3968        };
3969        let __s_b = self.gpu.stream();
3970        let mut b = __s_b.launch_builder(&f);
3971        b.arg(&mut *x)
3972            .arg(hist)
3973            .arg(&nh)
3974            .arg(&rep)
3975            .arg(&freq)
3976            .arg(&present)
3977            .arg(&ni)
3978            .arg(&nr);
3979        unsafe {
3980            b.launch(cfg)?;
3981        }
3982        Ok(())
3983    }
3984
3985    /// Heterogeneous serving-batch penalties over host-maintained sparse window counts.
3986    /// `offsets[r]..offsets[r+1]` indexes the unique positive-count `(id,count)` entries for logits row
3987    /// `rows[r]`; each row may carry independent repetition/frequency/presence coefficients.
3988    /// One thread owns one distinct logit, so the kernel needs neither atomics nor the
3989    /// history-squared dedup scan used by the speculative raw-history oracle.
3990    #[allow(clippy::too_many_arguments)]
3991    pub fn penalize_logits_sparse_rows(
3992        &self,
3993        x: &mut CudaSlice<f32>,
3994        ids: &[u32],
3995        counts: &[u32],
3996        offsets: &[i32],
3997        rows: &[i32],
3998        reps: &[f32],
3999        freqs: &[f32],
4000        presents: &[f32],
4001        n: usize,
4002    ) -> Result<(), Box<dyn std::error::Error>> {
4003        let nrow = rows.len();
4004        if nrow == 0 {
4005            return Ok(());
4006        }
4007        let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
4008        let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
4009        let entry_count =
4010            i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
4011        if ids.len() != counts.len()
4012            || offsets.len() != nrow + 1
4013            || reps.len() != nrow
4014            || freqs.len() != nrow
4015            || presents.len() != nrow
4016            || offsets.first().copied() != Some(0)
4017            || offsets.last().copied() != Some(entry_count)
4018        {
4019            return Err("sparse penalty row metadata shape mismatch".into());
4020        }
4021        if counts.contains(&0) {
4022            return Err("sparse penalty counts must be positive".into());
4023        }
4024        let mut max_len = 0usize;
4025        for pair in offsets.windows(2) {
4026            if pair[0] < 0 || pair[1] < pair[0] {
4027                return Err("sparse penalty offsets must be monotonic".into());
4028            }
4029            max_len = max_len.max((pair[1] - pair[0]) as usize);
4030        }
4031        if max_len == 0 {
4032            return Ok(());
4033        }
4034
4035        let mut seen = std::collections::HashSet::with_capacity(ids.len());
4036        for (r, &row) in rows.iter().enumerate() {
4037            if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
4038                return Err("sparse penalty row index exceeds logits shape".into());
4039            }
4040            let begin = offsets[r] as usize;
4041            let end = offsets[r + 1] as usize;
4042            for &id in &ids[begin..end] {
4043                if id as usize >= n {
4044                    return Err("sparse penalty token id exceeds logits row".into());
4045                }
4046                if !seen.insert((row, id)) {
4047                    return Err("sparse penalty entries must be unique per logits row".into());
4048                }
4049            }
4050        }
4051
4052        // SAFETY: the checks above establish every invariant of the launch-only helper.
4053        unsafe {
4054            self.penalize_logits_sparse_rows_unchecked(
4055                x, ids, counts, offsets, rows, reps, freqs, presents, n,
4056            )
4057        }
4058    }
4059
4060    /// Launch-only form for the serving hot path, whose `HashMap`-backed producer already
4061    /// guarantees unique ids and whose rows are enumerated from the live batch.
4062    ///
4063    /// # Safety
4064    ///
4065    /// Shapes must match the safe wrapper, offsets must be monotonic and in bounds, every row
4066    /// must index `x`, and each `(row,id)` pair must occur at most once. Token ids outside the
4067    /// logits row are safe no-ops because the kernel bounds-checks them before computing `x`.
4068    #[allow(clippy::too_many_arguments)]
4069    pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
4070        &self,
4071        x: &mut CudaSlice<f32>,
4072        ids: &[u32],
4073        counts: &[u32],
4074        offsets: &[i32],
4075        rows: &[i32],
4076        reps: &[f32],
4077        freqs: &[f32],
4078        presents: &[f32],
4079        n: usize,
4080    ) -> Result<(), Box<dyn std::error::Error>> {
4081        let nrow = rows.len();
4082        if nrow == 0 {
4083            return Ok(());
4084        }
4085        let max_len = offsets
4086            .windows(2)
4087            .map(|pair| (pair[1] - pair[0]) as usize)
4088            .max()
4089            .unwrap_or(0);
4090        if max_len == 0 {
4091            return Ok(());
4092        }
4093        let ids_d = self.htod_u32_v(ids)?;
4094        let counts_d = self.htod_u32_v(counts)?;
4095        let offsets_d = self.htod_i32(offsets)?;
4096        let rows_d = self.htod_i32(rows)?;
4097        let reps_d = self.htod(reps)?;
4098        let freqs_d = self.htod(freqs)?;
4099        let presents_d = self.htod(presents)?;
4100        let f = self.func("penalize_logits_sparse_rows_f32");
4101        let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
4102        let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
4103        let cfg = LaunchConfig {
4104            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
4105            block_dim: (128, 1, 1),
4106            shared_mem_bytes: 0,
4107        };
4108        let __s_b = self.gpu.stream();
4109        let mut b = __s_b.launch_builder(&f);
4110        b.arg(&mut *x)
4111            .arg(&ids_d)
4112            .arg(&counts_d)
4113            .arg(&offsets_d)
4114            .arg(&rows_d)
4115            .arg(&reps_d)
4116            .arg(&freqs_d)
4117            .arg(&presents_d)
4118            .arg(&ni)
4119            .arg(&nr);
4120        unsafe {
4121            b.launch(cfg)?;
4122        }
4123        Ok(())
4124    }
4125
4126    /// ROW-INCREMENTAL penalties (dspark penalized-sampled admission): row r of `x`
4127    /// penalizes over the last `min(win, n_hist0 + r)` entries of `hist[..n_hist0 + r]`,
4128    /// where `hist` = [session window (n_hist0) ++ per-row drafted tokens (nrow-1)]. This
4129    /// is the within-round evolving penalty state block drafting needs: verify row r's
4130    /// target is penalized by every token committed before it INCLUDING same-round
4131    /// accepts — `penalize_logits_rows` (one shared window) is the frozen-window
4132    /// approximation this exists to replace on the dspark route.
4133    #[allow(clippy::too_many_arguments)]
4134    pub fn penalize_logits_rows_inc(
4135        &self,
4136        x: &mut CudaSlice<f32>,
4137        hist: &CudaSlice<u32>,
4138        n_hist0: usize,
4139        rep: f32,
4140        freq: f32,
4141        present: f32,
4142        n: usize,
4143        nrow: usize,
4144        win: usize,
4145    ) -> Result<(), Box<dyn std::error::Error>> {
4146        if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
4147            return Ok(());
4148        }
4149        debug_assert!(
4150            hist.len() >= n_hist0 + nrow - 1,
4151            "rows-inc hist must carry n_hist0 + nrow - 1 ids"
4152        );
4153        let f = self.func("penalize_logits_rows_inc_f32");
4154        let max_len = win.min(n_hist0 + nrow - 1).max(1);
4155        let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
4156        let cfg = LaunchConfig {
4157            grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
4158            block_dim: (128, 1, 1),
4159            shared_mem_bytes: 0,
4160        };
4161        let __s_b = self.gpu.stream();
4162        let mut b = __s_b.launch_builder(&f);
4163        b.arg(&mut *x)
4164            .arg(hist)
4165            .arg(&nh)
4166            .arg(&rep)
4167            .arg(&freq)
4168            .arg(&present)
4169            .arg(&ni)
4170            .arg(&nr)
4171            .arg(&wi);
4172        unsafe {
4173            b.launch(cfg)?;
4174        }
4175        Ok(())
4176    }
4177
4178    /// WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle
4179    /// window (the fa launch reads KV, not weights) prefetch the NEXT matvec's
4180    /// decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same
4181    /// class as prefetch_l2 (numerics untouched by construction). Wired only where it
4182    /// measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates),
4183    /// 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer
4184    /// floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.
4185    pub fn wpf_level() -> u32 {
4186        static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
4187        *ON.get_or_init(|| {
4188            std::env::var("MEMRA_WPF")
4189                .ok()
4190                .and_then(|v| v.parse().ok())
4191                .unwrap_or(1)
4192        })
4193    }
4194
4195    /// PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue
4196    /// kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the
4197    /// grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the
4198    /// kernels' entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes
4199    /// every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens);
4200    /// 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3
4201    /// gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under
4202    /// capture (capture encodes native programmatic edges — the post-capture edge-REWRITE
4203    /// arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return
4204    /// CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam.
4205    /// See the `verify_exact` field. Scoped by the dflash round around its t=16 verify.
4206    /// Prefer `exact_scope` — the RAII form — anywhere a `?` can exit the scope: a manual
4207    /// true/false pair leaves the flag LATCHED engine-wide when an error propagates
4208    /// between the two calls (hermes finding on dspark_spec_session_burst, fixed
4209    /// 2026-08-23), and every later request then runs the exact-GEMM program.
4210    pub fn set_verify_exact(&self, on: bool) {
4211        self.verify_exact
4212            .store(on, std::sync::atomic::Ordering::Relaxed);
4213    }
4214    pub(crate) fn verify_exact_on(&self) -> bool {
4215        self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
4216    }
4217
4218    /// RAII scope over `verify_exact`: sets the flag to `on` now and restores the
4219    /// PREVIOUS value on drop — unwind, early `return`, and every `?` exit included.
4220    /// This is the required form for any scope an error can leave (see
4221    /// `set_verify_exact`); dropping the guard early (`drop(scope)`) ends the scope
4222    /// exactly where the manual `set_verify_exact(false)` used to sit.
4223    pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
4224        ExactScope::set(&self.verify_exact, on)
4225    }
4226
4227    /// m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the
4228    /// fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).
4229    pub fn qkv_append_on() -> bool {
4230        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4231        *ON.get_or_init(|| {
4232            std::env::var("MEMRA_QKV_APPEND")
4233                .map(|v| v != "0")
4234                .unwrap_or(true)
4235        })
4236    }
4237
4238    /// PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32,
4239    /// add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.
4240    pub fn pdl_wb_on() -> bool {
4241        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4242        *ON.get_or_init(|| {
4243            std::env::var("MEMRA_PDL_WB")
4244                .map(|v| v != "0")
4245                .unwrap_or(true)
4246        })
4247    }
4248
4249    /// Trunk-kernels norm ILP seam (lane/dspark-trunk-kernels-20260820): the T-row verify
4250    /// norms (rms_norm_f32 / add_rms_norm_f32 at grid=T, block=256) are serial-latency
4251    /// chains — 20 strided scalar load->fma rounds measured 11.8-12.2us/inst (nsys-B verify
4252    /// scope: 130 inst/rd = 1.51 ms/rd). The `_v2` twins unroll the element loop 4-deep
4253    /// (independent loads in flight; SAME per-thread element order into ONE accumulator,
4254    /// reduce VERBATIM) — BIT-IDENTICAL per row at every (ncols, blockDim).
4255    /// MEMRA_NORM_ILP=0 reverts to the v1 kernels alone.
4256    pub fn norm_ilp_on() -> bool {
4257        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4258        *ON.get_or_init(|| {
4259            std::env::var("MEMRA_NORM_ILP")
4260                .map(|v| v != "0")
4261                .unwrap_or(true)
4262        })
4263    }
4264
4265    /// `MEMRA_NORM_ILP_ZQ8=0` reverts the `rms_norm_zq8_f32_v2` twin ALONE (default ON under
4266    /// `MEMRA_NORM_ILP`; lane/glm5-norm-zq8-ilp-20260904). The per-kernel seam exists so the
4267    /// box can price this twin against the v1 kernel on ONE binary with every other norm twin
4268    /// held, and so a rollback of this kernel does not drag `rms_norm_f32_v2` with it.
4269    pub fn norm_ilp_zq8_on() -> bool {
4270        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4271        *ON.get_or_init(|| {
4272            Self::norm_ilp_on()
4273                && std::env::var("MEMRA_NORM_ILP_ZQ8")
4274                    .map(|v| v != "0")
4275                    .unwrap_or(true)
4276        })
4277    }
4278
4279    /// Trunk-kernels FFN dual seam (lane/dspark-trunk-kernels-20260820): the qwen35
4280    /// t-parallel verify FFN pair rides the PROVEN dual gate+up doors
4281    /// (`matmul_decode_exact_dual_pre` + `silu_mul_scaled_q8_1`, the q27 verify shape —
4282    /// bit-identical per (tensor,token,row), kernel-check-pinned, MEMRA_SPEC_DUAL_T
4283    /// receipts) instead of two singles + silu_mul + a standalone quantize. The doors
4284    /// existed but the qwen35 body never called them (nsys-B verify scope: gate+up singles
4285    /// = 107 launches/rd at grid 4352). MEMRA_TK_FFN_DUAL=0 reverts to the singles chain.
4286    pub fn tk_ffn_dual_on() -> bool {
4287        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4288        *ON.get_or_init(|| {
4289            std::env::var("MEMRA_TK_FFN_DUAL")
4290                .map(|v| v != "0")
4291                .unwrap_or(true)
4292        })
4293    }
4294
4295    /// PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep
4296    /// their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the
4297    /// per-model no-harm bisect knob.
4298    pub fn pdl_mmvq_on() -> bool {
4299        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4300        *ON.get_or_init(|| {
4301            std::env::var("MEMRA_PDL_MMVQ")
4302                .map(|v| v != "0")
4303                .unwrap_or(true)
4304        })
4305    }
4306
4307    pub fn pdl_on() -> bool {
4308        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4309        *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
4310    }
4311
4312    /// PDL wave-B seam (gap-diagnosis arc, GAP-DIAGNOSIS.md verdict 8): the gemma
4313    /// NVFP4mix decode chain's hot kernels — nvfp4 fused2/mr2 and the q8_0 `_rp`
4314    /// singles — join the wave-A launch class. Scheduling-only (the entry macro waits
4315    /// on the producer before any read), bit-identical by construction.
4316    /// MEMRA_PDL_NVFP4=0 reverts wave-B alone.
4317    pub fn pdl_nvfp4q8_on() -> bool {
4318        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4319        *ON.get_or_init(|| {
4320            std::env::var("MEMRA_PDL_NVFP4")
4321                .map(|v| v != "0")
4322                .unwrap_or(true)
4323        })
4324    }
4325
4326    /// Raw CUfunction for a PDL-attributed launch: the SAME kernels.fatbin loaded once more
4327    /// through the raw driver API (cudarc hides its CUfunction handles; a duplicate module
4328    /// of tiny glue kernels is free). Resolved lazily per name, cached process-wide.
4329    /// Fused t=1 q4_0 mr policy: env MEMRA_Q40_MR wins (1/2); else the per-model
4330    /// FUSED_MR1_DEFAULT (dense gemma = mr1, MoE = mr2 — see the static's doc).
4331    fn q40_mr1_on() -> bool {
4332        static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
4333        match *Q40MR.get_or_init(|| {
4334            std::env::var("MEMRA_Q40_MR")
4335                .ok()
4336                .and_then(|v| v.parse().ok())
4337        }) {
4338            Some(v) => v == 1,
4339            None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4340        }
4341    }
4342
4343    /// PDL wave-B2: flash-module PDL functions. `g` selects the kf8vf8 flavor — the
4344    /// caller MUST pass the SAME flavor its builder launch would resolve (fa_func/func_g
4345    /// mirror); the flavors differ semantically (KV byte formats), a wrong-module launch
4346    /// writes wrong bytes silently.
4347    fn pdl_func_flash(
4348        &self,
4349        g: bool,
4350        name: &'static str,
4351    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4352        use cudarc::driver::sys as cu;
4353        // PER-CONTEXT caches (M1-PP2 cross-device fix, 8x box 2026-08-02): CUmodule and
4354        // CUfunction handles are CONTEXT-scoped, and a remote-stage Engine
4355        // (MEMRA_PP_DEVICES=a,b) lives in the other device's primary context. The old
4356        // process-wide OnceLock cache handed stage 1 the dev-a handles, so every stage-1
4357        // launch_pdl* died CUDA_ERROR_INVALID_HANDLE. Key module + function caches by
4358        // this engine's CUcontext; single-context runs behave exactly as before.
4359        static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
4360            std::sync::Mutex::new(None);
4361        #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4362        static FNS: std::sync::Mutex<
4363            Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
4364        > = std::sync::Mutex::new(None);
4365        let ctx_key = self.ctx().cu_ctx() as usize;
4366        if let Some(&f) = FNS
4367            .lock()
4368            .unwrap()
4369            .get_or_insert_with(Default::default)
4370            .get(&(ctx_key, g, name))
4371        {
4372            return Ok(f as cu::CUfunction);
4373        }
4374        let module = {
4375            let mut mods = MODS.lock().unwrap();
4376            let map = mods.get_or_insert_with(Default::default);
4377            match map.get(&(ctx_key, g)) {
4378                Some(&m) => m,
4379                None => {
4380                    let m = self.pdl_load_module_in_ctx(if g {
4381                        FLASH_FATBIN_KF8VF8
4382                    } else {
4383                        FLASH_FATBIN
4384                    })?;
4385                    map.insert((ctx_key, g), m);
4386                    m
4387                }
4388            }
4389        };
4390        let cname = std::ffi::CString::new(name)?;
4391        let mut f: cu::CUfunction = std::ptr::null_mut();
4392        let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
4393        if r != cu::CUresult::CUDA_SUCCESS {
4394            return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
4395        }
4396        FNS.lock()
4397            .unwrap()
4398            .get_or_insert_with(Default::default)
4399            .insert((ctx_key, g, name), f as usize);
4400        Ok(f)
4401    }
4402
4403    /// Load a fatbin as a raw CUmodule IN THIS ENGINE'S CONTEXT. `cuModuleLoadData` binds
4404    /// the module to the thread's CURRENT context — a remote-stage engine must not
4405    /// inherit the primary's (the INVALID_HANDLE class above). Restores the caller's
4406    /// current context before returning.
4407    fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
4408        use cudarc::driver::sys as cu;
4409        let mut prev: cu::CUcontext = std::ptr::null_mut();
4410        unsafe {
4411            cu::cuCtxGetCurrent(&mut prev).result()?;
4412        }
4413        self.ctx().bind_to_thread()?;
4414        let mut m: cu::CUmodule = std::ptr::null_mut();
4415        let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
4416        let restore = if prev.is_null() {
4417            cu::CUresult::CUDA_SUCCESS
4418        } else {
4419            unsafe { cu::cuCtxSetCurrent(prev) }
4420        };
4421        if r != cu::CUresult::CUDA_SUCCESS {
4422            return Err(format!("pdl module load: {r:?}").into());
4423        }
4424        if restore != cu::CUresult::CUDA_SUCCESS {
4425            return Err(format!("pdl module load: ctx restore {restore:?}").into());
4426        }
4427        Ok(m as usize)
4428    }
4429
4430    /// Raw CUfunction for prebuilt-args dispatch experiments (M4 probe): same duplicate
4431    /// raw-module loading as the PDL path, WITHOUT the PDL launch attribute.
4432    pub fn raw_kernel_function(
4433        &self,
4434        name: &'static str,
4435    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4436        self.pdl_func(name)
4437    }
4438
4439    fn pdl_func(
4440        &self,
4441        name: &'static str,
4442    ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4443        use cudarc::driver::sys as cu;
4444        // PER-CONTEXT caches — same M1-PP2 cross-device fix as pdl_func_flash (handles
4445        // are context-scoped; key everything by this engine's CUcontext).
4446        static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
4447            std::sync::Mutex::new(None);
4448        // PDL wave-A: the mmvq kernels live in the qmatvec fatbin, not kernels.cu — second
4449        // duplicate module, loaded lazily on the first kernels-module miss.
4450        static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
4451            std::sync::Mutex::new(None);
4452        static FNS: std::sync::Mutex<
4453            Option<std::collections::HashMap<(usize, &'static str), usize>>,
4454        > = std::sync::Mutex::new(None);
4455        let ctx_key = self.ctx().cu_ctx() as usize;
4456        if let Some(&f) = FNS
4457            .lock()
4458            .unwrap()
4459            .get_or_insert_with(Default::default)
4460            .get(&(ctx_key, name))
4461        {
4462            return Ok(f as cu::CUfunction);
4463        }
4464        let module = {
4465            let mut mods = MODULES.lock().unwrap();
4466            let map = mods.get_or_insert_with(Default::default);
4467            match map.get(&ctx_key) {
4468                Some(&m) => m,
4469                None => {
4470                    let m = self.pdl_load_module_in_ctx(FATBIN)?;
4471                    map.insert(ctx_key, m);
4472                    m
4473                }
4474            }
4475        };
4476        let cname = std::ffi::CString::new(name)?;
4477        let mut f: cu::CUfunction = std::ptr::null_mut();
4478        let mut r =
4479            unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
4480        if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
4481            let qmodule = {
4482                let mut mods = QMODULES.lock().unwrap();
4483                let map = mods.get_or_insert_with(Default::default);
4484                match map.get(&ctx_key) {
4485                    Some(&m) => m,
4486                    None => {
4487                        let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
4488                        map.insert(ctx_key, m);
4489                        m
4490                    }
4491                }
4492            };
4493            r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
4494        }
4495        if r != cu::CUresult::CUDA_SUCCESS {
4496            return Err(format!("pdl_func {name}: {r:?}").into());
4497        }
4498        FNS.lock()
4499            .unwrap()
4500            .get_or_insert_with(Default::default)
4501            .insert((ctx_key, name), f as usize);
4502        Ok(f)
4503    }
4504
4505    /// cuLaunchKernelEx with CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION on the
4506    /// compute stream. ONLY legal for kernels whose entry carries MEMRA_PDL_ENTRY.
4507    ///
4508    /// # Safety
4509    /// `params` must match the kernel's exact parameter list (order, types, count) —
4510    /// a mismatch corrupts the launch silently.
4511    /// Flash-module twin of `launch_pdl` — `g` picks the kf8vf8 flavor (must mirror the
4512    /// builder path's fa_func/func_g choice exactly).
4513    ///
4514    /// # Safety
4515    /// Same contract as `launch_pdl`.
4516    unsafe fn launch_pdl_flash(
4517        &self,
4518        g: bool,
4519        name: &'static str,
4520        grid: (u32, u32, u32),
4521        block: (u32, u32, u32),
4522        smem: u32,
4523        params: &mut [*mut std::ffi::c_void],
4524    ) -> Result<(), Box<dyn std::error::Error>> {
4525        use cudarc::driver::sys as cu;
4526        let f = self.pdl_func_flash(g, name)?;
4527        if smem > 0 {
4528            // mirror the builder path's opt-in ceiling (idempotent host-side set).
4529            let r =
4530                unsafe {
4531                    cu::cuFuncSetAttribute(f,
4532                cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
4533                smem as i32)
4534                };
4535            if r != cu::CUresult::CUDA_SUCCESS {
4536                return Err(format!("pdl smem attr {name}: {r:?}").into());
4537            }
4538        }
4539        let mut attr = cu::CUlaunchAttribute {
4540            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
4541            pad: [0; 4],
4542            value: cu::CUlaunchAttributeValue {
4543                programmaticStreamSerializationAllowed: 1,
4544            },
4545        };
4546        let cfg = cu::CUlaunchConfig {
4547            gridDimX: grid.0,
4548            gridDimY: grid.1,
4549            gridDimZ: grid.2,
4550            blockDimX: block.0,
4551            blockDimY: block.1,
4552            blockDimZ: block.2,
4553            sharedMemBytes: smem,
4554            hStream: self.gpu.stream().cu_stream(),
4555            attrs: &mut attr,
4556            numAttrs: 1,
4557        };
4558        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
4559        if r != cu::CUresult::CUDA_SUCCESS {
4560            return Err(format!("launch_pdl_flash {name}: {r:?}").into());
4561        }
4562        Ok(())
4563    }
4564
4565    unsafe fn launch_pdl(
4566        &self,
4567        name: &'static str,
4568        grid: (u32, u32, u32),
4569        block: (u32, u32, u32),
4570        params: &mut [*mut std::ffi::c_void],
4571    ) -> Result<(), Box<dyn std::error::Error>> {
4572        use cudarc::driver::sys as cu;
4573        let f = self.pdl_func(name)?;
4574        let mut attr = cu::CUlaunchAttribute {
4575            id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
4576            pad: [0; 4],
4577            value: cu::CUlaunchAttributeValue {
4578                programmaticStreamSerializationAllowed: 1,
4579            },
4580        };
4581        let cfg = cu::CUlaunchConfig {
4582            gridDimX: grid.0,
4583            gridDimY: grid.1,
4584            gridDimZ: grid.2,
4585            blockDimX: block.0,
4586            blockDimY: block.1,
4587            blockDimZ: block.2,
4588            sharedMemBytes: 0,
4589            hStream: self.gpu.stream().cu_stream(),
4590            attrs: &mut attr,
4591            numAttrs: 1,
4592        };
4593        let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
4594        if r != cu::CUresult::CUDA_SUCCESS {
4595            return Err(format!("launch_pdl {name}: {r:?}").into());
4596        }
4597        Ok(())
4598    }
4599
4600    /// L2-prefetch a quant weight's DECODE plane (the rp4 split-plane mirror when present —
4601    /// that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.
4602    pub fn prefetch_weight_l2(
4603        &self,
4604        w: &crate::model::GpuTensor,
4605    ) -> Result<(), Box<dyn std::error::Error>> {
4606        if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
4607            let p = rp4.as_ref().unwrap_or(bytes);
4608            self.prefetch_l2(p, p.len())?;
4609        }
4610        Ok(())
4611    }
4612
4613    /// DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table
4614    /// by the DEVICE token id at tok[idx] into f32.
4615    pub fn gather_row_bf16(
4616        &self,
4617        table: &CudaSlice<u8>,
4618        tok: &CudaSlice<u32>,
4619        idx: usize,
4620        dst: &mut CudaSlice<f32>,
4621        ncols: usize,
4622    ) -> Result<(), Box<dyn std::error::Error>> {
4623        let f = self.func("gather_row_bf16_f32");
4624        let cfg = LaunchConfig {
4625            grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
4626            block_dim: (256, 1, 1),
4627            shared_mem_bytes: 0,
4628        };
4629        let (nc, ix) = (ncols as i32, idx as i32);
4630        let __s_b = self.gpu.stream();
4631        let mut b = __s_b.launch_builder(&f);
4632        b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
4633        unsafe {
4634            b.launch(cfg)?;
4635        }
4636        Ok(())
4637    }
4638
4639    /// DFlash2 grouped dynamic causal conv (dflash lane, DFLASH2-EVAL-20260820.md):
4640    /// out[p,c] = sum_{o<ksize, o<=p} (base[half][o][c] + dyn[p][half][o][group(c)])
4641    /// * x[p-o][c]. `dyn_` is the kernel_projection GEMM output [rows, 2*ksize*groups];
4642    ///   `base` is base_kernel [2, ksize, hidden] flattened; `half` picks prepare(0) /
4643    ///   finish(1).
4644    #[allow(clippy::too_many_arguments)]
4645    pub fn dflash2_dynconv(
4646        &self,
4647        x: &CudaSlice<f32>,
4648        dyn_: &CudaSlice<f32>,
4649        base: &CudaSlice<f32>,
4650        out: &mut CudaSlice<f32>,
4651        rows: usize,
4652        hidden: usize,
4653        group_size: usize,
4654        ksize: usize,
4655        half: usize,
4656    ) -> Result<(), Box<dyn std::error::Error>> {
4657        assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
4658        let f = self.func("dflash2_dynconv_f32");
4659        let n = rows * hidden;
4660        let cfg = LaunchConfig {
4661            grid_dim: (n.div_ceil(256) as u32, 1, 1),
4662            block_dim: (256, 1, 1),
4663            shared_mem_bytes: 0,
4664        };
4665        let (ri, hi, gi, ki, hf) = (
4666            rows as i32,
4667            hidden as i32,
4668            group_size as i32,
4669            ksize as i32,
4670            half as i32,
4671        );
4672        let __s_b = self.gpu.stream();
4673        let mut b = __s_b.launch_builder(&f);
4674        b.arg(x)
4675            .arg(dyn_)
4676            .arg(base)
4677            .arg(out)
4678            .arg(&ri)
4679            .arg(&hi)
4680            .arg(&gi)
4681            .arg(&ki)
4682            .arg(&hf);
4683        unsafe {
4684            b.launch(cfg)?;
4685        }
4686        Ok(())
4687    }
4688
4689    /// Per-row top-k (k <= 32) over a [n_rows, n_cols] logits matrix (DFlash2
4690    /// candidate selector). Returns (values [n_rows, k], column indices [n_rows, k]),
4691    /// value-descending, ties to the lower index.
4692    pub fn topk_rows(
4693        &self,
4694        logits: &CudaSlice<f32>,
4695        n_rows: usize,
4696        n_cols: usize,
4697        k: usize,
4698    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
4699        assert!((1..=32).contains(&k), "topk_rows supports 1..=32, got {k}");
4700        assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
4701        // MEMRA_TOPK_SHARDS (lane/glm5-matvec door K, default ON since 2026-08-31): the exact two-launch
4702        // shard split — n_rows*16 partial blocks + a per-row merge — instead of n_rows
4703        // blocks total (the DFlash2 selector: 15 blocks on the whole card, 7 GB/s). Top-k
4704        // under (value desc, column asc) is discrete selection: output-identical by
4705        // construction, gated by glm5_matvec_doors_gpu. Small columns fall through (the
4706        // shard overhead would dominate and the standing grid is already wide enough).
4707        if topk_shards_on() && n_cols >= 16 * 1024 && k <= n_cols / 16 {
4708            if TOPK_SHARDS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
4709                eprintln!(
4710                    "[topk-shards] engaged: rows={n_rows} cols={n_cols} k={k} shards=16 \
4711                     (MEMRA_TOPK_SHARDS=1)"
4712                );
4713            }
4714            return self.topk_rows_sharded(logits, n_rows, n_cols, k, 16);
4715        }
4716        let f = self.func("topk_rows_f32");
4717        let nth = 256usize;
4718        let mut vals = self.uninit(n_rows * k)?;
4719        let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
4720        let cfg = LaunchConfig {
4721            grid_dim: (n_rows as u32, 1, 1),
4722            block_dim: (nth as u32, 1, 1),
4723            shared_mem_bytes: (nth * k * 8) as u32,
4724        };
4725        let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
4726        let __s_b = self.gpu.stream();
4727        let mut b = __s_b.launch_builder(&f);
4728        b.arg(logits)
4729            .arg(&nr)
4730            .arg(&nc)
4731            .arg(&ki)
4732            .arg(&mut vals)
4733            .arg(&mut idxs);
4734        unsafe {
4735            b.launch(cfg)?;
4736        }
4737        Ok((vals, idxs))
4738    }
4739
4740    /// The exact two-launch shard split behind `MEMRA_TOPK_SHARDS` (see [`Self::topk_rows`]):
4741    /// per-(row, shard) partial top-k with the standing kernel's insertion/tie rules on
4742    /// global column indices, then a per-row k-way merge with the standing kernel's merge
4743    /// rules. Output-identical to `topk_rows_f32` by construction (discrete selection under
4744    /// the total order value-desc/index-asc); gated by `glm5_matvec_doors_gpu`.
4745    fn topk_rows_sharded(
4746        &self,
4747        logits: &CudaSlice<f32>,
4748        n_rows: usize,
4749        n_cols: usize,
4750        k: usize,
4751        n_shards: usize,
4752    ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
4753        assert!((1..=64).contains(&n_shards), "shard merge head cap is 64");
4754        let nth = 256usize;
4755        let mut pvals = self.uninit(n_rows * n_shards * k)?;
4756        let mut pidxs = self.alloc_uninit::<u32>(n_rows * n_shards * k)?;
4757        let f1 = self.func("topk_rows_shard_f32");
4758        let cfg1 = LaunchConfig {
4759            grid_dim: (n_rows as u32, n_shards as u32, 1),
4760            block_dim: (nth as u32, 1, 1),
4761            shared_mem_bytes: (nth * k * 8) as u32,
4762        };
4763        let (nr, nc, ki, ns) = (n_rows as i32, n_cols as i32, k as i32, n_shards as i32);
4764        {
4765            let __s_b = self.gpu.stream();
4766            let mut b = __s_b.launch_builder(&f1);
4767            b.arg(logits)
4768                .arg(&nr)
4769                .arg(&nc)
4770                .arg(&ki)
4771                .arg(&ns)
4772                .arg(&mut pvals)
4773                .arg(&mut pidxs);
4774            unsafe {
4775                b.launch(cfg1)?;
4776            }
4777        }
4778        let mut vals = self.uninit(n_rows * k)?;
4779        let mut idxs = self.alloc_uninit::<u32>(n_rows * k)?;
4780        let f2 = self.func("topk_rows_shard_merge_f32");
4781        let cfg2 = LaunchConfig {
4782            grid_dim: (n_rows as u32, 1, 1),
4783            block_dim: (32, 1, 1),
4784            shared_mem_bytes: 0,
4785        };
4786        let __s_b = self.gpu.stream();
4787        let mut b = __s_b.launch_builder(&f2);
4788        b.arg(&pvals)
4789            .arg(&pidxs)
4790            .arg(&nr)
4791            .arg(&ns)
4792            .arg(&ki)
4793            .arg(&mut vals)
4794            .arg(&mut idxs);
4795        unsafe {
4796            b.launch(cfg2)?;
4797        }
4798        Ok((vals, idxs))
4799    }
4800
4801    /// logits[row_off .. row_off+n] += bias[0..n] (in place, one row).
4802    pub fn add_row_inplace(
4803        &self,
4804        logits: &mut CudaSlice<f32>,
4805        bias: &CudaSlice<f32>,
4806        n: usize,
4807        row_off: usize,
4808    ) -> Result<(), Box<dyn std::error::Error>> {
4809        let f = self.func("add_row_inplace_f32");
4810        let cfg = LaunchConfig {
4811            grid_dim: (n.div_ceil(256) as u32, 1, 1),
4812            block_dim: (256, 1, 1),
4813            shared_mem_bytes: 0,
4814        };
4815        let (ni, off) = (n as i32, row_off as i64);
4816        let __s_b = self.gpu.stream();
4817        let mut b = __s_b.launch_builder(&f);
4818        b.arg(logits).arg(bias).arg(&ni).arg(&off);
4819        unsafe {
4820            b.launch(cfg)?;
4821        }
4822        Ok(())
4823    }
4824
4825    /// L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).
4826    pub fn prefetch_l2(
4827        &self,
4828        p: &CudaSlice<u8>,
4829        n: usize,
4830    ) -> Result<(), Box<dyn std::error::Error>> {
4831        let f = self.func("prefetch_l2_bytes");
4832        let lines = n.div_ceil(128);
4833        let ni = n as i64;
4834        let cfg = LaunchConfig {
4835            grid_dim: (lines.div_ceil(256) as u32, 1, 1),
4836            block_dim: (256, 1, 1),
4837            shared_mem_bytes: 0,
4838        };
4839        let __s_b = self.gpu.stream();
4840        let mut b = __s_b.launch_builder(&f);
4841        b.arg(p).arg(&ni);
4842        unsafe {
4843            b.launch(cfg)?;
4844        }
4845        Ok(())
4846    }
4847
4848    /// MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot.
4849    /// Different FP order than the cuBLAS path it replaces — battery-gated numeric config.
4850    pub fn router_gemv(
4851        &self,
4852        w: &CudaSlice<f32>,
4853        x: &CudaSlice<f32>,
4854        n_embd: usize,
4855        n_experts: usize,
4856        t: usize,
4857    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4858        // float4 v2 probed 2026-07-14: +0.25% but flips near-tie routing (new FP order,
4859        // stream differs) — too small to justify a numeric config change; deleted.
4860        // w8 twin (2026-07-31): on the 132-SM H100 the lone-warp form is 14.8% of the q35
4861        // decode step (latency-bound) — the calculus flipped. MEMRA_ROUTER_V2=0 reverts to
4862        // the warp form (rollback seam; new FP order, battery-arbitrated per model).
4863        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
4864            Ok("0") => false,
4865            Ok(_) => true,
4866            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4867        };
4868        // FAST-ROUTER batch twin (lane/fast-router, 2026-08-02): at prefill m the per-(e,tok)
4869        // w8 form re-streams both operand rows per output (GEMV program at GEMM shape — the
4870        // concat-prime exactness fix paid -10% q35 board-2048 prefill through it). The batch
4871        // twin (8x8 expert-x-token register tile) is BIT-IDENTICAL per row (same k order,
4872        // same tree, same fold — kernel-check sweeps m=1..2048 on real router weights), so
4873        // the crossover is pure perf, not a numeric config. MIN_T from the on-box sweep
4874        // (research/fast-router-20260802/crossover-router*.jsonl); decode t=1 and small-t
4875        // spec verify keep the plain w8 form. MEMRA_ROUTER_BATCH=0: rollback seam
4876        // (perf-only, bits equal).
4877        let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
4878        self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
4879    }
4880
4881    /// Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench
4882    /// force both forms; `batch` requires `w8`).
4883    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
4884    pub fn router_gemv_form(
4885        &self,
4886        w: &CudaSlice<f32>,
4887        x: &CudaSlice<f32>,
4888        n_embd: usize,
4889        n_experts: usize,
4890        t: usize,
4891        w8: bool,
4892        batch: bool,
4893    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4894        debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
4895        let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
4896        let f = if batch {
4897            self.func("router_gemv_f32_w8_batch")
4898        } else if w8 {
4899            self.func("router_gemv_f32_w8")
4900        } else {
4901            self.func("router_gemv_f32")
4902        };
4903        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
4904        let cfg = if batch {
4905            LaunchConfig {
4906                grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
4907                block_dim: (32, 8, 1),
4908                shared_mem_bytes: 0,
4909            }
4910        } else {
4911            LaunchConfig {
4912                grid_dim: (n_experts as u32, t as u32, 1),
4913                block_dim: (32, if w8 { 8 } else { 1 }, 1),
4914                shared_mem_bytes: 0,
4915            }
4916        };
4917        let __s_b = self.gpu.stream();
4918        let mut b = __s_b.launch_builder(&f);
4919        b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
4920        unsafe {
4921            b.launch(cfg)?;
4922        }
4923        Ok(y)
4924    }
4925
4926    /// `router_gemv` (decode form selection) writing into a caller-owned [t*n_experts]
4927    /// buffer — token-graph alloc-free.
4928    pub fn router_gemv_into(
4929        &self,
4930        w: &CudaSlice<f32>,
4931        x: &CudaSlice<f32>,
4932        y: &mut CudaSlice<f32>,
4933        n_embd: usize,
4934        n_experts: usize,
4935        t: usize,
4936    ) -> Result<(), Box<dyn std::error::Error>> {
4937        if y.len() < t * n_experts {
4938            return Err("router_gemv_into output too small".into());
4939        }
4940        let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
4941            Ok("0") => false,
4942            Ok(_) => true,
4943            Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4944        };
4945        let f = if w8 {
4946            self.func("router_gemv_f32_w8")
4947        } else {
4948            self.func("router_gemv_f32")
4949        };
4950        let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
4951        let cfg = LaunchConfig {
4952            grid_dim: (n_experts as u32, t as u32, 1),
4953            block_dim: (32, if w8 { 8 } else { 1 }, 1),
4954            shared_mem_bytes: 0,
4955        };
4956        let __s_b = self.gpu.stream();
4957        let mut b = __s_b.launch_builder(&f);
4958        b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
4959        unsafe {
4960            b.launch(cfg)?;
4961        }
4962        Ok(())
4963    }
4964
4965    /// f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).
4966    pub fn rows_permute(
4967        &self,
4968        src: &CudaSlice<f32>,
4969        idx: &CudaSlice<i32>,
4970        nrows: usize,
4971        ncols: usize,
4972    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4973        let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
4974        let f = self.func("rows_permute_f32");
4975        let (nc, nr) = (ncols as i32, nrows as i32);
4976        let cfg = LaunchConfig {
4977            grid_dim: (nrows as u32, 1, 1),
4978            block_dim: (256, 1, 1),
4979            shared_mem_bytes: 0,
4980        };
4981        let __s_b = self.gpu.stream();
4982        let mut b = __s_b.launch_builder(&f);
4983        b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
4984        unsafe {
4985            b.launch(cfg)?;
4986        }
4987        Ok(dst)
4988    }
4989
4990    /// shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer
4991    /// cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the
4992    /// splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1
4993    /// decode chain and the small-t spec-verify chain match per row by construction.
4994    pub fn sigmoid_dot_rows(
4995        &self,
4996        x: &CudaSlice<f32>,
4997        w: &CudaSlice<f32>,
4998        n_embd: usize,
4999        t: usize,
5000    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5001        // MEMRA_SHEXP_DOT=0: rollback seam to the cuBLASLt linear + sigmoid pair (numeric
5002        // config; same class as MEMRA_ROUTER_V2).
5003        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5004        if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
5005            let gs = self.linear(x, w, t, n_embd, 1)?;
5006            let mut g = self.uninit(t)?;
5007            self.sigmoid(&gs, &mut g, t)?;
5008            return Ok(g);
5009        }
5010        // FAST-ROUTER lane note (2026-08-02): a register-tiled 8-token batch twin of this
5011        // kernel was built, proven bit-identical, and measured SLOWER at every prefill t on
5012        // the 5090 (0.62-0.89x — launch-latency-bound op, ~7us/layer at m=2048;
5013        // research/fast-router-20260802/crossover-router.jsonl). Dispatch arm killed per
5014        // flags doctrine; this per-token form serves every t.
5015        let mut g = self.alloc_uninit::<f32>(t)?;
5016        let f = self.func("sigmoid_dot_rows_f32");
5017        let (ne, ti) = (n_embd as i32, t as i32);
5018        let cfg = LaunchConfig {
5019            grid_dim: (t as u32, 1, 1),
5020            block_dim: (32, 8, 1),
5021            shared_mem_bytes: 0,
5022        };
5023        let __s_b = self.gpu.stream();
5024        let mut b = __s_b.launch_builder(&f);
5025        b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
5026        unsafe {
5027            b.launch(cfg)?;
5028        }
5029        Ok(g)
5030    }
5031
5032    /// `sigmoid_dot_rows` writing into a caller-owned [t] buffer (token-graph alloc-free).
5033    pub fn sigmoid_dot_rows_into(
5034        &self,
5035        x: &CudaSlice<f32>,
5036        w: &CudaSlice<f32>,
5037        g: &mut CudaSlice<f32>,
5038        n_embd: usize,
5039        t: usize,
5040    ) -> Result<(), Box<dyn std::error::Error>> {
5041        if g.len() < t {
5042            return Err("sigmoid_dot_rows_into output too small".into());
5043        }
5044        let f = self.func("sigmoid_dot_rows_f32");
5045        let (ne, ti) = (n_embd as i32, t as i32);
5046        let cfg = LaunchConfig {
5047            grid_dim: (t as u32, 1, 1),
5048            block_dim: (32, 8, 1),
5049            shared_mem_bytes: 0,
5050        };
5051        let __s_b = self.gpu.stream();
5052        let mut b = __s_b.launch_builder(&f);
5053        b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
5054        unsafe {
5055            b.launch(cfg)?;
5056        }
5057        Ok(())
5058    }
5059
5060    /// ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.
5061    pub fn spec_rollback_stream(
5062        &self,
5063        len_ptrs: &CudaSlice<u64>,
5064        pos_start: &CudaSlice<i32>,
5065        acc: &CudaSlice<u32>,
5066        base: usize,
5067        n_rows: usize,
5068    ) -> Result<(), Box<dyn std::error::Error>> {
5069        let f = self.func("spec_rollback_stream");
5070        let (b, nr) = (base as i32, n_rows as i32);
5071        let cfg = LaunchConfig {
5072            grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
5073            block_dim: (64, 1, 1),
5074            shared_mem_bytes: 0,
5075        };
5076        let __s_bl = self.gpu.stream();
5077        let mut bl = __s_bl.launch_builder(&f);
5078        bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
5079        unsafe {
5080            bl.launch(cfg)?;
5081        }
5082        Ok(())
5083    }
5084
5085    /// PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].
5086    pub fn plain_tok_ring(
5087        &self,
5088        vam: &CudaSlice<u32>,
5089        pos_start: &CudaSlice<i32>,
5090        base: usize,
5091        ring: &mut CudaSlice<u32>,
5092    ) -> Result<(), Box<dyn std::error::Error>> {
5093        let f = self.func("plain_tok_ring");
5094        let (b, cap) = (base as i32, ring.len() as i32);
5095        let cfg = LaunchConfig {
5096            grid_dim: (1, 1, 1),
5097            block_dim: (32, 1, 1),
5098            shared_mem_bytes: 0,
5099        };
5100        let __s_bl = self.gpu.stream();
5101        let mut bl = __s_bl.launch_builder(&f);
5102        bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
5103        unsafe {
5104            bl.launch(cfg)?;
5105        }
5106        Ok(())
5107    }
5108
5109    /// ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.
5110    pub fn spec_ring_commit(
5111        &self,
5112        vtok: &CudaSlice<u32>,
5113        acc: &CudaSlice<u32>,
5114        brk: &CudaSlice<u32>,
5115        ring: &mut CudaSlice<u32>,
5116        pend: &mut CudaSlice<u32>,
5117    ) -> Result<(), Box<dyn std::error::Error>> {
5118        let f = self.func("spec_ring_commit");
5119        let cfg = LaunchConfig {
5120            grid_dim: (1, 1, 1),
5121            block_dim: (32, 1, 1),
5122            shared_mem_bytes: 0,
5123        };
5124        let __s_b = self.gpu.stream();
5125        let mut b = __s_b.launch_builder(&f);
5126        b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
5127        unsafe {
5128            b.launch(cfg)?;
5129        }
5130        Ok(())
5131    }
5132    pub fn i32_copy_add(
5133        &self,
5134        src: &CudaSlice<i32>,
5135        dst: &mut CudaSlice<i32>,
5136        delta: i32,
5137    ) -> Result<(), Box<dyn std::error::Error>> {
5138        let f = self.func("i32_copy_add");
5139        let cfg = LaunchConfig {
5140            grid_dim: (1, 1, 1),
5141            block_dim: (32, 1, 1),
5142            shared_mem_bytes: 0,
5143        };
5144        let __s_b = self.gpu.stream();
5145        let mut b = __s_b.launch_builder(&f);
5146        b.arg(src).arg(dst).arg(&delta);
5147        unsafe {
5148            b.launch(cfg)?;
5149        }
5150        Ok(())
5151    }
5152    pub fn u32_copy(
5153        &self,
5154        src: &CudaSlice<u32>,
5155        dst: &mut CudaSlice<u32>,
5156    ) -> Result<(), Box<dyn std::error::Error>> {
5157        let f = self.func("u32_copy");
5158        let cfg = LaunchConfig {
5159            grid_dim: (1, 1, 1),
5160            block_dim: (32, 1, 1),
5161            shared_mem_bytes: 0,
5162        };
5163        let __s_b = self.gpu.stream();
5164        let mut b = __s_b.launch_builder(&f);
5165        b.arg(src).arg(dst);
5166        unsafe {
5167            b.launch(cfg)?;
5168        }
5169        Ok(())
5170    }
5171
5172    /// ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host
5173    /// adaptive policy as a captured device op (policy-identical: the accept walk depth
5174    /// caps acceptance exactly like drafting fewer tokens).
5175    pub fn spec_adapt_k(
5176        &self,
5177        acc: &CudaSlice<u32>,
5178        brk: &mut CudaSlice<u32>,
5179        floor: usize,
5180        cap: usize,
5181    ) -> Result<(), Box<dyn std::error::Error>> {
5182        let f = self.func("spec_adapt_k");
5183        let (fl, cp) = (floor as i32, cap as i32);
5184        let cfg = LaunchConfig {
5185            grid_dim: (1, 1, 1),
5186            block_dim: (32, 1, 1),
5187            shared_mem_bytes: 0,
5188        };
5189        let __s_b = self.gpu.stream();
5190        let mut b = __s_b.launch_builder(&f);
5191        b.arg(acc).arg(brk).arg(&fl).arg(&cp);
5192        unsafe {
5193            b.launch(cfg)?;
5194        }
5195        Ok(())
5196    }
5197
5198    /// ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).
5199    pub fn spec_accept_greedy_dc(
5200        &self,
5201        preds: &CudaSlice<u32>,
5202        vtok: &CudaSlice<u32>,
5203        last_pred: &CudaSlice<u32>,
5204        brk: &CudaSlice<u32>,
5205        out: &mut CudaSlice<u32>,
5206    ) -> Result<(), Box<dyn std::error::Error>> {
5207        let f = self.func("spec_accept_greedy_dc");
5208        let cfg = LaunchConfig {
5209            grid_dim: (1, 1, 1),
5210            block_dim: (32, 1, 1),
5211            shared_mem_bytes: 0,
5212        };
5213        let __s_b = self.gpu.stream();
5214        let mut b = __s_b.launch_builder(&f);
5215        b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
5216        unsafe {
5217            b.launch(cfg)?;
5218        }
5219        Ok(())
5220    }
5221
5222    /// ROUND-STREAM stage (c) 2: verify-chain device-pos entries.
5223    pub fn pos_iota(
5224        &self,
5225        pos0: &CudaSlice<i32>,
5226        out: &mut CudaSlice<i32>,
5227        t: usize,
5228    ) -> Result<(), Box<dyn std::error::Error>> {
5229        let f = self.func("pos_iota_i32");
5230        let ti = t as i32;
5231        let cfg = LaunchConfig {
5232            grid_dim: (1, 1, 1),
5233            block_dim: (t.max(1) as u32, 1, 1),
5234            shared_mem_bytes: 0,
5235        };
5236        let __s_b = self.gpu.stream();
5237        let mut b = __s_b.launch_builder(&f);
5238        b.arg(pos0).arg(out).arg(&ti);
5239        unsafe {
5240            b.launch(cfg)?;
5241        }
5242        Ok(())
5243    }
5244    #[allow(clippy::too_many_arguments)]
5245    pub fn append_kv_quantized_rows_dc(
5246        &self,
5247        k_rows: &CudaSlice<f32>,
5248        v_rows: &CudaSlice<f32>,
5249        kc: &mut CudaSlice<u8>,
5250        vc: &mut CudaSlice<u8>,
5251        t0_dev: &CudaSlice<i32>,
5252        t: usize,
5253        kv_dim_k: usize,
5254        kv_dim_v: usize,
5255        k_tok_bytes: usize,
5256        v_tok_bytes: usize,
5257        g: bool,
5258    ) -> Result<(), Box<dyn std::error::Error>> {
5259        let f = if g {
5260            self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
5261        } else {
5262            self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
5263        };
5264        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
5265        let cfg = LaunchConfig {
5266            grid_dim: (nblk, t as u32, 1),
5267            block_dim: (32, 1, 1),
5268            shared_mem_bytes: 0,
5269        };
5270        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
5271        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5272        let __s_b = self.gpu.stream();
5273        let mut b = __s_b.launch_builder(&f);
5274        b.arg(k_rows)
5275            .arg(v_rows)
5276            .arg(kc)
5277            .arg(vc)
5278            .arg(t0_dev)
5279            .arg(&kdk)
5280            .arg(&kdv)
5281            .arg(&ktb)
5282            .arg(&vtb);
5283        unsafe {
5284            b.launch(cfg)?;
5285        }
5286        Ok(())
5287    }
5288
5289    /// t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces
5290    /// append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).
5291    #[allow(clippy::too_many_arguments)]
5292    pub fn append_kv_quantized_row_dc_inc(
5293        &self,
5294        k_row: &CudaSlice<f32>,
5295        v_row: &CudaSlice<f32>,
5296        kc: &mut CudaSlice<u8>,
5297        vc: &mut CudaSlice<u8>,
5298        t0_dev: &mut CudaSlice<i32>,
5299        kv_dim_k: usize,
5300        kv_dim_v: usize,
5301        k_tok_bytes: usize,
5302        v_tok_bytes: usize,
5303        g: bool,
5304    ) -> Result<(), Box<dyn std::error::Error>> {
5305        let f = if g {
5306            self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
5307        } else {
5308            self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
5309        };
5310        let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
5311        let cfg = LaunchConfig {
5312            grid_dim: (1, 1, 1),
5313            block_dim: (nthreads, 1, 1),
5314            shared_mem_bytes: 0,
5315        };
5316        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
5317        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5318        let __s_b = self.gpu.stream();
5319        let mut b = __s_b.launch_builder(&f);
5320        b.arg(k_row)
5321            .arg(v_row)
5322            .arg(kc)
5323            .arg(vc)
5324            .arg(t0_dev)
5325            .arg(&kdk)
5326            .arg(&kdv)
5327            .arg(&ktb)
5328            .arg(&vtb);
5329        unsafe {
5330            b.launch(cfg)?;
5331        }
5332        Ok(())
5333    }
5334
5335    /// ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).
5336    pub fn pack_tok_p(
5337        &self,
5338        tok: &CudaSlice<u32>,
5339        p: &CudaSlice<f32>,
5340        out: &mut CudaSlice<u32>,
5341        slot: usize,
5342    ) -> Result<(), Box<dyn std::error::Error>> {
5343        let f = self.func("pack_tok_p");
5344        let sl = slot as i32;
5345        let cfg = LaunchConfig {
5346            grid_dim: (1, 1, 1),
5347            block_dim: (32, 1, 1),
5348            shared_mem_bytes: 0,
5349        };
5350        let __s_b = self.gpu.stream();
5351        let mut b = __s_b.launch_builder(&f);
5352        b.arg(tok).arg(p).arg(out).arg(&sl);
5353        unsafe {
5354            b.launch(cfg)?;
5355        }
5356        Ok(())
5357    }
5358    pub fn tok_map_u32(
5359        &self,
5360        tok: &mut CudaSlice<u32>,
5361        map: &CudaSlice<u32>,
5362    ) -> Result<(), Box<dyn std::error::Error>> {
5363        let f = self.func("tok_map_u32");
5364        let cfg = LaunchConfig {
5365            grid_dim: (1, 1, 1),
5366            block_dim: (32, 1, 1),
5367            shared_mem_bytes: 0,
5368        };
5369        let __s_b = self.gpu.stream();
5370        let mut b = __s_b.launch_builder(&f);
5371        b.arg(tok).arg(map);
5372        unsafe {
5373            b.launch(cfg)?;
5374        }
5375        Ok(())
5376    }
5377
5378    /// ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.
5379    #[allow(clippy::too_many_arguments)]
5380    pub fn spec_assemble_verify(
5381        &self,
5382        tokp: &CudaSlice<u32>,
5383        pend: &CudaSlice<u32>,
5384        d2t: Option<&CudaSlice<u32>>,
5385        vtok: &mut CudaSlice<u32>,
5386        brk: &mut CudaSlice<u32>,
5387        p_min: f32,
5388        k: usize,
5389        pmin0: bool,
5390    ) -> Result<(), Box<dyn std::error::Error>> {
5391        let f = self.func("spec_assemble_verify");
5392        let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
5393        let cfg = LaunchConfig {
5394            grid_dim: (1, 1, 1),
5395            block_dim: (32, 1, 1),
5396            shared_mem_bytes: 0,
5397        };
5398        let __s_b = self.gpu.stream();
5399        let mut b = __s_b.launch_builder(&f);
5400        match d2t {
5401            Some(m) => {
5402                b.arg(tokp)
5403                    .arg(pend)
5404                    .arg(m)
5405                    .arg(vtok)
5406                    .arg(brk)
5407                    .arg(&p_min)
5408                    .arg(&ki)
5409                    .arg(&pm);
5410                unsafe {
5411                    b.launch(cfg)?;
5412                }
5413            }
5414            None => {
5415                let null: u64 = 0;
5416                b.arg(tokp)
5417                    .arg(pend)
5418                    .arg(&null)
5419                    .arg(vtok)
5420                    .arg(brk)
5421                    .arg(&p_min)
5422                    .arg(&ki)
5423                    .arg(&pm);
5424                unsafe {
5425                    b.launch(cfg)?;
5426                }
5427            }
5428        }
5429        Ok(())
5430    }
5431
5432    /// ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).
5433    #[allow(clippy::too_many_arguments)]
5434    pub fn ssm_conv_ring_rebuild_dc(
5435        &self,
5436        qkv_tm: &CudaSlice<f32>,
5437        ring_old: &CudaSlice<f32>,
5438        conv_state: &mut CudaSlice<f32>,
5439        conv_dim: usize,
5440        acc: &CudaSlice<u32>,
5441        base: usize,
5442        t_v: usize,
5443        d_conv: usize,
5444    ) -> Result<(), Box<dyn std::error::Error>> {
5445        let f = self.func("ssm_conv_ring_rebuild_f32_dc");
5446        let n = conv_dim * (d_conv - 1);
5447        let cfg = LaunchConfig::for_num_elems(n as u32);
5448        let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
5449        let __s_b = self.gpu.stream();
5450        let mut b = __s_b.launch_builder(&f);
5451        b.arg(qkv_tm)
5452            .arg(ring_old)
5453            .arg(conv_state)
5454            .arg(&cd)
5455            .arg(acc)
5456            .arg(&b0)
5457            .arg(&tv)
5458            .arg(&dc);
5459        unsafe {
5460            b.launch(cfg)?;
5461        }
5462        Ok(())
5463    }
5464    #[allow(clippy::too_many_arguments)]
5465    pub fn gdn_scan_s128_dc(
5466        &self,
5467        q: &CudaSlice<f32>,
5468        k: &CudaSlice<f32>,
5469        v: &CudaSlice<f32>,
5470        g: &CudaSlice<f32>,
5471        beta: &CudaSlice<f32>,
5472        state_in: &CudaSlice<f32>,
5473        state_out: &mut CudaSlice<f32>,
5474        o: &mut CudaSlice<f32>,
5475        n_head: usize,
5476        acc: &CudaSlice<u32>,
5477        base: usize,
5478        t_v: usize,
5479        scale: f32,
5480    ) -> Result<(), Box<dyn std::error::Error>> {
5481        let f = self.func("gdn_scan_s128_dc");
5482        const S_V: u32 = 128;
5483        const WARP: u32 = 32;
5484        const COLS_PER_BLOCK: u32 = 4;
5485        let cfg = LaunchConfig {
5486            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
5487            block_dim: (WARP, COLS_PER_BLOCK, 1),
5488            shared_mem_bytes: 0,
5489        };
5490        let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
5491        let __s_b = self.gpu.stream();
5492        let mut b = __s_b.launch_builder(&f);
5493        b.arg(q)
5494            .arg(k)
5495            .arg(v)
5496            .arg(g)
5497            .arg(beta)
5498            .arg(state_in)
5499            .arg(state_out)
5500            .arg(o)
5501            .arg(&h)
5502            .arg(acc)
5503            .arg(&b0)
5504            .arg(&tv)
5505            .arg(&scale);
5506        unsafe {
5507            b.launch(cfg)?;
5508        }
5509        Ok(())
5510    }
5511
5512    /// ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).
5513    pub fn spec_rollback_kv(
5514        &self,
5515        len_ptrs: &CudaSlice<u64>,
5516        saved: &CudaSlice<i32>,
5517        acc: &CudaSlice<u32>,
5518        base: usize,
5519        n_layer: usize,
5520    ) -> Result<(), Box<dyn std::error::Error>> {
5521        let f = self.func("spec_rollback_kv");
5522        let (b, nl) = (base as i32, n_layer as i32);
5523        let cfg = LaunchConfig {
5524            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
5525            block_dim: (64, 1, 1),
5526            shared_mem_bytes: 0,
5527        };
5528        let __s_bl = self.gpu.stream();
5529        let mut bl = __s_bl.launch_builder(&f);
5530        bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
5531        unsafe {
5532            bl.launch(cfg)?;
5533        }
5534        Ok(())
5535    }
5536
5537    /// OPTIPIPE increment 1: derive the K=1 successor-valid bit on device.
5538    pub fn spec_fork_valid(
5539        &self,
5540        acc: &CudaSlice<u32>,
5541        optimistic_pending: u32,
5542        valid: &mut CudaSlice<u32>,
5543    ) -> Result<(), Box<dyn std::error::Error>> {
5544        let f = self.func("spec_fork_valid");
5545        let cfg = LaunchConfig {
5546            grid_dim: (1, 1, 1),
5547            block_dim: (1, 1, 1),
5548            shared_mem_bytes: 0,
5549        };
5550        let __s_bl = self.gpu.stream();
5551        let mut bl = __s_bl.launch_builder(&f);
5552        bl.arg(acc).arg(&optimistic_pending).arg(valid);
5553        unsafe {
5554            bl.launch(cfg)?;
5555        }
5556        Ok(())
5557    }
5558
5559    /// OPTIPIPE increment 1: leave stage-local KV lengths on hit, restore them on miss.
5560    pub fn spec_fork_reconcile_kv(
5561        &self,
5562        len_ptrs: &CudaSlice<u64>,
5563        saved: &CudaSlice<i32>,
5564        acc: &CudaSlice<u32>,
5565        valid: &CudaSlice<u32>,
5566        base: usize,
5567        n_layer: usize,
5568    ) -> Result<(), Box<dyn std::error::Error>> {
5569        let f = self.func("spec_fork_reconcile_kv");
5570        let (b, nl) = (base as i32, n_layer as i32);
5571        let cfg = LaunchConfig {
5572            grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
5573            block_dim: (64, 1, 1),
5574            shared_mem_bytes: 0,
5575        };
5576        let __s_bl = self.gpu.stream();
5577        let mut bl = __s_bl.launch_builder(&f);
5578        bl.arg(len_ptrs)
5579            .arg(saved)
5580            .arg(acc)
5581            .arg(valid)
5582            .arg(&b)
5583            .arg(&nl);
5584        unsafe {
5585            bl.launch(cfg)?;
5586        }
5587        Ok(())
5588    }
5589
5590    /// OPTIPIPE increment 1: conditionally restore one stage-owned recurrent-state buffer.
5591    pub fn spec_fork_restore_f32(
5592        &self,
5593        snapshot: &CudaSlice<f32>,
5594        state: &mut CudaSlice<f32>,
5595        valid: &CudaSlice<u32>,
5596    ) -> Result<(), Box<dyn std::error::Error>> {
5597        assert_eq!(
5598            snapshot.len(),
5599            state.len(),
5600            "fork recurrent snapshot shape mismatch"
5601        );
5602        let f = self.func("spec_fork_restore_f32");
5603        let n = state.len() as i32;
5604        #[allow(clippy::manual_clamp)]
5605        // allow: the min/max chain mirrors the reference arithmetic order in pinned sizing/quant math
5606        let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
5607        let cfg = LaunchConfig {
5608            grid_dim: (blocks, 1, 1),
5609            block_dim: (256, 1, 1),
5610            shared_mem_bytes: 0,
5611        };
5612        let __s_bl = self.gpu.stream();
5613        let mut bl = __s_bl.launch_builder(&f);
5614        bl.arg(snapshot).arg(state).arg(valid).arg(&n);
5615        unsafe {
5616            bl.launch(cfg)?;
5617        }
5618        Ok(())
5619    }
5620
5621    /// ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header).
5622    /// Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).
5623    pub fn spec_seed_gather(
5624        &self,
5625        vx: &CudaSlice<f32>,
5626        fill_prev: &CudaSlice<f32>,
5627        acc: &CudaSlice<u32>,
5628        h_seed: &mut CudaSlice<f32>,
5629        base: usize,
5630        n_embd: usize,
5631    ) -> Result<(), Box<dyn std::error::Error>> {
5632        let f = self.func("spec_seed_gather");
5633        let (b, ne) = (base as i32, n_embd as i32);
5634        let cfg = LaunchConfig {
5635            grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
5636            block_dim: (256, 1, 1),
5637            shared_mem_bytes: 0,
5638        };
5639        let __s_bl = self.gpu.stream();
5640        let mut bl = __s_bl.launch_builder(&f);
5641        bl.arg(vx)
5642            .arg(fill_prev)
5643            .arg(acc)
5644            .arg(h_seed)
5645            .arg(&b)
5646            .arg(&ne);
5647        unsafe {
5648            bl.launch(cfg)?;
5649        }
5650        Ok(())
5651    }
5652
5653    /// ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).
5654    pub fn spec_accept_greedy(
5655        &self,
5656        preds: &CudaSlice<u32>,
5657        draft: &CudaSlice<u32>,
5658        last_pred: u32,
5659        base: usize,
5660        k_round: usize,
5661        out: &mut CudaSlice<u32>,
5662    ) -> Result<(), Box<dyn std::error::Error>> {
5663        let f = self.func("spec_accept_greedy");
5664        let (b, k) = (base as i32, k_round as i32);
5665        let cfg = LaunchConfig {
5666            grid_dim: (1, 1, 1),
5667            block_dim: (32, 1, 1),
5668            shared_mem_bytes: 0,
5669        };
5670        let __s_bl = self.gpu.stream();
5671        let mut bl = __s_bl.launch_builder(&f);
5672        bl.arg(preds)
5673            .arg(draft)
5674            .arg(&last_pred)
5675            .arg(&b)
5676            .arg(&k)
5677            .arg(out);
5678        unsafe {
5679            bl.launch(cfg)?;
5680        }
5681        Ok(())
5682    }
5683
5684    // ================= SAMPLED-SPEC PRIMITIVES (spec_sample.cu, piece A) =================
5685    // Counter-based randomness: every call takes (seed, stream_pos) — the caller owns the
5686    // event counter (one per sampled token). temp <= 0 arms are exact greedy limits.
5687
5688    /// y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y
5689    /// = one categorical sample at temperature `temp`). temp<=0: y = x (pure copy).
5690    pub fn gumbel_perturb(
5691        &self,
5692        x: &CudaSlice<f32>,
5693        y: &mut CudaSlice<f32>,
5694        n: usize,
5695        seed: u64,
5696        stream_pos: u32,
5697        temp: f32,
5698    ) -> Result<(), Box<dyn std::error::Error>> {
5699        let f = self.func("gumbel_perturb_f32");
5700        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5701        let cfg = LaunchConfig {
5702            grid_dim: (n.div_ceil(256) as u32, 1, 1),
5703            block_dim: (256, 1, 1),
5704            shared_mem_bytes: 0,
5705        };
5706        let __s_b = self.gpu.stream();
5707        let mut b = __s_b.launch_builder(&f);
5708        b.arg(x)
5709            .arg(&mut *y)
5710            .arg(&ni)
5711            .arg(&slo)
5712            .arg(&shi)
5713            .arg(&stream_pos)
5714            .arg(&temp);
5715        unsafe {
5716            b.launch(cfg)?;
5717        }
5718        Ok(())
5719    }
5720
5721    /// GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id
5722    /// whose bit is unset in the packed llguidance bitset, IN PLACE on row `col` of a stacked
5723    /// [B, n_vocab] logits buffer. `mask` = the SimpleVob u32 words H2D'd verbatim
5724    /// (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail)
5725    /// are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX ==
5726    /// the argmax/gumbel kernels' init sentinel, so a fully-banned tail can never win and
5727    /// ordering matches the host -inf mask bit-for-bit for every finite logit.
5728    pub fn mask_logits_col(
5729        &self,
5730        logits: &mut CudaSlice<f32>,
5731        mask: &CudaSlice<u32>,
5732        col: usize,
5733        n: usize,
5734        mask_words: usize,
5735    ) -> Result<(), Box<dyn std::error::Error>> {
5736        let f = self.func("mask_logits_f32");
5737        let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
5738        let cfg = LaunchConfig {
5739            grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
5740            block_dim: (256, 1, 1),
5741            shared_mem_bytes: 0,
5742        };
5743        let __s_b = self.gpu.stream();
5744        let mut b = __s_b.launch_builder(&f);
5745        b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
5746        unsafe {
5747            b.launch(cfg)?;
5748        }
5749        Ok(())
5750    }
5751
5752    /// Column-`col` twin of `gumbel_perturb` over stacked logits [B, n_vocab] (the batched
5753    /// serving tick's device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane).
5754    /// SAME kernel/Philox mapping as `gumbel_perturb` — bit-identical perturbation for the
5755    /// same (seed, stream_pos, temp) regardless of which batch column the row sits in
5756    /// (the lane index is the in-row position; `col` only moves the input pointer). That
5757    /// pointer-invariance IS the serving isolation contract for sampled rows.
5758    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5759    pub fn gumbel_perturb_col(
5760        &self,
5761        x: &CudaSlice<f32>,
5762        col: usize,
5763        y: &mut CudaSlice<f32>,
5764        n: usize,
5765        seed: u64,
5766        stream_pos: u32,
5767        temp: f32,
5768    ) -> Result<(), Box<dyn std::error::Error>> {
5769        let f = self.func("gumbel_perturb_f32");
5770        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5771        let col_view = x.slice(col * n..(col + 1) * n);
5772        let cfg = LaunchConfig {
5773            grid_dim: (n.div_ceil(256) as u32, 1, 1),
5774            block_dim: (256, 1, 1),
5775            shared_mem_bytes: 0,
5776        };
5777        let __s_b = self.gpu.stream();
5778        let mut b = __s_b.launch_builder(&f);
5779        b.arg(&col_view)
5780            .arg(&mut *y)
5781            .arg(&ni)
5782            .arg(&slo)
5783            .arg(&shi)
5784            .arg(&stream_pos)
5785            .arg(&temp);
5786        unsafe {
5787            b.launch(cfg)?;
5788        }
5789        Ok(())
5790    }
5791
5792    /// Filtered twin of `gumbel_perturb_col`: the per-row (row_max, th) floor comes from
5793    /// DEVICE buffers (`filter_stats` output slots at `stat_idx`) — one filtered draw from
5794    /// the top-k/top-p/min-p-truncated softmax with no stat D2H and no row copy. Same
5795    /// Philox mapping as every gumbel kernel (pointer-invariant across batch columns —
5796    /// the serving isolation contract for sampled rows).
5797    #[allow(clippy::too_many_arguments)]
5798    pub fn gumbel_perturb_filtered_col(
5799        &self,
5800        x: &CudaSlice<f32>,
5801        col: usize,
5802        y: &mut CudaSlice<f32>,
5803        n: usize,
5804        seed: u64,
5805        stream_pos: u32,
5806        temp: f32,
5807        stat_max: &CudaSlice<f32>,
5808        stat_th: &CudaSlice<f32>,
5809        stat_idx: usize,
5810    ) -> Result<(), Box<dyn std::error::Error>> {
5811        let f = self.func("gumbel_perturb_filtered_col_f32");
5812        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5813        let (ci, si) = (col as i32, stat_idx as i32);
5814        let cfg = LaunchConfig {
5815            grid_dim: (n.div_ceil(256) as u32, 1, 1),
5816            block_dim: (256, 1, 1),
5817            shared_mem_bytes: 0,
5818        };
5819        let __s_b = self.gpu.stream();
5820        let mut b = __s_b.launch_builder(&f);
5821        b.arg(x)
5822            .arg(&ci)
5823            .arg(&mut *y)
5824            .arg(&ni)
5825            .arg(&slo)
5826            .arg(&shi)
5827            .arg(&stream_pos)
5828            .arg(&temp)
5829            .arg(stat_max)
5830            .arg(stat_th)
5831            .arg(&si);
5832        unsafe {
5833            b.launch(cfg)?;
5834        }
5835        Ok(())
5836    }
5837
5838    /// In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled
5839    /// graph-draft chain replays with FIXED kernel args, so the Philox event counter must be
5840    /// DEVICE data — the host seeds it once per round; every replay bumps it before the perturb
5841    /// reads it (counter is data, not state — graph-replay-safe).
5842    pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
5843        let f = self.func("memra_sctr_inc");
5844        let cfg = LaunchConfig {
5845            grid_dim: (1, 1, 1),
5846            block_dim: (1, 1, 1),
5847            shared_mem_bytes: 0,
5848        };
5849        let __s_b = self.gpu.stream();
5850        let mut b = __s_b.launch_builder(&f);
5851        b.arg(&mut *ctr);
5852        unsafe {
5853            b.launch(cfg)?;
5854        }
5855        Ok(())
5856    }
5857
5858    /// Graph-capturable `gumbel_perturb`: the sampling-event counter comes from DEVICE memory
5859    /// (`ctr[0]`) instead of a host scalar. Identical math to `gumbel_perturb` at
5860    /// stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled
5861    /// chains produce bit-identical perturbations for the same (seed, counter, temp).
5862    pub fn gumbel_perturb_ctr(
5863        &self,
5864        x: &CudaSlice<f32>,
5865        y: &mut CudaSlice<f32>,
5866        n: usize,
5867        seed: u64,
5868        ctr: &CudaSlice<u32>,
5869        temp: f32,
5870    ) -> Result<(), Box<dyn std::error::Error>> {
5871        let f = self.func("gumbel_perturb_ctr_f32");
5872        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5873        let cfg = LaunchConfig {
5874            grid_dim: (n.div_ceil(256) as u32, 1, 1),
5875            block_dim: (256, 1, 1),
5876            shared_mem_bytes: 0,
5877        };
5878        let __s_b = self.gpu.stream();
5879        let mut b = __s_b.launch_builder(&f);
5880        b.arg(x)
5881            .arg(&mut *y)
5882            .arg(&ni)
5883            .arg(&slo)
5884            .arg(&shi)
5885            .arg(ctr)
5886            .arg(&temp);
5887        unsafe {
5888            b.launch(cfg)?;
5889        }
5890        Ok(())
5891    }
5892
5893    /// Graph-capturable `gumbel_perturb_filtered` (lane/step37-draft-graph-serving): the
5894    /// sampling-event counter comes from DEVICE memory (`ctr[0]`) and the filter stats
5895    /// (row_max, th) from DEVICE slots — the `filter_stats` outputs of the same captured
5896    /// body. Identical math (same Philox call, same lane mapping, same e0 filter test) to
5897    /// `gumbel_perturb_filtered` at stream_pos == ctr[0], row_max == mx[0], th == th_d[0]:
5898    /// the eager and graph FILTERED sampled chains produce bit-identical perturbations for
5899    /// the same (seed, counter, stats). Launch geometry mirrors the host-scalar wrapper.
5900    #[allow(clippy::too_many_arguments)]
5901    pub fn gumbel_perturb_filtered_ctr(
5902        &self,
5903        x: &CudaSlice<f32>,
5904        y: &mut CudaSlice<f32>,
5905        n: usize,
5906        seed: u64,
5907        ctr: &CudaSlice<u32>,
5908        temp: f32,
5909        stat_max: &CudaSlice<f32>,
5910        stat_th: &CudaSlice<f32>,
5911    ) -> Result<(), Box<dyn std::error::Error>> {
5912        let f = self.func("gumbel_perturb_filtered_ctr_f32");
5913        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5914        let cfg = LaunchConfig {
5915            grid_dim: (n.div_ceil(256) as u32, 1, 1),
5916            block_dim: (256, 1, 1),
5917            shared_mem_bytes: 0,
5918        };
5919        let __s_b = self.gpu.stream();
5920        let mut b = __s_b.launch_builder(&f);
5921        b.arg(x)
5922            .arg(&mut *y)
5923            .arg(&ni)
5924            .arg(&slo)
5925            .arg(&shi)
5926            .arg(ctr)
5927            .arg(&temp)
5928            .arg(stat_max)
5929            .arg(stat_th);
5930        unsafe {
5931            b.launch(cfg)?;
5932        }
5933        Ok(())
5934    }
5935
5936    /// out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index
5937    /// into x with `row_stride` f32s per row. temp<=0: out = 1.0 iff id is the row argmax
5938    /// (smallest-index tie-break — matches the argmax-gate contract).
5939    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5940    pub fn softmax_gather(
5941        &self,
5942        x: &CudaSlice<f32>,
5943        row_stride: usize,
5944        ids: &CudaSlice<u32>,
5945        rows: &CudaSlice<i32>,
5946        out: &mut CudaSlice<f32>,
5947        n: usize,
5948        npair: usize,
5949        temp: f32,
5950    ) -> Result<(), Box<dyn std::error::Error>> {
5951        let f = self.func("softmax_gather_f32");
5952        let (ni, rs) = (n as i32, row_stride as i64);
5953        let np = npair as i32;
5954        let cfg = LaunchConfig {
5955            grid_dim: (npair as u32, 1, 1),
5956            block_dim: (256, 1, 1),
5957            shared_mem_bytes: 0,
5958        };
5959        let __s_b = self.gpu.stream();
5960        let mut b = __s_b.launch_builder(&f);
5961        b.arg(x)
5962            .arg(&rs)
5963            .arg(ids)
5964            .arg(rows)
5965            .arg(&mut *out)
5966            .arg(&ni)
5967            .arg(&np)
5968            .arg(&temp);
5969        unsafe {
5970            b.launch(cfg)?;
5971        }
5972        Ok(())
5973    }
5974
5975    /// Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain
5976    /// categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed
5977    /// (softmax_gather's pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.
5978    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
5979    pub fn residual_sample(
5980        &self,
5981        p: &CudaSlice<f32>,
5982        q: Option<&CudaSlice<f32>>,
5983        n: usize,
5984        temp: f32,
5985        seed: u64,
5986        stream_pos: u32,
5987        out_tok: &mut CudaSlice<u32>,
5988    ) -> Result<(), Box<dyn std::error::Error>> {
5989        let f = self.func("residual_sample_f32");
5990        let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5991        let nth = 1024u32;
5992        let cfg = LaunchConfig {
5993            grid_dim: (1, 1, 1),
5994            block_dim: (nth, 1, 1),
5995            shared_mem_bytes: 0,
5996        };
5997        let has_q: i32 = q.is_some() as i32;
5998        let qbuf = q.unwrap_or(p); // dummy when absent; kernel gates on has_q
5999        let __s_b = self.gpu.stream();
6000        let mut b = __s_b.launch_builder(&f);
6001        b.arg(p)
6002            .arg(qbuf)
6003            .arg(&has_q)
6004            .arg(&ni)
6005            .arg(&temp)
6006            .arg(&slo)
6007            .arg(&shi)
6008            .arg(&stream_pos)
6009            .arg(&mut *out_tok);
6010        unsafe {
6011            b.launch(cfg)?;
6012        }
6013        Ok(())
6014    }
6015
6016    /// Access the shared MoE residency cache (EDGE-1 §B), building it on first use under
6017    /// MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not
6018    /// the GEMM. `max_block_bytes` sizes the slots (largest of gate/up/down). Returns the closure's
6019    /// result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).
6020    pub fn with_moe_cache<R>(
6021        &self,
6022        max_block_bytes: usize,
6023        f: impl FnOnce(
6024            &mut crate::moe_cache::MoeSlotCache,
6025            &Engine,
6026        ) -> Result<R, Box<dyn std::error::Error>>,
6027    ) -> Result<R, Box<dyn std::error::Error>> {
6028        let mut guard = self.moe_cache.lock().unwrap();
6029        if guard.is_none() {
6030            *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
6031        }
6032        let cache = guard.as_mut().unwrap();
6033        f(cache, self)
6034    }
6035
6036    /// Freeze the already-built MoE residency set. This never constructs a cache: callers use it
6037    /// only after a real prefill has populated the machine-specific CPU/GPU working set.
6038    pub fn freeze_moe_cache(&self) {
6039        if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
6040            cache.freeze();
6041        }
6042    }
6043
6044    /// The current residency set as (layer, proj, ex) triples, or None if no cache was built.
6045    /// Never constructs a cache.
6046    pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
6047        self.moe_cache
6048            .lock()
6049            .unwrap()
6050            .as_ref()
6051            .map(crate::moe_cache::MoeSlotCache::export_residency)
6052    }
6053
6054    pub(crate) fn moe_cache_frozen(&self) -> bool {
6055        self.moe_cache
6056            .lock()
6057            .unwrap()
6058            .as_ref()
6059            .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
6060    }
6061
6062    /// A frozen heterogeneous CPU/GPU expert split cannot use Hy3's ordinary batched prefill
6063    /// efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing
6064    /// expert through the GPU spill path. Replay the short prompt through decode after freezing,
6065    /// while leaving the profiling warmup's established batched behavior untouched.
6066    /// (`pub`: run-gen's #46 batched-prime gate skips itself when generation will take the
6067    /// tokenwise arm anyway.)
6068    pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
6069        crate::cpu_experts::configured()
6070            && self.moe_cache_frozen()
6071            && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
6072    }
6073
6074    /// Install the loaded model's exact retained expert-block inventory before lazy cache build.
6075    pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
6076        assert!(
6077            self.moe_cache.lock().unwrap().is_none(),
6078            "MoE cache layout configured after cache construction"
6079        );
6080        *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
6081    }
6082
6083    pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
6084        self.moe_cache_layout.lock().unwrap().clone()
6085    }
6086
6087    /// True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
6088    pub fn moe_cache_enabled() -> bool {
6089        std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
6090    }
6091
6092    /// Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate.
6093    /// Returns None if the cache was never built (disabled or no MoE forward ran).
6094    pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
6095        let guard = self.moe_cache.lock().unwrap();
6096        guard
6097            .as_ref()
6098            .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
6099    }
6100
6101    /// Experimental CPU expert backend counters: completed layer calls, experts served, and the
6102    /// sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses;
6103    /// callers compare a before/after snapshot around a decode window.
6104    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6105    pub fn cpu_expert_stats(
6106        &self,
6107    ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
6108        crate::cpu_experts::configured().then(crate::cpu_experts::stats)
6109    }
6110
6111    /// Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure
6112    /// the backend tail that resident-GPU expert work did not hide.
6113    pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
6114        crate::cpu_experts::predictor_stats()
6115    }
6116
6117    pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
6118        crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
6119    }
6120
6121    /// CPU-routed expert selections grouped by how many of their three projections were already
6122    /// resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.
6123    pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
6124        crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
6125    }
6126
6127    /// Positioned-read proof-backend counters:
6128    /// `(reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full)`.
6129    pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
6130        let guard = self.moe_cache.lock().unwrap();
6131        guard
6132            .as_ref()
6133            .and_then(|cache| cache.pread_stats())
6134            .map(|stats| {
6135                (
6136                    stats.reads,
6137                    stats.bytes,
6138                    stats.read_errors,
6139                    stats.short_reads,
6140                    stats.fallbacks,
6141                    stats.buffer_waits,
6142                    stats.ring_full,
6143                )
6144            })
6145    }
6146
6147    /// Spill configuration values that warned and substituted their documented defaults.
6148    pub fn spill_config_fallbacks(&self) -> u64 {
6149        crate::spill_pread::config_fallbacks()
6150    }
6151
6152    /// Reset the MoE cache perf counters (to separate warmup from steady-state windows).
6153    pub fn moe_cache_reset_counters(&self) {
6154        if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
6155            c.reset_counters();
6156        }
6157    }
6158
6159    #[track_caller]
6160    pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6161        crate::alloc_trace_hit(v.len());
6162        Ok(self.gpu.stream().clone_htod(v)?)
6163    }
6164
6165    /// `htod_bytes` with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B
6166    /// past the final q4_0 block through their aligned window — the bytes never reach a
6167    /// result (funnelshift discards them) but must be mapped memory.
6168    pub fn htod_bytes_padded(
6169        &self,
6170        v: &[u8],
6171        pad: usize,
6172    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6173        let mut d = self.alloc_u8_uninit(v.len() + pad)?;
6174        {
6175            let mut view = d.slice_mut(0..v.len());
6176            self.gpu.stream().memcpy_htod(v, &mut view)?;
6177        }
6178        Ok(d)
6179    }
6180
6181    /// Device-to-device copy of `src` into `dst[off..off+len]` (f32). For in-place KV append.
6182    pub fn copy_into(
6183        &self,
6184        dst: &mut CudaSlice<f32>,
6185        off: usize,
6186        src: &CudaSlice<f32>,
6187        len: usize,
6188    ) -> Result<(), Box<dyn std::error::Error>> {
6189        let mut view = dst.slice_mut(off..off + len);
6190        self.gpu
6191            .stream()
6192            .memcpy_dtod(&src.slice(0..len), &mut view)?;
6193        Ok(())
6194    }
6195
6196    /// D2D copy with an offset on BOTH sides. `copy_into` always reads the source from 0,
6197    /// which cannot express "copy the TAIL of this buffer" — the shape a sliding-window draft
6198    /// KV export needs (lane/dspark-draft-plane-20260827).
6199    pub fn copy_range_into(
6200        &self,
6201        dst: &mut CudaSlice<f32>,
6202        dst_off: usize,
6203        src: &CudaSlice<f32>,
6204        src_off: usize,
6205        len: usize,
6206    ) -> Result<(), Box<dyn std::error::Error>> {
6207        let mut view = dst.slice_mut(dst_off..dst_off + len);
6208        self.gpu
6209            .stream()
6210            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut view)?;
6211        Ok(())
6212    }
6213
6214    /// View a sub-range of a device buffer (for attending over [0..len) of a KV cache).
6215    /// u8 twin of copy_into (D2D byte-range copy at an offset).
6216    pub fn copy_u8_into(
6217        &self,
6218        dst: &mut CudaSlice<u8>,
6219        off: usize,
6220        src: &CudaSlice<u8>,
6221        len: usize,
6222    ) -> Result<(), Box<dyn std::error::Error>> {
6223        // try_slice_mut, not slice_mut: an out-of-bounds range here panics the GPU worker
6224        // thread and takes the whole server with it (2026-08-29 warm-turn-at-40k incident).
6225        // A bounds miss is a caller bug, but it must fail the request, not the fleet.
6226        let cap = dst.len();
6227        let mut view = dst.try_slice_mut(off..off + len).ok_or_else(|| {
6228            format!(
6229                "copy_u8_into dst range [{off},{}) exceeds capacity {cap}",
6230                off + len,
6231            )
6232        })?;
6233        self.gpu
6234            .stream()
6235            .memcpy_dtod(&src.slice(0..len), &mut view)?;
6236        Ok(())
6237    }
6238
6239    /// D2D byte-range copy with explicit source and destination offsets.
6240    pub fn copy_u8_range_into(
6241        &self,
6242        dst: &mut CudaSlice<u8>,
6243        dst_off: usize,
6244        src: &CudaSlice<u8>,
6245        src_off: usize,
6246        len: usize,
6247    ) -> Result<(), Box<dyn std::error::Error>> {
6248        // try_slice_mut for the same reason as copy_u8_into: bounds misses fail the request,
6249        // never panic the worker.
6250        let cap = dst.len();
6251        let mut dst_view = dst.try_slice_mut(dst_off..dst_off + len).ok_or_else(|| {
6252            format!(
6253                "copy_u8_range_into dst range [{dst_off},{}) exceeds capacity {cap}",
6254                dst_off + len,
6255            )
6256        })?;
6257        self.gpu
6258            .stream()
6259            .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
6260        Ok(())
6261    }
6262
6263    /// Resolve an absolute append slot to the Step35 SWA layer's physical rows. At wrap, copy
6264    /// only the aligned live prefix through temporary device storage and rebase it at row zero,
6265    /// keeping the audited attention range contiguous without changing its absolute start.
6266    /// #[track_caller]: every ring-backed append that REBASES sets the plane's `base`, and a
6267    /// later append or rewind that needs a lower row is then refused. Three attempts at the
6268    /// SWA-ring lap failed because the writer that actually moved `base` was never the site being
6269    /// patched — the bare "SWA ring lapped required rows" message named neither the caller nor
6270    /// what it retained. Cost of the annotation is nothing; cost of not having it was two wrong
6271    /// fixes on hardware.
6272    #[track_caller]
6273    pub fn prepare_kv_append(
6274        &self,
6275        kv: &mut crate::cache::KvLayer,
6276        retain_from: usize,
6277        append_rows: usize,
6278    ) -> Result<usize, Box<dyn std::error::Error>> {
6279        let caller = std::panic::Location::caller();
6280        let base_before = kv.ring.as_ref().map(|r| r.base());
6281        let Some(plan) = kv
6282            .ring
6283            .as_ref()
6284            .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
6285            .transpose()
6286            .map_err(|err| -> Box<dyn std::error::Error> {
6287                format!(
6288                    "{err} [append len={} retain_from={retain_from} append_rows={append_rows}                      base={base_before:?} called from {caller}]",
6289                    kv.len
6290                )
6291                .into()
6292            })?
6293        else {
6294            return Ok(kv.len);
6295        };
6296        match plan {
6297            crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
6298            crate::cache::KvRingAppend::Rebase {
6299                src_row,
6300                keep_rows,
6301                new_base,
6302                write_row,
6303            } => {
6304                if keep_rows > 0 {
6305                    let k_len = keep_rows * kv.k_tok_bytes;
6306                    let v_len = keep_rows * kv.v_tok_bytes;
6307                    let mut k_tmp = self.alloc_u8_uninit(k_len)?;
6308                    let mut v_tmp = self.alloc_u8_uninit(v_len)?;
6309                    self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
6310                    self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
6311                    self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
6312                    self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
6313                }
6314                // One line per distinct (caller, new_base) so the writers that move `base` are
6315                // enumerable from a single run instead of inferred from which error fires.
6316                if std::env::var("MEMRA_KV_REBASE_TRACE").as_deref() == Ok("1") {
6317                    eprintln!(
6318                        "[kv-rebase] new_base={new_base} keep_rows={keep_rows} len={} \
6319                         retain_from={retain_from} called from {caller}",
6320                        kv.len
6321                    );
6322                }
6323                kv.ring.as_mut().unwrap().apply_rebase(new_base);
6324                // The dcw draft arm's device mirror of the ring base (see KvLayer::base_d).
6325                // Rebase is the ONLY writer of `base`, and rebases run host-side outside any
6326                // captured region, so this one line keeps the device view exact.
6327                if let Some(base_d) = kv.base_d.as_mut() {
6328                    self.set_i32_one(base_d, new_base as i32)?;
6329                }
6330                Ok(write_row)
6331            }
6332        }
6333    }
6334
6335    /// H2D write of `src` into `dst[off..off+src.len()]` (u8). In-place row updates for the
6336    /// adaptive trim head: no realloc, so captured graphs keep their baked addresses.
6337    pub fn htod_u8_into(
6338        &self,
6339        dst: &mut CudaSlice<u8>,
6340        off: usize,
6341        src: &[u8],
6342    ) -> Result<(), Box<dyn std::error::Error>> {
6343        let mut view = dst.slice_mut(off..off + src.len());
6344        self.gpu.stream().memcpy_htod(src, &mut view)?;
6345        Ok(())
6346    }
6347
6348    pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
6349        b.slice(0..len)
6350    }
6351
6352    /// View the first `len` BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)).
6353    /// Byte-range view (gemma4 R6 window offset into the quantized KV stream).
6354    pub fn view_u8_range<'a>(
6355        &self,
6356        b: &'a CudaSlice<u8>,
6357        start: usize,
6358        end: usize,
6359    ) -> cudarc::driver::CudaView<'a, u8> {
6360        b.slice(start..end)
6361    }
6362    pub fn view_u8<'a>(
6363        &self,
6364        b: &'a CudaSlice<u8>,
6365        len: usize,
6366    ) -> cudarc::driver::CudaView<'a, u8> {
6367        b.slice(0..len)
6368    }
6369
6370    /// Append-quantize ONE token's post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at
6371    /// token index `t` (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes
6372    /// the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].
6373    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6374    pub fn append_kv_quantized(
6375        &self,
6376        k_row: &CudaSlice<f32>,
6377        v_row: &CudaSlice<f32>,
6378        kc: &mut CudaSlice<u8>,
6379        vc: &mut CudaSlice<u8>,
6380        t: usize,
6381        kv_dim_k: usize,
6382        kv_dim_v: usize,
6383        k_tok_bytes: usize,
6384        v_tok_bytes: usize,
6385        g: bool,
6386    ) -> Result<(), Box<dyn std::error::Error>> {
6387        let f = if g {
6388            self.func_g("append_quantize_kv_q8_0_q5_1")
6389        } else {
6390            self.func("append_quantize_kv_q8_0_q5_1")
6391        };
6392        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6393        let cfg = LaunchConfig {
6394            grid_dim: (nblk, 1, 1),
6395            block_dim: (32, 1, 1),
6396            shared_mem_bytes: 0,
6397        };
6398        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
6399        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6400        let __s_b = self.gpu.stream();
6401        let mut b = __s_b.launch_builder(&f);
6402        b.arg(k_row)
6403            .arg(v_row)
6404            .arg(kc)
6405            .arg(vc)
6406            .arg(&ti)
6407            .arg(&kdk)
6408            .arg(&kdv)
6409            .arg(&ktb)
6410            .arg(&vtb);
6411        unsafe {
6412            b.launch(cfg)?;
6413        }
6414        Ok(())
6415    }
6416
6417    /// Device-counter variant of `append_kv_quantized` (CUDA-GRAPH-PLAN Phase 2): the write slot
6418    /// `t` is read from `t_dev[0]` (a resident device i32[1]) instead of a host int arg, so the
6419    /// launch args are FIXED across decode steps (graph-capturable). Identical quant math.
6420    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6421    pub fn append_kv_quantized_dc(
6422        &self,
6423        k_row: &CudaSlice<f32>,
6424        v_row: &CudaSlice<f32>,
6425        kc: &mut CudaSlice<u8>,
6426        vc: &mut CudaSlice<u8>,
6427        t_dev: &CudaSlice<i32>,
6428        kv_dim_k: usize,
6429        kv_dim_v: usize,
6430        k_tok_bytes: usize,
6431        v_tok_bytes: usize,
6432        g: bool,
6433    ) -> Result<(), Box<dyn std::error::Error>> {
6434        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6435        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
6436        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6437        // PDL wave-B2: flash-module flavor mirrors the builder path's g flag exactly.
6438        if Self::pdl_on() && Self::pdl_wb_on() {
6439            use cudarc::driver::{DevicePtr, DevicePtrMut};
6440            let s = &self.gpu.stream();
6441            let (pk, _g0) = k_row.device_ptr(s);
6442            let (pv, _g1) = v_row.device_ptr(s);
6443            let (pkc, _g2) = kc.device_ptr_mut(s);
6444            let (pvc, _g3) = vc.device_ptr_mut(s);
6445            let (pt, _g4) = t_dev.device_ptr(s);
6446            let mut ps = [
6447                &pk as *const _ as *mut std::ffi::c_void,
6448                &pv as *const _ as *mut _,
6449                &pkc as *const _ as *mut _,
6450                &pvc as *const _ as *mut _,
6451                &pt as *const _ as *mut _,
6452                &kdk as *const _ as *mut _,
6453                &kdv as *const _ as *mut _,
6454                &ktb as *const _ as *mut _,
6455                &vtb as *const _ as *mut _,
6456            ];
6457            unsafe {
6458                self.launch_pdl_flash(
6459                    g,
6460                    "append_quantize_kv_q8_0_q5_1_dc",
6461                    (nblk, 1, 1),
6462                    (32, 1, 1),
6463                    0,
6464                    &mut ps,
6465                )?;
6466            }
6467            return Ok(());
6468        }
6469        let f = if g {
6470            self.func_g("append_quantize_kv_q8_0_q5_1_dc")
6471        } else {
6472            self.func("append_quantize_kv_q8_0_q5_1_dc")
6473        };
6474        let cfg = LaunchConfig {
6475            grid_dim: (nblk, 1, 1),
6476            block_dim: (32, 1, 1),
6477            shared_mem_bytes: 0,
6478        };
6479        let __s_b = self.gpu.stream();
6480        let mut b = __s_b.launch_builder(&f);
6481        b.arg(k_row)
6482            .arg(v_row)
6483            .arg(kc)
6484            .arg(vc)
6485            .arg(t_dev)
6486            .arg(&kdk)
6487            .arg(&kdv)
6488            .arg(&ktb)
6489            .arg(&vtb);
6490        unsafe {
6491            b.launch(cfg)?;
6492        }
6493        Ok(())
6494    }
6495
6496    /// Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are
6497    /// token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the
6498    /// batched `_rows` kernel: one (nblk, T) launch whose per-(block,token) warp program is the
6499    /// per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential
6500    /// `append_kv_quantized_view` calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1
6501    /// forces the T-launch per-row loop (the A/B seam that measured the launch overhead).
6502    #[allow(clippy::too_many_arguments)]
6503    pub fn append_kv_quantized_rows(
6504        &self,
6505        k_rows: &CudaSlice<f32>,
6506        v_rows: &CudaSlice<f32>,
6507        kc: &mut CudaSlice<u8>,
6508        vc: &mut CudaSlice<u8>,
6509        t0: usize,
6510        t: usize,
6511        kv_dim_k: usize,
6512        kv_dim_v: usize,
6513        k_tok_bytes: usize,
6514        v_tok_bytes: usize,
6515        g: bool,
6516    ) -> Result<(), Box<dyn std::error::Error>> {
6517        if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
6518            for i in 0..t {
6519                let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
6520                let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
6521                self.append_kv_quantized_view(
6522                    &k_row,
6523                    &v_row,
6524                    kc,
6525                    vc,
6526                    t0 + i,
6527                    kv_dim_k,
6528                    kv_dim_v,
6529                    k_tok_bytes,
6530                    v_tok_bytes,
6531                    g,
6532                )?;
6533            }
6534            return Ok(());
6535        }
6536        let f = if g {
6537            self.func_g("append_quantize_kv_q8_0_q5_1_rows")
6538        } else {
6539            self.func("append_quantize_kv_q8_0_q5_1_rows")
6540        };
6541        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6542        let cfg = LaunchConfig {
6543            grid_dim: (nblk, t as u32, 1),
6544            block_dim: (32, 1, 1),
6545            shared_mem_bytes: 0,
6546        };
6547        let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
6548        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6549        let __s_b = self.gpu.stream();
6550        let mut b = __s_b.launch_builder(&f);
6551        b.arg(k_rows)
6552            .arg(v_rows)
6553            .arg(kc)
6554            .arg(vc)
6555            .arg(&t0i)
6556            .arg(&kdk)
6557            .arg(&kdv)
6558            .arg(&ktb)
6559            .arg(&vtb);
6560        unsafe {
6561            b.launch(cfg)?;
6562        }
6563        Ok(())
6564    }
6565
6566    /// Increment a device i32[1] counter in place (p[0] += 1) via the resident `inc_i32` kernel.
6567    /// Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and,
6568    /// later, inside a captured graph) without a host round-trip.
6569    pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
6570        let f = self.func("inc_i32");
6571        let cfg = LaunchConfig {
6572            grid_dim: (1, 1, 1),
6573            block_dim: (1, 1, 1),
6574            shared_mem_bytes: 0,
6575        };
6576        let __s_b = self.gpu.stream();
6577        let mut b = __s_b.launch_builder(&f);
6578        b.arg(p);
6579        unsafe {
6580            b.launch(cfg)?;
6581        }
6582        Ok(())
6583    }
6584
6585    /// Like `append_kv_quantized` but k_row/v_row are CudaViews (one token's row sliced out of a
6586    /// token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).
6587    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6588    pub fn append_kv_quantized_view(
6589        &self,
6590        k_row: &cudarc::driver::CudaView<f32>,
6591        v_row: &cudarc::driver::CudaView<f32>,
6592        kc: &mut CudaSlice<u8>,
6593        vc: &mut CudaSlice<u8>,
6594        t: usize,
6595        kv_dim_k: usize,
6596        kv_dim_v: usize,
6597        k_tok_bytes: usize,
6598        v_tok_bytes: usize,
6599        g: bool,
6600    ) -> Result<(), Box<dyn std::error::Error>> {
6601        let stream = self.gpu.stream();
6602        ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
6603        ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
6604        ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
6605        ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
6606        let f = if g {
6607            self.func_g("append_quantize_kv_q8_0_q5_1")
6608        } else {
6609            self.func("append_quantize_kv_q8_0_q5_1")
6610        };
6611        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6612        let cfg = LaunchConfig {
6613            grid_dim: (nblk, 1, 1),
6614            block_dim: (32, 1, 1),
6615            shared_mem_bytes: 0,
6616        };
6617        let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
6618        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6619        let mut b = stream.launch_builder(&f);
6620        b.arg(k_row)
6621            .arg(v_row)
6622            .arg(kc)
6623            .arg(vc)
6624            .arg(&ti)
6625            .arg(&kdk)
6626            .arg(&kdv)
6627            .arg(&ktb)
6628            .arg(&vtb);
6629        unsafe {
6630            b.launch(cfg)?;
6631        }
6632        Ok(())
6633    }
6634
6635    /// Device-to-device copy of a CudaView `src` into `dst[off..off+len]` (f32). Like `copy_into`
6636    /// but the source is a sub-view (e.g. one column of a token-major activation buffer).
6637    pub fn copy_view_into(
6638        &self,
6639        dst: &mut CudaSlice<f32>,
6640        off: usize,
6641        src: &cudarc::driver::CudaView<f32>,
6642        len: usize,
6643    ) -> Result<(), Box<dyn std::error::Error>> {
6644        let mut view = dst.slice_mut(off..off + len);
6645        self.gpu
6646            .stream()
6647            .memcpy_dtod(&src.slice(0..len), &mut view)?;
6648        Ok(())
6649    }
6650
6651    /// Real device-to-device COPY of `src` into a freshly allocated buffer. Used for cache
6652    /// snapshots (MTP-PLAN §D.4), where a snapshot must not alias the live buffer.
6653    ///
6654    /// CORRECTION (memra-next#23, verified against the LOCKED cudarc 0.19.8): this comment used to say
6655    /// "`CudaSlice::clone()` only bumps a refcount and would alias the live buffer". That is
6656    /// FALSE and it propagated — `impl Clone for CudaSlice` is `try_clone().unwrap()`, and
6657    /// `try_clone` is `self.stream.clone_dtod(self)`, so a plain `.clone()` already allocates and
6658    /// copies. Code that wants real aliasing needs an `Arc<CudaSlice<T>>` (see
6659    /// `vision::EmbedOverlay::rows`).
6660    ///
6661    /// THE TWO ARE NOT INTERCHANGEABLE, AND THE DIFFERENCE IS NOT ONLY FALLIBILITY — a second
6662    /// correction, from the peer review of that first one, because getting this backwards is how
6663    /// a residency bug gets written. `CudaSlice::clone()` allocates on the SLICE's own stream, so
6664    /// the copy lands in the SOURCE's context. This method allocates on `self.gpu.stream()`,
6665    /// which is the thread-local pp stage stream whenever a stage scope is active — so under a
6666    /// stage scope THIS method is the one that lands in a foreign context. Choose by what you
6667    /// need: `try_clone()` for a fallible copy that stays with the source, this method for a copy
6668    /// deliberately placed on the calling engine's current stream (and check the landing context
6669    /// if residency matters). Minor: cudarc's path uses an uninitialized alloc, this one
6670    /// `alloc_zeros`, i.e. an extra full memset.
6671    pub fn clone_dtod(
6672        &self,
6673        src: &CudaSlice<f32>,
6674    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6675        let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
6676        self.gpu.stream().memcpy_dtod(src, &mut dst)?;
6677        Ok(dst)
6678    }
6679
6680    /// D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into `dst`.
6681    /// Stream-ordered, async — decode_batch's per-sequence row plumbing.
6682    pub fn dtod_copy_view(
6683        &self,
6684        src: &cudarc::driver::CudaView<f32>,
6685        dst: &mut CudaSlice<f32>,
6686    ) -> Result<(), Box<dyn std::error::Error>> {
6687        self.gpu.stream().memcpy_dtod(src, dst)?;
6688        Ok(())
6689    }
6690
6691    /// D2D i8 twin of `dtod_copy_view` (q8_1 activation rows).
6692    pub fn dtod_copy_view_i8(
6693        &self,
6694        src: &cudarc::driver::CudaView<i8>,
6695        dst: &mut CudaSlice<i8>,
6696    ) -> Result<(), Box<dyn std::error::Error>> {
6697        self.gpu.stream().memcpy_dtod(src, dst)?;
6698        Ok(())
6699    }
6700
6701    /// D2D row placement: copy `src` into `dst[offset .. offset+src.len()]`.
6702    pub fn dtod_copy_into(
6703        &self,
6704        src: &CudaSlice<f32>,
6705        dst: &mut CudaSlice<f32>,
6706        offset: usize,
6707    ) -> Result<(), Box<dyn std::error::Error>> {
6708        let n = src.len();
6709        let mut dv = dst.slice_mut(offset..offset + n);
6710        self.gpu.stream().memcpy_dtod(src, &mut dv)?;
6711        Ok(())
6712    }
6713
6714    /// Batched uniform-size D2D copy (engine-bundle slice 1, DSF-ROUNDCOST-20260820 §1.1):
6715    /// `n` disjoint regions of `words` f32 each; `table` = [src_0..src_{n-1}, dst_0..dst_{n-1}]
6716    /// raw device pointers. ONE kernel launch replaces `n` memcpy_dtod dispatches — the dspark
6717    /// round's snap/commit copy dribble (~0.9 ms/round of dispatch serialization measured).
6718    /// Bytes and stream order are identical to the memcpy sequence it replaces.
6719    pub fn copy_batch_uniform_f32(
6720        &self,
6721        table: &CudaSlice<u64>,
6722        n: usize,
6723        words: usize,
6724    ) -> Result<(), Box<dyn std::error::Error>> {
6725        if n == 0 || words == 0 {
6726            return Ok(());
6727        }
6728        debug_assert!(
6729            table.len() >= 2 * n,
6730            "pointer table must hold n srcs + n dsts"
6731        );
6732        let f = self.func("copy_batch_uniform_f32");
6733        // Enough blocks to stream a multi-MB region, few enough that (chunks x n) stays a
6734        // sane grid: 512K-word ssm regions get 48 grid-striding blocks each.
6735        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
6736        let (ni, wi) = (n as i32, words as i32);
6737        let cfg = LaunchConfig {
6738            grid_dim: (chunks, n as u32, 1),
6739            block_dim: (256, 1, 1),
6740            shared_mem_bytes: 0,
6741        };
6742        let __s = self.gpu.stream();
6743        let mut b = __s.launch_builder(&f);
6744        b.arg(table).arg(&ni).arg(&wi);
6745        unsafe {
6746            b.launch(cfg)?;
6747        }
6748        Ok(())
6749    }
6750
6751    /// Copy one uniform quantized K/V row range for each layer in `table` and publish every
6752    /// layer's device length in the same launch. Table layout is five pointer planes:
6753    /// K source, V source, K destination, V destination, and i32 length destination.
6754    #[allow(clippy::too_many_arguments)] // allow: row bytes and source strides are independent K/V geometry and collapsing them would hide the peer-layout contract
6755    pub fn copy_batch_uniform_kv_u8_set_len(
6756        &self,
6757        table: &CudaSlice<u64>,
6758        n: usize,
6759        rows: usize,
6760        k_row_bytes: usize,
6761        v_row_bytes: usize,
6762        k_src_stride: usize,
6763        v_src_stride: usize,
6764        logical_len: usize,
6765    ) -> Result<(), Box<dyn std::error::Error>> {
6766        if n == 0 || rows == 0 || (k_row_bytes == 0 && v_row_bytes == 0) {
6767            return Ok(());
6768        }
6769        if table.len() < 5 * n {
6770            return Err(format!(
6771                "TP KV repair table has {} words, expected at least {}",
6772                table.len(),
6773                5 * n
6774            )
6775            .into());
6776        }
6777        let ni = i32::try_from(n).map_err(|_| "TP KV repair layer count exceeds i32")?;
6778        let rows = i32::try_from(rows).map_err(|_| "TP KV repair rows exceed i32")?;
6779        let kb = i32::try_from(k_row_bytes).map_err(|_| "TP KV repair K bytes exceed i32")?;
6780        let vb = i32::try_from(v_row_bytes).map_err(|_| "TP KV repair V bytes exceed i32")?;
6781        let ks = i32::try_from(k_src_stride).map_err(|_| "TP KV repair K stride exceeds i32")?;
6782        let vs = i32::try_from(v_src_stride).map_err(|_| "TP KV repair V stride exceeds i32")?;
6783        let len = i32::try_from(logical_len).map_err(|_| "TP KV repair length exceeds i32")?;
6784        let f = self.func("copy_batch_uniform_kv_u8_set_len");
6785        let cfg = LaunchConfig {
6786            grid_dim: (n as u32, 1, 1),
6787            block_dim: (256, 1, 1),
6788            shared_mem_bytes: 0,
6789        };
6790        let stream = self.gpu.stream();
6791        let mut builder = stream.launch_builder(&f);
6792        builder
6793            .arg(table)
6794            .arg(&ni)
6795            .arg(&rows)
6796            .arg(&kb)
6797            .arg(&vb)
6798            .arg(&ks)
6799            .arg(&vs)
6800            .arg(&len);
6801        unsafe {
6802            builder.launch(cfg)?;
6803        }
6804        Ok(())
6805    }
6806
6807    /// H2D refresh of an EXISTING u64 pointer table IN PLACE (stable pointer — the batched
6808    /// state-copy tables are refreshed per round because the GDN ssm handles ping-pong).
6809    pub fn htod_u64_into(
6810        &self,
6811        v: &[u64],
6812        dst: &mut CudaSlice<u64>,
6813    ) -> Result<(), Box<dyn std::error::Error>> {
6814        let mut view = dst.slice_mut(0..v.len());
6815        self.gpu.stream().memcpy_htod(v, &mut view)?;
6816        Ok(())
6817    }
6818
6819    /// f32 twin of [`Self::htod_u64_into`] (the MoE vrows scale tables through the
6820    /// verify-walk workspace, door W).
6821    pub fn htod_f32_into(
6822        &self,
6823        v: &[f32],
6824        dst: &mut CudaSlice<f32>,
6825    ) -> Result<(), Box<dyn std::error::Error>> {
6826        let mut view = dst.slice_mut(0..v.len());
6827        self.gpu.stream().memcpy_htod(v, &mut view)?;
6828        Ok(())
6829    }
6830
6831    /// `htod_f32_into` landing at an element offset: `dst[off..off+v.len()] = v`. The EP
6832    /// dispatch-diet's bulk peer-row return lands the peer's compact block directly into the
6833    /// pair-slab tail with ONE upload instead of a per-row scatter.
6834    pub fn htod_f32_into_at(
6835        &self,
6836        v: &[f32],
6837        dst: &mut CudaSlice<f32>,
6838        off: usize,
6839    ) -> Result<(), Box<dyn std::error::Error>> {
6840        if off + v.len() > dst.len() {
6841            return Err(format!(
6842                "htod_f32_into_at range {}..{} exceeds dst {}",
6843                off,
6844                off + v.len(),
6845                dst.len()
6846            )
6847            .into());
6848        }
6849        let mut view = dst.slice_mut(off..off + v.len());
6850        self.gpu.stream().memcpy_htod(v, &mut view)?;
6851        Ok(())
6852    }
6853
6854    /// Indirect-source copy (engine-bundle slice 3): the src ADDRESS is loaded from a
6855    /// device pointer-table entry at run time, so a captured graph follows the gdn
6856    /// ping-pong through the same table its scan kernels read — a baked memcpy node
6857    /// would keep the capture-time physical buffer. `dst_off`/`words` in f32 elements.
6858    pub fn copy_indirect_src_f32(
6859        &self,
6860        src_entry: &cudarc::driver::CudaView<u64>,
6861        dst: &mut CudaSlice<f32>,
6862        dst_off: usize,
6863        words: usize,
6864    ) -> Result<(), Box<dyn std::error::Error>> {
6865        let f = self.func("copy_indirect_src_f32");
6866        let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
6867        let wi = words as i32;
6868        let cfg = LaunchConfig {
6869            grid_dim: (chunks, 1, 1),
6870            block_dim: (256, 1, 1),
6871            shared_mem_bytes: 0,
6872        };
6873        let mut dv = dst.slice_mut(dst_off..dst_off + words);
6874        let __s = self.gpu.stream();
6875        let mut b = __s.launch_builder(&f);
6876        b.arg(src_entry).arg(&mut dv).arg(&wi);
6877        unsafe {
6878            b.launch(cfg)?;
6879        }
6880        Ok(())
6881    }
6882
6883    /// Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
6884    #[track_caller]
6885    pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6886        self.alloc_uninit::<i8>(n)
6887    }
6888
6889    /// Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.
6890    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6891    pub fn qmatvec(
6892        &self,
6893        w: &CudaSlice<u8>,
6894        x: &CudaSlice<f32>,
6895        m: usize,
6896        in_f: usize,
6897        out_f: usize,
6898        qtype: i32,
6899        row_bytes: usize,
6900    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6901        let f = self.func("qmatvec_f32");
6902        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
6903        let cfg = LaunchConfig {
6904            grid_dim: (out_f as u32, m as u32, 1),
6905            block_dim: (256, 1, 1),
6906            shared_mem_bytes: 0,
6907        };
6908        let (inf, outf, mi, qt, rb) =
6909            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
6910        let __s_b = self.gpu.stream();
6911        let mut b = __s_b.launch_builder(&f);
6912        b.arg(w)
6913            .arg(x)
6914            .arg(&mut y)
6915            .arg(&inf)
6916            .arg(&outf)
6917            .arg(&mi)
6918            .arg(&qt)
6919            .arg(&rb);
6920        unsafe {
6921            b.launch(cfg)?;
6922        }
6923        Ok(y)
6924    }
6925
6926    /// Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
6927    #[track_caller]
6928    pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6929        crate::alloc_trace_hit(n);
6930        let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
6931        self.keep_if_capturing(&s);
6932        Ok(s)
6933    }
6934
6935    /// Uninitialized u8 scratch — skips alloc_zeros' memset. ONLY for staging buffers whose read
6936    /// range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE
6937    /// STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).
6938    #[track_caller]
6939    pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6940        crate::alloc_trace_hit(n);
6941        let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
6942        self.keep_if_capturing(&s);
6943        Ok(s)
6944    }
6945
6946    /// Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out
6947    /// memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).
6948    pub fn memset_zeros_view(
6949        &self,
6950        dst: &mut cudarc::driver::CudaViewMut<f32>,
6951    ) -> Result<(), Box<dyn std::error::Error>> {
6952        self.gpu.stream().memset_zeros(dst)?;
6953        Ok(())
6954    }
6955
6956    /// EDGE-1 staging: copy `host_bytes` (a sub-slice of a HostExps buffer) into `scratch`
6957    /// at byte offset `off` (async H2D on the default stream). Length is host_bytes.len().
6958    /// The qmatvec_view that reads `scratch[off..]` is enqueued on the SAME stream after this,
6959    /// so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd
6960    /// stream would require an event).
6961    pub fn stage_expert(
6962        &self,
6963        host_bytes: &[u8],
6964        scratch: &mut CudaSlice<u8>,
6965        off: usize,
6966    ) -> Result<(), Box<dyn std::error::Error>> {
6967        let mut dst = scratch.slice_mut(off..off + host_bytes.len()); // CudaViewMut<u8>
6968        self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; // accepts &[u8] HostSlice src
6969        Ok(())
6970    }
6971
6972    /// Split-plane repack of a whole resident NVFP4 expert slab (`n_expert` experts of `rows`
6973    /// rows, `nsb64` 64-wide blocks per row) on the device: returns the repacked slab (same
6974    /// length), the interleaved source is the caller's to drop. memra#147.
6975    pub fn nvfp4_expert_split_repack(
6976        &self,
6977        src: &CudaSlice<u8>,
6978        n_expert: usize,
6979        rows: usize,
6980        nsb64: usize,
6981    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6982        let need = n_expert * rows * nsb64 * 36;
6983        if src.len() < need {
6984            return Err(format!(
6985                "nvfp4_expert_split_repack: slab holds {} bytes, {n_expert} x {rows} x {nsb64} x 36 = {need} needed",
6986                src.len()
6987            )
6988            .into());
6989        }
6990        let f = self.func("nvfp4_expert_split_repack");
6991        let mut dst = self.alloc_u8_uninit(src.len())?; // every byte of the repacked region is written; the pad tail is never read
6992        let nblk = (n_expert * rows * nsb64) as u32;
6993        let cfg = LaunchConfig {
6994            grid_dim: (nblk.div_ceil(256), 1, 1),
6995            block_dim: (256, 1, 1),
6996            shared_mem_bytes: 0,
6997        };
6998        let (ne, nr, ns) = (n_expert as i32, rows as i32, nsb64 as i32);
6999        let __s_b = self.gpu.stream();
7000        let mut b = __s_b.launch_builder(&f);
7001        b.arg(src).arg(&mut dst).arg(&ne).arg(&nr).arg(&ns);
7002        unsafe {
7003            b.launch(cfg)?;
7004        }
7005        Ok(dst)
7006    }
7007
7008    /// EDGE-1 §A: fused MoE router. `logits` is the router output [t, n_expert] (device, f32, the
7009    /// `gate_inp @ z` result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k
7010    /// expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the
7011    /// host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298).
7012    /// One CTA per token row, 256 threads (one per expert).
7013    pub fn moe_router_topk(
7014        &self,
7015        logits: &CudaSlice<f32>,
7016        t: usize,
7017        n_expert: usize,
7018        n_used: usize,
7019    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7020        let f = self.func("moe_router_topk_f32");
7021        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; // kernel fully overwrites
7022        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; // kernel fully overwrites
7023        let cfg = LaunchConfig {
7024            grid_dim: (t as u32, 1, 1),
7025            block_dim: (n_expert as u32, 1, 1),
7026            shared_mem_bytes: 0,
7027        };
7028        let (ne, nu) = (n_expert as i32, n_used as i32);
7029        let __s_b = self.gpu.stream();
7030        let mut b = __s_b.launch_builder(&f);
7031        b.arg(logits)
7032            .arg(&mut sel_idx)
7033            .arg(&mut sel_w)
7034            .arg(&ne)
7035            .arg(&nu);
7036        unsafe {
7037            b.launch(cfg)?;
7038        }
7039        Ok((sel_idx, sel_w))
7040    }
7041
7042    /// gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the
7043    /// separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).
7044    pub fn moe_router_topk_scaled(
7045        &self,
7046        logits: &CudaSlice<f32>,
7047        t: usize,
7048        n_expert: usize,
7049        n_used: usize,
7050        ex_scale: &CudaSlice<f32>,
7051    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7052        // barrier-lean v2 twin (per-warp top-k + one-warp merge) FALSIFIED 2026-07-14:
7053        // bit-identical streams but −1.4% (26B plain N=3 interleaved) — at t=1 the grid is
7054        // ONE block, so the 6.6us is launch/dependency overhead, not the barrier chain;
7055        // fewer barriers bought nothing and the merge structure cost. jsonl is the record.
7056        let f = self.func("moe_router_topk_scaled_f32");
7057        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
7058        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
7059        let cfg = LaunchConfig {
7060            grid_dim: (t as u32, 1, 1),
7061            block_dim: (n_expert as u32, 1, 1),
7062            shared_mem_bytes: 0,
7063        };
7064        let (ne, nu) = (n_expert as i32, n_used as i32);
7065        let __s_b = self.gpu.stream();
7066        let mut b = __s_b.launch_builder(&f);
7067        b.arg(logits)
7068            .arg(&mut sel_idx)
7069            .arg(&mut sel_w)
7070            .arg(&ne)
7071            .arg(&nu)
7072            .arg(ex_scale);
7073        unsafe {
7074            b.launch(cfg)?;
7075        }
7076        Ok((sel_idx, sel_w))
7077    }
7078
7079    /// LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old
7080    /// MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then
7081    /// dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the
7082    /// host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully
7083    /// overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer
7084    /// (flags=0 — cacheable, NOT cudarc's WRITECOMBINED default, so the host-side reads of sel/w
7085    /// stay cached), then ONE synchronize. Numerics identical to `moe_router_topk` (same kernel).
7086    pub fn moe_router_topk_host(
7087        &self,
7088        logits: &CudaSlice<f32>,
7089        t: usize,
7090        n_expert: usize,
7091        n_used: usize,
7092    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
7093        let f = self.func("moe_router_topk_f32");
7094        let n = t * n_used;
7095        let mut sel_idx = self.alloc_uninit::<i32>(n)?;
7096        let mut sel_w = self.alloc_uninit::<f32>(n)?;
7097        let cfg = LaunchConfig {
7098            grid_dim: (t as u32, 1, 1),
7099            block_dim: (n_expert as u32, 1, 1),
7100            shared_mem_bytes: 0,
7101        };
7102        let (ne, nu) = (n_expert as i32, n_used as i32);
7103        let __s_b = self.gpu.stream();
7104        let mut b = __s_b.launch_builder(&f);
7105        b.arg(logits)
7106            .arg(&mut sel_idx)
7107            .arg(&mut sel_w)
7108            .arg(&ne)
7109            .arg(&nu);
7110        unsafe {
7111            b.launch(cfg)?;
7112        }
7113        // single-sync readback: sel (i32) at offset 0, w (f32) at offset n*4 of the pinned stage.
7114        let bytes = n * 8;
7115        let mut guard = self.router_stage.lock().unwrap();
7116        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
7117            *guard = Some(PinnedStage::new(bytes.max(4096))?);
7118        }
7119        let stage = guard.as_mut().unwrap();
7120        let (si, sw) = unsafe {
7121            (
7122                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
7123                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
7124            )
7125        };
7126        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; // async (pinned dst)
7127        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; // async (pinned dst)
7128        self.gpu.stream().synchronize()?; // ONE sync for both
7129        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
7130    }
7131
7132    /// Device sigmoid router for Step-3.7 / DeepSeek-V3-class MoEs. `correction_bias` is added
7133    /// only to the top-k key; returned weights use the un-biased sigmoid score. `active` masks
7134    /// original expert ids before top-k. Exact key ties choose the smaller original id.
7135    #[allow(clippy::too_many_arguments)]
7136    pub fn moe_router_sigmoid_topk(
7137        &self,
7138        logits: &CudaSlice<f32>,
7139        t: usize,
7140        n_expert: usize,
7141        n_used: usize,
7142        active_count: usize,
7143        correction_bias: &CudaSlice<f32>,
7144        active: &CudaSlice<u8>,
7145        scaling_factor: f32,
7146        route_norm: bool,
7147    ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7148        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
7149        if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
7150            return Err(format!(
7151                "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
7152            )
7153            .into());
7154        }
7155        if logits.len() < t * n_expert
7156            || correction_bias.len() != n_expert
7157            || active.len() != n_expert
7158        {
7159            return Err(format!(
7160                "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
7161                logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
7162            ).into());
7163        }
7164        let f = self.func(crate::sigmoid_topk_kernel(
7165            crate::sig_expf_dev_on(),
7166            crate::topk_fast_on(),
7167            n_used,
7168        ));
7169        let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
7170        let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
7171        let threads = n_expert.div_ceil(32) * 32;
7172        let cfg = LaunchConfig {
7173            grid_dim: (t as u32, 1, 1),
7174            block_dim: (threads as u32, 1, 1),
7175            shared_mem_bytes: 0,
7176        };
7177        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
7178        let __s_b = self.gpu.stream();
7179        let mut b = __s_b.launch_builder(&f);
7180        b.arg(logits)
7181            .arg(correction_bias)
7182            .arg(active)
7183            .arg(&mut sel_idx)
7184            .arg(&mut sel_w)
7185            .arg(&ne)
7186            .arg(&nu)
7187            .arg(&scaling_factor)
7188            .arg(&rn);
7189        unsafe {
7190            b.launch(cfg)?;
7191        }
7192        Ok((sel_idx, sel_w))
7193    }
7194
7195    /// `moe_router_sigmoid_topk` writing into caller-owned buffers (alloc-free: child graphs
7196    /// cannot contain mem nodes, so the token-graph e-sections pre-own every output).
7197    #[allow(clippy::too_many_arguments)]
7198    /// Ring a doorbell flag at a RAW device address (see `memra_ring_flag`): one store of
7199    /// `value`, fenced. Used by a peer rank to signal join readiness into root memory, where
7200    /// the model engine can wait on it with a same-device stream memop.
7201    pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
7202        if ptr == 0 {
7203            return Err("ring_flag_raw: unarmed flag".into());
7204        }
7205        let f = self.func("memra_ring_flag");
7206        let cfg = LaunchConfig {
7207            grid_dim: (1, 1, 1),
7208            block_dim: (32, 1, 1),
7209            shared_mem_bytes: 0,
7210        };
7211        let __s_b = self.gpu.stream();
7212        let mut b = __s_b.launch_builder(&f);
7213        b.arg(&ptr).arg(&value);
7214        unsafe {
7215            b.launch(cfg)?;
7216        }
7217        Ok(())
7218    }
7219
7220    /// One-launch mirror of a routed selection (`sel` int32 + `route_w` f32) — see
7221    /// `moe_sel_w_mirror`. Replaces the two tiny D2D copies the rank pull used to issue.
7222    pub fn moe_sel_w_mirror(
7223        &self,
7224        sel_src: &CudaSlice<i32>,
7225        w_src: &CudaSlice<f32>,
7226        sel_dst: &mut CudaSlice<i32>,
7227        w_dst: &mut CudaSlice<f32>,
7228        n: usize,
7229    ) -> Result<(), Box<dyn std::error::Error>> {
7230        if n == 0
7231            || n > i32::MAX as usize
7232            || sel_src.len() < n
7233            || w_src.len() < n
7234            || sel_dst.len() < n
7235            || w_dst.len() < n
7236        {
7237            return Err(format!("moe_sel_w_mirror geometry n={n}").into());
7238        }
7239        let f = self.func("moe_sel_w_mirror");
7240        let threads = if n <= 32 { 32 } else { 128 };
7241        let cfg = LaunchConfig {
7242            grid_dim: ((n as u32).div_ceil(threads), 1, 1),
7243            block_dim: (threads, 1, 1),
7244            shared_mem_bytes: 0,
7245        };
7246        let ni = n as i32;
7247        let __s_b = self.gpu.stream();
7248        let mut b = __s_b.launch_builder(&f);
7249        b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
7250        unsafe {
7251            b.launch(cfg)?;
7252        }
7253        Ok(())
7254    }
7255
7256    /// One-launch W4A16 EP staging: peer-read the active f32 input plus routed ids/weights from
7257    /// the root device, round the input directly into the rank-local BF16 buffer, and mirror the
7258    /// fixed route metadata. The caller orders root production with an entry event.
7259    #[allow(clippy::too_many_arguments)]
7260    pub fn nvfp4_ep_stage_inputs(
7261        &self,
7262        input_src: &CudaSlice<f32>,
7263        sel_src: &CudaSlice<i32>,
7264        w_src: &CudaSlice<f32>,
7265        input_bf16_dst: &mut CudaSlice<u8>,
7266        sel_dst: &mut CudaSlice<i32>,
7267        w_dst: &mut CudaSlice<f32>,
7268        input_values: usize,
7269        pairs: usize,
7270        copy_weights: bool,
7271    ) -> Result<(), Box<dyn std::error::Error>> {
7272        if input_values == 0
7273            || pairs == 0
7274            || input_src.len() < input_values
7275            || sel_src.len() < pairs
7276            || w_src.len() < pairs
7277            || input_bf16_dst.len() < 2 * input_values
7278            || sel_dst.len() < pairs
7279            || w_dst.len() < pairs
7280        {
7281            return Err(format!(
7282                "W4A16 EP stage geometry input={} sel={} weights={} input_bf16={} \
7283                 sel_dst={} weights_dst={} active={input_values} pairs={pairs}",
7284                input_src.len(),
7285                sel_src.len(),
7286                w_src.len(),
7287                input_bf16_dst.len(),
7288                sel_dst.len(),
7289                w_dst.len(),
7290            )
7291            .into());
7292        }
7293        let f = self.func("nvfp4_ep_stage_inputs");
7294        let n = input_values.max(pairs);
7295        let cfg = LaunchConfig::for_num_elems(n as u32);
7296        let (input_values, pairs, copy_weights) =
7297            (input_values as i32, pairs as i32, i32::from(copy_weights));
7298        let __s_b = self.gpu.stream();
7299        let mut b = __s_b.launch_builder(&f);
7300        b.arg(input_src)
7301            .arg(sel_src)
7302            .arg(w_src)
7303            .arg(input_bf16_dst)
7304            .arg(sel_dst)
7305            .arg(w_dst)
7306            .arg(&input_values)
7307            .arg(&pairs)
7308            .arg(&copy_weights);
7309        unsafe {
7310            b.launch(cfg)?;
7311        }
7312        Ok(())
7313    }
7314
7315    /// Capture-safe twin of `nvfp4_ep_stage_inputs`: the three sources are persistent raw
7316    /// device addresses owned by the root engine. Destinations remain rank-local typed slices.
7317    #[allow(clippy::too_many_arguments)]
7318    pub fn nvfp4_ep_stage_inputs_raw(
7319        &self,
7320        input_src: u64,
7321        sel_src: u64,
7322        w_src: u64,
7323        input_bf16_dst: &mut CudaSlice<u8>,
7324        sel_dst: &mut CudaSlice<i32>,
7325        w_dst: &mut CudaSlice<f32>,
7326        input_values: usize,
7327        pairs: usize,
7328        copy_weights: bool,
7329    ) -> Result<(), Box<dyn std::error::Error>> {
7330        if input_src == 0
7331            || sel_src == 0
7332            || w_src == 0
7333            || input_values == 0
7334            || pairs == 0
7335            || input_bf16_dst.len() < 2 * input_values
7336            || sel_dst.len() < pairs
7337            || w_dst.len() < pairs
7338        {
7339            return Err(format!(
7340                "W4A16 EP raw stage geometry input={input_src:#x} sel={sel_src:#x} \
7341                 weights={w_src:#x} input_bf16={} sel_dst={} weights_dst={} \
7342                 active={input_values} pairs={pairs}",
7343                input_bf16_dst.len(),
7344                sel_dst.len(),
7345                w_dst.len(),
7346            )
7347            .into());
7348        }
7349        let f = self.func("nvfp4_ep_stage_inputs");
7350        let n = input_values.max(pairs);
7351        let cfg = LaunchConfig::for_num_elems(n as u32);
7352        let (input_values, pairs, copy_weights) =
7353            (input_values as i32, pairs as i32, i32::from(copy_weights));
7354        let __s_b = self.gpu.stream();
7355        let mut b = __s_b.launch_builder(&f);
7356        b.arg(&input_src)
7357            .arg(&sel_src)
7358            .arg(&w_src)
7359            .arg(input_bf16_dst)
7360            .arg(sel_dst)
7361            .arg(w_dst)
7362            .arg(&input_values)
7363            .arg(&pairs)
7364            .arg(&copy_weights);
7365        unsafe {
7366            b.launch(cfg)?;
7367        }
7368        Ok(())
7369    }
7370
7371    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7372    pub fn moe_router_sigmoid_topk_into(
7373        &self,
7374        logits: &CudaSlice<f32>,
7375        t: usize,
7376        n_expert: usize,
7377        n_used: usize,
7378        active_count: usize,
7379        correction_bias: &CudaSlice<f32>,
7380        active: &CudaSlice<u8>,
7381        scaling_factor: f32,
7382        route_norm: bool,
7383        sel_idx: &mut CudaSlice<i32>,
7384        sel_w: &mut CudaSlice<f32>,
7385    ) -> Result<(), Box<dyn std::error::Error>> {
7386        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
7387        if n_expert == 0
7388            || n_expert > 1024
7389            || n_used == 0
7390            || n_used > 32 // the kernels' shared pick cache (s_pick_w) is sized 32
7391            || n_used > n_expert
7392            || logits.len() < t * n_expert
7393            || correction_bias.len() != n_expert
7394            || active.len() != n_expert
7395            || sel_idx.len() < t * n_used
7396            || sel_w.len() < t * n_used
7397        {
7398            return Err("sigmoid router _into geometry mismatch".into());
7399        }
7400        let f = self.func(crate::sigmoid_topk_kernel(
7401            crate::sig_expf_dev_on(),
7402            crate::topk_fast_on(),
7403            n_used,
7404        ));
7405        let threads = n_expert.div_ceil(32) * 32;
7406        let cfg = LaunchConfig {
7407            grid_dim: (t as u32, 1, 1),
7408            block_dim: (threads as u32, 1, 1),
7409            shared_mem_bytes: 0,
7410        };
7411        let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
7412        let __s_b = self.gpu.stream();
7413        let mut b = __s_b.launch_builder(&f);
7414        b.arg(logits)
7415            .arg(correction_bias)
7416            .arg(active)
7417            .arg(&mut *sel_idx)
7418            .arg(&mut *sel_w)
7419            .arg(&ne)
7420            .arg(&nu)
7421            .arg(&scaling_factor)
7422            .arg(&rn);
7423        unsafe {
7424            b.launch(cfg)?;
7425        }
7426        Ok(())
7427    }
7428
7429    /// Single-sync pinned readback twin of `moe_router_sigmoid_topk`. This preserves the existing
7430    /// grouped/staged dispatch contract while replacing the full-logit DtoH plus host sigmoid/sort.
7431    #[allow(clippy::too_many_arguments)]
7432    pub fn moe_router_sigmoid_topk_host(
7433        &self,
7434        logits: &CudaSlice<f32>,
7435        t: usize,
7436        n_expert: usize,
7437        n_used: usize,
7438        active_count: usize,
7439        correction_bias: &CudaSlice<f32>,
7440        active: &CudaSlice<u8>,
7441        scaling_factor: f32,
7442        route_norm: bool,
7443    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
7444        // FAIL LOUD INSIDE A CAPTURE REGION. Under `cudaStreamCaptureModeRelaxed` the DtoH below
7445        // is RECORDED, not executed: the call returns success, the pinned stage keeps whatever
7446        // bytes it already held, and the caller routes on an uninitialised selection whose expert
7447        // ids index the slab out of range. The graph then bakes those garbage pointers and every
7448        // replay reproduces them, with no error anywhere — which is exactly the shape twelve
7449        // glm5 decode-graph box takes chased: `TOKEN MISMATCH step 1: eager=437 graph=0`, the
7450        // same constant token at every step, replays engaged, nothing in the log.
7451        //
7452        // The decode-graph door admits a stage only when every captured layer's T=1 device-table
7453        // MoE arm will fire, so reaching here under an open capture means the admission predicate
7454        // and the dispatch predicate disagreed. That is a defect either way; refusing by name
7455        // turns it into a named capture failure, and the door's contract sends the token down the
7456        // byte-identical eager walk instead of serving the garbage.
7457        if glm5_graph_capture_open() {
7458            return Err("moe_router_sigmoid_topk_host was reached inside an open CUDA graph                         capture: the host readback would record an unexecuted DtoH and route on                         uninitialised memory. The captured layer's device-table MoE arm did not                         fire, so the capture-admission predicate (glm5_t1_dev_moe_ready) and the                         dispatch predicate (vrows_fires) disagree"
7459                .into());
7460        }
7461        let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
7462            logits,
7463            t,
7464            n_expert,
7465            n_used,
7466            active_count,
7467            correction_bias,
7468            active,
7469            scaling_factor,
7470            route_norm,
7471        )?;
7472        let n = t * n_used;
7473        let bytes = n * 8;
7474        let mut guard = self.router_stage.lock().unwrap();
7475        if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
7476            *guard = Some(PinnedStage::new(bytes.max(4096))?);
7477        }
7478        let stage = guard.as_mut().unwrap();
7479        let (si, sw) = unsafe {
7480            (
7481                std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
7482                std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
7483            )
7484        };
7485        self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
7486        self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
7487        self.gpu.stream().synchronize()?;
7488        Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
7489    }
7490
7491    /// EDGE-1 §C.2: async H2D of `host_bytes` into `scratch[off..]` on the COPY stream, returning a
7492    /// recorded event the compute stream can `wait` on before the dependent GEMM. Used for in-token
7493    /// expert prefetch (pipeline by one). `host_bytes` should be pinned for a true DMA (§C.1).
7494    pub fn stage_expert_async(
7495        &self,
7496        host_bytes: &[u8],
7497        scratch: &mut CudaSlice<u8>,
7498        off: usize,
7499    ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
7500        let mut dst = scratch.slice_mut(off..off + host_bytes.len());
7501        self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
7502        Ok(self.copy_stream.record_event(None)?)
7503    }
7504
7505    /// Make the compute stream wait for an async copy event (the consumer side of `stage_expert_async`).
7506    pub fn compute_wait(
7507        &self,
7508        ev: &cudarc::driver::CudaEvent,
7509    ) -> Result<(), Box<dyn std::error::Error>> {
7510        self.gpu.stream().wait(ev)?;
7511        Ok(())
7512    }
7513
7514    /// qmatvec over a byte sub-range of a (resident/scratch) CudaSlice<u8> holding ONE expert
7515    /// matrix. x is a CudaView<f32> (a sliced row of z, or a sliced activation). Reuses the
7516    /// validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The
7517    /// CudaView base+offset pointer is honored by the launch arg.
7518    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7519    pub fn qmatvec_view(
7520        &self,
7521        w: &CudaSlice<u8>,
7522        range: std::ops::Range<usize>,
7523        x: &cudarc::driver::CudaView<f32>,
7524        m: usize,
7525        in_f: usize,
7526        out_f: usize,
7527        qtype: i32,
7528        row_bytes: usize,
7529    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7530        self.qmatvec_view_inner(w, range, x, m, in_f, out_f, qtype, row_bytes)
7531    }
7532
7533    /// W4A16 expert matvec: round the floating activation to checkpoint BF16 before the
7534    /// existing f32-dequant weight dot. The output remains f32. This is selected per model by
7535    /// `MoeWeights`; it is not a process-global NVFP4 policy.
7536    #[allow(clippy::too_many_arguments)]
7537    pub fn qmatvec_view_bf16_activation(
7538        &self,
7539        w: &CudaSlice<u8>,
7540        range: std::ops::Range<usize>,
7541        x: &cudarc::driver::CudaView<f32>,
7542        m: usize,
7543        in_f: usize,
7544        out_f: usize,
7545        qtype: i32,
7546        row_bytes: usize,
7547    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7548        let n = m * in_f;
7549        if x.len() != n {
7550            return Err(format!(
7551                "W4A16 BF16 activation input length {} != {m}x{in_f}",
7552                x.len()
7553            )
7554            .into());
7555        }
7556        let mut x_bf16 = self.alloc_u8_uninit(n * 2)?;
7557        self.f32_to_bf16_v(x, &mut x_bf16, n)?;
7558        let x_f32 = self.bf16_to_f32(&x_bf16.slice(0..n * 2), n)?;
7559        self.qmatvec_view_inner(
7560            w,
7561            range,
7562            &x_f32.slice(0..n),
7563            m,
7564            in_f,
7565            out_f,
7566            qtype,
7567            row_bytes,
7568        )
7569    }
7570
7571    #[allow(clippy::too_many_arguments)]
7572    fn qmatvec_view_inner(
7573        &self,
7574        w: &CudaSlice<u8>,
7575        range: std::ops::Range<usize>,
7576        x: &cudarc::driver::CudaView<f32>,
7577        m: usize,
7578        in_f: usize,
7579        out_f: usize,
7580        qtype: i32,
7581        row_bytes: usize,
7582    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7583        let f = self.func("qmatvec_f32");
7584        let wv = w.slice(range); // CudaView<u8>, offset honored
7585        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
7586        let cfg = LaunchConfig {
7587            grid_dim: (out_f as u32, m as u32, 1),
7588            block_dim: (256, 1, 1),
7589            shared_mem_bytes: 0,
7590        };
7591        let (inf, outf, mi, qt, rb) =
7592            (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
7593        let __s_b = self.gpu.stream();
7594        let mut b = __s_b.launch_builder(&f);
7595        b.arg(&wv)
7596            .arg(x)
7597            .arg(&mut y)
7598            .arg(&inf)
7599            .arg(&outf)
7600            .arg(&mi)
7601            .arg(&qt)
7602            .arg(&rb);
7603        unsafe {
7604            b.launch(cfg)?;
7605        }
7606        Ok(y)
7607    }
7608
7609    /// STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer's gate+up+SiLU for all `n_used` routed
7610    /// experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x
7611    /// silu_mul = 24 launches). `gp`/`up` are the 8 expert weight-block device pointers (SLRU
7612    /// cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff].
7613    /// BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32's exact 256-thread
7614    /// reduction; the SiLU epilogue is silu_mul_f32's exact expression (see kernel header).
7615    #[allow(clippy::too_many_arguments)]
7616    /// dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but
7617    /// consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree)
7618    /// — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.
7619    pub fn moe_gate_up_silu8_q8(
7620        &self,
7621        gp: WPtr8,
7622        up: WPtr8,
7623        aq: &CudaSlice<i8>,
7624        ad: &CudaSlice<f32>,
7625        in_f: usize,
7626        n_ff: usize,
7627        n_used: usize,
7628        qt_g: i32,
7629        qt_u: i32,
7630        rb_g: usize,
7631        rb_u: usize,
7632    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7633        let f = self.func("moe_gate_up_silu8_q8");
7634        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7635        let cfg = LaunchConfig {
7636            grid_dim: (n_ff as u32, n_used as u32, 1),
7637            block_dim: (32, 1, 1),
7638            shared_mem_bytes: 0,
7639        };
7640        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7641        let __s_b = self.gpu.stream();
7642        let mut b = __s_b.launch_builder(&f);
7643        b.arg(&gp)
7644            .arg(&up)
7645            .arg(aq)
7646            .arg(ad)
7647            .arg(&mut act)
7648            .arg(&inf)
7649            .arg(&nff)
7650            .arg(&qt_g)
7651            .arg(&qt_u)
7652            .arg(&rbg)
7653            .arg(&rbu);
7654        unsafe {
7655            b.launch(cfg)?;
7656        }
7657        Ok(act)
7658    }
7659
7660    /// The PRE-clamped, macro-folding twin of [`Engine::moe_gate_up_silu8_q8`] — the kernel
7661    /// class for any MoE family whose activation clamps the gate BEFORE the silu (glm5_next is
7662    /// the first such family; the door names the arithmetic, not the family).
7663    ///
7664    /// Same grid/block/dots/warp reduction; the epilogue is
7665    /// `silu(min(gate*gs, limit)) * clamp(up*us, ±limit)` — `swiglu_preclamped_mul_scaled_f32`'s
7666    /// expression verbatim — and `gs`/`us` carry the SELECTED experts' NVFP4 `weight_scale_2`
7667    /// macro scales in router slot order (1.0 for a macro-free bank).
7668    ///
7669    /// `limit` must be live: at `limit == 0` every gate collapses to `silu(0) == 0`, so a caller
7670    /// with no clamp belongs on the plain-SiLU sibling, not here. Same contract as
7671    /// [`Engine::swiglu_preclamped_mul_scaled`].
7672    #[allow(clippy::too_many_arguments)]
7673    pub fn moe_gate_up_preclamp8_q8(
7674        &self,
7675        gp: WPtr8,
7676        up: WPtr8,
7677        aq: &CudaSlice<i8>,
7678        ad: &CudaSlice<f32>,
7679        gs: F32x8,
7680        us: F32x8,
7681        limit: f32,
7682        in_f: usize,
7683        n_ff: usize,
7684        n_used: usize,
7685        qt_g: i32,
7686        qt_u: i32,
7687        rb_g: usize,
7688        rb_u: usize,
7689    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7690        debug_assert!(
7691            limit > 1e-6,
7692            "moe_gate_up_preclamp8_q8 needs a live limit; use moe_gate_up_silu8_q8"
7693        );
7694        let f = self.func("moe_gate_up_preclamp8_q8");
7695        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7696        let cfg = LaunchConfig {
7697            grid_dim: (n_ff as u32, n_used as u32, 1),
7698            block_dim: (32, 1, 1),
7699            shared_mem_bytes: 0,
7700        };
7701        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7702        let __s_b = self.gpu.stream();
7703        let mut b = __s_b.launch_builder(&f);
7704        b.arg(&gp)
7705            .arg(&up)
7706            .arg(aq)
7707            .arg(ad)
7708            .arg(&gs)
7709            .arg(&us)
7710            .arg(&limit)
7711            .arg(&mut act)
7712            .arg(&inf)
7713            .arg(&nff)
7714            .arg(&qt_g)
7715            .arg(&qt_u)
7716            .arg(&rbg)
7717            .arg(&rbu);
7718        unsafe {
7719            b.launch(cfg)?;
7720        }
7721        Ok(act)
7722    }
7723
7724    #[allow(clippy::too_many_arguments)]
7725    pub fn moe_down8_fma_q8(
7726        &self,
7727        dp: WPtr8,
7728        w: F32x8,
7729        aq2: &CudaSlice<i8>,
7730        ad2: &CudaSlice<f32>,
7731        dst: &mut cudarc::driver::CudaViewMut<f32>,
7732        in_f: usize,
7733        out_f: usize,
7734        n_used: usize,
7735        qt: i32,
7736        rb: usize,
7737    ) -> Result<(), Box<dyn std::error::Error>> {
7738        let f = self.func("moe_down8_fma_q8");
7739        let cfg = LaunchConfig {
7740            grid_dim: (out_f as u32, 1, 1),
7741            block_dim: (32, 1, 1),
7742            shared_mem_bytes: 0,
7743        };
7744        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
7745        let __s_b = self.gpu.stream();
7746        let mut b = __s_b.launch_builder(&f);
7747        b.arg(&dp)
7748            .arg(&w)
7749            .arg(aq2)
7750            .arg(ad2)
7751            .arg(dst)
7752            .arg(&inf)
7753            .arg(&outf)
7754            .arg(&nu)
7755            .arg(&qt)
7756            .arg(&rbi);
7757        unsafe {
7758            b.launch(cfg)?;
7759        }
7760        Ok(())
7761    }
7762
7763    /// WARP-PACKED twin of [`Engine::moe_gate_up_preclamp8_q8`] (MEMRA_B200_MATVEC_ARM occupancy
7764    /// arm, lane/b200-matvec-occupancy-20260902): MEMRA_MMVQ_ROWS warps/block on threadIdx.y
7765    /// instead of one warp/block, same per-warp body -> bit-identical per (o,j). See
7766    /// `b200_matvec_arm_on` and docs/FLAGS.md for the door.
7767    #[allow(clippy::too_many_arguments)]
7768    pub fn moe_gate_up_preclamp8_q8_w4(
7769        &self,
7770        gp: WPtr8,
7771        up: WPtr8,
7772        aq: &CudaSlice<i8>,
7773        ad: &CudaSlice<f32>,
7774        gs: F32x8,
7775        us: F32x8,
7776        limit: f32,
7777        in_f: usize,
7778        n_ff: usize,
7779        n_used: usize,
7780        qt_g: i32,
7781        qt_u: i32,
7782        rb_g: usize,
7783        rb_u: usize,
7784    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7785        debug_assert!(
7786            limit > 1e-6,
7787            "moe_gate_up_preclamp8_q8_w4 needs a live limit; use moe_gate_up_silu8_q8"
7788        );
7789        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
7790        let f = self.func("moe_gate_up_preclamp8_q8_w4");
7791        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7792        let cfg = LaunchConfig {
7793            grid_dim: ((n_ff as u32).div_ceil(ROWS), n_used as u32, 1),
7794            block_dim: (32, ROWS, 1),
7795            shared_mem_bytes: 0,
7796        };
7797        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7798        let __s_b = self.gpu.stream();
7799        let mut b = __s_b.launch_builder(&f);
7800        b.arg(&gp)
7801            .arg(&up)
7802            .arg(aq)
7803            .arg(ad)
7804            .arg(&gs)
7805            .arg(&us)
7806            .arg(&limit)
7807            .arg(&mut act)
7808            .arg(&inf)
7809            .arg(&nff)
7810            .arg(&qt_g)
7811            .arg(&qt_u)
7812            .arg(&rbg)
7813            .arg(&rbu);
7814        unsafe {
7815            b.launch(cfg)?;
7816        }
7817        Ok(act)
7818    }
7819
7820    /// WARP-PACKED twin of [`Engine::moe_down8_fma_q8`] (MEMRA_B200_MATVEC_ARM occupancy arm) —
7821    /// see [`Engine::moe_gate_up_preclamp8_q8_w4`].
7822    #[allow(clippy::too_many_arguments)]
7823    pub fn moe_down8_fma_q8_w4(
7824        &self,
7825        dp: WPtr8,
7826        w: F32x8,
7827        aq2: &CudaSlice<i8>,
7828        ad2: &CudaSlice<f32>,
7829        dst: &mut cudarc::driver::CudaViewMut<f32>,
7830        in_f: usize,
7831        out_f: usize,
7832        n_used: usize,
7833        qt: i32,
7834        rb: usize,
7835    ) -> Result<(), Box<dyn std::error::Error>> {
7836        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
7837        let f = self.func("moe_down8_fma_q8_w4");
7838        let cfg = LaunchConfig {
7839            grid_dim: ((out_f as u32).div_ceil(ROWS), 1, 1),
7840            block_dim: (32, ROWS, 1),
7841            shared_mem_bytes: 0,
7842        };
7843        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
7844        let __s_b = self.gpu.stream();
7845        let mut b = __s_b.launch_builder(&f);
7846        b.arg(&dp)
7847            .arg(&w)
7848            .arg(aq2)
7849            .arg(ad2)
7850            .arg(dst)
7851            .arg(&inf)
7852            .arg(&outf)
7853            .arg(&nu)
7854            .arg(&qt)
7855            .arg(&rbi);
7856        unsafe {
7857            b.launch(cfg)?;
7858        }
7859        Ok(())
7860    }
7861
7862    /// DEVICE-SIDE build of the verify-rows pair's pointer/scale tables (door D,
7863    /// `MEMRA_MOE_VROWS_DEV_TABLES`) from the router's own device selection. Replaces the host
7864    /// loop plus its two pageable HtoD, and lets the caller skip the router's pinned readback
7865    /// and its full `cuStreamSynchronize` entirely. Arithmetic is term-for-term the host loop's
7866    /// (see the kernel comment in `qmatvec.cu`), so the tables — and therefore every downstream
7867    /// byte — are identical.
7868    ///
7869    /// `macros` is the model's immutable `(gate, up, down)` `weight_scale_2` host planes, or
7870    /// `None` for a non-macro bank (the kernel then takes 1.0f, `macro_scale`'s own answer).
7871    /// The planes get a resident device mirror keyed by `(il, plane)` on first use — uploading
7872    /// them per call would ADD three HtoD to a door whose purpose is removing two.
7873    #[allow(clippy::too_many_arguments)]
7874    // allow: the parameter list mirrors the kernel/FFI/call contract
7875    pub fn moe_vrows_tables_from_sel(
7876        &self,
7877        sel: &CudaSlice<i32>,
7878        selw: &CudaSlice<f32>,
7879        il: u16,
7880        macros: Option<(&[f32], &[f32], &[f32])>,
7881        (pg, pu, pd): (u64, u64, u64),
7882        (sg, su, sd): (usize, usize, usize),
7883        n_pairs: usize,
7884        ptrs: &mut CudaSlice<u64>,
7885        scl: &mut CudaSlice<f32>,
7886    ) -> Result<(), Box<dyn std::error::Error>> {
7887        debug_assert!(sel.len() >= n_pairs && selw.len() >= n_pairs);
7888        // `>=` not `==`: door E appends a fourth (expert-major order) plane to the same table.
7889        debug_assert!(ptrs.len() >= 3 * n_pairs);
7890        debug_assert_eq!(scl.len(), 3 * n_pairs);
7891        // Resident macro mirrors, uploaded once per (layer, plane). The guard is held across the
7892        // launch because `CudaSlice` is not clonable — the same shape as the w8-mirror sites.
7893        let mut mac = self
7894            .vrows_macro_dev
7895            .lock()
7896            .map_err(|_| "vrows macro mirror map is poisoned")?;
7897        if let Some((hg, hu, hd)) = macros {
7898            for (plane, host) in [(0u8, hg), (1u8, hu), (2u8, hd)] {
7899                // `entry` rather than contains_key+insert: the upload is fallible, so it lands in
7900                // the Vacant arm instead of an `or_insert_with` closure.
7901                if let std::collections::hash_map::Entry::Vacant(slot) = mac.entry((il, plane)) {
7902                    slot.insert(self.htod(host)?);
7903                }
7904            }
7905        }
7906        // Absent macro planes: the three kernel pointers must still be legal device addresses,
7907        // so the call aliases the selection weights and never dereferences them (have_macros=0).
7908        let (mg, mu, md, have) = match macros {
7909            Some(_) => (
7910                mac.get(&(il, 0)).expect("gate macro mirror built above"),
7911                mac.get(&(il, 1)).expect("up macro mirror built above"),
7912                mac.get(&(il, 2)).expect("down macro mirror built above"),
7913                1i32,
7914            ),
7915            None => (selw, selw, selw, 0i32),
7916        };
7917        let f = self.func("moe_vrows_tables_from_sel");
7918        let threads = 128u32;
7919        let cfg = LaunchConfig {
7920            grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
7921            block_dim: (threads, 1, 1),
7922            shared_mem_bytes: 0,
7923        };
7924        let (sgi, sui, sdi) = (sg as i64, su as i64, sd as i64);
7925        let (npi, havei) = (n_pairs as i32, have);
7926        let __s_b = self.gpu.stream();
7927        let mut b = __s_b.launch_builder(&f);
7928        b.arg(sel)
7929            .arg(selw)
7930            .arg(mg)
7931            .arg(mu)
7932            .arg(md)
7933            .arg(&mut *ptrs)
7934            .arg(&mut *scl)
7935            .arg(&pg)
7936            .arg(&pu)
7937            .arg(&pd)
7938            .arg(&sgi)
7939            .arg(&sui)
7940            .arg(&sdi)
7941            .arg(&npi)
7942            .arg(&havei);
7943        unsafe {
7944            b.launch(cfg)?;
7945        }
7946        Ok(())
7947    }
7948
7949    /// DEVICE-SIDE build of the verify-rows pair's EXPERT-MAJOR order plane (door E,
7950    /// `MEMRA_MOE_VROWS_DEDUP_ORDER`) from the router's own device selection, written into the
7951    /// pointer table's fourth plane `ptrs[3*n_pairs ..)`. Bit-identical to
7952    /// [`crate::vrows_expert_major_order`]: both are a stable order on `(expert id, pair index)`,
7953    /// the kernel by counting rank (see its comment in `qmatvec.cu`), the host by a stable sort.
7954    ///
7955    /// This launch exists ONLY in the door-D (device tables) arm — the host arm appends the plane
7956    /// to the vector it already uploads, so it costs zero extra transfers there. Cost in the
7957    /// device arm: 42 launches/round = ~0.093 ms at the box's 2.216 us eager-launch constant,
7958    /// against a predicted -2.17 ms/round; folding it into `moe_vrows_tables_from_sel` (same
7959    /// inputs, same one-thread-per-pair grid) is the named follow-up that recovers it.
7960    pub fn moe_vrows_order_from_sel(
7961        &self,
7962        sel: &CudaSlice<i32>,
7963        n_pairs: usize,
7964        ptrs: &mut CudaSlice<u64>,
7965    ) -> Result<(), Box<dyn std::error::Error>> {
7966        debug_assert!(sel.len() >= n_pairs);
7967        debug_assert!(
7968            ptrs.len() >= 4 * n_pairs,
7969            "the order plane lives at ptrs[3*n_pairs .. 4*n_pairs)"
7970        );
7971        let f = self.func("moe_vrows_order_from_sel");
7972        let threads = 128u32;
7973        let cfg = LaunchConfig {
7974            grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
7975            block_dim: (threads, 1, 1),
7976            shared_mem_bytes: 0,
7977        };
7978        let np = n_pairs as i32;
7979        let __s_b = self.gpu.stream();
7980        let mut b = __s_b.launch_builder(&f);
7981        b.arg(sel).arg(&mut *ptrs).arg(&np);
7982        unsafe {
7983            b.launch(cfg)?;
7984        }
7985        Ok(())
7986    }
7987
7988    /// Verify-rows twin of [`Self::moe_gate_up_preclamp8_q8`] (lane/glm5-vrest): one launch
7989    /// covers ALL `n_pairs = t * n_used` routed pairs of a spec-verify batch. `ptrs` /
7990    /// `scl` are the `[3 * n_pairs]` plane-major (gate | up | down) expert-pointer and
7991    /// scale tables (gs | us | w*macro_down); per pair the kernel body is the t=1 fused
7992    /// epilogue's verbatim, bit-gated per row vs the sequential chain.
7993    #[allow(clippy::too_many_arguments)]
7994    // allow: the parameter list mirrors the kernel/FFI/call contract
7995    pub fn moe_gate_up_preclamp8_q8_rows(
7996        &self,
7997        ptrs: &CudaSlice<u64>,
7998        scl: &CudaSlice<f32>,
7999        aq: &CudaSlice<i8>,
8000        ad: &CudaSlice<f32>,
8001        limit: f32,
8002        in_f: usize,
8003        n_ff: usize,
8004        n_used: usize,
8005        n_pairs: usize,
8006        qt_g: i32,
8007        qt_u: i32,
8008        rb_g: usize,
8009        rb_u: usize,
8010    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8011        debug_assert!(
8012            limit > 1e-6,
8013            "moe_gate_up_preclamp8_q8_rows needs a live limit; the kernel collapses every gate \
8014             to silu(0) at limit 0"
8015        );
8016        debug_assert!(ptrs.len() >= 3 * n_pairs);
8017        debug_assert_eq!(scl.len(), 3 * n_pairs);
8018        // MEMRA_MOE_VROWS_DEDUP_ORDER (lane/glm5-dedup door E, default OFF): the `_ord` twin —
8019        // pair index the FASTEST grid dimension, walked in expert-major order from the table's
8020        // fourth plane, so two verify rows sharing an expert read the identical gate/up rows in
8021        // adjacent blocks. `ptrs.len() >= 4*n_pairs` is a REQUIREMENT not a hint: the door engages
8022        // only when the caller actually built the order plane, so a direct launcher call with the
8023        // shipped 3-plane table (every standing gate) keeps the shipped program. Door M wins the
8024        // tie by being tested first — the two are refused together rather than crossed.
8025        let packed = moe_vrows_pack_on();
8026        let ordered =
8027            !packed && moe_vrows_dedup_order_on() && ptrs.len() >= 4 * n_pairs && n_ff <= 65535;
8028        // MEMRA_MOE_VROWS_ILP (lane/glm5-moe-rows-ilp-20260904, default OFF): the `_ilp` twins,
8029        // interleaved-NVFP4 only, composed with door M as `_w4_ilp`. Refuses by name otherwise.
8030        let ilp = moe_vrows_ilp_on()
8031            && if (qt_g == QT_NVFP4 && qt_u == QT_NVFP4)
8032                || (qt_g == QT_NVFP4_V2 && qt_u == QT_NVFP4_V2)
8033            {
8034                true
8035            } else {
8036                moe_vrows_ilp_refuse("gate/up", if qt_g == QT_NVFP4 { qt_u } else { qt_g });
8037                false
8038            };
8039        if ilp && MOE_VROWS_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
8040            eprintln!(
8041                "[moe-vrows-ilp] engaged: verify-rows MoE pair with four groups' loads per lane \
8042                 hoisted ahead of their math (MEMRA_MOE_VROWS_ILP=1, packed={packed})"
8043            );
8044        }
8045        let (f, cfg) = if packed {
8046            if MOE_VROWS_PACK_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
8047                eprintln!(
8048                    "[moe-vrows-pack] engaged: 4-warp blocks on the verify-rows MoE pair \
8049                     (MEMRA_MOE_VROWS_PACK=1)"
8050                );
8051            }
8052            (
8053                self.func(if ilp {
8054                    "moe_gate_up_preclamp8_q8_rows_w4_ilp"
8055                } else {
8056                    "moe_gate_up_preclamp8_q8_rows_w4"
8057                }),
8058                LaunchConfig {
8059                    grid_dim: ((n_ff as u32).div_ceil(4), n_pairs as u32, 1),
8060                    block_dim: (32, 4, 1),
8061                    shared_mem_bytes: 0,
8062                },
8063            )
8064        } else if ordered {
8065            if MOE_VROWS_DEDUP_ORDER_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8066                == 0
8067            {
8068                eprintln!(
8069                    "[moe-vrows-dedup-order] engaged: verify-rows gate/up walks the pair union \
8070                     EXPERT-MAJOR with the pair index as the fastest grid dimension, so the \
8071                     21.96%-measured repeat visits read a shared expert slab's rows in adjacent \
8072                     blocks (MEMRA_MOE_VROWS_DEDUP_ORDER=1)"
8073                );
8074            }
8075            (
8076                self.func("moe_gate_up_preclamp8_q8_rows_ord"),
8077                LaunchConfig {
8078                    grid_dim: (n_pairs as u32, n_ff as u32, 1),
8079                    block_dim: (32, 1, 1),
8080                    shared_mem_bytes: 0,
8081                },
8082            )
8083        } else {
8084            (
8085                self.func(if ilp {
8086                    "moe_gate_up_preclamp8_q8_rows_ilp"
8087                } else {
8088                    "moe_gate_up_preclamp8_q8_rows"
8089                }),
8090                LaunchConfig {
8091                    grid_dim: (n_ff as u32, n_pairs as u32, 1),
8092                    block_dim: (32, 1, 1),
8093                    shared_mem_bytes: 0,
8094                },
8095            )
8096        };
8097        // Door W: the vrows launcher is verify-walk-only; act is a pooled draw.
8098        let mut act = self.vws_uninit(n_pairs * n_ff)?;
8099        let (inf, nff, nu, np) = (in_f as i32, n_ff as i32, n_used as i32, n_pairs as i32);
8100        let (rbg, rbu) = (rb_g as i64, rb_u as i64);
8101        let __s_b = self.gpu.stream();
8102        let mut b = __s_b.launch_builder(&f);
8103        b.arg(ptrs)
8104            .arg(scl)
8105            .arg(aq)
8106            .arg(ad)
8107            .arg(&limit)
8108            .arg(&mut act)
8109            .arg(&inf)
8110            .arg(&nff)
8111            .arg(&nu)
8112            .arg(&np)
8113            .arg(&qt_g)
8114            .arg(&qt_u)
8115            .arg(&rbg)
8116            .arg(&rbu);
8117        unsafe {
8118            b.launch(cfg)?;
8119        }
8120        Ok(act)
8121    }
8122
8123    /// Verify-rows twin of [`Self::moe_down8_fma_q8`] (lane/glm5-vrest): every verify row's
8124    /// slot-ordered down+FMA chain in one launch. `dst` is `[t, out_f]`, fully overwritten;
8125    /// `ptrs`/`scl` are the same tables the gate/up rows launch consumed (down plane).
8126    #[allow(clippy::too_many_arguments)]
8127    // allow: the parameter list mirrors the kernel/FFI/call contract
8128    pub fn moe_down8_fma_q8_rows(
8129        &self,
8130        ptrs: &CudaSlice<u64>,
8131        scl: &CudaSlice<f32>,
8132        aq2: &CudaSlice<i8>,
8133        ad2: &CudaSlice<f32>,
8134        dst: &mut CudaSlice<f32>,
8135        in_f: usize,
8136        out_f: usize,
8137        n_used: usize,
8138        n_pairs: usize,
8139        qt: i32,
8140        rb: usize,
8141    ) -> Result<(), Box<dyn std::error::Error>> {
8142        debug_assert!(ptrs.len() >= 3 * n_pairs);
8143        debug_assert_eq!(scl.len(), 3 * n_pairs);
8144        debug_assert_eq!(n_pairs % n_used, 0, "pairs are dense slot-major");
8145        let t = n_pairs / n_used;
8146        debug_assert!(dst.len() >= t * out_f);
8147        // MEMRA_MOE_VROWS_PACK (door M): the _w4 twin, same packing as the gate/up launch.
8148        let packed = moe_vrows_pack_on();
8149        // MEMRA_MOE_VROWS_DOWN_TMAJ (door E-down): grid transposed to (t, out_f) — token fastest —
8150        // so the t verify rows at one output row are adjacent blocks and a repeated expert's down
8151        // row is read once for every token that shares it. The slot-ordered __fmaf_rn chain is
8152        // inside the block and keeps its ORIGINAL slot order; only the grid moves. Needs no table
8153        // plane (the down chain cannot be permuted), so it composes with either table provenance.
8154        let tmaj = !packed && moe_vrows_down_tmaj_on() && out_f <= 65535;
8155        // MEMRA_MOE_VROWS_ILP: the down `_ilp` twins, interleaved-NVFP4 only (see gate/up).
8156        let ilp = moe_vrows_ilp_on()
8157            && if qt == QT_NVFP4 || qt == QT_NVFP4_V2 {
8158                true
8159            } else {
8160                moe_vrows_ilp_refuse("down", qt);
8161                false
8162            };
8163        if ilp {
8164            MOE_VROWS_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8165        }
8166        let (f, cfg) = if packed {
8167            (
8168                self.func(if ilp {
8169                    "moe_down8_fma_q8_rows_w4_ilp"
8170                } else {
8171                    "moe_down8_fma_q8_rows_w4"
8172                }),
8173                LaunchConfig {
8174                    grid_dim: ((out_f as u32).div_ceil(4), t as u32, 1),
8175                    block_dim: (32, 4, 1),
8176                    shared_mem_bytes: 0,
8177                },
8178            )
8179        } else if tmaj {
8180            if MOE_VROWS_DOWN_TMAJ_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8181                == 0
8182            {
8183                eprintln!(
8184                    "[moe-vrows-down-tmaj] engaged: verify-rows down/FMA grid transposed to \
8185                     (t, out_f) so the verify rows at one output row are adjacent blocks; the \
8186                     slot-ordered FMA chain is unchanged (MEMRA_MOE_VROWS_DOWN_TMAJ=1)"
8187                );
8188            }
8189            (
8190                self.func("moe_down8_fma_q8_rows_tmaj"),
8191                LaunchConfig {
8192                    grid_dim: (t as u32, out_f as u32, 1),
8193                    block_dim: (32, 1, 1),
8194                    shared_mem_bytes: 0,
8195                },
8196            )
8197        } else {
8198            (
8199                self.func(if ilp {
8200                    "moe_down8_fma_q8_rows_ilp"
8201                } else {
8202                    "moe_down8_fma_q8_rows"
8203                }),
8204                LaunchConfig {
8205                    grid_dim: (out_f as u32, t as u32, 1),
8206                    block_dim: (32, 1, 1),
8207                    shared_mem_bytes: 0,
8208                },
8209            )
8210        };
8211        let (inf, outf, nu, np, rbi) = (
8212            in_f as i32,
8213            out_f as i32,
8214            n_used as i32,
8215            n_pairs as i32,
8216            rb as i64,
8217        );
8218        let __s_b = self.gpu.stream();
8219        let mut b = __s_b.launch_builder(&f);
8220        b.arg(ptrs)
8221            .arg(scl)
8222            .arg(aq2)
8223            .arg(ad2)
8224            .arg(dst)
8225            .arg(&inf)
8226            .arg(&outf)
8227            .arg(&nu)
8228            .arg(&np)
8229            .arg(&qt)
8230            .arg(&rbi);
8231        unsafe {
8232            b.launch(cfg)?;
8233        }
8234        Ok(())
8235    }
8236
8237    /// q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).
8238    #[allow(clippy::too_many_arguments)]
8239    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
8240    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
8241    pub fn qmatvec_expert_q8(
8242        &self,
8243        w: &CudaSlice<u8>,
8244        range: std::ops::Range<usize>,
8245        aq: &CudaSlice<i8>,
8246        ad: &CudaSlice<f32>,
8247        m: usize,
8248        in_f: usize,
8249        out_f: usize,
8250        qtype: i32,
8251        row_bytes: usize,
8252    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8253        let f = self.func("qmatvec_expert_q8");
8254        let wv = w.slice(range);
8255        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
8256        const ROWS: u32 = 4; // MEMRA_MMVQ_ROWS
8257        let cfg = LaunchConfig {
8258            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
8259            block_dim: (32, ROWS, 1),
8260            shared_mem_bytes: 0,
8261        };
8262        let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8263        let __s_b = self.gpu.stream();
8264        let mut b = __s_b.launch_builder(&f);
8265        b.arg(&wv)
8266            .arg(aq)
8267            .arg(ad)
8268            .arg(&mut y)
8269            .arg(&inf)
8270            .arg(&outf)
8271            .arg(&mi)
8272            .arg(&qtype)
8273            .arg(&rbi);
8274        unsafe {
8275            b.launch(cfg)?;
8276        }
8277        Ok(y)
8278    }
8279
8280    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
8281    pub fn moe_gate_up_silu8(
8282        &self,
8283        gp: WPtr8,
8284        up: WPtr8,
8285        x: &cudarc::driver::CudaView<f32>,
8286        in_f: usize,
8287        n_ff: usize,
8288        n_used: usize,
8289        qt_g: i32,
8290        qt_u: i32,
8291        rb_g: usize,
8292        rb_u: usize,
8293    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8294        let f = self.func("moe_gate_up_silu8_f32");
8295        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
8296        let cfg = LaunchConfig {
8297            grid_dim: (n_ff as u32, n_used as u32, 1),
8298            block_dim: (256, 1, 1),
8299            shared_mem_bytes: 0,
8300        };
8301        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
8302        let __s_b = self.gpu.stream();
8303        let mut b = __s_b.launch_builder(&f);
8304        b.arg(&gp)
8305            .arg(&up)
8306            .arg(x)
8307            .arg(&mut act)
8308            .arg(&inf)
8309            .arg(&nff)
8310            .arg(&qt_g)
8311            .arg(&qt_u)
8312            .arg(&rbg)
8313            .arg(&rbu);
8314        unsafe {
8315            b.launch(cfg)?;
8316        }
8317        Ok(act)
8318    }
8319
8320    /// STAGE-2 GROUPED DECODE: one MoE layer's down-proj + weighted accumulation for all `n_used`
8321    /// routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing
8322    /// the token's moe_out row DIRECTLY (`dst` is the zeroed row; the in-kernel slot-ordered
8323    /// __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the
8324    /// zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).
8325    #[allow(clippy::too_many_arguments)]
8326    pub fn moe_down8_fma_into(
8327        &self,
8328        dp: WPtr8,
8329        w: F32x8,
8330        act: &CudaSlice<f32>,
8331        dst: &mut cudarc::driver::CudaViewMut<f32>,
8332        in_f: usize,
8333        out_f: usize,
8334        n_used: usize,
8335        qt: i32,
8336        rb: usize,
8337    ) -> Result<(), Box<dyn std::error::Error>> {
8338        let f = self.func("moe_down8_fma_f32");
8339        let cfg = LaunchConfig {
8340            grid_dim: (out_f as u32, 1, 1),
8341            block_dim: (256, 1, 1),
8342            shared_mem_bytes: 0,
8343        };
8344        let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
8345        let __s_b = self.gpu.stream();
8346        let mut b = __s_b.launch_builder(&f);
8347        b.arg(&dp)
8348            .arg(&w)
8349            .arg(act)
8350            .arg(dst)
8351            .arg(&inf)
8352            .arg(&outf)
8353            .arg(&nu)
8354            .arg(&qt)
8355            .arg(&rbv);
8356        unsafe {
8357            b.launch(cfg)?;
8358        }
8359        Ok(())
8360    }
8361
8362    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_gate_up_silu8` for FULLY-RESIDENT
8363    /// layers. The expert ids come from the router kernel's DEVICE `sel` output (no DtoH) and the
8364    /// weight pointers from the per-layer device table `[3, n_expert]` of slot base addresses.
8365    /// BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs).
8366    #[allow(clippy::too_many_arguments)]
8367    /// dp4a q8 twin of the _dev pair (resident-experts arc).
8368    ///
8369    /// GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to
8370    /// the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row;
8371    /// down's FMA chain stays slot-ordered serial). Seams:
8372    ///   MEMRA_MOE_DEVQ8_GU   = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp)
8373    ///                       | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block)
8374    ///                       | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only)
8375    ///                       | u64 (nsb==64 unrolled ILP twin, geometry unchanged)
8376    ///   MEMRA_MOE_DEVQ8_WPB  = warps per block for _r twins / z-rows for s2z (default 4)
8377    ///   MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8%
8378    ///                       decode on 35B/rtx6000) | 0 (base one-warp serial-slot) | 1 | 2 | 4 ->
8379    ///                       _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16
8380    ///                       only) | w8h2 (h2 x slot-parallel)
8381    #[allow(clippy::too_many_arguments)]
8382    /// MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.
8383    #[allow(clippy::too_many_arguments)]
8384    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
8385    pub fn moe_pairs_matvec_q8(
8386        &self,
8387        table: &CudaSlice<u64>,
8388        proj: i32,
8389        pair_tok: &CudaSlice<i32>,
8390        pair_ex: &CudaSlice<i32>,
8391        aq: &CudaSlice<i8>,
8392        ad: &CudaSlice<f32>,
8393        in_f: usize,
8394        out_f: usize,
8395        n_expert: usize,
8396        n_pairs: usize,
8397        qtype: i32,
8398        row_bytes: usize,
8399    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8400        let f = self.func("moe_pairs_matvec_q8");
8401        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8402        const ROWS: u32 = 4;
8403        let cfg = LaunchConfig {
8404            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
8405            block_dim: (32, ROWS, 1),
8406            shared_mem_bytes: 0,
8407        };
8408        let (inf, outf, ne, np, rbi) = (
8409            in_f as i32,
8410            out_f as i32,
8411            n_expert as i32,
8412            n_pairs as i32,
8413            row_bytes as i64,
8414        );
8415        let __s_b = self.gpu.stream();
8416        let mut b = __s_b.launch_builder(&f);
8417        b.arg(table)
8418            .arg(&proj)
8419            .arg(pair_tok)
8420            .arg(pair_ex)
8421            .arg(aq)
8422            .arg(ad)
8423            .arg(&mut y)
8424            .arg(&inf)
8425            .arg(&outf)
8426            .arg(&ne)
8427            .arg(&np)
8428            .arg(&qtype)
8429            .arg(&rbi);
8430        unsafe {
8431            b.launch(cfg)?;
8432        }
8433        Ok(y)
8434    }
8435
8436    /// Expert-major pair matvec (weight-reuse across each expert's token group).
8437    #[allow(clippy::too_many_arguments)]
8438    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
8439    pub fn moe_pairs_matvec_q8_em(
8440        &self,
8441        table: &CudaSlice<u64>,
8442        proj: i32,
8443        ex_ids: &CudaSlice<i32>,
8444        ex_off: &CudaSlice<i32>,
8445        ex_pairs: &CudaSlice<i32>,
8446        pair_tok: &CudaSlice<i32>,
8447        aq: &CudaSlice<i8>,
8448        ad: &CudaSlice<f32>,
8449        in_f: usize,
8450        out_f: usize,
8451        n_expert: usize,
8452        n_active: usize,
8453        n_pairs: usize,
8454        qtype: i32,
8455        row_bytes: usize,
8456    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8457        let f = self.func("moe_pairs_matvec_q8_em");
8458        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8459        const ROWS: u32 = 4;
8460        let cfg = LaunchConfig {
8461            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
8462            block_dim: (32, ROWS, 1),
8463            shared_mem_bytes: 0,
8464        };
8465        let (inf, outf, ne, na, rbi) = (
8466            in_f as i32,
8467            out_f as i32,
8468            n_expert as i32,
8469            n_active as i32,
8470            row_bytes as i64,
8471        );
8472        let __s_b = self.gpu.stream();
8473        let mut b = __s_b.launch_builder(&f);
8474        b.arg(table)
8475            .arg(&proj)
8476            .arg(ex_ids)
8477            .arg(ex_off)
8478            .arg(ex_pairs)
8479            .arg(pair_tok)
8480            .arg(aq)
8481            .arg(ad)
8482            .arg(&mut y)
8483            .arg(&inf)
8484            .arg(&outf)
8485            .arg(&ne)
8486            .arg(&na)
8487            .arg(&qtype)
8488            .arg(&rbi);
8489        unsafe {
8490            b.launch(cfg)?;
8491        }
8492        Ok(y)
8493    }
8494
8495    // Decode-once expert-major MMQ (rung 3). Same CSR inputs/geometry as _em; kernel dequants each
8496    // weight group once per (row,group) then dp4a's across the expert's token group.
8497    #[allow(clippy::too_many_arguments)]
8498    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
8499    pub fn moe_pairs_matvec_q8_dec(
8500        &self,
8501        table: &CudaSlice<u64>,
8502        proj: i32,
8503        ex_ids: &CudaSlice<i32>,
8504        ex_off: &CudaSlice<i32>,
8505        ex_pairs: &CudaSlice<i32>,
8506        pair_tok: &CudaSlice<i32>,
8507        aq: &CudaSlice<i8>,
8508        ad: &CudaSlice<f32>,
8509        in_f: usize,
8510        out_f: usize,
8511        n_expert: usize,
8512        n_active: usize,
8513        n_pairs: usize,
8514        qtype: i32,
8515        row_bytes: usize,
8516    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8517        let f = self.func("moe_pairs_matvec_q8_dec");
8518        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8519        const ROWS: u32 = 4;
8520        let cfg = LaunchConfig {
8521            grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
8522            block_dim: (32, ROWS, 1),
8523            shared_mem_bytes: 0,
8524        };
8525        let (inf, outf, ne, na, rbi) = (
8526            in_f as i32,
8527            out_f as i32,
8528            n_expert as i32,
8529            n_active as i32,
8530            row_bytes as i64,
8531        );
8532        let __s_b = self.gpu.stream();
8533        let mut b = __s_b.launch_builder(&f);
8534        b.arg(table)
8535            .arg(&proj)
8536            .arg(ex_ids)
8537            .arg(ex_off)
8538            .arg(ex_pairs)
8539            .arg(pair_tok)
8540            .arg(aq)
8541            .arg(ad)
8542            .arg(&mut y)
8543            .arg(&inf)
8544            .arg(&outf)
8545            .arg(&ne)
8546            .arg(&na)
8547            .arg(&qtype)
8548            .arg(&rbi);
8549        unsafe {
8550            b.launch(cfg)?;
8551        }
8552        Ok(y)
8553    }
8554
8555    pub fn moe_pairs_gelu_mul(
8556        &self,
8557        gate: &CudaSlice<f32>,
8558        up: &CudaSlice<f32>,
8559        n: usize,
8560    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8561        let f = self.func("moe_pairs_gelu_mul");
8562        let mut act = self.alloc_uninit::<f32>(n)?;
8563        let cfg = LaunchConfig::for_num_elems(n as u32);
8564        let nl = n as i64;
8565        let __s_b = self.gpu.stream();
8566        let mut b = __s_b.launch_builder(&f);
8567        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
8568        unsafe {
8569            b.launch(cfg)?;
8570        }
8571        Ok(act)
8572    }
8573
8574    pub fn moe_pairs_silu_mul(
8575        &self,
8576        gate: &CudaSlice<f32>,
8577        up: &CudaSlice<f32>,
8578        n: usize,
8579    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8580        let f = self.func("moe_pairs_silu_mul");
8581        let mut act = self.alloc_uninit::<f32>(n)?;
8582        let cfg = LaunchConfig::for_num_elems(n as u32);
8583        let nl = n as i64;
8584        let __s_b = self.gpu.stream();
8585        let mut b = __s_b.launch_builder(&f);
8586        b.arg(gate).arg(up).arg(&mut act).arg(&nl);
8587        unsafe {
8588            b.launch(cfg)?;
8589        }
8590        Ok(act)
8591    }
8592
8593    #[allow(clippy::too_many_arguments)]
8594    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
8595    pub fn moe_pairs_scatter(
8596        &self,
8597        y_down: &CudaSlice<f32>,
8598        pair_w: &CudaSlice<f32>,
8599        tok_pair_off: &CudaSlice<i32>,
8600        tok_pair_ids: &CudaSlice<i32>,
8601        moe_out: &mut CudaSlice<f32>,
8602        t: usize,
8603        n_embd: usize,
8604    ) -> Result<(), Box<dyn std::error::Error>> {
8605        let f = self.func("moe_pairs_scatter");
8606        let cfg = LaunchConfig {
8607            grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
8608            block_dim: (256, 1, 1),
8609            shared_mem_bytes: 0,
8610        };
8611        let ne = n_embd as i32;
8612        let __s_b = self.gpu.stream();
8613        let mut b = __s_b.launch_builder(&f);
8614        b.arg(y_down)
8615            .arg(pair_w)
8616            .arg(tok_pair_off)
8617            .arg(tok_pair_ids)
8618            .arg(moe_out)
8619            .arg(&ne);
8620        unsafe {
8621            b.launch(cfg)?;
8622        }
8623        Ok(())
8624    }
8625
8626    /// gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2
8627    /// twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model
8628    /// decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).
8629    #[allow(clippy::too_many_arguments)]
8630    pub fn moe_gate_up_gelu8_dev_q8(
8631        &self,
8632        table: &CudaSlice<u64>,
8633        sel: &cudarc::driver::CudaView<i32>,
8634        aq: &CudaSlice<i8>,
8635        ad: &CudaSlice<f32>,
8636        in_f: usize,
8637        n_ff: usize,
8638        n_used: usize,
8639        n_expert: usize,
8640        qt_g: i32,
8641        qt_u: i32,
8642        rb_g: usize,
8643        rb_u: usize,
8644    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8645        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
8646        let (inf, nff, ne, rbg, rbu) = (
8647            in_f as i32,
8648            n_ff as i32,
8649            n_expert as i32,
8650            rb_g as i64,
8651            rb_u as i64,
8652        );
8653        let f = self.func("moe_gate_up_gelu8_dev_q8");
8654        let cfg = LaunchConfig {
8655            grid_dim: (n_ff as u32, n_used as u32, 1),
8656            block_dim: (32, 1, 1),
8657            shared_mem_bytes: 0,
8658        };
8659        let __s_b = self.gpu.stream();
8660        let mut b = __s_b.launch_builder(&f);
8661        b.arg(table)
8662            .arg(sel)
8663            .arg(aq)
8664            .arg(ad)
8665            .arg(&mut act)
8666            .arg(&inf)
8667            .arg(&nff)
8668            .arg(&ne)
8669            .arg(&qt_g)
8670            .arg(&qt_u)
8671            .arg(&rbg)
8672            .arg(&rbu);
8673        unsafe {
8674            b.launch(cfg)?;
8675        }
8676        Ok(act)
8677    }
8678
8679    /// gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).
8680    #[allow(clippy::too_many_arguments)]
8681    pub fn moe_gate_up_gelu8_dev_q8_rows(
8682        &self,
8683        table: &CudaSlice<u64>,
8684        sel: &CudaSlice<i32>,
8685        aq: &CudaSlice<i8>,
8686        ad: &CudaSlice<f32>,
8687        t: usize,
8688        in_f: usize,
8689        n_ff: usize,
8690        n_used: usize,
8691        n_expert: usize,
8692        qt_g: i32,
8693        qt_u: i32,
8694        rb_g: usize,
8695        rb_u: usize,
8696    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8697        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
8698        let (inf, nff, ne, rbg, rbu, nu) = (
8699            in_f as i32,
8700            n_ff as i32,
8701            n_expert as i32,
8702            rb_g as i64,
8703            rb_u as i64,
8704            n_used as i32,
8705        );
8706        let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
8707        let cfg = LaunchConfig {
8708            grid_dim: (n_ff as u32, n_used as u32, t as u32),
8709            block_dim: (32, 1, 1),
8710            shared_mem_bytes: 0,
8711        };
8712        let __s_b = self.gpu.stream();
8713        let mut b = __s_b.launch_builder(&f);
8714        b.arg(table)
8715            .arg(sel)
8716            .arg(aq)
8717            .arg(ad)
8718            .arg(&mut act)
8719            .arg(&inf)
8720            .arg(&nff)
8721            .arg(&ne)
8722            .arg(&qt_g)
8723            .arg(&qt_u)
8724            .arg(&rbg)
8725            .arg(&rbu)
8726            .arg(&nu);
8727        unsafe {
8728            b.launch(cfg)?;
8729        }
8730        Ok(act)
8731    }
8732
8733    /// gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).
8734    #[allow(clippy::too_many_arguments)]
8735    pub fn moe_gate_up_gelu8_dev_q8_csr(
8736        &self,
8737        table: &CudaSlice<u64>,
8738        sel: &CudaSlice<i32>,
8739        aq: &CudaSlice<i8>,
8740        ad: &CudaSlice<f32>,
8741        n_pairs: usize,
8742        in_f: usize,
8743        n_ff: usize,
8744        n_used: usize,
8745        n_expert: usize,
8746        qt_g: i32,
8747        qt_u: i32,
8748        rb_g: usize,
8749        rb_u: usize,
8750    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8751        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
8752        let (inf, nff, ne, rbg, rbu, nu, npi) = (
8753            in_f as i32,
8754            n_ff as i32,
8755            n_expert as i32,
8756            rb_g as i64,
8757            rb_u as i64,
8758            n_used as i32,
8759            n_pairs as i32,
8760        );
8761        let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
8762        let cfg = LaunchConfig {
8763            grid_dim: (n_ff as u32, n_pairs as u32, 1),
8764            block_dim: (32, 1, 1),
8765            shared_mem_bytes: 0,
8766        };
8767        let __s_b = self.gpu.stream();
8768        let mut b = __s_b.launch_builder(&f);
8769        b.arg(table)
8770            .arg(sel)
8771            .arg(aq)
8772            .arg(ad)
8773            .arg(&mut act)
8774            .arg(&inf)
8775            .arg(&nff)
8776            .arg(&ne)
8777            .arg(&qt_g)
8778            .arg(&qt_u)
8779            .arg(&rbg)
8780            .arg(&rbu)
8781            .arg(&nu)
8782            .arg(&npi);
8783        unsafe {
8784            b.launch(cfg)?;
8785        }
8786        Ok(act)
8787    }
8788
8789    /// gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).
8790    #[allow(clippy::too_many_arguments)]
8791    pub fn moe_down8_fma_dev_q8_rows_g(
8792        &self,
8793        table: &CudaSlice<u64>,
8794        sel: &CudaSlice<i32>,
8795        w: &CudaSlice<f32>,
8796        aq2: &CudaSlice<i8>,
8797        ad2: &CudaSlice<f32>,
8798        dst: &mut CudaSlice<f32>,
8799        t: usize,
8800        in_f: usize,
8801        out_f: usize,
8802        n_used: usize,
8803        n_expert: usize,
8804        qt: i32,
8805        rb: usize,
8806    ) -> Result<(), Box<dyn std::error::Error>> {
8807        let (inf, outf, nu, ne, rbi) = (
8808            in_f as i32,
8809            out_f as i32,
8810            n_used as i32,
8811            n_expert as i32,
8812            rb as i64,
8813        );
8814        // Exact Step-3.7 B=1 shape: expose the eight independent slot dots as
8815        // eight warps, then replay the original slot-ordered FMA chain. Every
8816        // other shape retains the generic one-warp rows kernel.
8817        let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
8818        let f = self.func(if step_b1_w8 {
8819            "moe_down8_fma_dev_q8_rows_w8"
8820        } else {
8821            "moe_down8_fma_dev_q8_rows_g"
8822        });
8823        let cfg = LaunchConfig {
8824            grid_dim: (out_f as u32, 1, t as u32),
8825            block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
8826            shared_mem_bytes: 0,
8827        };
8828        let __s_b = self.gpu.stream();
8829        let mut b = __s_b.launch_builder(&f);
8830        b.arg(table)
8831            .arg(sel)
8832            .arg(w)
8833            .arg(aq2)
8834            .arg(ad2)
8835            .arg(dst)
8836            .arg(&inf)
8837            .arg(&outf)
8838            .arg(&nu)
8839            .arg(&ne)
8840            .arg(&qt)
8841            .arg(&rbi);
8842        unsafe {
8843            b.launch(cfg)?;
8844        }
8845        Ok(())
8846    }
8847
8848    /// rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0
8849    /// split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise
8850    /// identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.
8851    pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
8852        let (out_f, in_f) = (2048usize, 2816usize);
8853        let nblk = in_f / 32;
8854        let mut seed = 0x9E3779B97F4A7C15u64;
8855        let mut rng = move || {
8856            seed = seed
8857                .wrapping_mul(6364136223846793005)
8858                .wrapping_add(1442695040888963407);
8859            (seed >> 33) as u8
8860        };
8861        let mut w = vec![0u8; out_f * nblk * 18];
8862        for b in w.iter_mut() {
8863            *b = rng();
8864        }
8865        for r in 0..out_f {
8866            for g in 0..nblk {
8867                let off = (r * nblk + g) * 18;
8868                w[off] = 0x00;
8869                w[off + 1] = 0x2C; // sane half d
8870            }
8871        }
8872        let qplane = out_f * nblk * 16;
8873        let mut wrp = vec![0u8; w.len()];
8874        for r in 0..out_f {
8875            for g in 0..nblk {
8876                let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
8877                wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
8878                    .copy_from_slice(&src[0..2]);
8879                wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
8880            }
8881        }
8882        let w_d = self.htod_bytes(&w)?;
8883        let wrp_d = self.htod_bytes(&wrp)?;
8884        let mut aq = vec![0i8; m * in_f];
8885        for v in aq.iter_mut() {
8886            *v = rng() as i8;
8887        }
8888        let aq_d = self.htod_i8(&aq)?;
8889        let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
8890        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
8891        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
8892        const RPB: u32 = 4;
8893        let cfg = LaunchConfig {
8894            grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
8895            block_dim: (32, RPB, 1),
8896            shared_mem_bytes: 0,
8897        };
8898        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
8899        let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
8900        let fb = self.func("qmatvec_q4_0_mmvq_b4");
8901        let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
8902        {
8903            let __s_b = self.gpu.stream();
8904            let mut b = __s_b.launch_builder(&fb);
8905            b.arg(&w_d)
8906                .arg(&aq_d)
8907                .arg(&ad_d)
8908                .arg(&mut y0)
8909                .arg(&inf)
8910                .arg(&outf)
8911                .arg(&mi)
8912                .arg(&rb);
8913            unsafe {
8914                b.launch(cfg)?;
8915            }
8916            let __s_b = self.gpu.stream();
8917            let mut b = __s_b.launch_builder(&fr);
8918            b.arg(&wrp_d)
8919                .arg(&aq_d)
8920                .arg(&ad_d)
8921                .arg(&mut y1)
8922                .arg(&inf)
8923                .arg(&outf)
8924                .arg(&mi)
8925                .arg(&qp);
8926            unsafe {
8927                b.launch(cfg)?;
8928            }
8929        }
8930        self.gpu.stream().synchronize()?;
8931        let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
8932        let nd = h0
8933            .iter()
8934            .zip(&h1)
8935            .filter(|(a, b)| a.to_bits() != b.to_bits())
8936            .count();
8937        if nd != 0 {
8938            return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
8939        }
8940        let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
8941            self.gpu.stream().synchronize()?;
8942            let t0 = std::time::Instant::now();
8943            for _ in 0..500 {
8944                if rp {
8945                    let __s_b = self.gpu.stream();
8946                    let mut b = __s_b.launch_builder(&fr);
8947                    b.arg(&wrp_d)
8948                        .arg(&aq_d)
8949                        .arg(&ad_d)
8950                        .arg(&mut y1)
8951                        .arg(&inf)
8952                        .arg(&outf)
8953                        .arg(&mi)
8954                        .arg(&qp);
8955                    unsafe {
8956                        b.launch(cfg)?;
8957                    }
8958                } else {
8959                    let __s_b = self.gpu.stream();
8960                    let mut b = __s_b.launch_builder(&fb);
8961                    b.arg(&w_d)
8962                        .arg(&aq_d)
8963                        .arg(&ad_d)
8964                        .arg(&mut y0)
8965                        .arg(&inf)
8966                        .arg(&outf)
8967                        .arg(&mi)
8968                        .arg(&rb);
8969                    unsafe {
8970                        b.launch(cfg)?;
8971                    }
8972                }
8973            }
8974            self.gpu.stream().synchronize()?;
8975            Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
8976        };
8977        let _ = time(false)?;
8978        let _ = time(true)?; // warm
8979        Ok((time(false)?, time(true)?))
8980    }
8981
8982    /// Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation,
8983    /// q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode
8984    /// dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent).
8985    /// VRAM cost == the tensor's weight size. MEMRA_Q4RP=0 disables at the call sites.
8986    pub fn build_q4_rp4(
8987        &self,
8988        t: &mut crate::model::GpuTensor,
8989    ) -> Result<(), Box<dyn std::error::Error>> {
8990        use crate::model::GpuTensor;
8991        let GpuTensor::Quant {
8992            bytes,
8993            qtype,
8994            row_bytes,
8995            ne,
8996            rp4,
8997            ..
8998        } = t
8999        else {
9000            return Ok(());
9001        };
9002        if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
9003            return Ok(());
9004        }
9005        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9006        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
9007            return Ok(());
9008        }
9009        let nblk = in_f / 32;
9010        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
9011        let f = self.func("q4_0_split_rp_build");
9012        let n = (out_f * nblk) as i32;
9013        let cfg = LaunchConfig {
9014            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
9015            block_dim: (256, 1, 1),
9016            shared_mem_bytes: 0,
9017        };
9018        let (of, nb) = (out_f as i32, nblk as i32);
9019        let _ = n;
9020        let __s_b = self.gpu.stream();
9021        let mut b = __s_b.launch_builder(&f);
9022        b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
9023        unsafe {
9024            b.launch(cfg)?;
9025        }
9026        *rp4 = Some(dst);
9027        Ok(())
9028    }
9029
9030    /// Q8_0 twin of `build_q4_rp4` (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride
9031    /// weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B
9032    /// ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout);
9033    /// the mmvq/batched decode arms prefer the mirror via `rp4`. Bit-identical outputs.
9034    pub fn build_q8_rp4(
9035        &self,
9036        t: &mut crate::model::GpuTensor,
9037    ) -> Result<(), Box<dyn std::error::Error>> {
9038        use crate::model::GpuTensor;
9039        let GpuTensor::Quant {
9040            bytes,
9041            qtype,
9042            row_bytes,
9043            ne,
9044            rp4,
9045            ..
9046        } = t
9047        else {
9048            return Ok(());
9049        };
9050        if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
9051            return Ok(());
9052        }
9053        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9054        if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
9055            return Ok(());
9056        }
9057        *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
9058        Ok(())
9059    }
9060
9061    /// Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane
9062    /// mirror without a GpuTensor (same kernel the loader path above uses).
9063    pub fn build_q8_rp4_raw(
9064        &self,
9065        bytes: &CudaSlice<u8>,
9066        in_f: usize,
9067        out_f: usize,
9068    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9069        assert!(in_f.is_multiple_of(32));
9070        let nblk = in_f / 32;
9071        let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
9072        let f = self.func("q8_0_split_rp_build");
9073        let cfg = LaunchConfig {
9074            grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
9075            block_dim: (256, 1, 1),
9076            shared_mem_bytes: 0,
9077        };
9078        let (of, nb) = (out_f as i32, nblk as i32);
9079        let __s_b = self.gpu.stream();
9080        let mut b = __s_b.launch_builder(&f);
9081        b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
9082        unsafe {
9083            b.launch(cfg)?;
9084        }
9085        Ok(dst)
9086    }
9087
9088    /// K-quant twins of `build_q8_rp4` (H100 K-quant coalescing fix, 2026-08-01 ncu on the
9089    /// q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with
9090    /// 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The
9091    /// mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++
9092    /// scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw
9093    /// bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched
9094    /// decode arms prefer the mirror via `rp4`. Bit-identical outputs.
9095    pub fn build_q4k_rp4(
9096        &self,
9097        t: &mut crate::model::GpuTensor,
9098    ) -> Result<(), Box<dyn std::error::Error>> {
9099        use crate::model::GpuTensor;
9100        let GpuTensor::Quant {
9101            bytes,
9102            qtype,
9103            row_bytes,
9104            ne,
9105            rp4,
9106            ..
9107        } = t
9108        else {
9109            return Ok(());
9110        };
9111        if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
9112            return Ok(());
9113        }
9114        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9115        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
9116            return Ok(());
9117        }
9118        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
9119        Ok(())
9120    }
9121
9122    pub fn build_q6k_rp4(
9123        &self,
9124        t: &mut crate::model::GpuTensor,
9125    ) -> Result<(), Box<dyn std::error::Error>> {
9126        use crate::model::GpuTensor;
9127        let GpuTensor::Quant {
9128            bytes,
9129            qtype,
9130            row_bytes,
9131            ne,
9132            rp4,
9133            ..
9134        } = t
9135        else {
9136            return Ok(());
9137        };
9138        if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
9139            return Ok(());
9140        }
9141        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9142        if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
9143            return Ok(());
9144        }
9145        *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
9146        Ok(())
9147    }
9148
9149    /// Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).
9150    pub fn build_kq_rp4_raw(
9151        &self,
9152        bytes: &CudaSlice<u8>,
9153        in_f: usize,
9154        out_f: usize,
9155        qtype: i32,
9156    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9157        assert!(in_f.is_multiple_of(256));
9158        let nsbk = in_f / 256;
9159        let (sb_bytes, kname) = match qtype {
9160            QT_Q4_K => (144usize, "q4_K_split_rp_build"),
9161            QT_Q6_K => (210usize, "q6_K_split_rp_build"),
9162            _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
9163        };
9164        let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
9165        let f = self.func(kname);
9166        let cfg = LaunchConfig {
9167            grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
9168            block_dim: (256, 1, 1),
9169            shared_mem_bytes: 0,
9170        };
9171        let (of, nb) = (out_f as i32, nsbk as i32);
9172        let __s_b = self.gpu.stream();
9173        let mut b = __s_b.launch_builder(&f);
9174        b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
9175        unsafe {
9176            b.launch(cfg)?;
9177        }
9178        Ok(dst)
9179    }
9180
9181    /// MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load.
9182    /// Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror
9183    /// VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).
9184    pub fn kqrp_enabled() -> bool {
9185        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9186        *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
9187            Ok("0") => false,
9188            Ok(_) => true,
9189            Err(_) => cfg!(memra_hopper_mma),
9190        })
9191    }
9192
9193    /// IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the
9194    /// GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor's size).
9195    /// The tensor's `rp` flag then routes every consumer (mmvq/batched `_rp` twins, the
9196    /// `qmatvec_gemm_q4_0_rp` prefill kernel). Callers gate on the fast path being active —
9197    /// the Stage-A f32 oracle (`MEMRA_FAST=0`) reads GGUF layout and must never see a swap.
9198    pub fn build_q4_rp_swap(
9199        &self,
9200        t: &mut crate::model::GpuTensor,
9201    ) -> Result<bool, Box<dyn std::error::Error>> {
9202        use crate::model::GpuTensor;
9203        // QTYPE GUARD IN THE SWAP ITSELF (gemma4 NVFP4mix prefill-NaN, 2026-08-17):
9204        // `rp4` is a SHARED field — the Q8RP walk parks Q8_0 MIRRORS there, and this
9205        // fn used to `take()` whatever it found. On a Q8_0-carrying gemma4-dense trunk
9206        // the swap hijacked those mirrors: `bytes` became split-plane in place, the
9207        // m<=16 `_rp` dispatch read them correctly (masking the corruption from every
9208        // decode pin), and every GGUF-layout prefill consumer (MMQ + GEMM) read the
9209        // fp16 d-plane as weights -> layer-0 NaN, <pad>-spam serving. Only a tensor
9210        // this fn's OWN builder serves may ever be swapped; everything else refuses
9211        // here, regardless of walk ordering.
9212        if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
9213            return Ok(false);
9214        }
9215        self.build_q4_rp4(t)?;
9216        self.gpu.stream().synchronize()?; // build kernel reads the GGUF bytes — drain BEFORE dropping them
9217        let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
9218            return Ok(false);
9219        };
9220        match rp4.take() {
9221            Some(split) => {
9222                *bytes = split; // the GGUF-layout buffer drops here
9223                *rp = true;
9224                Ok(true)
9225            }
9226            None => Ok(false),
9227        }
9228    }
9229
9230    /// MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
9231    pub fn q4rp_enabled() -> bool {
9232        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9233        *ON.get_or_init(|| {
9234            std::env::var("MEMRA_Q4RP")
9235                .map(|v| v != "0")
9236                .unwrap_or(true)
9237        })
9238    }
9239
9240    /// gemma4-E4B: dense [t][row_elems] gather of layer il's rows from the strided prologue
9241    /// buffer ([t][n_layer][n_epl]; off = il*n_epl, stride = n_layer*n_epl).
9242    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
9243    pub fn copy_rows_strided(
9244        &self,
9245        src: &CudaSlice<f32>,
9246        dst: &mut CudaSlice<f32>,
9247        row_elems: usize,
9248        n_rows: usize,
9249        src_stride: usize,
9250        src_off: usize,
9251    ) -> Result<(), Box<dyn std::error::Error>> {
9252        let f = self.func("copy_rows_strided_f32");
9253        let cfg = LaunchConfig {
9254            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
9255            block_dim: (256, 1, 1),
9256            shared_mem_bytes: 0,
9257        };
9258        let (re, nr) = (row_elems as i32, n_rows as i32);
9259        let (st, off) = (src_stride as i64, src_off as i64);
9260        let __s_b = self.gpu.stream();
9261        let mut b = __s_b.launch_builder(&f);
9262        b.arg(src)
9263            .arg(&mut *dst)
9264            .arg(&re)
9265            .arg(&nr)
9266            .arg(&st)
9267            .arg(&off);
9268        unsafe {
9269            b.launch(cfg)?;
9270        }
9271        Ok(())
9272    }
9273
9274    /// Place dense `[row][row_elems]` source rows into one column range of a strided destination.
9275    ///
9276    /// This is a byte-preserving layout operation. It exists so multi-GPU collectives can move
9277    /// one dense shard per rank and reconstruct the canonical token-major matrix without issuing
9278    /// one peer copy per token.
9279    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
9280    pub fn place_rows_strided(
9281        &self,
9282        src: &CudaSlice<f32>,
9283        dst: &mut CudaSlice<f32>,
9284        row_elems: usize,
9285        n_rows: usize,
9286        dst_stride: usize,
9287        dst_off: usize,
9288    ) -> Result<(), Box<dyn std::error::Error>> {
9289        if row_elems == 0 || n_rows == 0 {
9290            return Err("strided row placement requires nonzero rows and row width".into());
9291        }
9292        let src_len = n_rows
9293            .checked_mul(row_elems)
9294            .ok_or("strided row placement source size overflow")?;
9295        let dst_len = n_rows
9296            .checked_sub(1)
9297            .and_then(|rows| rows.checked_mul(dst_stride))
9298            .and_then(|base| base.checked_add(dst_off))
9299            .and_then(|base| base.checked_add(row_elems))
9300            .ok_or("strided row placement destination size overflow")?;
9301        let row_end = dst_off
9302            .checked_add(row_elems)
9303            .ok_or("strided row placement row size overflow")?;
9304        if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
9305            return Err(format!(
9306                "strided row placement geometry mismatch: src={} need_src={src_len} \
9307                 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
9308                 dst_stride={dst_stride} dst_off={dst_off}",
9309                src.len(),
9310                dst.len(),
9311            )
9312            .into());
9313        }
9314        if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
9315            return Err("strided row placement exceeds CUDA kernel geometry".into());
9316        }
9317        let f = self.func("place_rows_strided_f32");
9318        let cfg = LaunchConfig {
9319            grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
9320            block_dim: (256, 1, 1),
9321            shared_mem_bytes: 0,
9322        };
9323        let (re, nr) = (row_elems as i32, n_rows as i32);
9324        let (st, off) = (dst_stride as i64, dst_off as i64);
9325        let __s_b = self.gpu.stream();
9326        let mut b = __s_b.launch_builder(&f);
9327        b.arg(src)
9328            .arg(&mut *dst)
9329            .arg(&re)
9330            .arg(&nr)
9331            .arg(&st)
9332            .arg(&off);
9333        unsafe {
9334            b.launch(cfg)?;
9335        }
9336        Ok(())
9337    }
9338
9339    /// Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).
9340    pub fn u32_set_k(
9341        &self,
9342        dst: &mut CudaSlice<u32>,
9343        v: u32,
9344        idx: usize,
9345    ) -> Result<(), Box<dyn std::error::Error>> {
9346        let f = self.func("u32_set_k");
9347        let cfg = LaunchConfig {
9348            grid_dim: (1, 1, 1),
9349            block_dim: (1, 1, 1),
9350            shared_mem_bytes: 0,
9351        };
9352        let ii = idx as i32;
9353        let __s_b = self.gpu.stream();
9354        let mut b = __s_b.launch_builder(&f);
9355        b.arg(dst).arg(&v).arg(&ii);
9356        unsafe {
9357            b.launch(cfg)?;
9358        }
9359        Ok(())
9360    }
9361
9362    /// counter += v (device-slot append advance; the +1 twin is `inc_seqlen`).
9363    pub fn i32_add_k(
9364        &self,
9365        d: &mut CudaSlice<i32>,
9366        v: i32,
9367    ) -> Result<(), Box<dyn std::error::Error>> {
9368        let f = self.func("i32_add_k");
9369        let cfg = LaunchConfig {
9370            grid_dim: (1, 1, 1),
9371            block_dim: (32, 1, 1),
9372            shared_mem_bytes: 0,
9373        };
9374        let __s_b = self.gpu.stream();
9375        let mut b = __s_b.launch_builder(&f);
9376        b.arg(d).arg(&v);
9377        unsafe {
9378            b.launch(cfg)?;
9379        }
9380        Ok(())
9381    }
9382
9383    /// pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).
9384    pub fn i32_iota_from(
9385        &self,
9386        ctr: &CudaSlice<i32>,
9387        dst: &mut CudaSlice<i32>,
9388        n: usize,
9389    ) -> Result<(), Box<dyn std::error::Error>> {
9390        let f = self.func("i32_iota_from");
9391        let cfg = LaunchConfig::for_num_elems(n as u32);
9392        let ni = n as i32;
9393        let __s_b = self.gpu.stream();
9394        let mut b = __s_b.launch_builder(&f);
9395        b.arg(ctr).arg(dst).arg(&ni);
9396        unsafe {
9397            b.launch(cfg)?;
9398        }
9399        Ok(())
9400    }
9401
9402    /// In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).
9403    pub fn u32_map_k(
9404        &self,
9405        buf: &mut CudaSlice<u32>,
9406        map: &CudaSlice<u32>,
9407        idx: usize,
9408    ) -> Result<(), Box<dyn std::error::Error>> {
9409        let f = self.func("u32_map_k");
9410        let cfg = LaunchConfig {
9411            grid_dim: (1, 1, 1),
9412            block_dim: (1, 1, 1),
9413            shared_mem_bytes: 0,
9414        };
9415        let ii = idx as i32;
9416        let __s_b = self.gpu.stream();
9417        let mut b = __s_b.launch_builder(&f);
9418        b.arg(buf).arg(map).arg(&ii);
9419        unsafe {
9420            b.launch(cfg)?;
9421        }
9422        Ok(())
9423    }
9424
9425    /// Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).
9426    #[allow(clippy::too_many_arguments)]
9427    pub fn u32_pack2(
9428        &self,
9429        a: &CudaSlice<u32>,
9430        off_a: usize,
9431        n1: usize,
9432        b_in: &CudaSlice<u32>,
9433        n2: usize,
9434        out: &mut CudaSlice<u32>,
9435    ) -> Result<(), Box<dyn std::error::Error>> {
9436        let f = self.func("u32_pack2");
9437        let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
9438        let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
9439        let __s_b = self.gpu.stream();
9440        let mut b = __s_b.launch_builder(&f);
9441        b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
9442        unsafe {
9443            b.launch(cfg)?;
9444        }
9445        Ok(())
9446    }
9447
9448    /// gemma4 R3 device fold: w[i] *= s[sel[i]] over the router's [n] (sel, w) pair.
9449    pub fn moe_w_exscale(
9450        &self,
9451        w: &mut CudaSlice<f32>,
9452        sel: &CudaSlice<i32>,
9453        s: &CudaSlice<f32>,
9454        n: usize,
9455    ) -> Result<(), Box<dyn std::error::Error>> {
9456        let f = self.func("moe_w_exscale");
9457        let cfg = LaunchConfig::for_num_elems(n as u32);
9458        let ni = n as i32;
9459        let __s_b = self.gpu.stream();
9460        let mut b = __s_b.launch_builder(&f);
9461        b.arg(w).arg(sel).arg(s).arg(&ni);
9462        unsafe {
9463            b.launch(cfg)?;
9464        }
9465        Ok(())
9466    }
9467
9468    /// Down-projection macro fold: w[i] *= macros[2*n_expert + sel[i]] on the device router
9469    /// weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).
9470    pub fn moe_w_scale_by_expert(
9471        &self,
9472        w: &mut CudaSlice<f32>,
9473        sel: &CudaSlice<i32>,
9474        macros: &CudaSlice<f32>,
9475        n_expert: usize,
9476        n: usize,
9477    ) -> Result<(), Box<dyn std::error::Error>> {
9478        let f = self.func("moe_w_scale_by_expert");
9479        let cfg = LaunchConfig {
9480            grid_dim: (n.div_ceil(64) as u32, 1, 1),
9481            block_dim: (64, 1, 1),
9482            shared_mem_bytes: 0,
9483        };
9484        let (ne, nn) = (n_expert as i32, n as i32);
9485        let __s_b = self.gpu.stream();
9486        let mut b = __s_b.launch_builder(&f);
9487        b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
9488        unsafe {
9489            b.launch(cfg)?;
9490        }
9491        Ok(())
9492    }
9493
9494    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9495    pub fn moe_gate_up_silu8_dev_q8(
9496        &self,
9497        table: &CudaSlice<u64>,
9498        sel: &cudarc::driver::CudaView<i32>,
9499        aq: &CudaSlice<i8>,
9500        ad: &CudaSlice<f32>,
9501        in_f: usize,
9502        n_ff: usize,
9503        n_used: usize,
9504        n_expert: usize,
9505        qt_g: i32,
9506        qt_u: i32,
9507        rb_g: usize,
9508        rb_u: usize,
9509        macros: &CudaSlice<f32>,
9510    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9511        static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
9512        let (mode, wpb) = GU.get_or_init(|| {
9513            let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
9514            let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
9515                .ok()
9516                .and_then(|v| v.parse().ok())
9517                .unwrap_or(4u32)
9518                .clamp(1, 16);
9519            (mode, wpb)
9520        });
9521        let (mode, wpb) = (mode.as_str(), *wpb);
9522        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
9523        let (inf, nff, ne, rbg, rbu) = (
9524            in_f as i32,
9525            n_ff as i32,
9526            n_expert as i32,
9527            rb_g as i64,
9528            rb_u as i64,
9529        );
9530        let (f, cfg) = match mode {
9531            "1" | "2" | "4" => {
9532                let rpw: u32 = mode.parse().unwrap();
9533                let f = self.func(match rpw {
9534                    1 => "moe_gate_up_silu8_dev_q8_r1",
9535                    2 => "moe_gate_up_silu8_dev_q8_r2",
9536                    _ => "moe_gate_up_silu8_dev_q8_r4",
9537                });
9538                let rows_per_block = (rpw * wpb) as usize;
9539                let gx = n_ff.div_ceil(rows_per_block) as u32;
9540                (
9541                    f,
9542                    LaunchConfig {
9543                        grid_dim: (gx, n_used as u32, 1),
9544                        block_dim: (32, wpb, 1),
9545                        shared_mem_bytes: 0,
9546                    },
9547                )
9548            }
9549            "j8" if n_used <= 32 => (
9550                self.func("moe_gate_up_silu8_dev_q8_j8"),
9551                LaunchConfig {
9552                    grid_dim: (n_ff as u32, 1, 1),
9553                    block_dim: (32, n_used as u32, 1),
9554                    shared_mem_bytes: 0,
9555                },
9556            ),
9557            // SMEM-GRID twins (IQ3_S 2KB grid copied to shared, static smem — bit-identical dots)
9558            "vsm2" => {
9559                let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
9560                let sh = (rb_g + rb_u) as u32;
9561                use cudarc::driver::sys::CUfunction_attribute_enum as A;
9562                f.set_attribute(
9563                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
9564                    sh as i32,
9565                )?;
9566                (
9567                    f,
9568                    LaunchConfig {
9569                        grid_dim: (n_ff as u32, n_used as u32, 1),
9570                        block_dim: (32, 1, 1),
9571                        shared_mem_bytes: sh,
9572                    },
9573                )
9574            }
9575            "vsm" => {
9576                let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
9577                let sh = (rb_g + rb_u) as u32;
9578                use cudarc::driver::sys::CUfunction_attribute_enum as A;
9579                f.set_attribute(
9580                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
9581                    sh as i32,
9582                )?;
9583                (
9584                    f,
9585                    LaunchConfig {
9586                        grid_dim: (n_ff as u32, n_used as u32, 1),
9587                        block_dim: (32, 1, 1),
9588                        shared_mem_bytes: sh,
9589                    },
9590                )
9591            }
9592            "sg" => (
9593                self.func("moe_gate_up_silu8_dev_q8_sg"),
9594                LaunchConfig {
9595                    grid_dim: (n_ff as u32, n_used as u32, 1),
9596                    block_dim: (32, 1, 1),
9597                    shared_mem_bytes: 0,
9598                },
9599            ),
9600            "j8sg" if n_used <= 32 => (
9601                self.func("moe_gate_up_silu8_dev_q8_j8sg"),
9602                LaunchConfig {
9603                    grid_dim: (n_ff as u32, 1, 1),
9604                    block_dim: (32, n_used as u32, 1),
9605                    shared_mem_bytes: 0,
9606                },
9607            ),
9608            "u64" if in_f == 2048 => (
9609                self.func("moe_gate_up_silu8_dev_q8_u64"),
9610                LaunchConfig {
9611                    grid_dim: (n_ff as u32, n_used as u32, 1),
9612                    block_dim: (32, 1, 1),
9613                    shared_mem_bytes: 0,
9614                },
9615            ),
9616            "gs4" if in_f == 2048 => (
9617                self.func("moe_gate_up_silu8_dev_q8_gs4"),
9618                LaunchConfig {
9619                    grid_dim: (n_ff as u32, n_used as u32, 1),
9620                    block_dim: (32, 4, 1),
9621                    shared_mem_bytes: 0,
9622                },
9623            ),
9624            // _v twin (down8 lane 2026-07-08): wide-load IQ4_XS dot, base geometry, bit-identical.
9625            "v" | "" => (
9626                self.func("moe_gate_up_silu8_dev_q8_v"),
9627                LaunchConfig {
9628                    grid_dim: (n_ff as u32, n_used as u32, 1),
9629                    block_dim: (32, 1, 1),
9630                    shared_mem_bytes: 0,
9631                },
9632            ),
9633            "s2" => (
9634                self.func("moe_gate_up_silu8_dev_q8_s2"),
9635                LaunchConfig {
9636                    grid_dim: (n_ff as u32, n_used as u32, 1),
9637                    block_dim: (32, 2, 1),
9638                    shared_mem_bytes: 0,
9639                },
9640            ),
9641            "s2z" => {
9642                let rz = wpb.min(16); // s2z smem tile is [16][2]
9643                (
9644                    self.func("moe_gate_up_silu8_dev_q8_s2z"),
9645                    LaunchConfig {
9646                        grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
9647                        block_dim: (32, 2, rz),
9648                        shared_mem_bytes: 0,
9649                    },
9650                )
9651            }
9652            _ => (
9653                self.func("moe_gate_up_silu8_dev_q8"),
9654                LaunchConfig {
9655                    grid_dim: (n_ff as u32, n_used as u32, 1),
9656                    block_dim: (32, 1, 1),
9657                    shared_mem_bytes: 0,
9658                },
9659            ),
9660        };
9661        let __s_b = self.gpu.stream();
9662        let mut b = __s_b.launch_builder(&f);
9663        b.arg(table)
9664            .arg(sel)
9665            .arg(aq)
9666            .arg(ad)
9667            .arg(&mut act)
9668            .arg(&inf)
9669            .arg(&nff)
9670            .arg(&ne)
9671            .arg(&qt_g)
9672            .arg(&qt_u)
9673            .arg(&rbg)
9674            .arg(&rbu)
9675            .arg(macros);
9676        unsafe {
9677            b.launch(cfg)?;
9678        }
9679        Ok(act)
9680    }
9681
9682    #[allow(clippy::too_many_arguments)]
9683    pub fn moe_down8_fma_dev_q8(
9684        &self,
9685        table: &CudaSlice<u64>,
9686        sel: &cudarc::driver::CudaView<i32>,
9687        w: &cudarc::driver::CudaView<f32>,
9688        aq2: &CudaSlice<i8>,
9689        ad2: &CudaSlice<f32>,
9690        dst: &mut cudarc::driver::CudaViewMut<f32>,
9691        in_f: usize,
9692        out_f: usize,
9693        n_used: usize,
9694        n_expert: usize,
9695        qt: i32,
9696        rb: usize,
9697    ) -> Result<(), Box<dyn std::error::Error>> {
9698        static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
9699        let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
9700        let (inf, outf, nu, ne, rbi) = (
9701            in_f as i32,
9702            out_f as i32,
9703            n_used as i32,
9704            n_expert as i32,
9705            rb as i64,
9706        );
9707        // the w8 twins' smem tile is [RPW][8] — n_used must fit the 8-slot tile;
9708        // the h2 twins are nsb==16 (in_f==512) shape-gated.
9709        let (f, cfg) = match mode.as_str() {
9710            m @ ("1" | "2" | "4") if n_used <= 8 => {
9711                let rpw: usize = m.parse().unwrap();
9712                let f = self.func(match rpw {
9713                    1 => "moe_down8_fma_dev_q8_w8r1",
9714                    2 => "moe_down8_fma_dev_q8_w8r2",
9715                    _ => "moe_down8_fma_dev_q8_w8r4",
9716                });
9717                (
9718                    f,
9719                    LaunchConfig {
9720                        grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
9721                        block_dim: (32, n_used as u32, 1),
9722                        shared_mem_bytes: 0,
9723                    },
9724                )
9725            }
9726            "h2" if in_f == 512 => (
9727                self.func("moe_down8_fma_dev_q8_h2"),
9728                LaunchConfig {
9729                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9730                    block_dim: (32, 1, 1),
9731                    shared_mem_bytes: 0,
9732                },
9733            ),
9734            // "" = AUTO gemma shape (in_f==704): w8r2 measured +1 tok/s vs base (sweep
9735            // 1/2/4 -> 133.6/134.2/133.6, 2026-07-10); slot-ordered chain preserved.
9736            "" if in_f == 704 && n_used <= 8 => (
9737                self.func("moe_down8_fma_dev_q8_w8r2"),
9738                LaunchConfig {
9739                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9740                    block_dim: (32, n_used as u32, 1),
9741                    shared_mem_bytes: 0,
9742                },
9743            ),
9744            // "" = AUTO: the measured winner for the 35B expert shape (arc 2026-07-05, +3.8%);
9745            // any shape the h2 kernels can't take (nsb!=16 / n_used>8) falls to base via `_`.
9746            // _v twins (down8 lane 2026-07-08): wide-load IQ4_XS dot, bit-identical outputs.
9747            "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
9748                self.func("moe_down8_fma_dev_q8_w8h2v"),
9749                LaunchConfig {
9750                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9751                    block_dim: (32, n_used as u32, 1),
9752                    shared_mem_bytes: 0,
9753                },
9754            ),
9755            "w8h2r2v" if in_f == 512 && n_used <= 8 => (
9756                self.func("moe_down8_fma_dev_q8_w8h2r2v"),
9757                LaunchConfig {
9758                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
9759                    block_dim: (32, n_used as u32, 1),
9760                    shared_mem_bytes: 0,
9761                },
9762            ),
9763            "w8h2r2" if in_f == 512 && n_used <= 8 => (
9764                self.func("moe_down8_fma_dev_q8_w8h2r2"),
9765                LaunchConfig {
9766                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
9767                    block_dim: (32, n_used as u32, 1),
9768                    shared_mem_bytes: 0,
9769                },
9770            ),
9771            "w8h2" if in_f == 512 && n_used <= 8 => (
9772                self.func("moe_down8_fma_dev_q8_w8h2"),
9773                LaunchConfig {
9774                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9775                    block_dim: (32, n_used as u32, 1),
9776                    shared_mem_bytes: 0,
9777                },
9778            ),
9779            _ => (
9780                self.func("moe_down8_fma_dev_q8"),
9781                LaunchConfig {
9782                    grid_dim: (out_f as u32, 1, 1),
9783                    block_dim: (32, 1, 1),
9784                    shared_mem_bytes: 0,
9785                },
9786            ),
9787        };
9788        let __s_b = self.gpu.stream();
9789        let mut b = __s_b.launch_builder(&f);
9790        b.arg(table)
9791            .arg(sel)
9792            .arg(w)
9793            .arg(aq2)
9794            .arg(ad2)
9795            .arg(dst)
9796            .arg(&inf)
9797            .arg(&outf)
9798            .arg(&nu)
9799            .arg(&ne)
9800            .arg(&qt)
9801            .arg(&rbi);
9802        unsafe {
9803            b.launch(cfg)?;
9804        }
9805        Ok(())
9806    }
9807
9808    /// SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all `t` tokens
9809    /// of the spec verify's MoE dev gate/up (grid.z = token) — the _v geometry per token, with
9810    /// tok-offset sel/aq/ad/act pointers matching the serial loop's slices. BIT-IDENTICAL per
9811    /// token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows —
9812    /// quantize_q8_1's per-32-block program is row-independent, so batched rows == the serial
9813    /// loop's per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].
9814    #[allow(clippy::too_many_arguments)]
9815    pub fn moe_gate_up_silu8_dev_q8_rows(
9816        &self,
9817        table: &CudaSlice<u64>,
9818        sel: &CudaSlice<i32>,
9819        aq: &CudaSlice<i8>,
9820        ad: &CudaSlice<f32>,
9821        t: usize,
9822        in_f: usize,
9823        n_ff: usize,
9824        n_used: usize,
9825        n_expert: usize,
9826        qt_g: i32,
9827        qt_u: i32,
9828        rb_g: usize,
9829        rb_u: usize,
9830        macros: &CudaSlice<f32>,
9831    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9832        let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
9833        let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
9834        let cfg = LaunchConfig {
9835            grid_dim: (n_ff as u32, n_used as u32, t as u32),
9836            block_dim: (32, 1, 1),
9837            shared_mem_bytes: 0,
9838        };
9839        let (inf, nff, ne, nu, rbg, rbu) = (
9840            in_f as i32,
9841            n_ff as i32,
9842            n_expert as i32,
9843            n_used as i32,
9844            rb_g as i64,
9845            rb_u as i64,
9846        );
9847        let __s_b = self.gpu.stream();
9848        let mut b = __s_b.launch_builder(&f);
9849        b.arg(table)
9850            .arg(sel)
9851            .arg(aq)
9852            .arg(ad)
9853            .arg(&mut act)
9854            .arg(&inf)
9855            .arg(&nff)
9856            .arg(&ne)
9857            .arg(&qt_g)
9858            .arg(&qt_u)
9859            .arg(&rbg)
9860            .arg(&rbu)
9861            .arg(&nu)
9862            .arg(macros);
9863        unsafe {
9864            b.launch(cfg)?;
9865        }
9866        Ok(act)
9867    }
9868
9869    /// SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token
9870    /// axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as
9871    /// the AUTO dispatch in `moe_down8_fma_dev_q8`. aq2/ad2 = batched act quantize
9872    /// ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.
9873    #[allow(clippy::too_many_arguments)]
9874    pub fn moe_down8_fma_dev_q8_rows(
9875        &self,
9876        table: &CudaSlice<u64>,
9877        sel: &CudaSlice<i32>,
9878        w: &CudaSlice<f32>,
9879        aq2: &CudaSlice<i8>,
9880        ad2: &CudaSlice<f32>,
9881        dst: &mut CudaSlice<f32>,
9882        t: usize,
9883        in_f: usize,
9884        out_f: usize,
9885        n_used: usize,
9886        n_expert: usize,
9887        qt: i32,
9888        rb: usize,
9889    ) -> Result<(), Box<dyn std::error::Error>> {
9890        assert!(
9891            in_f == 512 && n_used <= 8,
9892            "down rows twin is w8h2v shape-gated"
9893        );
9894        let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
9895        let cfg = LaunchConfig {
9896            grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
9897            block_dim: (32, n_used as u32, 1),
9898            shared_mem_bytes: 0,
9899        };
9900        let (inf, outf, nu, ne, rbi) = (
9901            in_f as i32,
9902            out_f as i32,
9903            n_used as i32,
9904            n_expert as i32,
9905            rb as i64,
9906        );
9907        let __s_b = self.gpu.stream();
9908        let mut b = __s_b.launch_builder(&f);
9909        b.arg(table)
9910            .arg(sel)
9911            .arg(w)
9912            .arg(aq2)
9913            .arg(ad2)
9914            .arg(dst)
9915            .arg(&inf)
9916            .arg(&outf)
9917            .arg(&nu)
9918            .arg(&ne)
9919            .arg(&qt)
9920            .arg(&rbi);
9921        unsafe {
9922            b.launch(cfg)?;
9923        }
9924        Ok(())
9925    }
9926
9927    /// CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller
9928    /// gates), grid.y = pair index; the first pair of each expert serves all its pairs.
9929    /// Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).
9930    #[allow(clippy::too_many_arguments)]
9931    pub fn moe_gate_up_silu8_dev_q8_csr(
9932        &self,
9933        table: &CudaSlice<u64>,
9934        sel: &CudaSlice<i32>,
9935        aq: &CudaSlice<i8>,
9936        ad: &CudaSlice<f32>,
9937        n_pairs: usize,
9938        in_f: usize,
9939        n_ff: usize,
9940        n_used: usize,
9941        n_expert: usize,
9942        qt_g: i32,
9943        qt_u: i32,
9944        rb_g: usize,
9945        rb_u: usize,
9946    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9947        // NVFP4 experts take the NVFP4-specialized owner-scan twin (lane/moebatch-q35moe);
9948        // host gate guarantees qt_g == qt_u within a supported class.
9949        let f = if qt_g == crate::QT_NVFP4 {
9950            self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
9951        } else {
9952            self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
9953        };
9954        let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
9955        let cfg = LaunchConfig {
9956            grid_dim: (n_ff as u32, n_pairs as u32, 1),
9957            block_dim: (32, 1, 1),
9958            shared_mem_bytes: 0,
9959        };
9960        let (inf, nff, ne, nu, npi, rbg, rbu) = (
9961            in_f as i32,
9962            n_ff as i32,
9963            n_expert as i32,
9964            n_used as i32,
9965            n_pairs as i32,
9966            rb_g as i64,
9967            rb_u as i64,
9968        );
9969        let __s_b = self.gpu.stream();
9970        let mut b = __s_b.launch_builder(&f);
9971        b.arg(table)
9972            .arg(sel)
9973            .arg(aq)
9974            .arg(ad)
9975            .arg(&mut act)
9976            .arg(&inf)
9977            .arg(&nff)
9978            .arg(&ne)
9979            .arg(&qt_g)
9980            .arg(&qt_u)
9981            .arg(&rbg)
9982            .arg(&rbu)
9983            .arg(&nu)
9984            .arg(&npi);
9985        unsafe {
9986            b.launch(cfg)?;
9987        }
9988        Ok(act)
9989    }
9990
9991    /// TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its
9992    /// canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte-
9993    /// compare variants in one process. Variants: "base", "w8h2", "w8h2r2", "w8h2v", "w8h2r2v".
9994    #[allow(clippy::too_many_arguments)]
9995    pub fn moe_down8_fma_dev_q8_variant(
9996        &self,
9997        variant: &str,
9998        table: &CudaSlice<u64>,
9999        sel: &cudarc::driver::CudaView<i32>,
10000        w: &cudarc::driver::CudaView<f32>,
10001        aq2: &CudaSlice<i8>,
10002        ad2: &CudaSlice<f32>,
10003        dst: &mut cudarc::driver::CudaViewMut<f32>,
10004        in_f: usize,
10005        out_f: usize,
10006        n_used: usize,
10007        n_expert: usize,
10008        qt: i32,
10009        rb: usize,
10010    ) -> Result<(), Box<dyn std::error::Error>> {
10011        let (inf, outf, nu, ne, rbi) = (
10012            in_f as i32,
10013            out_f as i32,
10014            n_used as i32,
10015            n_expert as i32,
10016            rb as i64,
10017        );
10018        let (f, cfg) = match variant {
10019            "w8h2" | "w8h2v" => (
10020                self.func(if variant == "w8h2" {
10021                    "moe_down8_fma_dev_q8_w8h2"
10022                } else {
10023                    "moe_down8_fma_dev_q8_w8h2v"
10024                }),
10025                LaunchConfig {
10026                    grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
10027                    block_dim: (32, n_used as u32, 1),
10028                    shared_mem_bytes: 0,
10029                },
10030            ),
10031            "w8h2r2" | "w8h2r2v" => (
10032                self.func(if variant == "w8h2r2" {
10033                    "moe_down8_fma_dev_q8_w8h2r2"
10034                } else {
10035                    "moe_down8_fma_dev_q8_w8h2r2v"
10036                }),
10037                LaunchConfig {
10038                    grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
10039                    block_dim: (32, n_used as u32, 1),
10040                    shared_mem_bytes: 0,
10041                },
10042            ),
10043            _ => (
10044                self.func("moe_down8_fma_dev_q8"),
10045                LaunchConfig {
10046                    grid_dim: (out_f as u32, 1, 1),
10047                    block_dim: (32, 1, 1),
10048                    shared_mem_bytes: 0,
10049                },
10050            ),
10051        };
10052        let __s_b = self.gpu.stream();
10053        let mut b = __s_b.launch_builder(&f);
10054        b.arg(table)
10055            .arg(sel)
10056            .arg(w)
10057            .arg(aq2)
10058            .arg(ad2)
10059            .arg(dst)
10060            .arg(&inf)
10061            .arg(&outf)
10062            .arg(&nu)
10063            .arg(&ne)
10064            .arg(&qt)
10065            .arg(&rbi);
10066        unsafe {
10067            b.launch(cfg)?;
10068        }
10069        Ok(())
10070    }
10071
10072    /// TEST SEAM (down8 lane): gate_up twin of the above. Variants: "base", "v".
10073    #[allow(clippy::too_many_arguments)]
10074    pub fn moe_gate_up_silu8_dev_q8_variant(
10075        &self,
10076        variant: &str,
10077        table: &CudaSlice<u64>,
10078        sel: &cudarc::driver::CudaView<i32>,
10079        aq: &CudaSlice<i8>,
10080        ad: &CudaSlice<f32>,
10081        in_f: usize,
10082        n_ff: usize,
10083        n_used: usize,
10084        n_expert: usize,
10085        qt_g: i32,
10086        qt_u: i32,
10087        rb_g: usize,
10088        rb_u: usize,
10089    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10090        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
10091        let (inf, nff, ne, rbg, rbu) = (
10092            in_f as i32,
10093            n_ff as i32,
10094            n_expert as i32,
10095            rb_g as i64,
10096            rb_u as i64,
10097        );
10098        let f = self.func(if variant == "v" {
10099            "moe_gate_up_silu8_dev_q8_v"
10100        } else {
10101            "moe_gate_up_silu8_dev_q8"
10102        });
10103        let cfg = LaunchConfig {
10104            grid_dim: (n_ff as u32, n_used as u32, 1),
10105            block_dim: (32, 1, 1),
10106            shared_mem_bytes: 0,
10107        };
10108        let __s_b = self.gpu.stream();
10109        let mut b = __s_b.launch_builder(&f);
10110        b.arg(table)
10111            .arg(sel)
10112            .arg(aq)
10113            .arg(ad)
10114            .arg(&mut act)
10115            .arg(&inf)
10116            .arg(&nff)
10117            .arg(&ne)
10118            .arg(&qt_g)
10119            .arg(&qt_u)
10120            .arg(&rbg)
10121            .arg(&rbu);
10122        unsafe {
10123            b.launch(cfg)?;
10124        }
10125        Ok(act)
10126    }
10127
10128    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
10129    pub fn moe_gate_up_silu8_dev(
10130        &self,
10131        table: &CudaSlice<u64>,
10132        sel: &cudarc::driver::CudaView<i32>,
10133        x: &cudarc::driver::CudaView<f32>,
10134        in_f: usize,
10135        n_ff: usize,
10136        n_used: usize,
10137        n_expert: usize,
10138        qt_g: i32,
10139        qt_u: i32,
10140        rb_g: usize,
10141        rb_u: usize,
10142        macros: &CudaSlice<f32>,
10143    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10144        let f = self.func("moe_gate_up_silu8_dev");
10145        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; // fully overwritten
10146        let cfg = LaunchConfig {
10147            grid_dim: (n_ff as u32, n_used as u32, 1),
10148            block_dim: (256, 1, 1),
10149            shared_mem_bytes: 0,
10150        };
10151        let (inf, nff, ne, rbg, rbu) = (
10152            in_f as i32,
10153            n_ff as i32,
10154            n_expert as i32,
10155            rb_g as i64,
10156            rb_u as i64,
10157        );
10158        let __s_b = self.gpu.stream();
10159        let mut b = __s_b.launch_builder(&f);
10160        b.arg(table)
10161            .arg(sel)
10162            .arg(x)
10163            .arg(&mut act)
10164            .arg(&inf)
10165            .arg(&nff)
10166            .arg(&ne)
10167            .arg(&qt_g)
10168            .arg(&qt_u)
10169            .arg(&rbg)
10170            .arg(&rbu)
10171            .arg(macros);
10172        unsafe {
10173            b.launch(cfg)?;
10174        }
10175        Ok(act)
10176    }
10177
10178    /// LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of `moe_down8_fma_into` — expert ids AND
10179    /// renormalized weights read from the router kernel's device output. BIT-IDENTICAL chain.
10180    #[allow(clippy::too_many_arguments)]
10181    pub fn moe_down8_fma_dev(
10182        &self,
10183        table: &CudaSlice<u64>,
10184        sel: &cudarc::driver::CudaView<i32>,
10185        w: &cudarc::driver::CudaView<f32>,
10186        act: &CudaSlice<f32>,
10187        dst: &mut cudarc::driver::CudaViewMut<f32>,
10188        in_f: usize,
10189        out_f: usize,
10190        n_used: usize,
10191        n_expert: usize,
10192        qt: i32,
10193        rb: usize,
10194    ) -> Result<(), Box<dyn std::error::Error>> {
10195        let f = self.func("moe_down8_fma_dev");
10196        let cfg = LaunchConfig {
10197            grid_dim: (out_f as u32, 1, 1),
10198            block_dim: (256, 1, 1),
10199            shared_mem_bytes: 0,
10200        };
10201        let (inf, outf, nu, ne, rbv) = (
10202            in_f as i32,
10203            out_f as i32,
10204            n_used as i32,
10205            n_expert as i32,
10206            rb as i64,
10207        );
10208        let __s_b = self.gpu.stream();
10209        let mut b = __s_b.launch_builder(&f);
10210        b.arg(table)
10211            .arg(sel)
10212            .arg(w)
10213            .arg(act)
10214            .arg(dst)
10215            .arg(&inf)
10216            .arg(&outf)
10217            .arg(&nu)
10218            .arg(&ne)
10219            .arg(&qt)
10220            .arg(&rbv);
10221        unsafe {
10222            b.launch(cfg)?;
10223        }
10224        Ok(())
10225    }
10226
10227    /// dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).
10228    pub fn axpy_into(
10229        &self,
10230        src: &CudaSlice<f32>,
10231        alpha: f32,
10232        dst: &mut cudarc::driver::CudaViewMut<f32>,
10233        n: usize,
10234    ) -> Result<(), Box<dyn std::error::Error>> {
10235        let f = self.func("axpy_f32");
10236        let cfg = LaunchConfig::for_num_elems(n as u32);
10237        let (a, ni) = (alpha, n as i32);
10238        let __s_b = self.gpu.stream();
10239        let mut b = __s_b.launch_builder(&f);
10240        b.arg(src).arg(dst).arg(&a).arg(&ni);
10241        unsafe {
10242            b.launch(cfg)?;
10243        }
10244        Ok(())
10245    }
10246
10247    /// Host-oracle twin of `axpy_into` with separate RN multiply and add.
10248    pub fn axpy_host_into(
10249        &self,
10250        src: &cudarc::driver::CudaView<'_, f32>,
10251        alpha: f32,
10252        dst: &mut cudarc::driver::CudaViewMut<f32>,
10253        n: usize,
10254    ) -> Result<(), Box<dyn std::error::Error>> {
10255        let f = self.func("axpy_host_f32");
10256        let cfg = LaunchConfig::for_num_elems(n as u32);
10257        let (a, ni) = (alpha, n as i32);
10258        let __s_b = self.gpu.stream();
10259        let mut b = __s_b.launch_builder(&f);
10260        b.arg(src).arg(dst).arg(&a).arg(&ni);
10261        unsafe {
10262            b.launch(cfg)?;
10263        }
10264        Ok(())
10265    }
10266
10267    /// dst[r*ncols + c] += src[r*ncols + c] * scale[r]. Per-row scalar accumulate (shared expert).
10268    pub fn add_scaled_rows(
10269        &self,
10270        src: &CudaSlice<f32>,
10271        scale: &CudaSlice<f32>,
10272        dst: &mut CudaSlice<f32>,
10273        ncols: usize,
10274        nrows: usize,
10275    ) -> Result<(), Box<dyn std::error::Error>> {
10276        let f = self.func("add_scaled_rows_f32");
10277        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10278        let (nc, nr) = (ncols as i32, nrows as i32);
10279        let __s_b = self.gpu.stream();
10280        let mut b = __s_b.launch_builder(&f);
10281        b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
10282        unsafe {
10283            b.launch(cfg)?;
10284        }
10285        Ok(())
10286    }
10287
10288    /// `add_scaled_rows` with an all-ones scale drawn from the resident ones buffer (door H,
10289    /// `MEMRA_HTOD_DIET`) — the UNGATED shared-expert add, without re-uploading the
10290    /// constant every MoE layer-call. Same kernel, same values: the buffer may be longer than
10291    /// `nrows` because `add_scaled_rows_f32` reads only `scale[0..nrows]`.
10292    pub fn add_scaled_rows_ones(
10293        &self,
10294        src: &CudaSlice<f32>,
10295        dst: &mut CudaSlice<f32>,
10296        ncols: usize,
10297        nrows: usize,
10298    ) -> Result<(), Box<dyn std::error::Error>> {
10299        let mut guard = self
10300            .shexp_ones
10301            .lock()
10302            .map_err(|_| "shexp ones buffer is poisoned")?;
10303        if guard.as_ref().map(|b| b.len() < nrows).unwrap_or(true) {
10304            // One upload per process (or per growth step): the serving shapes are t <= 8 for the
10305            // verify walk and the prime's chunk width otherwise.
10306            *guard = Some(self.htod(&vec![1.0f32; nrows.max(64)])?);
10307        }
10308        let ones = guard.as_ref().expect("just ensured");
10309        let f = self.func("add_scaled_rows_f32");
10310        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10311        let (nc, nr) = (ncols as i32, nrows as i32);
10312        let __s_b = self.gpu.stream();
10313        let mut b = __s_b.launch_builder(&f);
10314        b.arg(src).arg(ones).arg(&mut *dst).arg(&nc).arg(&nr);
10315        unsafe {
10316            b.launch(cfg)?;
10317        }
10318        Ok(())
10319    }
10320
10321    /// The `len_d` i32 mirror store, door H aware (`MEMRA_HTOD_DIET`): the async
10322    /// [`Self::i32_set_k`] launch when the door is on, else the shipped synchronizing pageable
10323    /// `memcpy_htod`. Identical value into the identical slot, both stream-ordered.
10324    pub fn i32_mirror_store(
10325        &self,
10326        dst: &mut CudaSlice<i32>,
10327        v: i32,
10328    ) -> Result<(), Box<dyn std::error::Error>> {
10329        if crate::htod_diet_on() {
10330            HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10331            return self.i32_set_k(dst, v);
10332        }
10333        self.gpu.stream().memcpy_htod(&[v], dst)?;
10334        Ok(())
10335    }
10336
10337    /// y[r, :] *= s[r] in place (per-CSR-row macro scale for the grouped prime's gate/up —
10338    /// silu is nonlinear, so per-expert NVFP4 macros must land before it).
10339    pub fn scale_rows(
10340        &self,
10341        y: &mut CudaSlice<f32>,
10342        s: &CudaSlice<f32>,
10343        ncols: usize,
10344        nrows: usize,
10345    ) -> Result<(), Box<dyn std::error::Error>> {
10346        let f = self.func("scale_rows_f32");
10347        let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10348        let (nc, nr) = (ncols as i32, nrows as i32);
10349        let __s_b = self.gpu.stream();
10350        let mut b = __s_b.launch_builder(&f);
10351        b.arg(&mut *y).arg(s).arg(&nc).arg(&nr);
10352        unsafe {
10353            b.launch(cfg)?;
10354        }
10355        Ok(())
10356    }
10357
10358    /// Fused grouped-prime tail: join both rank partials (canonical shard order), permute
10359    /// CSR->pair via `inv`, weight, and scatter to tokens in one pass — replaces
10360    /// rows_permute + add + scatter and the three large temporaries they needed.
10361    #[allow(clippy::too_many_arguments)]
10362    pub fn moe_prime_join_scatter(
10363        &self,
10364        y0: &CudaSlice<f32>,
10365        y1: &CudaSlice<f32>,
10366        inv: &CudaSlice<i32>,
10367        w: &CudaSlice<f32>,
10368        out: &mut CudaSlice<f32>,
10369        ncols: usize,
10370        n_used: usize,
10371        t: usize,
10372    ) -> Result<(), Box<dyn std::error::Error>> {
10373        let f = self.func("moe_prime_join_scatter_f32");
10374        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10375        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10376        let __s_b = self.gpu.stream();
10377        let mut b = __s_b.launch_builder(&f);
10378        b.arg(y0)
10379            .arg(y1)
10380            .arg(inv)
10381            .arg(w)
10382            .arg(&mut *out)
10383            .arg(&nc)
10384            .arg(&nu)
10385            .arg(&ti);
10386        unsafe {
10387            b.launch(cfg)?;
10388        }
10389        Ok(())
10390    }
10391
10392    /// out[t, :] += sum_j w[t*n_used+j] * y[t*n_used+j, :], the j-sum sequential per thread —
10393    /// a pinned per-token reduction order, never atomics (the grouped prime's scatter).
10394    pub fn moe_pairs_weighted_scatter(
10395        &self,
10396        y: &CudaSlice<f32>,
10397        w: &CudaSlice<f32>,
10398        out: &mut CudaSlice<f32>,
10399        ncols: usize,
10400        n_used: usize,
10401        t: usize,
10402    ) -> Result<(), Box<dyn std::error::Error>> {
10403        let f = self.func("moe_pairs_weighted_scatter_f32");
10404        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10405        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10406        let __s_b = self.gpu.stream();
10407        let mut b = __s_b.launch_builder(&f);
10408        b.arg(y).arg(w).arg(&mut *out).arg(&nc).arg(&nu).arg(&ti);
10409        unsafe {
10410            b.launch(cfg)?;
10411        }
10412        Ok(())
10413    }
10414
10415    // ======== A2 GROUPED MoE PREFILL KERNELS ========
10416
10417    /// Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].
10418    pub fn gather_rows(
10419        &self,
10420        src: &CudaSlice<f32>,
10421        idx: &CudaSlice<i32>,
10422        dst: &mut CudaSlice<f32>,
10423        ncols: usize,
10424        m_e: usize,
10425    ) -> Result<(), Box<dyn std::error::Error>> {
10426        let f = self.func("gather_rows_f32");
10427        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
10428        let (nc, me) = (ncols as i32, m_e as i32);
10429        let __s_b = self.gpu.stream();
10430        let mut b = __s_b.launch_builder(&f);
10431        b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
10432        unsafe {
10433            b.launch(cfg)?;
10434        }
10435        Ok(())
10436    }
10437
10438    /// Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r].
10439    /// dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot.
10440    /// Scatter expert outputs into per-token slots (raw copy, no weight multiply).
10441    /// Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.
10442    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
10443    pub fn scatter_slot(
10444        &self,
10445        src: &CudaSlice<f32>,
10446        tok_idx: &CudaSlice<i32>,
10447        slot_idx: &CudaSlice<i32>,
10448        weight: &CudaSlice<f32>,
10449        dst: &mut CudaSlice<f32>,
10450        wbuf: &mut CudaSlice<f32>,
10451        ncols: usize,
10452        n_used: usize,
10453        m_e: usize,
10454    ) -> Result<(), Box<dyn std::error::Error>> {
10455        let f = self.func("scatter_add_slot_f32");
10456        let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
10457        let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
10458        let __s_b = self.gpu.stream();
10459        let mut b = __s_b.launch_builder(&f);
10460        b.arg(src)
10461            .arg(tok_idx)
10462            .arg(slot_idx)
10463            .arg(weight)
10464            .arg(dst)
10465            .arg(wbuf)
10466            .arg(&nc)
10467            .arg(&nu)
10468            .arg(&me);
10469        unsafe {
10470            b.launch(cfg)?;
10471        }
10472        Ok(())
10473    }
10474
10475    /// Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col].
10476    /// Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc).
10477    /// Uses FMA for bit-identity with the sequential axpy path.
10478    pub fn reduce_slots(
10479        &self,
10480        slots: &CudaSlice<f32>,
10481        wbuf: &CudaSlice<f32>,
10482        dst: &mut CudaSlice<f32>,
10483        ncols: usize,
10484        n_used: usize,
10485        t: usize,
10486    ) -> Result<(), Box<dyn std::error::Error>> {
10487        let f = self.func("reduce_slots_f32");
10488        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10489        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10490        let __s_b = self.gpu.stream();
10491        let mut b = __s_b.launch_builder(&f);
10492        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
10493        unsafe {
10494            b.launch(cfg)?;
10495        }
10496        Ok(())
10497    }
10498
10499    /// Canonical slot-order reduction with separately rounded multiply and add.
10500    ///
10501    /// This is the one-launch twin of repeated `axpy_host_into` calls. It preserves the official
10502    /// Step host-oracle arithmetic while allowing owner outputs to remain device-resident.
10503    pub fn reduce_slots_host(
10504        &self,
10505        slots: &CudaSlice<f32>,
10506        wbuf: &CudaSlice<f32>,
10507        dst: &mut CudaSlice<f32>,
10508        ncols: usize,
10509        n_used: usize,
10510        t: usize,
10511    ) -> Result<(), Box<dyn std::error::Error>> {
10512        let f = self.func("reduce_slots_host_f32");
10513        let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10514        let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10515        let __s_b = self.gpu.stream();
10516        let mut b = __s_b.launch_builder(&f);
10517        b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
10518        unsafe {
10519            b.launch(cfg)?;
10520        }
10521        Ok(())
10522    }
10523
10524    /// Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale).
10525    /// Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the
10526    /// forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up
10527    /// share `z`; q/k/v and wqkv/gate/beta/alpha share `h`) — quantize_q8_1 was 13.5% of decode
10528    /// GPU time, ~half of it redundant re-quantization of the same row.
10529    /// quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.
10530    pub fn quantize_q8_1_view(
10531        &self,
10532        x: &cudarc::driver::CudaView<f32>,
10533        m: usize,
10534        in_f: usize,
10535    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10536        let f = self.func("quantize_q8_1");
10537        let nblk = in_f / 32;
10538        let mut q = self.alloc_uninit::<i8>(m * in_f)?;
10539        let mut d = self.alloc_uninit::<f32>(m * nblk)?;
10540        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10541        let (inf, mi) = (in_f as i32, m as i32);
10542        let __s_b = self.gpu.stream();
10543        let mut b = __s_b.launch_builder(&f);
10544        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
10545        unsafe {
10546            b.launch(cfg)?;
10547        }
10548        Ok((q, d))
10549    }
10550
10551    pub fn quantize_q8_1(
10552        &self,
10553        x: &CudaSlice<f32>,
10554        m: usize,
10555        in_f: usize,
10556    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10557        let nblk = in_f / 32;
10558        let mut q = self.alloc_uninit::<i8>(m * in_f)?; // full-overwrite output: skip memset
10559        let mut d = self.alloc_uninit::<f32>(m * nblk)?; // full-overwrite output: skip memset
10560        // WARP-PER-BLOCK kernel: one warp per 32-block -> m*in_f threads total.
10561        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10562        let (inf, mi) = (in_f as i32, m as i32);
10563        if Self::pdl_on() && Self::pdl_wb_on() {
10564            {
10565                use cudarc::driver::{DevicePtr, DevicePtrMut};
10566                let s = &self.gpu.stream();
10567                let (px, _g0) = x.device_ptr(s);
10568                let (pq, _g1) = q.device_ptr_mut(s);
10569                let (pd, _g2) = d.device_ptr_mut(s);
10570                let mut ps = [
10571                    &px as *const _ as *mut std::ffi::c_void,
10572                    &pq as *const _ as *mut _,
10573                    &pd as *const _ as *mut _,
10574                    &inf as *const _ as *mut _,
10575                    &mi as *const _ as *mut _,
10576                ];
10577                unsafe {
10578                    self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10579                }
10580            }
10581            return Ok((q, d));
10582        }
10583        let f = self.func("quantize_q8_1");
10584        let __s_b = self.gpu.stream();
10585        let mut b = __s_b.launch_builder(&f);
10586        b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
10587        unsafe {
10588            b.launch(cfg)?;
10589        }
10590        Ok((q, d))
10591    }
10592
10593    /// Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16
10594    /// UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants.
10595    /// in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).
10596    pub fn quantize_fp4_act(
10597        &self,
10598        x: &CudaSlice<f32>,
10599        m: usize,
10600        in_f: usize,
10601    ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
10602        let f = self.func("quantize_fp4_act");
10603        let nb16 = in_f / 16;
10604        let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; // full-overwrite output: skip memset
10605        let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; // full-overwrite output: skip memset
10606        let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
10607        let (inf, mi) = (in_f as i32, m as i32);
10608        let __s_b = self.gpu.stream();
10609        let mut b = __s_b.launch_builder(&f);
10610        b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
10611        unsafe {
10612            b.launch(cfg)?;
10613        }
10614        Ok((aq4, ad4))
10615    }
10616
10617    /// Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw
10618    /// e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak,
10619    /// 3.5x int8). Activation `x` is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied
10620    /// post (scale==1.0 -> no-op). `bytes` = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.
10621    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
10622    pub fn qmatvec_gemm_nvfp4_fp4(
10623        &self,
10624        bytes: &CudaSlice<u8>,
10625        x: &CudaSlice<f32>,
10626        m: usize,
10627        in_f: usize,
10628        out_f: usize,
10629        row_bytes: usize,
10630        scale: f32,
10631    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10632        assert!(
10633            in_f.is_multiple_of(64),
10634            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
10635        );
10636        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
10637        let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
10638        if scale != 1.0 {
10639            self.scale_inplace(&mut y, scale, m * out_f)?;
10640        }
10641        Ok(y)
10642    }
10643
10644    /// Shared mxf4 GEMM launch (pre-quantized FP4 activation aq4/ad4). Same CTA tile as the int8 GEMM
10645    /// (BM=64 rows x BN=128 tokens, 4 warps). No macro-scale applied here.
10646    #[allow(clippy::too_many_arguments)]
10647    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
10648    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
10649    fn fp4_gemm_launch(
10650        &self,
10651        bytes: &CudaSlice<u8>,
10652        aq4: &CudaSlice<u32>,
10653        ad4: &CudaSlice<u8>,
10654        m: usize,
10655        in_f: usize,
10656        out_f: usize,
10657        row_bytes: usize,
10658    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10659        let f = self.func("qmatvec_gemm_nvfp4_fp4");
10660        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
10661        const BM: u32 = 64;
10662        const BN: u32 = 256;
10663        let cfg = LaunchConfig {
10664            grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
10665            block_dim: (32, 4, 1),
10666            shared_mem_bytes: 0,
10667        };
10668        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10669        let __s_b = self.gpu.stream();
10670        let mut b = __s_b.launch_builder(&f);
10671        b.arg(bytes)
10672            .arg(aq4)
10673            .arg(ad4)
10674            .arg(&mut y)
10675            .arg(&inf)
10676            .arg(&outf)
10677            .arg(&mi)
10678            .arg(&rb);
10679        unsafe {
10680            b.launch(cfg)?;
10681        }
10682        Ok(y)
10683    }
10684
10685    /// Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).
10686    pub fn qmatvec_gemm_nvfp4_fp4_raw(
10687        &self,
10688        bytes: &CudaSlice<u8>,
10689        x: &CudaSlice<f32>,
10690        m: usize,
10691        in_f: usize,
10692        out_f: usize,
10693        row_bytes: usize,
10694    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10695        assert!(
10696            in_f.is_multiple_of(64),
10697            "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
10698        );
10699        let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
10700        self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
10701    }
10702
10703    /// Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.
10704    pub fn qmatvec_q8_0_fast(
10705        &self,
10706        w: &CudaSlice<u8>,
10707        x: &CudaSlice<f32>,
10708        m: usize,
10709        in_f: usize,
10710        out_f: usize,
10711        row_bytes: usize,
10712    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10713        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10714        let f = self.func("qmatvec_q8_0_dp4a");
10715        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
10716        let cfg = LaunchConfig {
10717            grid_dim: (out_f as u32, m as u32, 1),
10718            block_dim: (128, 1, 1),
10719            shared_mem_bytes: 0,
10720        };
10721        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10722        let __s_b = self.gpu.stream();
10723        let mut b = __s_b.launch_builder(&f);
10724        b.arg(w)
10725            .arg(&aq)
10726            .arg(&ad)
10727            .arg(&mut y)
10728            .arg(&inf)
10729            .arg(&outf)
10730            .arg(&mi)
10731            .arg(&rb);
10732        unsafe {
10733            b.launch(cfg)?;
10734        }
10735        Ok(y)
10736    }
10737
10738    /// Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
10739    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
10740    pub fn qmatvec_q4_K_fast(
10741        &self,
10742        w: &CudaSlice<u8>,
10743        x: &CudaSlice<f32>,
10744        m: usize,
10745        in_f: usize,
10746        out_f: usize,
10747        row_bytes: usize,
10748    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10749        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10750        let f = self.func("qmatvec_q4_K_dp4a");
10751        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
10752        let cfg = LaunchConfig {
10753            grid_dim: (out_f as u32, m as u32, 1),
10754            block_dim: (128, 1, 1),
10755            shared_mem_bytes: 0,
10756        };
10757        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10758        let __s_b = self.gpu.stream();
10759        let mut b = __s_b.launch_builder(&f);
10760        b.arg(w)
10761            .arg(&aq)
10762            .arg(&ad)
10763            .arg(&mut y)
10764            .arg(&inf)
10765            .arg(&outf)
10766            .arg(&mi)
10767            .arg(&rb);
10768        unsafe {
10769            b.launch(cfg)?;
10770        }
10771        Ok(y)
10772    }
10773
10774    /// Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).
10775    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
10776    pub fn qmatvec_q6_K_fast(
10777        &self,
10778        w: &CudaSlice<u8>,
10779        x: &CudaSlice<f32>,
10780        m: usize,
10781        in_f: usize,
10782        out_f: usize,
10783        row_bytes: usize,
10784    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10785        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10786        let f = self.func("qmatvec_q6_K_dp4a");
10787        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
10788        let cfg = LaunchConfig {
10789            grid_dim: (out_f as u32, m as u32, 1),
10790            block_dim: (128, 1, 1),
10791            shared_mem_bytes: 0,
10792        };
10793        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10794        let __s_b = self.gpu.stream();
10795        let mut b = __s_b.launch_builder(&f);
10796        b.arg(w)
10797            .arg(&aq)
10798            .arg(&ad)
10799            .arg(&mut y)
10800            .arg(&inf)
10801            .arg(&outf)
10802            .arg(&mi)
10803            .arg(&rb);
10804        unsafe {
10805            b.launch(cfg)?;
10806        }
10807        Ok(y)
10808    }
10809
10810    /// Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.
10811    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
10812    pub fn qmatvec_q5_K_fast(
10813        &self,
10814        w: &CudaSlice<u8>,
10815        x: &CudaSlice<f32>,
10816        m: usize,
10817        in_f: usize,
10818        out_f: usize,
10819        row_bytes: usize,
10820    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10821        self.qmatvec_dp4a_named(
10822            "qmatvec_q5_K_dp4a",
10823            &w.slice(0..w.len()),
10824            x,
10825            m,
10826            in_f,
10827            out_f,
10828            row_bytes,
10829        )
10830    }
10831    /// Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).
10832    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
10833    pub fn qmatvec_q3_K_fast(
10834        &self,
10835        w: &CudaSlice<u8>,
10836        x: &CudaSlice<f32>,
10837        m: usize,
10838        in_f: usize,
10839        out_f: usize,
10840        row_bytes: usize,
10841    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10842        self.qmatvec_dp4a_named(
10843            "qmatvec_q3_K_dp4a",
10844            &w.slice(0..w.len()),
10845            x,
10846            m,
10847            in_f,
10848            out_f,
10849            row_bytes,
10850        )
10851    }
10852    /// A6 split-plane twin of `qmatvec_nvfp4_fast` (weights repacked; used by the rp gates).
10853    pub fn qmatvec_nvfp4_fast_rp(
10854        &self,
10855        w: &CudaSlice<u8>,
10856        x: &CudaSlice<f32>,
10857        m: usize,
10858        in_f: usize,
10859        out_f: usize,
10860        row_bytes: usize,
10861    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10862        assert!(
10863            in_f.is_multiple_of(64),
10864            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10865        );
10866        self.qmatvec_dp4a_named(
10867            "qmatvec_nvfp4_dp4a_rp",
10868            &w.slice(0..w.len()),
10869            x,
10870            m,
10871            in_f,
10872            out_f,
10873            row_bytes,
10874        )
10875    }
10876    /// Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).
10877    pub fn qmatvec_nvfp4_fast(
10878        &self,
10879        w: &cudarc::driver::CudaView<'_, u8>,
10880        x: &CudaSlice<f32>,
10881        m: usize,
10882        in_f: usize,
10883        out_f: usize,
10884        row_bytes: usize,
10885    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10886        // B1: the NVFP4 dp4a kernel maps two 32-elem q8_1 blocks onto one 64-elem block_nvfp4
10887        // (sblk = g >> 1). in_f must be a multiple of 64 or the last block reads a partial superblock.
10888        assert!(
10889            in_f.is_multiple_of(64),
10890            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10891        );
10892        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
10893    }
10894    /// Slot-major-layout twin of `qmatvec_nvfp4_fast`: bit-identical per row, coalesced
10895    /// reads. Since the 2026-08-29 `MEMRA_NVFP4_BANK_V2` door removal its only in-tree
10896    /// producer of slot-major banks is the EP2 whole-expert bank build; this is EP2's
10897    /// host-canonical oracle reader (plus offline harnesses like moe_tp2_repro).
10898    pub fn qmatvec_nvfp4_fast_v2(
10899        &self,
10900        w: &cudarc::driver::CudaView<'_, u8>,
10901        x: &CudaSlice<f32>,
10902        m: usize,
10903        in_f: usize,
10904        out_f: usize,
10905        row_bytes: usize,
10906    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10907        assert!(
10908            in_f.is_multiple_of(64),
10909            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10910        );
10911        self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
10912    }
10913    /// Stage-B (optional perf): IQ4_XS codebook int8 dp4a.
10914    #[allow(non_snake_case)] // Keep GGUF qtype spelling visible at the public kernel boundary.
10915    pub fn qmatvec_iq4_XS_fast(
10916        &self,
10917        w: &CudaSlice<u8>,
10918        x: &CudaSlice<f32>,
10919        m: usize,
10920        in_f: usize,
10921        out_f: usize,
10922        row_bytes: usize,
10923    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10924        self.qmatvec_dp4a_named(
10925            "qmatvec_iq4_XS_dp4a",
10926            &w.slice(0..w.len()),
10927            x,
10928            m,
10929            in_f,
10930            out_f,
10931            row_bytes,
10932        )
10933    }
10934
10935    /// Shared dp4a launcher: quantize_q8_1 then call the named kernel (grid (out,m), block 64).
10936    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
10937    fn qmatvec_dp4a_named(
10938        &self,
10939        name: &str,
10940        w: &cudarc::driver::CudaView<'_, u8>,
10941        x: &CudaSlice<f32>,
10942        m: usize,
10943        in_f: usize,
10944        out_f: usize,
10945        row_bytes: usize,
10946    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10947        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10948        let f = self.func(name);
10949        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
10950        let cfg = LaunchConfig {
10951            grid_dim: (out_f as u32, m as u32, 1),
10952            block_dim: (128, 1, 1),
10953            shared_mem_bytes: 0,
10954        };
10955        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10956        let __s_b = self.gpu.stream();
10957        let mut b = __s_b.launch_builder(&f);
10958        b.arg(w)
10959            .arg(&aq)
10960            .arg(&ad)
10961            .arg(&mut y)
10962            .arg(&inf)
10963            .arg(&outf)
10964            .arg(&mi)
10965            .arg(&rb);
10966        unsafe {
10967            b.launch(cfg)?;
10968        }
10969        Ok(y)
10970    }
10971
10972    /// NVFP4 dp4a matvec over PRE-QUANTIZED q8_1 activations, writing a caller-provided output.
10973    /// Same kernel and math as `qmatvec_nvfp4_fast` (which quantizes internally and allocates
10974    /// its output); this entry exists so a routed-expert program can quantize one activation
10975    /// ONCE and reuse it across every expert's gate/up, feed `silu_mul_scaled_q8_1`'s q8_1
10976    /// straight into down, and keep persistent output workspaces — zero per-expert allocations.
10977    #[allow(clippy::too_many_arguments)]
10978    pub fn qmatvec_nvfp4_fast_prequant_into(
10979        &self,
10980        w: &CudaSlice<u8>,
10981        aq: &CudaSlice<i8>,
10982        ad: &CudaSlice<f32>,
10983        y: &mut CudaSlice<f32>,
10984        m: usize,
10985        in_f: usize,
10986        out_f: usize,
10987        row_bytes: usize,
10988    ) -> Result<(), Box<dyn std::error::Error>> {
10989        assert!(
10990            in_f.is_multiple_of(64),
10991            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10992        );
10993        if y.len() < m * out_f {
10994            return Err(format!(
10995                "NVFP4 prequant output {} is shorter than {m}x{out_f}",
10996                y.len()
10997            )
10998            .into());
10999        }
11000        let f = self.func("qmatvec_nvfp4_dp4a");
11001        let cfg = LaunchConfig {
11002            grid_dim: (out_f as u32, m as u32, 1),
11003            block_dim: (128, 1, 1),
11004            shared_mem_bytes: 0,
11005        };
11006        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11007        let __s_b = self.gpu.stream();
11008        let mut b = __s_b.launch_builder(&f);
11009        b.arg(w)
11010            .arg(aq)
11011            .arg(ad)
11012            .arg(y)
11013            .arg(&inf)
11014            .arg(&outf)
11015            .arg(&mi)
11016            .arg(&rb);
11017        unsafe {
11018            b.launch(cfg)?;
11019        }
11020        Ok(())
11021    }
11022
11023    /// Fused QKV F32 matvec (one launch for all three rank-local projections; see the kernel
11024    /// doc for the numeric-class note). Requires `in_f % 4 == 0`.
11025    #[allow(clippy::too_many_arguments)]
11026    pub fn matvec_f32_qkv_into(
11027        &self,
11028        wq: &CudaSlice<f32>,
11029        wk: &CudaSlice<f32>,
11030        wv: &CudaSlice<f32>,
11031        wg: &CudaSlice<f32>,
11032        x: &CudaSlice<f32>,
11033        yq: &mut CudaSlice<f32>,
11034        yk: &mut CudaSlice<f32>,
11035        yv: &mut CudaSlice<f32>,
11036        yg: &mut CudaSlice<f32>,
11037        in_f: usize,
11038        out_q: usize,
11039        out_kv: usize,
11040        out_g: usize,
11041    ) -> Result<(), Box<dyn std::error::Error>> {
11042        if !in_f.is_multiple_of(4)
11043            || wq.len() != out_q * in_f
11044            || wk.len() != out_kv * in_f
11045            || wv.len() != out_kv * in_f
11046            || wg.len() < out_g * in_f
11047            || x.len() < in_f
11048            || yq.len() < out_q
11049            || yk.len() < out_kv
11050            || yv.len() < out_kv
11051            || (out_g > 0 && yg.len() < out_g)
11052        {
11053            return Err(format!(
11054                "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
11055                 wq={} wk={} wv={} wg={}",
11056                wq.len(),
11057                wk.len(),
11058                wv.len(),
11059                wg.len()
11060            )
11061            .into());
11062        }
11063        let f = self.func("matvec_f32_qkv");
11064        let cfg = LaunchConfig {
11065            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
11066            block_dim: (128, 1, 1),
11067            shared_mem_bytes: 0,
11068        };
11069        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
11070        let __s_b = self.gpu.stream();
11071        let mut b = __s_b.launch_builder(&f);
11072        b.arg(wq)
11073            .arg(wk)
11074            .arg(wv)
11075            .arg(wg)
11076            .arg(x)
11077            .arg(yq)
11078            .arg(yk)
11079            .arg(yv)
11080            .arg(yg)
11081            .arg(&inf)
11082            .arg(&oq)
11083            .arg(&okv)
11084            .arg(&og);
11085        unsafe {
11086            b.launch(cfg)?;
11087        }
11088        Ok(())
11089    }
11090
11091    /// PROGRAM 2 (`MEMRA_NVFP4_SEL_GU`, default OFF): the routed gate and up sweeps in ONE
11092    /// launch. The two sweeps share `sel`/`aq`/`ad` and have identical geometry, so blocks
11093    /// `[0,out_f)` run the exact `_sel_v2` body on the GATE bank and `[out_f,2*out_f)` on the UP
11094    /// bank — per-row BIT-IDENTICAL to two `qmatvec_nvfp4_sel_into` calls, with half the sweep
11095    /// launches and double the grid fill.
11096    ///
11097    /// SLOT-MAJOR ONLY, and the caller proves it: the kernel reads the slot-major byte map, so
11098    /// this refuses banks that do not carry it rather than trusting an env door. In the removed
11099    /// implementation this fusion auto-armed on `nvfp4_bank_v2_on()` with NO door of its own,
11100    /// which is one of the three programs that moved together behind one env var and made the
11101    /// 2026-08-29 bisect unable to name a mechanism (DIAGNOSIS.md).
11102    #[allow(clippy::too_many_arguments)]
11103    pub fn qmatvec_nvfp4_sel_gu_into(
11104        &self,
11105        gate_bank: &CudaSlice<u8>,
11106        up_bank: &CudaSlice<u8>,
11107        sel: &CudaSlice<i32>,
11108        aq: &CudaSlice<i8>,
11109        ad: &CudaSlice<f32>,
11110        yg: &mut CudaSlice<f32>,
11111        yu: &mut CudaSlice<f32>,
11112        n_sel: usize,
11113        in_f: usize,
11114        out_f: usize,
11115        row_bytes: usize,
11116        expert_stride: usize,
11117        slot_major: bool,
11118    ) -> Result<(), Box<dyn std::error::Error>> {
11119        assert!(
11120            in_f.is_multiple_of(64),
11121            "NVFP4 dp4a requires in_f % 64 == 0"
11122        );
11123        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
11124            return Err("NVFP4 gu sel geometry".into());
11125        }
11126        if !slot_major {
11127            return Err(
11128                "NVFP4 gu sel fusion reads slot-major rows: these banks are block_nvfp4 \
11129                        v1 (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
11130                    .into(),
11131            );
11132        }
11133        // MEMRA_NVFP4_SEL_GU_RPW=2|4 (sub-door, default OFF, UNPRICED): multirow twin — the
11134        // activation group is read once and reused across RPW rows' gate+up dots. Per-row
11135        // accumulation order and reduce tree are the base kernel's -> bit-identical.
11136        static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
11137        let rpw = *RPW.get_or_init(|| {
11138            std::env::var("MEMRA_NVFP4_SEL_GU_RPW")
11139                .ok()
11140                .and_then(|v| v.parse().ok())
11141                .filter(|r| *r == 2 || *r == 4)
11142                .unwrap_or(1)
11143        });
11144        let rpw = if out_f.is_multiple_of(rpw) { rpw } else { 1 };
11145        // MEMRA_NVFP4_SEL_GU_WPR=1 (sub-door, default OFF, UNPRICED): warp-per-row.
11146        // NUMERIC-CLASS — the per-row REDUCTION ORDER changes, so a bit tape cannot apply and
11147        // acceptance is the argmax gate plus the boot battery (the QKV_FUSED/BF16_MMV class).
11148        // It is deliberately NOT part of this lane's priced arms, which are all bit-gateable.
11149        static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11150        let wpr =
11151            *WPR.get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_GU_WPR").as_deref() == Ok("1"));
11152        let f = self.func(match (wpr, rpw) {
11153            (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
11154            (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
11155            (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
11156            _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
11157        });
11158        let cfg = LaunchConfig {
11159            grid_dim: if wpr {
11160                (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
11161            } else if rpw == 1 {
11162                ((2 * out_f) as u32, n_sel as u32, 1)
11163            } else {
11164                ((out_f / rpw) as u32, n_sel as u32, 1)
11165            },
11166            block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
11167            shared_mem_bytes: 0,
11168        };
11169        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11170        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11171        let (ars, adrs) = (0i64, 0i64);
11172        let __s_b = self.gpu.stream();
11173        let mut b = __s_b.launch_builder(&f);
11174        b.arg(gate_bank)
11175            .arg(up_bank)
11176            .arg(sel)
11177            .arg(aq)
11178            .arg(ad)
11179            .arg(yg)
11180            .arg(yu)
11181            .arg(&inf)
11182            .arg(&outf)
11183            .arg(&ns)
11184            .arg(&rb)
11185            .arg(&es)
11186            .arg(&ars)
11187            .arg(&adrs);
11188        unsafe {
11189            b.launch(cfg)?;
11190        }
11191        Ok(())
11192    }
11193
11194    /// PROGRAM 3 (`MEMRA_NVFP4_SEL_DOWN8`, **default ON since 2026-09-01**): the DOWN sweep and
11195    /// the route-weight
11196    /// combine in ONE launch (`qmatvec_nvfp4_dp4a_sel_v2_down8`, the q8 `down8 w8` occupancy arm
11197    /// ported to the NVFP4 banks). Block = `(32, n_sel)`: one warp per slot instead of one warp
11198    /// per (row, slot), and the `n_sel x out_f` partial buffer disappears. BIT-IDENTICAL to
11199    /// `qmatvec_nvfp4_sel_into` + `axpy_rows_seq_md_into` — same dot program, same reduce tree,
11200    /// same slot-ordered combine chain.
11201    ///
11202    /// Requires slot-major rows and `nsb <= 32` (the fit-block class the reduce identity is
11203    /// argued at). The removed implementation refused on `!nvfp4_bank_v2_on()`; it now refuses on
11204    /// the LAYOUT THE CALLER READ OFF THE BANK, so the guard cannot disagree with the bytes.
11205    #[allow(clippy::too_many_arguments)]
11206    pub fn qmatvec_nvfp4_sel_down8_into(
11207        &self,
11208        bank: &CudaSlice<u8>,
11209        sel: &CudaSlice<i32>,
11210        aq: &CudaSlice<i8>,
11211        ad: &CudaSlice<f32>,
11212        route_w: &CudaSlice<f32>,
11213        md: &CudaSlice<f32>,
11214        dst: &mut CudaSlice<f32>,
11215        n_sel: usize,
11216        in_f: usize,
11217        out_f: usize,
11218        row_bytes: usize,
11219        expert_stride: usize,
11220        act_row_stride: usize,
11221        ad_row_stride: usize,
11222        slot_major: bool,
11223    ) -> Result<(), Box<dyn std::error::Error>> {
11224        if !in_f.is_multiple_of(64)
11225            || n_sel == 0
11226            || n_sel > 8
11227            || (in_f >> 5) > 32
11228            || dst.len() < out_f
11229            || sel.len() < n_sel
11230            || route_w.len() < n_sel
11231        {
11232            return Err(format!(
11233                "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
11234                dst.len()
11235            )
11236            .into());
11237        }
11238        if !slot_major {
11239            return Err(
11240                "NVFP4 sel down8 reads slot-major rows: this shard is block_nvfp4 v1 \
11241                        (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
11242                    .into(),
11243            );
11244        }
11245        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
11246        let cfg = LaunchConfig {
11247            grid_dim: (out_f as u32, 1, 1),
11248            block_dim: (32, n_sel as u32, 1),
11249            shared_mem_bytes: 0,
11250        };
11251        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11252        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11253        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
11254        let __s_b = self.gpu.stream();
11255        let mut b = __s_b.launch_builder(&f);
11256        b.arg(bank)
11257            .arg(sel)
11258            .arg(aq)
11259            .arg(ad)
11260            .arg(route_w)
11261            .arg(md)
11262            .arg(dst)
11263            .arg(&inf)
11264            .arg(&outf)
11265            .arg(&ns)
11266            .arg(&rb)
11267            .arg(&es)
11268            .arg(&ars)
11269            .arg(&adrs);
11270        unsafe {
11271            b.launch(cfg)?;
11272        }
11273        Ok(())
11274    }
11275
11276    /// EP2 owner-guarded gate+up sweep: full-width rows, pairs whose expert this rank
11277    /// does not own exit immediately. Per-pair dot == the _sel_v2 gu body.
11278    #[allow(clippy::too_many_arguments)]
11279    pub fn qmatvec_nvfp4_sel_gu_ep_into(
11280        &self,
11281        gate_bank: &CudaSlice<u8>,
11282        up_bank: &CudaSlice<u8>,
11283        sel: &CudaSlice<i32>,
11284        aq: &CudaSlice<i8>,
11285        ad: &CudaSlice<f32>,
11286        yg: &mut CudaSlice<f32>,
11287        yu: &mut CudaSlice<f32>,
11288        n_sel: usize,
11289        in_f: usize,
11290        out_f: usize,
11291        row_bytes: usize,
11292        expert_stride: usize,
11293        owner: usize,
11294    ) -> Result<(), Box<dyn std::error::Error>> {
11295        assert!(
11296            in_f.is_multiple_of(64),
11297            "NVFP4 dp4a requires in_f % 64 == 0"
11298        );
11299        if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
11300            return Err("NVFP4 gu ep geometry".into());
11301        }
11302        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
11303        let cfg = LaunchConfig {
11304            grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
11305            block_dim: (128, 1, 1),
11306            shared_mem_bytes: 0,
11307        };
11308        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
11309        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11310        let (ars, adrs) = (0i64, 0i64);
11311        let __s_b = self.gpu.stream();
11312        let mut b = __s_b.launch_builder(&f);
11313        b.arg(gate_bank)
11314            .arg(up_bank)
11315            .arg(sel)
11316            .arg(aq)
11317            .arg(ad)
11318            .arg(yg)
11319            .arg(yu)
11320            .arg(&inf)
11321            .arg(&outf)
11322            .arg(&ns)
11323            .arg(&rb)
11324            .arg(&es)
11325            .arg(&ars)
11326            .arg(&adrs)
11327            .arg(&own);
11328        unsafe {
11329            b.launch(cfg)?;
11330        }
11331        Ok(())
11332    }
11333
11334    /// EP2 owner-guarded SwiGLU (q8_1 emission), clamped or plain by `limit`.
11335    #[allow(clippy::too_many_arguments)]
11336    pub fn silu_mul_scaled_q8_1_sel_ep_into(
11337        &self,
11338        gate: &CudaSlice<f32>,
11339        up: &CudaSlice<f32>,
11340        gmac: &CudaSlice<f32>,
11341        umac: &CudaSlice<f32>,
11342        sel: &CudaSlice<i32>,
11343        limit: Option<f32>,
11344        out_q: &mut CudaSlice<i8>,
11345        out_d: &mut CudaSlice<f32>,
11346        n_per: usize,
11347        n_sel: usize,
11348        owner: usize,
11349    ) -> Result<(), Box<dyn std::error::Error>> {
11350        if !n_per.is_multiple_of(32)
11351            || out_q.len() < n_sel * n_per
11352            || out_d.len() < n_sel * n_per / 32
11353        {
11354            return Err("NVFP4 silu ep geometry".into());
11355        }
11356        let f = self.func("silu_mul_scaled_q8_1_sel_ep");
11357        let warps = n_sel * n_per / 32;
11358        let cfg = LaunchConfig {
11359            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
11360            block_dim: (128, 1, 1),
11361            shared_mem_bytes: 0,
11362        };
11363        let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
11364        let (lim, has) = match limit {
11365            Some(l) => (l, 1i32),
11366            None => (0.0f32, 0i32),
11367        };
11368        let __s_b = self.gpu.stream();
11369        let mut b = __s_b.launch_builder(&f);
11370        b.arg(gate)
11371            .arg(up)
11372            .arg(gmac)
11373            .arg(umac)
11374            .arg(sel)
11375            .arg(&lim)
11376            .arg(&has)
11377            .arg(out_q)
11378            .arg(out_d)
11379            .arg(&np)
11380            .arg(&ns)
11381            .arg(&own);
11382        unsafe {
11383            b.launch(cfg)?;
11384        }
11385        Ok(())
11386    }
11387
11388    /// EP2 owner-guarded down + owned-slot combine in one launch (block `(32, n_sel)`).
11389    #[allow(clippy::too_many_arguments)]
11390    pub fn qmatvec_nvfp4_sel_down8_ep_into(
11391        &self,
11392        bank: &CudaSlice<u8>,
11393        sel: &CudaSlice<i32>,
11394        aq: &CudaSlice<i8>,
11395        ad: &CudaSlice<f32>,
11396        route_w: &CudaSlice<f32>,
11397        md: &CudaSlice<f32>,
11398        dst: &mut CudaSlice<f32>,
11399        n_sel: usize,
11400        in_f: usize,
11401        out_f: usize,
11402        row_bytes: usize,
11403        expert_stride: usize,
11404        act_row_stride: usize,
11405        ad_row_stride: usize,
11406        owner: usize,
11407    ) -> Result<(), Box<dyn std::error::Error>> {
11408        if !in_f.is_multiple_of(64)
11409            || n_sel == 0
11410            || n_sel > 8
11411            || (in_f >> 5) > 64
11412            || dst.len() < out_f
11413        {
11414            return Err("NVFP4 down8 ep geometry".into());
11415        }
11416        let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
11417        let cfg = LaunchConfig {
11418            grid_dim: (out_f as u32, 1, 1),
11419            block_dim: (32, n_sel as u32, 1),
11420            shared_mem_bytes: 0,
11421        };
11422        let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
11423        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11424        let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
11425        let __s_b = self.gpu.stream();
11426        let mut b = __s_b.launch_builder(&f);
11427        b.arg(bank)
11428            .arg(sel)
11429            .arg(aq)
11430            .arg(ad)
11431            .arg(route_w)
11432            .arg(md)
11433            .arg(dst)
11434            .arg(&inf)
11435            .arg(&outf)
11436            .arg(&ns)
11437            .arg(&rb)
11438            .arg(&es)
11439            .arg(&ars)
11440            .arg(&adrs)
11441            .arg(&own);
11442        unsafe {
11443            b.launch(cfg)?;
11444        }
11445        Ok(())
11446    }
11447
11448    /// Selected-experts batched twin of `qmatvec_nvfp4_fast_prequant_into`: one launch covers
11449    /// every selected expert, weights indexed `sel[t] * expert_stride` into a contiguous
11450    /// per-rank bank, activations advancing `act_row_stride`/`ad_row_stride` elements per
11451    /// selection (0 for a shared input). Per (expert, row) bit-identical to the per-expert
11452    /// kernel — the batching only removes host launch latency.
11453    ///
11454    /// `slot_major` names the LAYOUT OF THE BYTES AT `bank` and is REQUIRED, never defaulted:
11455    /// true routes the `_sel_v2` reader (slot g's 16 qs bytes at `g*16`, scale tail at
11456    /// `nslots*16`), false the block_nvfp4 v1 reader. The caller reads it off the resident bank
11457    /// (`ResidentNvfp4{Column,Row}BankRank::slot_major`) — never off an env door, and never with
11458    /// a default. A defaulted layout scalar in exactly this position is what produced the
11459    /// 2026-08-29 step37 corruption (`kq_fetch(..., int in_f = 0)`,
11460    /// research/step37-bankv3-20260901/DIAGNOSIS.md).
11461    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
11462    pub fn qmatvec_nvfp4_sel_into(
11463        &self,
11464        bank: &CudaSlice<u8>,
11465        sel: &CudaSlice<i32>,
11466        aq: &CudaSlice<i8>,
11467        ad: &CudaSlice<f32>,
11468        y: &mut CudaSlice<f32>,
11469        n_sel: usize,
11470        in_f: usize,
11471        out_f: usize,
11472        row_bytes: usize,
11473        expert_stride: usize,
11474        act_row_stride: usize,
11475        ad_row_stride: usize,
11476        slot_major: bool,
11477    ) -> Result<(), Box<dyn std::error::Error>> {
11478        assert!(
11479            in_f.is_multiple_of(64),
11480            "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
11481        );
11482        if y.len() < n_sel * out_f || sel.len() < n_sel {
11483            return Err(format!(
11484                "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
11485                y.len(),
11486                sel.len()
11487            )
11488            .into());
11489        }
11490        // MEMRA_SEL_MR=1: 4-concurrent-row-groups twin — per row bit-identical (same 128-thread
11491        // striding + reduction). MEASURED SLOWER on the 188-SM card (40.8 vs 42.9 tok/s e2e,
11492        // 2026-08-21: 512-thread blocks trade occupancy for launch-tail savings and lose; the
11493        // sequential-rows variant was flat). Default stays the single-row form.
11494        // MEMRA_SEL_STREAM=1: 16-rows-per-block streaming twin with next-row register
11495        // prefetch (bit-identical per row; one group per thread, so in_f <= 4096 only).
11496        static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
11497        let mode = *MR.get_or_init(|| {
11498            if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
11499                2
11500            } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
11501                1
11502            } else {
11503                0
11504            }
11505        });
11506        let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
11507        // Mode 3 is the SLOT-MAJOR reader, and it is chosen by the BANK's layout, not by an env
11508        // door: `MEMRA_SEL_MR`/`MEMRA_SEL_STREAM` are v1-layout probes and cannot read these
11509        // bytes at all, so the layout overrides them rather than racing them.
11510        let mode = if slot_major { 3 } else { mode };
11511        // MEMRA_NVFP4_SEL_SM_STREAM=1 (sub-door of MEMRA_NVFP4_BANK_SM, default OFF, UNPRICED):
11512        // 8 contiguous rows per block with next-row int4 prefetch. Needs 16B-aligned rows
11513        // (step37 gate/up 2304B yes, down 360B no -> single-row) and one slot per thread.
11514        static SM_STREAM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11515        let sm_stream = mode == 3
11516            && *SM_STREAM
11517                .get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_SM_STREAM").as_deref() == Ok("1"))
11518            && row_bytes.is_multiple_of(16)
11519            && in_f <= 4096;
11520        let kname = match (mode, sm_stream) {
11521            (3, true) => "qmatvec_nvfp4_dp4a_sel_v2s",
11522            (3, false) => "qmatvec_nvfp4_dp4a_sel_v2",
11523            (2, _) => "qmatvec_nvfp4_dp4a_sel_stream",
11524            (1, _) => "qmatvec_nvfp4_dp4a_sel_mr4",
11525            _ => "qmatvec_nvfp4_dp4a_sel",
11526        };
11527        // ENGAGEMENT RECEIPT for PROGRAM 1, one line per distinct (kernel, geometry) pair. The
11528        // door being SET in the environment does not prove the slot-major READER ran; only the
11529        // selected kernel name does. Without this, a pricing cell that reports a flat delta
11530        // cannot distinguish "the program is worth nothing" from "the program never ran" — the
11531        // defect the MEMRA_BF16_MMV lane hit when its engagement grep returned 0 in both arms.
11532        {
11533            static SEEN_SEL: std::sync::Mutex<Vec<(&'static str, usize, usize)>> =
11534                std::sync::Mutex::new(Vec::new());
11535            let combo = (kname, in_f, out_f);
11536            let mut seen = SEEN_SEL.lock().unwrap();
11537            if !seen.contains(&combo) {
11538                seen.push(combo);
11539                eprintln!(
11540                    "[nvfp4-sel] kernel={kname} slot_major={slot_major} in_f={in_f} \
11541                     out_f={out_f} nsb={} row_bytes={row_bytes}",
11542                    in_f >> 5
11543                );
11544            }
11545        }
11546        let f = self.func(kname);
11547        // Thread-fit block for narrow rows (the DOWN sweep: in_f=640 -> nsb=20 slots left
11548        // 108 of 128 threads idle AND thread-capped resident blocks). blockDim >= nsb keeps
11549        // thread g on slot g; the dropped threads contributed exact 0.0 partials to the
11550        // reduce, so the result bits are unchanged. Applies to the single-row forms only.
11551        let nsb = in_f >> 5;
11552        let fit_block: u32 = if (mode == 0 || mode == 3) && !sm_stream && nsb <= 32 {
11553            32
11554        } else if mode == 1 {
11555            512
11556        } else {
11557            128
11558        };
11559        let cfg = LaunchConfig {
11560            grid_dim: (
11561                if sm_stream {
11562                    (out_f as u32).div_ceil(8)
11563                } else {
11564                    match mode {
11565                        2 => (out_f as u32).div_ceil(16),
11566                        1 => (out_f as u32).div_ceil(4),
11567                        _ => out_f as u32,
11568                    }
11569                },
11570                n_sel as u32,
11571                1,
11572            ),
11573            block_dim: (fit_block, 1, 1),
11574            shared_mem_bytes: 0,
11575        };
11576        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11577        let (rb, es, ars, adrs) = (
11578            row_bytes as i64,
11579            expert_stride as i64,
11580            act_row_stride as i64,
11581            ad_row_stride as i64,
11582        );
11583        let __s_b = self.gpu.stream();
11584        let mut b = __s_b.launch_builder(&f);
11585        b.arg(bank)
11586            .arg(sel)
11587            .arg(aq)
11588            .arg(ad)
11589            .arg(y)
11590            .arg(&inf)
11591            .arg(&outf)
11592            .arg(&ns)
11593            .arg(&rb)
11594            .arg(&es)
11595            .arg(&ars)
11596            .arg(&adrs);
11597        unsafe {
11598            b.launch(cfg)?;
11599        }
11600        Ok(())
11601    }
11602
11603    /// W4A16 selected-expert gate+up pair. `x_bf16` contains checkpoint-rounded BF16
11604    /// activations; selected ids are local to the rank's contiguous expert bank.
11605    #[allow(clippy::too_many_arguments)]
11606    pub fn qmatvec_nvfp4_bf16_sel_dual_rows_into(
11607        &self,
11608        gate_bank: &CudaSlice<u8>,
11609        up_bank: &CudaSlice<u8>,
11610        sel: &CudaSlice<i32>,
11611        token_rows: &CudaSlice<i32>,
11612        x_bf16: &CudaSlice<u8>,
11613        gate_out: &mut CudaSlice<f32>,
11614        up_out: &mut CudaSlice<f32>,
11615        n_sel: usize,
11616        in_f: usize,
11617        out_f: usize,
11618        row_bytes: usize,
11619        expert_stride: usize,
11620        tokens: usize,
11621    ) -> Result<(), Box<dyn std::error::Error>> {
11622        if !in_f.is_multiple_of(64)
11623            || sel.len() < n_sel
11624            || token_rows.len() < n_sel
11625            || gate_out.len() < n_sel * out_f
11626            || up_out.len() < n_sel * out_f
11627            || x_bf16.len() < 2 * in_f * tokens
11628        {
11629            return Err(format!(
11630                "W4A16 NVFP4 dual selected rows geometry sel={} token_rows={} x={} gate={} up={} \
11631                 n_sel={n_sel} tokens={tokens} in={in_f} out={out_f}",
11632                sel.len(),
11633                token_rows.len(),
11634                x_bf16.len(),
11635                gate_out.len(),
11636                up_out.len(),
11637            )
11638            .into());
11639        }
11640        let adjacent_rows = tokens > 1;
11641        let f = if adjacent_rows {
11642            self.func("qmatvec_nvfp4_bf16_sel_quad_rows")
11643        } else {
11644            self.func("qmatvec_nvfp4_bf16_sel_dual_rows")
11645        };
11646        let cfg = LaunchConfig {
11647            grid_dim: (
11648                if adjacent_rows {
11649                    out_f.div_ceil(2) as u32
11650                } else {
11651                    (2 * out_f) as u32
11652                },
11653                n_sel as u32,
11654                1,
11655            ),
11656            block_dim: (256, 1, 1),
11657            shared_mem_bytes: 0,
11658        };
11659        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11660        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11661        let __s_b = self.gpu.stream();
11662        let mut b = __s_b.launch_builder(&f);
11663        b.arg(gate_bank)
11664            .arg(up_bank)
11665            .arg(sel)
11666            .arg(token_rows)
11667            .arg(x_bf16)
11668            .arg(gate_out)
11669            .arg(up_out)
11670            .arg(&inf)
11671            .arg(&outf)
11672            .arg(&ns)
11673            .arg(&rb)
11674            .arg(&es);
11675        unsafe {
11676            b.launch(cfg)?;
11677        }
11678        Ok(())
11679    }
11680
11681    /// Device-routed W4A16 gate+up over fixed token/slot rows. Selection ids remain global;
11682    /// each rank rejects non-owned slots and translates owned ids into its local expert bank.
11683    #[allow(clippy::too_many_arguments)]
11684    pub fn qmatvec_nvfp4_bf16_ep_dual_slots_into(
11685        &self,
11686        gate_bank: &CudaSlice<u8>,
11687        up_bank: &CudaSlice<u8>,
11688        sel: &CudaSlice<i32>,
11689        x_bf16: &CudaSlice<u8>,
11690        gate_out: &mut CudaSlice<f32>,
11691        up_out: &mut CudaSlice<f32>,
11692        n_pairs: usize,
11693        top_k: usize,
11694        in_f: usize,
11695        out_f: usize,
11696        owner_start: usize,
11697        owner_end: usize,
11698        row_bytes: usize,
11699        expert_stride: usize,
11700    ) -> Result<(), Box<dyn std::error::Error>> {
11701        let tokens = n_pairs.div_ceil(top_k);
11702        if top_k == 0
11703            || owner_start >= owner_end
11704            || !in_f.is_multiple_of(64)
11705            || sel.len() < n_pairs
11706            || gate_out.len() < n_pairs * out_f
11707            || up_out.len() < n_pairs * out_f
11708            || x_bf16.len() < 2 * in_f * tokens
11709        {
11710            return Err(format!(
11711                "W4A16 NVFP4 device EP dual-slot geometry sel={} x={} gate={} up={} \
11712                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11713                 owner={owner_start}..{owner_end}",
11714                sel.len(),
11715                x_bf16.len(),
11716                gate_out.len(),
11717                up_out.len(),
11718            )
11719            .into());
11720        }
11721        let pair_parallel = tokens > 1;
11722        let f = if pair_parallel {
11723            self.func("qmatvec_nvfp4_bf16_ep_quad_pairs")
11724        } else {
11725            self.func("qmatvec_nvfp4_bf16_ep_dual_slots")
11726        };
11727        let cfg = LaunchConfig {
11728            grid_dim: (
11729                if pair_parallel {
11730                    out_f.div_ceil(2) as u32
11731                } else {
11732                    (2 * out_f) as u32
11733                },
11734                if pair_parallel { n_pairs as u32 } else { 1 },
11735                1,
11736            ),
11737            block_dim: (256, 1, 1),
11738            shared_mem_bytes: 0,
11739        };
11740        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11741        let (os, oe) = (owner_start as i32, owner_end as i32);
11742        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11743        let __s_b = self.gpu.stream();
11744        let mut b = __s_b.launch_builder(&f);
11745        b.arg(gate_bank)
11746            .arg(up_bank)
11747            .arg(sel)
11748            .arg(x_bf16)
11749            .arg(gate_out)
11750            .arg(up_out)
11751            .arg(&inf)
11752            .arg(&outf)
11753            .arg(&np)
11754            .arg(&tk)
11755            .arg(&os)
11756            .arg(&oe)
11757            .arg(&rb)
11758            .arg(&es);
11759        unsafe {
11760            b.launch(cfg)?;
11761        }
11762        Ok(())
11763    }
11764
11765    /// Optional A8 t=1 gate+up program over global fixed slots.
11766    #[allow(clippy::too_many_arguments)]
11767    pub fn qmatvec_nvfp4_q8_ep_dual_slots_into(
11768        &self,
11769        gate_bank: &CudaSlice<u8>,
11770        up_bank: &CudaSlice<u8>,
11771        sel: &CudaSlice<i32>,
11772        aq: &CudaSlice<i8>,
11773        ad: &CudaSlice<f32>,
11774        gate_out: &mut CudaSlice<f32>,
11775        up_out: &mut CudaSlice<f32>,
11776        n_pairs: usize,
11777        top_k: usize,
11778        in_f: usize,
11779        out_f: usize,
11780        owner_start: usize,
11781        owner_end: usize,
11782        row_bytes: usize,
11783        expert_stride: usize,
11784    ) -> Result<(), Box<dyn std::error::Error>> {
11785        let tokens = n_pairs.div_ceil(top_k);
11786        if top_k == 0
11787            || owner_start >= owner_end
11788            || !in_f.is_multiple_of(64)
11789            || sel.len() < n_pairs
11790            || aq.len() < tokens * in_f
11791            || ad.len() < tokens * (in_f / 32)
11792            || gate_out.len() < n_pairs * out_f
11793            || up_out.len() < n_pairs * out_f
11794        {
11795            return Err(format!(
11796                "W4A8 NVFP4 device EP gate/up geometry sel={} aq={} ad={} gate={} up={} \
11797                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11798                 owner={owner_start}..{owner_end}",
11799                sel.len(),
11800                aq.len(),
11801                ad.len(),
11802                gate_out.len(),
11803                up_out.len(),
11804            )
11805            .into());
11806        }
11807        let f = self.func("qmatvec_nvfp4_q8_ep_dual_slots");
11808        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
11809        let cfg = LaunchConfig {
11810            grid_dim: (out_f as u32, 1, 1),
11811            block_dim: (threads, 1, 1),
11812            shared_mem_bytes: 0,
11813        };
11814        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11815        let (os, oe) = (owner_start as i32, owner_end as i32);
11816        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11817        let __s_b = self.gpu.stream();
11818        let mut b = __s_b.launch_builder(&f);
11819        b.arg(gate_bank)
11820            .arg(up_bank)
11821            .arg(sel)
11822            .arg(aq)
11823            .arg(ad)
11824            .arg(gate_out)
11825            .arg(up_out)
11826            .arg(&inf)
11827            .arg(&outf)
11828            .arg(&np)
11829            .arg(&tk)
11830            .arg(&os)
11831            .arg(&oe)
11832            .arg(&rb)
11833            .arg(&es);
11834        unsafe {
11835            b.launch(cfg)?;
11836        }
11837        Ok(())
11838    }
11839
11840    /// Known-good paired gate+up Q8 schedule: one CTA owns the same output row in both banks,
11841    /// shares the activation bytes, and retains one independent accumulator/reduction per bank.
11842    #[allow(clippy::too_many_arguments)]
11843    pub fn qmatvec_nvfp4_q8_ep_paired_slots_into(
11844        &self,
11845        gate_bank: &CudaSlice<u8>,
11846        up_bank: &CudaSlice<u8>,
11847        sel: &CudaSlice<i32>,
11848        aq: &CudaSlice<i8>,
11849        ad: &CudaSlice<f32>,
11850        gate_out: &mut CudaSlice<f32>,
11851        up_out: &mut CudaSlice<f32>,
11852        n_pairs: usize,
11853        top_k: usize,
11854        in_f: usize,
11855        out_f: usize,
11856        owner_start: usize,
11857        owner_end: usize,
11858        row_bytes: usize,
11859        expert_stride: usize,
11860    ) -> Result<(), Box<dyn std::error::Error>> {
11861        let tokens = n_pairs.div_ceil(top_k);
11862        if top_k == 0
11863            || owner_start >= owner_end
11864            || !in_f.is_multiple_of(64)
11865            || sel.len() < n_pairs
11866            || aq.len() < tokens * in_f
11867            || ad.len() < tokens * (in_f / 32)
11868            || gate_out.len() < n_pairs * out_f
11869            || up_out.len() < n_pairs * out_f
11870        {
11871            return Err(format!(
11872                "W4A8 NVFP4 paired gate/up geometry sel={} aq={} ad={} gate={} up={} \
11873                 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11874                 owner={owner_start}..{owner_end}",
11875                sel.len(),
11876                aq.len(),
11877                ad.len(),
11878                gate_out.len(),
11879                up_out.len(),
11880            )
11881            .into());
11882        }
11883        let f = self.func("qmatvec_nvfp4_q8_ep_paired_slots");
11884        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
11885        let cfg = LaunchConfig {
11886            grid_dim: (out_f as u32, 1, 1),
11887            block_dim: (threads, 1, 1),
11888            shared_mem_bytes: 0,
11889        };
11890        let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11891        let (os, oe) = (owner_start as i32, owner_end as i32);
11892        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11893        let __s_b = self.gpu.stream();
11894        let mut b = __s_b.launch_builder(&f);
11895        b.arg(gate_bank)
11896            .arg(up_bank)
11897            .arg(sel)
11898            .arg(aq)
11899            .arg(ad)
11900            .arg(gate_out)
11901            .arg(up_out)
11902            .arg(&inf)
11903            .arg(&outf)
11904            .arg(&np)
11905            .arg(&tk)
11906            .arg(&os)
11907            .arg(&oe)
11908            .arg(&rb)
11909            .arg(&es);
11910        unsafe {
11911            b.launch(cfg)?;
11912        }
11913        Ok(())
11914    }
11915
11916    /// W4A16 selected down rows scattered into canonical global pair positions on the root.
11917    #[allow(clippy::too_many_arguments)]
11918    pub fn qmatvec_nvfp4_bf16_sel_down_rows_raw(
11919        &self,
11920        bank: &CudaSlice<u8>,
11921        sel: &CudaSlice<i32>,
11922        global_pairs: &CudaSlice<i32>,
11923        activation_bf16: &CudaSlice<u8>,
11924        macros_down: &CudaSlice<f32>,
11925        dst_raw: u64,
11926        n_sel: usize,
11927        in_f: usize,
11928        out_f: usize,
11929        row_bytes: usize,
11930        expert_stride: usize,
11931        total_pairs: usize,
11932    ) -> Result<(), Box<dyn std::error::Error>> {
11933        if !in_f.is_multiple_of(64)
11934            || sel.len() < n_sel
11935            || global_pairs.len() < n_sel
11936            || activation_bf16.len() < 2 * n_sel * in_f
11937            || dst_raw == 0
11938        {
11939            return Err(format!(
11940                "W4A16 NVFP4 down rows geometry sel={} pairs={} act={} dst_raw={dst_raw:#x} \
11941                 n_sel={n_sel} total_pairs={total_pairs} in={in_f} out={out_f}",
11942                sel.len(),
11943                global_pairs.len(),
11944                activation_bf16.len(),
11945            )
11946            .into());
11947        }
11948        let f = self.func("qmatvec_nvfp4_bf16_sel_down_rows");
11949        let cfg = LaunchConfig {
11950            grid_dim: (out_f.div_ceil(2) as u32, n_sel as u32, 1),
11951            block_dim: (256, 1, 1),
11952            shared_mem_bytes: 0,
11953        };
11954        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11955        let (rb, es) = (row_bytes as i64, expert_stride as i64);
11956        let __s_b = self.gpu.stream();
11957        let mut b = __s_b.launch_builder(&f);
11958        b.arg(bank)
11959            .arg(sel)
11960            .arg(global_pairs)
11961            .arg(activation_bf16)
11962            .arg(macros_down)
11963            .arg(&dst_raw)
11964            .arg(&inf)
11965            .arg(&outf)
11966            .arg(&ns)
11967            .arg(&rb)
11968            .arg(&es);
11969        unsafe {
11970            b.launch(cfg)?;
11971        }
11972        Ok(())
11973    }
11974
11975    /// Device-routed W4A16 down rows. Exactly one owner rank writes each global token/slot row
11976    /// into the root device's peer-accessible slab.
11977    #[allow(clippy::too_many_arguments)]
11978    pub fn qmatvec_nvfp4_bf16_ep_down_slots_raw(
11979        &self,
11980        bank: &CudaSlice<u8>,
11981        sel: &CudaSlice<i32>,
11982        activation_bf16: &CudaSlice<u8>,
11983        macros_down: &CudaSlice<f32>,
11984        dst_raw: u64,
11985        n_pairs: usize,
11986        in_f: usize,
11987        out_f: usize,
11988        owner_start: usize,
11989        owner_end: usize,
11990        row_bytes: usize,
11991        expert_stride: usize,
11992    ) -> Result<(), Box<dyn std::error::Error>> {
11993        if owner_start >= owner_end
11994            || !in_f.is_multiple_of(64)
11995            || sel.len() < n_pairs
11996            || activation_bf16.len() < 2 * n_pairs * in_f
11997            || dst_raw == 0
11998        {
11999            return Err(format!(
12000                "W4A16 NVFP4 device EP down-slot geometry sel={} act={} dst_raw={dst_raw:#x} \
12001                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
12002                sel.len(),
12003                activation_bf16.len(),
12004            )
12005            .into());
12006        }
12007        let f = self.func("qmatvec_nvfp4_bf16_ep_down_slots");
12008        let cfg = LaunchConfig {
12009            grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
12010            block_dim: (256, 1, 1),
12011            shared_mem_bytes: 0,
12012        };
12013        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12014        let (os, oe) = (owner_start as i32, owner_end as i32);
12015        let (rb, es) = (row_bytes as i64, expert_stride as i64);
12016        let __s_b = self.gpu.stream();
12017        let mut b = __s_b.launch_builder(&f);
12018        b.arg(bank)
12019            .arg(sel)
12020            .arg(activation_bf16)
12021            .arg(macros_down)
12022            .arg(&dst_raw)
12023            .arg(&inf)
12024            .arg(&outf)
12025            .arg(&np)
12026            .arg(&os)
12027            .arg(&oe)
12028            .arg(&rb)
12029            .arg(&es);
12030        unsafe {
12031            b.launch(cfg)?;
12032        }
12033        Ok(())
12034    }
12035
12036    /// Pair-parallel multi-token twin of `qmatvec_nvfp4_bf16_ep_down_slots_raw`.
12037    #[allow(clippy::too_many_arguments)]
12038    pub fn qmatvec_nvfp4_bf16_ep_down_pairs_raw(
12039        &self,
12040        bank: &CudaSlice<u8>,
12041        sel: &CudaSlice<i32>,
12042        activation_bf16: &CudaSlice<u8>,
12043        macros_down: &CudaSlice<f32>,
12044        dst_raw: u64,
12045        n_pairs: usize,
12046        in_f: usize,
12047        out_f: usize,
12048        owner_start: usize,
12049        owner_end: usize,
12050        row_bytes: usize,
12051        expert_stride: usize,
12052    ) -> Result<(), Box<dyn std::error::Error>> {
12053        if owner_start >= owner_end
12054            || !in_f.is_multiple_of(64)
12055            || sel.len() < n_pairs
12056            || activation_bf16.len() < 2 * n_pairs * in_f
12057            || dst_raw == 0
12058        {
12059            return Err(format!(
12060                "W4A16 NVFP4 device EP down-pair geometry sel={} act={} dst_raw={dst_raw:#x} \
12061                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
12062                sel.len(),
12063                activation_bf16.len(),
12064            )
12065            .into());
12066        }
12067        let f = self.func("qmatvec_nvfp4_bf16_ep_down_pairs");
12068        let cfg = LaunchConfig {
12069            grid_dim: (out_f.div_ceil(2) as u32, n_pairs as u32, 1),
12070            block_dim: (256, 1, 1),
12071            shared_mem_bytes: 0,
12072        };
12073        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12074        let (os, oe) = (owner_start as i32, owner_end as i32);
12075        let (rb, es) = (row_bytes as i64, expert_stride as i64);
12076        let __s_b = self.gpu.stream();
12077        let mut b = __s_b.launch_builder(&f);
12078        b.arg(bank)
12079            .arg(sel)
12080            .arg(activation_bf16)
12081            .arg(macros_down)
12082            .arg(&dst_raw)
12083            .arg(&inf)
12084            .arg(&outf)
12085            .arg(&np)
12086            .arg(&os)
12087            .arg(&oe)
12088            .arg(&rb)
12089            .arg(&es);
12090        unsafe {
12091            b.launch(cfg)?;
12092        }
12093        Ok(())
12094    }
12095
12096    /// Host-expf W4A16 SwiGLU selected rows, rounded directly to BF16 for the down projection.
12097    #[allow(clippy::too_many_arguments)]
12098    pub fn silu_mul_scaled_host_expf_bf16_sel_into(
12099        &self,
12100        gate: &CudaSlice<f32>,
12101        up: &CudaSlice<f32>,
12102        gate_macros: &CudaSlice<f32>,
12103        up_macros: &CudaSlice<f32>,
12104        sel: &CudaSlice<i32>,
12105        limit: Option<f32>,
12106        output_bf16: &mut CudaSlice<u8>,
12107        n_per: usize,
12108        n_sel: usize,
12109    ) -> Result<(), Box<dyn std::error::Error>> {
12110        let n = n_per * n_sel;
12111        if sel.len() < n_sel || gate.len() < n || up.len() < n || output_bf16.len() < 2 * n {
12112            return Err(format!(
12113                "W4A16 selected activation geometry sel={} gate={} up={} out={} \
12114                 n_per={n_per} n_sel={n_sel}",
12115                sel.len(),
12116                gate.len(),
12117                up.len(),
12118                output_bf16.len(),
12119            )
12120            .into());
12121        }
12122        let (limit, has_limit) = match limit {
12123            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
12124            Some(limit) => {
12125                return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
12126            }
12127            None => (0.0f32, 0i32),
12128        };
12129        let f = self.func("silu_mul_scaled_host_expf_bf16_sel");
12130        let cfg = LaunchConfig::for_num_elems(n as u32);
12131        let (np, ns) = (n_per as i32, n_sel as i32);
12132        let __s_b = self.gpu.stream();
12133        let mut b = __s_b.launch_builder(&f);
12134        b.arg(gate)
12135            .arg(up)
12136            .arg(gate_macros)
12137            .arg(up_macros)
12138            .arg(sel)
12139            .arg(&limit)
12140            .arg(&has_limit)
12141            .arg(output_bf16)
12142            .arg(&np)
12143            .arg(&ns);
12144        unsafe {
12145            b.launch(cfg)?;
12146        }
12147        Ok(())
12148    }
12149
12150    /// Device-routed fixed token/slot W4A16 activation. Global expert ids are translated into
12151    /// rank-local macro rows only on the owning rank.
12152    #[allow(clippy::too_many_arguments)]
12153    pub fn silu_mul_scaled_host_expf_bf16_ep_slots_into(
12154        &self,
12155        gate: &CudaSlice<f32>,
12156        up: &CudaSlice<f32>,
12157        gate_macros: &CudaSlice<f32>,
12158        up_macros: &CudaSlice<f32>,
12159        sel: &CudaSlice<i32>,
12160        owner_start: usize,
12161        owner_end: usize,
12162        limit: Option<f32>,
12163        output_bf16: &mut CudaSlice<u8>,
12164        n_per: usize,
12165        n_pairs: usize,
12166    ) -> Result<(), Box<dyn std::error::Error>> {
12167        let n = n_per * n_pairs;
12168        if owner_start >= owner_end
12169            || sel.len() < n_pairs
12170            || gate.len() < n
12171            || up.len() < n
12172            || output_bf16.len() < 2 * n
12173        {
12174            return Err(format!(
12175                "W4A16 device EP activation geometry sel={} gate={} up={} out={} \
12176                 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
12177                sel.len(),
12178                gate.len(),
12179                up.len(),
12180                output_bf16.len(),
12181            )
12182            .into());
12183        }
12184        let (limit, has_limit) = match limit {
12185            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
12186            Some(limit) => {
12187                return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
12188            }
12189            None => (0.0f32, 0i32),
12190        };
12191        let f = self.func("silu_mul_scaled_host_expf_bf16_ep_slots");
12192        let cfg = LaunchConfig::for_num_elems(n as u32);
12193        let (np, pairs) = (n_per as i32, n_pairs as i32);
12194        let (os, oe) = (owner_start as i32, owner_end as i32);
12195        let __s_b = self.gpu.stream();
12196        let mut b = __s_b.launch_builder(&f);
12197        b.arg(gate)
12198            .arg(up)
12199            .arg(gate_macros)
12200            .arg(up_macros)
12201            .arg(sel)
12202            .arg(&limit)
12203            .arg(&has_limit)
12204            .arg(output_bf16)
12205            .arg(&np)
12206            .arg(&pairs)
12207            .arg(&os)
12208            .arg(&oe);
12209        unsafe {
12210            b.launch(cfg)?;
12211        }
12212        Ok(())
12213    }
12214
12215    /// Optional A8 host-expf SwiGLU over global fixed slots.
12216    #[allow(clippy::too_many_arguments)]
12217    pub fn silu_mul_scaled_host_expf_q8_ep_slots_into(
12218        &self,
12219        gate: &CudaSlice<f32>,
12220        up: &CudaSlice<f32>,
12221        gate_macros: &CudaSlice<f32>,
12222        up_macros: &CudaSlice<f32>,
12223        sel: &CudaSlice<i32>,
12224        owner_start: usize,
12225        owner_end: usize,
12226        limit: Option<f32>,
12227        output_q8: &mut CudaSlice<i8>,
12228        output_scales: &mut CudaSlice<f32>,
12229        n_per: usize,
12230        n_pairs: usize,
12231    ) -> Result<(), Box<dyn std::error::Error>> {
12232        let n = n_per * n_pairs;
12233        if owner_start >= owner_end
12234            || !n_per.is_multiple_of(32)
12235            || sel.len() < n_pairs
12236            || gate.len() < n
12237            || up.len() < n
12238            || output_q8.len() < n
12239            || output_scales.len() < n / 32
12240        {
12241            return Err(format!(
12242                "W4A8 device EP activation geometry sel={} gate={} up={} q8={} scales={} \
12243                 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
12244                sel.len(),
12245                gate.len(),
12246                up.len(),
12247                output_q8.len(),
12248                output_scales.len(),
12249            )
12250            .into());
12251        }
12252        let (limit, has_limit) = match limit {
12253            Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
12254            Some(limit) => {
12255                return Err(format!("W4A8 selected activation limit {limit} is invalid").into());
12256            }
12257            None => (0.0f32, 0i32),
12258        };
12259        let f = self.func("silu_mul_scaled_host_expf_q8_ep_slots");
12260        let warps = n / 32;
12261        let cfg = LaunchConfig {
12262            grid_dim: ((warps as u32).div_ceil(4), 1, 1),
12263            block_dim: (128, 1, 1),
12264            shared_mem_bytes: 0,
12265        };
12266        let (np, pairs) = (n_per as i32, n_pairs as i32);
12267        let (os, oe) = (owner_start as i32, owner_end as i32);
12268        let __s_b = self.gpu.stream();
12269        let mut b = __s_b.launch_builder(&f);
12270        b.arg(gate)
12271            .arg(up)
12272            .arg(gate_macros)
12273            .arg(up_macros)
12274            .arg(sel)
12275            .arg(&limit)
12276            .arg(&has_limit)
12277            .arg(output_q8)
12278            .arg(output_scales)
12279            .arg(&np)
12280            .arg(&pairs)
12281            .arg(&os)
12282            .arg(&oe);
12283        unsafe {
12284            b.launch(cfg)?;
12285        }
12286        Ok(())
12287    }
12288
12289    /// W4A16 selected-expert down projection plus owner-local route combine. The destination may
12290    /// reside in the model engine's peer-accessible root pool.
12291    #[allow(clippy::too_many_arguments)]
12292    pub fn qmatvec_nvfp4_bf16_sel_down_fma_into(
12293        &self,
12294        bank: &CudaSlice<u8>,
12295        sel: &CudaSlice<i32>,
12296        activation_bf16: &CudaSlice<u8>,
12297        route_weights: &CudaSlice<f32>,
12298        macros_down: &CudaSlice<f32>,
12299        dst: &mut cudarc::driver::CudaViewMut<f32>,
12300        n_sel: usize,
12301        in_f: usize,
12302        out_f: usize,
12303        row_bytes: usize,
12304        expert_stride: usize,
12305    ) -> Result<(), Box<dyn std::error::Error>> {
12306        if !in_f.is_multiple_of(64)
12307            || sel.len() < n_sel
12308            || route_weights.len() < n_sel
12309            || activation_bf16.len() < 2 * n_sel * in_f
12310            || dst.len() < out_f
12311        {
12312            return Err(format!(
12313                "W4A16 NVFP4 down selected geometry sel={} act={} weights={} dst={} \
12314                 n_sel={n_sel} in={in_f} out={out_f}",
12315                sel.len(),
12316                activation_bf16.len(),
12317                route_weights.len(),
12318                dst.len(),
12319            )
12320            .into());
12321        }
12322        let f = self.func("qmatvec_nvfp4_bf16_sel_down_fma");
12323        let cfg = LaunchConfig {
12324            grid_dim: (out_f as u32, 1, 1),
12325            block_dim: (256, 1, 1),
12326            shared_mem_bytes: 0,
12327        };
12328        let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
12329        let (rb, es) = (row_bytes as i64, expert_stride as i64);
12330        let __s_b = self.gpu.stream();
12331        let mut b = __s_b.launch_builder(&f);
12332        b.arg(bank)
12333            .arg(sel)
12334            .arg(activation_bf16)
12335            .arg(route_weights)
12336            .arg(macros_down)
12337            .arg(dst)
12338            .arg(&inf)
12339            .arg(&outf)
12340            .arg(&ns)
12341            .arg(&rb)
12342            .arg(&es);
12343        unsafe {
12344            b.launch(cfg)?;
12345        }
12346        Ok(())
12347    }
12348
12349    /// Device-routed t=1 W4A16 down projection plus owner-local weighted combine.
12350    #[allow(clippy::too_many_arguments)]
12351    pub fn qmatvec_nvfp4_bf16_ep_down_fma_into(
12352        &self,
12353        bank: &CudaSlice<u8>,
12354        sel: &CudaSlice<i32>,
12355        activation_bf16: &CudaSlice<u8>,
12356        route_weights: &CudaSlice<f32>,
12357        macros_down: &CudaSlice<f32>,
12358        dst: &mut cudarc::driver::CudaViewMut<f32>,
12359        n_pairs: usize,
12360        in_f: usize,
12361        out_f: usize,
12362        owner_start: usize,
12363        owner_end: usize,
12364        row_bytes: usize,
12365        expert_stride: usize,
12366    ) -> Result<(), Box<dyn std::error::Error>> {
12367        if owner_start >= owner_end
12368            || !in_f.is_multiple_of(64)
12369            || sel.len() < n_pairs
12370            || route_weights.len() < n_pairs
12371            || activation_bf16.len() < 2 * n_pairs * in_f
12372            || dst.len() < out_f
12373        {
12374            return Err(format!(
12375                "W4A16 device EP down-FMA geometry sel={} act={} weights={} dst={} \
12376                 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
12377                sel.len(),
12378                activation_bf16.len(),
12379                route_weights.len(),
12380                dst.len(),
12381            )
12382            .into());
12383        }
12384        let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
12385        let cfg = LaunchConfig {
12386            grid_dim: (out_f as u32, 1, 1),
12387            block_dim: (256, 1, 1),
12388            shared_mem_bytes: 0,
12389        };
12390        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12391        let (os, oe) = (owner_start as i32, owner_end as i32);
12392        let (rb, es) = (row_bytes as i64, expert_stride as i64);
12393        let __s_b = self.gpu.stream();
12394        let mut b = __s_b.launch_builder(&f);
12395        b.arg(bank)
12396            .arg(sel)
12397            .arg(activation_bf16)
12398            .arg(route_weights)
12399            .arg(macros_down)
12400            .arg(dst)
12401            .arg(&inf)
12402            .arg(&outf)
12403            .arg(&np)
12404            .arg(&os)
12405            .arg(&oe)
12406            .arg(&rb)
12407            .arg(&es);
12408        unsafe {
12409            b.launch(cfg)?;
12410        }
12411        Ok(())
12412    }
12413
12414    /// Capture-safe twin of `qmatvec_nvfp4_bf16_ep_down_fma_into`. The destination is one
12415    /// rank-owned row inside a persistent root-device slab.
12416    #[allow(clippy::too_many_arguments)]
12417    pub fn qmatvec_nvfp4_bf16_ep_down_fma_raw(
12418        &self,
12419        bank: &CudaSlice<u8>,
12420        sel: &CudaSlice<i32>,
12421        activation_bf16: &CudaSlice<u8>,
12422        route_weights: &CudaSlice<f32>,
12423        macros_down: &CudaSlice<f32>,
12424        dst_raw: u64,
12425        n_pairs: usize,
12426        in_f: usize,
12427        out_f: usize,
12428        owner_start: usize,
12429        owner_end: usize,
12430        row_bytes: usize,
12431        expert_stride: usize,
12432    ) -> Result<(), Box<dyn std::error::Error>> {
12433        if dst_raw == 0
12434            || owner_start >= owner_end
12435            || !in_f.is_multiple_of(64)
12436            || sel.len() < n_pairs
12437            || route_weights.len() < n_pairs
12438            || activation_bf16.len() < 2 * n_pairs * in_f
12439        {
12440            return Err(format!(
12441                "W4A16 device EP raw down-FMA geometry sel={} act={} weights={} \
12442                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12443                 owner={owner_start}..{owner_end}",
12444                sel.len(),
12445                activation_bf16.len(),
12446                route_weights.len(),
12447            )
12448            .into());
12449        }
12450        let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
12451        let cfg = LaunchConfig {
12452            grid_dim: (out_f as u32, 1, 1),
12453            block_dim: (256, 1, 1),
12454            shared_mem_bytes: 0,
12455        };
12456        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12457        let (os, oe) = (owner_start as i32, owner_end as i32);
12458        let (rb, es) = (row_bytes as i64, expert_stride as i64);
12459        let __s_b = self.gpu.stream();
12460        let mut b = __s_b.launch_builder(&f);
12461        b.arg(bank)
12462            .arg(sel)
12463            .arg(activation_bf16)
12464            .arg(route_weights)
12465            .arg(macros_down)
12466            .arg(&dst_raw)
12467            .arg(&inf)
12468            .arg(&outf)
12469            .arg(&np)
12470            .arg(&os)
12471            .arg(&oe)
12472            .arg(&rb)
12473            .arg(&es);
12474        unsafe {
12475            b.launch(cfg)?;
12476        }
12477        Ok(())
12478    }
12479
12480    /// Optional A8 fixed-slot down rows. Each owner rank writes its selected pair rows directly
12481    /// into the root slot slab; the root applies route weights in canonical token/slot order.
12482    #[allow(clippy::too_many_arguments)]
12483    pub fn qmatvec_nvfp4_q8_ep_down_slots_raw(
12484        &self,
12485        bank: &CudaSlice<u8>,
12486        sel: &CudaSlice<i32>,
12487        aq: &CudaSlice<i8>,
12488        ad: &CudaSlice<f32>,
12489        macros_down: &CudaSlice<f32>,
12490        dst_raw: u64,
12491        n_pairs: usize,
12492        in_f: usize,
12493        out_f: usize,
12494        owner_start: usize,
12495        owner_end: usize,
12496        row_bytes: usize,
12497        expert_stride: usize,
12498    ) -> Result<(), Box<dyn std::error::Error>> {
12499        if dst_raw == 0
12500            || owner_start >= owner_end
12501            || !in_f.is_multiple_of(64)
12502            || sel.len() < n_pairs
12503            || aq.len() < n_pairs * in_f
12504            || ad.len() < n_pairs * (in_f / 32)
12505        {
12506            return Err(format!(
12507                "W4A8 device EP raw down-slot geometry sel={} aq={} ad={} \
12508                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12509                 owner={owner_start}..{owner_end}",
12510                sel.len(),
12511                aq.len(),
12512                ad.len(),
12513            )
12514            .into());
12515        }
12516        let f = self.func("qmatvec_nvfp4_q8_ep_down_slots");
12517        let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
12518        let cfg = LaunchConfig {
12519            grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
12520            block_dim: (threads, 1, 1),
12521            shared_mem_bytes: 0,
12522        };
12523        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12524        let (os, oe) = (owner_start as i32, owner_end as i32);
12525        let (rb, es) = (row_bytes as i64, expert_stride as i64);
12526        let __s_b = self.gpu.stream();
12527        let mut b = __s_b.launch_builder(&f);
12528        b.arg(bank)
12529            .arg(sel)
12530            .arg(aq)
12531            .arg(ad)
12532            .arg(macros_down)
12533            .arg(&dst_raw)
12534            .arg(&inf)
12535            .arg(&outf)
12536            .arg(&np)
12537            .arg(&os)
12538            .arg(&oe)
12539            .arg(&rb)
12540            .arg(&es);
12541        unsafe {
12542            b.launch(cfg)?;
12543        }
12544        Ok(())
12545    }
12546
12547    /// Historical A8 t=1 down + owner-local route combine into a persistent root row.
12548    #[allow(clippy::too_many_arguments)]
12549    pub fn qmatvec_nvfp4_q8_ep_down_fma_raw(
12550        &self,
12551        bank: &CudaSlice<u8>,
12552        sel: &CudaSlice<i32>,
12553        aq: &CudaSlice<i8>,
12554        ad: &CudaSlice<f32>,
12555        route_weights: &CudaSlice<f32>,
12556        macros_down: &CudaSlice<f32>,
12557        dst_raw: u64,
12558        n_pairs: usize,
12559        in_f: usize,
12560        out_f: usize,
12561        owner_start: usize,
12562        owner_end: usize,
12563        row_bytes: usize,
12564        expert_stride: usize,
12565    ) -> Result<(), Box<dyn std::error::Error>> {
12566        if dst_raw == 0
12567            || owner_start >= owner_end
12568            || !in_f.is_multiple_of(64)
12569            || sel.len() < n_pairs
12570            || aq.len() < n_pairs * in_f
12571            || ad.len() < n_pairs * (in_f / 32)
12572            || route_weights.len() < n_pairs
12573        {
12574            return Err(format!(
12575                "W4A8 device EP raw down-FMA geometry sel={} aq={} ad={} weights={} \
12576                 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12577                 owner={owner_start}..{owner_end}",
12578                sel.len(),
12579                aq.len(),
12580                ad.len(),
12581                route_weights.len(),
12582            )
12583            .into());
12584        }
12585        let f = self.func("qmatvec_nvfp4_q8_ep_down_fma");
12586        let cfg = LaunchConfig {
12587            grid_dim: (out_f as u32, 1, 1),
12588            block_dim: (256, 1, 1),
12589            shared_mem_bytes: 0,
12590        };
12591        let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12592        let (os, oe) = (owner_start as i32, owner_end as i32);
12593        let (rb, es) = (row_bytes as i64, expert_stride as i64);
12594        let __s_b = self.gpu.stream();
12595        let mut b = __s_b.launch_builder(&f);
12596        b.arg(bank)
12597            .arg(sel)
12598            .arg(aq)
12599            .arg(ad)
12600            .arg(route_weights)
12601            .arg(macros_down)
12602            .arg(&dst_raw)
12603            .arg(&inf)
12604            .arg(&outf)
12605            .arg(&np)
12606            .arg(&os)
12607            .arg(&oe)
12608            .arg(&rb)
12609            .arg(&es);
12610        unsafe {
12611            b.launch(cfg)?;
12612        }
12613        Ok(())
12614    }
12615
12616    /// Selected-experts batched twin of `silu_mul_scaled_q8_1`: [n_sel, n_per] rows, macros
12617    /// from device arrays indexed via sel. Per expert row bit-identical to the scalar kernel.
12618    /// `limit` = the step35 routed SwiGLU clamp (min(silu, limit) * clamp(up, +-limit)); None
12619    /// takes the plain SiLU kernel.
12620    #[allow(clippy::too_many_arguments)]
12621    pub fn silu_mul_scaled_q8_1_sel_into(
12622        &self,
12623        gate: &CudaSlice<f32>,
12624        up: &CudaSlice<f32>,
12625        gmac: &CudaSlice<f32>,
12626        umac: &CudaSlice<f32>,
12627        sel: &CudaSlice<i32>,
12628        limit: Option<f32>,
12629        out_q: &mut CudaSlice<i8>,
12630        out_d: &mut CudaSlice<f32>,
12631        n_per: usize,
12632        n_sel: usize,
12633    ) -> Result<(), Box<dyn std::error::Error>> {
12634        let n = n_per * n_sel;
12635        if !n_per.is_multiple_of(32) || out_q.len() < n || out_d.len() < n / 32 {
12636            return Err(format!(
12637                "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
12638                out_q.len(),
12639                out_d.len()
12640            )
12641            .into());
12642        }
12643        if let Some(limit) = limit {
12644            if limit <= 1e-6 {
12645                return Err(format!(
12646                    "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
12647                )
12648                .into());
12649            }
12650            let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
12651            let cfg = LaunchConfig::for_num_elems(n as u32);
12652            let (np, ns) = (n_per as i32, n_sel as i32);
12653            let __s_b = self.gpu.stream();
12654            let mut b = __s_b.launch_builder(&f);
12655            b.arg(gate)
12656                .arg(up)
12657                .arg(gmac)
12658                .arg(umac)
12659                .arg(sel)
12660                .arg(&limit)
12661                .arg(out_q)
12662                .arg(out_d)
12663                .arg(&np)
12664                .arg(&ns);
12665            unsafe {
12666                b.launch(cfg)?;
12667            }
12668            return Ok(());
12669        }
12670        let f = self.func("silu_mul_scaled_q8_1_sel");
12671        let cfg = LaunchConfig::for_num_elems(n as u32);
12672        let (np, ns) = (n_per as i32, n_sel as i32);
12673        let __s_b = self.gpu.stream();
12674        let mut b = __s_b.launch_builder(&f);
12675        b.arg(gate)
12676            .arg(up)
12677            .arg(gmac)
12678            .arg(umac)
12679            .arg(sel)
12680            .arg(out_q)
12681            .arg(out_d)
12682            .arg(&np)
12683            .arg(&ns);
12684        unsafe {
12685            b.launch(cfg)?;
12686        }
12687        Ok(())
12688    }
12689
12690    #[track_caller]
12691    pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12692        crate::alloc_trace_hit(v.len() * 4);
12693        Ok(self.gpu.stream().clone_htod(v)?)
12694    }
12695    #[track_caller]
12696    pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
12697        crate::alloc_trace_hit(v.len() * 4);
12698        Ok(self.gpu.stream().clone_htod(v)?)
12699    }
12700    /// i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
12701    #[track_caller]
12702    pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
12703        crate::alloc_trace_hit(v.len());
12704        Ok(self.gpu.stream().clone_htod(v)?)
12705    }
12706    #[track_caller]
12707    pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
12708        crate::alloc_trace_hit(v.len() * 8);
12709        Ok(self.gpu.stream().clone_htod(v)?)
12710    }
12711    /// View twin of `dtoh` (lean-logits component 3: D2H one row of a [B, n_vocab] stack).
12712    #[track_caller]
12713    pub fn dtoh_view(
12714        &self,
12715        d: &cudarc::driver::CudaView<f32>,
12716    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
12717        crate::dtoh_trace_hit(d.len() * 4);
12718        let v = self.gpu.stream().clone_dtoh(d)?;
12719        self.gpu.stream().synchronize()?;
12720        Ok(v)
12721    }
12722    #[track_caller]
12723    pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
12724        crate::dtoh_trace_hit(d.len() * 4);
12725        let v = self.gpu.stream().clone_dtoh(d)?;
12726        self.gpu.stream().synchronize()?;
12727        Ok(v)
12728    }
12729    /// Queue two f32 device-to-host copies on the compute stream, then establish one host
12730    /// boundary for both. Hy3's CPU/GPU expert split needs the router logits and the MoE input;
12731    /// issuing them together avoids a second stream synchronization in every trunk layer.
12732    #[track_caller]
12733    pub fn dtoh_pair(
12734        &self,
12735        a: &CudaSlice<f32>,
12736        b: &CudaSlice<f32>,
12737    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12738        crate::dtoh_trace_hit((a.len() + b.len()) * 4);
12739        let av = self.gpu.stream().clone_dtoh(a)?;
12740        let bv = self.gpu.stream().clone_dtoh(b)?;
12741        self.gpu.stream().synchronize()?;
12742        Ok((av, bv))
12743    }
12744    /// View-scoped twin of `dtoh_pair` for reusable capacity buffers whose inactive tail must not
12745    /// cross a shape-sensitive host boundary.
12746    #[track_caller]
12747    pub fn dtoh_pair_views(
12748        &self,
12749        a: &cudarc::driver::CudaView<f32>,
12750        b: &cudarc::driver::CudaView<f32>,
12751    ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12752        crate::dtoh_trace_hit((a.len() + b.len()) * 4);
12753        let av = self.gpu.stream().clone_dtoh(a)?;
12754        let bv = self.gpu.stream().clone_dtoh(b)?;
12755        self.gpu.stream().synchronize()?;
12756        Ok((av, bv))
12757    }
12758    /// Device-to-host copy of an i32 buffer (fused-router sel_idx readback).
12759    #[track_caller]
12760    pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
12761        crate::dtoh_trace_hit(d.len() * 4);
12762        let v = self.gpu.stream().clone_dtoh(d)?;
12763        self.gpu.stream().synchronize()?;
12764        Ok(v)
12765    }
12766    /// Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).
12767    /// i8 twin of [`Self::dtoh_u8`], for the decode-graph trace's q8_1 activation checksum
12768    /// (`MEMRA_GLM5_GRAPH_TRACE`). Gate harness only: it synchronizes.
12769    #[track_caller]
12770    pub fn dtoh_i8(&self, d: &CudaSlice<i8>) -> Result<Vec<i8>, Box<dyn std::error::Error>> {
12771        crate::dtoh_trace_hit(d.len());
12772        let v = self.gpu.stream().clone_dtoh(d)?;
12773        self.gpu.stream().synchronize()?;
12774        Ok(v)
12775    }
12776    #[track_caller]
12777    pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
12778        crate::dtoh_trace_hit(d.len());
12779        let v = self.gpu.stream().clone_dtoh(d)?;
12780        self.gpu.stream().synchronize()?;
12781        Ok(v)
12782    }
12783    #[track_caller]
12784    pub fn dtoh_u8_view(
12785        &self,
12786        d: &cudarc::driver::CudaView<u8>,
12787    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
12788        crate::dtoh_trace_hit(d.len());
12789        let v = self.gpu.stream().clone_dtoh(d)?;
12790        self.gpu.stream().synchronize()?;
12791        Ok(v)
12792    }
12793    /// D2H copy of the first `n` bytes of `d` into a pinned CACHEABLE host buffer: the
12794    /// prefix-cache host-tier demote primitive (lane/kv-host-spill-20260830). Queued on the
12795    /// worker stream and synchronized before returning, exactly like `dtoh_u8`: v1 keeps every
12796    /// host-tier copy on the CUDA owner thread (the HY3 spill law). SEAM (named, not built): an
12797    /// overlapped copy-stream variant would queue this on a dedicated D2H stream with an event
12798    /// handshake against the compute stream; build it only with a tick-stall receipt that says
12799    /// the sync copy is the bottleneck.
12800    #[track_caller]
12801    pub fn dtoh_u8_into_pinned(
12802        &self,
12803        d: &CudaSlice<u8>,
12804        dst: &mut PinnedHostBuf,
12805        n: usize,
12806    ) -> Result<(), Box<dyn std::error::Error>> {
12807        crate::dtoh_trace_hit(d.len());
12808        if n > d.len() || n > dst.len() {
12809            return Err(format!(
12810                "dtoh_u8_into_pinned range {n} exceeds src {} or pinned dst {}",
12811                d.len(),
12812                dst.len(),
12813            )
12814            .into());
12815        }
12816        if n == 0 {
12817            return Ok(());
12818        }
12819        let host = &mut dst.as_mut_slice()[..n];
12820        self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
12821        self.gpu.stream().synchronize()?;
12822        Ok(())
12823    }
12824    /// f32 twin of [`Self::dtoh_u8_into_pinned`] (lane/spec-route-depth-20260902): D2H the
12825    /// first `n` floats of `d` into a pinned CACHEABLE host buffer, synchronized before
12826    /// returning (the caller CPU-reads the rows right after).
12827    #[track_caller]
12828    pub fn dtoh_f32_into_pinned(
12829        &self,
12830        d: &CudaSlice<f32>,
12831        dst: &mut PinnedHostBuf,
12832        n: usize,
12833    ) -> Result<(), Box<dyn std::error::Error>> {
12834        crate::dtoh_trace_hit(d.len() * 4);
12835        if n > d.len() || n * std::mem::size_of::<f32>() > dst.len() {
12836            return Err(format!(
12837                "dtoh_f32_into_pinned range {n} floats exceeds src {} or pinned dst {} bytes",
12838                d.len(),
12839                dst.len(),
12840            )
12841            .into());
12842        }
12843        if n == 0 {
12844            return Ok(());
12845        }
12846        // SAFETY: the pinned buffer is page-aligned (malloc_host) and holds >= n f32s.
12847        let host: &mut [f32] = unsafe {
12848            std::slice::from_raw_parts_mut(dst.as_mut_slice().as_mut_ptr() as *mut f32, n)
12849        };
12850        self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
12851        self.gpu.stream().synchronize()?;
12852        Ok(())
12853    }
12854
12855    /// ASYNC H2D of the first `n` floats of a pinned host buffer into a fresh device slice
12856    /// (lane/spec-route-depth-20260902: the chunked drafter prime's tap upload). Queued on
12857    /// the worker stream and NOT synchronized: the copy is DMA from page-locked memory, and
12858    /// every consumer is stream-ordered behind it. CONTRACT: the caller must not write
12859    /// `src` again until the stream has passed this copy (synchronize, or a later blocking
12860    /// readback on the same stream) — the chunked prime synchronizes at the end of each
12861    /// chunk's ingest before it refills the staging buffer.
12862    pub fn htod_f32_from_pinned_async(
12863        &self,
12864        src: &PinnedHostBuf,
12865        n: usize,
12866    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12867        if n * std::mem::size_of::<f32>() > src.len() {
12868            return Err(format!(
12869                "htod_f32_from_pinned_async range {n} floats exceeds pinned src {} bytes",
12870                src.len(),
12871            )
12872            .into());
12873        }
12874        let mut d = self.uninit(n)?;
12875        if n == 0 {
12876            return Ok(d);
12877        }
12878        // SAFETY: page-aligned pinned allocation with >= n f32s (checked above).
12879        let host: &[f32] =
12880            unsafe { std::slice::from_raw_parts(src.as_slice().as_ptr() as *const f32, n) };
12881        let stream = self.gpu.stream();
12882        let s: &CudaStream = &stream;
12883        {
12884            let (pd, _guard) = d.device_ptr_mut(s);
12885            // SAFETY: `pd` is a live device allocation of n f32s on this stream; `host` is
12886            // page-locked memory the caller keeps alive and unmodified until the stream
12887            // passes the copy (the documented contract).
12888            unsafe { cudarc::driver::result::memcpy_htod_async(pd, host, s.cu_stream())? };
12889        }
12890        Ok(d)
12891    }
12892
12893    /// 2D device-to-device copy (`cuMemcpy2DAsync`): `height` rows of `width_floats` floats
12894    /// from `src` (row pitch `src_pitch_floats`) into `dst` at float offset `dst_off_floats`
12895    /// (row pitch `dst_pitch_floats`), queued on the worker stream. The strided-scatter
12896    /// primitive the device-resident tap ingest uses to interleave per-slot tap planes into
12897    /// the drafter fc layout with no host bounce (lane/spec-route-depth-20260902). Same-
12898    /// device only; both slices must live on this engine's device.
12899    #[allow(clippy::too_many_arguments)]
12900    // allow: the parameter list IS the 2D copy descriptor (dst, dst offset, dst pitch,
12901    // src, src pitch, width, height); a struct would only rename the same seven fields
12902    pub fn copy_2d_dtod_async(
12903        &self,
12904        dst: &mut CudaSlice<f32>,
12905        dst_off_floats: usize,
12906        dst_pitch_floats: usize,
12907        src: &CudaSlice<f32>,
12908        src_pitch_floats: usize,
12909        width_floats: usize,
12910        height: usize,
12911    ) -> Result<(), Box<dyn std::error::Error>> {
12912        if height == 0 || width_floats == 0 {
12913            return Ok(());
12914        }
12915        let f = std::mem::size_of::<f32>();
12916        let dst_need = dst_off_floats + (height - 1) * dst_pitch_floats + width_floats;
12917        let src_need = (height - 1) * src_pitch_floats + width_floats;
12918        if width_floats > src_pitch_floats
12919            || width_floats > dst_pitch_floats
12920            || dst_need > dst.len()
12921            || src_need > src.len()
12922        {
12923            return Err(format!(
12924                "copy_2d_dtod_async out of bounds: dst needs {dst_need} of {}, src needs \
12925                 {src_need} of {}, width {width_floats} pitches {src_pitch_floats}/{dst_pitch_floats}",
12926                dst.len(),
12927                src.len(),
12928            )
12929            .into());
12930        }
12931        let stream = self.gpu.stream();
12932        let s: &CudaStream = &stream;
12933        let (sp, _g0) = src.device_ptr(s);
12934        let (dp, _g1) = dst.device_ptr_mut(s);
12935        let copy = cudarc::driver::sys::CUDA_MEMCPY2D_st {
12936            srcXInBytes: 0,
12937            srcY: 0,
12938            srcMemoryType: cudarc::driver::sys::CUmemorytype_enum::CU_MEMORYTYPE_DEVICE,
12939            srcHost: std::ptr::null(),
12940            srcDevice: sp,
12941            srcArray: std::ptr::null_mut(),
12942            srcPitch: src_pitch_floats * f,
12943            dstXInBytes: dst_off_floats * f,
12944            dstY: 0,
12945            dstMemoryType: cudarc::driver::sys::CUmemorytype_enum::CU_MEMORYTYPE_DEVICE,
12946            dstHost: std::ptr::null_mut(),
12947            dstDevice: dp,
12948            dstArray: std::ptr::null_mut(),
12949            dstPitch: dst_pitch_floats * f,
12950            WidthInBytes: width_floats * f,
12951            Height: height,
12952        };
12953        // SAFETY: both pointers are live device allocations on this engine's device, bounds
12954        // checked above; the copy is stream-ordered on the worker stream.
12955        unsafe { cudarc::driver::sys::cuMemcpy2DAsync_v2(&copy, s.cu_stream()).result()? };
12956        Ok(())
12957    }
12958
12959    /// Cross-device copy of `n` floats from `src` (on `src_engine`'s device) into `dst` (on
12960    /// this engine's device): `cudaMemcpyPeerAsync` with explicit contexts, issued on the
12961    /// SOURCE engine's stream (the pp.rs boundary transport's shape) so later writes on
12962    /// that stream are ordered behind it. NOT synchronized: the caller drains the source
12963    /// stream (or waits an event recorded on it) before consuming `dst` on this engine.
12964    /// Rebinds this engine's context on the calling thread before returning.
12965    pub fn copy_peer_from_async(
12966        &self,
12967        dst: &mut CudaSlice<f32>,
12968        src_engine: &Engine,
12969        src: &CudaSlice<f32>,
12970        n: usize,
12971    ) -> Result<(), Box<dyn std::error::Error>> {
12972        if n > src.len() || n > dst.len() {
12973            return Err(format!(
12974                "copy_peer_from_async range {n} exceeds src {} or dst {}",
12975                src.len(),
12976                dst.len(),
12977            )
12978            .into());
12979        }
12980        if n == 0 {
12981            return Ok(());
12982        }
12983        let stream = src_engine.gpu.stream();
12984        let s_src: &CudaStream = &stream;
12985        let (sp, _g0) = src.device_ptr(s_src);
12986        let (dp, _g1) = dst.device_ptr_mut(s_src);
12987        src_engine.ctx().bind_to_thread()?;
12988        // SAFETY: live allocations on the two named contexts, bounds checked above; the
12989        // copy is queued on the source stream with explicit src/dst contexts.
12990        let r = unsafe {
12991            cudarc::driver::result::memcpy_peer_async(
12992                self.ctx().cu_ctx(),
12993                dp,
12994                src_engine.ctx().cu_ctx(),
12995                sp,
12996                n * std::mem::size_of::<f32>(),
12997                s_src.cu_stream(),
12998            )
12999        };
13000        self.ctx().bind_to_thread()?;
13001        r?;
13002        Ok(())
13003    }
13004
13005    /// Free device memory on this engine's device, in MB (`cuMemGetInfo`; 0 on error).
13006    pub fn free_mem_mb(&self) -> u64 {
13007        self.ctx()
13008            .mem_get_info()
13009            .map(|(f, _)| f as u64 >> 20)
13010            .unwrap_or(0)
13011    }
13012
13013    #[track_caller]
13014    pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13015        crate::alloc_trace_hit(n * 4);
13016        SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13017        let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
13018        self.keep_if_capturing(&s);
13019        Ok(s)
13020    }
13021
13022    /// Take the pooled hc-glue decode workspace (MEMRA_HC_DECODE_WS) for one step's walk; put
13023    /// it back with [`Self::hyper_ws_put`]. A `None` here means another walk holds it (or it
13024    /// was never built) — the caller allocates fresh, which is always correct.
13025    pub(crate) fn hyper_ws_take(&self) -> Option<crate::hyper::HyperDecodeWs> {
13026        self.hyper_decode_ws.lock().unwrap().take()
13027    }
13028
13029    pub(crate) fn hyper_ws_put(&self, ws: crate::hyper::HyperDecodeWs) {
13030        *self.hyper_decode_ws.lock().unwrap() = Some(ws);
13031    }
13032
13033    /// Take the session's MLA segment workspace, creating it at this geometry on first use;
13034    /// `None` when the geometry differs from the held set (the caller then runs the owned path).
13035    pub(crate) fn mla_seg_ws_take(
13036        &self,
13037        nh: usize,
13038        dn: usize,
13039        dr: usize,
13040        r: usize,
13041        q_lora: usize,
13042    ) -> Result<crate::hybrid_forward::MlaSegWs, Box<dyn std::error::Error>> {
13043        let held = self.mla_seg_ws.lock().unwrap().take();
13044        match held {
13045            Some(ws) if ws.sig == (nh, dn, dr, r, q_lora) => Ok(ws),
13046            _ => crate::hybrid_forward::MlaSegWs::new(self, nh, dn, dr, r, q_lora),
13047        }
13048    }
13049
13050    pub(crate) fn mla_seg_ws_put(&self, ws: crate::hybrid_forward::MlaSegWs) {
13051        *self.mla_seg_ws.lock().unwrap() = Some(ws);
13052    }
13053
13054    /// `MEMRA_MLA_SEG_WS=1` (lane/glm5-mla-capture-20260904, default OFF): the T=1 MLA core runs
13055    /// its PRE segment into the session's stable buffers instead of fresh allocations, so the
13056    /// capture arc can hand a captured PRE graph's outputs to a captured POST graph. Read PER
13057    /// CALL. Byte-identical by construction (same kernels, same order, different address).
13058    pub(crate) fn mla_seg_ws_on() -> bool {
13059        std::env::var("MEMRA_MLA_SEG_WS").as_deref() == Ok("1")
13060    }
13061
13062    // ---- Verify-walk workspace (MEMRA_VERIFY_WS, door W — see VerifyWs). ----
13063    // take/recycle are no-ops with the door off, so every OFF-arm call site is byte-for-byte
13064    // the shipped program (fresh alloc, ordinary async free). All pooled sites are
13065    // verify-walk-only by construction (rows-exact matmuls, the KDA Rows stash arm, the MoE
13066    // vrows staging), and the pool is per-engine = per-stream: recycle-then-reuse carries the
13067    // same stream-ordering guarantee the async allocator's free-then-alloc does.
13068
13069    /// Pool-or-alloc f32 scratch for a verify-walk site (uninit contract unchanged).
13070    pub(crate) fn vws_uninit(
13071        &self,
13072        n: usize,
13073    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13074        if verify_ws_on() {
13075            let mut ws = self.verify_ws.lock().unwrap();
13076            let ws = &mut *ws;
13077            if let Some(s) = VerifyWs::take(&mut ws.f32_pool, &mut ws.held_bytes, n) {
13078                if VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
13079                    eprintln!(
13080                        "[glm5-verify-ws] engaged: verify-walk buffers recycling through \
13081                         the size-keyed pool (MEMRA_GLM5_VERIFY_WS=1)"
13082                    );
13083                }
13084                return Ok(s);
13085            }
13086        }
13087        self.alloc_uninit::<f32>(n)
13088    }
13089
13090    /// Pool-or-alloc i8 scratch (q8_1 activation planes).
13091    pub(crate) fn vws_uninit_i8(
13092        &self,
13093        n: usize,
13094    ) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
13095        if verify_ws_on() {
13096            let mut ws = self.verify_ws.lock().unwrap();
13097            let ws = &mut *ws;
13098            if let Some(s) = VerifyWs::take(&mut ws.i8_pool, &mut ws.held_bytes, n) {
13099                VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13100                return Ok(s);
13101            }
13102        }
13103        self.alloc_uninit::<i8>(n)
13104    }
13105
13106    /// Pool-or-alloc u64 scratch (the MoE vrows pointer tables).
13107    pub(crate) fn vws_uninit_u64(
13108        &self,
13109        n: usize,
13110    ) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
13111        if verify_ws_on() {
13112            let mut ws = self.verify_ws.lock().unwrap();
13113            let ws = &mut *ws;
13114            if let Some(s) = VerifyWs::take(&mut ws.u64_pool, &mut ws.held_bytes, n) {
13115                VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13116                return Ok(s);
13117            }
13118        }
13119        self.alloc_uninit::<u64>(n)
13120    }
13121
13122    /// The capture keeper the glm5 decode-graph door drains into its `RunGraph`. While
13123    /// [`glm5_graph_capture_open`] is set, `vws_recycle*` pushes here instead of returning the
13124    /// buffer to the verify workspace, so nothing a captured body baked can be re-issued to
13125    /// eager work between replays.
13126    #[allow(clippy::type_complexity)] // allow: mirrors the field's own type
13127    pub(crate) fn glm5_graph_keep(&self) -> &Mutex<Vec<Box<dyn std::any::Any + Send>>> {
13128        &self.capture_keep
13129    }
13130
13131    /// Return a dead verify-walk buffer to the pool (no-op with the door off: the buffer
13132    /// drops to the ordinary async free, the shipped program).
13133    pub(crate) fn vws_recycle(&self, s: CudaSlice<f32>) {
13134        if glm5_graph_capture_open() {
13135            self.capture_keep.lock().unwrap().push(Box::new(s));
13136            return;
13137        }
13138        if verify_ws_on() {
13139            let mut ws = self.verify_ws.lock().unwrap();
13140            let ws = &mut *ws;
13141            VerifyWs::put(&mut ws.f32_pool, &mut ws.held_bytes, s);
13142        }
13143    }
13144
13145    /// i8 twin of [`Self::vws_recycle`].
13146    pub(crate) fn vws_recycle_i8(&self, s: CudaSlice<i8>) {
13147        if glm5_graph_capture_open() {
13148            self.capture_keep.lock().unwrap().push(Box::new(s));
13149            return;
13150        }
13151        if verify_ws_on() {
13152            let mut ws = self.verify_ws.lock().unwrap();
13153            let ws = &mut *ws;
13154            VerifyWs::put(&mut ws.i8_pool, &mut ws.held_bytes, s);
13155        }
13156    }
13157
13158    /// u64 twin of [`Self::vws_recycle`].
13159    pub(crate) fn vws_recycle_u64(&self, s: CudaSlice<u64>) {
13160        if glm5_graph_capture_open() {
13161            self.capture_keep.lock().unwrap().push(Box::new(s));
13162            return;
13163        }
13164        if verify_ws_on() {
13165            let mut ws = self.verify_ws.lock().unwrap();
13166            let ws = &mut *ws;
13167            VerifyWs::put(&mut ws.u64_pool, &mut ws.held_bytes, s);
13168        }
13169    }
13170
13171    /// GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a
13172    /// resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread
13173    /// block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256
13174    /// blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host
13175    /// `argmax` (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read
13176    /// back (or kept resident for graph replay). Returns the device token buffer.
13177    /// Softmax probability of the (already-argmaxed) token `tok` under `logits` — the spec-decode
13178    /// p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.
13179    pub fn prob_of_token_device(
13180        &self,
13181        logits: &CudaSlice<f32>,
13182        tok: &CudaSlice<u32>,
13183        n_vocab: usize,
13184    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13185        let nb = ARGMAX_NB;
13186        let mut part = self.alloc_uninit::<f32>(nb)?;
13187        let mut p = self.alloc_uninit::<f32>(1)?;
13188        let f1 = self.func("prob_of_token_partial_f32");
13189        let cfg1 = LaunchConfig {
13190            grid_dim: (nb as u32, 1, 1),
13191            block_dim: (256, 1, 1),
13192            shared_mem_bytes: 0,
13193        };
13194        let nv = n_vocab as i32;
13195        let __s_b1 = self.gpu.stream();
13196        let mut b1 = __s_b1.launch_builder(&f1);
13197        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
13198        unsafe {
13199            b1.launch(cfg1)?;
13200        }
13201        let f2 = self.func("prob_of_token_final_f32");
13202        let cfg2 = LaunchConfig {
13203            grid_dim: (1, 1, 1),
13204            block_dim: (256, 1, 1),
13205            shared_mem_bytes: 0,
13206        };
13207        let nbi = nb as i32;
13208        let __s_b2 = self.gpu.stream();
13209        let mut b2 = __s_b2.launch_builder(&f2);
13210        b2.arg(&part).arg(&mut p).arg(&nbi);
13211        unsafe {
13212            b2.launch(cfg2)?;
13213        }
13214        Ok(p)
13215    }
13216
13217    /// Like `prob_of_token_device` but writes into a PERSISTENT `p_out` buffer (stable pointer).
13218    /// Required for CUDA-graph capture of the draft chain: the captured prob kernels must write
13219    /// where the host reads the p-min confidence between replays. Same kernels, same math.
13220    /// Slot-addressed twin of `prob_of_token_device_into`: token read from `tok_all[tok_idx]`
13221    /// (a view at the slot), probability written to `p_out[p_idx]` — same two kernels, the
13222    /// pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).
13223    pub fn prob_of_token_device_col(
13224        &self,
13225        logits: &CudaSlice<f32>,
13226        tok_all: &CudaSlice<u32>,
13227        tok_idx: usize,
13228        p_out: &mut CudaSlice<f32>,
13229        p_idx: usize,
13230        n_vocab: usize,
13231    ) -> Result<(), Box<dyn std::error::Error>> {
13232        let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
13233        let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
13234        let nb = ARGMAX_NB;
13235        let mut part = self.alloc_uninit::<f32>(nb)?;
13236        let f1 = self.func("prob_of_token_partial_f32");
13237        let cfg1 = LaunchConfig {
13238            grid_dim: (nb as u32, 1, 1),
13239            block_dim: (256, 1, 1),
13240            shared_mem_bytes: 0,
13241        };
13242        let nv = n_vocab as i32;
13243        let __s_b1 = self.gpu.stream();
13244        let mut b1 = __s_b1.launch_builder(&f1);
13245        b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
13246        unsafe {
13247            b1.launch(cfg1)?;
13248        }
13249        let f2 = self.func("prob_of_token_final_f32");
13250        let cfg2 = LaunchConfig {
13251            grid_dim: (1, 1, 1),
13252            block_dim: (256, 1, 1),
13253            shared_mem_bytes: 0,
13254        };
13255        let nbi = nb as i32;
13256        let __s_b2 = self.gpu.stream();
13257        let mut b2 = __s_b2.launch_builder(&f2);
13258        b2.arg(&part).arg(&mut p_v).arg(&nbi);
13259        unsafe {
13260            b2.launch(cfg2)?;
13261        }
13262        Ok(())
13263    }
13264
13265    pub fn prob_of_token_device_into(
13266        &self,
13267        logits: &CudaSlice<f32>,
13268        tok: &CudaSlice<u32>,
13269        p_out: &mut CudaSlice<f32>,
13270        n_vocab: usize,
13271    ) -> Result<(), Box<dyn std::error::Error>> {
13272        let nb = ARGMAX_NB;
13273        let mut part = self.alloc_uninit::<f32>(nb)?;
13274        let f1 = self.func("prob_of_token_partial_f32");
13275        let cfg1 = LaunchConfig {
13276            grid_dim: (nb as u32, 1, 1),
13277            block_dim: (256, 1, 1),
13278            shared_mem_bytes: 0,
13279        };
13280        let nv = n_vocab as i32;
13281        let __s_b1 = self.gpu.stream();
13282        let mut b1 = __s_b1.launch_builder(&f1);
13283        b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
13284        unsafe {
13285            b1.launch(cfg1)?;
13286        }
13287        let f2 = self.func("prob_of_token_final_f32");
13288        let cfg2 = LaunchConfig {
13289            grid_dim: (1, 1, 1),
13290            block_dim: (256, 1, 1),
13291            shared_mem_bytes: 0,
13292        };
13293        let nbi = nb as i32;
13294        let __s_b2 = self.gpu.stream();
13295        let mut b2 = __s_b2.launch_builder(&f2);
13296        b2.arg(&part).arg(p_out).arg(&nbi);
13297        unsafe {
13298            b2.launch(cfg2)?;
13299        }
13300        Ok(())
13301    }
13302
13303    /// Token-graph chunk loop: hist[idx] = *tok; idx += 1 — device-indexed history append
13304    /// (graph-constant params, device-varying index). Capture-safe.
13305    pub fn u32_hist_append(
13306        &self,
13307        tok: &CudaSlice<u32>,
13308        hist: &mut CudaSlice<u32>,
13309        idx: &mut CudaSlice<i32>,
13310    ) -> Result<(), Box<dyn std::error::Error>> {
13311        let f = self.func("u32_hist_append");
13312        let cfg = LaunchConfig {
13313            grid_dim: (1, 1, 1),
13314            block_dim: (32, 1, 1),
13315            shared_mem_bytes: 0,
13316        };
13317        let __s_b = self.gpu.stream();
13318        let mut b = __s_b.launch_builder(&f);
13319        b.arg(tok).arg(&mut *hist).arg(&mut *idx);
13320        unsafe {
13321            b.launch(cfg)?;
13322        }
13323        Ok(())
13324    }
13325
13326    pub fn argmax_token_device(
13327        &self,
13328        logits: &CudaSlice<f32>,
13329        n_vocab: usize,
13330    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13331        let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
13332        self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
13333        Ok(tok)
13334    }
13335    /// Like `argmax_token_device` but writes into a PERSISTENT `tok` buffer (stable pointer) instead
13336    /// of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the
13337    /// next token into the SAME device buffer the next replay's embed_gather reads, so the buffer
13338    /// pointer is baked once and the token id never round-trips to host inside steady state. The
13339    /// pass-1 partials scratch (`argmax_partials`) is also a resident stable-pointer buffer so both
13340    /// captured passes bake fixed addresses.
13341    pub fn argmax_token_device_into(
13342        &self,
13343        logits: &CudaSlice<f32>,
13344        tok: &mut CudaSlice<u32>,
13345        n_vocab: usize,
13346    ) -> Result<(), Box<dyn std::error::Error>> {
13347        let nb = ARGMAX_NB;
13348        let f1 = self.func("argmax_partial_f32");
13349        let f2 = self.func("argmax_final_f32");
13350        let mut guard = self.argmax_partials.lock().unwrap();
13351        if guard.is_none() {
13352            // allocate ONCE; under generate_graph this runs in the tracking-off prime window so the
13353            // buffers carry no cudarc events (illegal inside capture).
13354            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
13355            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
13356            *guard = Some((pv, pi));
13357        }
13358        let (part_v, part_i) = guard.as_mut().unwrap();
13359        let nv = n_vocab as i32;
13360        let nbi = nb as i32;
13361        // pass 1: NB blocks x 256 threads grid-stride scan -> per-block (val, idx) partials.
13362        let cfg1 = LaunchConfig {
13363            grid_dim: (nb as u32, 1, 1),
13364            block_dim: (256, 1, 1),
13365            shared_mem_bytes: 0,
13366        };
13367        let __s_b1 = self.gpu.stream();
13368        let mut b1 = __s_b1.launch_builder(&f1);
13369        b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
13370        unsafe {
13371            b1.launch(cfg1)?;
13372        }
13373        // pass 2: one block reduces NB partials -> token_out[0].
13374        let cfg2 = LaunchConfig {
13375            grid_dim: (1, 1, 1),
13376            block_dim: (256, 1, 1),
13377            shared_mem_bytes: 0,
13378        };
13379        let __s_b2 = self.gpu.stream();
13380        let mut b2 = __s_b2.launch_builder(&f2);
13381        b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
13382        unsafe {
13383            b2.launch(cfg2)?;
13384        }
13385        Ok(())
13386    }
13387    /// Column-`col` device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept
13388    /// walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels
13389    /// and tie-break contract as `argmax_token_device_into` (bit-identical to host argmax,
13390    /// argmax_gate-validated) — only the input pointer (a column view) and the output slot differ.
13391    /// Lets the accept walk read ONE [t] u32 instead of dtoh'ing the full [t, n_vocab] logits.
13392    pub fn argmax_token_device_col(
13393        &self,
13394        logits: &CudaSlice<f32>,
13395        col: usize,
13396        n_vocab: usize,
13397        toks: &mut CudaSlice<u32>,
13398        out_idx: usize,
13399    ) -> Result<(), Box<dyn std::error::Error>> {
13400        let nb = ARGMAX_NB;
13401        let f1 = self.func("argmax_partial_f32");
13402        let f2 = self.func("argmax_final_f32");
13403        let mut guard = self.argmax_partials.lock().unwrap();
13404        if guard.is_none() {
13405            let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
13406            let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
13407            *guard = Some((pv, pi));
13408        }
13409        let (part_v, part_i) = guard.as_mut().unwrap();
13410        let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
13411        let nv = n_vocab as i32;
13412        let nbi = nb as i32;
13413        let cfg1 = LaunchConfig {
13414            grid_dim: (nb as u32, 1, 1),
13415            block_dim: (256, 1, 1),
13416            shared_mem_bytes: 0,
13417        };
13418        let __s_b1 = self.gpu.stream();
13419        let mut b1 = __s_b1.launch_builder(&f1);
13420        b1.arg(&col_view)
13421            .arg(&mut *part_v)
13422            .arg(&mut *part_i)
13423            .arg(&nv);
13424        unsafe {
13425            b1.launch(cfg1)?;
13426        }
13427        let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
13428        let cfg2 = LaunchConfig {
13429            grid_dim: (1, 1, 1),
13430            block_dim: (256, 1, 1),
13431            shared_mem_bytes: 0,
13432        };
13433        let __s_b2 = self.gpu.stream();
13434        let mut b2 = __s_b2.launch_builder(&f2);
13435        b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
13436        unsafe {
13437            b2.launch(cfg2)?;
13438        }
13439        Ok(())
13440    }
13441    /// Read back a device u32 buffer (the spec accept walk's [t] per-column argmax tokens).
13442    pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13443        Ok(self.gpu.stream().clone_htod(v)?)
13444    }
13445    #[track_caller]
13446    pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
13447        crate::dtoh_trace_hit(d.len() * 8);
13448        let v = self.gpu.stream().clone_dtoh(d)?;
13449        self.gpu.stream().synchronize()?;
13450        Ok(v)
13451    }
13452
13453    #[track_caller]
13454    pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
13455        crate::dtoh_trace_hit(d.len() * 4);
13456        let v = self.gpu.stream().clone_dtoh(d)?;
13457        self.gpu.stream().synchronize()?;
13458        Ok(v)
13459    }
13460    /// Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots).
13461    /// H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload:
13462    /// contents change every step, the address must not, so a captured graph can read it).
13463    pub fn htod_u32_into(
13464        &self,
13465        dst: &mut CudaSlice<u32>,
13466        src: &[u32],
13467    ) -> Result<(), Box<dyn std::error::Error>> {
13468        let mut view = dst.slice_mut(0..src.len());
13469        self.gpu.stream().memcpy_htod(src, &mut view)?;
13470        Ok(())
13471    }
13472
13473    /// H2D into an existing i32 buffer. OPTIPIPE uses this to refresh a stage-local saved-len
13474    /// table without changing the device address its reconcile kernel consumes.
13475    pub fn htod_i32_into(
13476        &self,
13477        dst: &mut CudaSlice<i32>,
13478        src: &[i32],
13479    ) -> Result<(), Box<dyn std::error::Error>> {
13480        let mut view = dst.slice_mut(0..src.len());
13481        self.gpu.stream().memcpy_htod(src, &mut view)?;
13482        Ok(())
13483    }
13484
13485    #[track_caller]
13486    pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13487        crate::alloc_trace_hit(n * 4);
13488        let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
13489        self.keep_if_capturing(&s);
13490        Ok(s)
13491    }
13492    /// embed_gather into a PERSISTENT `x_out` buffer (stable pointer) for CUDA-graph capture (the
13493    /// embed output starts the per-step kernel chain and must be at a fixed address across replays).
13494    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
13495    pub fn embed_gather_device_into(
13496        &self,
13497        embd: &CudaSlice<u8>,
13498        token_d: &CudaSlice<u32>,
13499        x_out: &mut CudaSlice<f32>,
13500        n_embd: usize,
13501        qtype: i32,
13502        row_bytes: usize,
13503    ) -> Result<(), Box<dyn std::error::Error>> {
13504        let f = self.func("embed_gather_u32");
13505        let cfg = LaunchConfig {
13506            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
13507            block_dim: (256, 1, 1),
13508            shared_mem_bytes: 0,
13509        };
13510        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
13511        let __s_b = self.gpu.stream();
13512        let mut b = __s_b.launch_builder(&f);
13513        b.arg(embd)
13514            .arg(token_d)
13515            .arg(x_out)
13516            .arg(&ne)
13517            .arg(&qt)
13518            .arg(&rb);
13519        unsafe {
13520            b.launch(cfg)?;
13521        }
13522        Ok(())
13523    }
13524    /// Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.
13525    #[track_caller]
13526    pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
13527        crate::dtoh_trace_hit(d.len() * 4);
13528        let v = self.gpu.stream().clone_dtoh(d)?;
13529        self.gpu.stream().synchronize()?;
13530        Ok(v[0])
13531    }
13532    /// Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the
13533    /// graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores
13534    /// the counter value after the throwaway capture warmups corrupt it.
13535    /// ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync).
13536    /// The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable
13537    /// copy (fine at stream-idle boundaries, poison mid-round).
13538    pub fn i32_set_k(
13539        &self,
13540        dst: &mut CudaSlice<i32>,
13541        v: i32,
13542    ) -> Result<(), Box<dyn std::error::Error>> {
13543        let f = self.func("i32_set_k");
13544        let cfg = LaunchConfig {
13545            grid_dim: (1, 1, 1),
13546            block_dim: (1, 1, 1),
13547            shared_mem_bytes: 0,
13548        };
13549        let idx = 0i32;
13550        let __s_b = self.gpu.stream();
13551        let mut b = __s_b.launch_builder(&f);
13552        b.arg(dst).arg(&v).arg(&idx);
13553        unsafe {
13554            b.launch(cfg)?;
13555        }
13556        Ok(())
13557    }
13558
13559    pub fn set_i32_one(
13560        &self,
13561        d: &mut CudaSlice<i32>,
13562        v: i32,
13563    ) -> Result<(), Box<dyn std::error::Error>> {
13564        self.gpu.stream().memcpy_htod(&[v], d)?;
13565        Ok(())
13566    }
13567    /// Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident `token_d` counter
13568    /// during priming / capture-state restore.
13569    pub fn set_u32_one(
13570        &self,
13571        d: &mut CudaSlice<u32>,
13572        v: u32,
13573    ) -> Result<(), Box<dyn std::error::Error>> {
13574        self.gpu.stream().memcpy_htod(&[v], d)?;
13575        Ok(())
13576    }
13577    /// Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.
13578    #[track_caller]
13579    pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
13580        crate::dtoh_trace_hit(d.len() * 4);
13581        let v = self.gpu.stream().clone_dtoh(d)?;
13582        self.gpu.stream().synchronize()?;
13583        Ok(v[0])
13584    }
13585    /// Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).
13586    pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
13587        Ok(self.gpu.stream().clone_htod(bytes)?)
13588    }
13589    /// Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in
13590    /// `token_d[0]` from the resident embed table `embd` -> x_out[n_embd]. Bit-identical to host
13591    /// EmbedHost::gather (same per-dtype `deq`). No host round-trip of the token id.
13592    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
13593    pub fn embed_gather_device(
13594        &self,
13595        embd: &CudaSlice<u8>,
13596        token_d: &CudaSlice<u32>,
13597        n_embd: usize,
13598        qtype: i32,
13599        row_bytes: usize,
13600    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13601        let f = self.func("embed_gather_u32");
13602        let mut x = self.alloc_uninit::<f32>(n_embd)?;
13603        let cfg = LaunchConfig {
13604            grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
13605            block_dim: (256, 1, 1),
13606            shared_mem_bytes: 0,
13607        };
13608        let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
13609        let __s_b = self.gpu.stream();
13610        let mut b = __s_b.launch_builder(&f);
13611        b.arg(embd)
13612            .arg(token_d)
13613            .arg(&mut x)
13614            .arg(&ne)
13615            .arg(&qt)
13616            .arg(&rb);
13617        unsafe {
13618            b.launch(cfg)?;
13619        }
13620        Ok(x)
13621    }
13622
13623    /// T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod,
13624    /// rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + T*n_embd*4B htod
13625    /// (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).
13626    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
13627    pub fn embed_gather_device_t(
13628        &self,
13629        embd: &CudaSlice<u8>,
13630        tokens: &[u32],
13631        n_embd: usize,
13632        qtype: i32,
13633        row_bytes: usize,
13634    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13635        let t = tokens.len();
13636        let tok_d = self.gpu.stream().clone_htod(tokens)?;
13637        let f = self.func("embed_gather_u32_t");
13638        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13639        let cfg = LaunchConfig {
13640            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13641            block_dim: (256, 1, 1),
13642            shared_mem_bytes: 0,
13643        };
13644        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13645        let __s_b = self.gpu.stream();
13646        let mut b = __s_b.launch_builder(&f);
13647        b.arg(embd)
13648            .arg(&tok_d)
13649            .arg(&mut x)
13650            .arg(&ne)
13651            .arg(&qt)
13652            .arg(&rb)
13653            .arg(&ti);
13654        unsafe {
13655            b.launch(cfg)?;
13656        }
13657        Ok(x)
13658    }
13659
13660    /// T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens
13661    /// are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel
13662    /// as embed_gather_device_t — bit-identical rows.
13663    /// embed_gather over a token VIEW (spec round: tokens live in the round's batch buffer).
13664    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
13665    pub fn embed_gather_device_tv(
13666        &self,
13667        embd: &CudaSlice<u8>,
13668        tok_v: &cudarc::driver::CudaView<u32>,
13669        t: usize,
13670        n_embd: usize,
13671        qtype: i32,
13672        row_bytes: usize,
13673    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13674        let f = self.func("embed_gather_u32_t");
13675        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13676        let cfg = LaunchConfig {
13677            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13678            block_dim: (256, 1, 1),
13679            shared_mem_bytes: 0,
13680        };
13681        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13682        let __s_b = self.gpu.stream();
13683        let mut b = __s_b.launch_builder(&f);
13684        b.arg(embd)
13685            .arg(tok_v)
13686            .arg(&mut x)
13687            .arg(&ne)
13688            .arg(&qt)
13689            .arg(&rb)
13690            .arg(&ti);
13691        unsafe {
13692            b.launch(cfg)?;
13693        }
13694        Ok(x)
13695    }
13696
13697    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
13698    pub fn embed_gather_device_td(
13699        &self,
13700        embd: &CudaSlice<u8>,
13701        tok_d: &CudaSlice<u32>,
13702        t: usize,
13703        n_embd: usize,
13704        qtype: i32,
13705        row_bytes: usize,
13706    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13707        let f = self.func("embed_gather_u32_t");
13708        let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13709        let cfg = LaunchConfig {
13710            grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13711            block_dim: (256, 1, 1),
13712            shared_mem_bytes: 0,
13713        };
13714        let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13715        let __s_b = self.gpu.stream();
13716        let mut b = __s_b.launch_builder(&f);
13717        b.arg(embd)
13718            .arg(tok_d)
13719            .arg(&mut x)
13720            .arg(&ne)
13721            .arg(&qt)
13722            .arg(&rb)
13723            .arg(&ti);
13724        unsafe {
13725            b.launch(cfg)?;
13726        }
13727        Ok(x)
13728    }
13729
13730    /// Uninitialized device buffer — SKIPS the memset that `alloc_zeros` always issues. Decode
13731    /// profile (nsys): ~1050 memsets/token = 6.5% of decode GPU time + ~half the launch count, the
13732    /// dominant contributor to the 19% inter-kernel idle gap and a blocker for clean CUDA-graph
13733    /// capture. Use ONLY for buffers a kernel FULLY overwrites (every element written, no `+=`).
13734    /// SAFETY: caller guarantees the producing kernel writes every element before any read.
13735    #[inline]
13736    /// Keep an allocation alive for the current capture (no-op when retain mode is off).
13737    fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
13738        if self
13739            .capture_keep_on
13740            .load(std::sync::atomic::Ordering::Relaxed)
13741        {
13742            self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
13743        }
13744    }
13745
13746    #[track_caller]
13747    fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
13748        &self,
13749        n: usize,
13750    ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
13751        SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13752        crate::alloc_trace_hit(n * std::mem::size_of::<T>());
13753        let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
13754        // MEMRA_DEBUG_ZERO_ALLOCS=1 (task #14 defect hunt): memset EVERY engine allocation —
13755        // the global uninit-read discriminator (the prime-fn-scoped zeroing experiment could
13756        // not cover engine-internal buffers). Debug-only: massive launch overhead.
13757        {
13758            static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13759            if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
13760                // raw D8 memset (T lacks ValidAsZeroBits in the generic bound)
13761                use cudarc::driver::DevicePtrMut;
13762                let n_bytes = s.len() * std::mem::size_of::<T>();
13763                let stream = self.gpu.stream();
13764                let (p_, _g) = s.device_ptr_mut(&stream);
13765                unsafe {
13766                    cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
13767                        .result()?;
13768                }
13769            }
13770        }
13771        self.keep_if_capturing(&s);
13772        Ok(s)
13773    }
13774
13775    /// Public f32 uninitialized scratch (see `alloc_uninit`). For decode/forward scratch a kernel
13776    /// fully overwrites. SAFETY: producing kernel must write every element before any read.
13777    /// Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit
13778    /// consumers alloc through this (m=1 decode arms).
13779    #[track_caller]
13780    pub fn uninit_q8_pair(
13781        &self,
13782        n: usize,
13783    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13784        Ok((
13785            self.alloc_uninit::<i8>(n)?,
13786            self.alloc_uninit::<f32>(n / 32)?,
13787        ))
13788    }
13789
13790    #[track_caller]
13791    pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13792        self.alloc_uninit::<f32>(n)
13793    }
13794
13795    /// i8 uninitialized scratch (same contract as `uninit`).
13796    #[track_caller]
13797    pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
13798        self.alloc_uninit::<i8>(n)
13799    }
13800
13801    /// i32 uninitialized scratch (same contract as `uninit`) — the DSA indexer's position lists.
13802    #[track_caller]
13803    pub fn uninit_i32(&self, n: usize) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
13804        self.alloc_uninit::<i32>(n)
13805    }
13806
13807    /// RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads.
13808    /// gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights).
13809    /// Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).
13810    #[allow(clippy::too_many_arguments)]
13811    pub fn rms_norm3(
13812        &self,
13813        x: &CudaSlice<f32>,
13814        w0: &CudaSlice<f32>,
13815        w1: &CudaSlice<f32>,
13816        w2: &CudaSlice<f32>,
13817        d0: &mut CudaSlice<f32>,
13818        d1: &mut CudaSlice<f32>,
13819        d2: &mut CudaSlice<f32>,
13820        ncols: usize,
13821        nrows: usize,
13822        eps: f32,
13823    ) -> Result<(), Box<dyn std::error::Error>> {
13824        let f = self.func("rms_norm3_f32");
13825        let cfg = LaunchConfig {
13826            grid_dim: (nrows as u32, 1, 1),
13827            block_dim: (rms_block(), 1, 1),
13828            shared_mem_bytes: 0,
13829        };
13830        let (nc, e) = (ncols as i32, eps);
13831        let __s_b = self.gpu.stream();
13832        let mut b = __s_b.launch_builder(&f);
13833        b.arg(x)
13834            .arg(w0)
13835            .arg(w1)
13836            .arg(w2)
13837            .arg(d0)
13838            .arg(d1)
13839            .arg(d2)
13840            .arg(&nc)
13841            .arg(&e);
13842        unsafe {
13843            b.launch(cfg)?;
13844        }
13845        Ok(())
13846    }
13847
13848    /// gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim).
13849    #[allow(clippy::too_many_arguments)]
13850    /// True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane
13851    /// piggybacks on the same conditions.
13852    pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
13853        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13854        *WARP_ON.get_or_init(|| {
13855            std::env::var("MEMRA_QKVNORM_W")
13856                .map(|v| v != "0")
13857                .unwrap_or(true)
13858        }) && ncols.is_multiple_of(4)
13859            && rows >= 64
13860    }
13861
13862    /// w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as
13863    /// bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.
13864    #[allow(clippy::too_many_arguments)]
13865    pub fn rms_norm_qkv_w4b(
13866        &self,
13867        q: &CudaSlice<f32>,
13868        k: &CudaSlice<f32>,
13869        v: &CudaSlice<f32>,
13870        wq: &CudaSlice<f32>,
13871        wk: &CudaSlice<f32>,
13872        wv: &CudaSlice<f32>,
13873        dq: &mut CudaSlice<f32>,
13874        dk: &mut CudaSlice<f32>,
13875        dv: &mut CudaSlice<f32>,
13876        dvb: &mut CudaSlice<u8>,
13877        ncols: usize,
13878        rq: usize,
13879        rk: usize,
13880        eps: f32,
13881        vf16: bool,
13882    ) -> Result<(), Box<dyn std::error::Error>> {
13883        assert!(ncols.is_multiple_of(4) && rq + 2 * rk >= 64);
13884        let f = self.func("rms_norm_qkv_w4b_f32");
13885        let rows = (rq + 2 * rk) as u32;
13886        let cfg = LaunchConfig {
13887            grid_dim: (rows.div_ceil(8), 1, 1),
13888            block_dim: (256, 1, 1),
13889            shared_mem_bytes: 0,
13890        };
13891        let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
13892        let vf = vf16 as i32;
13893        let __s_b = self.gpu.stream();
13894        let mut b = __s_b.launch_builder(&f);
13895        b.arg(q)
13896            .arg(k)
13897            .arg(v)
13898            .arg(wq)
13899            .arg(wk)
13900            .arg(wv)
13901            .arg(dq)
13902            .arg(dk)
13903            .arg(dv)
13904            .arg(&mut *dvb)
13905            .arg(&nc)
13906            .arg(&rqi)
13907            .arg(&rki)
13908            .arg(&rvi)
13909            .arg(&e)
13910            .arg(&vf);
13911        unsafe {
13912            b.launch(cfg)?;
13913        }
13914        Ok(())
13915    }
13916
13917    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
13918    pub fn rms_norm_qkv(
13919        &self,
13920        q: &CudaSlice<f32>,
13921        k: &CudaSlice<f32>,
13922        v: &CudaSlice<f32>,
13923        wq: &CudaSlice<f32>,
13924        wk: &CudaSlice<f32>,
13925        wv: &CudaSlice<f32>,
13926        dq: &mut CudaSlice<f32>,
13927        dk: &mut CudaSlice<f32>,
13928        dv: &mut CudaSlice<f32>,
13929        ncols: usize,
13930        rq: usize,
13931        rk: usize,
13932        eps: f32,
13933    ) -> Result<(), Box<dyn std::error::Error>> {
13934        // Warp-per-row float4 twin (default; MEMRA_QKVNORM_W=0 reverts): the block-per-row form
13935        // spends 767us/launch on 17k+ 2KB rows at prefill depth (launch/reduce latency-bound,
13936        // ~92GB/s). Own numeric config (reduce order differs) — battery-gated.
13937        static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13938        let warp_on = *WARP_ON.get_or_init(|| {
13939            std::env::var("MEMRA_QKVNORM_W")
13940                .map(|v| v != "0")
13941                .unwrap_or(true)
13942        });
13943        // rows >= 64 keeps decode (nh + 2*nkv rows) on the block-tree kernel — decode/verify/
13944        // replay numerics are untouched on every model; only prefill depth takes the new config.
13945        if warp_on && ncols.is_multiple_of(4) && rq + 2 * rk >= 64 {
13946            let f = self.func("rms_norm_qkv_w4_f32");
13947            let rows = (rq + 2 * rk) as u32;
13948            let cfg = LaunchConfig {
13949                grid_dim: (rows.div_ceil(8), 1, 1),
13950                block_dim: (256, 1, 1),
13951                shared_mem_bytes: 0,
13952            };
13953            let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
13954            let __s_b = self.gpu.stream();
13955            let mut b = __s_b.launch_builder(&f);
13956            b.arg(q)
13957                .arg(k)
13958                .arg(v)
13959                .arg(wq)
13960                .arg(wk)
13961                .arg(wv)
13962                .arg(dq)
13963                .arg(dk)
13964                .arg(dv)
13965                .arg(&nc)
13966                .arg(&rqi)
13967                .arg(&rki)
13968                .arg(&rvi)
13969                .arg(&e);
13970            unsafe {
13971                b.launch(cfg)?;
13972            }
13973            return Ok(());
13974        }
13975        let f = self.func("rms_norm_qkv_f32");
13976        let grid = (rq + 2 * rk) as u32;
13977        let cfg = LaunchConfig {
13978            grid_dim: (grid, 1, 1),
13979            block_dim: (rms_block(), 1, 1),
13980            shared_mem_bytes: 0,
13981        };
13982        let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
13983        let __s_b = self.gpu.stream();
13984        let mut b = __s_b.launch_builder(&f);
13985        b.arg(q)
13986            .arg(k)
13987            .arg(v)
13988            .arg(wq)
13989            .arg(wk)
13990            .arg(wv)
13991            .arg(dq)
13992            .arg(dk)
13993            .arg(dv)
13994            .arg(&nc)
13995            .arg(&rqi)
13996            .arg(&rki)
13997            .arg(&e);
13998        unsafe {
13999            b.launch(cfg)?;
14000        }
14001        Ok(())
14002    }
14003
14004    /// gemma4 fused pair of rms_norms over two different inputs (same width).
14005    #[allow(clippy::too_many_arguments)]
14006    pub fn rms_norm2x(
14007        &self,
14008        a: &CudaSlice<f32>,
14009        bb: &CudaSlice<f32>,
14010        wa: &CudaSlice<f32>,
14011        wb: &CudaSlice<f32>,
14012        da: &mut CudaSlice<f32>,
14013        db: &mut CudaSlice<f32>,
14014        ncols: usize,
14015        nrows: usize,
14016        eps: f32,
14017    ) -> Result<(), Box<dyn std::error::Error>> {
14018        let f = self.func("rms_norm2x_f32");
14019        let cfg = LaunchConfig {
14020            grid_dim: (2 * nrows as u32, 1, 1),
14021            block_dim: (rms_block(), 1, 1),
14022            shared_mem_bytes: 0,
14023        };
14024        let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
14025        let __s_b = self.gpu.stream();
14026        let mut b = __s_b.launch_builder(&f);
14027        b.arg(a)
14028            .arg(bb)
14029            .arg(wa)
14030            .arg(wb)
14031            .arg(da)
14032            .arg(db)
14033            .arg(&nc)
14034            .arg(&nr)
14035            .arg(&e);
14036        unsafe {
14037            b.launch(cfg)?;
14038        }
14039        Ok(())
14040    }
14041
14042    /// gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).
14043    pub fn softcap(
14044        &self,
14045        y: &mut CudaSlice<f32>,
14046        cap: f32,
14047        n: usize,
14048    ) -> Result<(), Box<dyn std::error::Error>> {
14049        let f = self.func("softcap_f32");
14050        let cfg = LaunchConfig::for_num_elems(n as u32);
14051        let ni = n as i32;
14052        let __s_b = self.gpu.stream();
14053        let mut b = __s_b.launch_builder(&f);
14054        b.arg(y).arg(&cap).arg(&ni);
14055        unsafe {
14056            b.launch(cfg)?;
14057        }
14058        Ok(())
14059    }
14060
14061    /// gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch —
14062    /// graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).
14063    pub fn mask_ids_rows(
14064        &self,
14065        y: &mut CudaSlice<f32>,
14066        ids: &CudaSlice<i32>,
14067        n_ids: usize,
14068        n_vocab: usize,
14069        t: usize,
14070    ) -> Result<(), Box<dyn std::error::Error>> {
14071        let f = self.func("mask_ids_rows_f32");
14072        let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
14073        let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
14074        let __s_b = self.gpu.stream();
14075        let mut b = __s_b.launch_builder(&f);
14076        b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
14077        unsafe {
14078            b.launch(cfg)?;
14079        }
14080        Ok(())
14081    }
14082
14083    /// gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.
14084    #[allow(clippy::too_many_arguments)]
14085    pub fn add_scale_rms_norm(
14086        &self,
14087        a: &CudaSlice<f32>,
14088        b_in: &CudaSlice<f32>,
14089        c: f32,
14090        w: &CudaSlice<f32>,
14091        res: &mut CudaSlice<f32>,
14092        dst: &mut CudaSlice<f32>,
14093        ncols: usize,
14094        nrows: usize,
14095        eps: f32,
14096    ) -> Result<(), Box<dyn std::error::Error>> {
14097        let f = self.func("add_scale_rms_norm_f32");
14098        let cfg = LaunchConfig {
14099            grid_dim: (nrows as u32, 1, 1),
14100            block_dim: (rms_block(), 1, 1),
14101            shared_mem_bytes: 0,
14102        };
14103        let (nc, e2) = (ncols as i32, eps);
14104        let __s_b = self.gpu.stream();
14105        let mut b = __s_b.launch_builder(&f);
14106        b.arg(a)
14107            .arg(b_in)
14108            .arg(&c)
14109            .arg(w)
14110            .arg(res)
14111            .arg(dst)
14112            .arg(&nc)
14113            .arg(&e2);
14114        unsafe {
14115            b.launch(cfg)?;
14116        }
14117        Ok(())
14118    }
14119
14120    /// gemma4: res = (a+b)*c AND the next layer's attn_norm EMITTED q8_1 in one launch.
14121    /// Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).
14122    #[allow(clippy::too_many_arguments)]
14123    pub fn add_scale_rms_norm_q8_1(
14124        &self,
14125        a: &CudaSlice<f32>,
14126        b_in: &CudaSlice<f32>,
14127        c: f32,
14128        w: &CudaSlice<f32>,
14129        res: &mut CudaSlice<f32>,
14130        ncols: usize,
14131        nrows: usize,
14132        eps: f32,
14133    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14134        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14135        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14136        let (nc, e2) = (ncols as i32, eps);
14137        if Self::pdl_on() && Self::pdl_wb_on() {
14138            {
14139                use cudarc::driver::{DevicePtr, DevicePtrMut};
14140                let s = &self.gpu.stream();
14141                let (pa, _g0) = a.device_ptr(s);
14142                let (pb, _g1) = b_in.device_ptr(s);
14143                let (pw, _g2) = w.device_ptr(s);
14144                let (pr, _g3) = res.device_ptr_mut(s);
14145                let (pq, _g4) = out_q.device_ptr_mut(s);
14146                let (pd, _g5) = out_d.device_ptr_mut(s);
14147                let mut ps = [
14148                    &pa as *const _ as *mut std::ffi::c_void,
14149                    &pb as *const _ as *mut _,
14150                    &c as *const _ as *mut _,
14151                    &pw as *const _ as *mut _,
14152                    &pr as *const _ as *mut _,
14153                    &pq as *const _ as *mut _,
14154                    &pd as *const _ as *mut _,
14155                    &nc as *const _ as *mut _,
14156                    &e2 as *const _ as *mut _,
14157                ];
14158                unsafe {
14159                    self.launch_pdl(
14160                        "add_scale_rms_norm_q8_1",
14161                        (nrows as u32, 1, 1),
14162                        (rms_block(), 1, 1),
14163                        &mut ps,
14164                    )?;
14165                }
14166            }
14167            return Ok((out_q, out_d));
14168        }
14169        let f = self.func("add_scale_rms_norm_q8_1");
14170        let cfg = LaunchConfig {
14171            grid_dim: (nrows as u32, 1, 1),
14172            block_dim: (rms_block(), 1, 1),
14173            shared_mem_bytes: 0,
14174        };
14175        let __s_b = self.gpu.stream();
14176        let mut b = __s_b.launch_builder(&f);
14177        b.arg(a)
14178            .arg(b_in)
14179            .arg(&c)
14180            .arg(w)
14181            .arg(res)
14182            .arg(&mut out_q)
14183            .arg(&mut out_d)
14184            .arg(&nc)
14185            .arg(&e2);
14186        unsafe {
14187            b.launch(cfg)?;
14188        }
14189        Ok((out_q, out_d))
14190    }
14191
14192    /// Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).
14193    #[allow(clippy::too_many_arguments)]
14194    pub fn add_scale_rms_norm_q8_1_into(
14195        &self,
14196        a: &CudaSlice<f32>,
14197        b_in: &CudaSlice<f32>,
14198        c: f32,
14199        w: &CudaSlice<f32>,
14200        res: &mut CudaSlice<f32>,
14201        ncols: usize,
14202        nrows: usize,
14203        eps: f32,
14204        out_q: &mut CudaSlice<i8>,
14205        out_d: &mut CudaSlice<f32>,
14206    ) -> Result<(), Box<dyn std::error::Error>> {
14207        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
14208        let (nc, e2) = (ncols as i32, eps);
14209        if Self::pdl_on() && Self::pdl_wb_on() {
14210            use cudarc::driver::{DevicePtr, DevicePtrMut};
14211            let s = &self.gpu.stream();
14212            let (pa, _g0) = a.device_ptr(s);
14213            let (pb, _g1) = b_in.device_ptr(s);
14214            let (pw, _g2) = w.device_ptr(s);
14215            let (pr, _g3) = res.device_ptr_mut(s);
14216            let (pq, _g4) = out_q.device_ptr_mut(s);
14217            let (pd, _g5) = out_d.device_ptr_mut(s);
14218            let mut ps = [
14219                &pa as *const _ as *mut std::ffi::c_void,
14220                &pb as *const _ as *mut _,
14221                &c as *const _ as *mut _,
14222                &pw as *const _ as *mut _,
14223                &pr as *const _ as *mut _,
14224                &pq as *const _ as *mut _,
14225                &pd as *const _ as *mut _,
14226                &nc as *const _ as *mut _,
14227                &e2 as *const _ as *mut _,
14228            ];
14229            unsafe {
14230                self.launch_pdl(
14231                    "add_scale_rms_norm_q8_1",
14232                    (nrows as u32, 1, 1),
14233                    (rms_block(), 1, 1),
14234                    &mut ps,
14235                )?;
14236            }
14237            return Ok(());
14238        }
14239        let f = self.func("add_scale_rms_norm_q8_1");
14240        let cfg = LaunchConfig {
14241            grid_dim: (nrows as u32, 1, 1),
14242            block_dim: (rms_block(), 1, 1),
14243            shared_mem_bytes: 0,
14244        };
14245        let __s_b = self.gpu.stream();
14246        let mut b = __s_b.launch_builder(&f);
14247        b.arg(a)
14248            .arg(b_in)
14249            .arg(&c)
14250            .arg(w)
14251            .arg(res)
14252            .arg(&mut *out_q)
14253            .arg(&mut *out_d)
14254            .arg(&nc)
14255            .arg(&e2);
14256        unsafe {
14257            b.launch(cfg)?;
14258        }
14259        Ok(())
14260    }
14261
14262    /// E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch
14263    /// replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.
14264    #[allow(clippy::too_many_arguments)]
14265    pub fn rms_pre_add_scale_rms_norm_q8_1(
14266        &self,
14267        a: &CudaSlice<f32>,
14268        wa: &CudaSlice<f32>,
14269        b_in: &CudaSlice<f32>,
14270        c: f32,
14271        w: &CudaSlice<f32>,
14272        res: &mut CudaSlice<f32>,
14273        ncols: usize,
14274        nrows: usize,
14275        eps: f32,
14276    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14277        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14278        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14279        let (nc, e2) = (ncols as i32, eps);
14280        if Self::pdl_on() {
14281            {
14282                use cudarc::driver::{DevicePtr, DevicePtrMut};
14283                let s = &self.gpu.stream();
14284                let (pa, _g0) = a.device_ptr(s);
14285                let (pwa, _g1) = wa.device_ptr(s);
14286                let (pb, _g2) = b_in.device_ptr(s);
14287                let (pw, _g3) = w.device_ptr(s);
14288                let (pr, _g4) = res.device_ptr_mut(s);
14289                let (pq, _g5) = out_q.device_ptr_mut(s);
14290                let (pd, _g6) = out_d.device_ptr_mut(s);
14291                let mut ps = [
14292                    &pa as *const _ as *mut std::ffi::c_void,
14293                    &pwa as *const _ as *mut _,
14294                    &pb as *const _ as *mut _,
14295                    &c as *const _ as *mut _,
14296                    &pw as *const _ as *mut _,
14297                    &pr as *const _ as *mut _,
14298                    &pq as *const _ as *mut _,
14299                    &pd as *const _ as *mut _,
14300                    &nc as *const _ as *mut _,
14301                    &e2 as *const _ as *mut _,
14302                ];
14303                unsafe {
14304                    self.launch_pdl(
14305                        "rms_pre_add_scale_rms_norm_q8_1",
14306                        (nrows as u32, 1, 1),
14307                        (rms_block(), 1, 1),
14308                        &mut ps,
14309                    )?;
14310                }
14311            }
14312            return Ok((out_q, out_d));
14313        }
14314        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
14315        let cfg = LaunchConfig {
14316            grid_dim: (nrows as u32, 1, 1),
14317            block_dim: (rms_block(), 1, 1),
14318            shared_mem_bytes: 0,
14319        };
14320        let __s_b = self.gpu.stream();
14321        let mut b = __s_b.launch_builder(&f);
14322        b.arg(a)
14323            .arg(wa)
14324            .arg(b_in)
14325            .arg(&c)
14326            .arg(w)
14327            .arg(res)
14328            .arg(&mut out_q)
14329            .arg(&mut out_d)
14330            .arg(&nc)
14331            .arg(&e2);
14332        unsafe {
14333            b.launch(cfg)?;
14334        }
14335        Ok((out_q, out_d))
14336    }
14337
14338    /// GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the
14339    /// consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.
14340    pub fn gelu_tanh_mul_q8_1(
14341        &self,
14342        gate: &CudaSlice<f32>,
14343        up: &cudarc::driver::CudaView<f32>,
14344        act: &mut CudaSlice<f32>,
14345        ncols: usize,
14346        nrows: usize,
14347    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14348        debug_assert!(ncols.is_multiple_of(128));
14349        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14350        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14351        let nc = ncols as i32;
14352        if Self::pdl_on() {
14353            {
14354                use cudarc::driver::{DevicePtr, DevicePtrMut};
14355                let s = &self.gpu.stream();
14356                let (pg, _g0) = gate.device_ptr(s);
14357                let (pu, _g1) = up.device_ptr(s);
14358                let (pact, _g2) = act.device_ptr_mut(s);
14359                let (pq, _g3) = out_q.device_ptr_mut(s);
14360                let (pd, _g4) = out_d.device_ptr_mut(s);
14361                let mut ps = [
14362                    &pg as *const _ as *mut std::ffi::c_void,
14363                    &pu as *const _ as *mut _,
14364                    &pact as *const _ as *mut _,
14365                    &pq as *const _ as *mut _,
14366                    &pd as *const _ as *mut _,
14367                    &nc as *const _ as *mut _,
14368                ];
14369                unsafe {
14370                    self.launch_pdl(
14371                        "gelu_tanh_mul_q8_1",
14372                        (nrows as u32, 1, 1),
14373                        (rms_block(), 1, 1),
14374                        &mut ps,
14375                    )?;
14376                }
14377            }
14378            return Ok((out_q, out_d));
14379        }
14380        let f = self.func("gelu_tanh_mul_q8_1");
14381        let cfg = LaunchConfig {
14382            grid_dim: (nrows as u32, 1, 1),
14383            block_dim: (rms_block(), 1, 1),
14384            shared_mem_bytes: 0,
14385        };
14386        let __s_b = self.gpu.stream();
14387        let mut b = __s_b.launch_builder(&f);
14388        b.arg(gate)
14389            .arg(up)
14390            .arg(act)
14391            .arg(&mut out_q)
14392            .arg(&mut out_d)
14393            .arg(&nc);
14394        unsafe {
14395            b.launch(cfg)?;
14396        }
14397        Ok((out_q, out_d))
14398    }
14399
14400    /// Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).
14401    #[allow(clippy::too_many_arguments)]
14402    pub fn gelu_tanh_mul_q8_1_into(
14403        &self,
14404        gate: &CudaSlice<f32>,
14405        up: &cudarc::driver::CudaView<f32>,
14406        act: &mut CudaSlice<f32>,
14407        ncols: usize,
14408        nrows: usize,
14409        out_q: &mut CudaSlice<i8>,
14410        out_d: &mut CudaSlice<f32>,
14411    ) -> Result<(), Box<dyn std::error::Error>> {
14412        debug_assert!(ncols.is_multiple_of(128));
14413        debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
14414        let nc = ncols as i32;
14415        if Self::pdl_on() {
14416            use cudarc::driver::{DevicePtr, DevicePtrMut};
14417            let s = &self.gpu.stream();
14418            let (pg, _g0) = gate.device_ptr(s);
14419            let (pu, _g1) = up.device_ptr(s);
14420            let (pact, _g2) = act.device_ptr_mut(s);
14421            let (pq, _g3) = out_q.device_ptr_mut(s);
14422            let (pd, _g4) = out_d.device_ptr_mut(s);
14423            let mut ps = [
14424                &pg as *const _ as *mut std::ffi::c_void,
14425                &pu as *const _ as *mut _,
14426                &pact as *const _ as *mut _,
14427                &pq as *const _ as *mut _,
14428                &pd as *const _ as *mut _,
14429                &nc as *const _ as *mut _,
14430            ];
14431            unsafe {
14432                self.launch_pdl(
14433                    "gelu_tanh_mul_q8_1",
14434                    (nrows as u32, 1, 1),
14435                    (rms_block(), 1, 1),
14436                    &mut ps,
14437                )?;
14438            }
14439            return Ok(());
14440        }
14441        let f = self.func("gelu_tanh_mul_q8_1");
14442        let cfg = LaunchConfig {
14443            grid_dim: (nrows as u32, 1, 1),
14444            block_dim: (rms_block(), 1, 1),
14445            shared_mem_bytes: 0,
14446        };
14447        let __s_b = self.gpu.stream();
14448        let mut b = __s_b.launch_builder(&f);
14449        b.arg(gate)
14450            .arg(up)
14451            .arg(&mut *act)
14452            .arg(&mut *out_q)
14453            .arg(&mut *out_d)
14454            .arg(&nc);
14455        unsafe {
14456            b.launch(cfg)?;
14457        }
14458        Ok(())
14459    }
14460
14461    /// gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).
14462    #[allow(clippy::too_many_arguments)]
14463    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
14464    pub fn add_rms_norm3_q8z(
14465        &self,
14466        a: &CudaSlice<f32>,
14467        b_in: &CudaSlice<f32>,
14468        w0: &CudaSlice<f32>,
14469        w1: &CudaSlice<f32>,
14470        w2: &CudaSlice<f32>,
14471        res: &mut CudaSlice<f32>,
14472        out1: &mut CudaSlice<f32>,
14473        ncols: usize,
14474        nrows: usize,
14475        eps: f32,
14476    ) -> Result<
14477        (
14478            (CudaSlice<i8>, CudaSlice<f32>),
14479            (CudaSlice<i8>, CudaSlice<f32>),
14480        ),
14481        Box<dyn std::error::Error>,
14482    > {
14483        let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
14484        let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14485        let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
14486        let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14487        let f = self.func("add_rms_norm3_q8z_f32");
14488        let cfg = LaunchConfig {
14489            grid_dim: (nrows as u32, 1, 1),
14490            block_dim: (rms_block(), 1, 1),
14491            shared_mem_bytes: 0,
14492        };
14493        let (nc, e2) = (ncols as i32, eps);
14494        let __s_b = self.gpu.stream();
14495        let mut b = __s_b.launch_builder(&f);
14496        b.arg(a)
14497            .arg(b_in)
14498            .arg(w0)
14499            .arg(w1)
14500            .arg(w2)
14501            .arg(res)
14502            .arg(&mut q0)
14503            .arg(&mut d0)
14504            .arg(out1)
14505            .arg(&mut q2)
14506            .arg(&mut d2)
14507            .arg(&nc)
14508            .arg(&e2);
14509        unsafe {
14510            b.launch(cfg)?;
14511        }
14512        Ok(((q0, d0), (q2, d2)))
14513    }
14514
14515    /// gemma4: res = a+b AND the three rms_norms of res in one launch.
14516    #[allow(clippy::too_many_arguments)]
14517    pub fn add_rms_norm3(
14518        &self,
14519        a: &CudaSlice<f32>,
14520        b_in: &CudaSlice<f32>,
14521        w0: &CudaSlice<f32>,
14522        w1: &CudaSlice<f32>,
14523        w2: &CudaSlice<f32>,
14524        res: &mut CudaSlice<f32>,
14525        d0: &mut CudaSlice<f32>,
14526        d1: &mut CudaSlice<f32>,
14527        d2: &mut CudaSlice<f32>,
14528        ncols: usize,
14529        nrows: usize,
14530        eps: f32,
14531    ) -> Result<(), Box<dyn std::error::Error>> {
14532        let f = self.func("add_rms_norm3_f32");
14533        let cfg = LaunchConfig {
14534            grid_dim: (nrows as u32, 1, 1),
14535            block_dim: (rms_block(), 1, 1),
14536            shared_mem_bytes: 0,
14537        };
14538        let (nc, e2) = (ncols as i32, eps);
14539        let __s_b = self.gpu.stream();
14540        let mut b = __s_b.launch_builder(&f);
14541        b.arg(a)
14542            .arg(b_in)
14543            .arg(w0)
14544            .arg(w1)
14545            .arg(w2)
14546            .arg(res)
14547            .arg(d0)
14548            .arg(d1)
14549            .arg(d2)
14550            .arg(&nc)
14551            .arg(&e2);
14552        unsafe {
14553            b.launch(cfg)?;
14554        }
14555        Ok(())
14556    }
14557
14558    /// dst = (a + b) * c (residual add + layer scale, one launch).
14559    pub fn add_scale(
14560        &self,
14561        a: &CudaSlice<f32>,
14562        b_in: &CudaSlice<f32>,
14563        c: f32,
14564        dst: &mut CudaSlice<f32>,
14565        n: usize,
14566    ) -> Result<(), Box<dyn std::error::Error>> {
14567        let f = self.func("add_scale_f32");
14568        let cfg = LaunchConfig::for_num_elems(n as u32);
14569        let ni = n as i32;
14570        let __s_b = self.gpu.stream();
14571        let mut b = __s_b.launch_builder(&f);
14572        b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
14573        unsafe {
14574            b.launch(cfg)?;
14575        }
14576        Ok(())
14577    }
14578
14579    /// Vision-tower LayerNorm (with bias) over [nrows, ncols] — lane/vision.
14580    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
14581    pub fn layer_norm_bias(
14582        &self,
14583        x: &CudaSlice<f32>,
14584        w: &CudaSlice<f32>,
14585        b: &CudaSlice<f32>,
14586        dst: &mut CudaSlice<f32>,
14587        ncols: usize,
14588        nrows: usize,
14589        eps: f32,
14590    ) -> Result<(), Box<dyn std::error::Error>> {
14591        let f = self.func("layer_norm_bias_f32");
14592        let (nc, e) = (ncols as i32, eps);
14593        let cfg = LaunchConfig {
14594            grid_dim: (nrows as u32, 1, 1),
14595            block_dim: (256, 1, 1),
14596            shared_mem_bytes: 0,
14597        };
14598        let __s_b = self.gpu.stream();
14599        let mut lb = __s_b.launch_builder(&f);
14600        lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
14601        unsafe {
14602            lb.launch(cfg)?;
14603        }
14604        Ok(())
14605    }
14606
14607    /// gelu_pytorch_tanh elementwise (vision tower MLP activation).
14608    pub fn gelu_tanh(
14609        &self,
14610        x: &CudaSlice<f32>,
14611        dst: &mut CudaSlice<f32>,
14612        n: usize,
14613    ) -> Result<(), Box<dyn std::error::Error>> {
14614        let f = self.func("gelu_tanh_f32");
14615        let ni = n as i64;
14616        let cfg = LaunchConfig {
14617            grid_dim: (n.div_ceil(256) as u32, 1, 1),
14618            block_dim: (256, 1, 1),
14619            shared_mem_bytes: 0,
14620        };
14621        let __s_b = self.gpu.stream();
14622        let mut lb = __s_b.launch_builder(&f);
14623        lb.arg(x).arg(&mut *dst).arg(&ni);
14624        unsafe {
14625            lb.launch(cfg)?;
14626        }
14627        Ok(())
14628    }
14629
14630    /// In-place row softmax over [nrows, ncols] (bidirectional vision attention).
14631    pub fn row_softmax(
14632        &self,
14633        x: &mut CudaSlice<f32>,
14634        ncols: usize,
14635        nrows: usize,
14636    ) -> Result<(), Box<dyn std::error::Error>> {
14637        let f = self.func("row_softmax_f32");
14638        let nc = ncols as i32;
14639        let cfg = LaunchConfig {
14640            grid_dim: (nrows as u32, 1, 1),
14641            block_dim: (256, 1, 1),
14642            shared_mem_bytes: 0,
14643        };
14644        let __s_b = self.gpu.stream();
14645        let mut lb = __s_b.launch_builder(&f);
14646        lb.arg(&mut *x).arg(&nc);
14647        unsafe {
14648            lb.launch(cfg)?;
14649        }
14650        Ok(())
14651    }
14652
14653    pub fn rms_norm(
14654        &self,
14655        x: &CudaSlice<f32>,
14656        w: &CudaSlice<f32>,
14657        dst: &mut CudaSlice<f32>,
14658        ncols: usize,
14659        nrows: usize,
14660        eps: f32,
14661    ) -> Result<(), Box<dyn std::error::Error>> {
14662        let (nc, e) = (ncols as i32, eps);
14663        let kname = if Self::norm_ilp_on() {
14664            "rms_norm_f32_v2"
14665        } else {
14666            "rms_norm_f32"
14667        };
14668        if Self::pdl_on() && Self::pdl_wb_on() {
14669            use cudarc::driver::{DevicePtr, DevicePtrMut};
14670            let s = &self.gpu.stream();
14671            let (px, _g0) = x.device_ptr(s);
14672            let (pw, _g1) = w.device_ptr(s);
14673            let (pd, _g2) = dst.device_ptr_mut(s);
14674            let mut ps = [
14675                &px as *const _ as *mut std::ffi::c_void,
14676                &pw as *const _ as *mut _,
14677                &pd as *const _ as *mut _,
14678                &nc as *const _ as *mut _,
14679                &e as *const _ as *mut _,
14680            ];
14681            unsafe {
14682                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
14683            }
14684            return Ok(());
14685        }
14686        let f = self.func(kname);
14687        let cfg = LaunchConfig {
14688            grid_dim: (nrows as u32, 1, 1),
14689            block_dim: (rms_block(), 1, 1),
14690            shared_mem_bytes: 0,
14691        };
14692        let __s_b = self.gpu.stream();
14693        let mut b = __s_b.launch_builder(&f);
14694        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
14695        unsafe {
14696            b.launch(cfg)?;
14697        }
14698        Ok(())
14699    }
14700
14701    /// RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused `rms_norm_q8_1` and
14702    /// `add_rms_norm_q8_1` kernels' sum-of-squares reduction. The spec verify path MUST use this
14703    /// to match decode's FP accumulation order: the standard `rms_norm` at blockDim=256 has a
14704    /// different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a
14705    /// different shfl-tree reduction that can shift `scale = rsqrt(sum/n + eps)` by ULPs, causing
14706    /// divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying
14707    /// `rms_norm_f32` kernel supports any blockDim (generic reduce with shared[32]).
14708    pub fn rms_norm_decode(
14709        &self,
14710        x: &CudaSlice<f32>,
14711        w: &CudaSlice<f32>,
14712        dst: &mut CudaSlice<f32>,
14713        ncols: usize,
14714        nrows: usize,
14715        eps: f32,
14716    ) -> Result<(), Box<dyn std::error::Error>> {
14717        let f = self.func(if Self::norm_ilp_on() {
14718            "rms_norm_f32_v2"
14719        } else {
14720            "rms_norm_f32"
14721        });
14722        let cfg = LaunchConfig {
14723            grid_dim: (nrows as u32, 1, 1),
14724            block_dim: (1024, 1, 1),
14725            shared_mem_bytes: 0,
14726        };
14727        let (nc, e) = (ncols as i32, eps);
14728        let __s_b = self.gpu.stream();
14729        let mut b = __s_b.launch_builder(&f);
14730        b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
14731        unsafe {
14732            b.launch(cfg)?;
14733        }
14734        Ok(())
14735    }
14736
14737    /// DECODE GLUE-FUSION LEVER: `z = rms_norm(x)*w` emitted DIRECTLY as q8_1 (no f32 `z` materialized,
14738    /// no standalone quantize_q8_1 launch). Returns (out_q [nrows*ncols i8], out_d [nrows*nblk f32])
14739    /// ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.
14740    pub fn rms_norm_q8_1(
14741        &self,
14742        x: &CudaSlice<f32>,
14743        w: &CudaSlice<f32>,
14744        ncols: usize,
14745        nrows: usize,
14746        eps: f32,
14747    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14748        let nblk = ncols / 32;
14749        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
14750        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
14751        let (nc, e) = (ncols as i32, eps);
14752        if Self::pdl_on() {
14753            {
14754                use cudarc::driver::{DevicePtr, DevicePtrMut};
14755                let s = &self.gpu.stream();
14756                let (px, _g0) = x.device_ptr(s);
14757                let (pw, _g1) = w.device_ptr(s);
14758                let (pq, _g2) = q.device_ptr_mut(s);
14759                let (pd, _g3) = d.device_ptr_mut(s);
14760                let mut ps = [
14761                    &px as *const _ as *mut std::ffi::c_void,
14762                    &pw as *const _ as *mut _,
14763                    &pq as *const _ as *mut _,
14764                    &pd as *const _ as *mut _,
14765                    &nc as *const _ as *mut _,
14766                    &e as *const _ as *mut _,
14767                ];
14768                unsafe {
14769                    self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
14770                }
14771            }
14772            return Ok((q, d));
14773        }
14774        let f = self.func("rms_norm_q8_1");
14775        // 1024 threads: decode is nrows=1 -> ONE CTA; 32 warps hide the pass1->pass2 latency
14776        // (s[32] reduce already sized for 32 warps). Same shape math at any blockDim.
14777        let cfg = LaunchConfig {
14778            grid_dim: (nrows as u32, 1, 1),
14779            block_dim: (1024, 1, 1),
14780            shared_mem_bytes: 0,
14781        };
14782        let __s_b = self.gpu.stream();
14783        let mut b = __s_b.launch_builder(&f);
14784        b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
14785        unsafe {
14786            b.launch(cfg)?;
14787        }
14788        Ok((q, d))
14789    }
14790
14791    /// Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the
14792    /// PDL arm), caller-owned outputs.
14793    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
14794    pub fn rms_norm_q8_1_into(
14795        &self,
14796        x: &CudaSlice<f32>,
14797        w: &CudaSlice<f32>,
14798        ncols: usize,
14799        nrows: usize,
14800        eps: f32,
14801        q: &mut CudaSlice<i8>,
14802        d: &mut CudaSlice<f32>,
14803    ) -> Result<(), Box<dyn std::error::Error>> {
14804        let nblk = ncols / 32;
14805        debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
14806        let (nc, e) = (ncols as i32, eps);
14807        if Self::pdl_on() {
14808            use cudarc::driver::{DevicePtr, DevicePtrMut};
14809            let s = &self.gpu.stream();
14810            let (px, _g0) = x.device_ptr(s);
14811            let (pw, _g1) = w.device_ptr(s);
14812            let (pq, _g2) = q.device_ptr_mut(s);
14813            let (pd, _g3) = d.device_ptr_mut(s);
14814            let mut ps = [
14815                &px as *const _ as *mut std::ffi::c_void,
14816                &pw as *const _ as *mut _,
14817                &pq as *const _ as *mut _,
14818                &pd as *const _ as *mut _,
14819                &nc as *const _ as *mut _,
14820                &e as *const _ as *mut _,
14821            ];
14822            unsafe {
14823                self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
14824            }
14825            return Ok(());
14826        }
14827        let f = self.func("rms_norm_q8_1");
14828        let cfg = LaunchConfig {
14829            grid_dim: (nrows as u32, 1, 1),
14830            block_dim: (1024, 1, 1),
14831            shared_mem_bytes: 0,
14832        };
14833        let __s_b = self.gpu.stream();
14834        let mut b = __s_b.launch_builder(&f);
14835        b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
14836        unsafe {
14837            b.launch(cfg)?;
14838        }
14839        Ok(())
14840    }
14841
14842    /// Door `MEMRA_GLM5_Q8_FUSE` (lane/b200-q8-fuse-20260902): RMSNorm emitting BOTH the f32
14843    /// normed row `z` (needed by callers that also read the un-quantized row — the MoE router
14844    /// logits, an ungated shared-expert path) AND its q8_1 quantization (int8 qs + per-32 f32
14845    /// scale), in one launch. Removes the standalone `quantize_q8_1(&z, ...)` launch a caller
14846    /// would otherwise issue against the identical bytes. BIT-IDENTICAL to
14847    /// `rms_norm(x,w,&mut z,ncols,nrows,eps)` then `quantize_q8_1(&z,nrows,ncols)` — see
14848    /// `rms_norm_zq8_f32`'s header in cu/kernels.cu for the identity argument. MUST launch at
14849    /// `rms_block()`, the SAME blockDim `rms_norm` uses: the sum-of-squares block-reduce tree
14850    /// depends on blockDim (per-thread stride, shfl-tree depth), so a fixed 1024 would diverge
14851    /// from `rms_norm`'s actual per-model blockDim (256 by default; 1024 only where a loader
14852    /// overrides `RMS_BLOCK_DEFAULT`, e.g. gemma4) — caught by `q8_fuse_gate`'s ncols=1536 shape
14853    /// before this landed (README: keep this dynamic, never hardcode the block size again).
14854    pub fn rms_norm_zq8_f32(
14855        &self,
14856        x: &CudaSlice<f32>,
14857        w: &CudaSlice<f32>,
14858        z: &mut CudaSlice<f32>,
14859        ncols: usize,
14860        nrows: usize,
14861        eps: f32,
14862    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14863        self.rms_norm_zq8_f32_arm(x, w, z, ncols, nrows, eps, Self::norm_ilp_zq8_on())
14864    }
14865
14866    /// The arm-explicit form of [`Engine::rms_norm_zq8_f32`]: `ilp` selects the
14867    /// `rms_norm_zq8_f32_v2` twin (MEMRA_NORM_ILP and MEMRA_NORM_ILP_ZQ8, both default ON; four loads in flight per round
14868    /// in both passes, bit-identical by construction, see its header in cu/kernels.cu) over the
14869    /// v1 kernel. Public so the gate (`tests/norm_zq8_ilp_gpu.rs`) can run BOTH arms in one
14870    /// process; `norm_ilp_on()` is a process-lifetime latch and cannot be flipped mid-test.
14871    #[allow(clippy::too_many_arguments)]
14872    pub fn rms_norm_zq8_f32_arm(
14873        &self,
14874        x: &CudaSlice<f32>,
14875        w: &CudaSlice<f32>,
14876        z: &mut CudaSlice<f32>,
14877        ncols: usize,
14878        nrows: usize,
14879        eps: f32,
14880        ilp: bool,
14881    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14882        assert!(ncols.is_multiple_of(32));
14883        let nblk = ncols / 32;
14884        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?; // full-overwrite output
14885        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?; // full-overwrite output
14886        let f = self.func(if ilp {
14887            "rms_norm_zq8_f32_v2"
14888        } else {
14889            "rms_norm_zq8_f32"
14890        });
14891        // One block per row, blockDim = rms_block() — MUST match `rms_norm`'s launch exactly
14892        // (see the doc comment above); the kernel body is blockDim-generic (rms_norm_f32's
14893        // reduce shape), so this is the only thing that has to track it.
14894        let cfg = LaunchConfig {
14895            grid_dim: (nrows as u32, 1, 1),
14896            block_dim: (rms_block(), 1, 1),
14897            shared_mem_bytes: 0,
14898        };
14899        let (nc, ep) = (ncols as i32, eps);
14900        let __s_b = self.gpu.stream();
14901        let mut b = __s_b.launch_builder(&f);
14902        b.arg(x)
14903            .arg(w)
14904            .arg(&mut *z)
14905            .arg(&mut q)
14906            .arg(&mut d)
14907            .arg(&nc)
14908            .arg(&ep);
14909        unsafe {
14910            b.launch(cfg)?;
14911        }
14912        Ok((q, d))
14913    }
14914
14915    /// Slot-fed quantize_q8_1 twin (alloc-free capture lane).
14916    pub fn quantize_q8_1_into(
14917        &self,
14918        x: &CudaSlice<f32>,
14919        m: usize,
14920        in_f: usize,
14921        q: &mut CudaSlice<i8>,
14922        d: &mut CudaSlice<f32>,
14923    ) -> Result<(), Box<dyn std::error::Error>> {
14924        let nblk = in_f / 32;
14925        debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
14926        let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
14927        let (inf, mi) = (in_f as i32, m as i32);
14928        if Self::pdl_on() && Self::pdl_wb_on() {
14929            use cudarc::driver::{DevicePtr, DevicePtrMut};
14930            let s = &self.gpu.stream();
14931            let (px, _g0) = x.device_ptr(s);
14932            let (pq, _g1) = q.device_ptr_mut(s);
14933            let (pd, _g2) = d.device_ptr_mut(s);
14934            let mut ps = [
14935                &px as *const _ as *mut std::ffi::c_void,
14936                &pq as *const _ as *mut _,
14937                &pd as *const _ as *mut _,
14938                &inf as *const _ as *mut _,
14939                &mi as *const _ as *mut _,
14940            ];
14941            unsafe {
14942                self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
14943            }
14944            return Ok(());
14945        }
14946        let f = self.func("quantize_q8_1");
14947        let __s_b = self.gpu.stream();
14948        let mut b = __s_b.launch_builder(&f);
14949        b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
14950        unsafe {
14951            b.launch(cfg)?;
14952        }
14953        Ok(())
14954    }
14955
14956    /// DECODE GLUE-FUSION LEVER: `res = a+b; z = rms_norm(res)*w` with z emitted as q8_1. `res` is
14957    /// still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1.
14958    /// Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.
14959    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
14960    pub fn add_rms_norm_q8_1(
14961        &self,
14962        a: &CudaSlice<f32>,
14963        b_in: &CudaSlice<f32>,
14964        w: &CudaSlice<f32>,
14965        res: &mut CudaSlice<f32>,
14966        ncols: usize,
14967        nrows: usize,
14968        eps: f32,
14969    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14970        let nblk = ncols / 32;
14971        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
14972        let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
14973        let f = self.func("add_rms_norm_q8_1");
14974        // 1024 threads: same single-CTA-at-decode reasoning as rms_norm_q8_1.
14975        let cfg = LaunchConfig {
14976            grid_dim: (nrows as u32, 1, 1),
14977            block_dim: (1024, 1, 1),
14978            shared_mem_bytes: 0,
14979        };
14980        let (nc, e) = (ncols as i32, eps);
14981        let __s_bld = self.gpu.stream();
14982        let mut bld = __s_bld.launch_builder(&f);
14983        bld.arg(a)
14984            .arg(b_in)
14985            .arg(w)
14986            .arg(res)
14987            .arg(&mut q)
14988            .arg(&mut d)
14989            .arg(&nc)
14990            .arg(&e);
14991        unsafe {
14992            bld.launch(cfg)?;
14993        }
14994        Ok((q, d))
14995    }
14996
14997    /// RANK3 LEVER (add+rmsnorm fuse): `res = a + b; dst = rms_norm(res) * w` in ONE launch. Fuses
14998    /// e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per
14999    /// residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).
15000    /// O-PROJ TAIL FUSION M2: mixed = a0+a1 composed in-register, then the VERBATIM
15001    /// add_rms_norm program. Raw UVA pointers for the join partials (persistent ws rows).
15002    #[allow(clippy::too_many_arguments)]
15003    pub fn join_add_rms_norm_raw(
15004        &self,
15005        a0_raw: u64,
15006        a1_raw: u64,
15007        x: &CudaSlice<f32>,
15008        w: &CudaSlice<f32>,
15009        res: &mut CudaSlice<f32>,
15010        dst: &mut CudaSlice<f32>,
15011        ncols: usize,
15012        eps: f32,
15013    ) -> Result<(), Box<dyn std::error::Error>> {
15014        if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
15015            return Err("join_add_rms_norm geometry".into());
15016        }
15017        let f = self.func("join_add_rms_norm_f32");
15018        let cfg = LaunchConfig {
15019            grid_dim: (1, 1, 1),
15020            block_dim: (rms_block(), 1, 1),
15021            shared_mem_bytes: 0,
15022        };
15023        let (nc, e) = (ncols as i32, eps);
15024        let __s_b = self.gpu.stream();
15025        let mut b = __s_b.launch_builder(&f);
15026        b.arg(&a0_raw)
15027            .arg(&a1_raw)
15028            .arg(x)
15029            .arg(w)
15030            .arg(&mut *res)
15031            .arg(&mut *dst)
15032            .arg(&nc)
15033            .arg(&e);
15034        unsafe {
15035            b.launch(cfg)?;
15036        }
15037        Ok(())
15038    }
15039
15040    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
15041    pub fn add_rms_norm(
15042        &self,
15043        a: &CudaSlice<f32>,
15044        b: &CudaSlice<f32>,
15045        w: &CudaSlice<f32>,
15046        res: &mut CudaSlice<f32>,
15047        dst: &mut CudaSlice<f32>,
15048        ncols: usize,
15049        nrows: usize,
15050        eps: f32,
15051    ) -> Result<(), Box<dyn std::error::Error>> {
15052        let (nc, e) = (ncols as i32, eps);
15053        let kname = if Self::norm_ilp_on() {
15054            "add_rms_norm_f32_v2"
15055        } else {
15056            "add_rms_norm_f32"
15057        };
15058        if Self::pdl_on() && Self::pdl_wb_on() {
15059            use cudarc::driver::{DevicePtr, DevicePtrMut};
15060            let s = &self.gpu.stream();
15061            let (pa, _g0) = a.device_ptr(s);
15062            let (pb, _g1) = b.device_ptr(s);
15063            let (pw, _g2) = w.device_ptr(s);
15064            let (pr, _g3) = res.device_ptr_mut(s);
15065            let (pd, _g4) = dst.device_ptr_mut(s);
15066            let mut ps = [
15067                &pa as *const _ as *mut std::ffi::c_void,
15068                &pb as *const _ as *mut _,
15069                &pw as *const _ as *mut _,
15070                &pr as *const _ as *mut _,
15071                &pd as *const _ as *mut _,
15072                &nc as *const _ as *mut _,
15073                &e as *const _ as *mut _,
15074            ];
15075            unsafe {
15076                self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
15077            }
15078            return Ok(());
15079        }
15080        let f = self.func(kname);
15081        let cfg = LaunchConfig {
15082            grid_dim: (nrows as u32, 1, 1),
15083            block_dim: (rms_block(), 1, 1),
15084            shared_mem_bytes: 0,
15085        };
15086        let __s_b2 = self.gpu.stream();
15087        let mut b2 = __s_b2.launch_builder(&f);
15088        b2.arg(a)
15089            .arg(b)
15090            .arg(w)
15091            .arg(&mut *res)
15092            .arg(&mut *dst)
15093            .arg(&nc)
15094            .arg(&e);
15095        unsafe {
15096            b2.launch(cfg)?;
15097        }
15098        Ok(())
15099    }
15100
15101    /// E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into
15102    /// the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).
15103    #[allow(clippy::too_many_arguments)]
15104    pub fn rms_pre_add_rms_norm(
15105        &self,
15106        a: &CudaSlice<f32>,
15107        wa: &CudaSlice<f32>,
15108        b: &CudaSlice<f32>,
15109        w: &CudaSlice<f32>,
15110        res: &mut CudaSlice<f32>,
15111        dst: &mut CudaSlice<f32>,
15112        ncols: usize,
15113        nrows: usize,
15114        eps: f32,
15115    ) -> Result<(), Box<dyn std::error::Error>> {
15116        let f = self.func("rms_pre_add_rms_norm_f32");
15117        let cfg = LaunchConfig {
15118            grid_dim: (nrows as u32, 1, 1),
15119            block_dim: (rms_block(), 1, 1),
15120            shared_mem_bytes: 0,
15121        };
15122        let (nc, e) = (ncols as i32, eps);
15123        let __s_b2 = self.gpu.stream();
15124        let mut b2 = __s_b2.launch_builder(&f);
15125        b2.arg(a)
15126            .arg(wa)
15127            .arg(b)
15128            .arg(w)
15129            .arg(&mut *res)
15130            .arg(&mut *dst)
15131            .arg(&nc)
15132            .arg(&e);
15133        unsafe {
15134            b2.launch(cfg)?;
15135        }
15136        Ok(())
15137    }
15138
15139    /// wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).
15140    #[allow(clippy::too_many_arguments)]
15141    pub fn rms_pre_add_rms_norm_q8z(
15142        &self,
15143        a: &CudaSlice<f32>,
15144        wa: &CudaSlice<f32>,
15145        b: &CudaSlice<f32>,
15146        w: &CudaSlice<f32>,
15147        res: &mut CudaSlice<f32>,
15148        dst: &mut CudaSlice<f32>,
15149        ncols: usize,
15150        nrows: usize,
15151        eps: f32,
15152    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15153        debug_assert!(ncols.is_multiple_of(128));
15154        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
15155        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
15156        let (nc, e) = (ncols as i32, eps);
15157        if Self::pdl_on() {
15158            {
15159                use cudarc::driver::{DevicePtr, DevicePtrMut};
15160                let s = &self.gpu.stream();
15161                let (pa, _g0) = a.device_ptr(s);
15162                let (pwa, _g1) = wa.device_ptr(s);
15163                let (pb, _g2) = b.device_ptr(s);
15164                let (pw, _g3) = w.device_ptr(s);
15165                let (pr, _g4) = res.device_ptr_mut(s);
15166                let (pdst, _g5) = dst.device_ptr_mut(s);
15167                let (pq, _g6) = out_q.device_ptr_mut(s);
15168                let (pd, _g7) = out_d.device_ptr_mut(s);
15169                let mut ps = [
15170                    &pa as *const _ as *mut std::ffi::c_void,
15171                    &pwa as *const _ as *mut _,
15172                    &pb as *const _ as *mut _,
15173                    &pw as *const _ as *mut _,
15174                    &pr as *const _ as *mut _,
15175                    &pdst as *const _ as *mut _,
15176                    &pq as *const _ as *mut _,
15177                    &pd as *const _ as *mut _,
15178                    &nc as *const _ as *mut _,
15179                    &e as *const _ as *mut _,
15180                ];
15181                unsafe {
15182                    self.launch_pdl(
15183                        "rms_pre_add_rms_norm_q8z_f32",
15184                        (nrows as u32, 1, 1),
15185                        (rms_block(), 1, 1),
15186                        &mut ps,
15187                    )?;
15188                }
15189            }
15190            return Ok((out_q, out_d));
15191        }
15192        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
15193        let cfg = LaunchConfig {
15194            grid_dim: (nrows as u32, 1, 1),
15195            block_dim: (rms_block(), 1, 1),
15196            shared_mem_bytes: 0,
15197        };
15198        let __s_b2 = self.gpu.stream();
15199        let mut b2 = __s_b2.launch_builder(&f);
15200        b2.arg(a)
15201            .arg(wa)
15202            .arg(b)
15203            .arg(w)
15204            .arg(&mut *res)
15205            .arg(&mut *dst)
15206            .arg(&mut out_q)
15207            .arg(&mut out_d)
15208            .arg(&nc)
15209            .arg(&e);
15210        unsafe {
15211            b2.launch(cfg)?;
15212        }
15213        Ok((out_q, out_d))
15214    }
15215
15216    /// Slot-fed twin of `rms_pre_add_rms_norm_q8z` (gemma4 pn-fold, slotted/graph arm):
15217    /// identical kernel, caller-owned outputs, PLAIN launch only — the dc_slotted capture
15218    /// body must stay attribute-free (the fused2_into precedent).
15219    #[allow(clippy::too_many_arguments)]
15220    pub fn rms_pre_add_rms_norm_q8z_into(
15221        &self,
15222        a: &CudaSlice<f32>,
15223        wa: &CudaSlice<f32>,
15224        b: &CudaSlice<f32>,
15225        w: &CudaSlice<f32>,
15226        res: &mut CudaSlice<f32>,
15227        dst: &mut CudaSlice<f32>,
15228        ncols: usize,
15229        nrows: usize,
15230        eps: f32,
15231        out_q: &mut CudaSlice<i8>,
15232        out_d: &mut CudaSlice<f32>,
15233    ) -> Result<(), Box<dyn std::error::Error>> {
15234        debug_assert!(ncols.is_multiple_of(128));
15235        let (nc, e) = (ncols as i32, eps);
15236        let f = self.func("rms_pre_add_rms_norm_q8z_f32");
15237        let cfg = LaunchConfig {
15238            grid_dim: (nrows as u32, 1, 1),
15239            block_dim: (rms_block(), 1, 1),
15240            shared_mem_bytes: 0,
15241        };
15242        let __s_b = self.gpu.stream();
15243        let mut b2 = __s_b.launch_builder(&f);
15244        b2.arg(a)
15245            .arg(wa)
15246            .arg(b)
15247            .arg(w)
15248            .arg(&mut *res)
15249            .arg(&mut *dst)
15250            .arg(&mut *out_q)
15251            .arg(&mut *out_d)
15252            .arg(&nc)
15253            .arg(&e);
15254        unsafe {
15255            b2.launch(cfg)?;
15256        }
15257        Ok(())
15258    }
15259
15260    /// Slot-fed twin of `rms_pre_add_scale_rms_norm_q8_1` (gemma4 pn-fold exit, slotted
15261    /// arm): identical kernel, caller-owned outputs, PLAIN launch (capture-safe).
15262    #[allow(clippy::too_many_arguments)]
15263    pub fn rms_pre_add_scale_rms_norm_q8_1_into(
15264        &self,
15265        a: &CudaSlice<f32>,
15266        wa: &CudaSlice<f32>,
15267        b_in: &CudaSlice<f32>,
15268        c: f32,
15269        w: &CudaSlice<f32>,
15270        res: &mut CudaSlice<f32>,
15271        ncols: usize,
15272        nrows: usize,
15273        eps: f32,
15274        out_q: &mut CudaSlice<i8>,
15275        out_d: &mut CudaSlice<f32>,
15276    ) -> Result<(), Box<dyn std::error::Error>> {
15277        debug_assert!(ncols.is_multiple_of(128));
15278        let (nc, e2) = (ncols as i32, eps);
15279        let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
15280        let cfg = LaunchConfig {
15281            grid_dim: (nrows as u32, 1, 1),
15282            block_dim: (rms_block(), 1, 1),
15283            shared_mem_bytes: 0,
15284        };
15285        let __s_b = self.gpu.stream();
15286        let mut b2 = __s_b.launch_builder(&f);
15287        b2.arg(a)
15288            .arg(wa)
15289            .arg(b_in)
15290            .arg(&c)
15291            .arg(w)
15292            .arg(&mut *res)
15293            .arg(&mut *out_q)
15294            .arg(&mut *out_d)
15295            .arg(&nc)
15296            .arg(&e2);
15297        unsafe {
15298            b2.launch(cfg)?;
15299        }
15300        Ok(())
15301    }
15302
15303    /// gemma4 pn-fold seam (GAP-DIAGNOSIS verdict 7, the E4B glue backport): the dense
15304    /// decode/verify/slotted trio folds post_attn_norm into the tail entry
15305    /// (rms_pre_add_rms_norm[_q8z]) and post_ffw_norm into the residual exit
15306    /// (rms_pre_add_scale_rms_norm_q8_1). BITS-CHANGING vs the two-launch chain (the
15307    /// single-phase reduction's expansion rounding — E4B receipts); every arm moves
15308    /// together so decode == verify == graph parity holds BY CONSTRUCTION within either
15309    /// seam value. MEMRA_G4_PNFOLD=0 restores the unfused chain everywhere.
15310    pub fn g4_pnfold_on() -> bool {
15311        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15312        *ON.get_or_init(|| {
15313            std::env::var("MEMRA_G4_PNFOLD")
15314                .map(|v| v != "0")
15315                .unwrap_or(true)
15316        })
15317    }
15318
15319    /// wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent
15320    /// blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None
15321    /// off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).
15322    pub fn build_q4_out_concat3(
15323        &self,
15324        w0: &crate::model::GpuTensor,
15325        w1: &crate::model::GpuTensor,
15326        w2: &crate::model::GpuTensor,
15327    ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
15328        use crate::model::GpuTensor;
15329        let part = |w: &GpuTensor| -> Option<(usize, usize)> {
15330            match w {
15331                GpuTensor::Quant {
15332                    qtype,
15333                    row_bytes,
15334                    rp,
15335                    ..
15336                } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
15337                _ => None,
15338            }
15339        };
15340        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
15341        else {
15342            return Ok(None);
15343        };
15344        if rb0 != rb1
15345            || rb0 != rb2
15346            || w0.in_features() != w1.in_features()
15347            || w0.in_features() != w2.in_features()
15348        {
15349            return Ok(None);
15350        }
15351        fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
15352            match w {
15353                crate::model::GpuTensor::Quant { bytes, .. } => bytes,
15354                _ => unreachable!(),
15355            }
15356        }
15357        let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
15358        let total = rb0 * (o0 + o1 + o2);
15359        let mut cat = self.alloc_u8(total)?;
15360        self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
15361        self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
15362        self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
15363        Ok(Some(GpuTensor::Quant {
15364            bytes: cat,
15365            qtype: QT_Q4_0,
15366            row_bytes: rb0,
15367            ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
15368            scale: 1.0,
15369            rp: false,
15370            #[cfg(memra_cutlass)]
15371            cutlass: None,
15372            fp8: None,
15373            blk: None,
15374            rp4: None,
15375            f16: None,
15376        }))
15377    }
15378
15379    /// FULL-WIDTH-ROPE CONTRACT for the fused rms_norm+qkv+rope kernels
15380    /// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up).
15381    ///
15382    /// `rms_norm_qkv_rope_f32` / `_cat_f32` (`cu/kernels.cu`) and
15383    /// `rms_norm_qkv_rope_append_body` (`cu/flash_attn.cu`) take NO `n_dims`/`n_rot` argument.
15384    /// They compute `int half = ncols / 2` and rotate the FULL head width by construction — the
15385    /// standalone `rope_neox*` kernels take `n_dims` and early-return above it, these do not.
15386    ///
15387    /// Every call site today is a gemma-4 arm, where `n_rot == head_dim`, so nothing is wrong
15388    /// now. But this is the n_rot bug ONE FUSION away: fuse a partial-rotary arch onto these
15389    /// kernels (qwen3.5 = 64 rotary dims of a 256-wide head; step35 full-attn = 64 of 128) and
15390    /// 192 dims that must pass through unrotated get rotated silently — no error, no NaN, just a
15391    /// wrong model. The n_rot lane already paid for that class once, in the config derivation.
15392    ///
15393    /// So the fusions now take the layer's DERIVED rope width and refuse anything but full
15394    /// width. A future partial-rotary caller fails at its first launch with the geometry named
15395    /// instead of serving quietly wrong logits.
15396    fn full_width_rope_only(
15397        kernel: &str,
15398        n_rot: usize,
15399        head_dim: usize,
15400    ) -> Result<(), Box<dyn std::error::Error>> {
15401        if n_rot == head_dim {
15402            return Ok(());
15403        }
15404        Err(format!(
15405            "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
15406             rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
15407             width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
15408             through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
15409             n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
15410        )
15411        .into())
15412    }
15413
15414    /// wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.
15415    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
15416    /// ([`Engine::full_width_rope_only`]).
15417    #[allow(clippy::too_many_arguments)]
15418    pub fn rms_norm_qkv_rope_cat(
15419        &self,
15420        qkv: &CudaSlice<f32>,
15421        wq: &CudaSlice<f32>,
15422        wk: &CudaSlice<f32>,
15423        wv: &CudaSlice<f32>,
15424        q: &mut CudaSlice<f32>,
15425        k: &mut CudaSlice<f32>,
15426        v: &mut CudaSlice<f32>,
15427        head_dim: usize,
15428        n_rot: usize,
15429        rq: usize,
15430        rk: usize,
15431        pos: &CudaSlice<i32>,
15432        nh_q: usize,
15433        nh_k: usize,
15434        base: f32,
15435        freq_scale: f32,
15436        ff: Option<&CudaSlice<f32>>,
15437        eps: f32,
15438    ) -> Result<(), Box<dyn std::error::Error>> {
15439        Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
15440        let rows = rq + rk + rk;
15441        let theta_scale = base.powf(-2.0 / head_dim as f32);
15442        let (nc, rqi, rki, nhq, nhk) = (
15443            head_dim as i32,
15444            rq as i32,
15445            rk as i32,
15446            nh_q as i32,
15447            nh_k as i32,
15448        );
15449        if Self::pdl_on() {
15450            use cudarc::driver::{DevicePtr, DevicePtrMut};
15451            let s = &self.gpu.stream();
15452            let (pqkv, _g0) = qkv.device_ptr(s);
15453            let (pwq, _g1) = wq.device_ptr(s);
15454            let (pwk, _g2) = wk.device_ptr(s);
15455            let (pwv, _g3) = wv.device_ptr(s);
15456            let (pq, _g4) = q.device_ptr_mut(s);
15457            let (pk, _g5) = k.device_ptr_mut(s);
15458            let (pv, _g6) = v.device_ptr_mut(s);
15459            let (ppos, _g7) = pos.device_ptr(s);
15460            let (pff, _g8) = match ff {
15461                Some(t) => {
15462                    let (p, g) = t.device_ptr(s);
15463                    (p, Some(g))
15464                }
15465                None => (0, None),
15466            };
15467            let mut ps = [
15468                &pqkv as *const _ as *mut std::ffi::c_void,
15469                &pwq as *const _ as *mut _,
15470                &pwk as *const _ as *mut _,
15471                &pwv as *const _ as *mut _,
15472                &pq as *const _ as *mut _,
15473                &pk as *const _ as *mut _,
15474                &pv as *const _ as *mut _,
15475                &nc as *const _ as *mut _,
15476                &rqi as *const _ as *mut _,
15477                &rki as *const _ as *mut _,
15478                &ppos as *const _ as *mut _,
15479                &nhq as *const _ as *mut _,
15480                &nhk as *const _ as *mut _,
15481                &theta_scale as *const _ as *mut _,
15482                &freq_scale as *const _ as *mut _,
15483                &pff as *const _ as *mut _,
15484                &eps as *const _ as *mut _,
15485            ];
15486            unsafe {
15487                self.launch_pdl(
15488                    "rms_norm_qkv_rope_cat_f32",
15489                    (rows as u32, 1, 1),
15490                    (rms_block(), 1, 1),
15491                    &mut ps,
15492                )?;
15493            }
15494            return Ok(());
15495        }
15496        let f = self.func("rms_norm_qkv_rope_cat_f32");
15497        let cfg = LaunchConfig {
15498            grid_dim: (rows as u32, 1, 1),
15499            block_dim: (rms_block(), 1, 1),
15500            shared_mem_bytes: 0,
15501        };
15502        let __s_b = self.gpu.stream();
15503        let mut b = __s_b.launch_builder(&f);
15504        match ff {
15505            Some(t) => {
15506                b.arg(qkv)
15507                    .arg(wq)
15508                    .arg(wk)
15509                    .arg(wv)
15510                    .arg(&mut *q)
15511                    .arg(&mut *k)
15512                    .arg(&mut *v)
15513                    .arg(&nc)
15514                    .arg(&rqi)
15515                    .arg(&rki)
15516                    .arg(pos)
15517                    .arg(&nhq)
15518                    .arg(&nhk)
15519                    .arg(&theta_scale)
15520                    .arg(&freq_scale)
15521                    .arg(t)
15522                    .arg(&eps);
15523                unsafe {
15524                    b.launch(cfg)?;
15525                }
15526            }
15527            None => {
15528                let null: u64 = 0;
15529                b.arg(qkv)
15530                    .arg(wq)
15531                    .arg(wk)
15532                    .arg(wv)
15533                    .arg(&mut *q)
15534                    .arg(&mut *k)
15535                    .arg(&mut *v)
15536                    .arg(&nc)
15537                    .arg(&rqi)
15538                    .arg(&rki)
15539                    .arg(pos)
15540                    .arg(&nhq)
15541                    .arg(&nhk)
15542                    .arg(&theta_scale)
15543                    .arg(&freq_scale)
15544                    .arg(&null)
15545                    .arg(&eps);
15546                unsafe {
15547                    b.launch(cfg)?;
15548                }
15549            }
15550        }
15551        Ok(())
15552    }
15553
15554    /// wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).
15555    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
15556    /// ([`Engine::full_width_rope_only`]).
15557    #[allow(clippy::too_many_arguments)]
15558    pub fn rms_norm_qkv_rope(
15559        &self,
15560        q0: &CudaSlice<f32>,
15561        k0: &CudaSlice<f32>,
15562        v0: &CudaSlice<f32>,
15563        wq: &CudaSlice<f32>,
15564        wk: &CudaSlice<f32>,
15565        wv: &CudaSlice<f32>,
15566        q: &mut CudaSlice<f32>,
15567        k: &mut CudaSlice<f32>,
15568        v: &mut CudaSlice<f32>,
15569        head_dim: usize,
15570        n_rot: usize,
15571        rq: usize,
15572        rk: usize,
15573        pos: &CudaSlice<i32>,
15574        nh_q: usize,
15575        nh_k: usize,
15576        base: f32,
15577        freq_scale: f32,
15578        ff: Option<&CudaSlice<f32>>,
15579        eps: f32,
15580    ) -> Result<(), Box<dyn std::error::Error>> {
15581        Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
15582        let f = self.func("rms_norm_qkv_rope_f32");
15583        let rows = rq + rk + rk; // q rows + k rows + v rows (rk == rv)
15584        let cfg = LaunchConfig {
15585            grid_dim: (rows as u32, 1, 1),
15586            block_dim: (rms_block(), 1, 1),
15587            shared_mem_bytes: 0,
15588        };
15589        let theta_scale = base.powf(-2.0 / head_dim as f32);
15590        let (nc, rqi, rki, nhq, nhk) = (
15591            head_dim as i32,
15592            rq as i32,
15593            rk as i32,
15594            nh_q as i32,
15595            nh_k as i32,
15596        );
15597        let __s_b = self.gpu.stream();
15598        let mut b = __s_b.launch_builder(&f);
15599        match ff {
15600            Some(t) => {
15601                b.arg(q0)
15602                    .arg(k0)
15603                    .arg(v0)
15604                    .arg(wq)
15605                    .arg(wk)
15606                    .arg(wv)
15607                    .arg(&mut *q)
15608                    .arg(&mut *k)
15609                    .arg(&mut *v)
15610                    .arg(&nc)
15611                    .arg(&rqi)
15612                    .arg(&rki)
15613                    .arg(pos)
15614                    .arg(&nhq)
15615                    .arg(&nhk)
15616                    .arg(&theta_scale)
15617                    .arg(&freq_scale)
15618                    .arg(t)
15619                    .arg(&eps);
15620                unsafe {
15621                    b.launch(cfg)?;
15622                }
15623            }
15624            None => {
15625                let null: u64 = 0;
15626                b.arg(q0)
15627                    .arg(k0)
15628                    .arg(v0)
15629                    .arg(wq)
15630                    .arg(wk)
15631                    .arg(wv)
15632                    .arg(&mut *q)
15633                    .arg(&mut *k)
15634                    .arg(&mut *v)
15635                    .arg(&nc)
15636                    .arg(&rqi)
15637                    .arg(&rki)
15638                    .arg(pos)
15639                    .arg(&nhq)
15640                    .arg(&nhk)
15641                    .arg(&theta_scale)
15642                    .arg(&freq_scale)
15643                    .arg(&null)
15644                    .arg(&eps);
15645                unsafe {
15646                    b.launch(cfg)?;
15647                }
15648            }
15649        }
15650        Ok(())
15651    }
15652
15653    /// FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the
15654    /// rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins
15655    /// (format-flavored quant tail) — `g` must mirror the append path's flavor exactly.
15656    /// `n_rot` is the layer's derived rotary width and MUST equal `head_dim`
15657    /// ([`Engine::full_width_rope_only`]).
15658    #[allow(clippy::too_many_arguments)]
15659    pub fn rms_norm_qkv_rope_append_dc(
15660        &self,
15661        q0: &CudaSlice<f32>,
15662        k0: &CudaSlice<f32>,
15663        v0: &CudaSlice<f32>,
15664        wq: &CudaSlice<f32>,
15665        wk: &CudaSlice<f32>,
15666        wv: &CudaSlice<f32>,
15667        q: &mut CudaSlice<f32>,
15668        k: &mut CudaSlice<f32>,
15669        v: &mut CudaSlice<f32>,
15670        head_dim: usize,
15671        n_rot: usize,
15672        rq: usize,
15673        rk: usize,
15674        pos: &CudaSlice<i32>,
15675        nh_q: usize,
15676        nh_k: usize,
15677        base: f32,
15678        freq_scale: f32,
15679        ff: Option<&CudaSlice<f32>>,
15680        eps: f32,
15681        kc: &mut CudaSlice<u8>,
15682        vc: &mut CudaSlice<u8>,
15683        t_dev: &CudaSlice<i32>,
15684        k_tok_bytes: usize,
15685        v_tok_bytes: usize,
15686        g: bool,
15687    ) -> Result<(), Box<dyn std::error::Error>> {
15688        Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
15689        let rows = rq + rk + rk;
15690        let theta_scale = base.powf(-2.0 / head_dim as f32);
15691        let (nc, rqi, rki, nhq, nhk) = (
15692            head_dim as i32,
15693            rq as i32,
15694            rk as i32,
15695            nh_q as i32,
15696            nh_k as i32,
15697        );
15698        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15699        if Self::pdl_on() && Self::pdl_wb_on() {
15700            use cudarc::driver::{DevicePtr, DevicePtrMut};
15701            let s = &self.gpu.stream();
15702            let (p0, _a0) = q0.device_ptr(s);
15703            let (p1, _a1) = k0.device_ptr(s);
15704            let (p2, _a2) = v0.device_ptr(s);
15705            let (pwq, _a3) = wq.device_ptr(s);
15706            let (pwk, _a4) = wk.device_ptr(s);
15707            let (pwv, _a5) = wv.device_ptr(s);
15708            let (pq, _a6) = q.device_ptr_mut(s);
15709            let (pk, _a7) = k.device_ptr_mut(s);
15710            let (pv, _a8) = v.device_ptr_mut(s);
15711            let (pp, _a9) = pos.device_ptr(s);
15712            let pff: u64 = match ff {
15713                Some(t) => {
15714                    let (p, _gg) = t.device_ptr(s);
15715                    p
15716                }
15717                None => 0,
15718            };
15719            let (pkc, _a10) = kc.device_ptr_mut(s);
15720            let (pvc, _a11) = vc.device_ptr_mut(s);
15721            let (pt, _a12) = t_dev.device_ptr(s);
15722            let mut ps = [
15723                &p0 as *const _ as *mut std::ffi::c_void,
15724                &p1 as *const _ as *mut _,
15725                &p2 as *const _ as *mut _,
15726                &pwq as *const _ as *mut _,
15727                &pwk as *const _ as *mut _,
15728                &pwv as *const _ as *mut _,
15729                &pq as *const _ as *mut _,
15730                &pk as *const _ as *mut _,
15731                &pv as *const _ as *mut _,
15732                &nc as *const _ as *mut _,
15733                &rqi as *const _ as *mut _,
15734                &rki as *const _ as *mut _,
15735                &pp as *const _ as *mut _,
15736                &nhq as *const _ as *mut _,
15737                &nhk as *const _ as *mut _,
15738                &theta_scale as *const _ as *mut _,
15739                &freq_scale as *const _ as *mut _,
15740                &pff as *const _ as *mut _,
15741                &eps as *const _ as *mut _,
15742                &pkc as *const _ as *mut _,
15743                &pvc as *const _ as *mut _,
15744                &pt as *const _ as *mut _,
15745                &ktb as *const _ as *mut _,
15746                &vtb as *const _ as *mut _,
15747            ];
15748            unsafe {
15749                self.launch_pdl_flash(
15750                    g,
15751                    "rms_norm_qkv_rope_append_dc_f32",
15752                    (rows as u32, 1, 1),
15753                    (rms_block(), 1, 1),
15754                    0,
15755                    &mut ps,
15756                )?;
15757            }
15758            return Ok(());
15759        }
15760        let f = if g {
15761            self.func_g("rms_norm_qkv_rope_append_dc_f32")
15762        } else {
15763            self.func("rms_norm_qkv_rope_append_dc_f32")
15764        };
15765        let cfg = LaunchConfig {
15766            grid_dim: (rows as u32, 1, 1),
15767            block_dim: (rms_block(), 1, 1),
15768            shared_mem_bytes: 0,
15769        };
15770        let __s_b = self.gpu.stream();
15771        let mut b = __s_b.launch_builder(&f);
15772        match ff {
15773            Some(t) => {
15774                b.arg(q0)
15775                    .arg(k0)
15776                    .arg(v0)
15777                    .arg(wq)
15778                    .arg(wk)
15779                    .arg(wv)
15780                    .arg(&mut *q)
15781                    .arg(&mut *k)
15782                    .arg(&mut *v)
15783                    .arg(&nc)
15784                    .arg(&rqi)
15785                    .arg(&rki)
15786                    .arg(pos)
15787                    .arg(&nhq)
15788                    .arg(&nhk)
15789                    .arg(&theta_scale)
15790                    .arg(&freq_scale)
15791                    .arg(t)
15792                    .arg(&eps)
15793                    .arg(&mut *kc)
15794                    .arg(&mut *vc)
15795                    .arg(t_dev)
15796                    .arg(&ktb)
15797                    .arg(&vtb);
15798                unsafe {
15799                    b.launch(cfg)?;
15800                }
15801            }
15802            None => {
15803                let null: u64 = 0;
15804                b.arg(q0)
15805                    .arg(k0)
15806                    .arg(v0)
15807                    .arg(wq)
15808                    .arg(wk)
15809                    .arg(wv)
15810                    .arg(&mut *q)
15811                    .arg(&mut *k)
15812                    .arg(&mut *v)
15813                    .arg(&nc)
15814                    .arg(&rqi)
15815                    .arg(&rki)
15816                    .arg(pos)
15817                    .arg(&nhq)
15818                    .arg(&nhk)
15819                    .arg(&theta_scale)
15820                    .arg(&freq_scale)
15821                    .arg(&null)
15822                    .arg(&eps)
15823                    .arg(&mut *kc)
15824                    .arg(&mut *vc)
15825                    .arg(t_dev)
15826                    .arg(&ktb)
15827                    .arg(&vtb);
15828                unsafe {
15829                    b.launch(cfg)?;
15830                }
15831            }
15832        }
15833        Ok(())
15834    }
15835
15836    /// wave-2 fold: a + b with the sum emitted q8_1 alongside f32.
15837    /// Host-len twin of `rms_norm_qkv_rope_append_dc` (zoo-fusion arc): the eager decode
15838    /// arm tracks kv length on host (`kvl.len`), so the append slot rides the launch arg
15839    /// instead of the device counter. Kernel body is SHARED with the _dc entry (one
15840    /// inlined body — bit-identical to the rms_norm_qkv_rope + append pair it replaces,
15841    /// same law as the dc fold). `n_rot` is the layer's derived rotary width and MUST equal
15842    /// `head_dim` ([`Engine::full_width_rope_only`]).
15843    #[allow(clippy::too_many_arguments)]
15844    pub fn rms_norm_qkv_rope_append(
15845        &self,
15846        q0: &CudaSlice<f32>,
15847        k0: &CudaSlice<f32>,
15848        v0: &CudaSlice<f32>,
15849        wq: &CudaSlice<f32>,
15850        wk: &CudaSlice<f32>,
15851        wv: &CudaSlice<f32>,
15852        q: &mut CudaSlice<f32>,
15853        k: &mut CudaSlice<f32>,
15854        v: &mut CudaSlice<f32>,
15855        head_dim: usize,
15856        n_rot: usize,
15857        rq: usize,
15858        rk: usize,
15859        pos: &CudaSlice<i32>,
15860        nh_q: usize,
15861        nh_k: usize,
15862        base: f32,
15863        freq_scale: f32,
15864        ff: Option<&CudaSlice<f32>>,
15865        eps: f32,
15866        kc: &mut CudaSlice<u8>,
15867        vc: &mut CudaSlice<u8>,
15868        t: usize,
15869        k_tok_bytes: usize,
15870        v_tok_bytes: usize,
15871        g: bool,
15872    ) -> Result<(), Box<dyn std::error::Error>> {
15873        Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
15874        let rows = rq + rk + rk;
15875        let theta_scale = base.powf(-2.0 / head_dim as f32);
15876        let (nc, rqi, rki, nhq, nhk) = (
15877            head_dim as i32,
15878            rq as i32,
15879            rk as i32,
15880            nh_q as i32,
15881            nh_k as i32,
15882        );
15883        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15884        let ti = t as i32;
15885        if Self::pdl_on() && Self::pdl_wb_on() {
15886            use cudarc::driver::{DevicePtr, DevicePtrMut};
15887            let s = &self.gpu.stream();
15888            let (p0, _a0) = q0.device_ptr(s);
15889            let (p1, _a1) = k0.device_ptr(s);
15890            let (p2, _a2) = v0.device_ptr(s);
15891            let (pwq, _a3) = wq.device_ptr(s);
15892            let (pwk, _a4) = wk.device_ptr(s);
15893            let (pwv, _a5) = wv.device_ptr(s);
15894            let (pq, _a6) = q.device_ptr_mut(s);
15895            let (pk, _a7) = k.device_ptr_mut(s);
15896            let (pv, _a8) = v.device_ptr_mut(s);
15897            let (pp, _a9) = pos.device_ptr(s);
15898            let pff: u64 = match ff {
15899                Some(t) => {
15900                    let (p, _gg) = t.device_ptr(s);
15901                    p
15902                }
15903                None => 0,
15904            };
15905            let (pkc, _a10) = kc.device_ptr_mut(s);
15906            let (pvc, _a11) = vc.device_ptr_mut(s);
15907            let mut ps = [
15908                &p0 as *const _ as *mut std::ffi::c_void,
15909                &p1 as *const _ as *mut _,
15910                &p2 as *const _ as *mut _,
15911                &pwq as *const _ as *mut _,
15912                &pwk as *const _ as *mut _,
15913                &pwv as *const _ as *mut _,
15914                &pq as *const _ as *mut _,
15915                &pk as *const _ as *mut _,
15916                &pv as *const _ as *mut _,
15917                &nc as *const _ as *mut _,
15918                &rqi as *const _ as *mut _,
15919                &rki as *const _ as *mut _,
15920                &pp as *const _ as *mut _,
15921                &nhq as *const _ as *mut _,
15922                &nhk as *const _ as *mut _,
15923                &theta_scale as *const _ as *mut _,
15924                &freq_scale as *const _ as *mut _,
15925                &pff as *const _ as *mut _,
15926                &eps as *const _ as *mut _,
15927                &pkc as *const _ as *mut _,
15928                &pvc as *const _ as *mut _,
15929                &ti as *const _ as *mut _,
15930                &ktb as *const _ as *mut _,
15931                &vtb as *const _ as *mut _,
15932            ];
15933            unsafe {
15934                self.launch_pdl_flash(
15935                    g,
15936                    "rms_norm_qkv_rope_append_f32",
15937                    (rows as u32, 1, 1),
15938                    (rms_block(), 1, 1),
15939                    0,
15940                    &mut ps,
15941                )?;
15942            }
15943            return Ok(());
15944        }
15945        let f = if g {
15946            self.func_g("rms_norm_qkv_rope_append_f32")
15947        } else {
15948            self.func("rms_norm_qkv_rope_append_f32")
15949        };
15950        let cfg = LaunchConfig {
15951            grid_dim: (rows as u32, 1, 1),
15952            block_dim: (rms_block(), 1, 1),
15953            shared_mem_bytes: 0,
15954        };
15955        let __s_b = self.gpu.stream();
15956        let mut b = __s_b.launch_builder(&f);
15957        let null: u64 = 0;
15958        b.arg(q0)
15959            .arg(k0)
15960            .arg(v0)
15961            .arg(wq)
15962            .arg(wk)
15963            .arg(wv)
15964            .arg(&mut *q)
15965            .arg(&mut *k)
15966            .arg(&mut *v)
15967            .arg(&nc)
15968            .arg(&rqi)
15969            .arg(&rki)
15970            .arg(pos)
15971            .arg(&nhq)
15972            .arg(&nhk)
15973            .arg(&theta_scale)
15974            .arg(&freq_scale);
15975        match ff {
15976            Some(t) => {
15977                b.arg(t);
15978            }
15979            None => {
15980                b.arg(&null);
15981            }
15982        }
15983        b.arg(&eps)
15984            .arg(&mut *kc)
15985            .arg(&mut *vc)
15986            .arg(&ti)
15987            .arg(&ktb)
15988            .arg(&vtb);
15989        unsafe {
15990            b.launch(cfg)?;
15991        }
15992        Ok(())
15993    }
15994
15995    pub fn add_q8_1(
15996        &self,
15997        a: &CudaSlice<f32>,
15998        b: &CudaSlice<f32>,
15999        res: &mut CudaSlice<f32>,
16000        ncols: usize,
16001        nrows: usize,
16002    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16003        debug_assert!(ncols.is_multiple_of(128));
16004        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
16005        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
16006        let f = self.func("add_q8_1_f32");
16007        let cfg = LaunchConfig {
16008            grid_dim: (nrows as u32, 1, 1),
16009            block_dim: (rms_block(), 1, 1),
16010            shared_mem_bytes: 0,
16011        };
16012        let nc = ncols as i32;
16013        let __s_b2 = self.gpu.stream();
16014        let mut b2 = __s_b2.launch_builder(&f);
16015        b2.arg(a)
16016            .arg(b)
16017            .arg(&mut *res)
16018            .arg(&mut out_q)
16019            .arg(&mut out_d)
16020            .arg(&nc);
16021        unsafe {
16022            b2.launch(cfg)?;
16023        }
16024        Ok((out_q, out_d))
16025    }
16026
16027    /// E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair
16028    /// in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config
16029    /// as both parents (bit-identity: identical reduction + quad-walk quantize).
16030    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
16031    pub fn rms_pre_add_q8_1(
16032        &self,
16033        a: &CudaSlice<f32>,
16034        wa: &CudaSlice<f32>,
16035        b: &CudaSlice<f32>,
16036        res: &mut CudaSlice<f32>,
16037        ncols: usize,
16038        nrows: usize,
16039        eps: f32,
16040    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16041        debug_assert!(ncols.is_multiple_of(128));
16042        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
16043        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
16044        let f = self.func("rms_pre_add_q8_1_f32");
16045        let cfg = LaunchConfig {
16046            grid_dim: (nrows as u32, 1, 1),
16047            block_dim: (rms_block(), 1, 1),
16048            shared_mem_bytes: 0,
16049        };
16050        let (nc, ep) = (ncols as i32, eps);
16051        let __s_b2 = self.gpu.stream();
16052        let mut b2 = __s_b2.launch_builder(&f);
16053        b2.arg(a)
16054            .arg(wa)
16055            .arg(b)
16056            .arg(&mut *res)
16057            .arg(&mut out_q)
16058            .arg(&mut out_d)
16059            .arg(&nc)
16060            .arg(&ep);
16061        unsafe {
16062            b2.launch(cfg)?;
16063        }
16064        Ok((out_q, out_d))
16065    }
16066
16067    /// L2 norm per row (head_dim), no weight.
16068    /// PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config
16069    /// seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.
16070    pub fn l2_v2_on(ncols: usize) -> bool {
16071        ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
16072    }
16073
16074    pub fn l2_norm_pp(
16075        &self,
16076        x: &CudaSlice<f32>,
16077        dst: &mut CudaSlice<f32>,
16078        dst16: Option<&mut CudaSlice<u8>>,
16079        ncols: usize,
16080        nrows: usize,
16081        eps: f32,
16082    ) -> Result<(), Box<dyn std::error::Error>> {
16083        if Self::l2_v2_on(ncols) {
16084            let f = self.func("l2_norm_pp_v2_f32");
16085            let rows_per_block = 8u32; // 256 threads = 8 warps = 8 rows
16086            let cfg = LaunchConfig {
16087                grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
16088                block_dim: (256, 1, 1),
16089                shared_mem_bytes: 0,
16090            };
16091            let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
16092            // mirror-fold: bf16 twin address by value (0 = skip; matches the nullable param)
16093            let d16: u64 = match dst16 {
16094                Some(d) => self.addr_u8(d),
16095                None => 0,
16096            };
16097            let __s_b = self.gpu.stream();
16098            let mut b = __s_b.launch_builder(&f);
16099            b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
16100            unsafe {
16101                b.launch(cfg)?;
16102            }
16103            return Ok(());
16104        }
16105        self.l2_norm(x, dst, ncols, nrows, eps)
16106    }
16107
16108    pub fn l2_norm(
16109        &self,
16110        x: &CudaSlice<f32>,
16111        dst: &mut CudaSlice<f32>,
16112        ncols: usize,
16113        nrows: usize,
16114        eps: f32,
16115    ) -> Result<(), Box<dyn std::error::Error>> {
16116        let f = self.func("l2_norm_f32");
16117        let cfg = LaunchConfig {
16118            grid_dim: (nrows as u32, 1, 1),
16119            block_dim: (256, 1, 1),
16120            shared_mem_bytes: 0,
16121        };
16122        let (nc, e) = (ncols as i32, eps);
16123        let __s_b = self.gpu.stream();
16124        let mut b = __s_b.launch_builder(&f);
16125        b.arg(x).arg(dst).arg(&nc).arg(&e);
16126        unsafe {
16127            b.launch(cfg)?;
16128        }
16129        Ok(())
16130    }
16131
16132    /// L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32's
16133    /// per-warp L2 norm. The verify path MUST use this to match decode's FP accumulation order:
16134    /// l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element
16135    /// squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that
16136    /// propagate through gdn_scan and flip argmax on marginal logits.
16137    pub fn l2_norm_decode(
16138        &self,
16139        x: &CudaSlice<f32>,
16140        dst: &mut CudaSlice<f32>,
16141        ncols: usize,
16142        nrows: usize,
16143        eps: f32,
16144    ) -> Result<(), Box<dyn std::error::Error>> {
16145        let f = self.func("l2_norm_f32");
16146        let cfg = LaunchConfig {
16147            grid_dim: (nrows as u32, 1, 1),
16148            block_dim: (32, 1, 1),
16149            shared_mem_bytes: 0,
16150        };
16151        let (nc, e) = (ncols as i32, eps);
16152        let __s_b = self.gpu.stream();
16153        let mut b = __s_b.launch_builder(&f);
16154        b.arg(x).arg(dst).arg(&nc).arg(&e);
16155        unsafe {
16156            b.launch(cfg)?;
16157        }
16158        Ok(())
16159    }
16160
16161    /// RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].
16162    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
16163    pub fn rope_neox(
16164        &self,
16165        x: &mut CudaSlice<f32>,
16166        pos: &CudaSlice<i32>,
16167        head_dim: usize,
16168        n_dims: usize,
16169        n_heads: usize,
16170        n_tokens: usize,
16171        freq_base: f32,
16172        freq_scale: f32,
16173    ) -> Result<(), Box<dyn std::error::Error>> {
16174        let f = self.func("rope_neox_f32");
16175        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16176        let grid = (n_heads * n_tokens) as u32;
16177        let cfg = LaunchConfig {
16178            grid_dim: (grid, 1, 1),
16179            block_dim: ((head_dim / 2) as u32, 1, 1),
16180            shared_mem_bytes: 0,
16181        };
16182        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
16183        let __s_b = self.gpu.stream();
16184        let mut b = __s_b.launch_builder(&f);
16185        b.arg(x)
16186            .arg(pos)
16187            .arg(&hd)
16188            .arg(&nd)
16189            .arg(&nh)
16190            .arg(&theta_scale)
16191            .arg(&freq_scale);
16192        unsafe {
16193            b.launch(cfg)?;
16194        }
16195        Ok(())
16196    }
16197
16198    /// RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).
16199    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
16200    pub fn rope_neox_ff(
16201        &self,
16202        x: &mut CudaSlice<f32>,
16203        pos: &CudaSlice<i32>,
16204        head_dim: usize,
16205        n_dims: usize,
16206        n_heads: usize,
16207        n_tokens: usize,
16208        freq_base: f32,
16209        freq_scale: f32,
16210        ff: &CudaSlice<f32>,
16211    ) -> Result<(), Box<dyn std::error::Error>> {
16212        let f = self.func("rope_neox_ff_f32");
16213        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16214        let grid = (n_heads * n_tokens) as u32;
16215        let cfg = LaunchConfig {
16216            grid_dim: (grid, 1, 1),
16217            block_dim: ((head_dim / 2) as u32, 1, 1),
16218            shared_mem_bytes: 0,
16219        };
16220        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
16221        let __s_b = self.gpu.stream();
16222        let mut b = __s_b.launch_builder(&f);
16223        b.arg(x)
16224            .arg(pos)
16225            .arg(&hd)
16226            .arg(&nd)
16227            .arg(&nh)
16228            .arg(&theta_scale)
16229            .arg(&freq_scale)
16230            .arg(ff);
16231        unsafe {
16232            b.launch(cfg)?;
16233        }
16234        Ok(())
16235    }
16236
16237    /// RoPE NEOX with per-dim freq factors AND the YaRN attention factor on cos/sin
16238    /// (qwen4_exp yarn lane — `rope_neox_ffm_f32`; ff = yarn_frequency_divisors, mscale =
16239    /// yarn_attention_factor). Identity inputs (ones, 1.0) reproduce `rope_neox` bit-for-bit.
16240    #[allow(clippy::too_many_arguments)]
16241    pub fn rope_neox_ffm(
16242        &self,
16243        x: &mut CudaSlice<f32>,
16244        pos: &CudaSlice<i32>,
16245        head_dim: usize,
16246        n_dims: usize,
16247        n_heads: usize,
16248        n_tokens: usize,
16249        freq_base: f32,
16250        freq_scale: f32,
16251        ff: &CudaSlice<f32>,
16252        mscale: f32,
16253    ) -> Result<(), Box<dyn std::error::Error>> {
16254        let f = self.func("rope_neox_ffm_f32");
16255        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16256        let grid = (n_heads * n_tokens) as u32;
16257        let cfg = LaunchConfig {
16258            grid_dim: (grid, 1, 1),
16259            block_dim: ((head_dim / 2) as u32, 1, 1),
16260            shared_mem_bytes: 0,
16261        };
16262        let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
16263        let __s_b = self.gpu.stream();
16264        let mut b = __s_b.launch_builder(&f);
16265        b.arg(x)
16266            .arg(pos)
16267            .arg(&hd)
16268            .arg(&nd)
16269            .arg(&nh)
16270            .arg(&theta_scale)
16271            .arg(&freq_scale)
16272            .arg(ff)
16273            .arg(&mscale);
16274        unsafe {
16275            b.launch(cfg)?;
16276        }
16277        Ok(())
16278    }
16279
16280    /// gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).
16281    #[allow(clippy::too_many_arguments)]
16282    pub fn rope_neox2(
16283        &self,
16284        q: &mut CudaSlice<f32>,
16285        k: &mut CudaSlice<f32>,
16286        pos: &CudaSlice<i32>,
16287        head_dim: usize,
16288        n_dims: usize,
16289        nh_q: usize,
16290        nh_k: usize,
16291        n_tokens: usize,
16292        freq_base: f32,
16293        freq_scale: f32,
16294        ff: Option<&CudaSlice<f32>>,
16295    ) -> Result<(), Box<dyn std::error::Error>> {
16296        let f = self.func("rope_neox2_f32");
16297        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16298        let grid = ((nh_q + nh_k) * n_tokens) as u32;
16299        let cfg = LaunchConfig {
16300            grid_dim: (grid, 1, 1),
16301            block_dim: ((head_dim / 2) as u32, 1, 1),
16302            shared_mem_bytes: 0,
16303        };
16304        let (hd, nd, nq, nk, nt) = (
16305            head_dim as i32,
16306            n_dims as i32,
16307            nh_q as i32,
16308            nh_k as i32,
16309            n_tokens as i32,
16310        );
16311        let __s_b = self.gpu.stream();
16312        let mut b = __s_b.launch_builder(&f);
16313        b.arg(q)
16314            .arg(k)
16315            .arg(pos)
16316            .arg(&hd)
16317            .arg(&nd)
16318            .arg(&nq)
16319            .arg(&nk)
16320            .arg(&nt)
16321            .arg(&theta_scale)
16322            .arg(&freq_scale);
16323        match ff {
16324            Some(ffv) => {
16325                b.arg(ffv);
16326                unsafe {
16327                    b.launch(cfg)?;
16328                }
16329            }
16330            None => {
16331                let null: u64 = 0;
16332                b.arg(&null);
16333                unsafe {
16334                    b.launch(cfg)?;
16335                }
16336            }
16337        }
16338        Ok(())
16339    }
16340
16341    /// gemma4 R1: dst = GELU_tanh(gate) * up.
16342    pub fn gelu_tanh_mul(
16343        &self,
16344        gate: &CudaSlice<f32>,
16345        up: &CudaSlice<f32>,
16346        dst: &mut CudaSlice<f32>,
16347        n: usize,
16348    ) -> Result<(), Box<dyn std::error::Error>> {
16349        let f = self.func("gelu_tanh_mul_f32");
16350        let cfg = LaunchConfig::for_num_elems(n as u32);
16351        let ni = n as i32;
16352        let __s_b = self.gpu.stream();
16353        let mut b = __s_b.launch_builder(&f);
16354        b.arg(gate).arg(up).arg(dst).arg(&ni);
16355        unsafe {
16356            b.launch(cfg)?;
16357        }
16358        Ok(())
16359    }
16360
16361    pub fn silu_mul(
16362        &self,
16363        gate: &CudaSlice<f32>,
16364        up: &CudaSlice<f32>,
16365        dst: &mut CudaSlice<f32>,
16366        n: usize,
16367    ) -> Result<(), Box<dyn std::error::Error>> {
16368        let f = self.func("silu_mul_f32");
16369        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
16370        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16371        let ni = n as i32;
16372        let __s_b = self.gpu.stream();
16373        let mut b = __s_b.launch_builder(&f);
16374        b.arg(gate).arg(up).arg(dst).arg(&ni);
16375        unsafe {
16376            b.launch(cfg)?;
16377        }
16378        Ok(())
16379    }
16380
16381    /// SwiGLU twin using Memra's host-matching expf transcription.
16382    pub fn silu_mul_host_expf(
16383        &self,
16384        gate: &CudaSlice<f32>,
16385        up: &CudaSlice<f32>,
16386        dst: &mut CudaSlice<f32>,
16387        n: usize,
16388    ) -> Result<(), Box<dyn std::error::Error>> {
16389        let f = self.func("silu_mul_host_expf_f32");
16390        let cfg = LaunchConfig::for_num_elems(n as u32);
16391        let ni = n as i32;
16392        let __s_b = self.gpu.stream();
16393        let mut b = __s_b.launch_builder(&f);
16394        b.arg(gate).arg(up).arg(dst).arg(&ni);
16395        unsafe {
16396            b.launch(cfg)?;
16397        }
16398        Ok(())
16399    }
16400
16401    /// Step routed-expert clamp twin using Memra's host-matching expf transcription.
16402    pub fn silu_clamped_mul_host_expf(
16403        &self,
16404        gate: &CudaSlice<f32>,
16405        up: &CudaSlice<f32>,
16406        limit: f32,
16407        dst: &mut CudaSlice<f32>,
16408        n: usize,
16409    ) -> Result<(), Box<dyn std::error::Error>> {
16410        if !limit.is_finite() || limit <= 0.0 {
16411            return Err(
16412                format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
16413            );
16414        }
16415        let f = self.func("silu_clamped_mul_host_expf_f32");
16416        let cfg = LaunchConfig::for_num_elems(n as u32);
16417        let ni = n as i32;
16418        let __s_b = self.gpu.stream();
16419        let mut b = __s_b.launch_builder(&f);
16420        b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
16421        unsafe {
16422            b.launch(cfg)?;
16423        }
16424        Ok(())
16425    }
16426
16427    /// f16out twin of `silu_mul` (task #17): the epilogue also emits the fp16 GEMM operand
16428    /// for the down projection — kills the standalone convert pass. Bit-identical class.
16429    pub fn silu_mul_f16out(
16430        &self,
16431        gate: &CudaSlice<f32>,
16432        up: &CudaSlice<f32>,
16433        dst: &mut CudaSlice<f32>,
16434        dst16: &mut CudaSlice<u8>,
16435        n: usize,
16436    ) -> Result<(), Box<dyn std::error::Error>> {
16437        let f = self.func("silu_mul_f16out_f32");
16438        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16439        let ni = n as i32;
16440        let __s_b = self.gpu.stream();
16441        let mut b = __s_b.launch_builder(&f);
16442        b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
16443        unsafe {
16444            b.launch(cfg)?;
16445        }
16446        Ok(())
16447    }
16448
16449    /// FFN SwiGLU epilogue fusion (RANK3 LEVER 2): `dst = silu(gate*gs) * (up*us)` in ONE launch,
16450    /// folding the per-tensor NVFP4 macro-scale (`gs`,`us`) that would otherwise be two separate
16451    /// `scale_inplace` launches on the gate/up matmul outputs. BIT-IDENTICAL to
16452    /// scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in
16453    /// identical order. For non-NVFP4 weights gs==us==1.0 -> identical to `silu_mul`. Net: -2
16454    /// launches per dense FFN layer (the gate+up post-matmul scales).
16455    pub fn silu_mul_scaled(
16456        &self,
16457        gate: &CudaSlice<f32>,
16458        up: &CudaSlice<f32>,
16459        gs: f32,
16460        us: f32,
16461        dst: &mut CudaSlice<f32>,
16462        n: usize,
16463    ) -> Result<(), Box<dyn std::error::Error>> {
16464        let f = self.func("silu_mul_scaled_f32");
16465        let cfg = LaunchConfig::for_num_elems(n as u32);
16466        let ni = n as i32;
16467        let (gsf, usf) = (gs, us);
16468        let __s_b = self.gpu.stream();
16469        let mut b = __s_b.launch_builder(&f);
16470        b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
16471        unsafe {
16472            b.launch(cfg)?;
16473        }
16474        Ok(())
16475    }
16476
16477    /// swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp
16478    /// ggml_cuda_op_swiglu_oai_single. `dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us))`.
16479    /// gs/us fold the NVFP4 macro-scales exactly like `silu_mul_scaled`.
16480    #[allow(clippy::too_many_arguments)]
16481    pub fn swigluoai_mul_scaled(
16482        &self,
16483        gate: &CudaSlice<f32>,
16484        up: &CudaSlice<f32>,
16485        gs: f32,
16486        us: f32,
16487        alpha: f32,
16488        limit: f32,
16489        dst: &mut CudaSlice<f32>,
16490        n: usize,
16491    ) -> Result<(), Box<dyn std::error::Error>> {
16492        let f = self.func("swigluoai_mul_scaled_f32");
16493        let cfg = LaunchConfig::for_num_elems(n as u32);
16494        let ni = n as i32;
16495        let __s_b = self.gpu.stream();
16496        let mut b = __s_b.launch_builder(&f);
16497        b.arg(gate)
16498            .arg(up)
16499            .arg(&gs)
16500            .arg(&us)
16501            .arg(&alpha)
16502            .arg(&limit)
16503            .arg(dst)
16504            .arg(&ni);
16505        unsafe {
16506            b.launch(cfg)?;
16507        }
16508        Ok(())
16509    }
16510
16511    /// RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of `act`
16512    /// directly (aq int8 [n] + ad f32 [n/32]), so ffn_down's standalone `quantize_q8_1` launch is
16513    /// removed — the down-proj activation has one consumer, so the quant folds into the producer for
16514    /// free (no extra HBM read; no f32 `act` write). gs/us fold the gate/up NVFP4 macro-scales like
16515    /// `silu_mul_scaled`. BIT-IDENTICAL q8_1 to silu_mul_scaled(...) then quantize_q8_1(...). Only
16516    /// valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks `uses_q8_1_fast(ffn_down)`.
16517    /// n must be a multiple of 32 (n_ff always is).
16518    pub fn silu_mul_scaled_q8_1(
16519        &self,
16520        gate: &CudaSlice<f32>,
16521        up: &CudaSlice<f32>,
16522        gs: f32,
16523        us: f32,
16524        n: usize,
16525    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16526        let f = self.func("silu_mul_scaled_q8_1");
16527        let nblk = n / 32;
16528        let mut aq = self.alloc_uninit::<i8>(n)?; // full-overwrite output
16529        let mut ad = self.alloc_uninit::<f32>(nblk)?; // full-overwrite output
16530        // WARP-PER-BLOCK kernel: one warp (32 lanes) per 32-block -> n threads total.
16531        let cfg = LaunchConfig::for_num_elems(n as u32);
16532        let (gsf, usf, ni) = (gs, us, n as i32);
16533        let __s_b = self.gpu.stream();
16534        let mut b = __s_b.launch_builder(&f);
16535        b.arg(gate)
16536            .arg(up)
16537            .arg(&gsf)
16538            .arg(&usf)
16539            .arg(&mut aq)
16540            .arg(&mut ad)
16541            .arg(&ni);
16542        unsafe {
16543            b.launch(cfg)?;
16544        }
16545        Ok((aq, ad))
16546    }
16547
16548    pub fn add(
16549        &self,
16550        a: &CudaSlice<f32>,
16551        b_in: &CudaSlice<f32>,
16552        dst: &mut CudaSlice<f32>,
16553        n: usize,
16554    ) -> Result<(), Box<dyn std::error::Error>> {
16555        let f = self.func("add_f32");
16556        // float4 kernel: one thread per 4 elements (tail handled in-kernel)
16557        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16558        let ni = n as i32;
16559        let __s_bld = self.gpu.stream();
16560        let mut bld = __s_bld.launch_builder(&f);
16561        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
16562        unsafe {
16563            bld.launch(cfg)?;
16564        }
16565        Ok(())
16566    }
16567
16568    pub fn mul(
16569        &self,
16570        a: &CudaSlice<f32>,
16571        b_in: &CudaSlice<f32>,
16572        dst: &mut CudaSlice<f32>,
16573        n: usize,
16574    ) -> Result<(), Box<dyn std::error::Error>> {
16575        let f = self.func("mul_f32");
16576        let cfg = LaunchConfig::for_num_elems(n as u32);
16577        let ni = n as i32;
16578        let __s_bld = self.gpu.stream();
16579        let mut bld = __s_bld.launch_builder(&f);
16580        bld.arg(a).arg(b_in).arg(dst).arg(&ni);
16581        unsafe {
16582            bld.launch(cfg)?;
16583        }
16584        Ok(())
16585    }
16586
16587    /// Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and
16588    /// float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.
16589    pub fn matmul(
16590        &self,
16591        w: &crate::model::GpuTensor,
16592        x: &CudaSlice<f32>,
16593        m: usize,
16594    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16595        use crate::model::GpuTensor;
16596        let in_f = w.in_features();
16597        let out_f = w.out_features();
16598        // PREFILL (T>1) ROOT FIX: batched tensor-core int8 GEMM. Decodes each weight tile to int8
16599        // in smem ONCE and reuses across all tokens via mma — vs the dp4a matvec's per-token weight
16600        // re-read. Only the 4 daily-hot dtypes; m=1 decode keeps dp4a (it's bandwidth-bound, mma
16601        // gives nothing). Quantize the activation once here then call the GEMM.
16602        // m cutoff FIXED at 16: the m=4 MMA-verify A/B (2026-07-06, was MEMRA_GEMM_M) measured
16603        // NEGATIVE — the MMA tile grid starves at m=4 (BN=256 -> grid.y=1) and its FP order
16604        // shifted verify argmax at tight margins. Do not lower without re-running that battery.
16605        #[allow(non_snake_case)]
16606        // VERIFY-EXACT scope pushes the GEMM crossover out of reach (usize::MAX) — the
16607        // t>=16 dflash verify must ride the decode-exact batched class (parity law).
16608        let GEMM_M_THRESHOLD = if self.verify_exact_on() {
16609            usize::MAX
16610        } else {
16611            16usize
16612        };
16613
16614        // PREFILL GEMM (m>=16). ACCURACY-FIRST dispatch (2026-06-28, prefill-gemm-beat-research wf
16615        // wllbyo6vc step 1): the int8 W4A8 GEMM (qmatvec_gemm, q8_1 activation, s32 accumulate) is
16616        // ACCURATE (prefill logit maxdiff 0.159, < dp4a 0.55) and the default. The FP4 W4A4 mxf4 path
16617        // (try_fp4_gemm) quantizes the ACTIVATION to e2m1 4-bit (8 magnitude levels) -> maxdiff 1.0
16618        // when combined — a real accuracy loss, NOT a math bug. So FP4-W4A4 is taken ONLY under the
16619        // explicit MEMRA_FP4 opt-in AND it must come SECOND (int8 W4A8 is the correct default for NVFP4).
16620        // The workflow plan rebuilds the FP4 path (kill per-K repack, widen K, deepen pipeline, TMA) to
16621        // be both fast AND accurate; until then NVFP4 prefill defaults to the accurate int8 GEMM.
16622        // TINY-OUT_F GUARD (2026-06-28, ncu trace): the tiling GEMM's grid is (ceil(out_f/BM=64),
16623        // ceil(m/BN=256)). For tiny out_f (ssm_beta/ssm_alpha out_f=num_v_heads~32), grid.x=1 -> only
16624        // ceil(m/256) CTAs (e.g. 2 for m=512) on 82 SMs = 0.39% SM throughput, 852us EACH (measured
16625        // worst offender). The dp4a path grids (out_f, m) = far more CTAs, filling the GPU. So route
16626        // out_f < 2*BM to dp4a (skip the tiling GEMM which structurally can't fill the SMs here).
16627        const GEMM_MIN_OUT_F: usize = 128; // 2*BM; below this the GEMM grid.x starves the 82 SMs
16628        // VENDORED llama MMQ prefill GEMMs. NVFP4 W4A8 is DEFAULT-ON (2026-07-05 flip: same int8
16629        // accuracy class as the int8 GEMM below at ~1.9x pp512, rp-loader coexists with the A6
16630        // repack; MEMRA_MMQ_W4A8=0 = escape hatch). W4A4 mxf4nvf4 + Q4_K/Q5_K stay behind MEMRA_MMQ=1.
16631        // The env policy lives in mmq_supports/qmatvec_mmq. Feeds raw f32 activation `x` (the
16632        // launcher quantizes internally). out_f>=MMQ_Y/2 keeps the tile grid from starving the SMs.
16633        // FP8-ACT PREFILL (MEMRA_PP_FP8=1, probe verdict 2026-07-08): F8-E4M3-origin projections
16634        // carry their raw e4m3 device bytes (the `fp8` operand stashed at load next to the Q8_0
16635        // re-encode) — cuBLASLt FP8 TN at 620-795 TF vs 47-72 TF for this class's int8 GEMM.
16636        // Weight side EXACT (checkpoint bytes); activation rides ONE per-batch e4m3 scale
16637        // (amax/448) folded with weight_scale in-GEMM. Prefill only; decode keeps Q8_0 untouched.
16638        if m >= GEMM_M_THRESHOLD {
16639            if let Some(y) = self.try_fp8_gemm(w, x, m)? {
16640                return Ok(y);
16641            }
16642            // PER-BLOCK FP8 MMQ (lane/fp8-mmq): the block-128 class try_fp8_gemm skips (cuBLASLt
16643            // takes no block grid on sm_120). Exact per block — the checkpoint's e4m3 bytes and its
16644            // f32 grid go into the tile unchanged. TWO SOURCES, TWO DEFAULTS: the load-time stash is
16645            // opt-in (MEMRA_FP8_MMQ=1), the native-resident QT_F8_E4M3_BLK grid is DEFAULT ON
16646            // (MEMRA_FP8_MMQ=0 reverts it to dequant-per-call) — see fp8_ffi.rs for why the same
16647            // tile defaults differently by operand source.
16648            if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
16649                return Ok(y);
16650            }
16651            // FP16-mirror prefill (MEMRA_PP_F16=1, probe 2026-07-26: 3.2-3.7x the MMQ class).
16652            // Mirror presence IS the gate (only built under the env). Decode never reaches here.
16653            if let Some(y) = self.try_f16_gemm(w, x, m)? {
16654                return Ok(y);
16655            }
16656        }
16657        // F8-E4M3 BLOCK-128 (QT_F8_E4M3_BLK, lane/fp8-blk128-decode). TWO arms, split at the SAME
16658        // m threshold the rest of this method uses:
16659        //   * m >= threshold (prefill): dequant-per-call to the ARM B' Q8_0 slab and recurse, so
16660        //     prefill keeps the floor's kernels AND the floor's bits (try_e4m3_blk_prefill).
16661        //   * m <  threshold: the native per-block GEMV — m=1 decode and the m=2..15 verify tiers.
16662        //     grid.y=m runs the exact m=1 program per (token,row), so the decode-parity law holds
16663        //     across every tier by construction with no batched twin needed.
16664        //
16665        // NOT gated on `fast`: this dtype has no dp4a twin and no Stage-A f32-dequant oracle (the
16666        // generic `deq()` switch has no block-scale input), exactly as QT_F8_E4M3 has none, so
16667        // MEMRA_FAST=0 cannot route it anywhere else. Placed before every GEMM/MMQ arm below
16668        // because gemm_supports/mmq_supports/mmvq_supports all deliberately REFUSE this qtype —
16669        // reaching the generic tail would panic rather than produce wrong numbers, and this pair of
16670        // arms is what makes sure it never gets there.
16671        if let GpuTensor::Quant { qtype, .. } = w
16672            && *qtype == QT_F8_E4M3_BLK
16673        {
16674            if m >= GEMM_M_THRESHOLD
16675                && let Some(y) = self.try_e4m3_blk_prefill(w, x, m)?
16676            {
16677                return Ok(y);
16678            }
16679            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16680            if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
16681                return Ok(y);
16682            }
16683        }
16684        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
16685            return self.qmatvec_mmq(w, x, m);
16686        }
16687        if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
16688            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16689            return self.qmatvec_gemm(w, &aq, &ad, m);
16690        }
16691        // FP4 W4A4 only as an explicit speed/accuracy tradeoff opt-in, and only if the int8 GEMM
16692        // above didn't already handle this weight (e.g. NVFP4 with in_f%64!=0, or MEMRA_NO_GEMM set).
16693        if m >= GEMM_M_THRESHOLD
16694            && let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)?
16695        {
16696            return Ok(y);
16697        }
16698        // Stage-B fast int8 dp4a is the DEFAULT since 2026-07-08 (it has been the daily path
16699        // for weeks; the old opt-in flag was a silent-slow-path landmine). MEMRA_FAST=0 reverts
16700        // to Stage-A f32-dequant (the correctness oracle path).
16701        let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
16702        // PERF-3 decode-GEMV: m=1 warp-per-row MMVQ (MEMRA_MMVQ). The big decode matvecs reach
16703        // `matmul` directly (ffn_down, lm_head output, wo), so route them here too — not only the
16704        // matmul_pre siblings. qmatvec_mmvq_raw quantizes the activation internally (q8_1) like the
16705        // _fast paths; the NVFP4 macro-scale is applied by the `scale != 1.0` block below.
16706        if m == 1
16707            && fast
16708            && let GpuTensor::Quant {
16709                bytes,
16710                qtype,
16711                row_bytes,
16712                rp,
16713                rp4,
16714                scale,
16715                ..
16716            } = w
16717            && self.mmvq_supports(*qtype)
16718        {
16719            // NVFP4 macro-scale rides the kernel's fused epilogue arg (one launch total);
16720            // non-NVFP4 has scale==1.0 so qmatvec_mmvq skips scale_inplace either way.
16721            // Q4_0 split-plane mirror (rp4): the decode arm reads it via the _rp twins.
16722            let (bytes, rp) = match rp4 {
16723                Some(m4) => (m4, true),
16724                None => (bytes, *rp),
16725            };
16726            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16727            return self.qmatvec_mmvq(
16728                bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
16729            );
16730        }
16731        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward's ffn_down, wo, and
16732        // lm_head `output` reach `matmul` directly at m=T=2..4). Walks the weight ONCE, dp4a vs all m
16733        // activation columns -> 1 weight read for m tokens (vs grid.y=m re-reading m times below). Quant
16734        // the activation once here (q8_1) like the _fast paths; macro-scale applied via the scale!=1.0
16735        // block below. MEMRA_NO_BATCHED -> per-m path.
16736        //
16737        // DECODE-PARITY GATE (2026-07-07, the 9B synth K=3/4/6 spec FAIL root cause): the batched
16738        // kernels are bit-identical per (token,row) to MMVQ's 32-thread warp reduce, NOT to the
16739        // dp4a kernels' 128-thread two-level reduce. Without MEMRA_MMVQ the m=1 decode chain rides
16740        // dp4a, so a verify riding batched here has a DIFFERENT FP order than the decode it must
16741        // match bit-for-bit — greedy spec flips at tight-margin tokens (the old HANDOVER "ENV LAW:
16742        // FAST+MMVQ both required" footgun, closed here). Parity law: the m>1 kernel CLASS must be
16743        // a pure function of (dtype, env) equal to the m=1 class — batched iff MMVQ. Without MMVQ
16744        // the verify falls to the per-m grid.y=m dp4a path below (each column = the exact m=1
16745        // dp4a program). MEMRA_MMVQ=1 (the daily config) is dispatch-unchanged.
16746        if (2..=16).contains(&m)
16747            && fast
16748            && std::env::var("MEMRA_NO_BATCHED").is_err()
16749            && (m <= 4 || Self::b8_enabled())
16750        {
16751            // b16 tier (2026-07-11, spec K>7): Q4_0/Q6_K have base+_rp b16 kernels; Q8_0's
16752            // b16 exists only as the split-plane _rp twin, so it joins iff the q8rp mirror
16753            // is present (rp4) — the mirror pick below then routes to the _rp family.
16754            // QT_F8_E4M3 joins unconditionally (lane/rp-on-st): its b16 IS the base kernel,
16755            // because the native e4m3 row layout is already aligned and needs no mirror.
16756            // NVFP4/Q4_K/Q8_0 all join unconditionally now (lane/rp-on-st): each has base + _rp
16757            // b16 twins, so either residency layout has its aligned form at this width. Q8_0's
16758            // old `rp4.is_some()` precondition is GONE — the mirror is a bandwidth lever, not the
16759            // exact tier's admission ticket (it was refusing FP8-ST over 23.9 MiB of ssm_beta).
16760            let m_ok = m <= 8
16761                || matches!(w, GpuTensor::Quant { qtype, .. }
16762                if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
16763                    || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
16764            if m_ok
16765                && let GpuTensor::Quant {
16766                    bytes,
16767                    qtype,
16768                    row_bytes,
16769                    rp,
16770                    rp4,
16771                    ..
16772                } = w
16773                && self.batched_supports(*qtype)
16774                && self.mmvq_supports(*qtype)
16775            {
16776                let (bytes, rp) = match rp4 {
16777                    Some(m4) => (m4, true),
16778                    None => (bytes, *rp),
16779                };
16780                let mcols = Self::batched_mcols(m);
16781                let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16782                let mut y = self.qmatvec_mmvq_batched(
16783                    bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
16784                )?;
16785                if let GpuTensor::Quant { scale, .. } = w
16786                    && *scale != 1.0
16787                {
16788                    self.scale_inplace(&mut y, *scale, m * out_f)?;
16789                }
16790                return Ok(y);
16791            }
16792        }
16793        // F8-E4M3 (MEMRA_ST_E4M3) catch-all for the m<16 band the arms above didn't take (m=9..15,
16794        // the K=8 verify tier; or m=2..8 under MEMRA_NO_BATCHED/MEMRA_B8=0): grid.y=m e4m3 mmvq —
16795        // the SAME per-(token,row) program as the m=1 decode launch (bit-identical by construction),
16796        // weight re-read m times (rare tier; exactness over bandwidth here). There is no _dp4a twin
16797        // for this dtype, so the generic match below must never see it under `fast`.
16798        if fast
16799            && let GpuTensor::Quant {
16800                bytes,
16801                qtype,
16802                row_bytes,
16803                scale,
16804                ..
16805            } = w
16806            && *qtype == QT_F8_E4M3
16807        {
16808            let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16809            return self.qmatvec_mmvq(
16810                bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
16811            );
16812        }
16813        let mut y = match w {
16814            GpuTensor::Quant {
16815                bytes,
16816                qtype,
16817                row_bytes,
16818                ..
16819            } if fast && *qtype == QT_Q8_0 => {
16820                self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16821            }
16822            GpuTensor::Quant {
16823                bytes,
16824                qtype,
16825                row_bytes,
16826                ..
16827            } if fast && *qtype == QT_Q4_K => {
16828                self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16829            }
16830            GpuTensor::Quant {
16831                bytes,
16832                qtype,
16833                row_bytes,
16834                ..
16835            } if fast && *qtype == QT_Q6_K => {
16836                self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16837            }
16838            GpuTensor::Quant {
16839                bytes,
16840                qtype,
16841                row_bytes,
16842                ..
16843            } if fast && *qtype == QT_Q5_K => {
16844                self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16845            }
16846            GpuTensor::Quant {
16847                bytes,
16848                qtype,
16849                row_bytes,
16850                ..
16851            } if fast && *qtype == QT_Q3_K => {
16852                self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16853            }
16854            GpuTensor::Quant {
16855                bytes,
16856                qtype,
16857                row_bytes,
16858                rp,
16859                ..
16860            } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
16861                if *rp {
16862                    "qmatvec_nvfp4_dp4a_rp"
16863                } else {
16864                    "qmatvec_nvfp4_dp4a"
16865                },
16866                &bytes.slice(0..bytes.len()),
16867                x,
16868                m,
16869                in_f,
16870                out_f,
16871                *row_bytes,
16872            )?,
16873            // IQ4_XS trunk fast path — DEFAULT ON since 2026-08-02 (MEMRA_IQ_FAST=0 reverts to
16874            // Stage-A; see iq_fast_enabled). The old opt-in default was the KAT-Coder decode
16875            // anomaly (research/kat-anomaly-20260802/).
16876            GpuTensor::Quant {
16877                bytes,
16878                qtype,
16879                row_bytes,
16880                ..
16881            } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
16882                self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16883            }
16884            // B3: IQ3_S uses the Stage-A f32 dequant-in-kernel path. There is NO
16885            // qmatvec_iq3_s_dp4a kernel — do NOT add a `*qtype == QT_IQ3_S` fast guard here
16886            // without first writing the matching kernel, or func() will panic
16887            // "kernel ... not in any fatbin".
16888            GpuTensor::Quant {
16889                bytes,
16890                qtype,
16891                row_bytes,
16892                rp,
16893                ..
16894            } =>
16895            // Stage-A generic: repacked NVFP4 uses the device-side split-plane tag (the
16896            // deq(row,j) form cannot address the planes; same value/product order).
16897            {
16898                self.qmatvec(
16899                    bytes,
16900                    x,
16901                    m,
16902                    in_f,
16903                    out_f,
16904                    if *rp && *qtype == QT_NVFP4 {
16905                        QT_NVFP4_RP
16906                    } else {
16907                        *qtype
16908                    },
16909                    *row_bytes,
16910                )?
16911            }
16912            GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
16913            // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use to f32 scratch, then the same
16914            // cuBLASLt f32 GEMV as the Float arm.
16915            GpuTensor::FloatBf16 { data, .. } => {
16916                // DECODE-TIER ROWS FAST PATH (2 <= m <= 8, bf16-mmv class): the chunked
16917                // arm dequants the WHOLE weight to f32 scratch per call — 4.7 ms/call on
16918                // the 1.24 GB LM head (nsys: 8x591us bf16_to_f32 per batch tick / per
16919                // verify round). matvec_bf16_f32acc_x4_rows runs the t=1 decode head
16920                // program PER ROW (identical dot + reduce), so decode/verify tiers keep
16921                // the t=1 numeric class and skip the convert. Prefill (m>8) keeps GEMM.
16922                if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f.is_multiple_of(8) {
16923                    let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16924                    self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
16925                    y
16926                } else {
16927                    self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
16928                }
16929            }
16930        };
16931        // NVFP4 per-tensor macro-scale (post-matmul). scale==1.0 for all other quants/float -> no-op.
16932        if let GpuTensor::Quant { scale, .. } = w
16933            && *scale != 1.0
16934        {
16935            self.scale_inplace(&mut y, *scale, m * out_f)?;
16936        }
16937        Ok(y)
16938    }
16939
16940    /// True when `MEMRA_FAST=0`, i.e. the Stage-A f32 oracle is the requested arithmetic and every
16941    /// `matmul_pre` call will take the raw-f32 escape rather than the q8_1 pair.
16942    ///
16943    /// WHY THIS EXISTS AS ITS OWN PREDICATE: `uses_q8_1_fast` needs a weight, but the callers that
16944    /// have to DECIDE WHETHER TO MATERIALIZE an f32 activation sit one level above any weight — the
16945    /// gemma-4 decode trunk emits a q8_1 pair per layer and hands it down, so the f32 has to be
16946    /// produced (or not) by the loop that owns the residual. Cached in a OnceLock like every other
16947    /// `*_on()` flag here: this is read once per layer per token on the decode path, and the daily
16948    /// path must not pay an env lookup for a flag that is off.
16949    pub fn stage_a_raw_needed() -> bool {
16950        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16951        *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
16952    }
16953
16954    /// True if `w` would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be
16955    /// pre-quantized once and shared across sibling matmuls via `matmul_pre`).
16956    pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
16957        use crate::model::GpuTensor;
16958        if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
16959            return false;
16960        }
16961        match w {
16962            // QT_F8_E4M3_BLK is admitted for the same reason QT_F8_E4M3 is: its ONLY kernel class
16963            // takes the shared q8_1 activation, so callers may pre-quantize once and share it
16964            // across siblings. It is NOT admitted to any of the fused/dual epilogue doors those
16965            // siblings can then open (`q8_fused_params`, `e4m3_fused_params` and
16966            // `matmul_pre_dual_noscale` all match on their own qtype and refuse this one) — the
16967            // block class has no fused twin yet, so each of its projections takes its own launch.
16968            GpuTensor::Quant { qtype, .. } => {
16969                matches!(
16970                    *qtype,
16971                    QT_Q8_0
16972                        | QT_Q4_K
16973                        | QT_Q6_K
16974                        | QT_Q5_K
16975                        | QT_Q3_K
16976                        | QT_NVFP4
16977                        | QT_F8_E4M3
16978                        | QT_F8_E4M3_BLK
16979                        | QT_Q4_0
16980                ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
16981            }
16982            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16983        }
16984    }
16985
16986    /// matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from `quantize_q8_1`). Skips the
16987    /// per-matmul re-quantize so sibling matmuls that share an input (gate+up share `z`;
16988    /// q/k/v + wqkv/gate/beta/alpha share `h`) quantize ONCE. Caller MUST have checked
16989    /// `uses_q8_1_fast(w)`; falls back to plain `matmul` otherwise (Stage-A / Float / non-fast).
16990    pub fn matmul_pre(
16991        &self,
16992        w: &crate::model::GpuTensor,
16993        aq: &CudaSlice<i8>,
16994        ad: &CudaSlice<f32>,
16995        x_fallback: &CudaSlice<f32>,
16996        m: usize,
16997    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16998        use crate::model::GpuTensor;
16999        // Every raw-f32 arm below (fp8/f16/MMQ/fp4) reads m*in_f from x_fallback. Callers that
17000        // pre-quantized and dropped the f32 input pass an EMPTY x_fallback (E4B's fusion port:
17001        // h = zeros(0)) — the length guard keeps those on the aq/ad GEMM instead of feeding a
17002        // 0-byte buffer to a convert kernel (illegal address -> cublasLt status 13; the E4B
17003        // rc=30013 dig, 2026-07-31).
17004        let x_raw_ok = x_fallback.len() >= m * w.in_features();
17005        // FP8-ACT PREFILL (MEMRA_PP_FP8=1): same arm as `matmul` — the fp8 operand needs the RAW
17006        // f32 activation (per-batch e4m3 quant differs from q8_1), so x_fallback not aq/ad.
17007        if m >= 16 && x_raw_ok && !self.verify_exact_on() {
17008            if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
17009                return Ok(y);
17010            }
17011            // PER-BLOCK FP8 MMQ — same arm as `matmul` (stash opt-in, native-resident default ON);
17012            // its own quantizer wants the RAW f32 activation, so x_fallback not aq/ad.
17013            if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
17014                return Ok(y);
17015            }
17016            // FP16-mirror prefill (same arm as `matmul` — fp16 wants the RAW f32 activation).
17017            if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
17018                return Ok(y);
17019            }
17020        }
17021        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK) — the same two arms as `matmul`, split at the same m, and
17022        // placed at the same point in the order (after the prefill GEMM hooks, before every arm
17023        // that refuses this qtype). The prefill arm needs the RAW f32 activation for the Q8_0
17024        // dispatch it recurses into, so it takes x_fallback and is skipped when that is empty
17025        // (a pre-quantized caller that dropped its f32 input never runs at prefill m anyway).
17026        if m >= 16
17027            && x_raw_ok
17028            && !self.verify_exact_on()
17029            && let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)?
17030        {
17031            return Ok(y);
17032        }
17033        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
17034            return Ok(y);
17035        }
17036        // VENDORED llama MMQ prefill GEMMs (NVFP4 W4A8 default-on; W4A4/k-quant behind MEMRA_MMQ=1
17037        // — policy in mmq_supports) — use the RAW f32 activation (their own internal quant:
17038        // q8_1 D4 for NVFP4 W4A8, FP8/UE4M3 for W4A4, q8_1 DS4 for Q4_K/Q5_K), so x_fallback not
17039        // aq/ad.
17040        if m >= 16
17041            && w.out_features() >= 128
17042            && self.mmq_supports(w)
17043            && !self.verify_exact_on()
17044            && x_raw_ok
17045        {
17046            return self.qmatvec_mmq(w, x_fallback, m);
17047        }
17048        // Stage-C FP4 prefill (MEMRA_FP4): native mxf4 GEMM needs the f32 activation (FP4-quant differs
17049        // from q8_1), so re-quantize from x_fallback rather than reuse aq/ad. NVFP4 only, m>=16.
17050        if m >= 16
17051            && x_raw_ok
17052            && !self.verify_exact_on()
17053            && let Some(y) =
17054                self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
17055        {
17056            return Ok(y);
17057        }
17058        // Prefill GEMM root fix: if T>1 and the dtype has a GEMM kernel, batch via tensor cores
17059        // (reuses the already-quantized aq/ad — no extra quantize). m=1 falls through to dp4a.
17060        if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
17061            return self.qmatvec_gemm(w, aq, ad, m);
17062        }
17063        // THE SEVENTH ARM. Every raw-f32 arm above is guarded by `x_raw_ok`; this one — the
17064        // Stage-A / Float escape — was not, and it is the ONLY one `MEMRA_FAST=0` opens. So the
17065        // 2026-07-31 E4B rc=30013 fix (add the length guard, keep empty-fallback callers off the
17066        // raw-f32 arms) protected the six arms the FAST path can reach and missed the one the
17067        // ORACLE path reaches. Consequence before this guard: the gemma-4 decode arms, which pass
17068        // `e.zeros(0)` because their trunk emits only a q8_1 pair and never materializes an f32
17069        // attn-normed activation, fell through to Stage-A `qmatvec_f32` (cu/qmatvec.cu:5442),
17070        // which reads `m * in_f` floats out of a 0-byte allocation ->
17071        // CUDA_ERROR_ILLEGAL_ADDRESS at layer 0 of the first decode token. The fault is STICKY:
17072        // it poisons the context, so every LATER request in that process fails with an unrelated
17073        // message ("cache alloc failed: ...") and the true cause appears exactly once, in the
17074        // first failure. That is what made `MEMRA_FAST=0` — the reference named in
17075        // `tools/argmax-margin-gate.sh`'s own header — return an opaque HTTP 500 on the gemma
17076        // dense artifact and left the arm with no working truth instrument.
17077        //
17078        // Refuse loudly instead of reading out of bounds. A named error at the true call site is
17079        // strictly better than an illegal address surfacing later at an unrelated sync point, and
17080        // an oracle that cannot run must say so rather than corrupt the context it runs in.
17081        if !self.uses_q8_1_fast(w) {
17082            if !x_raw_ok {
17083                return Err(format!(
17084                    "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
17085                     (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
17086                     dropped the f32, so there is nothing to fall back to — pass the real f32 \
17087                     activation (see Engine::rms_norm_decode, which is bit-identical to \
17088                     rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
17089                    x_fallback.len(),
17090                    m,
17091                    w.in_features(),
17092                    m * w.in_features()
17093                )
17094                .into());
17095            }
17096            return self.matmul(w, x_fallback, m);
17097        }
17098        let in_f = w.in_features();
17099        let out_f = w.out_features();
17100        let (bytes, qtype, row_bytes, scale, rp) = match w {
17101            GpuTensor::Quant {
17102                bytes,
17103                qtype,
17104                row_bytes,
17105                scale,
17106                rp,
17107                ..
17108            } => (bytes, *qtype, *row_bytes, *scale, *rp),
17109            _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
17110        };
17111        // Q4_0 split-plane mirror: only the mmvq/batched decode arms read it (the _rp twins);
17112        // the dp4a/oracle tails below keep the raw GGUF bytes.
17113        let (mbytes, mrp) = match w {
17114            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
17115            _ => (bytes, rp),
17116        };
17117        // PERF-3 decode-GEMV: warp-per-row MMVQ for the m=1 decode arm, gated behind MEMRA_MMVQ.
17118        // Only the 4 daily-hot dtypes have an _mmvq kernel (Q8_0/Q4_K/Q6_K/NVFP4); Q5_K/Q3_K/IQ4_XS
17119        // keep _dp4a (the oracle/fallback). Bit-equivalent to _dp4a up to f32 reduction order.
17120        if m == 1 && self.mmvq_supports(qtype) {
17121            return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
17122        }
17123        // BATCHED weight-resident matvec for the m=2-4 band (the MTP/verify forward: full_attn_verify
17124        // and decode_step_t run their projections at m=T=k=2..4). The plain _dp4a path below launches
17125        // grid.y=m INDEPENDENT blocks per output row -> the weight row is re-read m times from HBM/L2.
17126        // The _b2/_b4 kernels walk the weight ONCE and dp4a vs all m activation columns, so m tokens
17127        // cost ~1 weight read instead of m (decode is weight-BW-bound). BIT-IDENTICAL per (token,row)
17128        // to the _mmvq path (32-thread warp reduce — NOT the dp4a 128-thread reduce below).
17129        // m=2 -> mcols=2; m∈{3,4} -> mcols=4; m∈{5..8} -> mcols=8 (kernel guards c>=m).
17130        // MEMRA_NO_BATCHED forces the per-m grid.y=m path (the A/B reference); MEMRA_B8=0 keeps
17131        // m=5..8 on the old per-m path (b8-tier-only seam).
17132        // DECODE-PARITY GATE (2026-07-07): batched iff mmvq_supports — see matmul's parity note.
17133        // Without MEMRA_MMVQ, m=1 decode rides dp4a (the arm below at m=1); the verify must ride
17134        // the SAME class per column (grid.y=m dp4a = the exact m=1 dp4a program per column).
17135        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
17136            && std::env::var("MEMRA_NO_BATCHED").is_err()
17137            && (m <= 4 || Self::b8_enabled())
17138            // b16 tier: every class routed here now has base + _rp b16 kernels (Q4_0/Q6_K
17139            // pre-existing; NVFP4/Q4_K/Q8_0-base/F8_E4M3 added lane/rp-on-st 2026-08-06), so
17140            // there is no mirror precondition left — `mrp` still selects the LAYOUT below.
17141            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
17142                || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
17143        {
17144            let mcols = Self::batched_mcols(m);
17145            return self.qmatvec_mmvq_batched(
17146                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
17147            );
17148        }
17149        // F8-E4M3 catch-all (m=9..15 / batched-disabled seams): grid.y=m e4m3 mmvq — this dtype
17150        // has NO _dp4a twin, and per (token,row) the mmvq body is the exact m=1 decode program.
17151        // Q4_0 joins the catch-all (2026-07-11): adaptive-K cap 8 makes verify t=9 reachable
17152        // for the first time (past the b8 tier) and Q4_0 has no dp4a twin either. The mirror
17153        // (mbytes/mrp) keeps the rp layout consistent with the m=1 decode program.
17154        if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
17155            let (b2, r2) = if qtype == QT_Q4_0 {
17156                (mbytes, mrp)
17157            } else {
17158                (bytes, rp)
17159            };
17160            return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
17161        }
17162        let name = match qtype {
17163            QT_Q8_0 => "qmatvec_q8_0_dp4a",
17164            QT_Q4_K => "qmatvec_q4_K_dp4a",
17165            QT_Q6_K => "qmatvec_q6_K_dp4a",
17166            QT_Q5_K => "qmatvec_q5_K_dp4a",
17167            QT_Q3_K => "qmatvec_q3_K_dp4a",
17168            QT_NVFP4 => {
17169                if rp {
17170                    "qmatvec_nvfp4_dp4a_rp"
17171                } else {
17172                    "qmatvec_nvfp4_dp4a"
17173                }
17174            }
17175            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
17176            _ => unreachable!(),
17177        };
17178        let f = self.func(name);
17179        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
17180        let cfg = LaunchConfig {
17181            grid_dim: (out_f as u32, m as u32, 1),
17182            block_dim: (128, 1, 1),
17183            shared_mem_bytes: 0,
17184        };
17185        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17186        let __s_b = self.gpu.stream();
17187        let mut b = __s_b.launch_builder(&f);
17188        b.arg(bytes)
17189            .arg(aq)
17190            .arg(ad)
17191            .arg(&mut y)
17192            .arg(&inf)
17193            .arg(&outf)
17194            .arg(&mi)
17195            .arg(&rb);
17196        unsafe {
17197            b.launch(cfg)?;
17198        }
17199        if scale != 1.0 {
17200            self.scale_inplace(&mut y, scale, m * out_f)?;
17201        }
17202        Ok(y)
17203    }
17204
17205    /// DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP
17206    /// accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST
17207    /// use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel
17208    /// (128 threads, two-level reduction) used by `matmul`/`matmul_pre` at m>=5 has a different
17209    /// shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips.
17210    /// The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp
17211    /// reduce as m=1); this method just forces that path unconditionally.
17212    pub fn matmul_decode_exact(
17213        &self,
17214        w: &crate::model::GpuTensor,
17215        x: &CudaSlice<f32>,
17216        m: usize,
17217    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17218        use crate::model::GpuTensor;
17219        // FLOAT tensors (35B ssm_beta/ssm_alpha on every linear layer, F32 ne=[2048,32]): the
17220        // generic path is cuBLASLt, whose reduction splits are n-DEPENDENT — m=1 vs m=2 col-0
17221        // outputs differ in every bit (probe 2026-07-06: 32/32 bit-diff, maxdiff 3.5e-3), which
17222        // shifted 35B verify logits 0.26-0.56 vs eager and flipped greedy at tight margins (the
17223        // p3 spec FAIL). Decode-exact contract: per-COLUMN m=1 cuBLASLt calls — each column's
17224        // reduction is the exact kernel the T=1 decode path runs, so verify==decode bit-for-bit.
17225        // m<=10 here (K+2 verify tier), so the extra launches are a handful of 4us gemvs.
17226        if let GpuTensor::Float { data, .. } = w {
17227            return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
17228        }
17229        // MEMRA_FULL_PREC bf16-resident weight: dequant-on-use, then the per-column decode-exact
17230        // float linear (same n-independent reduction contract as the Float arm above).
17231        if let GpuTensor::FloatBf16 { data, .. } = w {
17232            let (in_f, out_f) = (w.in_features(), w.out_features());
17233            // Rows fast path: per-row t=1 program (STRONGER than the chunked per-column
17234            // contract — the whole-weight f32 dequant disappears too).
17235            if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
17236                let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17237                self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
17238                return Ok(y);
17239            }
17240            return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
17241        }
17242        if !self.uses_q8_1_fast(w) {
17243            return self.matmul(w, x, m);
17244        }
17245        let in_f = w.in_features();
17246        let out_f = w.out_features();
17247        let (bytes, qtype, row_bytes, scale, rp) = match w {
17248            GpuTensor::Quant {
17249                bytes,
17250                qtype,
17251                row_bytes,
17252                scale,
17253                rp,
17254                ..
17255            } => (bytes, *qtype, *row_bytes, *scale, *rp),
17256            _ => return self.matmul(w, x, m),
17257        };
17258        // Q4_0 split-plane mirror for the mmvq/batched arms below (dp4a tail = matmul_pre,
17259        // which does its own mirror pick).
17260        let (bytes, rp) = match w {
17261            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
17262            _ => (bytes, rp),
17263        };
17264        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17265        // BLOCK-128 e4m3 (QT_F8_E4M3_BLK): the same single kernel every other entry dispatches, so
17266        // the decode-exact contract needs nothing special — grid.y=m runs the m=1 program per
17267        // (token,row) by construction, which is exactly what this method exists to guarantee.
17268        if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
17269            return Ok(y);
17270        }
17271        // Batched weight-resident matvec for m=2-8: BIT-IDENTICAL per (token,row) to MMVQ (exact
17272        // integer dp4a, same warp reduce — kernel-check gate rel=0.00e0), one weight read for m
17273        // tokens. The dispatch the divergence fix must avoid is dp4a's 128-thread two-level
17274        // reduce, NOT this. m=5..8 is the K=4..7 spec-verify tier (b8): pre-b8 T=5 fell to the
17275        // grid.y=m per-row MMVQ below = 5 full weight reads/launch — the measured 27B K=4 cliff.
17276        // DECODE-PARITY GATE (2026-07-07): batched (MMVQ-class order) only when the m=1 decode
17277        // chain rides MMVQ too — without MEMRA_MMVQ decode is dp4a, so the exact-contract here
17278        // must be per-column dp4a (matmul_pre fallthrough), not the MMVQ order.
17279        if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
17280            && std::env::var("MEMRA_NO_BATCHED").is_err()
17281            && (m <= 4 || Self::b8_enabled())
17282            // Every b16 class has base + _rp twins after lane/rp-on-st (see matmul_pre's note):
17283            // no mirror precondition, `rp` selects the layout only.
17284            && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
17285                || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
17286        {
17287            let mcols = Self::batched_mcols(m);
17288            return self.qmatvec_mmvq_batched(
17289                bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
17290            );
17291        }
17292        if self.mmvq_supports(qtype) {
17293            // MMVQ at grid.y=m: each row is processed by its own warp independently — same 32-thread
17294            // accumulation + warp_reduce_sum as m=1 decode. Bit-identical per row.
17295            return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
17296        }
17297        // Fallback for non-MMVQ quant types (Q5_K, Q3_K): use dp4a (the only available kernel).
17298        // These types are not used in the 27B's linear-attn NVFP4+Q4_K layers.
17299        self.matmul_pre(w, &aq, &ad, x, m)
17300    }
17301
17302    /// DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue
17303    /// re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT `matmul_decode_exact` dispatch for
17304    /// q8_1-fast Quant tensors, with the caller's (aq, ad) replacing the internal
17305    /// `quantize_q8_1`. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so
17306    /// sharing one quantize across sibling matmuls of the same activation — or consuming the
17307    /// q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 /
17308    /// silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to
17309    /// their unfused chains) — cannot change any dispatched kernel's input bytes.
17310    /// Caller MUST guarantee `uses_q8_1_fast(w)` (the fused epilogues only exist on that path).
17311    pub fn matmul_decode_exact_pre(
17312        &self,
17313        w: &crate::model::GpuTensor,
17314        aq: &CudaSlice<i8>,
17315        ad: &CudaSlice<f32>,
17316        m: usize,
17317    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17318        use crate::model::GpuTensor;
17319        debug_assert!(
17320            self.uses_q8_1_fast(w),
17321            "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
17322        );
17323        // BLOCK-128 e4m3: same single kernel, all m — see matmul_decode_exact's note.
17324        if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
17325            return Ok(y);
17326        }
17327        let in_f = w.in_features();
17328        let out_f = w.out_features();
17329        let (bytes, qtype, row_bytes, scale, rp) = match w {
17330            GpuTensor::Quant {
17331                bytes,
17332                qtype,
17333                row_bytes,
17334                scale,
17335                rp,
17336                ..
17337            } => (bytes, *qtype, *row_bytes, *scale, *rp),
17338            _ => {
17339                return Err(
17340                    "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
17341                );
17342            }
17343        };
17344        // Q4_0 split-plane mirror — same pick as matmul_decode_exact.
17345        let (bytes, rp) = match w {
17346            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
17347            _ => (bytes, rp),
17348        };
17349        // Dispatch mirror of matmul_decode_exact's q8_1-fast tail, condition for condition.
17350        if (2..=16).contains(&m)
17351            && self.batched_supports(qtype)
17352            && self.mmvq_supports(qtype)
17353            && std::env::var("MEMRA_NO_BATCHED").is_err()
17354            && (m <= 4 || Self::b8_enabled())
17355            && (m <= 8
17356                || qtype == QT_Q4_0
17357                || qtype == QT_Q6_K
17358                || qtype == QT_F8_E4M3
17359                || qtype == QT_NVFP4
17360                || qtype == QT_Q4_K
17361                || qtype == QT_Q5_K
17362                || qtype == QT_Q8_0)
17363        {
17364            let mcols = Self::batched_mcols(m);
17365            return self.qmatvec_mmvq_batched(
17366                bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
17367            );
17368        }
17369        if self.mmvq_supports(qtype) {
17370            return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
17371        }
17372        // Non-MMVQ quant types (Q5_K/Q3_K under MEMRA_MMVQ=0): dp4a via matmul_pre — the same
17373        // fallback matmul_decode_exact takes. m <= 16 on the verify tier never reads x_fallback.
17374        let x0 = self.zeros(0)?;
17375        self.matmul_pre(w, aq, ad, &x0, m)
17376    }
17377
17378    /// DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED
17379    /// (lane/vt-fixes fix 2): same eligibility as `matmul_decode_exact_dual`, but the caller's
17380    /// (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED
17381    /// instead of applied via two `scale_inplace` launches — the fused SwiGLU epilogue
17382    /// (`silu_mul_scaled_q8_1`) folds them, exactly like the m=1 decode chain does. Deferring
17383    /// is value-exact: `y[i]*s` inline in the epilogue is the same IEEE multiply scale_inplace
17384    /// would store (f32 store/load round-trips are exact). None -> caller falls back to the
17385    /// per-tensor path.
17386    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17387    pub fn matmul_decode_exact_dual_pre(
17388        &self,
17389        w0: &crate::model::GpuTensor,
17390        w1: &crate::model::GpuTensor,
17391        aq: &CudaSlice<i8>,
17392        ad: &CudaSlice<f32>,
17393        m: usize,
17394    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
17395    {
17396        use crate::model::GpuTensor;
17397        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17398        let on = *ON.get_or_init(|| {
17399            std::env::var("MEMRA_SPEC_DUAL_T")
17400                .map(|v| v != "0")
17401                .unwrap_or(true)
17402        });
17403        if !on
17404            || !(2..=7).contains(&m)
17405            || std::env::var("MEMRA_NO_BATCHED").is_ok()
17406            || !self.uses_q8_1_fast(w0)
17407            || !self.uses_q8_1_fast(w1)
17408        {
17409            return Ok(None);
17410        }
17411        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — the dual
17412        // kernels are the MMVQ warp-reduce family, and without MEMRA_MMVQ the m=1 decode
17413        // chain this verify must match bit-for-bit rides dp4a (see matmul_decode_exact's
17414        // note). The singles enforce this via `mmvq_supports`; the dual door skipped it.
17415        if !self.mmvq_supports(QT_NVFP4) {
17416            return Ok(None);
17417        }
17418        let (in_f, out_f) = (w0.in_features(), w0.out_features());
17419        if w1.in_features() != in_f || w1.out_features() != out_f {
17420            return Ok(None);
17421        }
17422        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
17423            (
17424                GpuTensor::Quant {
17425                    bytes: b0,
17426                    qtype: q0,
17427                    row_bytes: rb0,
17428                    scale: s0,
17429                    rp: rp0,
17430                    rp4: None,
17431                    ..
17432                },
17433                GpuTensor::Quant {
17434                    bytes: b1,
17435                    qtype: q1,
17436                    row_bytes: rb1,
17437                    scale: s1,
17438                    rp: rp1,
17439                    rp4: None,
17440                    ..
17441                },
17442            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
17443                (b0, b1, *rb0, *s0, *s1, *rp0)
17444            }
17445            _ => return Ok(None),
17446        };
17447        // m=5..7: only the exact-width rp duals exist (vt-fixes fix 1b); GGUF layout keeps
17448        // the singles. The b8 dual (MCOLS=8 at m=5..8) measured FLAT and stays dead.
17449        if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
17450        {
17451            return Ok(None);
17452        }
17453        let (y0, y1) =
17454            self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
17455        Ok(Some(((y0, s0), (y1, s1))))
17456    }
17457
17458    /// GROUP-4 GDN-tuple BATCHED matvec (trunk-kernels slice C): the qwen35 linear-layer
17459    /// in-projection 4-tuple (wqkv / wqkv_gate / ssm_beta / ssm_alpha) from ONE pre-quantized
17460    /// activation in ONE launch. Blocks map to the concatenated row space; every out_f must be
17461    /// a multiple of 8 (rows_per_block) so each warp's row pair resolves to one tensor; per
17462    /// (tensor, token, row) the kernel body is `nvfp4_mmvq_batched_rp` VERBATIM with the
17463    /// tensor's macro-scale fused at the write (== the conditional scale_inplace pass,
17464    /// bit-identical) -> BIT-IDENTICAL to the four single launches. Split-plane rp NVFP4 only,
17465    /// m=2..16 (exact-width MCOLS at m=5..7 mirroring the B567 law; m>4 requires b8_enabled
17466    /// like the singles; m=9..=16 rides the b16 form — the E4 width lift that lets
17467    /// `matmul_nvfp4_fused3/4` delegate that class here, lane/orndecode2).
17468    /// None -> caller runs the four singles. MEMRA_TK_GDN_GROUP=0 rollback.
17469    pub fn matmul_decode_exact_group4_pre(
17470        &self,
17471        ws: [&crate::model::GpuTensor; 4],
17472        aq: &CudaSlice<i8>,
17473        ad: &CudaSlice<f32>,
17474        m: usize,
17475    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17476        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17477        let on = *ON.get_or_init(|| {
17478            std::env::var("MEMRA_TK_GDN_GROUP")
17479                .map(|v| v != "0")
17480                .unwrap_or(true)
17481        });
17482        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
17483    }
17484
17485    /// GROUP-3 twin for the qwen35 full-attention q/k/v triple (trunk-kernels slice D):
17486    /// the SAME group4 kernels with n3=0 (blocks never reach the fourth range; W3/y3 are
17487    /// never dereferenced) — per (tensor, token, row) bit-identical to the three singles
17488    /// exactly as the group4 door is to its four. MEMRA_TK_FA_GROUP=0 rollback.
17489    pub fn matmul_decode_exact_group3_pre(
17490        &self,
17491        ws: [&crate::model::GpuTensor; 3],
17492        aq: &CudaSlice<i8>,
17493        ad: &CudaSlice<f32>,
17494        m: usize,
17495    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17496        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17497        let on = *ON.get_or_init(|| {
17498            std::env::var("MEMRA_TK_FA_GROUP")
17499                .map(|v| v != "0")
17500                .unwrap_or(true)
17501        });
17502        self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
17503    }
17504
17505    /// Shared core of the group3/group4 doors: eligibility mirror of the singles' batched
17506    /// dispatch, then ONE `qmatvec_nvfp4_mmvq_group4_b*_rp` launch over the concatenated
17507    /// row space (3-tensor callers ride n3=0). Returns one output per input tensor.
17508    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
17509    fn matmul_decode_exact_group_pre(
17510        &self,
17511        ws: &[&crate::model::GpuTensor],
17512        aq: &CudaSlice<i8>,
17513        ad: &CudaSlice<f32>,
17514        m: usize,
17515        on: bool,
17516        tag: &'static str,
17517    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17518        use crate::model::GpuTensor;
17519        if !on
17520            || !(2..=16).contains(&m)
17521            || std::env::var("MEMRA_NO_BATCHED").is_ok()
17522            || (m > 4 && !Self::b8_enabled())
17523            || !self.mmvq_supports(QT_NVFP4)
17524            || !self.batched_supports(QT_NVFP4)
17525        {
17526            return Ok(None);
17527        }
17528        let in_f = ws[0].in_features();
17529        let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
17530        for w in ws {
17531            if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
17532                return Ok(None);
17533            }
17534            match w {
17535                GpuTensor::Quant {
17536                    bytes,
17537                    qtype,
17538                    scale,
17539                    rp: true,
17540                    rp4: None,
17541                    ..
17542                } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
17543                    parts.push((bytes, w.out_features(), *scale));
17544                }
17545                _ => return Ok(None),
17546            }
17547        }
17548        // MCOLS tier mirrors the singles: batched_mcols + the B567 exact-width law at m=5..7.
17549        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17550        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
17551        let mcols = if (5..=7).contains(&m) && b567 {
17552            m
17553        } else {
17554            Self::batched_mcols(m)
17555        };
17556        let kname: &'static str = match mcols {
17557            2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
17558            4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
17559            5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
17560            6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
17561            7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
17562            8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
17563            16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
17564            _ => return Ok(None),
17565        };
17566        // Engagement receipt PER DOOR (dead-arm lesson): one shared Once here suppressed
17567        // the second door's print on the slice-D battery — key the once-set by tag.
17568        if std::env::var("MEMRA_DEBUG").is_ok() {
17569            use std::sync::Mutex;
17570            static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
17571            let mut seen = SEEN.lock().unwrap();
17572            if !seen.contains(&tag) {
17573                seen.push(tag);
17574                eprintln!("[memra] {tag} batched ENGAGED (m={m})");
17575            }
17576        }
17577        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
17578        let rows_per_block = ROWS_PER_BLOCK * 2; // WROWS=2 in the group kernel
17579        let total: usize = parts.iter().map(|p| p.1).sum();
17580        let three = parts.len() == 3;
17581        let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
17582        let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
17583        let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
17584        // 3-tensor callers: n3=0 means no block ever resolves to the fourth range — W3/y3
17585        // are never dereferenced; a 1-element dummy keeps the launch ABI without aliasing y0.
17586        let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
17587        let cfg = LaunchConfig {
17588            grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
17589            block_dim: (32, ROWS_PER_BLOCK, 1),
17590            shared_mem_bytes: 0,
17591        };
17592        let (inf, mi) = (in_f as i32, m as i32);
17593        let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
17594        let n3 = if three { 0i32 } else { parts[3].1 as i32 };
17595        let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
17596        let s3 = if three { 1.0f32 } else { parts[3].2 };
17597        let w3 = if three { parts[0].0 } else { parts[3].0 };
17598        let f = self.func(kname);
17599        let __s_b = self.gpu.stream();
17600        let mut b = __s_b.launch_builder(&f);
17601        b.arg(parts[0].0)
17602            .arg(parts[1].0)
17603            .arg(parts[2].0)
17604            .arg(w3)
17605            .arg(aq)
17606            .arg(ad)
17607            .arg(&mut y0)
17608            .arg(&mut y1)
17609            .arg(&mut y2)
17610            .arg(&mut y3)
17611            .arg(&inf)
17612            .arg(&n0)
17613            .arg(&n1)
17614            .arg(&n2)
17615            .arg(&n3)
17616            .arg(&mi)
17617            .arg(&s0)
17618            .arg(&s1)
17619            .arg(&s2)
17620            .arg(&s3);
17621        unsafe {
17622            b.launch(cfg)?;
17623        }
17624        Ok(Some(if three {
17625            vec![y0, y1, y2]
17626        } else {
17627            vec![y0, y1, y2, y3]
17628        }))
17629    }
17630
17631    /// DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE
17632    /// launch computes both FFN projections of a verify batch — same activation, same shape,
17633    /// blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single
17634    /// batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF:
17635    /// b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single `matmul_decode_exact`
17636    /// launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e).
17637    /// The one activation quantize replaces two IDENTICAL quantizes of the same `x` (same
17638    /// kernel, same input -> same q8_1 bytes), and the two independent weight streams in one
17639    /// grid restore the memory-level parallelism the two-launch form loses to tail drain +
17640    /// launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair).
17641    /// `Some((y0, y1))` only when both tensors are NVFP4, the SAME layout (both rp or both
17642    /// GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4
17643    /// (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured
17644    /// FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was
17645    /// killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.
17646    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17647    pub fn matmul_decode_exact_dual(
17648        &self,
17649        w0: &crate::model::GpuTensor,
17650        w1: &crate::model::GpuTensor,
17651        x: &CudaSlice<f32>,
17652        m: usize,
17653    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17654        use crate::model::GpuTensor;
17655        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17656        let on = *ON.get_or_init(|| {
17657            std::env::var("MEMRA_SPEC_DUAL_T")
17658                .map(|v| v != "0")
17659                .unwrap_or(true)
17660        });
17661        if !on
17662            || !(2..=4).contains(&m)
17663            || std::env::var("MEMRA_NO_BATCHED").is_ok()
17664            || !self.uses_q8_1_fast(w0)
17665            || !self.uses_q8_1_fast(w1)
17666        {
17667            return Ok(None);
17668        }
17669        // DECODE-PARITY GATE (lane/nvfp4-strict, 2026-08-05): batched iff MMVQ — same law as
17670        // the singles' `batched_supports && mmvq_supports` check in matmul_decode_exact,
17671        // which this dual door bypassed. Without MEMRA_MMVQ the m=1 decode is dp4a; the
17672        // verify must ride the per-column dp4a class, not the MMVQ-family dual.
17673        if !self.mmvq_supports(QT_NVFP4) {
17674            return Ok(None);
17675        }
17676        let (in_f, out_f) = (w0.in_features(), w0.out_features());
17677        if w1.in_features() != in_f || w1.out_features() != out_f {
17678            return Ok(None);
17679        }
17680        let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
17681            (
17682                GpuTensor::Quant {
17683                    bytes: b0,
17684                    qtype: q0,
17685                    row_bytes: rb0,
17686                    scale: s0,
17687                    rp: rp0,
17688                    rp4: None,
17689                    ..
17690                },
17691                GpuTensor::Quant {
17692                    bytes: b1,
17693                    qtype: q1,
17694                    row_bytes: rb1,
17695                    scale: s1,
17696                    rp: rp1,
17697                    rp4: None,
17698                    ..
17699                },
17700            ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
17701                (b0, b1, *rb0, *s0, *s1, *rp0)
17702            }
17703            _ => return Ok(None),
17704        };
17705        // Engagement receipt (MEMRA_DEBUG=1): the first dead-arm A/B lesson — a `rp: false`
17706        // gate silently no-op'd the whole experiment; prove the arm is live in the log.
17707        if std::env::var("MEMRA_DEBUG").is_ok() {
17708            static ONCE: std::sync::Once = std::sync::Once::new();
17709            ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
17710        }
17711        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17712        let (y0, y1) =
17713            self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
17714        let mut y0 = y0;
17715        let mut y1 = y1;
17716        if s0 != 1.0 {
17717            self.scale_inplace(&mut y0, s0, m * out_f)?;
17718        }
17719        if s1 != 1.0 {
17720            self.scale_inplace(&mut y1, s1, m * out_f)?;
17721        }
17722        Ok(Some((y0, y1)))
17723    }
17724
17725    /// Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1
17726    /// activation (kernel-check's bit-equivalence entry; matmul_decode_exact_dual's core).
17727    /// mcols tier = batched_mcols(m); macro-scale NOT applied. `rp` selects the split-plane
17728    /// twins (both buffers must be the repacked layout).
17729    #[allow(clippy::too_many_arguments)]
17730    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
17731    pub fn qmatvec_batched_dual_raw(
17732        &self,
17733        b0: &CudaSlice<u8>,
17734        b1: &CudaSlice<u8>,
17735        aq: &CudaSlice<i8>,
17736        ad: &CudaSlice<f32>,
17737        m: usize,
17738        in_f: usize,
17739        out_f: usize,
17740        row_bytes: usize,
17741        rp: bool,
17742    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17743        const ROWS_PER_BLOCK: u32 = 4;
17744        let mcols = Self::batched_mcols(m);
17745        // EXACT-WIDTH duals at m=5..7 (vt-fixes fix 1b): rp-only; bit-identical to the two
17746        // b5/b6/b7 singles (blockIdx.y selects the tensor, same template body).
17747        let tiny_rp1 = rp
17748            && mcols == 4
17749            && out_f <= 128
17750            && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
17751        let (name, rows_per_block) = if tiny_rp1 {
17752            ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
17753        } else {
17754            match (mcols, rp, m) {
17755                (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
17756                (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
17757                (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
17758                (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
17759                (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
17760                (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
17761                (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
17762                _ => {
17763                    return Err(
17764                        format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
17765                    );
17766                }
17767            }
17768        };
17769        let f = self.func(name);
17770        let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
17771        let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
17772        let cfg = LaunchConfig {
17773            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
17774            block_dim: (32, ROWS_PER_BLOCK, 1),
17775            shared_mem_bytes: 0,
17776        };
17777        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17778        let __s_b = self.gpu.stream();
17779        let mut b = __s_b.launch_builder(&f);
17780        b.arg(b0)
17781            .arg(b1)
17782            .arg(aq)
17783            .arg(ad)
17784            .arg(&mut y0)
17785            .arg(&mut y1)
17786            .arg(&inf)
17787            .arg(&outf)
17788            .arg(&mi)
17789            .arg(&rb);
17790        unsafe {
17791            b.launch(cfg)?;
17792        }
17793        Ok((y0, y1))
17794    }
17795
17796    /// Like `matmul_pre` but RETURNS THE RAW (un-macro-scaled) matmul output together with the
17797    /// per-tensor NVFP4 scale, instead of applying `scale_inplace` internally. Used by the fused
17798    /// SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one `silu_mul_scaled` launch.
17799    /// `Some((y_raw, scale))` only on the m==1 decode fast path (mmvq / dp4a) where the scale is a
17800    /// separate post-launch op we can defer; returns `None` for every other path (prefill GEMM, FP4
17801    /// GEMM, Stage-A, Float) so the caller falls back to the scaled `matmul_pre` + `silu_mul`.
17802    /// DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same
17803    /// activation, same shape) — grid.y selects the tensor. Bit-identical per element to two
17804    /// mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro
17805    /// scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless
17806    /// both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.
17807    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17808    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
17809    pub fn matmul_pre_dual_noscale(
17810        &self,
17811        w0: &crate::model::GpuTensor,
17812        w1: &crate::model::GpuTensor,
17813        aq: &CudaSlice<i8>,
17814        ad: &CudaSlice<f32>,
17815        m: usize,
17816    ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
17817    {
17818        use crate::model::GpuTensor;
17819        if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
17820            return Ok(None);
17821        }
17822        // FP-ORDER LAW (lane/nvfp4-strict, 2026-08-05): every kernel this door can dispatch
17823        // (q8_0 fused2, nvfp4 dual_mr2) is the MMVQ family — 32-thread warp reduce. Without
17824        // MEMRA_MMVQ the m=1 singles ride dp4a (128-thread two-level reduce), so fusing here
17825        // would mix dispatch families across the pair — the exact class `q8_fused_params`
17826        // already refuses for Q8_0. The NVFP4 arm lacked this check, which is why
17827        // decode-batch-gate `--mode strict`'s equalizing env (MEMRA_MMVQ=0) never pinned
17828        // NVFP4 models: decode_step_h kept riding dual_mr2 while the batched body fell to
17829        // dp4a (gate1 maxdiff 1.639e-1 / gate2 step-8 divergence at the 2026-08-05 train
17830        // HEAD, research/nvfp4-strict-20260805/). Default env (MMVQ on) is dispatch-unchanged.
17831        if !self.mmvq_supports(QT_NVFP4) {
17832            return Ok(None);
17833        }
17834        let (in_f, out_f) = (w0.in_features(), w0.out_features());
17835        if w1.in_features() != in_f || w1.out_features() != out_f {
17836            return Ok(None);
17837        }
17838        // Q8_0 ARM (lane/q27-deepdive, 2026-08-05): the dense-FFN gate+up pair on a Q8_0 trunk fell
17839        // through this NVFP4-only gate to two `matmul_pre_noscale` launches — measured 128 of the
17840        // 1015 launches/token on q27-Q8_0 decode, the single largest un-fused class in the tick
17841        // (nsys `research/q27-deepdive-20260805/nsys/`). `q8_fused2_core` already serves the same
17842        // pair shape for the shared-expert gate/up, and its kernel body is `qmatvec_q8_0_mmvq`
17843        // VERBATIM per (tensor,row) -> BIT-IDENTICAL to the two separate launches. Q8_0 carries no
17844        // macro-scale (q8_fused_params requires scale==1.0), so the noscale contract is satisfied
17845        // by returning 1.0 for both: the SwiGLU epilogue's fold becomes the identity it already is
17846        // on this dtype today. Seam: MEMRA_Q8_FFN_FUSE2=0 rolls back to the two-launch pair.
17847        // rp4 guard: with MEMRA_Q8RP the singles route to the `_rp` split-plane twin over the
17848        // mirror buffer; the fused2 kernel has no `_rp` form, so fusing there would swap
17849        // dispatch families mid-model. Bail and let the two singles run (mirror lane unchanged).
17850        let no_mirror =
17851            |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
17852        if self.q8_ffn_fuse2_on()
17853            && no_mirror(w0)
17854            && no_mirror(w1)
17855            && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
17856        {
17857            let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
17858            return Ok(Some(((y0, 1.0), (y1, 1.0))));
17859        }
17860        // F8-E4M3 ARM (lane/fp8-decode-v1, 2026-08-05): with native e4m3 residency the FFN gate+up
17861        // pair (and the ssm beta+alpha dual, which routes through this same entry) fell through
17862        // both the NVFP4 gate below and the Q8_0 arm above to two `matmul_pre_noscale` launches —
17863        // native residency was UN-FUSING the trunk relative to the Q8_0 slab it replaces. The
17864        // fused2 kernel body is `qmatvec_e4m3_mmvq` VERBATIM per (tensor,row). Contract match:
17865        // `matmul_pre_noscale` on e4m3 launches with scale 1.0 and RETURNS the per-tensor
17866        // weight_scale for the caller to fold, so we pass ws=1.0 here and return (s0,s1) — same
17867        // bits, and the two macro-scale multiplies still fold into the SwiGLU epilogue.
17868        // MEMRA_E4M3_DUAL=0 rolls back to the two-launch pair.
17869        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
17870            let (y0, y1) =
17871                self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
17872            return Ok(Some(((y0, p0.3), (y1, p1.3))));
17873        }
17874        let (b0, q0, rb0, s0, rp0) = match w0 {
17875            GpuTensor::Quant {
17876                bytes,
17877                qtype,
17878                row_bytes,
17879                scale,
17880                rp,
17881                ..
17882            } => (bytes, *qtype, *row_bytes, *scale, *rp),
17883            _ => return Ok(None),
17884        };
17885        let (b1, q1, rb1, s1, rp1) = match w1 {
17886            GpuTensor::Quant {
17887                bytes,
17888                qtype,
17889                row_bytes,
17890                scale,
17891                rp,
17892                ..
17893            } => (bytes, *qtype, *row_bytes, *scale, *rp),
17894            _ => return Ok(None),
17895        };
17896        if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
17897            return Ok(None);
17898        }
17899        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
17900        const RPW: u32 = 2;
17901        let rows_per_block = ROWS_PER_BLOCK * RPW;
17902        let f = self.func(if rp0 {
17903            "qmatvec_nvfp4_mmvq_dual_mr2_rp"
17904        } else {
17905            "qmatvec_nvfp4_mmvq_dual_mr2"
17906        });
17907        let mut y0 = self.alloc_uninit::<f32>(out_f)?;
17908        let mut y1 = self.alloc_uninit::<f32>(out_f)?;
17909        let cfg = LaunchConfig {
17910            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
17911            block_dim: (32, ROWS_PER_BLOCK, 1),
17912            shared_mem_bytes: 0,
17913        };
17914        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
17915        // noscale contract: the caller folds s0/s1 into the SwiGLU epilogue — the kernel's fused
17916        // yscale args stay 1.0 here (they exist for the single-tensor callers).
17917        let one = 1.0f32;
17918        let __s_b = self.gpu.stream();
17919        let mut b = __s_b.launch_builder(&f);
17920        b.arg(b0)
17921            .arg(b1)
17922            .arg(aq)
17923            .arg(ad)
17924            .arg(&mut y0)
17925            .arg(&mut y1)
17926            .arg(&inf)
17927            .arg(&outf)
17928            .arg(&mi)
17929            .arg(&rb)
17930            .arg(&one)
17931            .arg(&one);
17932        unsafe {
17933            b.launch(cfg)?;
17934        }
17935        Ok(Some(((y0, s0), (y1, s1))))
17936    }
17937
17938    /// FUSED NVFP4 matvec TRIPLE with unequal out_f (rig-native decode increment 1,
17939    /// lane/rig-native-nvfp4): wq+wk+wv in ONE launch via the q8_0 fused2 block-offset
17940    /// recipe. Per (tensor,row,t) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM ->
17941    /// bit-identical to three separate `matmul_pre` launches; yscales fold in-kernel exactly
17942    /// as the singles do. grid.y = m (the t-parallel verify rows ride the same launch).
17943    /// None when ineligible (not all rp NVFP4 / in_f mismatch / mmvq off) — callers fall
17944    /// back to the three singles.
17945    #[allow(clippy::too_many_arguments)]
17946    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17947    pub fn matmul_nvfp4_fused3(
17948        &self,
17949        w0: &crate::model::GpuTensor,
17950        w1: &crate::model::GpuTensor,
17951        w2: &crate::model::GpuTensor,
17952        aq: &CudaSlice<i8>,
17953        ad: &CudaSlice<f32>,
17954        m: usize,
17955    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17956    {
17957        use crate::model::GpuTensor;
17958        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
17959        // (lane/samplat — the fused4_b8 pattern): the seg body is nvfp4_mmvq_batched_rp_sc
17960        // verbatim, weight rows read once for all m columns, bit-identical per
17961        // (tensor,row,column) to the three bN_rpsc singles. The old "at m>1 the fused
17962        // segments would re-read the weight per row" note described the grid.y=m lift,
17963        // which this twin deliberately is NOT.
17964        if !self.mmvq_supports(QT_NVFP4)
17965            || !self.uses_q8_1_fast(w0)
17966            || !self.uses_q8_1_fast(w1)
17967            || !self.uses_q8_1_fast(w2)
17968        {
17969            return Ok(None);
17970        }
17971        // m = 9..=16 (lane/orndecode2): the exact-16 tier's trio width rides the GROUP3
17972        // door — same family and bit-identity law as the fused4 delegate above.
17973        if (9..=16).contains(&m) {
17974            return Ok(
17975                match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
17976                    Some(mut ys) => {
17977                        let y2 = ys.pop().unwrap();
17978                        let y1 = ys.pop().unwrap();
17979                        let y0 = ys.pop().unwrap();
17980                        Some((y0, y1, y2))
17981                    }
17982                    None => None,
17983                },
17984            );
17985        }
17986        if !(1..=8).contains(&m) {
17987            return Ok(None);
17988        }
17989        if m > 1 {
17990            let in_f = w0.in_features();
17991            if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
17992                || !self.batched_supports(QT_NVFP4)
17993                || std::env::var("MEMRA_NO_BATCHED").is_ok()
17994                || (m > 4 && !Self::b8_enabled())
17995                || !in_f.is_multiple_of(512)
17996                || in_f / 64 > 272
17997            {
17998                return Ok(None);
17999            }
18000        }
18001        let unpack = |w: &crate::model::GpuTensor| match w {
18002            GpuTensor::Quant {
18003                bytes,
18004                qtype,
18005                scale,
18006                rp,
18007                ..
18008            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18009            _ => None,
18010        };
18011        let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
18012            return Ok(None);
18013        };
18014        let in_f = w0.in_features();
18015        if w1.in_features() != in_f || w2.in_features() != in_f {
18016            return Ok(None);
18017        }
18018        let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
18019        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
18020        const RPW: u32 = 2;
18021        let rows_pb = ROWS_PER_BLOCK * RPW;
18022        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18023        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
18024        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
18025        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
18026        let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
18027        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
18028        // only dereferenced for the launch-arg build inside this call.
18029        let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
18030        if m > 1 {
18031            // batched twin has no in-kernel scale — refuse scale carriers (GGUF trunk = 1.0).
18032            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
18033                return Ok(None);
18034            }
18035            let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
18036            let cfg = LaunchConfig {
18037                grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
18038                block_dim: (32, ROWS_PER_BLOCK, 1),
18039                shared_mem_bytes: 0,
18040            };
18041            let __s_b = self.gpu.stream();
18042            let mut b = __s_b.launch_builder(&f);
18043            b.arg(b0)
18044                .arg(b1)
18045                .arg(b2)
18046                .arg(aq)
18047                .arg(ad)
18048                .arg(&mut y0)
18049                .arg(&mut y1)
18050                .arg(&mut y2)
18051                .arg(&inf)
18052                .arg(&oi0)
18053                .arg(&oi1)
18054                .arg(&oi2)
18055                .arg(&mi);
18056            unsafe {
18057                b.launch(cfg)?;
18058            }
18059            return Ok(Some((y0, y1, y2)));
18060        }
18061        let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
18062        let cfg = LaunchConfig {
18063            grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
18064            block_dim: (32, ROWS_PER_BLOCK, 1),
18065            shared_mem_bytes: 0,
18066        };
18067        let __s_b = self.gpu.stream();
18068        let mut b = __s_b.launch_builder(&f);
18069        b.arg(b0)
18070            .arg(b1)
18071            .arg(b2)
18072            .arg(aq)
18073            .arg(ad)
18074            .arg(&mut y0)
18075            .arg(&mut y1)
18076            .arg(&mut y2)
18077            .arg(&inf)
18078            .arg(&oi0)
18079            .arg(&oi1)
18080            .arg(&oi2)
18081            .arg(&mi)
18082            .arg(&p0.1)
18083            .arg(&p1.1)
18084            .arg(&p2.1);
18085        unsafe {
18086            b.launch(cfg)?;
18087        }
18088        Ok(Some((y0, y1, y2)))
18089    }
18090
18091    /// fused2 twin of `matmul_nvfp4_fused3`, for MIXED-type weight groups: the gemma4
18092    /// dense NVFP4mix recipe keeps `attn_v` and `ffn_down` at Q8_0 (full-NVFP4 was
18093    /// measured garbage on this dense class), so its q/k/v trio and gate/up/down never
18094    /// satisfy an all-NVFP4 fused3 — the pairs that ARE uniformly NVFP4 (q,k and
18095    /// gate,up) fuse here instead. m==1 only, same law as fused3/fused4: per
18096    /// (tensor,row) the kernel seg body is VERBATIM, so the fusion is bit-identical to
18097    /// two separate launches. `MEMRA_NVFP4_FUSED2=0` is the rollback seam and the
18098    /// same-binary interleaved A/B arm.
18099    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18100    pub fn matmul_nvfp4_fused2(
18101        &self,
18102        w0: &crate::model::GpuTensor,
18103        w1: &crate::model::GpuTensor,
18104        aq: &CudaSlice<i8>,
18105        ad: &CudaSlice<f32>,
18106        m: usize,
18107    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18108        use crate::model::GpuTensor;
18109        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18110        let off =
18111            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
18112        // m==1 ONLY: at m>1 the singles ride the _b16 weight-once column program (one weight
18113        // read serves all m rows); the fused segments would re-read the weight per row.
18114        if off
18115            || m != 1
18116            || !self.mmvq_supports(QT_NVFP4)
18117            || !self.uses_q8_1_fast(w0)
18118            || !self.uses_q8_1_fast(w1)
18119        {
18120            return Ok(None);
18121        }
18122        let unpack = |w: &crate::model::GpuTensor| match w {
18123            GpuTensor::Quant {
18124                bytes,
18125                qtype,
18126                scale,
18127                rp,
18128                ..
18129            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18130            _ => None,
18131        };
18132        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
18133            return Ok(None);
18134        };
18135        let in_f = w0.in_features();
18136        if w1.in_features() != in_f {
18137            return Ok(None);
18138        }
18139        let (o0, o1) = (w0.out_features(), w1.out_features());
18140        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
18141        let mut rpw: u32 = 2;
18142        let mut kname = "qmatvec_nvfp4_mmvq_fused2_rp";
18143        // B200 sub-wave grid-fill (MEMRA_B200_MATVEC_ARM occupancy arm, lane/b200-matvec-
18144        // occupancy-20260902): halves RPW to 1 (doubling the grid) when the RPW=2 grid would
18145        // leave B200's 148 SMs under a full wave, dispatching the `_g2` twin that
18146        // instantiates `nvfp4_mmvq_fused_seg_rp<1>` instead of `<2>`. Per (tensor,row) the
18147        // seg body is the same template body regardless of RPW -> bit-identical. Default
18148        // OFF; sm_120a keeps the measured RPW=2 default unconditionally.
18149        if b200_matvec_arm_on() {
18150            let waves_at_rpw2 =
18151                (o0 as u32).div_ceil(ROWS_PER_BLOCK * 2) + (o1 as u32).div_ceil(ROWS_PER_BLOCK * 2);
18152            if waves_at_rpw2 < 2 * self.sm_count() as u32 {
18153                rpw = 1;
18154                kname = "qmatvec_nvfp4_mmvq_fused2_rp_g2";
18155            }
18156        }
18157        let rows_pb = ROWS_PER_BLOCK * rpw;
18158        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18159        let f = self.func(kname);
18160        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
18161        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
18162        let cfg = LaunchConfig {
18163            grid_dim: (nb(o0) + nb(o1), m as u32, 1),
18164            block_dim: (32, ROWS_PER_BLOCK, 1),
18165            shared_mem_bytes: 0,
18166        };
18167        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
18168        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
18169        // only dereferenced for the launch-arg build inside this call.
18170        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
18171        // PDL wave-B: the kernel carries MEMRA_PDL_ENTRY — launch overlaps the producer's
18172        // drain (quantize/norm emitting aq/ad). Same math, same order, bit-identical.
18173        if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
18174            {
18175                use cudarc::driver::{DevicePtr, DevicePtrMut};
18176                let s = &self.gpu.stream();
18177                let (pw0, _g0) = b0.device_ptr(s);
18178                let (pw1, _g1) = b1.device_ptr(s);
18179                let (paq, _g2) = aq.device_ptr(s);
18180                let (pad, _g3) = ad.device_ptr(s);
18181                let (py0, _g4) = y0.device_ptr_mut(s);
18182                let (py1, _g5) = y1.device_ptr_mut(s);
18183                let (s0, s1) = (p0.1, p1.1);
18184                let mut ps = [
18185                    &pw0 as *const _ as *mut std::ffi::c_void,
18186                    &pw1 as *const _ as *mut _,
18187                    &paq as *const _ as *mut _,
18188                    &pad as *const _ as *mut _,
18189                    &py0 as *const _ as *mut _,
18190                    &py1 as *const _ as *mut _,
18191                    &inf as *const _ as *mut _,
18192                    &oi0 as *const _ as *mut _,
18193                    &oi1 as *const _ as *mut _,
18194                    &mi as *const _ as *mut _,
18195                    &s0 as *const _ as *mut _,
18196                    &s1 as *const _ as *mut _,
18197                ];
18198                unsafe {
18199                    self.launch_pdl(kname, cfg.grid_dim, cfg.block_dim, &mut ps)?;
18200                }
18201            }
18202            return Ok(Some((y0, y1)));
18203        }
18204        let __s_b = self.gpu.stream();
18205        let mut b = __s_b.launch_builder(&f);
18206        b.arg(b0)
18207            .arg(b1)
18208            .arg(aq)
18209            .arg(ad)
18210            .arg(&mut y0)
18211            .arg(&mut y1)
18212            .arg(&inf)
18213            .arg(&oi0)
18214            .arg(&oi1)
18215            .arg(&mi)
18216            .arg(&p0.1)
18217            .arg(&p1.1);
18218        unsafe {
18219            b.launch(cfg)?;
18220        }
18221        Ok(Some((y0, y1)))
18222    }
18223
18224    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch to
18225    /// `matmul_nvfp4_fused2`, caller-owned outputs — the gemma4 dc_slotted graph body
18226    /// needs zero mem nodes, so the allocating wrapper can't serve it. Returns false
18227    /// (decline) on any non-NVFP4/rp pair; the caller chains or refuses.
18228    pub fn matmul_nvfp4_fused2_into(
18229        &self,
18230        w0: &crate::model::GpuTensor,
18231        w1: &crate::model::GpuTensor,
18232        aq: &CudaSlice<i8>,
18233        ad: &CudaSlice<f32>,
18234        y0: &mut CudaSlice<f32>,
18235        y1: &mut CudaSlice<f32>,
18236    ) -> Result<bool, Box<dyn std::error::Error>> {
18237        use crate::model::GpuTensor;
18238        static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18239        let off =
18240            *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
18241        if off
18242            || !self.mmvq_supports(QT_NVFP4)
18243            || !self.uses_q8_1_fast(w0)
18244            || !self.uses_q8_1_fast(w1)
18245        {
18246            return Ok(false);
18247        }
18248        let unpack = |w: &crate::model::GpuTensor| match w {
18249            GpuTensor::Quant {
18250                bytes,
18251                qtype,
18252                scale,
18253                rp,
18254                ..
18255            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18256            _ => None,
18257        };
18258        let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
18259            return Ok(false);
18260        };
18261        let in_f = w0.in_features();
18262        if w1.in_features() != in_f {
18263            return Ok(false);
18264        }
18265        let (o0, o1) = (w0.out_features(), w1.out_features());
18266        if y0.len() < o0 || y1.len() < o1 {
18267            return Ok(false);
18268        }
18269        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
18270        let mut rpw: u32 = 2;
18271        let mut kname = "qmatvec_nvfp4_mmvq_fused2_rp";
18272        // B200 sub-wave grid-fill (MEMRA_B200_MATVEC_ARM occupancy arm) — see
18273        // `matmul_nvfp4_fused2` above for the full rationale; identical policy, alloc-free
18274        // caller.
18275        if b200_matvec_arm_on() {
18276            let waves_at_rpw2 =
18277                (o0 as u32).div_ceil(ROWS_PER_BLOCK * 2) + (o1 as u32).div_ceil(ROWS_PER_BLOCK * 2);
18278            if waves_at_rpw2 < 2 * self.sm_count() as u32 {
18279                rpw = 1;
18280                kname = "qmatvec_nvfp4_mmvq_fused2_rp_g2";
18281            }
18282        }
18283        let rows_pb = ROWS_PER_BLOCK * rpw;
18284        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18285        let f = self.func(kname);
18286        let cfg = LaunchConfig {
18287            grid_dim: (nb(o0) + nb(o1), 1, 1),
18288            block_dim: (32, ROWS_PER_BLOCK, 1),
18289            shared_mem_bytes: 0,
18290        };
18291        let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
18292        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
18293        // only dereferenced for the launch-arg build inside this call.
18294        let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
18295        let __s_b = self.gpu.stream();
18296        let mut b = __s_b.launch_builder(&f);
18297        b.arg(b0)
18298            .arg(b1)
18299            .arg(aq)
18300            .arg(ad)
18301            .arg(&mut *y0)
18302            .arg(&mut *y1)
18303            .arg(&inf)
18304            .arg(&oi0)
18305            .arg(&oi1)
18306            .arg(&mi)
18307            .arg(&p0.1)
18308            .arg(&p1.1);
18309        unsafe {
18310            b.launch(cfg)?;
18311        }
18312        Ok(true)
18313    }
18314
18315    /// fused4 twin of `matmul_nvfp4_fused3`: the Linear-mixer projection quartet
18316    /// (wqkv + wqkv_gate + ssm_beta + ssm_alpha) in one launch, m==1 only. Per
18317    /// (tensor,row) the kernel body is nvfp4_mmvq_multirow_rp VERBATIM — bit-identical
18318    /// to four separate launches (rig-native decode increment 2, RIG-NATIVE-DECODE.md).
18319    #[allow(clippy::type_complexity)]
18320    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
18321    pub fn matmul_nvfp4_fused4(
18322        &self,
18323        w0: &crate::model::GpuTensor,
18324        w1: &crate::model::GpuTensor,
18325        w2: &crate::model::GpuTensor,
18326        w3: &crate::model::GpuTensor,
18327        aq: &CudaSlice<i8>,
18328        ad: &CudaSlice<f32>,
18329        m: usize,
18330    ) -> Result<
18331        Option<(
18332            CudaSlice<f32>,
18333            CudaSlice<f32>,
18334            CudaSlice<f32>,
18335            CudaSlice<f32>,
18336        )>,
18337        Box<dyn std::error::Error>,
18338    > {
18339        use crate::model::GpuTensor;
18340        // MEMRA_NVFP4_FUSED4=0: rollback seam + the same-binary interleaved A/B arm.
18341        // m == 1 rides the original fused kernel; m = 2..=8 rides the BATCHED fused twin
18342        // (lane/samplat, 2026-08-21): same quartet-in-one-launch shape, seg body =
18343        // nvfp4_mmvq_batched_rp_sc verbatim (weight rows read once for all m columns) —
18344        // bit-identical per (tensor,row,column) to the four bN_rpsc singles it replaces.
18345        // Admission mirrors the singles' batched gates below.
18346        if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
18347            || !self.mmvq_supports(QT_NVFP4)
18348            || !self.uses_q8_1_fast(w0)
18349            || !self.uses_q8_1_fast(w1)
18350            || !self.uses_q8_1_fast(w2)
18351            || !self.uses_q8_1_fast(w3)
18352        {
18353            return Ok(None);
18354        }
18355        // m = 9..=16 (lane/orndecode2, the exact-16 tier's trunk): the rp-sc seg twins stop
18356        // at 8; this width class rides the GROUP4 door instead — nvfp4_mmvq_batched_rp<16,_>
18357        // body, the SAME family as the b16_rp singles the tier would otherwise launch four
18358        // times, bit-identical per (tensor,token,row) incl. the fused write-side scale.
18359        if (9..=16).contains(&m) {
18360            return Ok(
18361                match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
18362                    Some(mut ys) => {
18363                        let y3 = ys.pop().unwrap();
18364                        let y2 = ys.pop().unwrap();
18365                        let y1 = ys.pop().unwrap();
18366                        let y0 = ys.pop().unwrap();
18367                        Some((y0, y1, y2, y3))
18368                    }
18369                    None => None,
18370                },
18371            );
18372        }
18373        if !(1..=8).contains(&m) {
18374            return Ok(None);
18375        }
18376        if m > 1 {
18377            // the batched-twin gates: the bN_rpsc program this must stay byte-identical to
18378            // (matmul_pre's batched arm), plus the rp-sc dispatch shape requirements.
18379            let in_f = w0.in_features();
18380            if !self.batched_supports(QT_NVFP4)
18381                || std::env::var("MEMRA_NO_BATCHED").is_ok()
18382                || (m > 4 && !Self::b8_enabled())
18383                || !in_f.is_multiple_of(512)
18384                || in_f / 64 > 272
18385            {
18386                return Ok(None);
18387            }
18388        }
18389        let unpack = |w: &crate::model::GpuTensor| match w {
18390            GpuTensor::Quant {
18391                bytes,
18392                qtype,
18393                scale,
18394                rp,
18395                ..
18396            } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18397            _ => None,
18398        };
18399        let (Some(p0), Some(p1), Some(p2), Some(p3)) =
18400            (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
18401        else {
18402            return Ok(None);
18403        };
18404        let in_f = w0.in_features();
18405        if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
18406            return Ok(None);
18407        }
18408        let (o0, o1, o2, o3) = (
18409            w0.out_features(),
18410            w1.out_features(),
18411            w2.out_features(),
18412            w3.out_features(),
18413        );
18414        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
18415        const RPW: u32 = 2;
18416        let rows_pb = ROWS_PER_BLOCK * RPW;
18417        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18418        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
18419        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
18420        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
18421        let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
18422        let (inf, oi0, oi1, oi2, oi3, mi) = (
18423            in_f as i32,
18424            o0 as i32,
18425            o1 as i32,
18426            o2 as i32,
18427            o3 as i32,
18428            m as i32,
18429        );
18430        // SAFETY: the raw pointers come straight from the &GpuTensor borrows above and are
18431        // only dereferenced for the launch-arg build inside this call.
18432        let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
18433        if m > 1 {
18434            // Batched fused twin: no in-kernel scale (the bN_rpsc program has none) — refuse
18435            // scale-carrying tensors so the singles path keeps them (GGUF trunk scales are 1.0).
18436            if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
18437                return Ok(None);
18438            }
18439            let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
18440            let cfg = LaunchConfig {
18441                grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
18442                block_dim: (32, ROWS_PER_BLOCK, 1),
18443                shared_mem_bytes: 0,
18444            };
18445            let __s_b = self.gpu.stream();
18446            let mut b = __s_b.launch_builder(&f);
18447            b.arg(b0)
18448                .arg(b1)
18449                .arg(b2)
18450                .arg(b3)
18451                .arg(aq)
18452                .arg(ad)
18453                .arg(&mut y0)
18454                .arg(&mut y1)
18455                .arg(&mut y2)
18456                .arg(&mut y3)
18457                .arg(&inf)
18458                .arg(&oi0)
18459                .arg(&oi1)
18460                .arg(&oi2)
18461                .arg(&oi3)
18462                .arg(&mi);
18463            unsafe {
18464                b.launch(cfg)?;
18465            }
18466            return Ok(Some((y0, y1, y2, y3)));
18467        }
18468        let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
18469        let cfg = LaunchConfig {
18470            grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
18471            block_dim: (32, ROWS_PER_BLOCK, 1),
18472            shared_mem_bytes: 0,
18473        };
18474        let __s_b = self.gpu.stream();
18475        let mut b = __s_b.launch_builder(&f);
18476        b.arg(b0)
18477            .arg(b1)
18478            .arg(b2)
18479            .arg(b3)
18480            .arg(aq)
18481            .arg(ad)
18482            .arg(&mut y0)
18483            .arg(&mut y1)
18484            .arg(&mut y2)
18485            .arg(&mut y3)
18486            .arg(&inf)
18487            .arg(&oi0)
18488            .arg(&oi1)
18489            .arg(&oi2)
18490            .arg(&oi3)
18491            .arg(&mi)
18492            .arg(&p0.1)
18493            .arg(&p1.1)
18494            .arg(&p2.1)
18495            .arg(&p3.1);
18496        unsafe {
18497            b.launch(cfg)?;
18498        }
18499        Ok(Some((y0, y1, y2, y3)))
18500    }
18501
18502    /// FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two
18503    /// same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp
18504    /// 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the
18505    /// dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is
18506    /// qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when
18507    /// ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls
18508    /// back to the per-tensor path.
18509    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18510    pub fn matmul_q8_fused2(
18511        &self,
18512        w0: &crate::model::GpuTensor,
18513        w1: &crate::model::GpuTensor,
18514        aq: &CudaSlice<i8>,
18515        ad: &CudaSlice<f32>,
18516    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18517        // e4m3 twin (lane/fp8-decode-v1): this entry is the trunk's generic m=1 pair door
18518        // (wqkv+wqkv_gate, ssm_beta+alpha, gate_shexp+up_shexp), so admitting QT_F8_E4M3 here
18519        // fuses the NATIVE-RESIDENCY FP8 trunk at every existing call site with no call-site
18520        // change. Scale is folded in-kernel per range -> the returned buffers are already scaled,
18521        // exactly like the per-tensor `matmul_pre` e4m3 dispatch this replaces.
18522        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
18523            return Ok(Some(self.e4m3_fused2_core(
18524                p0.0,
18525                p1.0,
18526                aq,
18527                ad,
18528                w0.in_features(),
18529                p0.1,
18530                p1.1,
18531                p0.2,
18532                p0.3,
18533                p1.3,
18534            )?));
18535        }
18536        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
18537            return Ok(None);
18538        };
18539        Ok(Some(self.q8_fused2_core(
18540            p0.0,
18541            p1.0,
18542            aq,
18543            ad,
18544            w0.in_features(),
18545            p0.1,
18546            p1.1,
18547            p0.2,
18548        )?))
18549    }
18550
18551    #[allow(clippy::too_many_arguments)]
18552    fn q8_fused2_core(
18553        &self,
18554        b0: &CudaSlice<u8>,
18555        b1: &CudaSlice<u8>,
18556        aq: &CudaSlice<i8>,
18557        ad: &CudaSlice<f32>,
18558        in_f: usize,
18559        out0: usize,
18560        out1: usize,
18561        row_bytes: usize,
18562    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18563        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
18564        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18565        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18566        let f = self.func("qmatvec_q8_0_mmvq_fused2");
18567        let mut y0 = self.alloc_uninit::<f32>(out0)?;
18568        let mut y1 = self.alloc_uninit::<f32>(out1)?;
18569        let cfg = LaunchConfig {
18570            grid_dim: (nb0 + nb1, 1, 1),
18571            block_dim: (32, ROWS_PER_BLOCK, 1),
18572            shared_mem_bytes: 0,
18573        };
18574        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
18575        let __s_b = self.gpu.stream();
18576        let mut b = __s_b.launch_builder(&f);
18577        b.arg(b0)
18578            .arg(b1)
18579            .arg(aq)
18580            .arg(ad)
18581            .arg(&mut y0)
18582            .arg(&mut y1)
18583            .arg(&inf)
18584            .arg(&o0)
18585            .arg(&o1)
18586            .arg(&rbl);
18587        unsafe {
18588            b.launch(cfg)?;
18589        }
18590        Ok((y0, y1))
18591    }
18592
18593    /// f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused
18594    /// launch — replaces two `matmul(w, x, 1)` calls that would each re-quantize the same x
18595    /// (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is
18596    /// deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the
18597    /// callers' m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).
18598    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18599    pub fn matmul_q8_fused2_x(
18600        &self,
18601        w0: &crate::model::GpuTensor,
18602        w1: &crate::model::GpuTensor,
18603        x: &CudaSlice<f32>,
18604    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18605        if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
18606            return Ok(None);
18607        }
18608        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
18609            let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
18610            return Ok(Some(self.e4m3_fused2_core(
18611                p0.0,
18612                p1.0,
18613                &aq,
18614                &ad,
18615                w0.in_features(),
18616                p0.1,
18617                p1.1,
18618                p0.2,
18619                p0.3,
18620                p1.3,
18621            )?));
18622        }
18623        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
18624            return Ok(None);
18625        };
18626        let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
18627        Ok(Some(self.q8_fused2_core(
18628            p0.0,
18629            p1.0,
18630            &aq,
18631            &ad,
18632            w0.in_features(),
18633            p0.1,
18634            p1.1,
18635            p0.2,
18636        )?))
18637    }
18638
18639    /// Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes,
18640    /// quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).
18641    #[allow(clippy::too_many_arguments)]
18642    pub fn qmatvec_q8_fused2_raw(
18643        &self,
18644        b0: &CudaSlice<u8>,
18645        b1: &CudaSlice<u8>,
18646        x: &CudaSlice<f32>,
18647        in_f: usize,
18648        out0: usize,
18649        out1: usize,
18650        row_bytes: usize,
18651    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18652        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
18653        self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
18654    }
18655
18656    /// FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512).
18657    /// Same block-offset recipe as `matmul_q8_fused2` with three ranges. BIT-IDENTICAL per
18658    /// (tensor,row) to three separate m=1 MMVQ launches.
18659    /// FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain
18660    /// identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.
18661    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18662    pub fn matmul_q4_fused3(
18663        &self,
18664        w0: &crate::model::GpuTensor,
18665        w1: &crate::model::GpuTensor,
18666        w2: &crate::model::GpuTensor,
18667        aq: &CudaSlice<i8>,
18668        ad: &CudaSlice<f32>,
18669    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
18670    {
18671        use crate::model::GpuTensor;
18672        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18673            match w {
18674                GpuTensor::Quant {
18675                    qtype, row_bytes, ..
18676                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18677                _ => None,
18678            }
18679        };
18680        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
18681            return Ok(None);
18682        };
18683        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
18684            return Ok(None);
18685        }
18686        // Effective (bytes, rp) per tensor: mirror (rp4) OR the in-place swap (rp flag,
18687        // bytes already split). Mixed layouts cannot share one fused launch -> fall back to
18688        // the separate matvecs (each routes its own rp).
18689        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18690            match w {
18691                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18692                    Some(m) => (m, true),
18693                    None => (bytes, *rp),
18694                },
18695                _ => unreachable!(),
18696            }
18697        }
18698        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
18699        if rp0 != rp1 || rp1 != rp2 {
18700            return Ok(None);
18701        }
18702        let rp = rp0;
18703        let rpb: u32 = 4;
18704        // mr1 (one row/warp, 2026-07-14): follows the singles' MEMRA_Q40_MR default — the
18705        // fused t=1 kernels were left on mr2 when the singles flipped (DRAM-duty map:
18706        // fused3 57% / fused2 86%; small qkv segments starve under mr2's half grid).
18707        let mr1 = rp && Self::q40_mr1_on();
18708        let nb = |o: usize| {
18709            if mr1 {
18710                (o as u32).div_ceil(rpb)
18711            } else {
18712                (o as u32).div_ceil(2).div_ceil(rpb)
18713            }
18714        };
18715        let grid = nb(o0) + nb(o1) + nb(o2);
18716        let mut y0 = self.alloc_uninit::<f32>(o0)?;
18717        let mut y1 = self.alloc_uninit::<f32>(o1)?;
18718        let mut y2 = self.alloc_uninit::<f32>(o2)?;
18719        let f = self.func(if mr1 {
18720            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
18721        } else if rp {
18722            "qmatvec_q4_0_mmvq_fused3_rp"
18723        } else {
18724            "qmatvec_q4_0_mmvq_fused3"
18725        });
18726        let cfg = LaunchConfig {
18727            grid_dim: (grid, 1, 1),
18728            block_dim: (32, rpb, 1),
18729            shared_mem_bytes: 0,
18730        };
18731        let inf = w0.in_features() as i32;
18732        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
18733        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
18734        // PDL wave-A (2026-07-23): the mr1 kernel carries MEMRA_PDL_ENTRY; only that
18735        // variant may take the programmatic-serialization launch.
18736        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18737            {
18738                use cudarc::driver::{DevicePtr, DevicePtrMut};
18739                let s = &self.gpu.stream();
18740                let (p0, _g0) = b0.device_ptr(s);
18741                let (p1, _g1) = b1.device_ptr(s);
18742                let (p2, _g2) = b2.device_ptr(s);
18743                let (paq, _g3) = aq.device_ptr(s);
18744                let (pad, _g4) = ad.device_ptr(s);
18745                let (py0, _g5) = y0.device_ptr_mut(s);
18746                let (py1, _g6) = y1.device_ptr_mut(s);
18747                let (py2, _g7) = y2.device_ptr_mut(s);
18748                let mut ps = [
18749                    &p0 as *const _ as *mut std::ffi::c_void,
18750                    &p1 as *const _ as *mut _,
18751                    &p2 as *const _ as *mut _,
18752                    &paq as *const _ as *mut _,
18753                    &pad as *const _ as *mut _,
18754                    &py0 as *const _ as *mut _,
18755                    &py1 as *const _ as *mut _,
18756                    &py2 as *const _ as *mut _,
18757                    &inf as *const _ as *mut _,
18758                    &oo0 as *const _ as *mut _,
18759                    &oo1 as *const _ as *mut _,
18760                    &oo2 as *const _ as *mut _,
18761                    &r0 as *const _ as *mut _,
18762                    &r1 as *const _ as *mut _,
18763                    &r2 as *const _ as *mut _,
18764                ];
18765                unsafe {
18766                    self.launch_pdl(
18767                        "qmatvec_q4_0_mmvq_fused3_mr1_rp",
18768                        (grid, 1, 1),
18769                        (32, rpb, 1),
18770                        &mut ps,
18771                    )?;
18772                }
18773            }
18774            return Ok(Some((y0, y1, y2)));
18775        }
18776        let __s_b = self.gpu.stream();
18777        let mut b = __s_b.launch_builder(&f);
18778        b.arg(b0)
18779            .arg(b1)
18780            .arg(b2)
18781            .arg(aq)
18782            .arg(ad)
18783            .arg(&mut y0)
18784            .arg(&mut y1)
18785            .arg(&mut y2)
18786            .arg(&inf)
18787            .arg(&oo0)
18788            .arg(&oo1)
18789            .arg(&oo2)
18790            .arg(&r0)
18791            .arg(&r1)
18792            .arg(&r2);
18793        unsafe {
18794            b.launch(cfg)?;
18795        }
18796        Ok(Some((y0, y1, y2)))
18797    }
18798
18799    /// Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs.
18800    /// Returns Ok(false) when the fused path is unavailable (caller falls back).
18801    #[allow(clippy::too_many_arguments)]
18802    pub fn matmul_q4_fused3_into(
18803        &self,
18804        w0: &crate::model::GpuTensor,
18805        w1: &crate::model::GpuTensor,
18806        w2: &crate::model::GpuTensor,
18807        aq: &CudaSlice<i8>,
18808        ad: &CudaSlice<f32>,
18809        y0: &mut CudaSlice<f32>,
18810        y1: &mut CudaSlice<f32>,
18811        y2: &mut CudaSlice<f32>,
18812    ) -> Result<bool, Box<dyn std::error::Error>> {
18813        use crate::model::GpuTensor;
18814        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18815            match w {
18816                GpuTensor::Quant {
18817                    qtype, row_bytes, ..
18818                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18819                _ => None,
18820            }
18821        };
18822        let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
18823            return Ok(false);
18824        };
18825        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
18826            return Ok(false);
18827        }
18828        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18829            match w {
18830                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18831                    Some(m) => (m, true),
18832                    None => (bytes, *rp),
18833                },
18834                _ => unreachable!(),
18835            }
18836        }
18837        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
18838        if rp0 != rp1 || rp1 != rp2 {
18839            return Ok(false);
18840        }
18841        let rp = rp0;
18842        let rpb: u32 = 4;
18843        let mr1 = rp && Self::q40_mr1_on();
18844        let nb = |o: usize| {
18845            if mr1 {
18846                (o as u32).div_ceil(rpb)
18847            } else {
18848                (o as u32).div_ceil(2).div_ceil(rpb)
18849            }
18850        };
18851        let grid = nb(o0) + nb(o1) + nb(o2);
18852        debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
18853        let f = self.func(if mr1 {
18854            "qmatvec_q4_0_mmvq_fused3_mr1_rp"
18855        } else if rp {
18856            "qmatvec_q4_0_mmvq_fused3_rp"
18857        } else {
18858            "qmatvec_q4_0_mmvq_fused3"
18859        });
18860        let cfg = LaunchConfig {
18861            grid_dim: (grid, 1, 1),
18862            block_dim: (32, rpb, 1),
18863            shared_mem_bytes: 0,
18864        };
18865        let inf = w0.in_features() as i32;
18866        let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
18867        let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
18868        // PDL wave-A: identical to the owned twin (capture-lane parity).
18869        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18870            use cudarc::driver::{DevicePtr, DevicePtrMut};
18871            let s = &self.gpu.stream();
18872            let (p0, _g0) = b0.device_ptr(s);
18873            let (p1, _g1) = b1.device_ptr(s);
18874            let (p2, _g2) = b2.device_ptr(s);
18875            let (paq, _g3) = aq.device_ptr(s);
18876            let (pad, _g4) = ad.device_ptr(s);
18877            let (py0, _g5) = y0.device_ptr_mut(s);
18878            let (py1, _g6) = y1.device_ptr_mut(s);
18879            let (py2, _g7) = y2.device_ptr_mut(s);
18880            let mut ps = [
18881                &p0 as *const _ as *mut std::ffi::c_void,
18882                &p1 as *const _ as *mut _,
18883                &p2 as *const _ as *mut _,
18884                &paq as *const _ as *mut _,
18885                &pad as *const _ as *mut _,
18886                &py0 as *const _ as *mut _,
18887                &py1 as *const _ as *mut _,
18888                &py2 as *const _ as *mut _,
18889                &inf as *const _ as *mut _,
18890                &oo0 as *const _ as *mut _,
18891                &oo1 as *const _ as *mut _,
18892                &oo2 as *const _ as *mut _,
18893                &r0 as *const _ as *mut _,
18894                &r1 as *const _ as *mut _,
18895                &r2 as *const _ as *mut _,
18896            ];
18897            unsafe {
18898                self.launch_pdl(
18899                    "qmatvec_q4_0_mmvq_fused3_mr1_rp",
18900                    (grid, 1, 1),
18901                    (32, rpb, 1),
18902                    &mut ps,
18903                )?;
18904            }
18905            return Ok(true);
18906        }
18907        let __s_b = self.gpu.stream();
18908        let mut b = __s_b.launch_builder(&f);
18909        b.arg(b0)
18910            .arg(b1)
18911            .arg(b2)
18912            .arg(aq)
18913            .arg(ad)
18914            .arg(&mut *y0)
18915            .arg(&mut *y1)
18916            .arg(&mut *y2)
18917            .arg(&inf)
18918            .arg(&oo0)
18919            .arg(&oo1)
18920            .arg(&oo2)
18921            .arg(&r0)
18922            .arg(&r1)
18923            .arg(&r2);
18924        unsafe {
18925            b.launch(cfg)?;
18926        }
18927        Ok(true)
18928    }
18929
18930    /// FUSED Q4_0 m=1 PAIR (gemma shared gate+up).
18931    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18932    pub fn matmul_q4_fused2(
18933        &self,
18934        w0: &crate::model::GpuTensor,
18935        w1: &crate::model::GpuTensor,
18936        aq: &CudaSlice<i8>,
18937        ad: &CudaSlice<f32>,
18938    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18939        use crate::model::GpuTensor;
18940        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18941            match w {
18942                GpuTensor::Quant {
18943                    qtype, row_bytes, ..
18944                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18945                _ => None,
18946            }
18947        };
18948        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
18949            return Ok(None);
18950        };
18951        if w0.in_features() != w1.in_features() {
18952            return Ok(None);
18953        }
18954        // Effective (bytes, rp) per tensor (mirror or in-place swap); mixed -> separate matvecs.
18955        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18956            match w {
18957                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18958                    Some(m) => (m, true),
18959                    None => (bytes, *rp),
18960                },
18961                _ => unreachable!(),
18962            }
18963        }
18964        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
18965        if rp0 != rp1 {
18966            return Ok(None);
18967        }
18968        let rp = rp0;
18969        let rpb: u32 = 4;
18970        // mr1 twin — see matmul_q4_fused3.
18971        let mr1 = rp && Self::q40_mr1_on();
18972        let nb = |o: usize| {
18973            if mr1 {
18974                (o as u32).div_ceil(rpb)
18975            } else {
18976                (o as u32).div_ceil(2).div_ceil(rpb)
18977            }
18978        };
18979        let grid = nb(o0) + nb(o1);
18980        let mut y0 = self.alloc_uninit::<f32>(o0)?;
18981        let mut y1 = self.alloc_uninit::<f32>(o1)?;
18982        let f = self.func(if mr1 {
18983            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
18984        } else if rp {
18985            "qmatvec_q4_0_mmvq_fused2_rp"
18986        } else {
18987            "qmatvec_q4_0_mmvq_fused2"
18988        });
18989        let cfg = LaunchConfig {
18990            grid_dim: (grid, 1, 1),
18991            block_dim: (32, rpb, 1),
18992            shared_mem_bytes: 0,
18993        };
18994        let inf = w0.in_features() as i32;
18995        let (oo0, oo1) = (o0 as i32, o1 as i32);
18996        let (r0, r1) = (rb0 as i64, rb1 as i64);
18997        // PDL wave-A: mr1 kernel carries MEMRA_PDL_ENTRY.
18998        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18999            {
19000                use cudarc::driver::{DevicePtr, DevicePtrMut};
19001                let s = &self.gpu.stream();
19002                let (p0, _g0) = b0.device_ptr(s);
19003                let (p1, _g1) = b1.device_ptr(s);
19004                let (paq, _g2) = aq.device_ptr(s);
19005                let (pad, _g3) = ad.device_ptr(s);
19006                let (py0, _g4) = y0.device_ptr_mut(s);
19007                let (py1, _g5) = y1.device_ptr_mut(s);
19008                let mut ps = [
19009                    &p0 as *const _ as *mut std::ffi::c_void,
19010                    &p1 as *const _ as *mut _,
19011                    &paq as *const _ as *mut _,
19012                    &pad as *const _ as *mut _,
19013                    &py0 as *const _ as *mut _,
19014                    &py1 as *const _ as *mut _,
19015                    &inf as *const _ as *mut _,
19016                    &oo0 as *const _ as *mut _,
19017                    &oo1 as *const _ as *mut _,
19018                    &r0 as *const _ as *mut _,
19019                    &r1 as *const _ as *mut _,
19020                ];
19021                unsafe {
19022                    self.launch_pdl(
19023                        "qmatvec_q4_0_mmvq_fused2_mr1_rp",
19024                        (grid, 1, 1),
19025                        (32, rpb, 1),
19026                        &mut ps,
19027                    )?;
19028                }
19029            }
19030            return Ok(Some((y0, y1)));
19031        }
19032        let __s_b = self.gpu.stream();
19033        let mut b = __s_b.launch_builder(&f);
19034        b.arg(b0)
19035            .arg(b1)
19036            .arg(aq)
19037            .arg(ad)
19038            .arg(&mut y0)
19039            .arg(&mut y1)
19040            .arg(&inf)
19041            .arg(&oo0)
19042            .arg(&oo1)
19043            .arg(&r0)
19044            .arg(&r1);
19045        unsafe {
19046            b.launch(cfg)?;
19047        }
19048        Ok(Some((y0, y1)))
19049    }
19050
19051    /// Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.
19052    pub fn matmul_q4_fused2_into(
19053        &self,
19054        w0: &crate::model::GpuTensor,
19055        w1: &crate::model::GpuTensor,
19056        aq: &CudaSlice<i8>,
19057        ad: &CudaSlice<f32>,
19058        y0: &mut CudaSlice<f32>,
19059        y1: &mut CudaSlice<f32>,
19060    ) -> Result<bool, Box<dyn std::error::Error>> {
19061        use crate::model::GpuTensor;
19062        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
19063            match w {
19064                GpuTensor::Quant {
19065                    qtype, row_bytes, ..
19066                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
19067                _ => None,
19068            }
19069        };
19070        let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
19071            return Ok(false);
19072        };
19073        if w0.in_features() != w1.in_features() {
19074            return Ok(false);
19075        }
19076        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
19077            match w {
19078                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
19079                    Some(m) => (m, true),
19080                    None => (bytes, *rp),
19081                },
19082                _ => unreachable!(),
19083            }
19084        }
19085        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
19086        if rp0 != rp1 {
19087            return Ok(false);
19088        }
19089        let rp = rp0;
19090        let rpb: u32 = 4;
19091        let mr1 = rp && Self::q40_mr1_on();
19092        let nb = |o: usize| {
19093            if mr1 {
19094                (o as u32).div_ceil(rpb)
19095            } else {
19096                (o as u32).div_ceil(2).div_ceil(rpb)
19097            }
19098        };
19099        let grid = nb(o0) + nb(o1);
19100        debug_assert!(y0.len() >= o0 && y1.len() >= o1);
19101        let f = self.func(if mr1 {
19102            "qmatvec_q4_0_mmvq_fused2_mr1_rp"
19103        } else if rp {
19104            "qmatvec_q4_0_mmvq_fused2_rp"
19105        } else {
19106            "qmatvec_q4_0_mmvq_fused2"
19107        });
19108        let cfg = LaunchConfig {
19109            grid_dim: (grid, 1, 1),
19110            block_dim: (32, rpb, 1),
19111            shared_mem_bytes: 0,
19112        };
19113        let inf = w0.in_features() as i32;
19114        let (oo0, oo1) = (o0 as i32, o1 as i32);
19115        let (r0, r1) = (rb0 as i64, rb1 as i64);
19116        // PDL wave-A: identical to the owned twin (capture-lane parity).
19117        if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
19118            use cudarc::driver::{DevicePtr, DevicePtrMut};
19119            let s = &self.gpu.stream();
19120            let (p0, _g0) = b0.device_ptr(s);
19121            let (p1, _g1) = b1.device_ptr(s);
19122            let (paq, _g2) = aq.device_ptr(s);
19123            let (pad, _g3) = ad.device_ptr(s);
19124            let (py0, _g4) = y0.device_ptr_mut(s);
19125            let (py1, _g5) = y1.device_ptr_mut(s);
19126            let mut ps = [
19127                &p0 as *const _ as *mut std::ffi::c_void,
19128                &p1 as *const _ as *mut _,
19129                &paq as *const _ as *mut _,
19130                &pad as *const _ as *mut _,
19131                &py0 as *const _ as *mut _,
19132                &py1 as *const _ as *mut _,
19133                &inf as *const _ as *mut _,
19134                &oo0 as *const _ as *mut _,
19135                &oo1 as *const _ as *mut _,
19136                &r0 as *const _ as *mut _,
19137                &r1 as *const _ as *mut _,
19138            ];
19139            unsafe {
19140                self.launch_pdl(
19141                    "qmatvec_q4_0_mmvq_fused2_mr1_rp",
19142                    (grid, 1, 1),
19143                    (32, rpb, 1),
19144                    &mut ps,
19145                )?;
19146            }
19147            return Ok(true);
19148        }
19149        let __s_b = self.gpu.stream();
19150        let mut b = __s_b.launch_builder(&f);
19151        b.arg(b0)
19152            .arg(b1)
19153            .arg(aq)
19154            .arg(ad)
19155            .arg(&mut *y0)
19156            .arg(&mut *y1)
19157            .arg(&inf)
19158            .arg(&oo0)
19159            .arg(&oo1)
19160            .arg(&r0)
19161            .arg(&r1);
19162        unsafe {
19163            b.launch(cfg)?;
19164        }
19165        Ok(true)
19166    }
19167
19168    /// BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in
19169    /// ONE segmented-grid launch — the up segment fills SMs as the gate segment drains
19170    /// (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical
19171    /// per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).
19172    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19173    pub fn matmul_q4_fused2_batched(
19174        &self,
19175        w0: &crate::model::GpuTensor,
19176        w1: &crate::model::GpuTensor,
19177        aq: &CudaSlice<i8>,
19178        ad: &CudaSlice<f32>,
19179        m: usize,
19180    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
19181        use crate::model::GpuTensor;
19182        if !(2..=8).contains(&m) {
19183            return Ok(None);
19184        }
19185        let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
19186            match w {
19187                GpuTensor::Quant {
19188                    qtype, row_bytes, ..
19189                } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
19190                _ => None,
19191            }
19192        };
19193        let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
19194            return Ok(None);
19195        };
19196        if w0.in_features() != w1.in_features() {
19197            return Ok(None);
19198        }
19199        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
19200            match w {
19201                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
19202                    Some(mr) => (mr, true),
19203                    None => (bytes, *rp),
19204                },
19205                _ => unreachable!(),
19206            }
19207        }
19208        let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
19209        if !rp0 || !rp1 {
19210            return Ok(None);
19211        }
19212        let mcols = Self::batched_mcols(m);
19213        let rpb: u32 = 4;
19214        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
19215        let grid = nb(o0) + nb(o1);
19216        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
19217        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
19218        let f = self.func(match mcols {
19219            2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
19220            4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
19221            _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
19222        });
19223        let cfg = LaunchConfig {
19224            grid_dim: (grid, 1, 1),
19225            block_dim: (32, rpb, 1),
19226            shared_mem_bytes: 0,
19227        };
19228        let inf = w0.in_features() as i32;
19229        let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
19230        let rb = rb0 as i64;
19231        let __s_b = self.gpu.stream();
19232        let mut b = __s_b.launch_builder(&f);
19233        b.arg(b0)
19234            .arg(b1)
19235            .arg(aq)
19236            .arg(ad)
19237            .arg(&mut y0)
19238            .arg(&mut y1)
19239            .arg(&inf)
19240            .arg(&oo0)
19241            .arg(&oo1)
19242            .arg(&mi)
19243            .arg(&rb);
19244        unsafe {
19245            b.launch(cfg)?;
19246        }
19247        Ok(Some((y0, y1)))
19248    }
19249
19250    /// BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the
19251    /// verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.
19252    #[allow(clippy::too_many_arguments)]
19253    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19254    pub fn matmul_q4_fused3_batched(
19255        &self,
19256        w0: &crate::model::GpuTensor,
19257        w1: &crate::model::GpuTensor,
19258        w2: &crate::model::GpuTensor,
19259        aq: &CudaSlice<i8>,
19260        ad: &CudaSlice<f32>,
19261        m: usize,
19262    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
19263    {
19264        use crate::model::GpuTensor;
19265        if !(2..=8).contains(&m) {
19266            return Ok(None);
19267        }
19268        let q4 = |w: &GpuTensor| -> Option<usize> {
19269            match w {
19270                GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
19271                _ => None,
19272            }
19273        };
19274        let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
19275            return Ok(None);
19276        };
19277        if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
19278            return Ok(None);
19279        }
19280        fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
19281            match w {
19282                GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
19283                    Some(mr) => (mr, true),
19284                    None => (bytes, *rp),
19285                },
19286                _ => unreachable!(),
19287            }
19288        }
19289        let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
19290        if !rp0 || !rp1 || !rp2 {
19291            return Ok(None);
19292        }
19293        let mcols = Self::batched_mcols(m);
19294        let rpb: u32 = 4;
19295        let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
19296        let grid = nb(o0) + nb(o1) + nb(o2);
19297        let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
19298        let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
19299        let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
19300        let f = self.func(match mcols {
19301            2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
19302            4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
19303            _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
19304        });
19305        let cfg = LaunchConfig {
19306            grid_dim: (grid, 1, 1),
19307            block_dim: (32, rpb, 1),
19308            shared_mem_bytes: 0,
19309        };
19310        let inf = w0.in_features() as i32;
19311        let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
19312        let rb = 0i64;
19313        let __s_b = self.gpu.stream();
19314        let mut b = __s_b.launch_builder(&f);
19315        b.arg(b0)
19316            .arg(b1)
19317            .arg(b2)
19318            .arg(aq)
19319            .arg(ad)
19320            .arg(&mut y0)
19321            .arg(&mut y1)
19322            .arg(&mut y2)
19323            .arg(&inf)
19324            .arg(&oo0)
19325            .arg(&oo1)
19326            .arg(&oo2)
19327            .arg(&mi)
19328            .arg(&rb);
19329        unsafe {
19330            b.launch(cfg)?;
19331        }
19332        Ok(Some((y0, y1, y2)))
19333    }
19334
19335    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19336    pub fn matmul_q8_fused3(
19337        &self,
19338        w0: &crate::model::GpuTensor,
19339        w1: &crate::model::GpuTensor,
19340        w2: &crate::model::GpuTensor,
19341        aq: &CudaSlice<i8>,
19342        ad: &CudaSlice<f32>,
19343    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
19344    {
19345        // e4m3 twin (lane/fp8-decode-v1): the full-attn wq/wk/wv triple — on the NV-27B those three
19346        // are per-tensor FP8, so native residency without this arm meant three separate launches.
19347        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
19348            return Ok(Some(self.e4m3_fused3_core(
19349                p0.0,
19350                p1.0,
19351                p2.0,
19352                aq,
19353                ad,
19354                w0.in_features(),
19355                p0.1,
19356                p1.1,
19357                p2.1,
19358                p0.2,
19359                p0.3,
19360                p1.3,
19361                p2.3,
19362            )?));
19363        }
19364        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
19365            return Ok(None);
19366        };
19367        Ok(Some(self.q8_fused3_core(
19368            p0.0,
19369            p1.0,
19370            p2.0,
19371            aq,
19372            ad,
19373            w0.in_features(),
19374            p0.1,
19375            p1.1,
19376            p2.1,
19377            p0.2,
19378        )?))
19379    }
19380
19381    #[allow(clippy::too_many_arguments)]
19382    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19383    fn q8_fused3_core(
19384        &self,
19385        b0: &CudaSlice<u8>,
19386        b1: &CudaSlice<u8>,
19387        b2: &CudaSlice<u8>,
19388        aq: &CudaSlice<i8>,
19389        ad: &CudaSlice<f32>,
19390        in_f: usize,
19391        out0: usize,
19392        out1: usize,
19393        out2: usize,
19394        row_bytes: usize,
19395    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19396        const ROWS_PER_BLOCK: u32 = 4;
19397        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19398        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19399        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19400        let f = self.func("qmatvec_q8_0_mmvq_fused3");
19401        let mut y0 = self.alloc_uninit::<f32>(out0)?;
19402        let mut y1 = self.alloc_uninit::<f32>(out1)?;
19403        let mut y2 = self.alloc_uninit::<f32>(out2)?;
19404        let cfg = LaunchConfig {
19405            grid_dim: (nb0 + nb1 + nb2, 1, 1),
19406            block_dim: (32, ROWS_PER_BLOCK, 1),
19407            shared_mem_bytes: 0,
19408        };
19409        let (inf, o0, o1, o2, rbl) = (
19410            in_f as i32,
19411            out0 as i32,
19412            out1 as i32,
19413            out2 as i32,
19414            row_bytes as i64,
19415        );
19416        let __s_b = self.gpu.stream();
19417        let mut b = __s_b.launch_builder(&f);
19418        b.arg(b0)
19419            .arg(b1)
19420            .arg(b2)
19421            .arg(aq)
19422            .arg(ad)
19423            .arg(&mut y0)
19424            .arg(&mut y1)
19425            .arg(&mut y2)
19426            .arg(&inf)
19427            .arg(&o0)
19428            .arg(&o1)
19429            .arg(&o2)
19430            .arg(&rbl);
19431        unsafe {
19432            b.launch(cfg)?;
19433        }
19434        Ok((y0, y1, y2))
19435    }
19436
19437    /// Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).
19438    #[allow(clippy::too_many_arguments)]
19439    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19440    pub fn qmatvec_q8_fused3_raw(
19441        &self,
19442        b0: &CudaSlice<u8>,
19443        b1: &CudaSlice<u8>,
19444        b2: &CudaSlice<u8>,
19445        x: &CudaSlice<f32>,
19446        in_f: usize,
19447        out0: usize,
19448        out1: usize,
19449        out2: usize,
19450        row_bytes: usize,
19451    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19452        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
19453        self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
19454    }
19455
19456    /// BATCHED twin of `matmul_q8_fused2` for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call
19457    /// sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens.
19458    /// Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical
19459    /// row mapping (Q8_0's batched_variant is always "base") -> BIT-IDENTICAL to the two
19460    /// per-tensor _b2/_b4 launches `matmul_decode_exact` dispatches at m=2-4, with the caller's
19461    /// single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is
19462    /// deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 /
19463    /// in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps
19464    /// dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused
19465    /// twin must not introduce a batched program the reference path would not run).
19466    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19467    pub fn matmul_q8_fused2_t(
19468        &self,
19469        w0: &crate::model::GpuTensor,
19470        w1: &crate::model::GpuTensor,
19471        aq: &CudaSlice<i8>,
19472        ad: &CudaSlice<f32>,
19473        m: usize,
19474    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
19475        // m<=8 (lane/q27-deepdive, 2026-08-05): was 2..=4 (the verify tier's mcols 2/4). The
19476        // serving tick's mcols-8 tier now has its fused2_b8 wrapper, so c=5..8 batched decode
19477        // fuses too — same template body, still bit-identical to the two _b8 launches.
19478        if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
19479            return Ok(None);
19480        }
19481        // e4m3 twin: MEMRA_B8 parity — without it m=5..8 e4m3 decode runs the per-m grid.y=m path,
19482        // so the fused b8 launch would introduce a batched program the reference path would not run.
19483        if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
19484            if m > 4 && !Self::b8_enabled() {
19485                return Ok(None);
19486            }
19487            return Ok(Some(self.e4m3_fused2_t_core(
19488                p0.0,
19489                p1.0,
19490                aq,
19491                ad,
19492                m,
19493                w0.in_features(),
19494                p0.1,
19495                p1.1,
19496                p0.2,
19497                p0.3,
19498                p1.3,
19499            )?));
19500        }
19501        let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
19502            return Ok(None);
19503        };
19504        Ok(Some(self.q8_fused2_t_core(
19505            p0.0,
19506            p1.0,
19507            aq,
19508            ad,
19509            m,
19510            w0.in_features(),
19511            p0.1,
19512            p1.1,
19513            p0.2,
19514        )?))
19515    }
19516
19517    #[allow(clippy::too_many_arguments)]
19518    fn q8_fused2_t_core(
19519        &self,
19520        b0: &CudaSlice<u8>,
19521        b1: &CudaSlice<u8>,
19522        aq: &CudaSlice<i8>,
19523        ad: &CudaSlice<f32>,
19524        m: usize,
19525        in_f: usize,
19526        out0: usize,
19527        out1: usize,
19528        row_bytes: usize,
19529    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19530        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
19531        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19532        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19533        let f = self.func(match Self::batched_mcols(m) {
19534            2 => "qmatvec_q8_0_mmvq_fused2_b2",
19535            4 => "qmatvec_q8_0_mmvq_fused2_b4",
19536            // b8 = the SERVING tier (lane/q27-deepdive): c=5..8 batched decode.
19537            _ => "qmatvec_q8_0_mmvq_fused2_b8",
19538        });
19539        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
19540        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
19541        let cfg = LaunchConfig {
19542            grid_dim: (nb0 + nb1, 1, 1),
19543            block_dim: (32, ROWS_PER_BLOCK, 1),
19544            shared_mem_bytes: 0,
19545        };
19546        let (inf, o0, o1, mi, rbl) = (
19547            in_f as i32,
19548            out0 as i32,
19549            out1 as i32,
19550            m as i32,
19551            row_bytes as i64,
19552        );
19553        let __s_b = self.gpu.stream();
19554        let mut b = __s_b.launch_builder(&f);
19555        b.arg(b0)
19556            .arg(b1)
19557            .arg(aq)
19558            .arg(ad)
19559            .arg(&mut y0)
19560            .arg(&mut y1)
19561            .arg(&inf)
19562            .arg(&o0)
19563            .arg(&o1)
19564            .arg(&mi)
19565            .arg(&rbl);
19566        unsafe {
19567            b.launch(cfg)?;
19568        }
19569        Ok((y0, y1))
19570    }
19571
19572    /// Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal
19573    /// q8_1 quant of the [m, in_f] activation), no env gating.
19574    #[allow(clippy::too_many_arguments)]
19575    pub fn qmatvec_q8_fused2_t_raw(
19576        &self,
19577        b0: &CudaSlice<u8>,
19578        b1: &CudaSlice<u8>,
19579        x: &CudaSlice<f32>,
19580        m: usize,
19581        in_f: usize,
19582        out0: usize,
19583        out1: usize,
19584        row_bytes: usize,
19585    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19586        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19587        self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
19588    }
19589
19590    /// BATCHED twin of `matmul_q8_fused3` (wq+wk+wv at verify t=2-4). Same contract as
19591    /// `matmul_q8_fused2_t` with three ranges.
19592    #[allow(clippy::too_many_arguments)]
19593    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19594    pub fn matmul_q8_fused3_t(
19595        &self,
19596        w0: &crate::model::GpuTensor,
19597        w1: &crate::model::GpuTensor,
19598        w2: &crate::model::GpuTensor,
19599        aq: &CudaSlice<i8>,
19600        ad: &CudaSlice<f32>,
19601        m: usize,
19602    ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
19603    {
19604        if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
19605            return Ok(None);
19606        }
19607        if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
19608            return Ok(Some(self.e4m3_fused3_t_core(
19609                p0.0,
19610                p1.0,
19611                p2.0,
19612                aq,
19613                ad,
19614                m,
19615                w0.in_features(),
19616                p0.1,
19617                p1.1,
19618                p2.1,
19619                p0.2,
19620                p0.3,
19621                p1.3,
19622                p2.3,
19623            )?));
19624        }
19625        let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
19626            return Ok(None);
19627        };
19628        Ok(Some(self.q8_fused3_t_core(
19629            p0.0,
19630            p1.0,
19631            p2.0,
19632            aq,
19633            ad,
19634            m,
19635            w0.in_features(),
19636            p0.1,
19637            p1.1,
19638            p2.1,
19639            p0.2,
19640        )?))
19641    }
19642
19643    #[allow(clippy::too_many_arguments)]
19644    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19645    fn q8_fused3_t_core(
19646        &self,
19647        b0: &CudaSlice<u8>,
19648        b1: &CudaSlice<u8>,
19649        b2: &CudaSlice<u8>,
19650        aq: &CudaSlice<i8>,
19651        ad: &CudaSlice<f32>,
19652        m: usize,
19653        in_f: usize,
19654        out0: usize,
19655        out1: usize,
19656        out2: usize,
19657        row_bytes: usize,
19658    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19659        const ROWS_PER_BLOCK: u32 = 4;
19660        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19661        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19662        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19663        let f = self.func(if Self::batched_mcols(m) == 2 {
19664            "qmatvec_q8_0_mmvq_fused3_b2"
19665        } else {
19666            "qmatvec_q8_0_mmvq_fused3_b4"
19667        });
19668        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
19669        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
19670        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
19671        let cfg = LaunchConfig {
19672            grid_dim: (nb0 + nb1 + nb2, 1, 1),
19673            block_dim: (32, ROWS_PER_BLOCK, 1),
19674            shared_mem_bytes: 0,
19675        };
19676        let (inf, o0, o1, o2, mi, rbl) = (
19677            in_f as i32,
19678            out0 as i32,
19679            out1 as i32,
19680            out2 as i32,
19681            m as i32,
19682            row_bytes as i64,
19683        );
19684        let __s_b = self.gpu.stream();
19685        let mut b = __s_b.launch_builder(&f);
19686        b.arg(b0)
19687            .arg(b1)
19688            .arg(b2)
19689            .arg(aq)
19690            .arg(ad)
19691            .arg(&mut y0)
19692            .arg(&mut y1)
19693            .arg(&mut y2)
19694            .arg(&inf)
19695            .arg(&o0)
19696            .arg(&o1)
19697            .arg(&o2)
19698            .arg(&mi)
19699            .arg(&rbl);
19700        unsafe {
19701            b.launch(cfg)?;
19702        }
19703        Ok((y0, y1, y2))
19704    }
19705
19706    /// Test entry for the kernel_check gate: fused3 batched from raw weight bytes.
19707    #[allow(clippy::too_many_arguments)]
19708    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19709    pub fn qmatvec_q8_fused3_t_raw(
19710        &self,
19711        b0: &CudaSlice<u8>,
19712        b1: &CudaSlice<u8>,
19713        b2: &CudaSlice<u8>,
19714        x: &CudaSlice<f32>,
19715        m: usize,
19716        in_f: usize,
19717        out0: usize,
19718        out1: usize,
19719        out2: usize,
19720        row_bytes: usize,
19721    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19722        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19723        self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
19724    }
19725
19726    /// Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in `matmul_pre_dual_noscale`
19727    /// (lane/q27-deepdive, 2026-08-05). Default ON; `MEMRA_Q8_FFN_FUSE2=0` restores the
19728    /// two-`matmul_pre_noscale` pair. Read once — the dispatch must not vary within a run.
19729    pub fn q8_ffn_fuse2_on(&self) -> bool {
19730        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19731        *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
19732    }
19733
19734    /// Eligibility + param extraction for the fused q8_0 launches: every tensor must be Quant Q8_0
19735    /// with macro-scale 1.0 (always true for GGUF q8_0; only NVFP4 carries scale) and share w[0]'s
19736    /// in_f (q8_0 row_bytes is a pure function of in_f, so equal in_f => equal row_bytes). MEMRA_MMVQ
19737    /// must be on: the fused body is the MMVQ kernel; without it decode m=1 runs dp4a and fusing
19738    /// would mix dispatch families (FP-order law). MEMRA_Q8_DUAL=0 = rollback seam.
19739    #[allow(clippy::type_complexity)]
19740    fn q8_fused_params<'w, const N: usize>(
19741        &self,
19742        ws: &[&'w crate::model::GpuTensor; N],
19743    ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
19744        use crate::model::GpuTensor;
19745        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
19746            return None;
19747        }
19748        if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
19749            return None;
19750        }
19751        let in_f = ws[0].in_features();
19752        let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
19753        for (i, w) in ws.iter().enumerate() {
19754            match w {
19755                GpuTensor::Quant {
19756                    bytes,
19757                    qtype,
19758                    row_bytes,
19759                    scale,
19760                    ..
19761                } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
19762                    out[i] = Some((bytes, w.out_features(), *row_bytes))
19763                }
19764                _ => return None,
19765            }
19766        }
19767        Some(out.map(|o| o.unwrap()))
19768    }
19769
19770    /// Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05).
19771    /// Default ON; `MEMRA_E4M3_DUAL=0` restores the per-tensor m=1/batched launches.
19772    pub fn e4m3_dual_on(&self) -> bool {
19773        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19774        *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
19775    }
19776
19777    /// Eligibility + param extraction for the FUSED e4m3 launches — the QT_F8_E4M3 twin of
19778    /// `q8_fused_params`. Differences that are inherent to the dtype, not policy:
19779    ///   * each tensor carries its OWN per-tensor `weight_scale` (returned as the 4th field);
19780    ///     Q8_0 hard-requires scale==1.0 because it has no macro-scale at all.
19781    ///   * no MEMRA_MMVQ gate: `mmvq_supports` exempts QT_F8_E4M3 (the e4m3 mmvq family is that
19782    ///     dtype's ONLY int8-act kernel class), so the per-tensor fallback these fused kernels
19783    ///     replace is ALWAYS the same mmvq body under every env — the FP-order law holds.
19784    ///   * `row_bytes == in_f` is asserted rather than derived: the native-residency load arm keeps
19785    ///     the checkpoint's raw [out_f, in_f] rows, and a re-encoded slab must never reach here.
19786    ///     Rejects any split-plane mirror (`rp`/`rp4`): there is no `_rp` e4m3 fused form, so fusing
19787    ///     there would swap dispatch families mid-model. MEMRA_E4M3_DUAL=0 = rollback seam.
19788    #[allow(clippy::type_complexity)]
19789    fn e4m3_fused_params<'w, const N: usize>(
19790        &self,
19791        ws: &[&'w crate::model::GpuTensor; N],
19792    ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
19793        use crate::model::GpuTensor;
19794        if !self.e4m3_dual_on() {
19795            return None;
19796        }
19797        let in_f = ws[0].in_features();
19798        let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
19799        for (i, w) in ws.iter().enumerate() {
19800            match w {
19801                GpuTensor::Quant {
19802                    bytes,
19803                    qtype,
19804                    row_bytes,
19805                    scale,
19806                    rp,
19807                    rp4,
19808                    ..
19809                } if *qtype == QT_F8_E4M3
19810                    && w.in_features() == in_f
19811                    && *row_bytes == in_f
19812                    && !*rp
19813                    && rp4.is_none() =>
19814                {
19815                    out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
19816                }
19817                _ => return None,
19818            }
19819        }
19820        Some(out.map(|o| o.unwrap()))
19821    }
19822
19823    /// FUSED e4m3 m=1 PAIR. Block-offset split (`qmatvec_e4m3_mmvq_fused2`), per-tensor
19824    /// weight_scale folded at the write like the single-tensor `qmatvec_e4m3_mmvq` — so per
19825    /// (tensor,row) this is BIT-IDENTICAL to two separate m=1 launches, scale included.
19826    #[allow(clippy::too_many_arguments)]
19827    fn e4m3_fused2_core(
19828        &self,
19829        b0: &CudaSlice<u8>,
19830        b1: &CudaSlice<u8>,
19831        aq: &CudaSlice<i8>,
19832        ad: &CudaSlice<f32>,
19833        in_f: usize,
19834        out0: usize,
19835        out1: usize,
19836        row_bytes: usize,
19837        ws0: f32,
19838        ws1: f32,
19839    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19840        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
19841        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19842        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19843        let f = self.func("qmatvec_e4m3_mmvq_fused2");
19844        let mut y0 = self.alloc_uninit::<f32>(out0)?;
19845        let mut y1 = self.alloc_uninit::<f32>(out1)?;
19846        let cfg = LaunchConfig {
19847            grid_dim: (nb0 + nb1, 1, 1),
19848            block_dim: (32, ROWS_PER_BLOCK, 1),
19849            shared_mem_bytes: 0,
19850        };
19851        let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
19852        let __s_b = self.gpu.stream();
19853        let mut b = __s_b.launch_builder(&f);
19854        b.arg(b0)
19855            .arg(b1)
19856            .arg(aq)
19857            .arg(ad)
19858            .arg(&mut y0)
19859            .arg(&mut y1)
19860            .arg(&inf)
19861            .arg(&o0)
19862            .arg(&o1)
19863            .arg(&rbl)
19864            .arg(&ws0)
19865            .arg(&ws1);
19866        unsafe {
19867            b.launch(cfg)?;
19868        }
19869        Ok((y0, y1))
19870    }
19871
19872    /// FUSED e4m3 m=1 TRIPLE (`qmatvec_e4m3_mmvq_fused3`). Same contract as the pair.
19873    #[allow(clippy::too_many_arguments)]
19874    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19875    fn e4m3_fused3_core(
19876        &self,
19877        b0: &CudaSlice<u8>,
19878        b1: &CudaSlice<u8>,
19879        b2: &CudaSlice<u8>,
19880        aq: &CudaSlice<i8>,
19881        ad: &CudaSlice<f32>,
19882        in_f: usize,
19883        out0: usize,
19884        out1: usize,
19885        out2: usize,
19886        row_bytes: usize,
19887        ws0: f32,
19888        ws1: f32,
19889        ws2: f32,
19890    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19891        const ROWS_PER_BLOCK: u32 = 4;
19892        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19893        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19894        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19895        let f = self.func("qmatvec_e4m3_mmvq_fused3");
19896        let mut y0 = self.alloc_uninit::<f32>(out0)?;
19897        let mut y1 = self.alloc_uninit::<f32>(out1)?;
19898        let mut y2 = self.alloc_uninit::<f32>(out2)?;
19899        let cfg = LaunchConfig {
19900            grid_dim: (nb0 + nb1 + nb2, 1, 1),
19901            block_dim: (32, ROWS_PER_BLOCK, 1),
19902            shared_mem_bytes: 0,
19903        };
19904        let (inf, o0, o1, o2, rbl) = (
19905            in_f as i32,
19906            out0 as i32,
19907            out1 as i32,
19908            out2 as i32,
19909            row_bytes as i64,
19910        );
19911        let __s_b = self.gpu.stream();
19912        let mut b = __s_b.launch_builder(&f);
19913        b.arg(b0)
19914            .arg(b1)
19915            .arg(b2)
19916            .arg(aq)
19917            .arg(ad)
19918            .arg(&mut y0)
19919            .arg(&mut y1)
19920            .arg(&mut y2)
19921            .arg(&inf)
19922            .arg(&o0)
19923            .arg(&o1)
19924            .arg(&o2)
19925            .arg(&rbl)
19926            .arg(&ws0)
19927            .arg(&ws1)
19928            .arg(&ws2);
19929        unsafe {
19930            b.launch(cfg)?;
19931        }
19932        Ok((y0, y1, y2))
19933    }
19934
19935    /// FUSED e4m3 m=1 SIX-GROUP (`qmatvec_e4m3_mmvq_fused6`) — the KDA six-projection group on a
19936    /// uniformly-e4m3 checkpoint. Same contract as the pair and the triple: one shared q8_1
19937    /// activation, one shared `row_bytes` (e4m3 rows are `in_f` bytes), a per-range weight scale,
19938    /// and per (range,row) output BITS identical to six separate m=1 launches.
19939    ///
19940    /// Writes into caller-owned outputs so the KDA door can keep its existing allocation shape.
19941    #[allow(clippy::too_many_arguments)]
19942    pub fn e4m3_fused6_into(
19943        &self,
19944        w: [&CudaSlice<u8>; 6],
19945        aq: &CudaSlice<i8>,
19946        ad: &CudaSlice<f32>,
19947        in_f: usize,
19948        dims: [usize; 6],
19949        row_bytes: usize,
19950        ws: [f32; 6],
19951        outs: &mut [CudaSlice<f32>; 6],
19952    ) -> Result<(), Box<dyn std::error::Error>> {
19953        self.e4m3_fused6_into_arm(
19954            w,
19955            aq,
19956            ad,
19957            in_f,
19958            dims,
19959            row_bytes,
19960            ws,
19961            outs,
19962            e4m3_row_ilp_level(),
19963        )
19964    }
19965
19966    /// The six-group launch with the arm chosen EXPLICITLY rather than from the door. The gate
19967    /// drives every arm in one process, which a `OnceLock`-backed flag read cannot express;
19968    /// keeping the policy in the wrapper above means the gate still exercises the served program.
19969    /// `arm`: 0 = serial, 1 = ILP.
19970    #[allow(clippy::too_many_arguments)]
19971    pub fn e4m3_fused6_into_arm(
19972        &self,
19973        w: [&CudaSlice<u8>; 6],
19974        aq: &CudaSlice<i8>,
19975        ad: &CudaSlice<f32>,
19976        in_f: usize,
19977        dims: [usize; 6],
19978        row_bytes: usize,
19979        ws: [f32; 6],
19980        outs: &mut [CudaSlice<f32>; 6],
19981        arm: u32,
19982    ) -> Result<(), Box<dyn std::error::Error>> {
19983        const ROWS_PER_BLOCK: u32 = 4;
19984        let blocks: u32 = dims
19985            .iter()
19986            .map(|&o| (o as u32).div_ceil(ROWS_PER_BLOCK))
19987            .sum();
19988        // arms 2/3/4 are the ILP-depth sweep (2, 8, 16); the profile says loads in flight, not
19989        // bandwidth, is what this family is short of. Bench-only until one earns a receipt.
19990        let f = self.func(match arm {
19991            0 => "qmatvec_e4m3_mmvq_fused6",
19992            1 => "qmatvec_e4m3_mmvq_fused6_ilp",
19993            2 => "qmatvec_e4m3_mmvq_fused6_ilp2",
19994            3 => "qmatvec_e4m3_mmvq_fused6_ilp8",
19995            _ => "qmatvec_e4m3_mmvq_fused6_ilp16",
19996        });
19997        let cfg = LaunchConfig {
19998            grid_dim: (blocks, 1, 1),
19999            block_dim: (32, ROWS_PER_BLOCK, 1),
20000            shared_mem_bytes: 0,
20001        };
20002        let inf = in_f as i32;
20003        let o: [i32; 6] = std::array::from_fn(|i| dims[i] as i32);
20004        let rbl = row_bytes as i64;
20005        let (o0, o1) = outs.split_at_mut(1);
20006        let (o1, o2) = o1.split_at_mut(1);
20007        let (o2, o3) = o2.split_at_mut(1);
20008        let (o3, o4) = o3.split_at_mut(1);
20009        let (o4, o5) = o4.split_at_mut(1);
20010        let __s_b = self.gpu.stream();
20011        let mut b = __s_b.launch_builder(&f);
20012        b.arg(w[0])
20013            .arg(w[1])
20014            .arg(w[2])
20015            .arg(w[3])
20016            .arg(w[4])
20017            .arg(w[5])
20018            .arg(aq)
20019            .arg(ad)
20020            .arg(&mut o0[0])
20021            .arg(&mut o1[0])
20022            .arg(&mut o2[0])
20023            .arg(&mut o3[0])
20024            .arg(&mut o4[0])
20025            .arg(&mut o5[0])
20026            .arg(&inf)
20027            .arg(&o[0])
20028            .arg(&o[1])
20029            .arg(&o[2])
20030            .arg(&o[3])
20031            .arg(&o[4])
20032            .arg(&o[5])
20033            .arg(&rbl)
20034            .arg(&ws[0])
20035            .arg(&ws[1])
20036            .arg(&ws[2])
20037            .arg(&ws[3])
20038            .arg(&ws[4])
20039            .arg(&ws[5]);
20040        unsafe {
20041            b.launch(cfg)?;
20042        }
20043        Ok(())
20044    }
20045
20046    /// BATCHED FUSED e4m3 pair (m=2..8). The batched kernels carry no `ws` arg (every batched
20047    /// kernel in the tree is scale-free), so each output takes its own `scale_inplace` — the
20048    /// SAME post-op the per-tensor batched dispatch applies, hence still bit-identical.
20049    #[allow(clippy::too_many_arguments)]
20050    fn e4m3_fused2_t_core(
20051        &self,
20052        b0: &CudaSlice<u8>,
20053        b1: &CudaSlice<u8>,
20054        aq: &CudaSlice<i8>,
20055        ad: &CudaSlice<f32>,
20056        m: usize,
20057        in_f: usize,
20058        out0: usize,
20059        out1: usize,
20060        row_bytes: usize,
20061        ws0: f32,
20062        ws1: f32,
20063    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20064        const ROWS_PER_BLOCK: u32 = 4;
20065        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
20066        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
20067        let f = self.func(match Self::batched_mcols(m) {
20068            2 => "qmatvec_e4m3_mmvq_fused2_b2",
20069            4 => "qmatvec_e4m3_mmvq_fused2_b4",
20070            _ => "qmatvec_e4m3_mmvq_fused2_b8",
20071        });
20072        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
20073        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
20074        let cfg = LaunchConfig {
20075            grid_dim: (nb0 + nb1, 1, 1),
20076            block_dim: (32, ROWS_PER_BLOCK, 1),
20077            shared_mem_bytes: 0,
20078        };
20079        let (inf, o0, o1, mi, rbl) = (
20080            in_f as i32,
20081            out0 as i32,
20082            out1 as i32,
20083            m as i32,
20084            row_bytes as i64,
20085        );
20086        let __s_b = self.gpu.stream();
20087        let mut b = __s_b.launch_builder(&f);
20088        b.arg(b0)
20089            .arg(b1)
20090            .arg(aq)
20091            .arg(ad)
20092            .arg(&mut y0)
20093            .arg(&mut y1)
20094            .arg(&inf)
20095            .arg(&o0)
20096            .arg(&o1)
20097            .arg(&mi)
20098            .arg(&rbl);
20099        unsafe {
20100            b.launch(cfg)?;
20101        }
20102        if ws0 != 1.0 {
20103            self.scale_inplace(&mut y0, ws0, m * out0)?;
20104        }
20105        if ws1 != 1.0 {
20106            self.scale_inplace(&mut y1, ws1, m * out1)?;
20107        }
20108        Ok((y0, y1))
20109    }
20110
20111    /// BATCHED FUSED e4m3 triple (m=2..4). Same contract as the batched pair.
20112    #[allow(clippy::too_many_arguments)]
20113    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
20114    fn e4m3_fused3_t_core(
20115        &self,
20116        b0: &CudaSlice<u8>,
20117        b1: &CudaSlice<u8>,
20118        b2: &CudaSlice<u8>,
20119        aq: &CudaSlice<i8>,
20120        ad: &CudaSlice<f32>,
20121        m: usize,
20122        in_f: usize,
20123        out0: usize,
20124        out1: usize,
20125        out2: usize,
20126        row_bytes: usize,
20127        ws0: f32,
20128        ws1: f32,
20129        ws2: f32,
20130    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20131        const ROWS_PER_BLOCK: u32 = 4;
20132        let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
20133        let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
20134        let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
20135        let f = self.func(if Self::batched_mcols(m) == 2 {
20136            "qmatvec_e4m3_mmvq_fused3_b2"
20137        } else {
20138            "qmatvec_e4m3_mmvq_fused3_b4"
20139        });
20140        let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
20141        let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
20142        let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
20143        let cfg = LaunchConfig {
20144            grid_dim: (nb0 + nb1 + nb2, 1, 1),
20145            block_dim: (32, ROWS_PER_BLOCK, 1),
20146            shared_mem_bytes: 0,
20147        };
20148        let (inf, o0, o1, o2, mi, rbl) = (
20149            in_f as i32,
20150            out0 as i32,
20151            out1 as i32,
20152            out2 as i32,
20153            m as i32,
20154            row_bytes as i64,
20155        );
20156        let __s_b = self.gpu.stream();
20157        let mut b = __s_b.launch_builder(&f);
20158        b.arg(b0)
20159            .arg(b1)
20160            .arg(b2)
20161            .arg(aq)
20162            .arg(ad)
20163            .arg(&mut y0)
20164            .arg(&mut y1)
20165            .arg(&mut y2)
20166            .arg(&inf)
20167            .arg(&o0)
20168            .arg(&o1)
20169            .arg(&o2)
20170            .arg(&mi)
20171            .arg(&rbl);
20172        unsafe {
20173            b.launch(cfg)?;
20174        }
20175        if ws0 != 1.0 {
20176            self.scale_inplace(&mut y0, ws0, m * out0)?;
20177        }
20178        if ws1 != 1.0 {
20179            self.scale_inplace(&mut y1, ws1, m * out1)?;
20180        }
20181        if ws2 != 1.0 {
20182            self.scale_inplace(&mut y2, ws2, m * out2)?;
20183        }
20184        Ok((y0, y1, y2))
20185    }
20186
20187    /// BLOCK-128 e4m3 MMVQ launcher (`qmatvec_e4m3_blk_mmvq`, lane/fp8-blk128-decode 2026-08-05).
20188    /// The per-block-dequant twin of `qmatvec_mmvq`'s QT_F8_E4M3 arm: same grid/block decomposition
20189    /// (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but
20190    /// the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel
20191    /// instead of one scalar folded at the write. It cannot share `qmatvec_mmvq`'s body because
20192    /// that launcher's arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).
20193    ///
20194    /// `mr` and `rp` have no analogue here (no split-plane e4m3 layout exists), so there is exactly
20195    /// one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.
20196    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20197    pub fn qmatvec_e4m3_blk_mmvq(
20198        &self,
20199        bytes: &CudaSlice<u8>,
20200        aq: &CudaSlice<i8>,
20201        ad: &CudaSlice<f32>,
20202        scales: &CudaSlice<f32>,
20203        m: usize,
20204        in_f: usize,
20205        out_f: usize,
20206        row_bytes: usize,
20207        scale_cols: usize,
20208    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20209        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite output: skip memset
20210        self.qmatvec_e4m3_blk_mmvq_into(
20211            bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
20212        )?;
20213        Ok(y)
20214    }
20215
20216    /// Slot-fed twin of `qmatvec_e4m3_blk_mmvq` (caller-owned output; the alloc-free capture lane).
20217    #[allow(clippy::too_many_arguments)]
20218    pub fn qmatvec_e4m3_blk_mmvq_into(
20219        &self,
20220        bytes: &CudaSlice<u8>,
20221        aq: &CudaSlice<i8>,
20222        ad: &CudaSlice<f32>,
20223        scales: &CudaSlice<f32>,
20224        m: usize,
20225        in_f: usize,
20226        out_f: usize,
20227        row_bytes: usize,
20228        scale_cols: usize,
20229        y: &mut CudaSlice<f32>,
20230    ) -> Result<(), Box<dyn std::error::Error>> {
20231        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
20232        let f = self.func("qmatvec_e4m3_blk_mmvq");
20233        let cfg = LaunchConfig {
20234            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
20235            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row
20236            shared_mem_bytes: 0,                // warp-only reduce
20237        };
20238        let (inf, outf, mi, rb, sc) = (
20239            in_f as i32,
20240            out_f as i32,
20241            m as i32,
20242            row_bytes as i64,
20243            scale_cols as i32,
20244        );
20245        let __s_b = self.gpu.stream();
20246        let mut b = __s_b.launch_builder(&f);
20247        b.arg(bytes)
20248            .arg(aq)
20249            .arg(ad)
20250            .arg(scales)
20251            .arg(&mut *y)
20252            .arg(&inf)
20253            .arg(&outf)
20254            .arg(&mi)
20255            .arg(&rb)
20256            .arg(&sc);
20257        unsafe {
20258            b.launch(cfg)?;
20259        }
20260        Ok(())
20261    }
20262
20263    /// BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
20264    /// `qmatvec_e4m3_blk_mmvq` for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
20265    /// (same fmaf chain, same per-k32 `s * ad` fold, same warp reduce), so it inherits the
20266    /// decode-exactness contract while reading the weight ONCE for up to `mcols` columns instead
20267    /// of `m` times. `mcols` must be one of {2,4,8,16} and satisfy `mcols >= m`.
20268    #[allow(clippy::too_many_arguments)]
20269    pub fn qmatvec_e4m3_blk_mmvq_batched(
20270        &self,
20271        bytes: &CudaSlice<u8>,
20272        aq: &CudaSlice<i8>,
20273        ad: &CudaSlice<f32>,
20274        scales: &CudaSlice<f32>,
20275        m: usize,
20276        in_f: usize,
20277        out_f: usize,
20278        row_bytes: usize,
20279        scale_cols: usize,
20280        mcols: usize,
20281    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20282        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
20283        debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
20284        let name = match mcols {
20285            2 => "qmatvec_e4m3_blk_mmvq_b2",
20286            4 => "qmatvec_e4m3_blk_mmvq_b4",
20287            8 => "qmatvec_e4m3_blk_mmvq_b8",
20288            16 => "qmatvec_e4m3_blk_mmvq_b16",
20289            _ => {
20290                return Err(
20291                    format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
20292                );
20293            }
20294        };
20295        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20296        let f = self.func(name);
20297        let cfg = LaunchConfig {
20298            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
20299            block_dim: (32, ROWS_PER_BLOCK, 1),
20300            shared_mem_bytes: 0,
20301        };
20302        let (inf, outf, mi, rb, sc) = (
20303            in_f as i32,
20304            out_f as i32,
20305            m as i32,
20306            row_bytes as i64,
20307            scale_cols as i32,
20308        );
20309        let __s_b = self.gpu.stream();
20310        let mut b = __s_b.launch_builder(&f);
20311        b.arg(bytes)
20312            .arg(aq)
20313            .arg(ad)
20314            .arg(scales)
20315            .arg(&mut y)
20316            .arg(&inf)
20317            .arg(&outf)
20318            .arg(&mi)
20319            .arg(&rb)
20320            .arg(&sc);
20321        unsafe {
20322            b.launch(cfg)?;
20323        }
20324        Ok(y)
20325    }
20326
20327    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
20328    /// bytes with an internal q8_1 quantize (mirrors `qmatvec_batched_raw`).
20329    #[allow(clippy::too_many_arguments)]
20330    pub fn qmatvec_e4m3_blk_batched_raw(
20331        &self,
20332        bytes: &CudaSlice<u8>,
20333        x: &CudaSlice<f32>,
20334        scales: &CudaSlice<f32>,
20335        m: usize,
20336        in_f: usize,
20337        out_f: usize,
20338        row_bytes: usize,
20339        scale_cols: usize,
20340        mcols: usize,
20341    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20342        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20343        self.qmatvec_e4m3_blk_mmvq_batched(
20344            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
20345        )
20346    }
20347
20348    /// Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with
20349    /// an internal q8_1 quantize (mirrors `qmatvec_mmvq_raw`).
20350    #[allow(clippy::too_many_arguments)]
20351    pub fn qmatvec_e4m3_blk_mmvq_raw(
20352        &self,
20353        bytes: &CudaSlice<u8>,
20354        x: &CudaSlice<f32>,
20355        scales: &CudaSlice<f32>,
20356        m: usize,
20357        in_f: usize,
20358        out_f: usize,
20359        row_bytes: usize,
20360        scale_cols: usize,
20361    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20362        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20363        self.qmatvec_e4m3_blk_mmvq(
20364            bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
20365        )
20366    }
20367
20368    /// Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight
20369    /// bytes with internal q8_1 quantize, no env gating (mirrors `qmatvec_q8_fused*_raw`).
20370    #[allow(clippy::too_many_arguments)]
20371    pub fn qmatvec_e4m3_fused2_raw(
20372        &self,
20373        b0: &CudaSlice<u8>,
20374        b1: &CudaSlice<u8>,
20375        x: &CudaSlice<f32>,
20376        in_f: usize,
20377        out0: usize,
20378        out1: usize,
20379        row_bytes: usize,
20380        ws0: f32,
20381        ws1: f32,
20382    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20383        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20384        self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
20385    }
20386
20387    #[allow(clippy::too_many_arguments)]
20388    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
20389    pub fn qmatvec_e4m3_fused3_raw(
20390        &self,
20391        b0: &CudaSlice<u8>,
20392        b1: &CudaSlice<u8>,
20393        b2: &CudaSlice<u8>,
20394        x: &CudaSlice<f32>,
20395        in_f: usize,
20396        out0: usize,
20397        out1: usize,
20398        out2: usize,
20399        row_bytes: usize,
20400        ws0: f32,
20401        ws1: f32,
20402        ws2: f32,
20403    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20404        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20405        self.e4m3_fused3_core(
20406            b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
20407        )
20408    }
20409
20410    /// Raw six-group entry (`qmatvec_e4m3_mmvq_fused6`): quantizes the activation once and
20411    /// launches all six ranges. The gate drives mutations through here so a mutated program is
20412    /// the exact one the KDA door serves.
20413    #[allow(clippy::too_many_arguments)]
20414    pub fn qmatvec_e4m3_fused6_raw(
20415        &self,
20416        w: [&CudaSlice<u8>; 6],
20417        x: &CudaSlice<f32>,
20418        in_f: usize,
20419        dims: [usize; 6],
20420        row_bytes: usize,
20421        ws: [f32; 6],
20422        arm: u32,
20423    ) -> Result<[CudaSlice<f32>; 6], Box<dyn std::error::Error>> {
20424        let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20425        let mut outs = [
20426            self.alloc_uninit::<f32>(dims[0])?,
20427            self.alloc_uninit::<f32>(dims[1])?,
20428            self.alloc_uninit::<f32>(dims[2])?,
20429            self.alloc_uninit::<f32>(dims[3])?,
20430            self.alloc_uninit::<f32>(dims[4])?,
20431            self.alloc_uninit::<f32>(dims[5])?,
20432        ];
20433        self.e4m3_fused6_into_arm(w, &aq, &ad, in_f, dims, row_bytes, ws, &mut outs, arm)?;
20434        Ok(outs)
20435    }
20436
20437    #[allow(clippy::too_many_arguments)]
20438    pub fn qmatvec_e4m3_fused2_t_raw(
20439        &self,
20440        b0: &CudaSlice<u8>,
20441        b1: &CudaSlice<u8>,
20442        x: &CudaSlice<f32>,
20443        m: usize,
20444        in_f: usize,
20445        out0: usize,
20446        out1: usize,
20447        row_bytes: usize,
20448        ws0: f32,
20449        ws1: f32,
20450    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20451        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20452        self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
20453    }
20454
20455    #[allow(clippy::too_many_arguments)]
20456    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
20457    pub fn qmatvec_e4m3_fused3_t_raw(
20458        &self,
20459        b0: &CudaSlice<u8>,
20460        b1: &CudaSlice<u8>,
20461        b2: &CudaSlice<u8>,
20462        x: &CudaSlice<f32>,
20463        m: usize,
20464        in_f: usize,
20465        out0: usize,
20466        out1: usize,
20467        out2: usize,
20468        row_bytes: usize,
20469        ws0: f32,
20470        ws1: f32,
20471        ws2: f32,
20472    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20473        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20474        self.e4m3_fused3_t_core(
20475            b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
20476        )
20477    }
20478
20479    /// THE single dispatch point for `QT_F8_E4M3_BLK` from a PRE-QUANTIZED q8_1 activation
20480    /// (lane/fp8-blk128-decode). Every `matmul_pre`-family entry calls this first, so the block-128
20481    /// class has exactly ONE code path across `matmul`, `matmul_pre`, `matmul_pre_noscale`,
20482    /// `matmul_decode_exact` and `matmul_decode_exact_pre` — the same kernel at the same grid for
20483    /// every m, which is what makes verify == decode bit-for-bit at every tier for free.
20484    ///
20485    /// Returns None for any other qtype (the caller continues its normal dispatch). The `blk: Some`
20486    /// pattern is part of the match, not an unwrap: qtype and grid presence are set together in the
20487    /// one residency arm that builds this tensor, and a qtype-without-grid would be a construction
20488    /// bug — better to fall through and hit a loud refusal than to unwrap a None here.
20489    fn try_e4m3_blk_pre(
20490        &self,
20491        w: &crate::model::GpuTensor,
20492        aq: &CudaSlice<i8>,
20493        ad: &CudaSlice<f32>,
20494        m: usize,
20495    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20496        use crate::model::GpuTensor;
20497        if let GpuTensor::Quant {
20498            bytes,
20499            qtype,
20500            row_bytes,
20501            blk: Some(g),
20502            ..
20503        } = w
20504            && *qtype == QT_F8_E4M3_BLK
20505        {
20506            // BATCHED tier m=2..16 (lane/rp-on-st): weight read ONCE for up to mcols columns
20507            // instead of m grid.y re-reads. Bit-identical per (token,row) to the grid.y=m form
20508            // below, so the decode-exactness contract is preserved at every width. Gated by
20509            // the same seams the other batched families honor (MEMRA_NO_BATCHED, MEMRA_B8) so
20510            // one rollback door covers every dtype's batched tier.
20511            if (2..=16).contains(&m)
20512                && std::env::var("MEMRA_NO_BATCHED").is_err()
20513                && (m <= 4 || Self::b8_enabled())
20514            {
20515                let mcols = Self::batched_mcols(m);
20516                return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
20517                    bytes,
20518                    aq,
20519                    ad,
20520                    &g.scales,
20521                    m,
20522                    w.in_features(),
20523                    w.out_features(),
20524                    *row_bytes,
20525                    g.cols,
20526                    mcols,
20527                )?));
20528            }
20529            return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
20530                bytes,
20531                aq,
20532                ad,
20533                &g.scales,
20534                m,
20535                w.in_features(),
20536                w.out_features(),
20537                *row_bytes,
20538                g.cols,
20539            )?));
20540        }
20541        Ok(None)
20542    }
20543
20544    /// PREFILL (m >= GEMM_M_THRESHOLD) for `QT_F8_E4M3_BLK` — DEQUANT-PER-CALL to the Q8_0 slab
20545    /// this class's residency replaced, then the ordinary Q8_0 prefill dispatch on the transient.
20546    ///
20547    /// WHY THIS EXISTS AT ALL, i.e. the regression it prevents: the decode kernel is a warp-per-row
20548    /// GEMV. At grid.y=m it re-reads the whole weight once PER TOKEN, so letting a 512-token prefill
20549    /// chunk reach it would be a ~500x weight-traffic blowup on the single most bandwidth-bound part
20550    /// of the forward. Native residency is a DECODE win and must not be paid for in prefill, so
20551    /// prefill keeps the floor's arithmetic and the floor's kernels.
20552    ///
20553    /// WHY DEQUANT-PER-CALL rather than a second resident slab: a resident slab is dual residency —
20554    /// it gives back the entire 1.0-vs-1.0625 B/weight win this lane exists to capture (and then
20555    /// some, since the e4m3 copy stays too). The transient costs one linear device pass per
20556    /// (projection, prefill call) and frees immediately.
20557    ///
20558    /// NUMERICALLY IT IS THE FLOOR, EXACTLY: `fp8_blk_dequant_q8_0` is the merged ARM B' kernel,
20559    /// gate-proven BYTE-IDENTICAL to the host dequant+re-encode (kernel-check `fp8-blk-gpu`). So the
20560    /// slab these bytes form is bit-for-bit the slab the `MEMRA_ST_E4M3_BLK=0` arm makes resident,
20561    /// and every prefill kernel downstream sees identical input — prefill logits under this lane are
20562    /// bit-identical to prefill logits under the floor, which is what makes the decode A/B a clean
20563    /// single-variable comparison instead of a two-variable one.
20564    ///
20565    /// WHAT IT COSTS, MEASURED, AND WHY THAT COST IS MOSTLY STRUCTURAL (27B block-128 ckpt, pp512,
20566    /// this rig = RTX 5090 Laptop, ~896 GB/s GDDR7). This arm makes prefill move the weight THREE
20567    /// times instead of once: read 6.88 GB of e4m3, write 7.31 GB of Q8_0, then the MMQ reads that
20568    /// 7.31 GB back. The two extra passes are 14.19 GB = 15.8 ms at this card's roofline against a
20569    /// ~332 ms pp512, i.e. **~-4.5% pp is a floor no kernel tuning can remove** — only deleting the
20570    /// dequant can. Measured: the dequant kernel costs 27.9 ms/pass (nsys, 208 projections) after
20571    /// the 2026-08-05 vector rewrite (was 66.5 ms at one byte per thread), and e2e pp512 is
20572    /// 1451.4 vs the slab arm's 1541.6 tok/s = -5.8% (N=3 interleaved pairs). So ~1.3pp of the
20573    /// -5.8% is residual kernel inefficiency and ~4.5pp is the extra traffic itself.
20574    ///
20575    /// SO THE DEQUANT IS NO LONGER THE DEFAULT ROUTE — it is the FALLBACK. The per-block FP8 MMQ
20576    /// tile (`try_fp8_blk_mmq`) consumes the resident e4m3 bytes + grid DIRECTLY, deleting both extra
20577    /// passes, and since 2026-08-05 it runs FIRST and by default for the native-resident source
20578    /// (`fp8_blk_mmq_native_enabled`; `MEMRA_FP8_MMQ=0` is the seam back to this dequant). On paper
20579    /// the trade was unassumable — lane/fp8-mmq-v2 measured that tile at 0.85-1.09x the Q8_0 MMQ
20580    /// floor GEMM-only, so it swapped a -4.5% traffic cost for a 0-to-15% GEMM cost of unknown sign.
20581    /// Measured on the 27B (3 arms interleaved, N=3, research/fp8blk-20260805/VERDICT.md): slab
20582    /// 1540.5 / this dequant 1449.1 / the tile 1553.3 tok/s, min(tile) > max(slab). The tile wins
20583    /// because v2's denominator had its slab already resident while this class's floor must build it
20584    /// every call; same tile, opposite sign, because the question changed.
20585    ///
20586    /// THIS ARM STILL RUNS, and is not dead code: every `try_fp8_blk_mmq` precondition (in_f % 16,
20587    /// grid dims vs shape, per-tensor scale == 1.0, the e4m3-NaN scan) refuses by falling through to
20588    /// here, so a checkpoint the tile cannot take keeps exact prefill on the floor's own bits rather
20589    /// than losing the class. It is also what `MEMRA_FP8_MMQ=0` reverts to.
20590    fn try_e4m3_blk_prefill(
20591        &self,
20592        w: &crate::model::GpuTensor,
20593        x: &CudaSlice<f32>,
20594        m: usize,
20595    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20596        use crate::model::GpuTensor;
20597        let GpuTensor::Quant {
20598            bytes,
20599            qtype,
20600            blk: Some(g),
20601            ..
20602        } = w
20603        else {
20604            return Ok(None);
20605        };
20606        if *qtype != QT_F8_E4M3_BLK {
20607            return Ok(None);
20608        }
20609        // NO-DEQUANT ROUTE, THE DEFAULT (MEMRA_FP8_MMQ=0 reverts): the per-block MMQ tile eats the
20610        // resident e4m3 bytes and grid as-is, so neither extra weight pass happens. Its own
20611        // preconditions (in_f % 16, grid dims, scale == 1.0, no e4m3 NaN code) can refuse — fall
20612        // through to the dequant below when they do, never silently produce nothing.
20613        if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
20614            return Ok(Some(y));
20615        }
20616        let (in_f, out_f) = (w.in_features(), w.out_features());
20617        let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
20618        let tmp = GpuTensor::Quant {
20619            bytes: slab,
20620            qtype: QT_Q8_0,
20621            row_bytes: in_f / 32 * 34,
20622            ne: vec![in_f as u64, out_f as u64],
20623            scale: 1.0,
20624            rp: false,
20625            #[cfg(memra_cutlass)]
20626            cutlass: None,
20627            fp8: None,
20628            blk: None,
20629            f16: None,
20630            rp4: None,
20631        };
20632        // Recursion terminates: `tmp` is QT_Q8_0 with `blk: None`, so it cannot re-enter this arm.
20633        Ok(Some(self.matmul(&tmp, x, m)?))
20634    }
20635
20636    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
20637    pub fn matmul_pre_noscale(
20638        &self,
20639        w: &crate::model::GpuTensor,
20640        aq: &CudaSlice<i8>,
20641        ad: &CudaSlice<f32>,
20642        m: usize,
20643    ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
20644        use crate::model::GpuTensor;
20645        // BLOCK-128 e4m3: every scale factor is folded inside the kernel per k128, so the
20646        // "separable post-op scale" this entry exists to defer is 1.0 — return it explicitly
20647        // rather than let the tail below refuse and cost the caller a re-dispatch.
20648        if m == 1
20649            && let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)?
20650        {
20651            return Ok(Some((y, 1.0)));
20652        }
20653        // Only the m==1 fast path applies the scale as a separable post-op; bail everywhere else.
20654        if m != 1 || !self.uses_q8_1_fast(w) {
20655            return Ok(None);
20656        }
20657        let in_f = w.in_features();
20658        let out_f = w.out_features();
20659        let (bytes, qtype, row_bytes, scale, rp) = match w {
20660            GpuTensor::Quant {
20661                bytes,
20662                qtype,
20663                row_bytes,
20664                scale,
20665                rp,
20666                ..
20667            } => (bytes, *qtype, *row_bytes, *scale, *rp),
20668            _ => return Ok(None),
20669        };
20670        // MMVQ warp-per-row (scale==1.0 passed -> kernel skips its internal scale; we return scale).
20671        if self.mmvq_supports(qtype) {
20672            // Q4_0 split-plane mirror (dp4a fallback below keeps the raw GGUF bytes).
20673            let (mbytes, mrp) = match w {
20674                GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
20675                _ => (bytes, rp),
20676            };
20677            let y = self.qmatvec_mmvq(
20678                mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, /*scale*/ 1.0, mrp,
20679            )?;
20680            return Ok(Some((y, scale)));
20681        }
20682        // dp4a fallback: same launch as matmul_pre but WITHOUT the post scale_inplace.
20683        let name = match qtype {
20684            QT_Q8_0 => "qmatvec_q8_0_dp4a",
20685            QT_Q4_K => "qmatvec_q4_K_dp4a",
20686            QT_Q6_K => "qmatvec_q6_K_dp4a",
20687            QT_Q5_K => "qmatvec_q5_K_dp4a",
20688            QT_Q3_K => "qmatvec_q3_K_dp4a",
20689            QT_NVFP4 => {
20690                if rp {
20691                    "qmatvec_nvfp4_dp4a_rp"
20692                } else {
20693                    "qmatvec_nvfp4_dp4a"
20694                }
20695            }
20696            QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
20697            _ => return Ok(None),
20698        };
20699        let f = self.func(name);
20700        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20701        let cfg = LaunchConfig {
20702            grid_dim: (out_f as u32, m as u32, 1),
20703            block_dim: (128, 1, 1),
20704            shared_mem_bytes: 0,
20705        };
20706        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20707        let __s_b = self.gpu.stream();
20708        let mut b = __s_b.launch_builder(&f);
20709        b.arg(bytes)
20710            .arg(aq)
20711            .arg(ad)
20712            .arg(&mut y)
20713            .arg(&inf)
20714            .arg(&outf)
20715            .arg(&mi)
20716            .arg(&rb);
20717        unsafe {
20718            b.launch(cfg)?;
20719        }
20720        Ok(Some((y, scale)))
20721    }
20722
20723    /// True if `qtype` has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4
20724    /// daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).
20725    pub fn mmvq_supports(&self, qtype: i32) -> bool {
20726        // DEFAULT ON since 2026-07-08 (MEMRA_MMVQ=0 reverts to the _dp4a matvec class).
20727        // QT_F8_E4M3 is exempt from the MEMRA_MMVQ=0 escape: the e4m3 mmvq family is that dtype's
20728        // ONLY int8-act kernel class (there is no _dp4a twin), so its m=1/verify/batched dispatch
20729        // is a pure function of the dtype — the decode-parity law holds under every env.
20730        if qtype == QT_F8_E4M3 {
20731            return true;
20732        }
20733        if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
20734            return false;
20735        }
20736        matches!(
20737            qtype,
20738            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
20739        )
20740    }
20741
20742    /// PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1):
20743    /// one warp owns one output row, warp-only __shfl reduction (no smem barrier). Bit-equivalent
20744    /// to qmatvec_*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4
20745    /// per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).
20746    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20747    pub fn qmatvec_mmvq(
20748        &self,
20749        bytes: &CudaSlice<u8>,
20750        aq: &CudaSlice<i8>,
20751        ad: &CudaSlice<f32>,
20752        m: usize,
20753        in_f: usize,
20754        out_f: usize,
20755        qtype: i32,
20756        row_bytes: usize,
20757        scale: f32,
20758        rp: bool,
20759    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20760        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
20761        self.qmatvec_mmvq_into(
20762            bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
20763        )?;
20764        Ok(y)
20765    }
20766
20767    /// Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.
20768    #[allow(clippy::too_many_arguments)]
20769    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
20770    pub fn qmatvec_mmvq_into(
20771        &self,
20772        bytes: &CudaSlice<u8>,
20773        aq: &CudaSlice<i8>,
20774        ad: &CudaSlice<f32>,
20775        m: usize,
20776        in_f: usize,
20777        out_f: usize,
20778        qtype: i32,
20779        row_bytes: usize,
20780        scale: f32,
20781        rp: bool,
20782        y: &mut CudaSlice<f32>,
20783    ) -> Result<(), Box<dyn std::error::Error>> {
20784        debug_assert!(y.len() >= m * out_f);
20785        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
20786        // SMALL-SHAPE GRID FILL (H100 lane, 2026-07-26 microbench: attn qkv out_f=2048 =
20787        // 0.97 waves at the 4-warp block -> 66% of peak). The g2 twin (2 warps/block)
20788        // doubles the grid when the 4-warp launch would be sub-wave; per-row program
20789        // identical -> bit-identical. MEMRA_Q80_G2=0 reverts.
20790        if qtype == QT_Q8_0
20791            && rp
20792            && m == 1
20793            && out_f >= 64
20794            && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
20795            && {
20796                static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20797                *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
20798            }
20799        {
20800            let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
20801            let cfg = LaunchConfig {
20802                grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
20803                block_dim: (32, 2, 1),
20804                shared_mem_bytes: 0,
20805            };
20806            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
20807            let __s_b = self.gpu.stream();
20808            let mut b = __s_b.launch_builder(&f);
20809            b.arg(bytes)
20810                .arg(aq)
20811                .arg(ad)
20812                .arg(&mut *y)
20813                .arg(&inf)
20814                .arg(&outf)
20815                .arg(&mi)
20816                .arg(&rb);
20817            unsafe {
20818                b.launch(cfg)?;
20819            }
20820            if scale != 1.0 {
20821                self.scale_inplace(y, scale, out_f)?;
20822            }
20823            return Ok(());
20824        }
20825        // Multi-row-per-warp (mr2) policy, fixed since the 2026-07 sweeps (the MEMRA_MMVQ_MR
20826        // override + mr4 kernel were retired 2026-07-08 — mr4 regressed on register pressure and
20827        // crashed under rp; q4_K/q6_K mr2 measured flat, "no gain = no change"):
20828        //   NVFP4 m=1 -> mr2 (clean +1-2% on 9B: RPW acc chains hide the weight-load latency
20829        //     that pins the single-row kernel at 30-46% DRAM). Bit-identical per row.
20830        //   Q5_K m=1 -> mr2 (2026-07-05: the FR-Spec trimmed draft head is Q5_K 32768 rows = 8%
20831        //     of the 27B p3 spec wall; latency-bound like the other k-quants pre-fix).
20832        //   Q4_K/Q6_K m=1 -> single-row (mr2 measured +0.7% / flat — weight-bandwidth-bound).
20833        let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
20834            2
20835        } else {
20836            1
20837        };
20838        // Q4_0 mr (gemma trunk): DEFAULT 1 since 2026-07-13 (MEMRA_Q40_MR=2 reverts) — the
20839        // mr1 rp twin doubles the block count and wins the tail-quantization/latency battle
20840        // on every gemma model (E4B +3.75%: 198.9 vs 191.7; 26B +0.7%; 31B +0.9%; N=2-3
20841        // valid-window interleaved, bit-identical per row — same dot program).
20842        if m == 1 && qtype == QT_Q4_0 {
20843            static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
20844            // shape policy PROBED NEGATIVE (2026-07-13): tall-only mr1 197.2 vs
20845            // mr1-everywhere 198.7 — mr1 wins wide-output shapes too; arm removed.
20846            mr = *Q40MR.get_or_init(|| {
20847                std::env::var("MEMRA_Q40_MR")
20848                    .ok()
20849                    .and_then(|v| v.parse().ok())
20850                    .unwrap_or(1)
20851            });
20852        }
20853        // q5issue lane (2026-07-08): MEMRA_Q5K_ISSUE swaps the q5_K m=1 mmvq kernels for the
20854        // issue-reduced `_il` bodies (uint4 header/qh/qs loads + branchless scale decode —
20855        // cuts ~34 LDG.U16 + ~5 LDG.U8 + a warp-divergent scale branch per 32-elem group-row
20856        // to 5 LDG.128). Bit-identical per (token,row) to the reference kernels.
20857        // `1` = shape-aware policy (N=3 clock-locked micro-bench, mem P0, synthetic real shapes):
20858        //   out_f <= 65536 (trunk/frspec regime): il at the default mr — mr2_il -9.5%/-10.5%
20859        //     on 4096x4096/4096x8192, -3.1% on the 32768 frspec head vs the mr2-ref default;
20860        //   out_f > 65536 (the 248320-row 27B lm_head, already ~97% of the mem wall): mr2_il
20861        //     REGRESSES +22% there but mr1_il wins -2.1% vs the mr2-ref default -> force mr=1.
20862        // `2` = force il at the current mr for EVERY shape (A/B probe seam). Default OFF.
20863        let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
20864        let q5_force = q5_mode.as_deref() == Some("2");
20865        // DEFAULT ON since 2026-07-08 (MEMRA_Q5K_ISSUE=0 reverts): +1.8% 9B plain e2e N=3
20866        // (128.2 -> 130.4), 27B flat (its big head is already at the mem wall), all gates green.
20867        let q5_il = qtype == QT_Q5_K
20868            && m == 1
20869            && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
20870        if q5_il && !q5_force && out_f > 65536 {
20871            mr = 1;
20872        }
20873        // Q4_0 split-plane rp: mr2 default; MEMRA_Q40_MR=1 reaches the mr1 rp twin
20874        // (2026-07-13 — the tall-input/short-output tail-quantization probe).
20875        if qtype == QT_Q4_0 && rp && mr != 1 {
20876            mr = 2;
20877        }
20878        // Q8_0 rp (H100 lane): mr1 default — the q4_0 mr2 recipe MEASURED NEGATIVE on H100
20879        // (2026-07-26 N=3: mr1 186.2 vs mr2 171.5 tok/s; halving the grid on 132 SMs costs
20880        // more than 2-row ILP buys). mr2 kernel stays behind MEMRA_Q80_MR=2 for the corpus.
20881        if qtype == QT_Q8_0 && rp {
20882            static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
20883            mr = *Q80MR.get_or_init(|| {
20884                std::env::var("MEMRA_Q80_MR")
20885                    .ok()
20886                    .and_then(|v| v.parse().ok())
20887                    .unwrap_or(1)
20888            });
20889        }
20890        // B200 sub-wave grid-fill (MEMRA_B200_MATVEC_ARM occupancy arm, lane/b200-matvec-
20891        // occupancy-20260902): the mr2_rp/RPW=2 default halves the grid vs mr1, which the
20892        // decode-kernel census (2026-09-02) found under-filling B200's 148 SMs for the
20893        // NVFP4 m=1 decode shapes (qmatvec_nvfp4_mmvq_mr2_rp: 5.2% of GPU time, 11.5us avg
20894        // over 30,400 launches). Forcing mr=1 here reuses the ALREADY-SHIPPED
20895        // `qmatvec_nvfp4_mmvq_rp` kernel and doubles the grid for the same output rows —
20896        // exactly the Q8_0 g2 "SMALL-SHAPE GRID FILL" recipe above, mapped onto NVFP4. Per
20897        // row the seg body is IDENTICAL between mr1 and mr2 (same dequant/dp4a/reduce
20898        // chain), so this changes zero output bits, only which warp computes which row.
20899        // Gated on sm_100a builds only (`b200_matvec_arm_on`); sm_120a keeps its measured
20900        // mr2 default unconditionally. Default OFF pending the B200 A/B (docs/FLAGS.md).
20901        if qtype == QT_NVFP4
20902            && mr == 2
20903            && rp
20904            && m == 1
20905            && b200_matvec_arm_on()
20906            && (out_f as u32).div_ceil(ROWS_PER_BLOCK * 2)
20907                < b200_mr1_fill() * self.sm_count() as u32
20908        {
20909            mr = 1;
20910        }
20911        // MEMRA_NVFP4_ROW_ILP (lane/glm5-nvfp4-row-ilp-20260904, default OFF): the `_ilp` twins
20912        // of the two split-plane NVFP4 trunk kernels, same grid, same per-row program.
20913        let nv_ilp = qtype == QT_NVFP4 && rp && nvfp4_row_ilp_on();
20914        if nv_ilp
20915            && NVFP4_ROW_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0
20916        {
20917            eprintln!(
20918                "[nvfp4-row-ilp] engaged: split-plane NVFP4 trunk matvec with four groups' loads \
20919                 per lane ahead of the dp4a chains (MEMRA_NVFP4_ROW_ILP=1, mr={mr})"
20920            );
20921        }
20922        let name = match (qtype, mr, rp) {
20923            (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
20924            (QT_NVFP4, 2, true) if nv_ilp => "qmatvec_nvfp4_mmvq_mr2_rp_ilp",
20925            (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
20926            (QT_NVFP4, _, true) if nv_ilp => "qmatvec_nvfp4_mmvq_rp_ilp",
20927            (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
20928            (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
20929            (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
20930            (QT_Q5_K, 2, _) => {
20931                if q5_il {
20932                    "qmatvec_q5_K_mmvq_mr2_il"
20933                } else {
20934                    "qmatvec_q5_K_mmvq_mr2"
20935                }
20936            }
20937            (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
20938            // rpca (cp.async-staged weight ring): MEASURED NEGATIVE on H100 for Q8_0
20939            // (2026-07-26 N=3: 181.8 vs plain rp 185.5 — the smem round-trip exceeds the
20940            // latency it hides for 8-bit direct-dp4a; the NVFP4 win case overlaps table
20941            // decode with half the bytes). OPT-IN via MEMRA_Q80_CA=1 for the corpus.
20942            (QT_Q8_0, _, true)
20943                if in_f.is_multiple_of(1024) && {
20944                    static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20945                    *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
20946                } =>
20947            {
20948                "qmatvec_q8_0_mmvq_rpca"
20949            }
20950            (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
20951            (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
20952            // K-quant split-plane twins (H100 K-quant coalescing fix, 2026-08-01): the rp4
20953            // mirror routes here; GGUF layout keeps the plain kernels. rp bytes MUST never
20954            // reach a GGUF-layout kernel or vice versa.
20955            (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
20956            (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
20957            (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
20958            (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
20959            (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
20960            (QT_Q5_K, _, _) => {
20961                if q5_il {
20962                    "qmatvec_q5_K_mmvq_il"
20963                } else {
20964                    "qmatvec_q5_K_mmvq"
20965                }
20966            }
20967            (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
20968            (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
20969            (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
20970            _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
20971        };
20972        let f = self.func(name);
20973        // each block still has ROWS_PER_BLOCK warps; with mr rows/warp it covers ROWS_PER_BLOCK*mr rows.
20974        let rows_per_block = ROWS_PER_BLOCK * mr;
20975        let cfg = LaunchConfig {
20976            grid_dim: (
20977                (out_f as u32 + rows_per_block - 1) / rows_per_block,
20978                m as u32,
20979                1,
20980            ),
20981            block_dim: (32, ROWS_PER_BLOCK, 1), // warp-per-row (x mr rows each)
20982            shared_mem_bytes: 0,                // warp-only reduce at m=1
20983        };
20984        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20985        let __s_b = self.gpu.stream();
20986        let mut b = __s_b.launch_builder(&f);
20987        // NVFP4 + e4m3 mmvq kernels take the macro-scale as a fused epilogue arg (applied at the
20988        // write — bit-identical to the old separate scale_inplace pass, minus one launch per matvec:
20989        // 53 scale launches/token on the 9B; for e4m3 the scale is the checkpoint's per-tensor f32
20990        // weight_scale). Other mmvq kernels keep the 8-arg signature.
20991        if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
20992            // PDL wave-B: the nvfp4 mr2_rp single (gemma wo / generic rp singles) joins
20993            // the wave-A launch class — 9-arg flavor (fused macro-scale epilogue).
20994            // The four split-plane NVFP4 trunk kernels all carry MEMRA_PDL_ENTRY (the mr1
20995            // kernel since lane/glm5-nvfp4-row-ilp-20260904), so the grid-fill and ILP twins
20996            // ride the same PDL launch class as the shipped mr2 kernel.
20997            if Self::pdl_on()
20998                && Self::pdl_mmvq_on()
20999                && Self::pdl_nvfp4q8_on()
21000                && matches!(
21001                    name,
21002                    "qmatvec_nvfp4_mmvq_mr2_rp"
21003                        | "qmatvec_nvfp4_mmvq_mr2_rp_ilp"
21004                        | "qmatvec_nvfp4_mmvq_rp"
21005                        | "qmatvec_nvfp4_mmvq_rp_ilp"
21006                )
21007            {
21008                use cudarc::driver::{DevicePtr, DevicePtrMut};
21009                let s = &self.gpu.stream();
21010                let (pw, _g0) = bytes.device_ptr(s);
21011                let (paq, _g1) = aq.device_ptr(s);
21012                let (pad, _g2) = ad.device_ptr(s);
21013                let (py, _g3) = y.device_ptr_mut(s);
21014                let mut ps = [
21015                    &pw as *const _ as *mut std::ffi::c_void,
21016                    &paq as *const _ as *mut _,
21017                    &pad as *const _ as *mut _,
21018                    &py as *const _ as *mut _,
21019                    &inf as *const _ as *mut _,
21020                    &outf as *const _ as *mut _,
21021                    &mi as *const _ as *mut _,
21022                    &rb as *const _ as *mut _,
21023                    &scale as *const _ as *mut _,
21024                ];
21025                unsafe {
21026                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
21027                }
21028                return Ok(());
21029            }
21030            b.arg(bytes)
21031                .arg(aq)
21032                .arg(ad)
21033                .arg(&mut *y)
21034                .arg(&inf)
21035                .arg(&outf)
21036                .arg(&mi)
21037                .arg(&rb)
21038                .arg(&scale);
21039            unsafe {
21040                b.launch(cfg)?;
21041            }
21042        } else if Self::pdl_on()
21043            && Self::pdl_mmvq_on()
21044            && (matches!(
21045                name,
21046                "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
21047            ) || (Self::pdl_nvfp4q8_on()
21048                && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
21049        {
21050            // PDL wave-A (2026-07-23): the two decode-hot single-matvec kernels carry
21051            // MEMRA_PDL_ENTRY — grid launches while the producer drains. ONLY the marked
21052            // names may take this launch (unmarked kernels would read unordered).
21053            {
21054                use cudarc::driver::{DevicePtr, DevicePtrMut};
21055                let s = &self.gpu.stream();
21056                let (pw, _g0) = bytes.device_ptr(s);
21057                let (paq, _g1) = aq.device_ptr(s);
21058                let (pad, _g2) = ad.device_ptr(s);
21059                let (py, _g3) = y.device_ptr_mut(s);
21060                let mut ps = [
21061                    &pw as *const _ as *mut std::ffi::c_void,
21062                    &paq as *const _ as *mut _,
21063                    &pad as *const _ as *mut _,
21064                    &py as *const _ as *mut _,
21065                    &inf as *const _ as *mut _,
21066                    &outf as *const _ as *mut _,
21067                    &mi as *const _ as *mut _,
21068                    &rb as *const _ as *mut _,
21069                ];
21070                unsafe {
21071                    self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
21072                }
21073            }
21074            if scale != 1.0 {
21075                self.scale_inplace(y, scale, m * out_f)?;
21076            }
21077        } else {
21078            b.arg(bytes)
21079                .arg(aq)
21080                .arg(ad)
21081                .arg(&mut *y)
21082                .arg(&inf)
21083                .arg(&outf)
21084                .arg(&mi)
21085                .arg(&rb);
21086            unsafe {
21087                b.launch(cfg)?;
21088            }
21089            if scale != 1.0 {
21090                self.scale_inplace(y, scale, m * out_f)?;
21091            }
21092        }
21093        Ok(())
21094    }
21095
21096    /// Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly
21097    /// from raw weight bytes (quantize the f32 activation `x` to q8_1 internally). NVFP4 per-tensor
21098    /// macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.
21099    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
21100    pub fn qmatvec_mmvq_raw(
21101        &self,
21102        bytes: &CudaSlice<u8>,
21103        x: &CudaSlice<f32>,
21104        m: usize,
21105        in_f: usize,
21106        out_f: usize,
21107        qtype: i32,
21108        row_bytes: usize,
21109        rp: bool,
21110    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21111        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
21112        self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
21113    }
21114
21115    /// Bench-only direct arm selector for the NVFP4 split-plane rp m=1 decode pair
21116    /// (b200_matvec_bench.rs). Bypasses `qmatvec_mmvq_into`'s policy AND the
21117    /// `MEMRA_B200_MATVEC_ARM` env door (which is read once into a process-wide `OnceLock` and
21118    /// so cannot flip mid-process for an in-process A/B) — `use_arm=false` launches
21119    /// `qmatvec_nvfp4_mmvq_mr2_rp` (shipped default, RPW=2, half the grid); `true` launches the
21120    /// already-shipped `qmatvec_nvfp4_mmvq_rp` (RPW=1, full grid) that the B200 grid-fill arm
21121    /// dispatches to. Per-row arithmetic is IDENTICAL between the two (see qmatvec.cu) — only
21122    /// the row/warp mapping and the resulting grid size differ.
21123    #[allow(clippy::too_many_arguments)]
21124    pub fn qmatvec_nvfp4_rp_arm_raw(
21125        &self,
21126        bytes: &CudaSlice<u8>,
21127        aq: &CudaSlice<i8>,
21128        ad: &CudaSlice<f32>,
21129        m: usize,
21130        in_f: usize,
21131        out_f: usize,
21132        row_bytes: usize,
21133        yscale: f32,
21134        use_arm: bool,
21135    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21136        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
21137        let mr: u32 = if use_arm { 1 } else { 2 };
21138        let name = if use_arm {
21139            "qmatvec_nvfp4_mmvq_rp"
21140        } else {
21141            "qmatvec_nvfp4_mmvq_mr2_rp"
21142        };
21143        let f = self.func(name);
21144        let rows_per_block = ROWS_PER_BLOCK * mr;
21145        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
21146        let cfg = LaunchConfig {
21147            grid_dim: ((out_f as u32).div_ceil(rows_per_block), m as u32, 1),
21148            block_dim: (32, ROWS_PER_BLOCK, 1),
21149            shared_mem_bytes: 0,
21150        };
21151        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
21152        let __s_b = self.gpu.stream();
21153        let mut b = __s_b.launch_builder(&f);
21154        b.arg(bytes)
21155            .arg(aq)
21156            .arg(ad)
21157            .arg(&mut y)
21158            .arg(&inf)
21159            .arg(&outf)
21160            .arg(&mi)
21161            .arg(&rb)
21162            .arg(&yscale);
21163        unsafe {
21164            b.launch(cfg)?;
21165        }
21166        Ok(y)
21167    }
21168
21169    /// Bench-only arm selector with the `MEMRA_NVFP4_ROW_ILP` twin as a third axis: `ilp`
21170    /// picks `_ilp` for whichever of mr1/mr2 `use_arm` chose.
21171    #[allow(clippy::too_many_arguments)]
21172    pub fn qmatvec_nvfp4_rp_arm_raw_ilp(
21173        &self,
21174        bytes: &CudaSlice<u8>,
21175        aq: &CudaSlice<i8>,
21176        ad: &CudaSlice<f32>,
21177        m: usize,
21178        in_f: usize,
21179        out_f: usize,
21180        row_bytes: usize,
21181        yscale: f32,
21182        use_arm: bool,
21183        ilp: bool,
21184    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21185        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
21186        let mr: u32 = if use_arm { 1 } else { 2 };
21187        let name = match (use_arm, ilp) {
21188            (true, false) => "qmatvec_nvfp4_mmvq_rp",
21189            (true, true) => "qmatvec_nvfp4_mmvq_rp_ilp",
21190            (false, false) => "qmatvec_nvfp4_mmvq_mr2_rp",
21191            (false, true) => "qmatvec_nvfp4_mmvq_mr2_rp_ilp",
21192        };
21193        let f = self.func(name);
21194        let rows_per_block = ROWS_PER_BLOCK * mr;
21195        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
21196        let cfg = LaunchConfig {
21197            grid_dim: ((out_f as u32).div_ceil(rows_per_block), m as u32, 1),
21198            block_dim: (32, ROWS_PER_BLOCK, 1),
21199            shared_mem_bytes: 0,
21200        };
21201        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
21202        let __s_b = self.gpu.stream();
21203        let mut b = __s_b.launch_builder(&f);
21204        b.arg(bytes)
21205            .arg(aq)
21206            .arg(ad)
21207            .arg(&mut y)
21208            .arg(&inf)
21209            .arg(&outf)
21210            .arg(&mi)
21211            .arg(&rb)
21212            .arg(&yscale);
21213        unsafe {
21214            b.launch(cfg)?;
21215        }
21216        Ok(y)
21217    }
21218
21219    /// Bench-only direct arm selector for `qmatvec_nvfp4_mmvq_fused2_rp` (b200_matvec_bench.rs) —
21220    /// same rationale as `qmatvec_nvfp4_rp_arm_raw`. `use_arm=false` launches the shipped
21221    /// RPW=2 kernel; `true` launches the RPW=1 `_g2` twin (instantiates the same
21222    /// `nvfp4_mmvq_fused_seg_rp` template at RPW=1). Per (tensor,row) bit-identical.
21223    #[allow(clippy::too_many_arguments)]
21224    pub fn qmatvec_nvfp4_fused2_rp_arm_raw(
21225        &self,
21226        w0: &CudaSlice<u8>,
21227        w1: &CudaSlice<u8>,
21228        aq: &CudaSlice<i8>,
21229        ad: &CudaSlice<f32>,
21230        in_f: usize,
21231        out0: usize,
21232        out1: usize,
21233        s0: f32,
21234        s1: f32,
21235        use_arm: bool,
21236    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21237        const ROWS_PER_BLOCK: u32 = 4; // MEMRA_MMVQ_ROWS
21238        let rpw: u32 = if use_arm { 1 } else { 2 };
21239        let name = if use_arm {
21240            "qmatvec_nvfp4_mmvq_fused2_rp_g2"
21241        } else {
21242            "qmatvec_nvfp4_mmvq_fused2_rp"
21243        };
21244        let rows_pb = ROWS_PER_BLOCK * rpw;
21245        let nb = |o: usize| (o as u32).div_ceil(rows_pb);
21246        let f = self.func(name);
21247        let mut y0 = self.alloc_uninit::<f32>(out0)?;
21248        let mut y1 = self.alloc_uninit::<f32>(out1)?;
21249        let cfg = LaunchConfig {
21250            grid_dim: (nb(out0) + nb(out1), 1, 1),
21251            block_dim: (32, ROWS_PER_BLOCK, 1),
21252            shared_mem_bytes: 0,
21253        };
21254        let (inf, oi0, oi1, mi) = (in_f as i32, out0 as i32, out1 as i32, 1i32);
21255        let __s_b = self.gpu.stream();
21256        let mut b = __s_b.launch_builder(&f);
21257        b.arg(w0)
21258            .arg(w1)
21259            .arg(aq)
21260            .arg(ad)
21261            .arg(&mut y0)
21262            .arg(&mut y1)
21263            .arg(&inf)
21264            .arg(&oi0)
21265            .arg(&oi1)
21266            .arg(&mi)
21267            .arg(&s0)
21268            .arg(&s1);
21269        unsafe {
21270            b.launch(cfg)?;
21271        }
21272        Ok((y0, y1))
21273    }
21274
21275    /// Bench-only direct arm selector for `matvec_bf16_f32acc_x4_rows` (b200_matvec_bench.rs) —
21276    /// same rationale as `qmatvec_nvfp4_rp_arm_raw`: bypasses `matvec_bf16_rows_into`'s policy
21277    /// stack and the memoized `MEMRA_B200_MATVEC_ARM` door. `use_arm=false` launches the shipped
21278    /// kernel; `true` launches the software-pipelined `_pf` twin. Bit-identical per (row,token).
21279    #[allow(clippy::too_many_arguments)]
21280    pub fn matvec_bf16_f32acc_x4_rows_arm_raw(
21281        &self,
21282        w: &CudaSlice<u8>,
21283        x: &CudaSlice<f32>,
21284        y: &mut CudaSlice<f32>,
21285        in_f: usize,
21286        out_f: usize,
21287        t: usize,
21288        use_arm: bool,
21289    ) -> Result<(), Box<dyn std::error::Error>> {
21290        let name = if use_arm {
21291            "matvec_bf16_f32acc_x4_rows_pf"
21292        } else {
21293            "matvec_bf16_f32acc_x4_rows"
21294        };
21295        let f = self.func(name);
21296        let cfg = LaunchConfig {
21297            grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
21298            block_dim: (mmv_block(), 1, 1),
21299            shared_mem_bytes: 0,
21300        };
21301        let (ini, outi) = (in_f as i32, out_f as i32);
21302        let __s_b = self.gpu.stream();
21303        let mut b = __s_b.launch_builder(&f);
21304        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21305        unsafe {
21306            b.launch(cfg)?;
21307        }
21308        Ok(())
21309    }
21310
21311    /// Split-K degree for the v2 bf16 GEMV at this shape, on this device. Returns 1 (the
21312    /// BIT-IDENTICAL single-pass kernel) whenever the row grid already covers two waves of CTAs
21313    /// over `sm_count()`; otherwise the smallest split that does, capped so every K chunk still
21314    /// gives each thread at least one full 8-element step. A returned value > 1 selects the
21315    /// NAMED numeric class `bf16_gemv_v2_splitk`.
21316    pub fn gemv_v2_ksplit(&self, in_f: usize, out_f: usize, t: usize) -> usize {
21317        let blocks = out_f.div_ceil(GEMV_V2_ROWS) * t.max(1);
21318        let want = 2 * self.sm_count().max(1) as usize;
21319        if blocks >= want || blocks == 0 {
21320            return 1;
21321        }
21322        let per_step = mmv_block() as usize * 8;
21323        let max_ks = (in_f / per_step).max(1);
21324        want.div_ceil(blocks).clamp(1, max_ks)
21325    }
21326
21327    /// The v2 bf16 GEMV launch (`matvec_bf16_v2`, or the split-K pair when `ksplit > 1`).
21328    /// POLICY-FREE: `ksplit` is the caller's choice so a gate or bench can drive both classes
21329    /// through one entry. `block_dim` is `mmv_block()` — the same blockDim the shipped kernel
21330    /// pins, because the reduction tree's shape (and therefore its bits) is a function of it.
21331    #[allow(clippy::too_many_arguments)]
21332    pub fn matvec_bf16_v2_raw(
21333        &self,
21334        w: &CudaSlice<u8>,
21335        x: &CudaSlice<f32>,
21336        y: &mut CudaSlice<f32>,
21337        in_f: usize,
21338        out_f: usize,
21339        t: usize,
21340        ksplit: usize,
21341    ) -> Result<(), Box<dyn std::error::Error>> {
21342        if t == 0 || !in_f.is_multiple_of(8) || x.len() < t * in_f || y.len() < t * out_f {
21343            return Err("matvec_bf16_v2 geometry".into());
21344        }
21345        let nb = mmv_block();
21346        let smem = (GEMV_V2_ROWS as u32) * nb * 4;
21347        let rows = out_f.div_ceil(GEMV_V2_ROWS) as u32;
21348        let (ini, outi) = (in_f as i32, out_f as i32);
21349        if ksplit <= 1 {
21350            let f = self.func("matvec_bf16_v2");
21351            let cfg = LaunchConfig {
21352                grid_dim: (rows, t as u32, 1),
21353                block_dim: (nb, 1, 1),
21354                shared_mem_bytes: smem,
21355            };
21356            let __s_b = self.gpu.stream();
21357            let mut b = __s_b.launch_builder(&f);
21358            b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21359            unsafe {
21360                b.launch(cfg)?;
21361            }
21362            return Ok(());
21363        }
21364        // Chunks are a multiple of 8 elements so every thread's 16 B loads stay aligned. The
21365        // EFFECTIVE split is recomputed from the chunk size: a requested ksplit whose last
21366        // plane would start past in_f must not be launched, or the combine would sum a plane
21367        // the kernel never wrote.
21368        let chunk = in_f.div_ceil(8).div_ceil(ksplit) * 8;
21369        let eff = in_f.div_ceil(chunk).max(1);
21370        let mut part = self.alloc_uninit::<f32>(eff * t * out_f)?;
21371        let ks = eff as i32;
21372        {
21373            let f = self.func("matvec_bf16_v2_sk");
21374            let cfg = LaunchConfig {
21375                grid_dim: (rows, t as u32, eff as u32),
21376                block_dim: (nb, 1, 1),
21377                shared_mem_bytes: smem,
21378            };
21379            let __s_b = self.gpu.stream();
21380            let mut b = __s_b.launch_builder(&f);
21381            b.arg(w).arg(x).arg(&mut part).arg(&ini).arg(&outi).arg(&ks);
21382            unsafe {
21383                b.launch(cfg)?;
21384            }
21385        }
21386        let f = self.func("matvec_bf16_v2_sk_combine");
21387        let n = (out_f * t) as u32;
21388        let cfg = LaunchConfig {
21389            grid_dim: (n.div_ceil(256), 1, 1),
21390            block_dim: (256, 1, 1),
21391            shared_mem_bytes: 0,
21392        };
21393        let ti = t as i32;
21394        let __s_b = self.gpu.stream();
21395        let mut b = __s_b.launch_builder(&f);
21396        b.arg(&part).arg(&mut *y).arg(&outi).arg(&ti).arg(&ks);
21397        unsafe {
21398            b.launch(cfg)?;
21399        }
21400        Ok(())
21401    }
21402
21403    /// Whether a v3 launch fits the 48 KB default dynamic-shared-memory cap at the current
21404    /// `mmv_block()`. Exposed so a bench or gate can skip the arm explicitly instead of
21405    /// discovering the decline as a launch error.
21406    pub fn gemv_v3_fits(&self) -> bool {
21407        gemv_v3_fits()
21408    }
21409
21410    /// The v3 bf16 GEMV launch (`matvec_bf16_v3`): the v2 kernel with its weight tiles staged
21411    /// through shared memory by `cp.async` instead of held in registers, which is the only one
21412    /// of this lane's three named next-levers that changes the in-flight-bytes arithmetic (see
21413    /// the kernel comment and the lane doc's section 9). Bit-identical to `matvec_bf16_v2`, and
21414    /// therefore to the shipped kernel: the chunk size is pinned to the shipped per-thread
21415    /// stride so a row's accumulation order is unchanged. No split-K twin — v3 is for the wide
21416    /// shapes; the caller falls back to v2 when a shape wants a split.
21417    pub fn matvec_bf16_v3_raw(
21418        &self,
21419        w: &CudaSlice<u8>,
21420        x: &CudaSlice<f32>,
21421        y: &mut CudaSlice<f32>,
21422        in_f: usize,
21423        out_f: usize,
21424        t: usize,
21425    ) -> Result<(), Box<dyn std::error::Error>> {
21426        if t == 0 || !in_f.is_multiple_of(8) || x.len() < t * in_f || y.len() < t * out_f {
21427            return Err("matvec_bf16_v3 geometry".into());
21428        }
21429        let nb = mmv_block();
21430        let smem = gemv_v3_smem_bytes(nb as usize);
21431        if smem > 48 * 1024 {
21432            return Err("matvec_bf16_v3 smem over the 48 KB default cap".into());
21433        }
21434        let f = self.func("matvec_bf16_v3");
21435        let cfg = LaunchConfig {
21436            grid_dim: (out_f.div_ceil(GEMV_V2_ROWS) as u32, t as u32, 1),
21437            block_dim: (nb, 1, 1),
21438            shared_mem_bytes: smem as u32,
21439        };
21440        let (ini, outi) = (in_f as i32, out_f as i32);
21441        let __s_b = self.gpu.stream();
21442        let mut b = __s_b.launch_builder(&f);
21443        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21444        unsafe {
21445            b.launch(cfg)?;
21446        }
21447        Ok(())
21448    }
21449
21450    /// v2 twin of `qmatvec_q8_0_mmvq_rp` — the kernel that actually serves the t=1 decode
21451    /// trunk in the `MEMRA_GLM5_W8` posture (`matvec_bf16_via_q8_mirror` ->
21452    /// `qmatvec_mmvq_into(.., QT_Q8_0, rp=true)`). Stages the q8_1 activation into shared
21453    /// memory once per CTA (the shipped kernel re-reads 36 B of activation per 34 B of weight,
21454    /// per lane, per block-iteration), packs 8 warps per block, and unrolls the block walk by
21455    /// two. BIT-IDENTICAL per output row: the per-row warp program is untouched.
21456    #[allow(clippy::too_many_arguments)]
21457    pub fn qmatvec_q8_0_rp_v2_raw(
21458        &self,
21459        w: &CudaSlice<u8>,
21460        aq: &CudaSlice<i8>,
21461        ad: &CudaSlice<f32>,
21462        y: &mut CudaSlice<f32>,
21463        in_f: usize,
21464        out_f: usize,
21465        t: usize,
21466    ) -> Result<(), Box<dyn std::error::Error>> {
21467        self.qmatvec_q8_0_rp_v2_raw_arm(w, aq, ad, y, in_f, out_f, t, q8_row_ilp_on())
21468    }
21469
21470    /// Arm-explicit form of [`Engine::qmatvec_q8_0_rp_v2_raw`]: `ilp` selects the
21471    /// `MEMRA_Q8_ROW_ILP` twin regardless of the door (the bench prices both on one process).
21472    #[allow(clippy::too_many_arguments)]
21473    pub fn qmatvec_q8_0_rp_v2_raw_arm(
21474        &self,
21475        w: &CudaSlice<u8>,
21476        aq: &CudaSlice<i8>,
21477        ad: &CudaSlice<f32>,
21478        y: &mut CudaSlice<f32>,
21479        in_f: usize,
21480        out_f: usize,
21481        t: usize,
21482        ilp: bool,
21483    ) -> Result<(), Box<dyn std::error::Error>> {
21484        if t == 0 || !in_f.is_multiple_of(32) || y.len() < t * out_f {
21485            return Err("qmatvec_q8_0_rp_v2 geometry".into());
21486        }
21487        let smem = Self::q8_v2_smem_bytes(in_f);
21488        if smem > 48 * 1024 {
21489            return Err("qmatvec_q8_0_rp_v2 smem over the 48 KB default cap".into());
21490        }
21491        if ilp {
21492            q8_row_ilp_note("qmatvec_q8_0_mmvq_rp_v2");
21493        }
21494        let f = self.func(if ilp {
21495            "qmatvec_q8_0_mmvq_rp_v2_ilp"
21496        } else {
21497            "qmatvec_q8_0_mmvq_rp_v2"
21498        });
21499        let cfg = LaunchConfig {
21500            grid_dim: ((out_f as u32).div_ceil(Q8_V2_ROWS), t as u32, 1),
21501            block_dim: (32, Q8_V2_ROWS, 1),
21502            shared_mem_bytes: smem as u32,
21503        };
21504        let (ini, outi, mi, rb) = (in_f as i32, out_f as i32, t as i32, 0i64);
21505        let __s_b = self.gpu.stream();
21506        let mut b = __s_b.launch_builder(&f);
21507        b.arg(w)
21508            .arg(aq)
21509            .arg(ad)
21510            .arg(&mut *y)
21511            .arg(&ini)
21512            .arg(&outi)
21513            .arg(&mi)
21514            .arg(&rb);
21515        unsafe {
21516            b.launch(cfg)?;
21517        }
21518        Ok(())
21519    }
21520
21521    /// v2 twin of `qmatvec_q8_0_rows_tw` — the VERIFY-width (t <= 8) W8 kernel reached through
21522    /// `matvec_bf16_via_q8_mirror_t` under `MEMRA_Q8T_WONCE`. The weight-once t-column structure
21523    /// is the shipped kernel's, so the activation is NOT staged (t*in_f would be 32 KB at t=8);
21524    /// the levers are the 8-warp packing and the block walk unrolled by two. BIT-IDENTICAL per
21525    /// (row, column).
21526    #[allow(clippy::too_many_arguments)]
21527    pub fn qmatvec_q8_0_rows_tw_v2_raw(
21528        &self,
21529        w: &CudaSlice<u8>,
21530        aq: &CudaSlice<i8>,
21531        ad: &CudaSlice<f32>,
21532        y: &mut CudaSlice<f32>,
21533        in_f: usize,
21534        out_f: usize,
21535        t: usize,
21536    ) -> Result<(), Box<dyn std::error::Error>> {
21537        if t == 0 || t > 8 || !in_f.is_multiple_of(32) || y.len() < t * out_f {
21538            return Err("qmatvec_q8_0_rows_tw_v2 geometry".into());
21539        }
21540        let f = self.func("qmatvec_q8_0_rows_tw_v2");
21541        let cfg = LaunchConfig {
21542            grid_dim: ((out_f as u32).div_ceil(Q8_V2_ROWS), 1, 1),
21543            block_dim: (32, Q8_V2_ROWS, 1),
21544            shared_mem_bytes: 0,
21545        };
21546        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
21547        let __s_b = self.gpu.stream();
21548        let mut b = __s_b.launch_builder(&f);
21549        b.arg(w)
21550            .arg(aq)
21551            .arg(ad)
21552            .arg(&mut *y)
21553            .arg(&ini)
21554            .arg(&outi)
21555            .arg(&ti);
21556        unsafe {
21557            b.launch(cfg)?;
21558        }
21559        Ok(())
21560    }
21561
21562    /// Bench-only direct arm selector for the W8 verify-width kernel (`b200_matvec_bench`,
21563    /// the `_arm_raw` precedent): `use_v2=false` launches the shipped `qmatvec_q8_0_rows_tw`
21564    /// (4 warps/block), `true` the v2 twin. Bypasses `matvec_bf16_via_q8_mirror_t`'s policy
21565    /// stack and the memoized door.
21566    #[allow(clippy::too_many_arguments)]
21567    pub fn qmatvec_q8_0_rows_tw_arm_raw(
21568        &self,
21569        w: &CudaSlice<u8>,
21570        aq: &CudaSlice<i8>,
21571        ad: &CudaSlice<f32>,
21572        y: &mut CudaSlice<f32>,
21573        in_f: usize,
21574        out_f: usize,
21575        t: usize,
21576        use_v2: bool,
21577    ) -> Result<(), Box<dyn std::error::Error>> {
21578        if use_v2 {
21579            return self.qmatvec_q8_0_rows_tw_v2_raw(w, aq, ad, y, in_f, out_f, t);
21580        }
21581        let f = self.func("qmatvec_q8_0_rows_tw");
21582        let cfg = LaunchConfig {
21583            grid_dim: ((out_f as u32).div_ceil(4), 1, 1),
21584            block_dim: (32, 4, 1),
21585            shared_mem_bytes: 0,
21586        };
21587        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
21588        let __s_b = self.gpu.stream();
21589        let mut b = __s_b.launch_builder(&f);
21590        b.arg(w)
21591            .arg(aq)
21592            .arg(ad)
21593            .arg(&mut *y)
21594            .arg(&ini)
21595            .arg(&outi)
21596            .arg(&ti);
21597        unsafe {
21598            b.launch(cfg)?;
21599        }
21600        Ok(())
21601    }
21602
21603    /// The FUSED six-projection KDA group for the `MEMRA_GLM5_W8` posture: one launch replaces
21604    /// the six separate `matvec_bf16_via_q8_mirror` calls (and their six redundant activation
21605    /// quantizes) the W8 path makes today. `bq`/`bk`/`bv` are the BF16 sources — they are the
21606    /// mirror cache keys, not the operands; their q8_0 rp4 mirrors are built on first use
21607    /// exactly as `matvec_bf16_via_q8_mirror` builds them, so nothing about residency changes.
21608    ///
21609    /// The W8 path had NO fused twin: `qmatvec_kda6_q8f32_mmvq` addresses interleaved 34 B
21610    /// blocks (a resident plain-layout Q8_0 tensor) while the W8 mirror is the split-plane rp4
21611    /// form, and `MEMRA_KDA_FUSED_PROJ`'s bf16 arm declines outright whenever W8 is on.
21612    ///
21613    /// NUMERIC CLASSES, unchanged from the sibling fused kernels: the three mirrored ranges are
21614    /// bit-identical to `qmatvec_q8_0_mmvq_rp` per row; the three f32 low-rank/beta ranges
21615    /// replace cuBLASLt with the same deterministic warp tree the q8 arm of
21616    /// `MEMRA_KDA_FUSED_PROJ` already ships and has pinned.
21617    #[allow(clippy::too_many_arguments)]
21618    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
21619    pub fn kda_proj_fused6_q8rp_raw(
21620        &self,
21621        bq: &CudaSlice<u8>,
21622        bk: &CudaSlice<u8>,
21623        bv: &CudaSlice<u8>,
21624        wfa: &CudaSlice<f32>,
21625        wga: &CudaSlice<f32>,
21626        wb: &CudaSlice<f32>,
21627        x: &CudaSlice<f32>,
21628        outs: &mut [CudaSlice<f32>; 6],
21629        in_f: usize,
21630        dims: [usize; 6],
21631        t: usize,
21632    ) -> Result<(), Box<dyn std::error::Error>> {
21633        self.kda_proj_fused6_q8rp_raw_arm(
21634            bq,
21635            bk,
21636            bv,
21637            wfa,
21638            wga,
21639            wb,
21640            x,
21641            outs,
21642            in_f,
21643            dims,
21644            t,
21645            q8_row_ilp_on(),
21646        )
21647    }
21648
21649    /// Arm-explicit form of [`Engine::kda_proj_fused6_q8rp_raw`]: `ilp` selects the
21650    /// `MEMRA_Q8_ROW_ILP` twin of the fused kernel regardless of the door.
21651    #[allow(clippy::too_many_arguments)]
21652    pub fn kda_proj_fused6_q8rp_raw_arm(
21653        &self,
21654        bq: &CudaSlice<u8>,
21655        bk: &CudaSlice<u8>,
21656        bv: &CudaSlice<u8>,
21657        wfa: &CudaSlice<f32>,
21658        wga: &CudaSlice<f32>,
21659        wb: &CudaSlice<f32>,
21660        x: &CudaSlice<f32>,
21661        outs: &mut [CudaSlice<f32>; 6],
21662        in_f: usize,
21663        dims: [usize; 6],
21664        t: usize,
21665        ilp: bool,
21666    ) -> Result<(), Box<dyn std::error::Error>> {
21667        self.kda_proj_fused6_q8rp_raw_pre(
21668            bq, bk, bv, wfa, wga, wb, x, outs, in_f, dims, t, ilp, None,
21669        )
21670    }
21671
21672    /// [`Engine::kda_proj_fused6_q8rp_raw_arm`] with an optional PRE-QUANTIZED activation
21673    /// (`pre_q8 = Some((aq, ad))`, the q8_1 view of `x` some producer already emitted, e.g.
21674    /// `rms_norm_zq8_f32` under `MEMRA_GLM5_Q8_FUSE_ATTN`): the launcher then skips its own
21675    /// `quantize_q8_1_into` and reads those planes. `quantize_q8_1_into` is `quantize_q8_1`
21676    /// verbatim and `rms_norm_zq8_f32` is `rms_norm` then `quantize_q8_1` bitwise, so the
21677    /// bytes are the ones this launcher would have produced (gate `tests/kda_fused_proj_gpu.rs`).
21678    #[allow(clippy::too_many_arguments)]
21679    pub fn kda_proj_fused6_q8rp_raw_pre(
21680        &self,
21681        bq: &CudaSlice<u8>,
21682        bk: &CudaSlice<u8>,
21683        bv: &CudaSlice<u8>,
21684        wfa: &CudaSlice<f32>,
21685        wga: &CudaSlice<f32>,
21686        wb: &CudaSlice<f32>,
21687        x: &CudaSlice<f32>,
21688        outs: &mut [CudaSlice<f32>; 6],
21689        in_f: usize,
21690        dims: [usize; 6],
21691        t: usize,
21692        ilp: bool,
21693        pre_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
21694    ) -> Result<(), Box<dyn std::error::Error>> {
21695        use cudarc::driver::DevicePtr;
21696        if t == 0 || t > 32 || !in_f.is_multiple_of(32) || x.len() < t * in_f {
21697            return Err("kda_proj_fused6_q8rp geometry".into());
21698        }
21699        let smem = Self::q8_v2_smem_bytes(in_f);
21700        if smem > 48 * 1024 {
21701            return Err("kda_proj_fused6_q8rp smem over the 48 KB default cap".into());
21702        }
21703        for (i, (o, want)) in outs.iter().zip(dims).enumerate() {
21704            if o.len() < t * want {
21705                return Err(format!("kda_proj_fused6_q8rp: output {i} too small").into());
21706            }
21707        }
21708        let nblk = in_f / 32;
21709        // Mirror keys, and the get-or-build, are `matvec_bf16_via_q8_mirror`'s verbatim.
21710        let keys: Vec<(u64, u32, u32)> = {
21711            let s = self.gpu.stream();
21712            [(bq, dims[0]), (bk, dims[1]), (bv, dims[2])]
21713                .iter()
21714                .map(|(d, out)| {
21715                    let (p, _g) = d.device_ptr(&s);
21716                    (p, in_f as u32, *out as u32)
21717                })
21718                .collect()
21719        };
21720        {
21721            let mut mirrors = self
21722                .w8_mirrors
21723                .lock()
21724                .map_err(|_| "w8 mirror map is poisoned")?;
21725            for ((d, out), key) in [(bq, dims[0]), (bk, dims[1]), (bv, dims[2])]
21726                .iter()
21727                .zip(&keys)
21728            {
21729                if !mirrors.contains_key(key) {
21730                    let mut interleaved = self.alloc_u8_uninit(out * Self::q8_0_row_bytes(in_f))?;
21731                    self.encode_q8_0_from_bf16(d, &mut interleaved, in_f, *out)?;
21732                    let planar = self.build_q8_rp4_raw(&interleaved, in_f, *out)?;
21733                    mirrors.insert(*key, planar);
21734                }
21735            }
21736        }
21737        // ONE activation quantize for all six projections. The unfused W8 path runs this once
21738        // per projection on the same `x` — six identical launches per layer.
21739        let akey = in_f * 64 + t.min(32);
21740        if pre_q8.is_none() {
21741            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21742            if let std::collections::hash_map::Entry::Vacant(slot) = act.entry(akey) {
21743                let aq = self.alloc_i8_uninit(32 * in_f)?;
21744                let ad = self.alloc_uninit::<f32>(32 * nblk)?;
21745                slot.insert((aq, ad));
21746            }
21747            let (aq, ad) = act.get_mut(&akey).expect("just inserted");
21748            self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
21749        }
21750        let mirrors = self
21751            .w8_mirrors
21752            .lock()
21753            .map_err(|_| "w8 mirror map is poisoned")?;
21754        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21755        let (aq, ad): (&CudaSlice<i8>, &CudaSlice<f32>) = match pre_q8 {
21756            Some((aq, ad)) => {
21757                if aq.len() < t * in_f || ad.len() < t * nblk {
21758                    return Err("kda_proj_fused6_q8rp: pre-quantized activation too small".into());
21759                }
21760                (aq, ad)
21761            }
21762            None => {
21763                let (aq, ad) = act.get(&akey).expect("built above");
21764                (aq, ad)
21765            }
21766        };
21767        let m0 = mirrors.get(&keys[0]).expect("built above");
21768        let m1 = mirrors.get(&keys[1]).expect("built above");
21769        let m2 = mirrors.get(&keys[2]).expect("built above");
21770        let blocks: usize = dims.iter().map(|d| d.div_ceil(Q8_V2_ROWS as usize)).sum();
21771        if ilp {
21772            q8_row_ilp_note("qmatvec_kda6_q8f32_rp_v2");
21773        }
21774        let f = self.func(if ilp {
21775            "qmatvec_kda6_q8f32_rp_v2_ilp"
21776        } else {
21777            "qmatvec_kda6_q8f32_rp_v2"
21778        });
21779        let cfg = LaunchConfig {
21780            grid_dim: (blocks as u32, t as u32, 1),
21781            block_dim: (32, Q8_V2_ROWS, 1),
21782            shared_mem_bytes: smem as u32,
21783        };
21784        let inf = in_f as i32;
21785        let d = dims.map(|v| v as i32);
21786        let mi = t as i32;
21787        let [o0, o1, o2, o3, o4, o5] = outs;
21788        let stream = self.gpu.stream();
21789        let mut b = stream.launch_builder(&f);
21790        b.arg(m0)
21791            .arg(m1)
21792            .arg(m2)
21793            .arg(wfa)
21794            .arg(wga)
21795            .arg(wb)
21796            .arg(aq)
21797            .arg(ad)
21798            .arg(x)
21799            .arg(&mut *o0)
21800            .arg(&mut *o1)
21801            .arg(&mut *o2)
21802            .arg(&mut *o3)
21803            .arg(&mut *o4)
21804            .arg(&mut *o5)
21805            .arg(&inf)
21806            .arg(&d[0])
21807            .arg(&d[1])
21808            .arg(&d[2])
21809            .arg(&d[3])
21810            .arg(&d[4])
21811            .arg(&d[5])
21812            .arg(&mi);
21813        unsafe {
21814            b.launch(cfg)?;
21815        }
21816        Ok(())
21817    }
21818
21819    /// v2 twin of [`Engine::moe_gate_up_preclamp8_q8`] (`MEMRA_B200_GEMV_V2`): 8 warps/block on
21820    /// `threadIdx.y` and a g-walk unrolled by two so both groups' weight/scale/activation loads
21821    /// issue before either dp4a chain runs. Per-warp arithmetic is the shipped kernel's, in the
21822    /// shipped per-accumulator order -> bit-identical per (o, j).
21823    #[allow(clippy::too_many_arguments)]
21824    pub fn moe_gate_up_preclamp8_q8_v2(
21825        &self,
21826        gp: WPtr8,
21827        up: WPtr8,
21828        aq: &CudaSlice<i8>,
21829        ad: &CudaSlice<f32>,
21830        gs: F32x8,
21831        us: F32x8,
21832        limit: f32,
21833        in_f: usize,
21834        n_ff: usize,
21835        n_used: usize,
21836        qt_g: i32,
21837        qt_u: i32,
21838        rb_g: usize,
21839        rb_u: usize,
21840    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21841        debug_assert!(
21842            limit > 1e-6,
21843            "moe_gate_up_preclamp8_q8_v2 needs a live limit; use moe_gate_up_silu8_q8"
21844        );
21845        const ROWS: u32 = 8;
21846        let f = self.func("moe_gate_up_preclamp8_q8_v2");
21847        let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
21848        let cfg = LaunchConfig {
21849            grid_dim: ((n_ff as u32).div_ceil(ROWS), n_used as u32, 1),
21850            block_dim: (32, ROWS, 1),
21851            shared_mem_bytes: 0,
21852        };
21853        let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
21854        let __s_b = self.gpu.stream();
21855        let mut b = __s_b.launch_builder(&f);
21856        b.arg(&gp)
21857            .arg(&up)
21858            .arg(aq)
21859            .arg(ad)
21860            .arg(&gs)
21861            .arg(&us)
21862            .arg(&limit)
21863            .arg(&mut act)
21864            .arg(&inf)
21865            .arg(&nff)
21866            .arg(&qt_g)
21867            .arg(&qt_u)
21868            .arg(&rbg)
21869            .arg(&rbu);
21870        unsafe {
21871            b.launch(cfg)?;
21872        }
21873        Ok(act)
21874    }
21875
21876    /// v2 twin of [`Engine::moe_down8_fma_q8`] (`MEMRA_B200_GEMV_V2`): ONE BLOCK per output row
21877    /// with warp `j` owning expert slot `j`, so the launch is `out_f * n_used` warps wide
21878    /// instead of `out_f` and the eight experts' bytes are in flight together. The slot chain
21879    /// is still one thread walking `j` ascending with `__fmaf_rn` on the same per-expert warp
21880    /// partials -> bit-identical per output row.
21881    #[allow(clippy::too_many_arguments)]
21882    pub fn moe_down8_fma_q8_v2(
21883        &self,
21884        dp: WPtr8,
21885        w: F32x8,
21886        aq2: &CudaSlice<i8>,
21887        ad2: &CudaSlice<f32>,
21888        dst: &mut cudarc::driver::CudaViewMut<f32>,
21889        in_f: usize,
21890        out_f: usize,
21891        n_used: usize,
21892        qt: i32,
21893        rb: usize,
21894    ) -> Result<(), Box<dyn std::error::Error>> {
21895        let f = self.func("moe_down8_fma_q8_v2");
21896        let cfg = LaunchConfig {
21897            grid_dim: (out_f as u32, 1, 1),
21898            block_dim: (32, 8, 1),
21899            shared_mem_bytes: 0,
21900        };
21901        let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
21902        let __s_b = self.gpu.stream();
21903        let mut b = __s_b.launch_builder(&f);
21904        b.arg(&dp)
21905            .arg(&w)
21906            .arg(aq2)
21907            .arg(ad2)
21908            .arg(dst)
21909            .arg(&inf)
21910            .arg(&outf)
21911            .arg(&nu)
21912            .arg(&qt)
21913            .arg(&rbi);
21914        unsafe {
21915            b.launch(cfg)?;
21916        }
21917        Ok(())
21918    }
21919
21920    /// True if `qtype` has a batched weight-resident (`_b2`/`_b4`) matvec kernel. These mirror the
21921    /// `_mmvq` kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave
21922    /// HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.
21923    pub fn batched_supports(&self, qtype: i32) -> bool {
21924        matches!(
21925            qtype,
21926            QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
21927        )
21928    }
21929
21930    /// IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A
21931    /// f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old
21932    /// opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/
21933    /// ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 ->
21934    /// 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The
21935    /// supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam),
21936    /// so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.
21937    pub fn iq_fast_enabled() -> bool {
21938        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21939        *ON.get_or_init(|| {
21940            std::env::var("MEMRA_IQ_FAST")
21941                .map(|v| v != "0")
21942                .unwrap_or(true)
21943        })
21944    }
21945
21946    /// b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch
21947    /// unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.
21948    pub fn b8_enabled() -> bool {
21949        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21950        *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
21951    }
21952
21953    /// Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
21954    pub fn batched_mcols(m: usize) -> usize {
21955        if m == 2 {
21956            2
21957        } else if m <= 4 {
21958            4
21959        } else if m <= 8 {
21960            8
21961        } else {
21962            16
21963        }
21964    }
21965
21966    /// Kernel name for the batched matvec of `(qtype, mcols)`. mcols ∈ {2,4,8}. The b8 tier is the
21967    /// K=4..7 spec-verify fix (T=5..8): pre-b8 those T fell to grid.y=m per-row MMVQ = m full
21968    /// weight reads/launch — the measured 27B K=4 cliff (101 -> 73 tok/s at p3 despite acceptance
21969    /// holding 54%). One b8 launch reads the weight ONCE for up to 8 columns (c >= m masked).
21970    fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
21971        Some(match (qtype, mcols) {
21972            (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
21973            (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
21974            (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
21975            // b16 now has BOTH forms (lane/rp-on-st, 2026-08-06). It used to be rp-ONLY, which
21976            // made the q8rp mirror the exact-16 tier's admission ticket for any model carrying a
21977            // single Q8_0 matmul — measured as the FP8-ST refusal (`L0.ssm_beta qtype=0
21978            // rp4=false`, 96 t / 23.9 MiB = 0.143% of resident weight). The mirror stays a
21979            // BANDWIDTH lever on Q8_0-dominant GGUFs; it is no longer a correctness prerequisite.
21980            (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
21981            (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
21982            (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
21983            (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
21984            // b16 base + _rp (lane/rp-on-st): the 9B NVFP4 GGUF's blocker — real NVFP4 GGUFs keep
21985            // Q4_K attention next to NVFP4 MLP, and the tier's predicate is an ALL.
21986            (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
21987            (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
21988            (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
21989            (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
21990            // b16 base only (lane/rp-on-st): Q5_K has no rp twins at any width, so there is
21991            // nothing to mirror. Named by the diagnostic as `L0.wqkv_gate qtype=3` on the 9B.
21992            (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
21993            (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
21994            (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
21995            (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
21996            (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
21997            (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
21998            (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
21999            (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
22000            // b16 (lane/rp-on-st): no mirror needed — NVFP4's 36 B/k32 block is already the
22001            // aligned form its own kernel walks. Unlocks the exact-16 tier for every NVFP4 model
22002            // AND for the mixed FP8-ST artifact, whose 193 NVFP4 tensors were refusing it.
22003            (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
22004            (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
22005            (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
22006            (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
22007            // b16 tier (lane/rp-on-st): e4m3 needs NO split-plane mirror to reach it — its native
22008            // row-major layout is already 32B-aligned per k32 block, so the base kernel IS the
22009            // aligned form. Contrast Q8_0, whose b16 exists only as the `_rp` twin (hence q8rp).
22010            (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
22011            (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
22012            (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
22013            (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
22014            (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
22015            _ => return None,
22016        })
22017    }
22018
22019    /// BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP
22020    /// win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2
22021    /// traffic 1x for m tokens (vs grid.y=m re-reading it m times). `mcols` ∈ {2,4,8} is the
22022    /// compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is
22023    /// [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post
22024    /// (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.
22025    ///
22026    /// NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real
22027    /// 27B verify (ncu --set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7
22028    /// for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG
22029    /// weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the
22030    /// column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy
22031    /// msweep on all six 27B shapes (2026-07-03):
22032    ///   `pf` = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere
22033    ///          it applies for b4 (-3..-14%), never loses;
22034    ///   `r2` = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but
22035    ///          wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a
22036    ///          fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on
22037    ///          a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%),
22038    ///          while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2:
22039    ///          -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144:
22040    ///          -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative.
22041    /// b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf.
22042    /// b2: in_f>=6144 -> r2, else base.
22043    /// MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam).
22044    /// All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce —
22045    /// only load issue time and the row->warp mapping change (kernel-check gates all of them).
22046    /// `rp` = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same
22047    /// wave-aware auto rule applies, mapped onto the `_rp` twins (rp/rpr2/rpr2w8 mirror
22048    /// pf/r2/r2w8 — regs 44/67/64 land in the same residency classes).
22049    /// The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so
22050    /// gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family
22051    /// (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required).
22052    /// Device SM count (cached) — grid-fill policy input.
22053    pub fn sm_count(&self) -> i32 {
22054        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
22055        *SMS.get_or_init(|| {
22056            use cudarc::driver::sys::CUdevice_attribute_enum as A;
22057            self.gpu
22058                .ctx
22059                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
22060                .unwrap_or(82)
22061        })
22062    }
22063
22064    #[allow(clippy::too_many_arguments)]
22065    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
22066    #[allow(clippy::if_same_then_else)] // allow: a fallback mapping table; distinct inputs deliberately share a target arm
22067    pub fn batched_variant(
22068        &self,
22069        _m: usize,
22070        in_f: usize,
22071        out_f: usize,
22072        qtype: i32,
22073        row_bytes: usize,
22074        mcols: usize,
22075        rp: bool,
22076    ) -> &'static str {
22077        // Q8_0 never joined the auto variant machinery (on sm_120 its only batched shapes
22078        // were tiny aux tensors). On Q8_0-trunk models the layout is the whole game: the
22079        // split-plane mirror (rp) routes to the _rp twins (H100 coalescing fix, 2026-07-26);
22080        // GGUF layout stays "base". rp bytes MUST never reach the base kernel or vice versa.
22081        if qtype == QT_Q8_0 {
22082            return if rp { "rp" } else { "base" };
22083        }
22084        static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
22085        let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
22086            Ok("base") => "base",
22087            Ok("pf") => "pf",
22088            Ok("r2") => "r2",
22089            Ok("r2w8") => "r2w8",
22090            Ok("pfr2") => "pfr2",
22091            Ok("ca") => "ca",
22092            Ok("car2") => "car2",
22093            // rp* = SPLIT-PLANE REPACKED layout kernels (A6 prototype): W must already be the
22094            // repacked buffer (msweep MSWEEP_RP harness) — never valid on GGUF-layout weights.
22095            Ok("rp") => "rp",
22096            Ok("rpr2") => "rpr2",
22097            Ok("rpr2w8") => "rpr2w8",
22098            // rpca* = cp.async software-pipelined split-plane (2026-07-05): hides the _rp
22099            // long_scoreboard load stall. rp-layout only; b4/b2 (no b8 twin).
22100            Ok("rpca") => "rpca",
22101            Ok("rpcar2") => "rpcar2",
22102            // 2026-07-06 m-small latency arc: rpsc = rpr2 + per-warp smem scale prestage (kills
22103            // the scale-plane global dependency, zero reg growth); rpms/rpmsc = m-split x2
22104            // across warp pairs (2x blocks of rpr2, column halves per warp, BIT-identical to
22105            // _rp); rpks/rpksc = k-split x2 (fastest microbench cells but k-reduce-order-shifted:
22106            // run-spec self-consistency FAILED on the 27B daily driver — verify logits must be
22107            // bit-identical to the decode path — measurement corpus ONLY, never auto).
22108            Ok("rpsc") => "rpsc",
22109            Ok("rpms") => "rpms",
22110            Ok("rpmsc") => "rpmsc",
22111            Ok("rpks") => "rpks",
22112            Ok("rpksc") => "rpksc",
22113            _ => "auto",
22114        });
22115        // cp.async ring variants need 16B-aligned rows (in_f%256==0 -> (in_f/64)*36 % 16 == 0)
22116        // and whole 32-group warp iterations (nsb%32==0 <=> in_f%1024==0). All 27B/9B trunk
22117        // shapes qualify; anything else falls back to the register variants.
22118        let ca_ok = qtype == QT_NVFP4 && row_bytes.is_multiple_of(16) && in_f.is_multiple_of(1024);
22119        // rpsc: smem scale plane fits (nsb64 <= 272) + int4-aligned staging (nsb64 % 4 == 0).
22120        // rpks/rpksc: half-plane staging alignment needs nsb64 % 8 == 0 (in_f % 512 == 0).
22121        // MEMRA_KS=0 removes the 2026-07-06 rpsc/rpks/rpksc entries from AUTO (rollback seam;
22122        // forced MEMRA_MMVQ_BV values still work).
22123        static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22124        let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
22125        let sc_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(256) && (in_f / 64 <= 272);
22126        let ks_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(512) && (in_f / 64 <= 272);
22127        static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
22128        let sms = *SMS.get_or_init(|| {
22129            use cudarc::driver::sys::CUdevice_attribute_enum as A;
22130            self.gpu
22131                .ctx
22132                .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
22133                .unwrap_or(82)
22134        });
22135        // k-quant r2 port (2026-07-04): q4_K/q5_K/q6_K have _r2/_r2w8 twins. ncu on the DRAM-cold
22136        // 9B msweep showed q4_K/q5_K b4 memory-latency bound like NVFP4 pre-fix (long_scoreboard
22137        // 19.6/16.4 per issue, DRAM 47.7/38.2%, L2 weight hit ~13%); q6_K lm_head is the exception
22138        // at DRAM 90-91% = wall-bound (yet r2 still wins -8%: deeper MLP raises achieved DRAM).
22139        // No _pf port (a k-quant group stages 10+ words vs NVFP4's 5 — register cost outweighs;
22140        // r2 covers the same MLP) and no rp (GGUF layout only). Q8_0 stays base: its only real
22141        // batched shapes are the tiny out_f=32 ssm_alpha/beta (8-block grids never fill one SM).
22142        // AUTO RULE = the measured winners table (differs from NVFP4's!):
22143        //   r2w8 NEVER in auto — the reg squeeze (72 -> 64 regs = stack spill) loses to unbounded
22144        //     r2 on every measured k-quant cell, incl. the wave-crossing lm_heads (q6_K 1316 vs
22145        //     r2 1258us) — kernels kept behind the force seam for the corpus;
22146        //   q4_K: r2 whenever the halved grid fills the SMs (blocks >= 4*SMs), INCLUDING the
22147        //     1.05-2.0 straggler window where NVFP4's r2 lost (qkv 1.78 waves: r2 -15% here; the
22148        //     k-quant base kernel leaves more latency on the table than a straggler wave costs);
22149        //   q5_K/q6_K: r2 only at waves >= 2 (the 248320-row lm_heads, 48+ waves: q6_K -8%, q5_K
22150        //     -2%); mid shapes measured base-or-flat (q5_K qkv 49.1 base vs 49.7 r2, attn_gate
22151        //     flat, attn_k base) — the 5/6-bit two-stream unpack makes r2's staging pricier.
22152        //   b2 same table with 8-row blocks: q4_K r2 when filled (-3..-22% all measured shapes),
22153        //     q5_K/q6_K r2 at waves >= 2 (27B lm_head -2.9%; 9B q6_K flat, harmless).
22154        let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
22155        // MEMRA_KQ_BV=base|r2|r2w8 forces the k-quant variant WITHOUT touching the NVFP4 dispatch
22156        // (MEMRA_MMVQ_BV is global — an interleaved k-quant-only e2e A/B needs this narrower seam).
22157        static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
22158        let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
22159            Ok("base") => "base",
22160            Ok("r2") => "r2",
22161            Ok("r2w8") => "r2w8",
22162            _ => "auto",
22163        });
22164        let variant: &'static str = if qtype == QT_Q4_0 {
22165            // Q4_0 r2 (gemma verify trunk, 2026-07-10): shared activation loads + the
22166            // row-independent ones-sum computed once per (col,group) for 2 rows. Same
22167            // fill rule as q4_K: r2 when the halved grid still fills the SMs.
22168            static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
22169            let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
22170                // ms/sm/la = force-only measurement seams (ALL FLAT/NEGATIVE 2026-07-13,
22171                // never auto): m-split flat (nvcc keeps 72 regs); smem-slab −11% (staging
22172                // + syncs cost more than the stalls, bank-pad made no difference);
22173                // register load-ahead flat (nvcc already reorders). The b-tier limiter
22174                // is still unidentified — see the jsonl row.
22175                Ok("base") => "base",
22176                Ok("r2") => "r2",
22177                Ok("ms") => "ms",
22178                Ok("sm") => "sm",
22179                Ok("la") => "la",
22180                _ => "auto",
22181            });
22182            let v = if q40 != "auto" {
22183                q40
22184            } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
22185                "r2"
22186            } else {
22187                "base"
22188            };
22189            // split-plane mirror twins (2026-07-10): same fill rule, _rp names.
22190            // (m-split r2 pair twin PROBED FLAT 2026-07-13 — nvcc kept 72 regs either way
22191            // and the limiter is the per-column activation load chain (long_scoreboard
22192            // 42.5%), not occupancy; arm killed per doctrine, jsonl row is the record.)
22193            if rp {
22194                match v {
22195                    "ms" => "r2ms_rp",
22196                    "sm" => "r2sm_rp",
22197                    "la" => "r2la_rp",
22198                    "r2" => "r2_rp",
22199                    _ => "rp",
22200                }
22201            } else if matches!(v, "ms" | "sm" | "la") {
22202                "r2"
22203            } else {
22204                v
22205            }
22206        } else if qtype != QT_NVFP4 && !kq_r2 {
22207            "base"
22208        } else if kq_r2 && rp {
22209            // K-quant split-plane mirror (2026-08-01): only the plain _rp batched twins are
22210            // compiled for q4_K/q6_K — rp is a LAYOUT, it must survive every heuristic
22211            // (split-plane bytes through a GGUF-layout kernel = NaN). q5_K never mirrors.
22212            "rp"
22213        } else if kq_r2 {
22214            // k-quant r2w8 only exists at b4 (b2_r2 already 8-resident; b8 has no w8 twin) ->
22215            // mcols != 4 forced r2w8 falls to unbounded r2.
22216            if kq_bv != "auto" {
22217                if kq_bv == "r2w8" && mcols != 4 {
22218                    "r2"
22219                } else {
22220                    kq_bv
22221                }
22222            } else if bv != "auto" {
22223                match bv {
22224                    "r2" | "pfr2" | "rpr2" | "car2" => "r2",
22225                    "r2w8" | "rpr2w8" => {
22226                        if mcols != 4 {
22227                            "r2"
22228                        } else {
22229                            "r2w8"
22230                        }
22231                    }
22232                    _ => "base", // base/pf/ca/rp forced -> base (no such k-quant kernels)
22233                }
22234            } else {
22235                #[allow(clippy::manual_div_ceil)]
22236                // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
22237                let blocks = (out_f + 7) / 8;
22238                let waves = blocks as f64 / (7 * sms as usize) as f64;
22239                let filled = blocks >= 4 * sms as usize;
22240                let use_r2 = if qtype == QT_Q4_K {
22241                    filled
22242                } else {
22243                    waves >= 2.0
22244                };
22245                if use_r2 { "r2" } else { "base" }
22246            }
22247        } else if bv != "auto" {
22248            // r2w8 only exists for b4/b8 (the b2_r2 kernel is already 8-blocks-resident at 60 regs).
22249            // ca/car2 need the alignment gate AND have no b8 twins; pfr2 has no b8 twin either —
22250            // unsupported (shape, mcols) combos fall back to pf/r2.
22251            // On rp buffers, forced legacy names map to their rp twins (layout law).
22252            let v = if bv == "r2w8" && mcols == 2 {
22253                "r2"
22254            } else if bv == "ca" && (!ca_ok || mcols == 8) {
22255                "pf"
22256            } else if bv == "car2" && (!ca_ok || mcols == 8) {
22257                "r2"
22258            } else if bv == "pfr2" && mcols == 8 {
22259                "r2"
22260            } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
22261                "rpr2"
22262            }
22263            // rpca* has no b8 twin (falls to rpr2w8/rpr2); needs the ca alignment gate.
22264            else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
22265                if mcols == 8 { "rpr2w8" } else { "rpr2" }
22266            } else if bv == "rpcar2" && mcols == 2 {
22267                "rpca"
22268            }
22269            // rpsc/rpmsc/rpks* gate on smem-fit + alignment; fall to rpr2 outside it
22270            // (rpms has no smem and no alignment need — always valid on rp buffers).
22271            else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
22272                "rpr2"
22273            } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
22274                "rpr2"
22275            } else {
22276                bv
22277            };
22278            if rp {
22279                match v {
22280                    "base" | "pf" | "ca" | "rp" => "rp",
22281                    "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
22282                    "r2w8" | "rpr2w8" => {
22283                        if mcols == 2 {
22284                            "rpr2"
22285                        } else {
22286                            "rpr2w8"
22287                        }
22288                    }
22289                    other => other, // rpca/rpcar2/rpsc/rpks/rpksc pass through (already rp-layout)
22290                }
22291            } else {
22292                v
22293            }
22294        } else if mcols == 8 {
22295            // b8 AUTO (2026-07-06 m-small latency arc, rtx6000 DRAM-cold rp msweep m=5/6/8 all five
22296            // 27B shapes): rpsc — the rpr2w8 schedule with the warp's scale rows prestaged to
22297            // smem, leaving ONE global dependency (the quant stream) in the k-loop at zero reg
22298            // growth. BIT-identical to rpr2w8 and wins or ties EVERY b8 cell: ffn_gate m5
22299            // 50.7->46.9 m8 64.1->57.1 (-11%), qkv m8 34.6->33.0, ssm_out m8 29.7->28.8,
22300            // attn_gate m8 26.9->26.1, ffn_down m5 58.2->56.9. The faster split-grid twins are
22301            // OUT: rpksc (k-split, ffn_down m5 -21%) broke run-spec self-consistency (k-reduce
22302            // order shifts verify argmax at tie margins — verify must stay bit-identical to the
22303            // m=1 decode chain); rpmsc (m-split, bit-identical) measured NEGATIVE everywhere
22304            // (twin warp's duplicated weight stream: ffn_down m5 85.7 vs 56.9).
22305            if rp {
22306                if sc_ok { "rpsc" } else { "rpr2w8" }
22307            } else {
22308                "r2w8"
22309            }
22310        } else if mcols >= 4 {
22311            // r2 runs 7 resident blocks/SM (67 regs); its __launch_bounds__(128,8) twin `r2w8`
22312            // (64 regs) runs 8. grid = ceil(out_f/8) for both. rp twins land in the same
22313            // residency classes (rp 44 regs ~ pf-class occupancy, rpr2 67, rpr2w8 64).
22314            #[allow(clippy::manual_div_ceil)]
22315            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
22316            let blocks = (out_f + 7) / 8;
22317            let r7 = 7 * sms as usize;
22318            let r8 = 8 * sms as usize;
22319            let waves = blocks as f64 / r7 as f64;
22320            let filled = blocks >= 4 * sms as usize;
22321            // 2026-07-06 m-small latency arc: b4 keeps the wave rule (rpms/rpmsc measured
22322            // flat-to-negative at m=3/4 on every shape — the m-split twin duplicates the weight
22323            // stream; rpsc b4 also negative on r2-class picks, ffn_down m4 51.1 vs 46.5).
22324            if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
22325                // the extra residency drops the INTEGER wave count -> the straggler wave a
22326                // latency-bound kernel pays in full disappears (ffn_down 1.11 -> 0.98 waves:
22327                // 112.5 -> 81.6us, beats pf 90.1; qkv 2.23 -> 1.95: 58.1 -> 51.1).
22328                if rp { "rpr2w8" } else { "r2w8" }
22329            } else if waves >= 2.0 || (waves <= 1.0 && filled) {
22330                // tail amortized (>=2 waves) or single wave: unbounded r2 (no reg-squeeze tax —
22331                // gate/up 81.1 vs 83.9 bounded, attn_q 61.0 vs 63.4).
22332                if rp { "rpr2" } else { "r2" }
22333            } else {
22334                // fractional straggler-wave window with no crossing, or grid too small to fill
22335                // the SMs (tiny out_f<=1024 shapes want max row-parallelism): prefetch variant
22336                // (rp = the r1 split-plane twin — measured the attn_gate winner, 35.4 vs pf 36.4).
22337                if rp { "rp" } else { "pf" }
22338            }
22339        } else if in_f >= 6144 {
22340            // b2 deep-k (2026-07-06): every new twin measured flat-to-negative here (rpms 44.1
22341            // vs rpr2 40.8 ffn_down; rpsc 43.6; the winning rpks is banned on k-order) — rpr2
22342            // stays.
22343            if rp { "rpr2" } else { "r2" }
22344        } else if rp {
22345            // b2 shallow-k: qkv (out_f=10240, 0.97 waves at 7-resident) is the one measured cell
22346            // where the r2-schedule scale-prestage twin beats the r1 rp pick (24.7 vs 28.9us
22347            // -15%); the wider (ffn_gate 1.65 waves) and smaller (attn_gate 0.58) shapes LOSE
22348            // (41.8 vs 38.2 / 16.6 vs 14.6) — gate on the single-wave window.
22349            #[allow(clippy::manual_div_ceil)]
22350            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
22351            let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
22352            if sc_ok && (0.9..=1.1).contains(&waves) {
22353                "rpsc"
22354            } else {
22355                "rp"
22356            }
22357        } else {
22358            "base"
22359        };
22360        variant
22361    }
22362
22363    #[allow(clippy::too_many_arguments)]
22364    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
22365    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
22366    pub fn qmatvec_mmvq_batched(
22367        &self,
22368        bytes: &CudaSlice<u8>,
22369        aq: &CudaSlice<i8>,
22370        ad: &CudaSlice<f32>,
22371        m: usize,
22372        in_f: usize,
22373        out_f: usize,
22374        qtype: i32,
22375        row_bytes: usize,
22376        mcols: usize,
22377        scale: f32,
22378        rp: bool,
22379    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22380        const ROWS_PER_BLOCK: u32 = 4;
22381        // TUNE SEAM (H100 lane): MEMRA_BVAR forces the batched-variant pick for the whole
22382        // process — the auto heuristics were tuned on sm_120 (82 SMs / 858 GB/s) and the
22383        // sm_90a re-tune sweeps this seam empirically. Layout variants stay safe: an rp
22384        // weight keeps its rp-layout kernel family regardless of the override.
22385        let forced: Option<&'static str> = {
22386            static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
22387            V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
22388                .as_deref()
22389                .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
22390        };
22391        let variant = match forced {
22392            Some(v) if !rp || v.contains("rp") => v,
22393            _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
22394        };
22395        let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
22396            format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
22397        })?;
22398        // b16 tier (t=9..16 verify): only base/_rp b16 kernels are compiled — the b2..b8
22399        // per-shape perf variants (r2/pf/...) do not apply at this width. rp is a LAYOUT,
22400        // not a perf variant: it must survive (base kernel on split-plane bytes = NaN).
22401        let variant = if mcols == 16 {
22402            if rp { "rp" } else { "base" }
22403        } else {
22404            variant
22405        };
22406        // EXACT-WIDTH b5/b6/b7 twins (lane/vt-fixes fix 1, 2026-08-03): the b8 kernels
22407        // allocate acc[WROWS][8] at ANY m, so T=5..7 verify paid the full 8-wide register
22408        // tax — the measured T=4->5 cliff. The same template at MCOLS=m runs the identical
22409        // per-(token,row) chain (columns c >= m never execute in either form) ->
22410        // BIT-IDENTICAL to the b8 launch. NVFP4 split-plane only (the sm_120 default trunk);
22411        // covers both b8 auto schedules (rpsc, rpr2w8). MEMRA_B567=0 rollback.
22412        static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22413        let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
22414        if b567
22415            && qtype == QT_NVFP4
22416            && rp
22417            && mcols == 8
22418            && (5..=7).contains(&m)
22419            && matches!(variant, "rpsc" | "rpr2w8")
22420        {
22421            let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
22422            let rows_per_block = ROWS_PER_BLOCK * 2; // r2-class schedules: 2 rows/warp
22423            let mut y = self.alloc_uninit::<f32>(m * out_f)?;
22424            let cfg = LaunchConfig {
22425                grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
22426                block_dim: (32, ROWS_PER_BLOCK, 1),
22427                shared_mem_bytes: 0,
22428            };
22429            let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22430            let __s_b = self.gpu.stream();
22431            let mut b = __s_b.launch_builder(&f);
22432            b.arg(bytes)
22433                .arg(aq)
22434                .arg(ad)
22435                .arg(&mut y)
22436                .arg(&inf)
22437                .arg(&outf)
22438                .arg(&mi)
22439                .arg(&rb);
22440            unsafe {
22441                b.launch(cfg)?;
22442            }
22443            if scale != 1.0 {
22444                self.scale_inplace(&mut y, scale, m * out_f)?;
22445            }
22446            return Ok(y);
22447        }
22448        let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
22449            "base" => (base_name.into(), ROWS_PER_BLOCK),
22450            "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
22451            "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
22452            "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
22453            "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), // 1 row/warp cp.async
22454            // split families: 2 warp-pairs x 2 rows = 4 rows/block (the k-range or column set
22455            // splits across the pair's two warps; grid.x doubles vs rpr2 at the same regs).
22456            "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
22457            "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
22458            "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
22459            "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
22460            "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
22461            "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
22462            "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
22463            v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), // r2-class: 2 rows/warp
22464        };
22465        debug_assert!(
22466            !rp || name.contains("_rp"),
22467            "rp weight dispatched to a GGUF-layout kernel"
22468        );
22469        let f = self.func(&name);
22470        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
22471        // r2sm_rp: [MCOLS][32 blk][8 int] activation slab + [MCOLS][32] f32 scales.
22472        let smem = if name.contains("_r2sm_rp") {
22473            (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
22474        } else {
22475            0
22476        };
22477        let cfg = LaunchConfig {
22478            grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
22479            block_dim: (32, ROWS_PER_BLOCK, 1),
22480            shared_mem_bytes: smem,
22481        };
22482        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22483        let __s_b = self.gpu.stream();
22484        let mut b = __s_b.launch_builder(&f);
22485        b.arg(bytes)
22486            .arg(aq)
22487            .arg(ad)
22488            .arg(&mut y)
22489            .arg(&inf)
22490            .arg(&outf)
22491            .arg(&mi)
22492            .arg(&rb);
22493        unsafe {
22494            b.launch(cfg)?;
22495        }
22496        if scale != 1.0 {
22497            self.scale_inplace(&mut y, scale, m * out_f)?;
22498        }
22499        Ok(y)
22500    }
22501
22502    /// BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation `x` to
22503    /// q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the
22504    /// kernel_check bit-equivalence gate. `mcols` ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.
22505    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
22506    pub fn qmatvec_batched_raw(
22507        &self,
22508        bytes: &CudaSlice<u8>,
22509        x: &CudaSlice<f32>,
22510        m: usize,
22511        in_f: usize,
22512        out_f: usize,
22513        qtype: i32,
22514        row_bytes: usize,
22515        mcols: usize,
22516        rp: bool,
22517    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22518        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
22519        self.qmatvec_mmvq_batched(
22520            bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
22521        )
22522    }
22523
22524    /// Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.
22525    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
22526    pub fn qmatvec_nvfp4_batched_raw(
22527        &self,
22528        bytes: &CudaSlice<u8>,
22529        x: &CudaSlice<f32>,
22530        m: usize,
22531        in_f: usize,
22532        out_f: usize,
22533        row_bytes: usize,
22534        mcols: usize,
22535        rp: bool,
22536    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22537        self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
22538    }
22539
22540    /// Stage-C FP4 gate (MEMRA_FP4): if `w` is an NVFP4 weight with in_f%64==0, run the native mxf4
22541    /// block-scale GEMM and apply the per-tensor macro-scale, returning Some(y). Else None (caller
22542    /// falls through to the int8 GEMM / dp4a). Strict opt-in over the proven int8 path; m>=16 only.
22543    fn try_fp4_gemm(
22544        &self,
22545        w: &crate::model::GpuTensor,
22546        x: &CudaSlice<f32>,
22547        m: usize,
22548        in_f: usize,
22549        out_f: usize,
22550    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22551        use crate::model::GpuTensor;
22552        if cfg!(memra_portable_cuda) {
22553            return Ok(None);
22554        }
22555        // MEMRA_FP4 reaches qmatvec_gemm_nvfp4_fp4, which ONLY the sm_120a fatbin contains:
22556        // cu/qmatvec_gemm.cu omits it on portable builds (MEMRA_PORTABLE_CUDA) AND on sm_100a
22557        // (build.rs passes -DMEMRA_DISABLE_NATIVE_FP4=1 there — the mxf4 block-scale MMA is an
22558        // sm_120a instruction encoding). Refuse at the door on EVERY build that lacks it. The
22559        // portable refusal alone was an enumeration, not a property: a 100a build is not
22560        // portable, so `MEMRA_FP4=1` sailed past it into Engine::func's "kernel not in any
22561        // fatbin" panic — found by the 100a fatbin-lookup census, lane/glm5-b200-prep-20260901
22562        // (same enumeration-vs-property class as the 2026-08-23 stub-polarity fixes in build.rs).
22563        if std::env::var("MEMRA_FP4").is_ok() {
22564            refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
22565            assert!(
22566                konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a"),
22567                "MEMRA_FP4 forces the native mxf4 block-scale GEMM (qmatvec_gemm_nvfp4_fp4), \
22568                 which only the sm_120a fatbin contains — this is an sm_{} build. Unset \
22569                 MEMRA_FP4; the W4A8 int8 path is the correct default for NVFP4 weights.",
22570                env!("MEMRA_BUILT_CUDA_ARCH")
22571            );
22572        }
22573        if std::env::var("MEMRA_FP4").is_err() {
22574            return Ok(None);
22575        }
22576        // CUTLASS prefill branch (m>=128 + MEMRA_FP4_CUTLASS + a repacked CutlassWeight present): route
22577        // to the CUTLASS sm120 NVFP4 GEMM, folding the per-tensor macro-scale into the epilogue alpha
22578        // (1/scale) — no post-matmul scale_inplace. Decode (m<128) and the m∈[16,128) middle band keep
22579        // the hand-roll below: CUTLASS's 128-row M-tile wastes work under 128.
22580        // The hand-roll applies the per-tensor macro-scale as a POST-matmul MULTIPLY (scale_inplace(y,
22581        // scale)); CUTLASS's epilogue does D = alpha * (A@B^T), so alpha == scale reproduces it exactly
22582        // (NOT 1/scale — the plan sketch had this inverted; the kernel_check arm gates it). scale==1.0
22583        // for the common no-macro-scale case.
22584        #[cfg(memra_cutlass)]
22585        if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
22586            if let GpuTensor::Quant {
22587                bytes,
22588                qtype,
22589                scale,
22590                row_bytes,
22591                cutlass,
22592                ..
22593            } = w
22594            {
22595                if *qtype == QT_NVFP4 && in_f % 64 == 0 {
22596                    if let Some(cw) = cutlass {
22597                        // Resident fast path: load-time-repacked B + swizzled SFB (no per-call repack).
22598                        let y = self.cutlass_fp4_gemm(
22599                            &cw.b_packed,
22600                            &cw.sfb_swizzled,
22601                            x,
22602                            *scale,
22603                            m,
22604                            out_f,
22605                            in_f,
22606                        )?;
22607                        return Ok(Some(y));
22608                    } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
22609                        // On-the-fly repack (MEMRA_FP4_CUTLASS_OTF): de-interleave + swizzle the B operand
22610                        // from raw bytes per prefill call. No resident doubling of the NVFP4 weight VRAM
22611                        // (the load-time repack ~doubles it) — needed for models that don't fit the
22612                        // resident path (e.g. the 27B on 24GB). Slower (per-call repack) but argmax-exact.
22613                        let (b_packed, sfb_sw) =
22614                            self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
22615                        let y =
22616                            self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
22617                        return Ok(Some(y));
22618                    }
22619                }
22620            }
22621        }
22622        if let GpuTensor::Quant {
22623            bytes,
22624            qtype,
22625            row_bytes,
22626            scale,
22627            rp,
22628            ..
22629        } = w
22630        {
22631            // A6: the hand-rolled W4A4 mxf4 GEMM reads 36B GGUF blocks — no rp port (MEMRA_FP4 is
22632            // an opt-in accuracy tradeoff); repacked tensors fall through to the int8 GEMM.
22633            if *qtype == QT_NVFP4 && in_f.is_multiple_of(64) && !*rp {
22634                let y =
22635                    self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
22636                return Ok(Some(y));
22637            }
22638        }
22639        Ok(None)
22640    }
22641
22642    /// rms_norm + fused fp16 twin (task #14): f32 output verbatim `rms_norm` + the fp16
22643    /// copy the f16-mirror GEMM group would otherwise produce with a standalone convert
22644    /// launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).
22645    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
22646    pub fn rms_norm_f16out(
22647        &self,
22648        x: &CudaSlice<f32>,
22649        w: &CudaSlice<f32>,
22650        dst: &mut CudaSlice<f32>,
22651        dst16: &mut CudaSlice<u8>,
22652        ncols: usize,
22653        nrows: usize,
22654        eps: f32,
22655    ) -> Result<(), Box<dyn std::error::Error>> {
22656        let f = self.func("rms_norm_f16out_f32");
22657        let cfg = LaunchConfig {
22658            grid_dim: (nrows as u32, 1, 1),
22659            block_dim: (rms_block(), 1, 1),
22660            shared_mem_bytes: 0,
22661        };
22662        let (nc, e) = (ncols as i32, eps);
22663        let __s_b = self.gpu.stream();
22664        let mut b = __s_b.launch_builder(&f);
22665        b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
22666        unsafe {
22667            b.launch(cfg)?;
22668        }
22669        Ok(())
22670    }
22671
22672    /// add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent —
22673    /// bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out's.
22674    #[allow(clippy::too_many_arguments)]
22675    pub fn add_rms_norm_f16out(
22676        &self,
22677        a: &CudaSlice<f32>,
22678        b: &CudaSlice<f32>,
22679        w: &CudaSlice<f32>,
22680        res: &mut CudaSlice<f32>,
22681        dst: &mut CudaSlice<f32>,
22682        dst16: &mut CudaSlice<u8>,
22683        ncols: usize,
22684        nrows: usize,
22685        eps: f32,
22686    ) -> Result<(), Box<dyn std::error::Error>> {
22687        let f = self.func("add_rms_norm_f16out_f32");
22688        let cfg = LaunchConfig {
22689            grid_dim: (nrows as u32, 1, 1),
22690            block_dim: (rms_block(), 1, 1),
22691            shared_mem_bytes: 0,
22692        };
22693        let (nc, e) = (ncols as i32, eps);
22694        let __s_lb = self.gpu.stream();
22695        let mut lb = __s_lb.launch_builder(&f);
22696        lb.arg(a)
22697            .arg(b)
22698            .arg(w)
22699            .arg(res)
22700            .arg(dst)
22701            .arg(dst16)
22702            .arg(&nc)
22703            .arg(&e);
22704        unsafe {
22705            lb.launch(cfg)?;
22706        }
22707        Ok(())
22708    }
22709
22710    /// matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused
22711    /// the convert). Mirror-less members fall back to `matmul` on the f32 activation.
22712    pub fn matmul_group_xh(
22713        &self,
22714        ws: &[&crate::model::GpuTensor],
22715        x: &CudaSlice<f32>,
22716        xh: &CudaSlice<u8>,
22717        m: usize,
22718    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22719        let mut out = Vec::with_capacity(ws.len());
22720        let in_f = ws[0].in_features();
22721        for w in ws {
22722            if w.in_features() == in_f
22723                && m >= 16
22724                && !self.verify_exact_on()
22725                && let Some(y) = self.try_f16_gemm_pre(w, xh, m)?
22726            {
22727                out.push(y);
22728                continue;
22729            }
22730            out.push(self.matmul(w, x, m)?);
22731        }
22732        Ok(out)
22733    }
22734
22735    /// task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity
22736    /// GDN steps). Layouts [T, H].
22737    pub fn gdn_pad_mask(
22738        &self,
22739        beta: &mut CudaSlice<f32>,
22740        g_log: &mut CudaSlice<f32>,
22741        len_d: &CudaSlice<i32>,
22742        h: usize,
22743        t: usize,
22744    ) -> Result<(), Box<dyn std::error::Error>> {
22745        let f = self.func("gdn_pad_mask_f32");
22746        let cfg = LaunchConfig::for_num_elems((t * h) as u32);
22747        let (hi, ti) = (h as i32, t as i32);
22748        let __s_b = self.gpu.stream();
22749        let mut b = __s_b.launch_builder(&f);
22750        b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
22751        unsafe {
22752            b.launch(cfg)?;
22753        }
22754        Ok(())
22755    }
22756
22757    /// task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row
22758    /// gather for the padded prime graph's h_seed/hlast.
22759    pub fn row_gather_dev(
22760        &self,
22761        src: &CudaSlice<f32>,
22762        dst: &mut CudaSlice<f32>,
22763        len_d: &CudaSlice<i32>,
22764        ncols: usize,
22765    ) -> Result<(), Box<dyn std::error::Error>> {
22766        let f = self.func("row_gather_dev_f32");
22767        let cfg = LaunchConfig::for_num_elems(ncols as u32);
22768        let nc = ncols as i32;
22769        let __s_b = self.gpu.stream();
22770        let mut b = __s_b.launch_builder(&f);
22771        b.arg(src).arg(dst).arg(len_d).arg(&nc);
22772        unsafe {
22773            b.launch(cfg)?;
22774        }
22775        Ok(())
22776    }
22777
22778    /// Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN
22779    /// 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to
22780    /// calling `matmul` per weight; the f16-mirror arm converts the activation ONCE for the
22781    /// whole group instead of once per GEMM (the standalone converts were ~250 launches/prime
22782    /// of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a
22783    /// different in_f) falls back to its own `matmul` — behavior unchanged.
22784    pub fn matmul_group(
22785        &self,
22786        ws: &[&crate::model::GpuTensor],
22787        x: &CudaSlice<f32>,
22788        m: usize,
22789    ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22790        use crate::model::GpuTensor;
22791        let mut out = Vec::with_capacity(ws.len());
22792        let any_mirror = ws
22793            .iter()
22794            .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
22795        if m >= 16 && any_mirror && !self.verify_exact_on() {
22796            let in_f = ws[0].in_features();
22797            let xh = self.f16_act(x, m * in_f, in_f)?;
22798            for w in ws {
22799                if w.in_features() == in_f
22800                    && let Some(y) = self.try_f16_gemm_pre(w, &xh, m)?
22801                {
22802                    out.push(y);
22803                    continue;
22804                }
22805                out.push(self.matmul(w, x, m)?);
22806            }
22807            return Ok(out);
22808        }
22809        for w in ws {
22810            out.push(self.matmul(w, x, m)?);
22811        }
22812        Ok(out)
22813    }
22814
22815    /// Cross-request grouped matmul (task #13): run ONE projection group over the
22816    /// CONCATENATION of several sequences' activations (m = sum of per-seq rows — the
22817    /// GEMM-batch win vLLM gets from continuous batching), then split each output back
22818    /// into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D
22819    /// copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K
22820    /// differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.
22821    pub fn matmul_group_multi(
22822        &self,
22823        ws: &[&crate::model::GpuTensor],
22824        xs: &[&CudaSlice<f32>],
22825        ms: &[usize],
22826    ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
22827        assert_eq!(xs.len(), ms.len());
22828        let in_f = ws[0].in_features();
22829        let total: usize = ms.iter().sum();
22830        let mut xcat = self.uninit(total * in_f)?;
22831        let mut off = 0usize;
22832        for (x, &m) in xs.iter().zip(ms) {
22833            self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
22834            off += m;
22835        }
22836        let ys = self.matmul_group(ws, &xcat, total)?;
22837        let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
22838        for (w, y) in ws.iter().zip(ys) {
22839            let out_f = w.out_features();
22840            let mut off = 0usize;
22841            for (s, &m) in ms.iter().enumerate() {
22842                let mut ys_s = self.uninit(m * out_f)?;
22843                let src = y.slice(off * out_f..(off + m) * out_f);
22844                self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
22845                out[s].push(ys_s);
22846                off += m;
22847            }
22848        }
22849        Ok(out)
22850    }
22851
22852    /// True if `w`'s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix).
22853    /// Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0.
22854    /// DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the
22855    /// dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax
22856    /// MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate
22857    /// shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only
22858    /// (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it
22859    /// at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a,
22860    /// MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).
22861    pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
22862        use crate::model::GpuTensor;
22863        if !legacy_quant_gemm_allowed(
22864            cfg!(memra_portable_cuda),
22865            cfg!(memra_hopper_mma),
22866            std::env::var_os("MEMRA_NO_GEMM").is_some(),
22867        ) {
22868            return false;
22869        }
22870        match w {
22871            GpuTensor::Quant { qtype, .. } => {
22872                matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
22873                    || (*qtype == QT_NVFP4 && w.in_features().is_multiple_of(64))
22874            }
22875            GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
22876        }
22877    }
22878
22879    /// Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill
22880    /// (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile,
22881    /// K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight
22882    /// read/decode N-fold (vs the dp4a matvec's per-token re-read). s32 accumulate is exact vs
22883    /// dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked
22884    /// `gemm_supports(w)`. y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.
22885    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
22886    pub fn qmatvec_gemm(
22887        &self,
22888        w: &crate::model::GpuTensor,
22889        aq: &CudaSlice<i8>,
22890        ad: &CudaSlice<f32>,
22891        m: usize,
22892    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22893        use crate::model::GpuTensor;
22894        let in_f = w.in_features();
22895        let out_f = w.out_features();
22896        let (bytes, qtype, row_bytes, scale, rp) = match w {
22897            GpuTensor::Quant {
22898                bytes,
22899                qtype,
22900                row_bytes,
22901                scale,
22902                rp,
22903                ..
22904            } => (bytes, *qtype, *row_bytes, *scale, *rp),
22905            _ => unreachable!("gemm_supports guaranteed Quant"),
22906        };
22907        // wgmma arm (sm_90a, task 8): the m64n64k32 warpgroup kernel reads the rp4 split-plane
22908        // mirror AS-IS (qplane rows = its A operand, the half dplane its scales) and the same
22909        // (aq, ad) activation planes. Same numeric class as the mma kernel below (exact s32 per
22910        // 32-block, one f32 scale fold per block, ascending K) — argmax/tolerance gated like
22911        // every prefill GEMM, not bit-gated. MEMRA_WGMMA=0 restores the portable kernel.
22912        if cfg!(memra_hopper_mma)
22913            && qtype == QT_Q8_0
22914            && out_f.is_multiple_of(64)
22915            && wgmma_gemm_enabled()
22916            && let GpuTensor::Quant { rp4: Some(m4), .. } = w
22917        {
22918            let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
22919            if scale != 1.0 {
22920                self.scale_inplace(&mut y, scale, m * out_f)?;
22921            }
22922            return Ok(y);
22923        }
22924        let name = match qtype {
22925            QT_Q8_0 => "qmatvec_gemm_q8_0",
22926            QT_Q4_K => "qmatvec_gemm_q4_K",
22927            QT_Q4_0 => {
22928                if rp {
22929                    "qmatvec_gemm_q4_0_rp"
22930                } else {
22931                    "qmatvec_gemm_q4_0"
22932                }
22933            }
22934            QT_Q5_K => "qmatvec_gemm_q5_K",
22935            QT_Q6_K => "qmatvec_gemm_q6_K",
22936            QT_NVFP4 => {
22937                if rp {
22938                    "qmatvec_gemm_nvfp4_rp"
22939                } else {
22940                    "qmatvec_gemm_nvfp4"
22941                }
22942            }
22943            _ => unreachable!(),
22944        };
22945        let f = self.func(name);
22946        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
22947        // CTA tile MUST match the .cu per-kernel tile. MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) runs llama's
22948        // 128x128 SQUARE tile (K1_BM=128 x K1_BN=128, 8 warps); kernel2 (Q6_K/NVFP4) keeps 64x256, 4 warps
22949        // (the macro BM/BN in the .cu). Grid dims are selected by qtype so each launches its own tile.
22950        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
22951        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
22952        let k1_tile = if is_k1 {
22953            k1_launch_override().unwrap_or((128, 128, 8))
22954        } else {
22955            (128, 128, 8)
22956        };
22957        let (bm, bn): (u32, u32) = if is_k1 {
22958            (k1_tile.0, k1_tile.1)
22959        } else {
22960            (64, 256)
22961        };
22962        let warps: u32 = if is_k1 {
22963            k1_tile.2
22964        } else {
22965            match qtype {
22966                QT_NVFP4 => 8,
22967                _ => 4,
22968            }
22969        };
22970        let cfg = LaunchConfig {
22971            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
22972            block_dim: (32, warps, 1),
22973            shared_mem_bytes: 0,
22974        };
22975        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22976        let __s_b = self.gpu.stream();
22977        let mut b = __s_b.launch_builder(&f);
22978        b.arg(bytes)
22979            .arg(aq)
22980            .arg(ad)
22981            .arg(&mut y)
22982            .arg(&inf)
22983            .arg(&outf)
22984            .arg(&mi)
22985            .arg(&rb);
22986        unsafe {
22987            b.launch(cfg)?;
22988        }
22989        if scale != 1.0 {
22990            self.scale_inplace(&mut y, scale, m * out_f)?;
22991        }
22992        Ok(y)
22993    }
22994
22995    /// Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes
22996    /// the f32 activation `x` to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor
22997    /// macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by
22998    /// kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.
22999    #[allow(clippy::too_many_arguments)]
23000    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23001    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
23002    pub fn qmatvec_gemm_raw(
23003        &self,
23004        bytes: &CudaSlice<u8>,
23005        x: &CudaSlice<f32>,
23006        m: usize,
23007        in_f: usize,
23008        out_f: usize,
23009        qtype: i32,
23010        row_bytes: usize,
23011    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23012        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
23013        let name = match qtype {
23014            QT_Q8_0 => "qmatvec_gemm_q8_0",
23015            QT_Q4_K => "qmatvec_gemm_q4_K",
23016            QT_Q4_0 => "qmatvec_gemm_q4_0",
23017            QT_Q5_K => "qmatvec_gemm_q5_K",
23018            QT_Q6_K => "qmatvec_gemm_q6_K",
23019            QT_NVFP4 => "qmatvec_gemm_nvfp4",
23020            QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
23021            _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
23022        };
23023        let f = self.func(name);
23024        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output: skip memset
23025        // MMQ-PORT: kernel1 (Q8_0/Q4_K/Q5_K) = llama 128x128 tile, 8 warps; kernel2 (Q6_K/NVFP4) = 64x256,
23026        // 4/8 warps. Grid tile per qtype (must match the .cu K1_BM/K1_BN vs BM/BN). KEEP IN SYNC w/ qmatvec_gemm.
23027        let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
23028        // TUNE SEAM: MEMRA_GEMM_K1_LAUNCH overrides kernel1's launch tile to match a -D-swept fatbin.
23029        let k1_tile = if is_k1 {
23030            k1_launch_override().unwrap_or((128, 128, 8))
23031        } else {
23032            (128, 128, 8)
23033        };
23034        let (bm, bn): (u32, u32) = if is_k1 {
23035            (k1_tile.0, k1_tile.1)
23036        } else {
23037            (64, 256)
23038        };
23039        let warps: u32 = if is_k1 {
23040            k1_tile.2
23041        } else {
23042            match qtype {
23043                QT_NVFP4 | QT_NVFP4_RP => 8,
23044                _ => 4,
23045            }
23046        };
23047        let cfg = LaunchConfig {
23048            grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
23049            block_dim: (32, warps, 1),
23050            shared_mem_bytes: 0,
23051        };
23052        let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
23053        let __s_b = self.gpu.stream();
23054        let mut b = __s_b.launch_builder(&f);
23055        b.arg(bytes)
23056            .arg(&aq)
23057            .arg(&ad)
23058            .arg(&mut y)
23059            .arg(&inf)
23060            .arg(&outf)
23061            .arg(&mi)
23062            .arg(&rb);
23063        unsafe {
23064            b.launch(cfg)?;
23065        }
23066        Ok(y)
23067    }
23068
23069    /// H100 warpgroup GEMM raw entry (task 8): launch `qmatvec_gemm_q8_0_wgmma` on an rp4
23070    /// split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr)
23071    /// owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads
23072    /// and dplane scale reads are unguarded); the token edge is guarded in-kernel.
23073    /// Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05
23074    /// vs CPU ref, 179us vs the portable mma kernel's 688us (3.84x, unpipelined).
23075    pub fn qmatvec_gemm_q8_0_wgmma_raw(
23076        &self,
23077        rp4: &CudaSlice<u8>,
23078        aq: &CudaSlice<i8>,
23079        ad: &CudaSlice<f32>,
23080        m: usize,
23081        in_f: usize,
23082        out_f: usize,
23083    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23084        assert!(
23085            out_f.is_multiple_of(64) && in_f.is_multiple_of(32),
23086            "wgmma GEMM needs out_f%64==0, in_f%32==0"
23087        );
23088        let f = self.func("qmatvec_gemm_q8_0_wgmma");
23089        let mut y = self.alloc_uninit::<f32>(m * out_f)?; // full-overwrite GEMM output
23090        let cfg = LaunchConfig {
23091            grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
23092            block_dim: (128, 1, 1),
23093            shared_mem_bytes: 0,
23094        };
23095        let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
23096        let __s_b = self.gpu.stream();
23097        let mut b = __s_b.launch_builder(&f);
23098        b.arg(rp4)
23099            .arg(aq)
23100            .arg(ad)
23101            .arg(&mut y)
23102            .arg(&inf)
23103            .arg(&outf)
23104            .arg(&mi);
23105        unsafe {
23106            b.launch(cfg)?;
23107        }
23108        Ok(y)
23109    }
23110
23111    /// y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.
23112    pub fn scale_inplace(
23113        &self,
23114        y: &mut CudaSlice<f32>,
23115        s: f32,
23116        n: usize,
23117    ) -> Result<(), Box<dyn std::error::Error>> {
23118        let f = self.func("scale_f32");
23119        let cfg = LaunchConfig::for_num_elems(n as u32);
23120        let (sf, ni) = (s, n as i32);
23121        let __s_b = self.gpu.stream();
23122        let mut b = __s_b.launch_builder(&f);
23123        b.arg(y).arg(&sf).arg(&ni);
23124        unsafe {
23125            b.launch(cfg)?;
23126        }
23127        Ok(())
23128    }
23129
23130    /// MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (`GpuTensor::FloatBf16`, raw
23131    /// bf16 bytes) to a transient f32 scratch of `n` elements, which then feeds the existing f32
23132    /// cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16
23133    /// weights + ONE (largest) weight's f32 expansion + activations. SLOW IS FINE (research mode).
23134    pub fn bf16_to_f32(
23135        &self,
23136        data: &cudarc::driver::CudaView<'_, u8>,
23137        n: usize,
23138    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23139        let mut out = self.alloc_uninit::<f32>(n)?;
23140        let f = self.func("bf16_to_f32");
23141        let cfg = LaunchConfig::for_num_elems(n as u32);
23142        let ni = n as i32;
23143        let __s_b = self.gpu.stream();
23144        let mut b = __s_b.launch_builder(&f);
23145        b.arg(data).arg(&mut out).arg(&ni);
23146        unsafe {
23147            b.launch(cfg)?;
23148        }
23149        Ok(out)
23150    }
23151
23152    /// Chunked bf16 linear (MEMRA_FULL_PREC): y[m,out] = x @ W_bf16^T with the f32 dequant scratch
23153    /// bounded to CHUNK_ROWS rows (256MB at in_f=4096) instead of the whole weight — the 4GB
23154    /// lm_head expansion OOM'd the 24GB budget. Row-chunking partitions OUTPUT rows; each row's
23155    /// dot is computed by the identical kernel on identical bytes, so per-(token,row) results are
23156    /// bit-identical to the unchunked form. `exact` selects linear_decode_exact (per-column m=1
23157    /// calls, the spec-verify contract) vs plain linear.
23158    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23159    fn linear_bf16_chunked(
23160        &self,
23161        x: &CudaSlice<f32>,
23162        data: &CudaSlice<u8>,
23163        m: usize,
23164        in_f: usize,
23165        out_f: usize,
23166        exact: bool,
23167        canonical_chunk_rows: Option<usize>,
23168    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23169        // MEMRA_STEP_TP_TIMING=1: cumulative cost of the per-call BF16->F32 expansion class
23170        // (alloc + convert kernel + f32 cuBLASLt = ~5x weight traffic). Prints every 1024 calls.
23171        // The stream sync per call perturbs wall; diagnostic only, never in a receipts run.
23172        static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23173        static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23174        static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23175        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
23176        let started = timing.then(std::time::Instant::now);
23177        let result =
23178            self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
23179        if let Some(started) = started {
23180            use std::sync::atomic::Ordering;
23181            self.stream().synchronize()?;
23182            let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
23183                + started.elapsed().as_nanos() as u64;
23184            let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
23185                + (in_f * out_f * 2) as u64;
23186            let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
23187            if calls.is_multiple_of(1024) {
23188                eprintln!(
23189                    "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
23190                     weight_gb={:.2}",
23191                    ns as f64 / 1.0e6,
23192                    ns as f64 / calls as f64 / 1.0e3,
23193                    wb as f64 / 1.0e9,
23194                );
23195            }
23196        }
23197        result
23198    }
23199
23200    /// MEMRA_BF16_MMV=1: decode-time (m=1) BF16-resident matvec door. Numeric class change vs
23201    /// the expansion path (single-kernel deterministic tree reduce vs f32 cuBLASLt), so it is
23202    /// default OFF and gated by the run-gen argmax gate + boot battery like the other
23203    /// numeric-class doors (DEV_ROUTES precedent).
23204    pub(crate) fn bf16_mmv_on() -> bool {
23205        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23206        *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
23207    }
23208
23209    /// One-block-per-row BF16 matvec: y[out_f] = W_bf16[out_f, in_f] @ x[in_f], f32 accumulate.
23210    /// Weights read once as raw bf16 (same bits<<16 expansion contract as `deq`'s QT_BF16 arm).
23211    fn matvec_bf16(
23212        &self,
23213        data: &CudaSlice<u8>,
23214        x: &CudaSlice<f32>,
23215        in_f: usize,
23216        out_f: usize,
23217    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23218        if data.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) {
23219            return Err(format!(
23220                "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
23221                data.len(),
23222                x.len()
23223            )
23224            .into());
23225        }
23226        let mut y = self.alloc_uninit::<f32>(out_f)?;
23227        let f = self.func("matvec_bf16_f32acc");
23228        let cfg = LaunchConfig {
23229            grid_dim: (out_f as u32, 1, 1),
23230            block_dim: (mmv_block(), 1, 1),
23231            shared_mem_bytes: 0,
23232        };
23233        let ini = in_f as i32;
23234        let __s_bld = self.gpu.stream();
23235        let mut bld = __s_bld.launch_builder(&f);
23236        bld.arg(data).arg(x).arg(&mut y).arg(&ini);
23237        unsafe {
23238            bld.launch(cfg)?;
23239        }
23240        Ok(y)
23241    }
23242
23243    /// Fused QK rms-norm + neox rope (t=1): one launch per rank replaces two rms_norm
23244    /// launches, a position upload, and the rope launch; the position is read directly from
23245    /// the caller's device counter (UVA). Numeric-class door (see the kernel doc).
23246    #[allow(clippy::too_many_arguments)]
23247    /// FUSION #1: qk norms + rope + dcw KV append + last-block len inc, one launch
23248    /// (replaces qk_norm_rope_into + append_kv_quantized_dcw + inc_i32 on the dcw path).
23249    /// Bit-identical to the split kernels; requires head_dim == 128 and
23250    /// kv_dim_v == kv_dim_k == nh_k * head_dim (caller-guarded fallback otherwise).
23251    #[allow(clippy::too_many_arguments)]
23252    /// T-ROW twin of `qk_norm_rope_append_inc_dcw` over a per-row session table (six u64
23253    /// words per row: K plane, V plane, len_ptr, base_ptr, done_ctr, pos_ptr). Raw q/k/v
23254    /// come from the [t, dim] tcol slabs; roped q lands in the [t, nh_q*head_dim] slab.
23255    /// Per-(row, head) block program == the t=1 kernel — bit-identical per row.
23256    #[allow(clippy::too_many_arguments)]
23257    pub fn qk_norm_rope_append_inc_dcw_rows(
23258        &self,
23259        q_raw_t: &CudaSlice<f32>,
23260        k_raw_t: &CudaSlice<f32>,
23261        v_raw_t: &CudaSlice<f32>,
23262        qw: &CudaSlice<f32>,
23263        kw: &CudaSlice<f32>,
23264        q_out_t: &mut CudaSlice<f32>,
23265        k_out_t: &mut CudaSlice<f32>,
23266        tab: &CudaSlice<u64>,
23267        pos_t: &CudaSlice<i32>,
23268        same_session: bool,
23269        t: usize,
23270        kv_dim_k: usize,
23271        kv_dim_v: usize,
23272        k_tok_bytes: usize,
23273        v_tok_bytes: usize,
23274        head_dim: usize,
23275        n_dims: usize,
23276        nh_q: usize,
23277        nh_k: usize,
23278        eps: f32,
23279        freq_base: f32,
23280        freq_scale: f32,
23281        ff: Option<&CudaSlice<f32>>,
23282    ) -> Result<(), Box<dyn std::error::Error>> {
23283        if head_dim != 128
23284            || kv_dim_v != kv_dim_k
23285            || kv_dim_k != nh_k * head_dim
23286            || t == 0
23287            || t > 32
23288            || tab.len() < t * 6
23289            || pos_t.len() < t
23290            || q_raw_t.len() < t * nh_q * head_dim
23291            || k_raw_t.len() < t * nh_k * head_dim
23292            || v_raw_t.len() < t * kv_dim_v
23293            || q_out_t.len() < t * nh_q * head_dim
23294            || k_out_t.len() < t * nh_k * head_dim
23295        {
23296            return Err(format!(
23297                "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
23298                 nh_q={nh_q} nh_k={nh_k}"
23299            )
23300            .into());
23301        }
23302        let f = self.func("qk_norm_rope_append_inc_dcw_rows");
23303        let same_t: i32 = if same_session { t as i32 } else { 0 };
23304        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
23305        let cfg = LaunchConfig {
23306            grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
23307            block_dim: (128, 1, 1),
23308            shared_mem_bytes: 0,
23309        };
23310        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
23311        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23312        let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
23313        let null: u64 = 0;
23314        let __s_b = self.gpu.stream();
23315        let mut b = __s_b.launch_builder(&f);
23316        b.arg(q_raw_t)
23317            .arg(k_raw_t)
23318            .arg(v_raw_t)
23319            .arg(qw)
23320            .arg(kw)
23321            .arg(q_out_t)
23322            .arg(k_out_t)
23323            .arg(tab)
23324            .arg(pos_t)
23325            .arg(&same_t)
23326            .arg(&kvk)
23327            .arg(&kvv)
23328            .arg(&ktb)
23329            .arg(&vtb)
23330            .arg(&hd)
23331            .arg(&nd)
23332            .arg(&nq)
23333            .arg(&nk)
23334            .arg(&eps)
23335            .arg(&theta_scale)
23336            .arg(&freq_scale);
23337        match ff {
23338            Some(freqs) => {
23339                b.arg(freqs);
23340            }
23341            None => {
23342                b.arg(&null);
23343            }
23344        }
23345        unsafe {
23346            b.launch(cfg)?;
23347        }
23348        Ok(())
23349    }
23350
23351    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23352    pub fn qk_norm_rope_append_inc_dcw(
23353        &self,
23354        q_raw: &CudaSlice<f32>,
23355        k_raw: &CudaSlice<f32>,
23356        v_raw: &CudaSlice<f32>,
23357        qw: &CudaSlice<f32>,
23358        kw: &CudaSlice<f32>,
23359        q_out: &mut CudaSlice<f32>,
23360        k_out: &mut CudaSlice<f32>,
23361        pos: &CudaSlice<i32>,
23362        k_plane: &mut CudaSlice<u8>,
23363        v_plane: &mut CudaSlice<u8>,
23364        // Shared ref by the planes_and_counters_mut split-borrow contract; the kernel is the
23365        // (single) writer, exactly like the split append+inc pair it replaces.
23366        len_dev: &CudaSlice<i32>,
23367        base_dev: Option<&CudaSlice<i32>>,
23368        done_ctr: &mut CudaSlice<u32>,
23369        kv_dim_k: usize,
23370        kv_dim_v: usize,
23371        k_tok_bytes: usize,
23372        v_tok_bytes: usize,
23373        head_dim: usize,
23374        n_dims: usize,
23375        nh_q: usize,
23376        nh_k: usize,
23377        eps: f32,
23378        freq_base: f32,
23379        freq_scale: f32,
23380        ff: Option<&CudaSlice<f32>>,
23381    ) -> Result<(), Box<dyn std::error::Error>> {
23382        if head_dim != 128
23383            || kv_dim_v != kv_dim_k
23384            || kv_dim_k != nh_k * head_dim
23385            || q_raw.len() < nh_q * head_dim
23386            || k_raw.len() < nh_k * head_dim
23387            || v_raw.len() < kv_dim_v
23388            || q_out.len() < nh_q * head_dim
23389            || k_out.len() < nh_k * head_dim
23390            || pos.is_empty()
23391            || done_ctr.is_empty()
23392        {
23393            return Err(format!(
23394                "qk_norm_rope_append_inc geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}                  kv_k={kv_dim_k} kv_v={kv_dim_v}"
23395            )
23396            .into());
23397        }
23398        let f = self.func("qk_norm_rope_append_inc_dcw");
23399        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
23400        let cfg = LaunchConfig {
23401            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
23402            block_dim: (128, 1, 1),
23403            shared_mem_bytes: 0,
23404        };
23405        let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
23406        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23407        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
23408        let null: u64 = 0;
23409        let __s_b = self.gpu.stream();
23410        let mut b = __s_b.launch_builder(&f);
23411        b.arg(q_raw)
23412            .arg(k_raw)
23413            .arg(v_raw)
23414            .arg(qw)
23415            .arg(kw)
23416            .arg(q_out)
23417            .arg(k_out)
23418            .arg(pos)
23419            .arg(&mut *k_plane)
23420            .arg(&mut *v_plane)
23421            .arg(len_dev);
23422        match base_dev {
23423            Some(base) => {
23424                b.arg(base);
23425            }
23426            None => {
23427                b.arg(&null);
23428            }
23429        }
23430        b.arg(&mut *done_ctr)
23431            .arg(&kvk)
23432            .arg(&kvv)
23433            .arg(&ktb)
23434            .arg(&vtb)
23435            .arg(&hd)
23436            .arg(&nd)
23437            .arg(&nq)
23438            .arg(&eps)
23439            .arg(&theta_scale)
23440            .arg(&freq_scale);
23441        match ff {
23442            Some(freqs) => {
23443                b.arg(freqs);
23444            }
23445            None => {
23446                b.arg(&null);
23447            }
23448        }
23449        unsafe {
23450            b.launch(cfg)?;
23451        }
23452        Ok(())
23453    }
23454
23455    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23456    pub fn qk_norm_rope_into(
23457        &self,
23458        q_raw: &CudaSlice<f32>,
23459        k_raw: &CudaSlice<f32>,
23460        qw: &CudaSlice<f32>,
23461        kw: &CudaSlice<f32>,
23462        q_out: &mut CudaSlice<f32>,
23463        k_out: &mut CudaSlice<f32>,
23464        pos: &CudaSlice<i32>,
23465        head_dim: usize,
23466        n_dims: usize,
23467        nh_q: usize,
23468        nh_k: usize,
23469        eps: f32,
23470        freq_base: f32,
23471        freq_scale: f32,
23472        ff: Option<&CudaSlice<f32>>,
23473    ) -> Result<(), Box<dyn std::error::Error>> {
23474        if head_dim > 512
23475            || q_raw.len() < nh_q * head_dim
23476            || k_raw.len() < nh_k * head_dim
23477            || q_out.len() < nh_q * head_dim
23478            || k_out.len() < nh_k * head_dim
23479            || qw.len() < head_dim
23480            || kw.len() < head_dim
23481            || pos.is_empty()
23482        {
23483            return Err(format!(
23484                "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
23485            )
23486            .into());
23487        }
23488        let f = self.func("qk_norm_rope_f32");
23489        let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
23490        let cfg = LaunchConfig {
23491            grid_dim: ((nh_q + nh_k) as u32, 1, 1),
23492            block_dim: (128, 1, 1),
23493            shared_mem_bytes: 0,
23494        };
23495        let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
23496        let __s_b = self.gpu.stream();
23497        let mut b = __s_b.launch_builder(&f);
23498        b.arg(q_raw)
23499            .arg(k_raw)
23500            .arg(qw)
23501            .arg(kw)
23502            .arg(q_out)
23503            .arg(k_out)
23504            .arg(pos)
23505            .arg(&hd)
23506            .arg(&nd)
23507            .arg(&nq)
23508            .arg(&eps)
23509            .arg(&theta_scale)
23510            .arg(&freq_scale);
23511        match ff {
23512            Some(ffv) => {
23513                b.arg(ffv);
23514                unsafe {
23515                    b.launch(cfg)?;
23516                }
23517            }
23518            None => {
23519                let null: u64 = 0;
23520                b.arg(&null);
23521                unsafe {
23522                    b.launch(cfg)?;
23523                }
23524            }
23525        }
23526        Ok(())
23527    }
23528
23529    /// Four-block F32 matvec with in-order block accumulation (see the kernel doc): one
23530    /// launch computes a rank's whole O partial from its four canonical column blocks.
23531    #[allow(clippy::too_many_arguments)]
23532    pub fn matvec_f32_b4_into(
23533        &self,
23534        w: [&CudaSlice<f32>; 4],
23535        x: &CudaSlice<f32>,
23536        y: &mut CudaSlice<f32>,
23537        block_cols: usize,
23538        out_f: usize,
23539    ) -> Result<(), Box<dyn std::error::Error>> {
23540        if !block_cols.is_multiple_of(4)
23541            || x.len() < 4 * block_cols
23542            || y.len() < out_f
23543            || w.iter().any(|w| w.len() != out_f * block_cols)
23544        {
23545            return Err(format!(
23546                "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
23547                x.len()
23548            )
23549            .into());
23550        }
23551        let f = self.func("matvec_f32_b4");
23552        let cfg = LaunchConfig {
23553            grid_dim: (out_f as u32, 1, 1),
23554            block_dim: (128, 1, 1),
23555            shared_mem_bytes: 0,
23556        };
23557        let (bc, of) = (block_cols as i32, out_f as i32);
23558        let __s_b = self.gpu.stream();
23559        let mut b = __s_b.launch_builder(&f);
23560        b.arg(w[0])
23561            .arg(w[1])
23562            .arg(w[2])
23563            .arg(w[3])
23564            .arg(x)
23565            .arg(y)
23566            .arg(&bc)
23567            .arg(&of);
23568        unsafe {
23569            b.launch(cfg)?;
23570        }
23571        Ok(())
23572    }
23573
23574    /// Sequential weighted row-sum: y[i] = sum_p w[p] * x[p*width+i] in row order — the exact
23575    /// per-element FP chain of a zero-reset plus n_rows sequential axpy launches.
23576    pub fn axpy_rows_seq_into(
23577        &self,
23578        x: &CudaSlice<f32>,
23579        w: &CudaSlice<f32>,
23580        y: &mut CudaSlice<f32>,
23581        width: usize,
23582        n_rows: usize,
23583    ) -> Result<(), Box<dyn std::error::Error>> {
23584        if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
23585            return Err(format!(
23586                "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
23587                x.len(),
23588                w.len(),
23589                y.len()
23590            )
23591            .into());
23592        }
23593        let f = self.func("axpy_rows_seq_f32");
23594        let cfg = LaunchConfig::for_num_elems(width as u32);
23595        let (wi, nr) = (width as i32, n_rows as i32);
23596        let __s_b = self.gpu.stream();
23597        let mut b = __s_b.launch_builder(&f);
23598        b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
23599        unsafe {
23600            b.launch(cfg)?;
23601        }
23602        Ok(())
23603    }
23604
23605    /// Token-major sequential weighted row sums. Each token reduces exactly `slots` rows in
23606    /// canonical route order.
23607    pub fn axpy_rows_seq_tokens_into(
23608        &self,
23609        x: &CudaSlice<f32>,
23610        w: &CudaSlice<f32>,
23611        y: &mut CudaSlice<f32>,
23612        width: usize,
23613        slots: usize,
23614        tokens: usize,
23615    ) -> Result<(), Box<dyn std::error::Error>> {
23616        let rows = slots
23617            .checked_mul(tokens)
23618            .ok_or("axpy_rows_seq_tokens row count overflow")?;
23619        if x.len() < rows * width || w.len() < rows || y.len() < tokens * width {
23620            return Err(format!(
23621                "axpy_rows_seq_tokens geometry x={} w={} y={} width={width} \
23622                 slots={slots} tokens={tokens}",
23623                x.len(),
23624                w.len(),
23625                y.len()
23626            )
23627            .into());
23628        }
23629        let f = self.func("axpy_rows_seq_tokens_f32");
23630        let block = 256u32;
23631        let cfg = LaunchConfig {
23632            grid_dim: ((width as u32).div_ceil(block), tokens as u32, 1),
23633            block_dim: (block, 1, 1),
23634            shared_mem_bytes: 0,
23635        };
23636        let (wi, sl, tk) = (width as i32, slots as i32, tokens as i32);
23637        let __s_b = self.gpu.stream();
23638        let mut b = __s_b.launch_builder(&f);
23639        b.arg(x).arg(w).arg(y).arg(&wi).arg(&sl).arg(&tk);
23640        unsafe {
23641            b.launch(cfg)?;
23642        }
23643        Ok(())
23644    }
23645
23646    /// Row-offset twin of `axpy_rows_seq_md_into` (spec verify t-column combine): the
23647    /// accumulation runs over rows [row0, row0+n_rows) of a taller partial slab — the
23648    /// exact sequential FP chain of the base kernel over that window.
23649    #[allow(clippy::too_many_arguments)]
23650    pub fn axpy_rows_seq_md_off_into(
23651        &self,
23652        x: &CudaSlice<f32>,
23653        w_route: &CudaSlice<f32>,
23654        md: &CudaSlice<f32>,
23655        sel: &CudaSlice<i32>,
23656        y: &mut CudaSlice<f32>,
23657        width: usize,
23658        n_rows: usize,
23659        row0: usize,
23660    ) -> Result<(), Box<dyn std::error::Error>> {
23661        if x.len() < (row0 + n_rows) * width
23662            || w_route.len() < row0 + n_rows
23663            || sel.len() < row0 + n_rows
23664            || y.len() < width
23665        {
23666            return Err(format!(
23667                "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
23668                 rows={n_rows} row0={row0}",
23669                x.len(),
23670                w_route.len(),
23671                sel.len(),
23672                y.len()
23673            )
23674            .into());
23675        }
23676        let f = self.func("axpy_rows_seq_md_off_f32");
23677        let cfg = LaunchConfig::for_num_elems(width as u32);
23678        let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
23679        let __s_b = self.gpu.stream();
23680        let mut b = __s_b.launch_builder(&f);
23681        b.arg(x)
23682            .arg(w_route)
23683            .arg(md)
23684            .arg(sel)
23685            .arg(y)
23686            .arg(&wi)
23687            .arg(&nr)
23688            .arg(&r0);
23689        unsafe {
23690            b.launch(cfg)?;
23691        }
23692        Ok(())
23693    }
23694
23695    /// Device-routed twin of `axpy_rows_seq_into`: the per-row weight folds in-kernel
23696    /// (w_route[p] * md[sel[p]] — the same single f32 multiply the host fold performs).
23697    #[allow(clippy::too_many_arguments)]
23698    pub fn axpy_rows_seq_md_into(
23699        &self,
23700        x: &CudaSlice<f32>,
23701        w_route: &CudaSlice<f32>,
23702        md: &CudaSlice<f32>,
23703        sel: &CudaSlice<i32>,
23704        y: &mut CudaSlice<f32>,
23705        width: usize,
23706        n_rows: usize,
23707    ) -> Result<(), Box<dyn std::error::Error>> {
23708        if x.len() < n_rows * width
23709            || w_route.len() < n_rows
23710            || sel.len() < n_rows
23711            || y.len() < width
23712        {
23713            return Err(format!(
23714                "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
23715                x.len(),
23716                w_route.len(),
23717                sel.len(),
23718                y.len()
23719            )
23720            .into());
23721        }
23722        let f = self.func("axpy_rows_seq_md_f32");
23723        let cfg = LaunchConfig::for_num_elems(width as u32);
23724        let (wi, nr) = (width as i32, n_rows as i32);
23725        let __s_b = self.gpu.stream();
23726        let mut b = __s_b.launch_builder(&f);
23727        b.arg(x)
23728            .arg(w_route)
23729            .arg(md)
23730            .arg(sel)
23731            .arg(y)
23732            .arg(&wi)
23733            .arg(&nr);
23734        unsafe {
23735            b.launch(cfg)?;
23736        }
23737        Ok(())
23738    }
23739
23740    /// BF16 twin of `matvec_f32_qkv_into` (weights as raw checkpoint bf16 bytes).
23741    #[allow(clippy::too_many_arguments)]
23742    /// T-COLUMN twin of `matvec_bf16_qkvg_into` (spec verify): weights read once, T input
23743    /// columns accumulated with per-column FP order identical to the t=1 kernel. Outputs
23744    /// land column-major-of-rows: yq[c*out_q + row] etc.
23745    #[allow(clippy::too_many_arguments)]
23746    pub fn matvec_bf16_qkvg_tcol_into(
23747        &self,
23748        wq: &CudaSlice<u8>,
23749        wk: &CudaSlice<u8>,
23750        wv: &CudaSlice<u8>,
23751        wg: &CudaSlice<u8>,
23752        x_t: &CudaSlice<f32>,
23753        yq: &mut CudaSlice<f32>,
23754        yk: &mut CudaSlice<f32>,
23755        yv: &mut CudaSlice<f32>,
23756        yg: &mut CudaSlice<f32>,
23757        in_f: usize,
23758        out_q: usize,
23759        out_kv: usize,
23760        out_g: usize,
23761        t: usize,
23762    ) -> Result<(), Box<dyn std::error::Error>> {
23763        if t == 0
23764            || t > 8
23765            || !in_f.is_multiple_of(8)
23766            || x_t.len() < t * in_f
23767            || yq.len() < t * out_q
23768            || yk.len() < t * out_kv
23769            || yv.len() < t * out_kv
23770            || (out_g > 0 && yg.len() < t * out_g)
23771        {
23772            return Err("matvec_bf16_qkvg_tcol geometry".into());
23773        }
23774        let grid = out_q + 2 * out_kv + out_g;
23775        let cfg = LaunchConfig {
23776            grid_dim: (grid as u32, 1, 1),
23777            block_dim: (mmv_block(), 1, 1),
23778            shared_mem_bytes: 0,
23779        };
23780        let (ini, oq, okv, og, ti) = (
23781            in_f as i32,
23782            out_q as i32,
23783            out_kv as i32,
23784            out_g as i32,
23785            t as i32,
23786        );
23787        let __s_b = self.gpu.stream();
23788        // One runtime-T program for every live width. The compile-time 2/4/8 twins are
23789        // retained in the fatbin as research controls, but dispatching them by the current
23790        // batch width changes kernels inside a request when peers arrive or retire. That is
23791        // a load-history numeric-program switch, and their pre-twin TOKFP receipts did not
23792        // qualify it (Hermes `64fa2b55baf0d887`).
23793        let f = self.func("matvec_bf16_qkvg_tcol");
23794        let mut b = __s_b.launch_builder(&f);
23795        b.arg(wq)
23796            .arg(wk)
23797            .arg(wv)
23798            .arg(wg)
23799            .arg(x_t)
23800            .arg(yq)
23801            .arg(yk)
23802            .arg(yv)
23803            .arg(yg)
23804            .arg(&ini)
23805            .arg(&oq)
23806            .arg(&okv)
23807            .arg(&og)
23808            .arg(&ti);
23809        unsafe {
23810            b.launch(cfg)?;
23811        }
23812        Ok(())
23813    }
23814
23815    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
23816    pub fn matvec_bf16_qkvg_into(
23817        &self,
23818        wq: &CudaSlice<u8>,
23819        wk: &CudaSlice<u8>,
23820        wv: &CudaSlice<u8>,
23821        wg: &CudaSlice<u8>,
23822        x: &CudaSlice<f32>,
23823        yq: &mut CudaSlice<f32>,
23824        yk: &mut CudaSlice<f32>,
23825        yv: &mut CudaSlice<f32>,
23826        yg: &mut CudaSlice<f32>,
23827        in_f: usize,
23828        out_q: usize,
23829        out_kv: usize,
23830        out_g: usize,
23831    ) -> Result<(), Box<dyn std::error::Error>> {
23832        if !in_f.is_multiple_of(8)
23833            || wq.len() != out_q * in_f * 2
23834            || wk.len() != out_kv * in_f * 2
23835            || wv.len() != out_kv * in_f * 2
23836            || wg.len() < out_g * in_f * 2
23837            || x.len() < in_f
23838            || yq.len() < out_q
23839            || yk.len() < out_kv
23840            || yv.len() < out_kv
23841            || (out_g > 0 && yg.len() < out_g)
23842        {
23843            return Err(format!(
23844                "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
23845            )
23846            .into());
23847        }
23848        let f = self.func("matvec_bf16_qkvg");
23849        let cfg = LaunchConfig {
23850            grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
23851            block_dim: (mmv_block(), 1, 1),
23852            shared_mem_bytes: 0,
23853        };
23854        let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
23855        let __s_b = self.gpu.stream();
23856        let mut b = __s_b.launch_builder(&f);
23857        b.arg(wq)
23858            .arg(wk)
23859            .arg(wv)
23860            .arg(wg)
23861            .arg(x)
23862            .arg(yq)
23863            .arg(yk)
23864            .arg(yv)
23865            .arg(yg)
23866            .arg(&inf)
23867            .arg(&oq)
23868            .arg(&okv)
23869            .arg(&og);
23870        unsafe {
23871            b.launch(cfg)?;
23872        }
23873        Ok(())
23874    }
23875
23876    /// BF16 twin of `matvec_f32_b4_into` (weights as raw checkpoint bf16 bytes).
23877    pub fn matvec_bf16_b4_into(
23878        &self,
23879        w: [&CudaSlice<u8>; 4],
23880        x: &CudaSlice<f32>,
23881        y: &mut CudaSlice<f32>,
23882        block_cols: usize,
23883        out_f: usize,
23884    ) -> Result<(), Box<dyn std::error::Error>> {
23885        if !block_cols.is_multiple_of(8)
23886            || x.len() < 4 * block_cols
23887            || y.len() < out_f
23888            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
23889        {
23890            return Err(format!(
23891                "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
23892                x.len()
23893            )
23894            .into());
23895        }
23896        // MEMRA_B4_X2=1: the #2b grid-halving twin — half the blocks, two rows each,
23897        // bit-identical per row (the second row's stream hides the first's reduce tail).
23898        static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23899        let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
23900        let f = self.func(if x2 {
23901            "matvec_bf16_b4_x2"
23902        } else {
23903            "matvec_bf16_b4"
23904        });
23905        let grid = if x2 { out_f.div_ceil(2) } else { out_f };
23906        let cfg = LaunchConfig {
23907            grid_dim: (grid as u32, 1, 1),
23908            block_dim: (mmv_block(), 1, 1),
23909            shared_mem_bytes: 0,
23910        };
23911        let (bc, of) = (block_cols as i32, out_f as i32);
23912        let __s_b = self.gpu.stream();
23913        let mut b = __s_b.launch_builder(&f);
23914        b.arg(w[0])
23915            .arg(w[1])
23916            .arg(w[2])
23917            .arg(w[3])
23918            .arg(x)
23919            .arg(y)
23920            .arg(&bc)
23921            .arg(&of);
23922        unsafe {
23923            b.launch(cfg)?;
23924        }
23925        Ok(())
23926    }
23927
23928    /// T-COLUMN twin of `matvec_bf16_b4_into` (spec verify o_proj): weights read once, T
23929    /// gated rows (each 4*block_cols wide) accumulated with per-column FP order identical
23930    /// to the t=1 kernel. Outputs land y[c*out_f + row]. Same blockDim as the t=1 launch —
23931    /// the shared-memory reduce order depends on it. Refuses under MEMRA_B4_X2 (different
23932    /// t=1 program).
23933    pub fn matvec_bf16_b4_tcol_into(
23934        &self,
23935        w: [&CudaSlice<u8>; 4],
23936        x_t: &CudaSlice<f32>,
23937        y_t: &mut CudaSlice<f32>,
23938        block_cols: usize,
23939        out_f: usize,
23940        t: usize,
23941    ) -> Result<(), Box<dyn std::error::Error>> {
23942        if !block_cols.is_multiple_of(8)
23943            || t == 0
23944            || t > 8
23945            || x_t.len() < t * 4 * block_cols
23946            || y_t.len() < t * out_f
23947            || w.iter().any(|w| w.len() != out_f * block_cols * 2)
23948        {
23949            return Err(format!(
23950                "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
23951                x_t.len()
23952            )
23953            .into());
23954        }
23955        if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
23956            return Err(
23957                "b4 tcol verify is qualified against the plain b4 kernel only \
23958                        (MEMRA_B4_X2=1 is a different t=1 program)"
23959                    .into(),
23960            );
23961        }
23962        // Keep one runtime-T program at every live width. Compile-time twins remain research
23963        // controls only; selecting them from the changing batch width switches programs
23964        // mid-request.
23965        let cfg = LaunchConfig {
23966            grid_dim: (out_f as u32, 1, 1),
23967            block_dim: (mmv_block(), 1, 1),
23968            shared_mem_bytes: 0,
23969        };
23970        let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
23971        let __s_b = self.gpu.stream();
23972        let f = self.func("matvec_bf16_b4_tcol");
23973        let mut b = __s_b.launch_builder(&f);
23974        b.arg(w[0])
23975            .arg(w[1])
23976            .arg(w[2])
23977            .arg(w[3])
23978            .arg(x_t)
23979            .arg(y_t)
23980            .arg(&bc)
23981            .arg(&of)
23982            .arg(&ti);
23983        unsafe {
23984            b.launch(cfg)?;
23985        }
23986        Ok(())
23987    }
23988
23989    /// `matvec_bf16` writing into a caller-owned output (persistent-workspace form).
23990    /// q8_0 row bytes for an `in_f`-wide weight row: one 34-byte block per 32 weights.
23991    pub fn q8_0_row_bytes(in_f: usize) -> usize {
23992        in_f / 32 * 34
23993    }
23994
23995    /// Dynamic shared memory one q8_0 v2 CTA needs: the staged q8_1 activation, `in_f` int8
23996    /// plus `in_f/32` f32 scales. 4.6 KB at in_f=4096. Mirrors the layout in cu/qmatvec.cu's
23997    /// `q8_0_stage_act`.
23998    pub fn q8_v2_smem_bytes(in_f: usize) -> usize {
23999        in_f + (in_f / 32) * 4
24000    }
24001
24002    /// Encode a resident bf16 weight slab into its q8_0 mirror (MEMRA_STEP_TP_W8). Runs once
24003    /// per matrix at load; the block program is the one `quant_K_block` writes for the KV
24004    /// cache, so the two formats cannot drift apart.
24005    pub fn encode_q8_0_from_bf16(
24006        &self,
24007        w_bf16: &CudaSlice<u8>,
24008        out: &mut CudaSlice<u8>,
24009        in_f: usize,
24010        out_f: usize,
24011    ) -> Result<(), Box<dyn std::error::Error>> {
24012        if !in_f.is_multiple_of(32)
24013            || w_bf16.len() < in_f * out_f * 2
24014            || out.len() < out_f * Self::q8_0_row_bytes(in_f)
24015        {
24016            return Err(format!(
24017                "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
24018                w_bf16.len(),
24019                out.len()
24020            )
24021            .into());
24022        }
24023        let f = self.func("encode_q8_0_rows_from_bf16");
24024        // Flat 1D grid of (row, 32-block) pairs, 4 pairs per block: rows on grid.y would cap
24025        // at 65535 and the LM head has 128896 rows.
24026        const PAIRS_PER_BLOCK: u32 = 4;
24027        let pairs = (out_f * (in_f / 32)) as u64;
24028        let cfg = LaunchConfig {
24029            grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
24030            block_dim: (32, PAIRS_PER_BLOCK, 1),
24031            shared_mem_bytes: 0,
24032        };
24033        let (ini, outi) = (in_f as i32, out_f as i32);
24034        let __s_b = self.gpu.stream();
24035        let mut b = __s_b.launch_builder(&f);
24036        b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
24037        unsafe {
24038            b.launch(cfg)?;
24039        }
24040        Ok(())
24041    }
24042
24043    /// ROW-RANGE-VIEW twin of `encode_q8_0_from_bf16`. Identical kernel, identical launch
24044    /// geometry, identical per-row program: only the operand type differs, because the split
24045    /// decode paths hold their rows as a `CudaView` of the resident slab, not as an owned slab.
24046    pub fn encode_q8_0_from_bf16_view(
24047        &self,
24048        w_bf16: &cudarc::driver::CudaView<'_, u8>,
24049        out: &mut CudaSlice<u8>,
24050        in_f: usize,
24051        out_f: usize,
24052    ) -> Result<(), Box<dyn std::error::Error>> {
24053        if !in_f.is_multiple_of(32)
24054            || w_bf16.len() < in_f * out_f * 2
24055            || out.len() < out_f * Self::q8_0_row_bytes(in_f)
24056        {
24057            return Err(format!(
24058                "encode_q8_0_from_bf16_view geometry in={in_f} out={out_f} src={} dst={}",
24059                w_bf16.len(),
24060                out.len()
24061            )
24062            .into());
24063        }
24064        let f = self.func("encode_q8_0_rows_from_bf16");
24065        const PAIRS_PER_BLOCK: u32 = 4;
24066        let pairs = (out_f * (in_f / 32)) as u64;
24067        let cfg = LaunchConfig {
24068            grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
24069            block_dim: (32, PAIRS_PER_BLOCK, 1),
24070            shared_mem_bytes: 0,
24071        };
24072        let (ini, outi) = (in_f as i32, out_f as i32);
24073        let __s_b = self.gpu.stream();
24074        let mut b = __s_b.launch_builder(&f);
24075        b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
24076        unsafe {
24077            b.launch(cfg)?;
24078        }
24079        Ok(())
24080    }
24081
24082    /// Fused q8_0 QKV against a q8_1 activation (MEMRA_STEP_TP_W8): one launch over the
24083    /// stacked q/k/v rows, each row running the exact `qmatvec_q8_0_mmvq_rp` per-row program.
24084    /// Bit-identical to three per-matrix mmvq calls; it exists because those three launches
24085    /// plus the activation quantize measured SLOWER than the bf16 fused kernel.
24086    #[allow(clippy::too_many_arguments)]
24087    pub fn qmatvec_q8_0_qkv_rp_into(
24088        &self,
24089        wq: &CudaSlice<u8>,
24090        wk: &CudaSlice<u8>,
24091        wv: &CudaSlice<u8>,
24092        aq: &CudaSlice<i8>,
24093        ad: &CudaSlice<f32>,
24094        yq: &mut CudaSlice<f32>,
24095        yk: &mut CudaSlice<f32>,
24096        yv: &mut CudaSlice<f32>,
24097        in_f: usize,
24098        out_q: usize,
24099        out_kv: usize,
24100    ) -> Result<(), Box<dyn std::error::Error>> {
24101        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
24102        let rows = out_q + 2 * out_kv;
24103        let nblk = in_f / 32;
24104        if !in_f.is_multiple_of(32)
24105            || aq.len() < in_f
24106            || ad.len() < nblk
24107            || yq.len() < out_q
24108            || yk.len() < out_kv
24109            || yv.len() < out_kv
24110            || wq.len() < out_q * nblk * 34
24111            || wk.len() < out_kv * nblk * 34
24112            || wv.len() < out_kv * nblk * 34
24113        {
24114            return Err(
24115                format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
24116            );
24117        }
24118        let f = self.func("qmatvec_q8_0_qkv_rp");
24119        let cfg = LaunchConfig {
24120            grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24121            block_dim: (32, ROWS_PER_BLOCK, 1),
24122            shared_mem_bytes: 0,
24123        };
24124        let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
24125        let __s_b = self.gpu.stream();
24126        let mut b = __s_b.launch_builder(&f);
24127        b.arg(wq)
24128            .arg(wk)
24129            .arg(wv)
24130            .arg(aq)
24131            .arg(ad)
24132            .arg(yq)
24133            .arg(yk)
24134            .arg(yv)
24135            .arg(&ini)
24136            .arg(&oq)
24137            .arg(&okv);
24138        unsafe {
24139            b.launch(cfg)?;
24140        }
24141        Ok(())
24142    }
24143
24144    /// Fused q8_0 O projection over the four HEAD_SPLIT blocks (MEMRA_STEP_TP_W8): one
24145    /// launch, one warp per output row, per-block reduce then add — the same shape
24146    /// `matvec_bf16_b4` uses, against a q8_1 activation.
24147    #[allow(clippy::too_many_arguments)]
24148    pub fn qmatvec_q8_0_b4_rp_into(
24149        &self,
24150        w: [&CudaSlice<u8>; 4],
24151        aq: &CudaSlice<i8>,
24152        ad: &CudaSlice<f32>,
24153        y: &mut CudaSlice<f32>,
24154        block_cols: usize,
24155        out_f: usize,
24156    ) -> Result<(), Box<dyn std::error::Error>> {
24157        const ROWS_PER_BLOCK: u32 = 4; // matches MEMRA_MMVQ_ROWS in qmatvec.cu
24158        let nblk = block_cols / 32;
24159        if !block_cols.is_multiple_of(32)
24160            || aq.len() < 4 * block_cols
24161            || ad.len() < 4 * nblk
24162            || y.len() < out_f
24163            || w.iter().any(|p| p.len() < out_f * nblk * 34)
24164        {
24165            return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
24166        }
24167        let f = self.func("qmatvec_q8_0_b4_rp");
24168        let cfg = LaunchConfig {
24169            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24170            block_dim: (32, ROWS_PER_BLOCK, 1),
24171            shared_mem_bytes: 0,
24172        };
24173        let (bc, of) = (block_cols as i32, out_f as i32);
24174        let __s_b = self.gpu.stream();
24175        let mut b = __s_b.launch_builder(&f);
24176        b.arg(w[0])
24177            .arg(w[1])
24178            .arg(w[2])
24179            .arg(w[3])
24180            .arg(aq)
24181            .arg(ad)
24182            .arg(y)
24183            .arg(&bc)
24184            .arg(&of);
24185        unsafe {
24186            b.launch(cfg)?;
24187        }
24188        Ok(())
24189    }
24190
24191    /// T-column twin of `matvec_bf16_via_q8_mirror`: one q8 launch over all t rows, sharing the
24192    /// same pointer-keyed mirror cache and a t-wide q8_1 activation.
24193    #[allow(clippy::map_entry)] // allow: the init bodies are fallible (`?`); Entry::or_insert_with cannot propagate errors
24194    fn matvec_bf16_via_q8_mirror_t(
24195        &self,
24196        data: &CudaSlice<u8>,
24197        x: &CudaSlice<f32>,
24198        y: &mut CudaSlice<f32>,
24199        in_f: usize,
24200        out_f: usize,
24201        t: usize,
24202    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
24203        use cudarc::driver::DevicePtr;
24204        let key = {
24205            let s = self.gpu.stream();
24206            let (p, _g) = data.device_ptr(&s);
24207            (p, in_f as u32, out_f as u32)
24208        };
24209        {
24210            let mut mirrors = self
24211                .w8_mirrors
24212                .lock()
24213                .map_err(|_| "w8 mirror map is poisoned")?;
24214            if !mirrors.contains_key(&key) {
24215                // memra#131: the mirror is built on FIRST DECODE USE. Under an open CUDA graph capture its
24216                // quantize and repack kernels would be RECORDED, never executed, the entry inserted as
24217                // built, and every later reader (the eager walk included) would read an uninitialised
24218                // mirror: that was the all-NaN KDA mixer at the first captured MoE-stage layer. Refuse by
24219                // name; the door warms every captured run before it captures, so this only fires when
24220                // something reaches an unwarmed weight inside a capture.
24221                if crate::glm5_graph_capture_open() {
24222                    return Err(format!(
24223                        "MEMRA_GLM5_W8: the q8_0 mirror for a {in_f}x{out_f} weight would be built inside an \
24224                         open CUDA graph capture (its quantize kernels recorded, never executed; memra#131). \
24225                         Warm the walk before capturing."
24226                    )
24227                    .into());
24228                }
24229                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
24230                self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
24231                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
24232                mirrors.insert(key, planar);
24233            }
24234        }
24235        let nblk = in_f / 32;
24236        // The t-wide activation scratch is keyed by (in_f, t-cap) so a wider walk regrows it.
24237        let akey = in_f * 64 + t.min(32);
24238        {
24239            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24240            if let std::collections::hash_map::Entry::Vacant(slot) = act.entry(akey) {
24241                let aq = self.alloc_i8_uninit(32 * in_f)?;
24242                let ad = self.alloc_uninit::<f32>(32 * nblk)?;
24243                slot.insert((aq, ad));
24244            }
24245            let (aq, ad) = act.get_mut(&akey).expect("just inserted");
24246            self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
24247        }
24248        let mirrors = self
24249            .w8_mirrors
24250            .lock()
24251            .map_err(|_| "w8 mirror map is poisoned")?;
24252        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24253        let mirror = mirrors.get(&key).expect("built above");
24254        let (aq, ad) = act.get(&akey).expect("built above");
24255        const ROWS_PER_BLOCK: u32 = 4;
24256        let (ini, of) = (in_f as i32, out_f as i32);
24257        // MEMRA_Q8T_WONCE=1: the weight-once twin — one row grid, each weight int4 loaded once
24258        // and dotted against all t columns. The `_t` form re-streams the shared weights per
24259        // column through __ldcs (measured 1.43-1.67x a single-column call for 2 columns).
24260        // MEMRA_B200_GEMV_V2, verify width: the t<=8 weight-once kernel is what the spec walk
24261        // runs in the W8 posture. Same weight-once structure, 8 warps per block and the block
24262        // walk unrolled by two; bit-identical per (row, column). t in 9..=32 keeps `_tw32`.
24263        if b200_gemv_v2_on() && q8t_wonce_on() && t <= 8 {
24264            if GEMV_V2_Q8_ROWS_TW_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0
24265            {
24266                eprintln!(
24267                    "[b200-gemv-v2] engaged arm=q8_rows_tw_v2 t={t} in_f={in_f} out_f={out_f} \
24268                     (W8 verify walk; MEMRA_B200_GEMV_V2=1)"
24269                );
24270            }
24271            self.qmatvec_q8_0_rows_tw_v2_raw(mirror, aq, ad, y, in_f, out_f, t)?;
24272            return Ok(Some(()));
24273        }
24274        if q8t_wonce_on() && t <= 32 {
24275            let f = self.func(if t <= 8 {
24276                "qmatvec_q8_0_rows_tw"
24277            } else {
24278                "qmatvec_q8_0_rows_tw32"
24279            });
24280            let cfg = LaunchConfig {
24281                grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24282                block_dim: (32, ROWS_PER_BLOCK, 1),
24283                shared_mem_bytes: 0,
24284            };
24285            let ti = t as i32;
24286            let __s_b = self.gpu.stream();
24287            let mut b = __s_b.launch_builder(&f);
24288            b.arg(mirror)
24289                .arg(aq)
24290                .arg(ad)
24291                .arg(&mut *y)
24292                .arg(&ini)
24293                .arg(&of)
24294                .arg(&ti);
24295            unsafe {
24296                b.launch(cfg)?;
24297            }
24298            return Ok(Some(()));
24299        }
24300        let f = self.func("qmatvec_q8_0_rows_t");
24301        let cfg = LaunchConfig {
24302            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
24303            block_dim: (32, ROWS_PER_BLOCK, 1),
24304            shared_mem_bytes: 0,
24305        };
24306        let __s_b = self.gpu.stream();
24307        let mut b = __s_b.launch_builder(&f);
24308        b.arg(mirror)
24309            .arg(aq)
24310            .arg(ad)
24311            .arg(&mut *y)
24312            .arg(&ini)
24313            .arg(&of);
24314        unsafe {
24315            b.launch(cfg)?;
24316        }
24317        Ok(Some(()))
24318    }
24319
24320    /// Get-or-build this bf16 weight's q8_0 mirror and run the GEMV through it. Returns
24321    /// `None` when the shape has no mirror form, so the caller falls back to bf16.
24322    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
24323    fn matvec_bf16_via_q8_mirror(
24324        &self,
24325        data: &CudaSlice<u8>,
24326        x: &CudaSlice<f32>,
24327        y: &mut CudaSlice<f32>,
24328        in_f: usize,
24329        out_f: usize,
24330    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
24331        use cudarc::driver::DevicePtr;
24332        let key = {
24333            let s = self.gpu.stream();
24334            let (p, _g) = data.device_ptr(&s);
24335            (p, in_f as u32, out_f as u32)
24336        };
24337        {
24338            let mut mirrors = self
24339                .w8_mirrors
24340                .lock()
24341                .map_err(|_| "w8 mirror map is poisoned")?;
24342            if !mirrors.contains_key(&key) {
24343                // memra#131: the mirror is built on FIRST DECODE USE. Under an open CUDA graph capture its
24344                // quantize and repack kernels would be RECORDED, never executed, the entry inserted as
24345                // built, and every later reader (the eager walk included) would read an uninitialised
24346                // mirror: that was the all-NaN KDA mixer at the first captured MoE-stage layer. Refuse by
24347                // name; the door warms every captured run before it captures, so this only fires when
24348                // something reaches an unwarmed weight inside a capture.
24349                if crate::glm5_graph_capture_open() {
24350                    return Err(format!(
24351                        "MEMRA_GLM5_W8: the q8_0 mirror for a {in_f}x{out_f} weight would be built inside an \
24352                         open CUDA graph capture (its quantize kernels recorded, never executed; memra#131). \
24353                         Warm the walk before capturing."
24354                    )
24355                    .into());
24356                }
24357                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
24358                self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
24359                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
24360                mirrors.insert(key, planar);
24361                // Which weights this half actually covers is not obvious from the call graph:
24362                // the head and the shared expert may reach the GPU through the rows fast path
24363                // or the fused dual-silu launcher instead of here. One line per mirror answers
24364                // that without a profiler (the hybrid half measured +0.1% and this is how we
24365                // find out whether it even fired).
24366                if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
24367                    eprintln!(
24368                        "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
24369                        mirrors.len()
24370                    );
24371                }
24372            }
24373        }
24374        let nblk = in_f / 32;
24375        {
24376            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24377            if !act.contains_key(&in_f) {
24378                let aq = self.alloc_uninit::<i8>(in_f)?;
24379                let ad = self.alloc_uninit::<f32>(nblk)?;
24380                act.insert(in_f, (aq, ad));
24381            }
24382            let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
24383            self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
24384        }
24385        let mirrors = self
24386            .w8_mirrors
24387            .lock()
24388            .map_err(|_| "w8 mirror map is poisoned")?;
24389        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24390        let mirror = mirrors.get(&key).expect("built above");
24391        let (aq, ad) = act.get(&in_f).expect("built above");
24392        // MEMRA_B200_GEMV_V2 (lane/b200-gemv-hbm-20260902 round 3). THIS is the t=1 decode
24393        // kernel in the posture we actually serve: the bf16 arms further up
24394        // `matvec_bf16_rows_into` are unreachable once MEMRA_GLM5_W8 reroutes here, which is
24395        // why serving A/B pair 1 moved nothing (49.2 -> 49.3 tok/s, no engagement line). The v2
24396        // twin stages the q8_1 activation into shared memory once per CTA — the shipped
24397        // `qmatvec_q8_0_mmvq_rp` re-reads 36 B of activation per 34 B of weight, per lane, per
24398        // block-iteration — packs 8 warps per block and unrolls the block walk by two.
24399        // Bit-identical per output row; declines to the shipped dispatch when the staged
24400        // activation would not fit the 48 KB default smem cap.
24401        if b200_gemv_v2_on() && in_f.is_multiple_of(32) && Self::q8_v2_smem_bytes(in_f) <= 48 * 1024
24402        {
24403            if GEMV_V2_Q8_RP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
24404                eprintln!(
24405                    "[b200-gemv-v2] engaged arm=q8_rp_v2 t=1 in_f={in_f} out_f={out_f} \
24406                     (W8 posture; MEMRA_B200_GEMV_V2=1)"
24407                );
24408            }
24409            self.qmatvec_q8_0_rp_v2_raw(mirror, aq, ad, y, in_f, out_f, 1)?;
24410            return Ok(Some(()));
24411        }
24412        self.qmatvec_mmvq_into(
24413            mirror,
24414            aq,
24415            ad,
24416            1,
24417            in_f,
24418            out_f,
24419            QT_Q8_0,
24420            Self::q8_0_row_bytes(in_f),
24421            1.0,
24422            true,
24423            y,
24424        )?;
24425        Ok(Some(()))
24426    }
24427
24428    /// T-column q8_0 QKV for the VERIFY walk (MEMRA_STEP_TP_W8). nsys put the bf16 twin
24429    /// `matvec_bf16_qkvg_tcol` at 12.3% of spec GPU time and `matvec_bf16_b4_tcol` at 24.8%:
24430    /// the W8 door had replaced only the decode kernels, so 37% of the verify still streamed
24431    /// bf16. Bit-identical to `t` separate `qmatvec_q8_0_qkv_rp` calls.
24432    #[allow(clippy::too_many_arguments)]
24433    pub fn qmatvec_q8_0_qkv_rp_t_into(
24434        &self,
24435        wq: &CudaSlice<u8>,
24436        wk: &CudaSlice<u8>,
24437        wv: &CudaSlice<u8>,
24438        aq: &CudaSlice<i8>,
24439        ad: &CudaSlice<f32>,
24440        yq: &mut CudaSlice<f32>,
24441        yk: &mut CudaSlice<f32>,
24442        yv: &mut CudaSlice<f32>,
24443        in_f: usize,
24444        out_q: usize,
24445        out_kv: usize,
24446        t: usize,
24447    ) -> Result<(), Box<dyn std::error::Error>> {
24448        const ROWS_PER_BLOCK: u32 = 4;
24449        let rows = out_q + 2 * out_kv;
24450        let nblk = in_f / 32;
24451        if !in_f.is_multiple_of(32)
24452            || t == 0
24453            || aq.len() < t * in_f
24454            || ad.len() < t * nblk
24455            || yq.len() < t * out_q
24456            || yk.len() < t * out_kv
24457            || yv.len() < t * out_kv
24458        {
24459            return Err(format!("q8_0 qkv rp_t geometry in={in_f} t={t}").into());
24460        }
24461        let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
24462        // MEMRA_Q8T_WONCE=1: weight-once twin — see qmatvec.cu's `_tw` block for why the `_t`
24463        // form re-streams the fully-shared QKV weights per column (__ldcs + column grid axis;
24464        // measured 1.67x a single-column call for 2 columns).
24465        if q8t_wonce_on() && t <= 32 {
24466            let f = self.func(if t <= 8 {
24467                "qmatvec_q8_0_qkv_rp_tw"
24468            } else {
24469                "qmatvec_q8_0_qkv_rp_tw32"
24470            });
24471            let cfg = LaunchConfig {
24472                grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24473                block_dim: (32, ROWS_PER_BLOCK, 1),
24474                shared_mem_bytes: 0,
24475            };
24476            let ti = t as i32;
24477            let __s_b = self.gpu.stream();
24478            let mut b = __s_b.launch_builder(&f);
24479            b.arg(wq)
24480                .arg(wk)
24481                .arg(wv)
24482                .arg(aq)
24483                .arg(ad)
24484                .arg(yq)
24485                .arg(yk)
24486                .arg(yv)
24487                .arg(&ini)
24488                .arg(&oq)
24489                .arg(&okv)
24490                .arg(&ti);
24491            unsafe {
24492                b.launch(cfg)?;
24493            }
24494            return Ok(());
24495        }
24496        let f = self.func("qmatvec_q8_0_qkv_rp_t");
24497        let cfg = LaunchConfig {
24498            grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
24499            block_dim: (32, ROWS_PER_BLOCK, 1),
24500            shared_mem_bytes: 0,
24501        };
24502        let __s_b = self.gpu.stream();
24503        let mut b = __s_b.launch_builder(&f);
24504        b.arg(wq)
24505            .arg(wk)
24506            .arg(wv)
24507            .arg(aq)
24508            .arg(ad)
24509            .arg(yq)
24510            .arg(yk)
24511            .arg(yv)
24512            .arg(&ini)
24513            .arg(&oq)
24514            .arg(&okv);
24515        unsafe {
24516            b.launch(cfg)?;
24517        }
24518        Ok(())
24519    }
24520
24521    /// T-column q8_0 o_proj over the four HEAD_SPLIT blocks (MEMRA_STEP_TP_W8, verify walk).
24522    /// Bit-identical to `t` separate `qmatvec_q8_0_b4_rp` calls.
24523    #[allow(clippy::too_many_arguments)]
24524    pub fn qmatvec_q8_0_b4_rp_t_into(
24525        &self,
24526        w: [&CudaSlice<u8>; 4],
24527        aq: &CudaSlice<i8>,
24528        ad: &CudaSlice<f32>,
24529        y: &mut CudaSlice<f32>,
24530        block_cols: usize,
24531        out_f: usize,
24532        t: usize,
24533    ) -> Result<(), Box<dyn std::error::Error>> {
24534        const ROWS_PER_BLOCK: u32 = 4;
24535        let nblk = block_cols / 32;
24536        if !block_cols.is_multiple_of(32)
24537            || t == 0
24538            || aq.len() < t * 4 * block_cols
24539            || ad.len() < t * 4 * nblk
24540            || y.len() < t * out_f
24541        {
24542            return Err(format!("q8_0 b4 rp_t geometry cols={block_cols} t={t}").into());
24543        }
24544        let (bc, of) = (block_cols as i32, out_f as i32);
24545        // MEMRA_Q8T_WONCE=1: weight-once twin (see qmatvec.cu; `_t` measured 1.43x for 2 columns
24546        // on fully-shared o_proj weights).
24547        if q8t_wonce_on() && t <= 32 {
24548            let f = self.func(if t <= 8 {
24549                "qmatvec_q8_0_b4_rp_tw"
24550            } else {
24551                "qmatvec_q8_0_b4_rp_tw32"
24552            });
24553            let cfg = LaunchConfig {
24554                grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24555                block_dim: (32, ROWS_PER_BLOCK, 1),
24556                shared_mem_bytes: 0,
24557            };
24558            let ti = t as i32;
24559            let __s_b = self.gpu.stream();
24560            let mut b = __s_b.launch_builder(&f);
24561            b.arg(w[0])
24562                .arg(w[1])
24563                .arg(w[2])
24564                .arg(w[3])
24565                .arg(aq)
24566                .arg(ad)
24567                .arg(y)
24568                .arg(&bc)
24569                .arg(&of)
24570                .arg(&ti);
24571            unsafe {
24572                b.launch(cfg)?;
24573            }
24574            return Ok(());
24575        }
24576        let f = self.func("qmatvec_q8_0_b4_rp_t");
24577        let cfg = LaunchConfig {
24578            grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
24579            block_dim: (32, ROWS_PER_BLOCK, 1),
24580            shared_mem_bytes: 0,
24581        };
24582        let __s_b = self.gpu.stream();
24583        let mut b = __s_b.launch_builder(&f);
24584        b.arg(w[0])
24585            .arg(w[1])
24586            .arg(w[2])
24587            .arg(w[3])
24588            .arg(aq)
24589            .arg(ad)
24590            .arg(y)
24591            .arg(&bc)
24592            .arg(&of);
24593        unsafe {
24594            b.launch(cfg)?;
24595        }
24596        Ok(())
24597    }
24598
24599    /// MEMRA_W8_VIEW: the q8_0 mirror for a bf16 GEMV whose weight is a ROW-RANGE VIEW.
24600    /// `MEMRA_W8_HYBRID` hangs off `matvec_bf16_into`, and the two split decode paths pinned in
24601    /// the step37 serving env send only their HI half there: HEAD_SPLIT runs
24602    /// `rank1.matvec_bf16_into(head_hi)` beside `e.matvec_bf16_view_into(head_lo)`, and
24603    /// SHEXP_OVERLAP does the same with the shared-expert down rows. The view launcher had no
24604    /// mirror, so the lo half kept streaming 2 B/w while its twin ran at 1.0625, and because the
24605    /// halves execute CONCURRENTLY on the two cards the critical path is the SLOW half.
24606    /// NUMERIC CLASS: identical to the rest of `MEMRA_STEP_TP_W8`, so it carries that argmax
24607    /// acceptance and that maxdiff class, not a new one. Default OFF until measured.
24608    #[allow(clippy::map_entry)] // allow: the init bodies are fallible (`?`); Entry::or_insert_with cannot propagate errors
24609    fn matvec_bf16_view_via_q8_mirror(
24610        &self,
24611        data: &cudarc::driver::CudaView<'_, u8>,
24612        x: &CudaSlice<f32>,
24613        y: &mut CudaSlice<f32>,
24614        in_f: usize,
24615        out_f: usize,
24616    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
24617        use cudarc::driver::DevicePtr;
24618        let key = {
24619            let s = self.gpu.stream();
24620            let (p, _g) = data.device_ptr(&s);
24621            (p, in_f as u32, out_f as u32)
24622        };
24623        {
24624            let mut mirrors = self
24625                .w8_mirrors
24626                .lock()
24627                .map_err(|_| "w8 mirror map is poisoned")?;
24628            if !mirrors.contains_key(&key) {
24629                let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
24630                self.encode_q8_0_from_bf16_view(data, &mut interleaved, in_f, out_f)?;
24631                let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
24632                mirrors.insert(key, planar);
24633                // Unconditional, once per distinct shape: a door with no announce cannot be read
24634                // in BOTH directions, and this lane was already burned once by a sweep that
24635                // inferred "never engages" from a log line that did not exist in the tree.
24636                eprintln!("[w8-view] mirror built in_f={in_f} out_f={out_f}");
24637            }
24638        }
24639        let nblk = in_f / 32;
24640        {
24641            let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24642            if !act.contains_key(&in_f) {
24643                let aq = self.alloc_uninit::<i8>(in_f)?;
24644                let ad = self.alloc_uninit::<f32>(nblk)?;
24645                act.insert(in_f, (aq, ad));
24646            }
24647            let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
24648            self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
24649        }
24650        let mirrors = self
24651            .w8_mirrors
24652            .lock()
24653            .map_err(|_| "w8 mirror map is poisoned")?;
24654        let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24655        let mirror = mirrors.get(&key).expect("built above");
24656        let (aq, ad) = act.get(&in_f).expect("built above");
24657        self.qmatvec_mmvq_into(
24658            mirror,
24659            aq,
24660            ad,
24661            1,
24662            in_f,
24663            out_f,
24664            QT_Q8_0,
24665            Self::q8_0_row_bytes(in_f),
24666            1.0,
24667            true,
24668            y,
24669        )?;
24670        Ok(Some(()))
24671    }
24672
24673    pub fn matvec_bf16_into(
24674        &self,
24675        data: &CudaSlice<u8>,
24676        x: &CudaSlice<f32>,
24677        y: &mut CudaSlice<f32>,
24678        in_f: usize,
24679        out_f: usize,
24680    ) -> Result<(), Box<dyn std::error::Error>> {
24681        if data.len() != in_f * out_f * 2
24682            || x.len() < in_f
24683            || !in_f.is_multiple_of(8)
24684            || y.len() < out_f
24685        {
24686            return Err(format!(
24687                "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24688                data.len(),
24689                x.len(),
24690                y.len()
24691            )
24692            .into());
24693        }
24694        // MEMRA_STEP_TP_W8, hybrid half: route this GEMV through a q8_0 mirror of the same
24695        // weight. Covers exactly the bf16 GEMVs that are NOT in a TP resident bank — the LM
24696        // head (324.4 -> 163.7 us measured), the shared-expert down rows (13.0 -> 5.6 us) and
24697        // the dense-FFN layers. Same numeric class as the QKV/o_proj arms (int8 dp4a with
24698        // per-32 scales), so it rides the same argmax acceptance; the bf16 slab stays resident
24699        // for prefill. The mirror builds on first use and is keyed by the slab's pointer.
24700        if step_tp_w8_on()
24701            && w8_hybrid_on()
24702            && in_f.is_multiple_of(32)
24703            && out_f >= 64
24704            && let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)?
24705        {
24706            return Ok(());
24707        }
24708        // MEMRA_DOWN_X4=1 (short-row shapes, in_f<=2048): four sequential rows per
24709        // block, exact f32acc per-row program — cures the 1-iteration latency
24710        // starvation (shexp down measured 420GB/s at in_f=1280).
24711        static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24712        let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
24713            && in_f <= 2048;
24714        if x4 {
24715            let f = self.func("matvec_bf16_f32acc_x4");
24716            let cfg = LaunchConfig {
24717                grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
24718                block_dim: (mmv_block(), 1, 1),
24719                shared_mem_bytes: 0,
24720            };
24721            let (ini, outi) = (in_f as i32, out_f as i32);
24722            let __s_b = self.gpu.stream();
24723            let mut b = __s_b.launch_builder(&f);
24724            b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
24725            unsafe {
24726                b.launch(cfg)?;
24727            }
24728            return Ok(());
24729        }
24730        let f = self.func("matvec_bf16_f32acc");
24731        let cfg = LaunchConfig {
24732            grid_dim: (out_f as u32, 1, 1),
24733            block_dim: (mmv_block(), 1, 1),
24734            shared_mem_bytes: 0,
24735        };
24736        let ini = in_f as i32;
24737        let __s_b = self.gpu.stream();
24738        let mut b = __s_b.launch_builder(&f);
24739        b.arg(data).arg(x).arg(y).arg(&ini);
24740        unsafe {
24741            b.launch(cfg)?;
24742        }
24743        Ok(())
24744    }
24745
24746    /// BF16 matvec over activation/output views. Automatic TP4 attention keeps each rank's
24747    /// O-projection input and canonical partial inside persistent slabs, so copying either view
24748    /// into a temporary allocation would give back the bandwidth and allocator win that TP is
24749    /// meant to provide.
24750    pub fn matvec_bf16_views_into(
24751        &self,
24752        data: &CudaSlice<u8>,
24753        x: &cudarc::driver::CudaView<'_, f32>,
24754        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
24755        in_f: usize,
24756        out_f: usize,
24757    ) -> Result<(), Box<dyn std::error::Error>> {
24758        if data.len() != in_f * out_f * 2
24759            || x.len() < in_f
24760            || !in_f.is_multiple_of(8)
24761            || y.len() < out_f
24762        {
24763            return Err(format!(
24764                "matvec_bf16_views_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24765                data.len(),
24766                x.len(),
24767                y.len()
24768            )
24769            .into());
24770        }
24771        let f = self.func("matvec_bf16_f32acc");
24772        let cfg = LaunchConfig {
24773            grid_dim: (out_f as u32, 1, 1),
24774            block_dim: (mmv_block(), 1, 1),
24775            shared_mem_bytes: 0,
24776        };
24777        let ini = in_f as i32;
24778        let __s_b = self.gpu.stream();
24779        let mut b = __s_b.launch_builder(&f);
24780        b.arg(data).arg(x).arg(y).arg(&ini);
24781        unsafe {
24782            b.launch(cfg)?;
24783        }
24784        Ok(())
24785    }
24786
24787    /// `matvec_bf16_into` over a WEIGHT VIEW (row-range slice of a bf16 tensor): the head-split
24788    /// door feeds each device its half of the lm-head rows. Same kernel, same per-row program.
24789    pub fn matvec_bf16_view_into(
24790        &self,
24791        data: &cudarc::driver::CudaView<'_, u8>,
24792        x: &CudaSlice<f32>,
24793        y: &mut CudaSlice<f32>,
24794        in_f: usize,
24795        out_f: usize,
24796    ) -> Result<(), Box<dyn std::error::Error>> {
24797        if data.len() != in_f * out_f * 2
24798            || x.len() < in_f
24799            || !in_f.is_multiple_of(8)
24800            || y.len() < out_f
24801        {
24802            return Err(format!(
24803                "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24804                data.len(),
24805                x.len(),
24806                y.len()
24807            )
24808            .into());
24809        }
24810        if w8_view_on()
24811            && step_tp_w8_on()
24812            && w8_hybrid_on()
24813            && in_f.is_multiple_of(32)
24814            && out_f >= 64
24815            && let Some(()) = self.matvec_bf16_view_via_q8_mirror(data, x, y, in_f, out_f)?
24816        {
24817            return Ok(());
24818        }
24819        let f = self.func("matvec_bf16_f32acc");
24820        let cfg = LaunchConfig {
24821            grid_dim: (out_f as u32, 1, 1),
24822            block_dim: (mmv_block(), 1, 1),
24823            shared_mem_bytes: 0,
24824        };
24825        let ini = in_f as i32;
24826        let __s_b = self.gpu.stream();
24827        let mut b = __s_b.launch_builder(&f);
24828        b.arg(data).arg(x).arg(y).arg(&ini);
24829        unsafe {
24830            b.launch(cfg)?;
24831        }
24832        Ok(())
24833    }
24834
24835    /// `matvec_bf16_into` with a RAW u64 output pointer (UVA — the dev1-shexp down row
24836    /// lands root-resident over P2P). Same kernel, same per-row program: bit-identical.
24837    pub fn matvec_bf16_raw_out(
24838        &self,
24839        w: &CudaSlice<u8>,
24840        x: &CudaSlice<f32>,
24841        y_raw: u64,
24842        in_f: usize,
24843        out_f: usize,
24844    ) -> Result<(), Box<dyn std::error::Error>> {
24845        if w.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) || y_raw == 0 {
24846            return Err("matvec_bf16_raw_out geometry".into());
24847        }
24848        let f = self.func("matvec_bf16_f32acc");
24849        let cfg = LaunchConfig {
24850            grid_dim: (out_f as u32, 1, 1),
24851            block_dim: (mmv_block(), 1, 1),
24852            shared_mem_bytes: 0,
24853        };
24854        let ini = in_f as i32;
24855        let __s_b = self.gpu.stream();
24856        let mut b = __s_b.launch_builder(&f);
24857        b.arg(w).arg(x).arg(&y_raw).arg(&ini);
24858        unsafe {
24859            b.launch(cfg)?;
24860        }
24861        Ok(())
24862    }
24863
24864    /// MOE TAIL FUSION M1: dst = (a + b) + sh*scale[0] in one launch (sh/scale as RAW
24865    /// UVA pointers so the caller passes persistent-static rows without holding locks).
24866    /// Exact per-element sequence of the split add + add_scaled_rows pair.
24867    pub fn add3_raw(
24868        &self,
24869        a: &CudaSlice<f32>,
24870        b: &CudaSlice<f32>,
24871        sh_raw: u64,
24872        scale_raw: u64,
24873        dst: &mut CudaSlice<f32>,
24874        n: usize,
24875    ) -> Result<(), Box<dyn std::error::Error>> {
24876        if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
24877            return Err("add3_raw geometry".into());
24878        }
24879        let f = self.func("add3_f32");
24880        let cfg = LaunchConfig {
24881            grid_dim: ((n as u32).div_ceil(256), 1, 1),
24882            block_dim: (256, 1, 1),
24883            shared_mem_bytes: 0,
24884        };
24885        let ni = n as i32;
24886        let __s_b = self.gpu.stream();
24887        let mut bld = __s_b.launch_builder(&f);
24888        bld.arg(a)
24889            .arg(b)
24890            .arg(&sh_raw)
24891            .arg(&scale_raw)
24892            .arg(dst)
24893            .arg(&ni);
24894        unsafe {
24895            bld.launch(cfg)?;
24896        }
24897        Ok(())
24898    }
24899
24900    /// FUSION #2e: shexp down matvec + scaled accumulate (dst[r] += dot_r * scale[0]),
24901    /// one launch replacing matvec_bf16_into + the ownership copy + add_scaled_rows.
24902    pub fn matvec_bf16_down_addscale_into(
24903        &self,
24904        w: &CudaSlice<u8>,
24905        x: &CudaSlice<f32>,
24906        scale: &CudaSlice<f32>,
24907        dst: &mut CudaSlice<f32>,
24908        in_f: usize,
24909        out_f: usize,
24910    ) -> Result<(), Box<dyn std::error::Error>> {
24911        if w.len() != in_f * out_f * 2
24912            || x.len() < in_f
24913            || !in_f.is_multiple_of(8)
24914            || dst.len() < out_f
24915            || scale.is_empty()
24916        {
24917            return Err("matvec_bf16_down_addscale geometry".into());
24918        }
24919        let f = self.func("matvec_bf16_down_addscale");
24920        let cfg = LaunchConfig {
24921            grid_dim: (out_f as u32, 1, 1),
24922            block_dim: (mmv_block(), 1, 1),
24923            shared_mem_bytes: 0,
24924        };
24925        let ini = in_f as i32;
24926        let __s_b = self.gpu.stream();
24927        let mut b = __s_b.launch_builder(&f);
24928        b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
24929        unsafe {
24930            b.launch(cfg)?;
24931        }
24932        Ok(())
24933    }
24934
24935    /// FUSION #2b: shexp dual matvec + SwiGLU act, one launch (bit-identical to
24936    /// matvec_bf16_dual_into + ffn_act_lim at gs=us=1; limit=None takes plain silu).
24937    /// T-ROW twin of `matvec_bf16_dual_silu_into` (per-row program identical).
24938    #[allow(clippy::too_many_arguments)]
24939    pub fn matvec_bf16_dual_silu_rows_into(
24940        &self,
24941        wg: &CudaSlice<u8>,
24942        wu: &CudaSlice<u8>,
24943        x: &CudaSlice<f32>,
24944        act: &mut CudaSlice<f32>,
24945        in_f: usize,
24946        out_f: usize,
24947        limit: Option<f32>,
24948        t: usize,
24949    ) -> Result<(), Box<dyn std::error::Error>> {
24950        if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
24951            return Err("matvec_bf16_dual_silu_rows geometry".into());
24952        }
24953        let f = self.func("matvec_bf16_dual_silu_rows");
24954        let cfg = LaunchConfig {
24955            grid_dim: (out_f as u32, t as u32, 1),
24956            block_dim: (mmv_block(), 1, 1),
24957            shared_mem_bytes: 0,
24958        };
24959        let (ini, outi) = (in_f as i32, out_f as i32);
24960        let lim = limit.unwrap_or(0.0);
24961        let __s_b = self.gpu.stream();
24962        let mut b = __s_b.launch_builder(&f);
24963        b.arg(wg)
24964            .arg(wu)
24965            .arg(x)
24966            .arg(&mut *act)
24967            .arg(&ini)
24968            .arg(&outi)
24969            .arg(&lim);
24970        unsafe {
24971            b.launch(cfg)?;
24972        }
24973        Ok(())
24974    }
24975
24976    /// T-ROW twin of the bf16 f32acc-x4 matvec (per-row program identical).
24977    pub fn matvec_bf16_rows_into(
24978        &self,
24979        w: &CudaSlice<u8>,
24980        x: &CudaSlice<f32>,
24981        y: &mut CudaSlice<f32>,
24982        in_f: usize,
24983        out_f: usize,
24984        t: usize,
24985    ) -> Result<(), Box<dyn std::error::Error>> {
24986        if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || !in_f.is_multiple_of(8)
24987        {
24988            return Err("matvec_bf16_rows geometry".into());
24989        }
24990        // MEMRA_STEP_TP_W8 + MEMRA_W8_HYBRID, t > 1: the VERIFY walk's shexp/dense rows land
24991        // here too (`matvec_bf16_f32acc_x4_rows` was 78 launches/round at 56.5 us in a spec
24992        // capture, ~162 ms of GPU over 37 rounds), and the t==1 gate below skipped them. The
24993        // t-column q8 kernel is bit-identical to t single-row calls.
24994        if (2..=32).contains(&t)
24995            && step_tp_w8_on()
24996            && w8_hybrid_on()
24997            && in_f.is_multiple_of(32)
24998            && out_f >= 64
24999            && let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)?
25000        {
25001            return Ok(());
25002        }
25003        // MEMRA_STEP_TP_W8: the LM head reaches the device HERE, not through
25004        // matvec_bf16_into — the W8 trace showed the hybrid half building mirrors only for
25005        // in_f=1280 out_f=4096 (the shared-expert down rows, which SHEXP_OVERLAP already
25006        // hides, hence its +0.1%). Route the t=1 decode row through the q8 mirror; wider t
25007        // (the verify walk) keeps bf16 so the prefill class is untouched.
25008        if t == 1
25009            && step_tp_w8_on()
25010            && w8_hybrid_on()
25011            && in_f.is_multiple_of(32)
25012            && out_f >= 64
25013            && let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)?
25014        {
25015            return Ok(());
25016        }
25017        // MEMRA_GLM5_W8, t in 2..=32: the glm5 verify-rows walk's KDA/MLA projections route
25018        // through the SAME t-column q8 mirror the step37 hybrid arm above uses (independent
25019        // door, independent predicate — the owner wants this receipted on its own, not folded
25020        // into the step37 lane). Bit-identical to t single-row q8_0 mirror calls by the
25021        // mirror's own construction (matvec_bf16_via_q8_mirror_t's contract).
25022        if (2..=32).contains(&t)
25023            && glm5_w8_on()
25024            && in_f.is_multiple_of(32)
25025            && out_f >= 64
25026            && let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)?
25027        {
25028            if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25029                eprintln!(
25030                    "[glm5-w8] engaged t={t} in_f={in_f} out_f={out_f} (q8_0 mirror, \
25031                     MEMRA_GLM5_W8=1)"
25032                );
25033            }
25034            return Ok(());
25035        }
25036        // MEMRA_GLM5_W8, t == 1: the decode-tier q8 mirror for the glm5_next KDA/MLA trunk.
25037        // Same building block MEMRA_STEP_TP_W8's hybrid half calls two blocks up
25038        // (matvec_bf16_via_q8_mirror); independent door, independent receipts.
25039        if t == 1
25040            && glm5_w8_on()
25041            && in_f.is_multiple_of(32)
25042            && out_f >= 64
25043            && let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)?
25044        {
25045            if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25046                eprintln!(
25047                    "[glm5-w8] engaged t=1 in_f={in_f} out_f={out_f} (q8_0 mirror, \
25048                     MEMRA_GLM5_W8=1)"
25049                );
25050            }
25051            return Ok(());
25052        }
25053        // MEMRA_BF16_TCOLS_WIDE (lane/glm5-matvec door T, default ON since 2026-08-31): t=2..=16 rides the
25054        // weight-once t-column class instead of the grid.y=t per-token weight re-read below.
25055        // Placed AFTER the W8-mirror intercepts (their precedence unchanged). Bit-identical
25056        // per (row, token) to the _rows kernel by the tcols class's standing construction
25057        // (order-pinned per-token chains + the identical red[256] tree); the motivating call
25058        // is the DFlash2 drafter's t=15 block-head matmul, which re-read the 1.269 GB lm
25059        // head 15x per spec round. Rollback seam: unset or =0 falls through unchanged.
25060        if (2..=16).contains(&t) && bf16_tcols_wide_on() {
25061            if BF16_TCOLS_WIDE_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25062                eprintln!(
25063                    "[bf16-tcols-wide] engaged: t={t} in_f={in_f} out_f={out_f} rides the \
25064                     weight-once tcols class (MEMRA_BF16_TCOLS_WIDE=1)"
25065                );
25066            }
25067            if t <= 8 {
25068                return self.matvec_bf16_tcols_into(w, x, y, in_f, out_f, t);
25069            }
25070            return self.matvec_bf16_tcols16_into(w, x, y, in_f, out_f, t);
25071        }
25072        // MEMRA_B200_GEMV_V2 (lane/b200-gemv-hbm-20260902): the HBM-speed rewrite. Same
25073        // arithmetic as the shipped kernel, rescheduled for bytes-in-flight (8 rows per block
25074        // accumulated concurrently on one activation load, 10 independent 16 B loads issued
25075        // before the first fma, one barrier chain per block instead of four). BIT-IDENTICAL per
25076        // (row, token) at ksplit=1, which is what every GLM-5.3 decode shape picks; a shape too
25077        // narrow to cover two CTA waves takes the named `bf16_gemv_v2_splitk` class instead.
25078        // Placed BEFORE the cuBLASLt reference door so that with both set the memra kernel wins
25079        // (the LT door is an instrument, never the product).
25080        if b200_gemv_v2_on() {
25081            let ksplit = self.gemv_v2_ksplit(in_f, out_f, t);
25082            // v3 (level 2) is the cp.async-staged walk. It has no split-K twin and needs its
25083            // 36 KB of dynamic smem to fit the 48 KB default cap, so it declines per call
25084            // rather than per process and v2 takes those shapes.
25085            let v3 = b200_gemv_v2_level() >= 2 && ksplit == 1 && gemv_v3_fits();
25086            if GEMV_V2_BF16_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25087                let arm = if v3 { "v3 (cp.async staged)" } else { "v2" };
25088                eprintln!(
25089                    "[b200-gemv-v2] engaged arm={arm} t={t} in_f={in_f} out_f={out_f} \
25090                     ksplit={ksplit} (MEMRA_B200_GEMV_V2={})",
25091                    b200_gemv_v2_level()
25092                );
25093            }
25094            if v3 {
25095                return self.matvec_bf16_v3_raw(w, x, y, in_f, out_f, t);
25096            }
25097            return self.matvec_bf16_v2_raw(w, x, y, in_f, out_f, t, ksplit);
25098        }
25099        // MEMRA_B200_BF16_GEMV_LT (lane/b200-gemv-hbm-20260902): the cuBLASLt REFERENCE door.
25100        // Routes this row matvec through the vendor library's m=t bf16 GEMV so a box can
25101        // measure what a tuned library reaches on these bytes on sm_100a. NAMED NUMERIC CLASS
25102        // `bf16_gemv_lt` (activation cast to bf16 + library summation order), default OFF,
25103        // reference only — see `b200_bf16_gemv_lt_on`. Placed AFTER the W8-mirror and
25104        // tcols intercepts so their precedence is unchanged, and a cuBLASLt decline falls
25105        // through to the shipped kernel below.
25106        if b200_bf16_gemv_lt_on() && self.bf16_gemv_lt_into(w, x, y, in_f, out_f, t)? {
25107            return Ok(());
25108        }
25109        // MEMRA_B200_MATVEC_ARM occupancy arm (lane/b200-matvec-occupancy-20260902): the
25110        // software-pipelined `_pf` twin double-buffers the K-loop's weight/activation loads
25111        // (next iteration's loads issue before the current iteration's fma chain runs) to
25112        // hide DRAM latency on B200's narrower SM/bandwidth shape. Same grid/block/reduction
25113        // tree, same per-thread fma order for the same i -> bit-identical per (row,token).
25114        // Default OFF; sm_120a keeps the shipped kernel unconditionally.
25115        let kname = if b200_matvec_arm_on() {
25116            "matvec_bf16_f32acc_x4_rows_pf"
25117        } else {
25118            "matvec_bf16_f32acc_x4_rows"
25119        };
25120        let f = self.func(kname);
25121        let cfg = LaunchConfig {
25122            grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
25123            block_dim: (mmv_block(), 1, 1),
25124            shared_mem_bytes: 0,
25125        };
25126        let (ini, outi) = (in_f as i32, out_f as i32);
25127        let __s_b = self.gpu.stream();
25128        let mut b = __s_b.launch_builder(&f);
25129        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
25130        unsafe {
25131            b.launch(cfg)?;
25132        }
25133        Ok(())
25134    }
25135
25136    /// One aligned K-range partial of the t=1 BF16 row matvec. This is the row-parallel TP
25137    /// building block: every rank computes all output rows over a disjoint K range from its
25138    /// compact local activation shard, then the persistent replicated-row collective sums the
25139    /// partials. The kernel preserves the unsplit program inside each range; the cross-rank
25140    /// association is separately gated.
25141    #[allow(clippy::too_many_arguments)]
25142    pub fn matvec_bf16_col_range_into(
25143        &self,
25144        w: &CudaSlice<u8>,
25145        x: &CudaSlice<f32>,
25146        y: &mut CudaSlice<f32>,
25147        in_f: usize,
25148        out_f: usize,
25149        k_start: usize,
25150        k_len: usize,
25151    ) -> Result<(), Box<dyn std::error::Error>> {
25152        let weight_bytes = in_f
25153            .checked_mul(out_f)
25154            .and_then(|elements| elements.checked_mul(2))
25155            .ok_or("BF16 column-range matvec geometry")?;
25156        let k_end = k_start
25157            .checked_add(k_len)
25158            .ok_or("BF16 column-range matvec geometry")?;
25159        if w.len() < weight_bytes
25160            || x.len() < k_len
25161            || y.len() < out_f
25162            || in_f > i32::MAX as usize
25163            || out_f > i32::MAX as usize
25164            || k_start > i32::MAX as usize
25165            || k_len > i32::MAX as usize
25166            || in_f == 0
25167            || out_f == 0
25168            || k_len == 0
25169            || !in_f.is_multiple_of(8)
25170            || !k_start.is_multiple_of(8)
25171            || !k_len.is_multiple_of(8)
25172            || k_end > in_f
25173        {
25174            return Err("BF16 column-range matvec geometry".into());
25175        }
25176        let function = self.func("matvec_bf16_f32acc_x4_range");
25177        let config = LaunchConfig {
25178            grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
25179            block_dim: (mmv_block(), 1, 1),
25180            shared_mem_bytes: 0,
25181        };
25182        let (in_f, out_f, k_start, k_len) =
25183            (in_f as i32, out_f as i32, k_start as i32, k_len as i32);
25184        let stream = self.gpu.stream();
25185        let mut launch = stream.launch_builder(&function);
25186        launch
25187            .arg(w)
25188            .arg(x)
25189            .arg(y)
25190            .arg(&in_f)
25191            .arg(&out_f)
25192            .arg(&k_start)
25193            .arg(&k_len);
25194        unsafe {
25195            launch.launch(config)?;
25196        }
25197        Ok(())
25198    }
25199
25200    /// T-COLUMN twin of the bf16 rows matvec (lane/glm5-verify-batch, the
25201    /// varlen-batched-cores pattern): one block owns 4 output rows for ALL t tokens, so the
25202    /// weight pack is read ONCE and reused across tokens — vs the `_rows` twin's grid.y=t
25203    /// per-token weight re-read. Per-(row,token) BIT-IDENTICAL to the t=1 program by
25204    /// construction (order-pinned single-chain accumulators, identical shared-tree reduce
25205    /// per token — LAW:vl-bit-identity-order-pinning); the `glm5_verify_batch_gpu` tcols
25206    /// bit-gate holds it. t is bounded by the kernel's MEMRA_BF16_TCOLS_MAX = 8.
25207    pub fn matvec_bf16_tcols_into(
25208        &self,
25209        w: &CudaSlice<u8>,
25210        x: &CudaSlice<f32>,
25211        y: &mut CudaSlice<f32>,
25212        in_f: usize,
25213        out_f: usize,
25214        t: usize,
25215    ) -> Result<(), Box<dyn std::error::Error>> {
25216        if x.len() < t * in_f
25217            || y.len() < t * out_f
25218            || !(2..=8).contains(&t)
25219            || !in_f.is_multiple_of(8)
25220        {
25221            return Err("matvec_bf16_tcols geometry".into());
25222        }
25223        // MEMRA_BF16_TCOLS_X1 (lane/glm5-matvec door X, default ON since 2026-08-31): one row per block
25224        // (grid.x = out_f) — 4x the wave count on the ~one-wave trunk grids (census: same
25225        // kernel runs 59% of peak at 512..2048 blocks, 80% at 38720). Per-row body and
25226        // reduce tree verbatim — bit-identical per (row, token). Rollback: unset or =0.
25227        // MEMRA_BF16_TCOLS_RED_FUSED (lane/glm5-door-r door R, default OFF): the chosen grid
25228        // form takes its `_rf` fused-reduce-tail twin — one barrier sequence shared by the t
25229        // columns plus intra-warp shuffles at the identical pairing (9t -> 3 barriers per
25230        // block). Composes with door X (grid choice first, tail twin second). Requires a
25231        // power-of-two block (the fused tail must pass exactly through s=32); any other
25232        // MEMRA_MMV_BLOCK falls through to the standing tree. Rollback: unset or =0.
25233        let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
25234        if rf
25235            && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
25236                == 0
25237        {
25238            eprintln!(
25239                "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
25240                 shared across the t token columns + intra-warp shuffles at the identical \
25241                 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
25242            );
25243        }
25244        let x1 = bf16_tcols_x1_on();
25245        let (fname, grid_x) = match (x1, rf) {
25246            (true, true) => ("matvec_bf16_f32acc_x1_tcols_rf", out_f as u32),
25247            (true, false) => ("matvec_bf16_f32acc_x1_tcols", out_f as u32),
25248            (false, true) => ("matvec_bf16_f32acc_x4_tcols_rf", out_f.div_ceil(4) as u32),
25249            (false, false) => ("matvec_bf16_f32acc_x4_tcols", out_f.div_ceil(4) as u32),
25250        };
25251        if x1 && BF16_TCOLS_X1_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25252            eprintln!(
25253                "[bf16-tcols-x1] engaged: one-row-per-block tcols grid \
25254                 (MEMRA_BF16_TCOLS_X1=1)"
25255            );
25256        }
25257        let f = self.func(fname);
25258        let cfg = LaunchConfig {
25259            grid_dim: (grid_x, 1, 1),
25260            block_dim: (mmv_block(), 1, 1),
25261            shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
25262        };
25263        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
25264        let __s_b = self.gpu.stream();
25265        let mut b = __s_b.launch_builder(&f);
25266        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
25267        unsafe {
25268            b.launch(cfg)?;
25269        }
25270        Ok(())
25271    }
25272
25273    /// WIDE-T twin of [`Self::matvec_bf16_tcols_into`] (lane/glm5-matvec door T,
25274    /// `MEMRA_BF16_TCOLS_WIDE`): t = 9..=16 through the SEPARATE `..._tcols16` kernel — its
25275    /// acc[16] register footprint never touches the priced t<=8 class (the qmatvec `_tw32`
25276    /// acc-sizing lesson). Bit-identical per (row, token) to the t=1 program by the same
25277    /// order-pinned construction; gated by `glm5_matvec_doors_gpu`.
25278    pub fn matvec_bf16_tcols16_into(
25279        &self,
25280        w: &CudaSlice<u8>,
25281        x: &CudaSlice<f32>,
25282        y: &mut CudaSlice<f32>,
25283        in_f: usize,
25284        out_f: usize,
25285        t: usize,
25286    ) -> Result<(), Box<dyn std::error::Error>> {
25287        if x.len() < t * in_f
25288            || y.len() < t * out_f
25289            || !(9..=16).contains(&t)
25290            || !in_f.is_multiple_of(8)
25291        {
25292            return Err("matvec_bf16_tcols16 geometry".into());
25293        }
25294        // MEMRA_BF16_TCOLS_RED_FUSED (lane/glm5-door-r door R, default OFF): the wide-t twin
25295        // takes its `_rf` fused tail too — the drafter head's t=15 is the extreme case (135
25296        // barriers -> 6 per block). Same power-of-two block guard as the t<=8 dispatch.
25297        let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
25298        if rf
25299            && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
25300                == 0
25301        {
25302            eprintln!(
25303                "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
25304                 shared across the t token columns + intra-warp shuffles at the identical \
25305                 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
25306            );
25307        }
25308        let f = self.func(if rf {
25309            "matvec_bf16_f32acc_x4_tcols16_rf"
25310        } else {
25311            "matvec_bf16_f32acc_x4_tcols16"
25312        });
25313        let cfg = LaunchConfig {
25314            grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
25315            block_dim: (mmv_block(), 1, 1),
25316            shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
25317        };
25318        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
25319        let __s_b = self.gpu.stream();
25320        let mut b = __s_b.launch_builder(&f);
25321        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
25322        unsafe {
25323            b.launch(cfg)?;
25324        }
25325        Ok(())
25326    }
25327
25328    /// GATE-ONLY launcher for door R arms no route dispatches (`glm5_matvec_doors_gpu`):
25329    /// the shifted-pairing RED twin (`matvec_bf16_f32acc_x1_tcols_rf_redshift`, the arm that
25330    /// proves the bit bar can see an association change) and the `_rf` twins at t=1 (the
25331    /// routed launchers refuse t<2; the door-R bar covers t=1..=16, so the degenerate
25332    /// column-loop bounds are gated here). `kernel` is an ALLOWLIST, not a name proxy.
25333    #[allow(clippy::too_many_arguments)]
25334    pub fn matvec_bf16_tcols_gate_kernel_into(
25335        &self,
25336        kernel: &str,
25337        w: &CudaSlice<u8>,
25338        x: &CudaSlice<f32>,
25339        y: &mut CudaSlice<f32>,
25340        in_f: usize,
25341        out_f: usize,
25342        t: usize,
25343    ) -> Result<(), Box<dyn std::error::Error>> {
25344        let (grid_x, t_max) = match kernel {
25345            "matvec_bf16_f32acc_x1_tcols_rf" | "matvec_bf16_f32acc_x1_tcols_rf_redshift" => {
25346                (out_f as u32, 8usize)
25347            }
25348            "matvec_bf16_f32acc_x4_tcols_rf" => (out_f.div_ceil(4) as u32, 8usize),
25349            "matvec_bf16_f32acc_x4_tcols16_rf" => (out_f.div_ceil(4) as u32, 16usize),
25350            _ => return Err("matvec_bf16_tcols_gate_kernel_into: unknown kernel".into()),
25351        };
25352        if x.len() < t * in_f
25353            || y.len() < t * out_f
25354            || !(1..=t_max).contains(&t)
25355            || !in_f.is_multiple_of(8)
25356            || !mmv_block().is_power_of_two()
25357        {
25358            return Err("matvec_bf16_tcols_gate_kernel geometry".into());
25359        }
25360        let f = self.func(kernel);
25361        let cfg = LaunchConfig {
25362            grid_dim: (grid_x, 1, 1),
25363            block_dim: (mmv_block(), 1, 1),
25364            shared_mem_bytes: (t as u32) * mmv_block() * 4,
25365        };
25366        let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
25367        let __s_b = self.gpu.stream();
25368        let mut b = __s_b.launch_builder(&f);
25369        b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
25370        unsafe {
25371            b.launch(cfg)?;
25372        }
25373        Ok(())
25374    }
25375
25376    /// DECODE-EXACT matmul for the glm5 verify-batch walk (lane/glm5-verify-batch): the
25377    /// exact `matmul_decode_exact` dispatch with ONE addition — FloatBf16 weights at
25378    /// t=2..=8 under `MEMRA_BF16_MMV` ride the tcols twin above (weight read once for all
25379    /// t rows). Refused back to `matmul_decode_exact` whenever the t=1 decode chain would
25380    /// ride the W8 q8-mirror class instead of the bf16 rows kernel (the decode-parity law:
25381    /// the m>1 class must equal the m=1 class). Every class stays per-row bit-exact vs
25382    /// the t=1 chain; only the verify-batch walk calls this.
25383    pub fn matmul_rows_exact(
25384        &self,
25385        w: &crate::model::GpuTensor,
25386        x: &CudaSlice<f32>,
25387        m: usize,
25388    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25389        use crate::model::GpuTensor;
25390        // MEMRA_GLM5_W8: the glm5 verify-rows walk's KDA/MLA projections take the SAME q8_0
25391        // mirror the plain t=1/small-t decode arm uses in `matvec_bf16_rows_into` — placed
25392        // BEFORE the tcols check below so the door's own class (not the bf16 tcols class) wins
25393        // when it engages. Independent of MEMRA_STEP_TP_W8/MEMRA_W8_HYBRID.
25394        if let GpuTensor::FloatBf16 { data, .. } = w
25395            && (2..=32).contains(&m)
25396            && Self::bf16_mmv_on()
25397            && w.in_features().is_multiple_of(32)
25398            && w.out_features() >= 64
25399            && glm5_w8_on()
25400        {
25401            let (in_f, out_f) = (w.in_features(), w.out_features());
25402            let mut y = self.vws_uninit(m * out_f)?;
25403            if let Some(()) = self.matvec_bf16_via_q8_mirror_t(data, x, &mut y, in_f, out_f, m)? {
25404                if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25405                    eprintln!(
25406                        "[glm5-w8] engaged rows-exact m={m} in_f={in_f} out_f={out_f} \
25407                         (q8_0 mirror, MEMRA_GLM5_W8=1)"
25408                    );
25409                }
25410                return Ok(y);
25411            }
25412        }
25413        if let GpuTensor::FloatBf16 { data, .. } = w
25414            && (2..=8).contains(&m)
25415            && Self::bf16_mmv_on()
25416            && w.in_features().is_multiple_of(8)
25417            && !(step_tp_w8_on() && w8_hybrid_on())
25418            && !glm5_w8_on()
25419        {
25420            let (in_f, out_f) = (w.in_features(), w.out_features());
25421            // Door W: rows-exact is verify-walk-only by contract, so its y is a pooled
25422            // draw (vws_uninit == alloc_uninit with the door off).
25423            let mut y = self.vws_uninit(m * out_f)?;
25424            self.matvec_bf16_tcols_into(data, x, &mut y, in_f, out_f, m)?;
25425            return Ok(y);
25426        }
25427        self.matmul_decode_exact(w, x, m)
25428    }
25429
25430    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
25431    pub fn matvec_bf16_dual_silu_into(
25432        &self,
25433        wg: &CudaSlice<u8>,
25434        wu: &CudaSlice<u8>,
25435        x: &CudaSlice<f32>,
25436        act: &mut CudaSlice<f32>,
25437        in_f: usize,
25438        out_f: usize,
25439        limit: Option<f32>,
25440    ) -> Result<(), Box<dyn std::error::Error>> {
25441        if wg.len() != in_f * out_f * 2
25442            || wu.len() != in_f * out_f * 2
25443            || x.len() < in_f
25444            || !in_f.is_multiple_of(8)
25445            || act.len() < out_f
25446        {
25447            return Err("matvec_bf16_dual_silu geometry".into());
25448        }
25449        let f = self.func("matvec_bf16_dual_silu");
25450        let cfg = LaunchConfig {
25451            grid_dim: (out_f as u32, 1, 1),
25452            block_dim: (mmv_block(), 1, 1),
25453            shared_mem_bytes: 0,
25454        };
25455        let (ini, outi) = (in_f as i32, out_f as i32);
25456        let lim = limit.unwrap_or(0.0);
25457        let __s_b = self.gpu.stream();
25458        let mut b = __s_b.launch_builder(&f);
25459        b.arg(wg)
25460            .arg(wu)
25461            .arg(x)
25462            .arg(act)
25463            .arg(&ini)
25464            .arg(&outi)
25465            .arg(&lim);
25466        unsafe {
25467            b.launch(cfg)?;
25468        }
25469        Ok(())
25470    }
25471
25472    /// `matvec_bf16_dual_into` over WEIGHT VIEWS (row-range slices): the shexp row-split
25473    /// door feeds each device its half of the gate/up rows. Same kernel, same per-row program.
25474    #[allow(clippy::too_many_arguments)]
25475    pub fn matvec_bf16_dual_view_into(
25476        &self,
25477        wg: &cudarc::driver::CudaView<'_, u8>,
25478        wu: &cudarc::driver::CudaView<'_, u8>,
25479        x: &CudaSlice<f32>,
25480        yg: &mut CudaSlice<f32>,
25481        yu: &mut CudaSlice<f32>,
25482        in_f: usize,
25483        out_f: usize,
25484    ) -> Result<(), Box<dyn std::error::Error>> {
25485        if wg.len() != in_f * out_f * 2
25486            || wu.len() != in_f * out_f * 2
25487            || x.len() < in_f
25488            || !in_f.is_multiple_of(8)
25489            || yg.len() < out_f
25490            || yu.len() < out_f
25491        {
25492            return Err(format!(
25493                "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
25494                wg.len(),
25495                wu.len(),
25496                x.len()
25497            )
25498            .into());
25499        }
25500        let f = self.func("matvec_bf16_dual");
25501        let cfg = LaunchConfig {
25502            grid_dim: ((2 * out_f) as u32, 1, 1),
25503            block_dim: (mmv_block(), 1, 1),
25504            shared_mem_bytes: 0,
25505        };
25506        let (ini, outi) = (in_f as i32, out_f as i32);
25507        let __s_b = self.gpu.stream();
25508        let mut b = __s_b.launch_builder(&f);
25509        b.arg(wg)
25510            .arg(wu)
25511            .arg(x)
25512            .arg(yg)
25513            .arg(yu)
25514            .arg(&ini)
25515            .arg(&outi);
25516        unsafe {
25517            b.launch(cfg)?;
25518        }
25519        Ok(())
25520    }
25521
25522    /// `matvec_bf16_dual` writing into caller-owned outputs (persistent-workspace form).
25523    #[allow(clippy::too_many_arguments)]
25524    pub fn matvec_bf16_dual_into(
25525        &self,
25526        wg: &CudaSlice<u8>,
25527        wu: &CudaSlice<u8>,
25528        x: &CudaSlice<f32>,
25529        yg: &mut CudaSlice<f32>,
25530        yu: &mut CudaSlice<f32>,
25531        in_f: usize,
25532        out_f: usize,
25533    ) -> Result<(), Box<dyn std::error::Error>> {
25534        if wg.len() != in_f * out_f * 2
25535            || wu.len() != in_f * out_f * 2
25536            || x.len() < in_f
25537            || !in_f.is_multiple_of(8)
25538            || yg.len() < out_f
25539            || yu.len() < out_f
25540        {
25541            return Err(format!(
25542                "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
25543                wg.len(),
25544                wu.len(),
25545                x.len()
25546            )
25547            .into());
25548        }
25549        let f = self.func("matvec_bf16_dual");
25550        let cfg = LaunchConfig {
25551            grid_dim: ((2 * out_f) as u32, 1, 1),
25552            block_dim: (mmv_block(), 1, 1),
25553            shared_mem_bytes: 0,
25554        };
25555        let (ini, outi) = (in_f as i32, out_f as i32);
25556        let __s_b = self.gpu.stream();
25557        let mut b = __s_b.launch_builder(&f);
25558        b.arg(wg)
25559            .arg(wu)
25560            .arg(x)
25561            .arg(yg)
25562            .arg(yu)
25563            .arg(&ini)
25564            .arg(&outi);
25565        unsafe {
25566            b.launch(cfg)?;
25567        }
25568        Ok(())
25569    }
25570
25571    /// Dual bf16 matvec: gate/up (same shape) from one shared input in one launch. Per row
25572    /// bit-identical to two `matvec_bf16` launches. Returns (gate, up).
25573    #[allow(dead_code)] // allow: base form of the matvec_bf16_dual_* family; kept as the reference entry point
25574    pub(crate) fn matvec_bf16_dual(
25575        &self,
25576        wg: &CudaSlice<u8>,
25577        wu: &CudaSlice<u8>,
25578        x: &CudaSlice<f32>,
25579        in_f: usize,
25580        out_f: usize,
25581    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25582        if wg.len() != in_f * out_f * 2
25583            || wu.len() != in_f * out_f * 2
25584            || x.len() < in_f
25585            || !in_f.is_multiple_of(8)
25586        {
25587            return Err(format!(
25588                "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
25589                wg.len(),
25590                wu.len(),
25591                x.len()
25592            )
25593            .into());
25594        }
25595        let mut yg = self.alloc_uninit::<f32>(out_f)?;
25596        let mut yu = self.alloc_uninit::<f32>(out_f)?;
25597        let f = self.func("matvec_bf16_dual");
25598        let cfg = LaunchConfig {
25599            grid_dim: ((2 * out_f) as u32, 1, 1),
25600            block_dim: (mmv_block(), 1, 1),
25601            shared_mem_bytes: 0,
25602        };
25603        let (ini, outi) = (in_f as i32, out_f as i32);
25604        let __s_b = self.gpu.stream();
25605        let mut b = __s_b.launch_builder(&f);
25606        b.arg(wg)
25607            .arg(wu)
25608            .arg(x)
25609            .arg(&mut yg)
25610            .arg(&mut yu)
25611            .arg(&ini)
25612            .arg(&outi);
25613        unsafe {
25614            b.launch(cfg)?;
25615        }
25616        Ok((yg, yu))
25617    }
25618
25619    #[allow(clippy::too_many_arguments)]
25620    #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
25621    fn linear_bf16_chunked_inner(
25622        &self,
25623        x: &CudaSlice<f32>,
25624        data: &CudaSlice<u8>,
25625        m: usize,
25626        in_f: usize,
25627        out_f: usize,
25628        exact: bool,
25629        canonical_chunk_rows: Option<usize>,
25630    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25631        const CHUNK_BYTES: usize = 256 << 20;
25632        // canonical_chunk_rows callers are the step TP resident ranks whose cross-topology
25633        // equality program is defined in cuBLASLt chunk shapes — the door leaves them alone.
25634        if m == 1
25635            && !exact
25636            && canonical_chunk_rows.is_none()
25637            && in_f.is_multiple_of(8)
25638            && Self::bf16_mmv_on()
25639        {
25640            return self.matvec_bf16(data, x, in_f, out_f);
25641        }
25642        // MEMRA_PP_BF16: prefill on the RESIDENT bf16 bytes through cuBLASLt tensor cores.
25643        // Below this door the whole weight is dequanted to f32 and multiplied without tensor
25644        // cores — the step37 prime's 14x gap to vLLM. `exact` and canonical-chunk callers are
25645        // numerical programs with their own equality gates and are left alone.
25646        if m >= 16
25647            && !exact
25648            && canonical_chunk_rows.is_none()
25649            && data.len() == in_f * out_f * 2
25650            && crate::f16_ffi::pp_bf16_enabled()
25651        {
25652            // None = cuBLASLt declined this shape (it announced which one); fall through to the
25653            // f32 dequant GEMM below, which is always correct.
25654            if let Some(y) = self.bf16_tc_gemm(data, x, m, in_f, out_f)? {
25655                return Ok(y);
25656            }
25657        }
25658        let row_bytes = in_f
25659            .checked_mul(std::mem::size_of::<f32>())
25660            .ok_or("BF16 chunk row byte count overflow")?;
25661        if row_bytes == 0 || out_f == 0 {
25662            return Err("BF16 chunk dimensions must be nonzero".into());
25663        }
25664        let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
25665        let chunk_rows = match canonical_chunk_rows {
25666            Some(0) => {
25667                return Err("canonical BF16 chunk rows must be nonzero".into());
25668            }
25669            Some(rows) if rows > max_chunk_rows => {
25670                return Err(format!(
25671                    "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
25672                )
25673                .into());
25674            }
25675            Some(rows) if out_f % rows != 0 => {
25676                return Err(format!(
25677                    "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
25678                )
25679                .into());
25680            }
25681            Some(rows) => rows,
25682            None => max_chunk_rows,
25683        };
25684        if chunk_rows >= out_f {
25685            let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
25686            return if exact {
25687                self.linear_decode_exact(x, &wf32, m, in_f, out_f)
25688            } else {
25689                self.linear(x, &wf32, m, in_f, out_f)
25690            };
25691        }
25692        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
25693        let mut r0 = 0usize;
25694        while r0 < out_f {
25695            let rows = chunk_rows.min(out_f - r0);
25696            let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
25697            let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
25698            let yc = if exact {
25699                self.linear_decode_exact(x, &wf32, m, in_f, rows)?
25700            } else {
25701                self.linear(x, &wf32, m, in_f, rows)?
25702            };
25703            // scatter [m, rows] into y[m, out_f] at column offset r0 (m is tiny in decode/verify)
25704            for mi in 0..m {
25705                let src = yc.slice(mi * rows..(mi + 1) * rows);
25706                let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
25707                self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
25708            }
25709            r0 += rows;
25710        }
25711        Ok(y)
25712    }
25713
25714    /// Execute an already resident BF16 projection. This is the model-faithful substrate used by
25715    /// Step tensor-parallel correctness ranks; it preserves checkpoint bytes and the existing
25716    /// chunked BF16 numerical program instead of re-encoding the weight.
25717    pub fn linear_bf16_resident(
25718        &self,
25719        x: &CudaSlice<f32>,
25720        data: &CudaSlice<u8>,
25721        m: usize,
25722        in_f: usize,
25723        out_f: usize,
25724    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25725        if data.len() != in_f * out_f * 2 {
25726            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
25727        }
25728        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
25729    }
25730
25731    /// Execute a resident BF16 projection as fixed-width output-row chunks.
25732    ///
25733    /// Tensor-parallel ranks use this to give TP1/TP2/TP4/TP8 the same cuBLASLt problem shape
25734    /// for every checkpoint row. Callers must derive `canonical_chunk_rows` from the registered
25735    /// model topology rather than the active rank count.
25736    pub fn linear_bf16_resident_canonical_rows(
25737        &self,
25738        x: &CudaSlice<f32>,
25739        data: &CudaSlice<u8>,
25740        m: usize,
25741        in_f: usize,
25742        out_f: usize,
25743        canonical_chunk_rows: usize,
25744    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25745        if data.len() != in_f * out_f * 2 {
25746            return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
25747        }
25748        self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
25749    }
25750
25751    /// Execute a load-time F32 mirror with the same fixed output-row chunks as the BF16 path.
25752    ///
25753    /// Expanding the checkpoint bytes once changes residency, not arithmetic: every cuBLASLt
25754    /// call receives the same F32 values and problem shape as `linear_bf16_chunked`.
25755    pub fn linear_f32_resident_canonical_rows(
25756        &self,
25757        x: &CudaSlice<f32>,
25758        data: &CudaSlice<f32>,
25759        m: usize,
25760        in_f: usize,
25761        out_f: usize,
25762        canonical_chunk_rows: usize,
25763    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25764        self.linear_f32_resident_canonical_rows_inner(
25765            x,
25766            data,
25767            m,
25768            in_f,
25769            out_f,
25770            canonical_chunk_rows,
25771            false,
25772        )
25773    }
25774
25775    /// Execute fixed output-row chunks and assemble them with one strided placement per chunk.
25776    ///
25777    /// The projection shapes and values are identical to
25778    /// [`Self::linear_f32_resident_canonical_rows`]. Only the byte-preserving output layout step
25779    /// changes, replacing one device copy per token with one placement kernel per output chunk.
25780    pub fn linear_f32_resident_canonical_rows_strided(
25781        &self,
25782        x: &CudaSlice<f32>,
25783        data: &CudaSlice<f32>,
25784        m: usize,
25785        in_f: usize,
25786        out_f: usize,
25787        canonical_chunk_rows: usize,
25788    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25789        self.linear_f32_resident_canonical_rows_inner(
25790            x,
25791            data,
25792            m,
25793            in_f,
25794            out_f,
25795            canonical_chunk_rows,
25796            true,
25797        )
25798    }
25799
25800    #[allow(clippy::too_many_arguments)]
25801    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
25802    #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
25803    fn linear_f32_resident_canonical_rows_inner(
25804        &self,
25805        x: &CudaSlice<f32>,
25806        data: &CudaSlice<f32>,
25807        m: usize,
25808        in_f: usize,
25809        out_f: usize,
25810        canonical_chunk_rows: usize,
25811        strided_output: bool,
25812    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25813        if data.len() != in_f * out_f {
25814            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
25815        }
25816        if canonical_chunk_rows == 0
25817            || canonical_chunk_rows > out_f
25818            || out_f % canonical_chunk_rows != 0
25819        {
25820            return Err(format!(
25821                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
25822            )
25823            .into());
25824        }
25825        if canonical_chunk_rows == out_f {
25826            return self.linear(x, data, m, in_f, out_f);
25827        }
25828
25829        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
25830        let input = x.slice(0..x.len());
25831        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
25832            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
25833            if m == 1 {
25834                let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
25835                self.linear_device_into(
25836                    &input,
25837                    &weights,
25838                    &mut destination,
25839                    1,
25840                    in_f,
25841                    canonical_chunk_rows,
25842                )?;
25843                continue;
25844            }
25845            let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
25846            if strided_output {
25847                self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
25848            } else {
25849                for token in 0..m {
25850                    let source = chunk
25851                        .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
25852                    let mut destination =
25853                        y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
25854                    self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
25855                }
25856            }
25857        }
25858        Ok(y)
25859    }
25860
25861    /// One-token twin of `linear_f32_resident_canonical_rows` writing into a caller-owned
25862    /// output. Same cuBLASLt calls, values, and chunk order as the allocating variant at
25863    /// `m == 1`; only the output residency changes (persistent workspace instead of a fresh
25864    /// allocation per call). This is the projection substrate of the v2 Step TP decode driver.
25865    #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
25866    pub fn linear_f32_resident_canonical_rows_t1_into(
25867        &self,
25868        x: &CudaSlice<f32>,
25869        data: &CudaSlice<f32>,
25870        y: &mut CudaSlice<f32>,
25871        in_f: usize,
25872        out_f: usize,
25873        canonical_chunk_rows: usize,
25874    ) -> Result<(), Box<dyn std::error::Error>> {
25875        if data.len() != in_f * out_f {
25876            return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
25877        }
25878        if y.len() != out_f || x.len() != in_f {
25879            return Err(format!(
25880                "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
25881                x.len(),
25882                y.len()
25883            )
25884            .into());
25885        }
25886        if canonical_chunk_rows == 0
25887            || canonical_chunk_rows > out_f
25888            || out_f % canonical_chunk_rows != 0
25889        {
25890            return Err(format!(
25891                "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
25892            )
25893            .into());
25894        }
25895        let input = x.slice(0..x.len());
25896        for r0 in (0..out_f).step_by(canonical_chunk_rows) {
25897            let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
25898            let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
25899            self.linear_device_into(
25900                &input,
25901                &weights,
25902                &mut destination,
25903                1,
25904                in_f,
25905                canonical_chunk_rows,
25906            )?;
25907        }
25908        Ok(())
25909    }
25910
25911    /// One-token view-to-view linear into a caller-owned destination — the `linear` twin
25912    /// without the allocation, for workspace-resident operands.
25913    pub fn linear_t1_into(
25914        &self,
25915        x: &cudarc::driver::CudaView<'_, f32>,
25916        w: &cudarc::driver::CudaView<'_, f32>,
25917        y: &mut cudarc::driver::CudaViewMut<'_, f32>,
25918        in_f: usize,
25919        out_f: usize,
25920    ) -> Result<(), Box<dyn std::error::Error>> {
25921        self.linear_device_into(x, w, y, 1, in_f, out_f)
25922    }
25923
25924    /// On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml).
25925    /// cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation).
25926    /// DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt's reduction split is
25927    /// n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches
25928    /// must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE
25929    /// router/shexp sites and matmul_decode_exact's Float arm.
25930    pub fn linear_decode_exact(
25931        &self,
25932        x: &CudaSlice<f32>,
25933        w: &CudaSlice<f32>,
25934        m_tokens: usize,
25935        in_f: usize,
25936        out_f: usize,
25937    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25938        if m_tokens == 1 {
25939            return self.linear(x, w, 1, in_f, out_f);
25940        }
25941        let xv = self.view(x, m_tokens * in_f);
25942        let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
25943        for t in 0..m_tokens {
25944            let row = xv.slice(t * in_f..(t + 1) * in_f);
25945            let mut xr = self.alloc_uninit::<f32>(in_f)?;
25946            self.copy_view_into(&mut xr, 0, &row, in_f)?;
25947            let yr = self.linear(&xr, w, 1, in_f, out_f)?;
25948            self.copy_into(&mut y, t * out_f, &yr, out_f)?;
25949        }
25950        Ok(y)
25951    }
25952
25953    pub fn linear(
25954        &self,
25955        x: &CudaSlice<f32>,
25956        w: &CudaSlice<f32>,
25957        m_tokens: usize,
25958        in_f: usize,
25959        out_f: usize,
25960    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25961        self.linear_device(x, w, m_tokens, in_f, out_f)
25962    }
25963
25964    fn linear_device<I>(
25965        &self,
25966        x: &I,
25967        w: &I,
25968        m_tokens: usize,
25969        in_f: usize,
25970        out_f: usize,
25971    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
25972    where
25973        I: cudarc::driver::DevicePtr<f32>,
25974    {
25975        let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; // cuBLASLt beta=0: C fully written
25976        self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
25977        Ok(c)
25978    }
25979
25980    fn linear_device_into<I, O>(
25981        &self,
25982        x: &I,
25983        w: &I,
25984        c: &mut O,
25985        m_tokens: usize,
25986        in_f: usize,
25987        out_f: usize,
25988    ) -> Result<(), Box<dyn std::error::Error>>
25989    where
25990        I: cudarc::driver::DevicePtr<f32>,
25991        O: cudarc::driver::DevicePtrMut<f32>,
25992    {
25993        use cudarc::cublaslt::{Matmul, MatmulConfig};
25994        let cfg = MatmulConfig {
25995            transa: true,
25996            transb: false,
25997            transc: false,
25998            m: out_f as u64,
25999            n: m_tokens as u64,
26000            k: in_f as u64,
26001            alpha: 1.0,
26002            lda: in_f as i64,
26003            ldb: in_f as i64,
26004            beta: 0.0,
26005            ldc: out_f as i64,
26006            stride_a: None,
26007            stride_b: None,
26008            stride_c: None,
26009            stride_bias: None,
26010            batch_size: None,
26011        };
26012        let blas = self.gpu.blas();
26013        unsafe {
26014            blas.matmul(cfg, w, x, c, None, None)?;
26015        }
26016        Ok(())
26017    }
26018
26019    /// Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].
26020    ///
26021    /// LONG-CTX DISPATCH (lane/hermes-perf-fixes, 2026-08-23): the smem kernel's `T_kv*4`
26022    /// dynamic shared memory exceeds the 48KB launch bound past T_kv=12288 — the plain
26023    /// full-attn sibling of the DFlash2 B2 crash the windowed layers fixed with
26024    /// `sdpa_naive_w_lo`. Past the bound this transparently takes the byte-identical
26025    /// gmem-scores twin (`sdpa_naive_gmem`, kernel_check-pinned) instead of returning the
26026    /// launch error mid-request.
26027    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
26028    pub fn sdpa_naive(
26029        &self,
26030        q: &CudaSlice<f32>,
26031        k: &CudaSlice<f32>,
26032        v: &CudaSlice<f32>,
26033        o: &mut CudaSlice<f32>,
26034        head_dim: usize,
26035        n_head: usize,
26036        n_head_kv: usize,
26037        t: usize,
26038        t_kv: usize,
26039        scale: f32,
26040        causal: bool,
26041    ) -> Result<(), Box<dyn std::error::Error>> {
26042        if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
26043            return self.sdpa_naive_gmem(
26044                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26045            );
26046        }
26047        let f = self.func("sdpa_naive_f32");
26048        let cfg = LaunchConfig {
26049            grid_dim: (n_head as u32, t as u32, 1),
26050            block_dim: (128, 1, 1),
26051            shared_mem_bytes: (t_kv * 4) as u32,
26052        };
26053        let (hd, nh, nhkv, ti, tkvi, cz) = (
26054            head_dim as i32,
26055            n_head as i32,
26056            n_head_kv as i32,
26057            t as i32,
26058            t_kv as i32,
26059            causal as i32,
26060        );
26061        let __s_b = self.gpu.stream();
26062        let mut b = __s_b.launch_builder(&f);
26063        b.arg(q)
26064            .arg(k)
26065            .arg(v)
26066            .arg(o)
26067            .arg(&hd)
26068            .arg(&nh)
26069            .arg(&nhkv)
26070            .arg(&ti)
26071            .arg(&tkvi)
26072            .arg(&scale)
26073            .arg(&cz);
26074        unsafe {
26075            b.launch(cfg)?;
26076        }
26077        Ok(())
26078    }
26079
26080    /// Global-memory-scores twin of [`Self::sdpa_naive`] (lane/hermes-perf-fixes, 2026-08-23).
26081    /// Same kernel body with the per-(head, query) scores row in a device workspace instead
26082    /// of dynamic shared memory: identical loop structure and reduction order, so the output
26083    /// is BYTE-IDENTICAL to the smem kernel wherever both launch (kernel_check
26084    /// `sdpa_naive_gmem` pins bit-identity plus the >12k arm where the smem kernel MUST
26085    /// fail). O(n_head * T * T_kv * 4) workspace — fine for the tall-KV block shapes that
26086    /// hit the bound (dspark/dflash full-attn: T <= block size), guarded so a square
26087    /// T==T_kv caller cannot silently allocate tens of GB.
26088    #[allow(clippy::too_many_arguments)]
26089    pub fn sdpa_naive_gmem(
26090        &self,
26091        q: &CudaSlice<f32>,
26092        k: &CudaSlice<f32>,
26093        v: &CudaSlice<f32>,
26094        o: &mut CudaSlice<f32>,
26095        head_dim: usize,
26096        n_head: usize,
26097        n_head_kv: usize,
26098        t: usize,
26099        t_kv: usize,
26100        scale: f32,
26101        causal: bool,
26102    ) -> Result<(), Box<dyn std::error::Error>> {
26103        let ws_len = n_head
26104            .checked_mul(t)
26105            .and_then(|x| x.checked_mul(t_kv))
26106            .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
26107        let ws_bytes = ws_len
26108            .checked_mul(std::mem::size_of::<f32>())
26109            .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
26110        if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
26111            return Err(format!(
26112                "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
26113                 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
26114                 needs a tiled/flash kernel, not the naive oracle"
26115            )
26116            .into());
26117        }
26118        let mut scores = self.uninit(ws_len)?;
26119        let f = self.func("sdpa_naive_gmem_f32");
26120        let cfg = LaunchConfig {
26121            grid_dim: (n_head as u32, t as u32, 1),
26122            block_dim: (128, 1, 1),
26123            shared_mem_bytes: 0,
26124        };
26125        let (hd, nh, nhkv, ti, tkvi, cz) = (
26126            head_dim as i32,
26127            n_head as i32,
26128            n_head_kv as i32,
26129            t as i32,
26130            t_kv as i32,
26131            causal as i32,
26132        );
26133        let __s_b = self.gpu.stream();
26134        let mut b = __s_b.launch_builder(&f);
26135        b.arg(q)
26136            .arg(k)
26137            .arg(v)
26138            .arg(o)
26139            .arg(&mut scores)
26140            .arg(&hd)
26141            .arg(&nh)
26142            .arg(&nhkv)
26143            .arg(&ti)
26144            .arg(&tkvi)
26145            .arg(&scale)
26146            .arg(&cz);
26147        unsafe {
26148            b.launch(cfg)?;
26149        }
26150        Ok(())
26151    }
26152
26153    /// Island twin (lane/gemma-vision): causal + sliding-window attention with
26154    /// bidirectional image islands. `span_id` labels each absolute kv position
26155    /// (-1 text, >=0 island id); same-island keys are visible unconditionally,
26156    /// reproducing the reference's non-causal image batch. window 0 = no window.
26157    #[allow(clippy::too_many_arguments)]
26158    pub fn sdpa_naive_island(
26159        &self,
26160        q: &CudaSlice<f32>,
26161        k: &CudaSlice<f32>,
26162        v: &CudaSlice<f32>,
26163        o: &mut CudaSlice<f32>,
26164        span_id: &CudaSlice<i32>,
26165        head_dim: usize,
26166        n_head: usize,
26167        n_head_kv: usize,
26168        t: usize,
26169        t_kv: usize,
26170        scale: f32,
26171        window: usize,
26172    ) -> Result<(), Box<dyn std::error::Error>> {
26173        let f = self.func("sdpa_naive_island_f32");
26174        let cfg = LaunchConfig {
26175            grid_dim: (n_head as u32, t as u32, 1),
26176            block_dim: (128, 1, 1),
26177            shared_mem_bytes: (t_kv * 4) as u32,
26178        };
26179        let (hd, nh, nhkv, ti, tkvi, wi) = (
26180            head_dim as i32,
26181            n_head as i32,
26182            n_head_kv as i32,
26183            t as i32,
26184            t_kv as i32,
26185            window as i32,
26186        );
26187        let __s_b = self.gpu.stream();
26188        let mut b = __s_b.launch_builder(&f);
26189        b.arg(q)
26190            .arg(k)
26191            .arg(v)
26192            .arg(o)
26193            .arg(span_id)
26194            .arg(&hd)
26195            .arg(&nh)
26196            .arg(&nhkv)
26197            .arg(&ti)
26198            .arg(&tkvi)
26199            .arg(&scale)
26200            .arg(&wi);
26201        unsafe {
26202            b.launch(cfg)?;
26203        }
26204        Ok(())
26205    }
26206
26207    /// Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).
26208    #[allow(clippy::too_many_arguments)]
26209    pub fn sdpa_naive_w(
26210        &self,
26211        q: &CudaSlice<f32>,
26212        k: &CudaSlice<f32>,
26213        v: &CudaSlice<f32>,
26214        o: &mut CudaSlice<f32>,
26215        head_dim: usize,
26216        n_head: usize,
26217        n_head_kv: usize,
26218        t: usize,
26219        t_kv: usize,
26220        scale: f32,
26221        causal: bool,
26222        window: usize,
26223    ) -> Result<(), Box<dyn std::error::Error>> {
26224        let f = self.func("sdpa_naive_w_f32");
26225        let cfg = LaunchConfig {
26226            grid_dim: (n_head as u32, t as u32, 1),
26227            block_dim: (128, 1, 1),
26228            shared_mem_bytes: (t_kv * 4) as u32,
26229        };
26230        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26231            head_dim as i32,
26232            n_head as i32,
26233            n_head_kv as i32,
26234            t as i32,
26235            t_kv as i32,
26236            causal as i32,
26237            window as i32,
26238        );
26239        let __s_b = self.gpu.stream();
26240        let mut b = __s_b.launch_builder(&f);
26241        b.arg(q)
26242            .arg(k)
26243            .arg(v)
26244            .arg(o)
26245            .arg(&hd)
26246            .arg(&nh)
26247            .arg(&nhkv)
26248            .arg(&ti)
26249            .arg(&tkvi)
26250            .arg(&scale)
26251            .arg(&cz)
26252            .arg(&wi);
26253        unsafe {
26254            b.launch(cfg)?;
26255        }
26256        Ok(())
26257    }
26258
26259    /// Lo-clipped windowed sdpa_naive twin (lane/dflash2-longctx, DFLASH2-EVAL §10.6(c)).
26260    /// Same mask law as `sdpa_naive_w`, but keys below every query's window floor are never
26261    /// read: kv_lo = max(0, (t_kv - t) + 1 - window) — the oldest key visible to the OLDEST
26262    /// query row (q_pos = t_kv - t). Dynamic shared memory shrinks from t_kv*4 bytes (which
26263    /// blows the 48KB launch bound at ~12k rows — the B2 ctx crash) to (t_kv - kv_lo)*4 =
26264    /// (window - 1 + t)*4, and the key scan drops from O(t_kv) to O(window + t). Output is
26265    /// byte-identical to `sdpa_naive_w` (masked keys contribute exact zeros to same-order
26266    /// reductions; kernel_check `sdpa_naive_w_lo` pins it). window == 0 (no window) keeps
26267    /// kv_lo = 0 and is then shape-identical to the legacy kernel, including its bound.
26268    ///
26269    /// `kv_floor` (lane/spec-exclusions-20260902, the DFlash2 COLD-DRAFTER arm): the first
26270    /// key row that EXISTS. `kv_lo` is raised to it, so keys below the floor are never read
26271    /// or scored, i.e. the queries attend to a context that is simply shorter than the
26272    /// window (the same program a short prompt runs). `0` = the pre-lane clip exactly. The
26273    /// floor never clips the query rows themselves: callers pass a floor `<= t_kv - t`.
26274    #[allow(clippy::too_many_arguments)]
26275    pub fn sdpa_naive_w_lo(
26276        &self,
26277        q: &CudaSlice<f32>,
26278        k: &CudaSlice<f32>,
26279        v: &CudaSlice<f32>,
26280        o: &mut CudaSlice<f32>,
26281        head_dim: usize,
26282        n_head: usize,
26283        n_head_kv: usize,
26284        t: usize,
26285        t_kv: usize,
26286        scale: f32,
26287        causal: bool,
26288        window: usize,
26289        kv_floor: usize,
26290    ) -> Result<(), Box<dyn std::error::Error>> {
26291        if kv_floor > t_kv - t {
26292            return Err(format!(
26293                "sdpa_naive_w_lo: kv_floor {kv_floor} would clip the query rows themselves \
26294                 (t_kv {t_kv} - t {t})"
26295            )
26296            .into());
26297        }
26298        let kv_lo = if window > 0 {
26299            (t_kv - t + 1).saturating_sub(window)
26300        } else {
26301            0
26302        }
26303        .max(kv_floor);
26304        let smem = (t_kv - kv_lo) * 4;
26305        if smem > 48 * 1024 {
26306            return Err(format!(
26307                "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
26308                 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
26309                 a window this wide needs the multi-pass long-ctx kernel"
26310            )
26311            .into());
26312        }
26313        let f = self.func("sdpa_naive_w_lo_f32");
26314        let cfg = LaunchConfig {
26315            grid_dim: (n_head as u32, t as u32, 1),
26316            block_dim: (128, 1, 1),
26317            shared_mem_bytes: smem as u32,
26318        };
26319        let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
26320            head_dim as i32,
26321            n_head as i32,
26322            n_head_kv as i32,
26323            t as i32,
26324            t_kv as i32,
26325            causal as i32,
26326            window as i32,
26327            kv_lo as i32,
26328        );
26329        let __s_b = self.gpu.stream();
26330        let mut b = __s_b.launch_builder(&f);
26331        b.arg(q)
26332            .arg(k)
26333            .arg(v)
26334            .arg(o)
26335            .arg(&hd)
26336            .arg(&nh)
26337            .arg(&nhkv)
26338            .arg(&ti)
26339            .arg(&tkvi)
26340            .arg(&scale)
26341            .arg(&cz)
26342            .arg(&wi)
26343            .arg(&lo);
26344        unsafe {
26345            b.launch(cfg)?;
26346        }
26347        Ok(())
26348    }
26349
26350    /// SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).
26351    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
26352    pub fn sdpa_naive_view(
26353        &self,
26354        q: &CudaSlice<f32>,
26355        k: &cudarc::driver::CudaView<f32>,
26356        v: &cudarc::driver::CudaView<f32>,
26357        o: &mut CudaSlice<f32>,
26358        head_dim: usize,
26359        n_head: usize,
26360        n_head_kv: usize,
26361        t: usize,
26362        t_kv: usize,
26363        scale: f32,
26364        causal: bool,
26365    ) -> Result<(), Box<dyn std::error::Error>> {
26366        let f = self.func("sdpa_naive_f32");
26367        let cfg = LaunchConfig {
26368            grid_dim: (n_head as u32, t as u32, 1),
26369            block_dim: (128, 1, 1),
26370            shared_mem_bytes: (t_kv * 4) as u32,
26371        };
26372        let (hd, nh, nhkv, ti, tkvi, cz) = (
26373            head_dim as i32,
26374            n_head as i32,
26375            n_head_kv as i32,
26376            t as i32,
26377            t_kv as i32,
26378            causal as i32,
26379        );
26380        let __s_b = self.gpu.stream();
26381        let mut b = __s_b.launch_builder(&f);
26382        b.arg(q)
26383            .arg(k)
26384            .arg(v)
26385            .arg(o)
26386            .arg(&hd)
26387            .arg(&nh)
26388            .arg(&nhkv)
26389            .arg(&ti)
26390            .arg(&tkvi)
26391            .arg(&scale)
26392            .arg(&cz);
26393        unsafe {
26394            b.launch(cfg)?;
26395        }
26396        Ok(())
26397    }
26398
26399    /// Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32
26400    /// workspaces, then calls `sdpa_naive`. This is an explicit API: the optimized prefill view
26401    /// dispatch remains unchanged, so callers can use it as a reference or compatibility path.
26402    /// Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch).
26403    /// `g` picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/
26404    /// fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the
26405    /// f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer's quantized rows.
26406    #[allow(clippy::too_many_arguments)]
26407    pub fn fa_dequant_kv_view_f32(
26408        &self,
26409        k: &cudarc::driver::CudaView<u8>,
26410        v: &cudarc::driver::CudaView<u8>,
26411        kf: &mut CudaSlice<f32>,
26412        vf: &mut CudaSlice<f32>,
26413        kv_dim_k: usize,
26414        kv_dim_v: usize,
26415        t_kv: usize,
26416        k_tok_bytes: usize,
26417        v_tok_bytes: usize,
26418        g: bool,
26419    ) -> Result<(), Box<dyn std::error::Error>> {
26420        let f = if g {
26421            self.func_g("fa_dequant_kv_ws_f32")
26422        } else {
26423            self.func("fa_dequant_kv_ws_f32")
26424        };
26425        let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
26426        #[allow(clippy::manual_div_ceil)]
26427        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26428        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26429        let cfg = LaunchConfig {
26430            grid_dim: (nblk.max(1), 1, 1),
26431            block_dim: (256, 1, 1),
26432            shared_mem_bytes: 0,
26433        };
26434        let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
26435        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26436        let __s_b = self.gpu.stream();
26437        let mut b = __s_b.launch_builder(&f);
26438        b.arg(k)
26439            .arg(v)
26440            .arg(&mut *kf)
26441            .arg(&mut *vf)
26442            .arg(&kdk)
26443            .arg(&kdv)
26444            .arg(&tkvi)
26445            .arg(&ktb)
26446            .arg(&vtb);
26447        unsafe {
26448            b.launch(cfg)?;
26449        }
26450        Ok(())
26451    }
26452
26453    #[allow(clippy::too_many_arguments)]
26454    pub fn sdpa_naive_quantized_view(
26455        &self,
26456        q: &CudaSlice<f32>,
26457        k: &cudarc::driver::CudaView<u8>,
26458        v: &cudarc::driver::CudaView<u8>,
26459        o: &mut CudaSlice<f32>,
26460        head_dim: usize,
26461        n_head: usize,
26462        n_head_kv: usize,
26463        t: usize,
26464        t_kv: usize,
26465        scale: f32,
26466        causal: bool,
26467        k_tok_bytes: usize,
26468        v_tok_bytes: usize,
26469    ) -> Result<(), Box<dyn std::error::Error>> {
26470        let kv_dim = n_head_kv * head_dim;
26471        let mut kf = self.uninit(t_kv * kv_dim)?;
26472        let mut vf = self.uninit(t_kv * kv_dim)?;
26473        let f = self.func("fa_dequant_kv_ws_f32");
26474        let total = (2 * t_kv * kv_dim) as u64;
26475        #[allow(clippy::manual_div_ceil)]
26476        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26477        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26478        let cfg = LaunchConfig {
26479            grid_dim: (nblk.max(1), 1, 1),
26480            block_dim: (256, 1, 1),
26481            shared_mem_bytes: 0,
26482        };
26483        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
26484        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
26485        let __s_b = self.gpu.stream();
26486        let mut b = __s_b.launch_builder(&f);
26487        b.arg(k)
26488            .arg(v)
26489            .arg(&mut kf)
26490            .arg(&mut vf)
26491            .arg(&kv_dim_i)
26492            .arg(&kv_dim_i)
26493            .arg(&t_kv_i)
26494            .arg(&k_tok_bytes_i)
26495            .arg(&v_tok_bytes_i);
26496        unsafe { b.launch(cfg)? };
26497        self.sdpa_naive(
26498            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26499        )
26500    }
26501
26502    /// WINDOWED twin of `sdpa_naive_quantized_view` (step35 SWA prefill): dequant the KV byte
26503    /// view into f32 workspaces with the SAME `fa_dequant_kv_ws_f32` launch, then run
26504    /// `sdpa_naive_w` instead of `sdpa_naive`. `window == 0` is the unwindowed form (the kernel
26505    /// treats a non-positive window as "no window mask"), so this is a strict superset of the
26506    /// unwindowed function above and produces bit-identical output at window == 0.
26507    ///
26508    /// Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
26509    /// only (`fa_prefill_w_f32` == `fa_prefill_f32_body<256>`, and the quantized-view windowed
26510    /// twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
26511    /// FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
26512    /// unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
26513    #[allow(clippy::too_many_arguments)]
26514    pub fn sdpa_naive_w_quantized_view(
26515        &self,
26516        q: &CudaSlice<f32>,
26517        k: &cudarc::driver::CudaView<u8>,
26518        v: &cudarc::driver::CudaView<u8>,
26519        o: &mut CudaSlice<f32>,
26520        head_dim: usize,
26521        n_head: usize,
26522        n_head_kv: usize,
26523        t: usize,
26524        t_kv: usize,
26525        scale: f32,
26526        causal: bool,
26527        window: usize,
26528        k_tok_bytes: usize,
26529        v_tok_bytes: usize,
26530    ) -> Result<(), Box<dyn std::error::Error>> {
26531        let kv_dim = n_head_kv * head_dim;
26532        let mut kf = self.uninit(t_kv * kv_dim)?;
26533        let mut vf = self.uninit(t_kv * kv_dim)?;
26534        let f = self.func("fa_dequant_kv_ws_f32");
26535        let total = (2 * t_kv * kv_dim) as u64;
26536        #[allow(clippy::manual_div_ceil)]
26537        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26538        let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26539        let cfg = LaunchConfig {
26540            grid_dim: (nblk.max(1), 1, 1),
26541            block_dim: (256, 1, 1),
26542            shared_mem_bytes: 0,
26543        };
26544        let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
26545        let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
26546        let __s_b = self.gpu.stream();
26547        let mut b = __s_b.launch_builder(&f);
26548        b.arg(k)
26549            .arg(v)
26550            .arg(&mut kf)
26551            .arg(&mut vf)
26552            .arg(&kv_dim_i)
26553            .arg(&kv_dim_i)
26554            .arg(&t_kv_i)
26555            .arg(&k_tok_bytes_i)
26556            .arg(&v_tok_bytes_i);
26557        unsafe { b.launch(cfg)? };
26558        self.sdpa_naive_w(
26559            q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
26560        )
26561    }
26562
26563    /// Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync,
26564    /// head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1.
26565    /// Q/K/V/O [head_dim, n_head(_kv), T].
26566    #[allow(clippy::too_many_arguments)]
26567    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
26568    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26569    pub fn fa_prefill(
26570        &self,
26571        q: &CudaSlice<f32>,
26572        k: &CudaSlice<f32>,
26573        v: &CudaSlice<f32>,
26574        o: &mut CudaSlice<f32>,
26575        head_dim: usize,
26576        n_head: usize,
26577        n_head_kv: usize,
26578        t: usize,
26579        t_kv: usize,
26580        scale: f32,
26581        causal: bool,
26582    ) -> Result<(), Box<dyn std::error::Error>> {
26583        if portable_mma_gated() {
26584            return self.sdpa_naive(
26585                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26586            );
26587        }
26588        // FA3 v10 arm (task #20, OPT-IN MEMRA_FA3=1 — harness-proven 883us vs the shipped
26589        // kernel's 993us at T=2048): TMA-swizzled wgmma FA, fresh causal hd256 only.
26590        // NEW NUMERIC CONFIG (GDN-mma precedent): online softmax / bf16-P class — the
26591        // run-gen argmax + greedy-stream batteries arbitrate; not bit-paired.
26592        // PROMOTED default-ON hopper (2026-07-27): 3-seed 2048-prime -> 128-decode
26593        // streams MATCH vs mma, full battery green, lane interleaved 5/5 (+2.4%).
26594        // MEMRA_FA3=0 reverts; kernel-check pins the mma config regardless.
26595        let fa3_on = head_dim == 256
26596            && causal
26597            && t == t_kv
26598            && match std::env::var("MEMRA_FA3").as_deref() {
26599                Ok("0") => false,
26600                // The force arm consults the arch now: the bf16 stage below calls
26601                // f32_to_bf16_into -> func("f32_to_bf16_bulk"), which cu/hybrid.cu:1623 omits on
26602                // a portable build. Refuse at the switch, not at the lookup.
26603                Ok("1") => {
26604                    refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
26605                    true
26606                }
26607                _ => cfg!(memra_hopper_mma),
26608            };
26609        if fa3_on {
26610            let n = t * n_head * head_dim;
26611            let nkv = t * n_head_kv * head_dim;
26612            let mut q16 = self.alloc_u8_uninit(n * 2)?;
26613            let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
26614            let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
26615            self.f32_to_bf16_into(q, &mut q16, n)?;
26616            self.f32_to_bf16_into(k, &mut k16, nkv)?;
26617            self.f32_to_bf16_into(v, &mut v16, nkv)?;
26618            let rc = {
26619                use cudarc::driver::{DevicePtr, DevicePtrMut};
26620                let stream = self.gpu.stream();
26621                let (qp, _g1) = q16.device_ptr(&stream);
26622                let (kp, _g2) = k16.device_ptr(&stream);
26623                let (vp, _g3) = v16.device_ptr(&stream);
26624                let (op, _g4) = o.device_ptr_mut(&stream);
26625                unsafe {
26626                    memra_fa3_prefill(
26627                        qp as *const core::ffi::c_void,
26628                        kp as *const core::ffi::c_void,
26629                        vp as *const core::ffi::c_void,
26630                        op as *mut f32,
26631                        t as i32,
26632                        n_head as i32,
26633                        n_head_kv as i32,
26634                        head_dim as i32,
26635                        scale,
26636                        stream.cu_stream() as *mut core::ffi::c_void,
26637                    )
26638                }
26639            };
26640            if rc != 0 {
26641                return Err(format!("memra_fa3_prefill rc={rc}").into());
26642            }
26643            return Ok(());
26644        }
26645        // FLOOR PORT (P2+P0a+P0b+P1): 4 warps/CTA, BLOCK_Q=64 query rows, BK=32 KV tile,
26646        // Q-in-reg + register-O, grid.y=n_head_kv (4 Q-heads share staged K/V).
26647        // P1 plain arm (MEMRA_FA_P1=1 opt-in until the qwen battery): the engine-study body
26648        // (FA2 schedule + boundary split + swizzle) on the non-windowed lane. bf16 pre-convert.
26649        static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26650        let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
26651        if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
26652            const BLOCK_Q: usize = 64;
26653            const BKX: usize = 32;
26654            let f = self.func("fa_prefill_bf16_p1");
26655            let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
26656                + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
26657            use cudarc::driver::sys::CUfunction_attribute_enum as A;
26658            f.set_attribute(
26659                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26660                shmem as i32,
26661            )?;
26662            let cfg = LaunchConfig {
26663                grid_dim: (
26664                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
26665                    n_head as u32,
26666                    1,
26667                ),
26668                block_dim: (32, 4, 1),
26669                shared_mem_bytes: shmem,
26670            };
26671            let (hd, nh, nhkv, ti, tkvi, cz) = (
26672                head_dim as i32,
26673                n_head as i32,
26674                n_head_kv as i32,
26675                t as i32,
26676                t_kv as i32,
26677                causal as i32,
26678            );
26679            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
26680            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
26681            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
26682            let __s_b = self.gpu.stream();
26683            let mut b = __s_b.launch_builder(&f);
26684            b.arg(&qb)
26685                .arg(&kb)
26686                .arg(&vb)
26687                .arg(o)
26688                .arg(&hd)
26689                .arg(&nh)
26690                .arg(&nhkv)
26691                .arg(&ti)
26692                .arg(&tkvi)
26693                .arg(&scale)
26694                .arg(&cz);
26695            unsafe {
26696                b.launch(cfg)?;
26697            }
26698            return Ok(());
26699        }
26700        // Edge 5a (DEFAULT): fa_prefill_f32_pp — register-resident softmax (no sSw smem
26701        // round-trip), the FA3 softmax-GEMM overlap variant. ncu (pp512): short_scoreboard
26702        // 4.32->3.47, wait 1.99->1.45, per-call ~577us->~440us (1.31x) at flat 12.1% warps /
26703        // 255 regs / 2 CTAs (occupancy preserved). Bit-safe: 9B+27B argmax MATCH, rel 2.55e-3
26704        // vs floor 3.03e-3. MEMRA_FA_FLOOR reverts to the serialized-softmax floor kernel.
26705        const BK: usize = 32;
26706        // W2 lane (MEMRA_FA_PP_W2=1, ncu 2026-07-26): 2-warp/32-row CTA tile doubles grid.x —
26707        // bit-identical per-row math, pure coverage trade for the 6.25%-occupancy starvation.
26708        let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
26709        let (block_q, warps, w2_sfx): (usize, u32, &str) =
26710            if w2 { (32, 2, "_w2") } else { (64, 4, "") };
26711        // hd128 twins (2026-07-07): the prefill kernels are template-stamped at 256 (original
26712        // names, dispatch unchanged) and 128 (`_hd128`, the MiniMax-M3 class). Callers gate
26713        // other head_dims to sdpa_naive before reaching here.
26714        let hd_sfx = fa_hd_suffix(head_dim)?;
26715        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
26716        // BF16-KV staging lane (2026-07-26, default ON): the kernel converts K/V to bf16
26717        // during staging anyway — pre-converting to bf16 mirrors is BIT-IDENTICAL (same
26718        // __float2bfloat16 values into the same mma) and turns the 67%-of-stalls scalar
26719        // staging into int4 vector copies. MEMRA_FA_BF16KV=0 reverts.
26720        let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
26721        let (kb16, vb16) = if bf16kv {
26722            let n = t_kv * n_head_kv * head_dim;
26723            let mut kb = self.alloc_u8_uninit(n * 2)?;
26724            let mut vb = self.alloc_u8_uninit(n * 2)?;
26725            let fcv = self.func("f32_to_bf16_bulk");
26726            let ni = n as i64;
26727            let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26728            let __s_b = self.gpu.stream();
26729            let mut b = __s_b.launch_builder(&fcv);
26730            b.arg(k).arg(&mut kb).arg(&ni);
26731            unsafe {
26732                b.launch(cfgc)?;
26733            }
26734            let __s_b = self.gpu.stream();
26735            let mut b = __s_b.launch_builder(&fcv);
26736            b.arg(v).arg(&mut vb).arg(&ni);
26737            unsafe {
26738                b.launch(cfgc)?;
26739            }
26740            (Some(kb), Some(vb))
26741        } else {
26742            (None, None)
26743        };
26744        let f = self.func(&if bf16kv {
26745            format!("fa_prefill_bf16kv_pp{hd_sfx}")
26746        } else {
26747            format!(
26748                "fa_prefill_f32{}{}{hd_sfx}",
26749                if floor { "" } else { "_pp" },
26750                if floor { "" } else { w2_sfx }
26751            )
26752        });
26753        // persistent smem: bf16*(KV_STAGES*(sK + sV) + sP) + f32*(sS + sM + sL);
26754        // the bf16kv ring doubles the K/V stages (KV_STAGES=2).
26755        let kv_stages = if bf16kv { 2 } else { 1 };
26756        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
26757            + 4 * (block_q * BK + 2 * block_q)) as u32;
26758        use cudarc::driver::sys::CUfunction_attribute_enum as A;
26759        f.set_attribute(
26760            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26761            shmem as i32,
26762        )?;
26763        let cfg = LaunchConfig {
26764            grid_dim: (
26765                (t as u32 + block_q as u32 - 1) / block_q as u32,
26766                n_head as u32,
26767                1,
26768            ),
26769            block_dim: (32, warps, 1),
26770            shared_mem_bytes: shmem,
26771        };
26772        let (hd, nh, nhkv, ti, tkvi, cz) = (
26773            head_dim as i32,
26774            n_head as i32,
26775            n_head_kv as i32,
26776            t as i32,
26777            t_kv as i32,
26778            causal as i32,
26779        );
26780        let __s_b = self.gpu.stream();
26781        let mut b = __s_b.launch_builder(&f);
26782        b.arg(q);
26783        match (&kb16, &vb16) {
26784            (Some(kb), Some(vb)) => {
26785                b.arg(kb).arg(vb);
26786            }
26787            _ => {
26788                b.arg(k).arg(v);
26789            }
26790        }
26791        b.arg(o)
26792            .arg(&hd)
26793            .arg(&nh)
26794            .arg(&nhkv)
26795            .arg(&ti)
26796            .arg(&tkvi)
26797            .arg(&scale)
26798            .arg(&cz);
26799        unsafe {
26800            b.launch(cfg)?;
26801        }
26802        Ok(())
26803    }
26804
26805    /// Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill's
26806    /// exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile
26807    /// skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.
26808    #[allow(clippy::too_many_arguments)]
26809    pub fn fa_prefill_w(
26810        &self,
26811        q: &CudaSlice<f32>,
26812        k: &CudaSlice<f32>,
26813        v: &CudaSlice<f32>,
26814        o: &mut CudaSlice<f32>,
26815        head_dim: usize,
26816        n_head: usize,
26817        n_head_kv: usize,
26818        t: usize,
26819        t_kv: usize,
26820        scale: f32,
26821        causal: bool,
26822        window: usize,
26823    ) -> Result<(), Box<dyn std::error::Error>> {
26824        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — the raw
26825        // portable_cuda gate was stale-conservative on Hopper; fa_prefill already flipped).
26826        if portable_mma_gated() {
26827            return self.sdpa_naive_w(
26828                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
26829            );
26830        }
26831        // Default: bf16-prestaged twin (same treatment as hd512 — Q/K/V pre-converted once,
26832        // int4 stage copies; bit-identical, kernel_check-gated). MEMRA_FAW_STAGE=f32 reverts;
26833        // MEMRA_FA_FLOOR keeps the f32 floor stamp untouched.
26834        static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26835        let faw_f32 =
26836            *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
26837        let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
26838        self.fa_prefill_w_arm(
26839            q,
26840            k,
26841            v,
26842            o,
26843            head_dim,
26844            n_head,
26845            n_head_kv,
26846            t,
26847            t_kv,
26848            scale,
26849            causal,
26850            window,
26851            floor || faw_f32,
26852            floor,
26853        )
26854    }
26855
26856    /// Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane).
26857    /// Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.
26858    #[allow(clippy::too_many_arguments)]
26859    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26860    pub fn fa_prefill_w_pre(
26861        &self,
26862        qb: &CudaSlice<u8>,
26863        kb: &CudaSlice<u8>,
26864        vb: &CudaSlice<u8>,
26865        o: &mut CudaSlice<f32>,
26866        head_dim: usize,
26867        n_head: usize,
26868        n_head_kv: usize,
26869        t: usize,
26870        t_kv: usize,
26871        scale: f32,
26872        causal: bool,
26873        window: usize,
26874        v_f16: bool,
26875    ) -> Result<(), Box<dyn std::error::Error>> {
26876        const BLOCK_Q: usize = 64;
26877        const BK: usize = 32;
26878        debug_assert_eq!(head_dim, 256);
26879        let hp = fa_f16pv_on()
26880            && faw_hp_on()
26881            && n_head.is_multiple_of(2)
26882            && (n_head / n_head_kv).is_multiple_of(2);
26883        debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
26884        if hp {
26885            const BLOCK_QH: usize = 32;
26886            // V bytes must be f16 for the h2 stamp; producer normally emits f16 (v_f16),
26887            // else re-encode through the pooled scratch (stream-ordered reuse).
26888            let mut vguard = self.fa_vf16_scratch.lock().unwrap();
26889            let vh: &CudaSlice<u8> = if v_f16 {
26890                vb
26891            } else {
26892                let n = t_kv * n_head_kv * head_dim;
26893                if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
26894                    *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
26895                }
26896                self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
26897                vguard.as_ref().unwrap()
26898            };
26899            let f = self.func("fa_prefill_w_bf16_p1h2");
26900            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
26901            use cudarc::driver::sys::CUfunction_attribute_enum as A;
26902            f.set_attribute(
26903                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26904                shmem as i32,
26905            )?;
26906            let cfg = LaunchConfig {
26907                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
26908                block_dim: (32, 4, 1),
26909                shared_mem_bytes: shmem,
26910            };
26911            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26912                head_dim as i32,
26913                n_head as i32,
26914                n_head_kv as i32,
26915                t as i32,
26916                t_kv as i32,
26917                causal as i32,
26918                window as i32,
26919            );
26920            let __s_b = self.gpu.stream();
26921            let mut b = __s_b.launch_builder(&f);
26922            b.arg(qb)
26923                .arg(kb)
26924                .arg(vh)
26925                .arg(o)
26926                .arg(&hd)
26927                .arg(&nh)
26928                .arg(&nhkv)
26929                .arg(&ti)
26930                .arg(&tkvi)
26931                .arg(&scale)
26932                .arg(&cz)
26933                .arg(&wi);
26934            unsafe {
26935                b.launch(cfg)?;
26936            }
26937            return Ok(());
26938        }
26939        let f = self.func("fa_prefill_w_bf16_p1");
26940        let shmem =
26941            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
26942        use cudarc::driver::sys::CUfunction_attribute_enum as A;
26943        f.set_attribute(
26944            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26945            shmem as i32,
26946        )?;
26947        let cfg = LaunchConfig {
26948            grid_dim: (
26949                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
26950                n_head as u32,
26951                1,
26952            ),
26953            block_dim: (32, 4, 1),
26954            shared_mem_bytes: shmem,
26955        };
26956        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26957            head_dim as i32,
26958            n_head as i32,
26959            n_head_kv as i32,
26960            t as i32,
26961            t_kv as i32,
26962            causal as i32,
26963            window as i32,
26964        );
26965        let __s_b = self.gpu.stream();
26966        let mut b = __s_b.launch_builder(&f);
26967        b.arg(qb)
26968            .arg(kb)
26969            .arg(vb)
26970            .arg(o)
26971            .arg(&hd)
26972            .arg(&nh)
26973            .arg(&nhkv)
26974            .arg(&ti)
26975            .arg(&tkvi)
26976            .arg(&scale)
26977            .arg(&cz)
26978            .arg(&wi);
26979        unsafe {
26980            b.launch(cfg)?;
26981        }
26982        Ok(())
26983    }
26984
26985    /// Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.
26986    #[allow(clippy::too_many_arguments)]
26987    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
26988    pub fn fa_prefill_w_arm(
26989        &self,
26990        q: &CudaSlice<f32>,
26991        k: &CudaSlice<f32>,
26992        v: &CudaSlice<f32>,
26993        o: &mut CudaSlice<f32>,
26994        head_dim: usize,
26995        n_head: usize,
26996        n_head_kv: usize,
26997        t: usize,
26998        t_kv: usize,
26999        scale: f32,
27000        causal: bool,
27001        window: usize,
27002        f32_stage: bool,
27003        floor: bool,
27004    ) -> Result<(), Box<dyn std::error::Error>> {
27005        const BLOCK_Q: usize = 64;
27006        const BK: usize = 32;
27007        debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
27008        // P1 (2026-07-22 engine study): per-head Br=64 stamp with the FA2 schedule (V-copy
27009        // over GEMM0, next-K over softmax+GEMM1) + boundary/interior mask split. FP order
27010        // preserved -> bit-identical (gated). MEMRA_FAW_P1=0 reverts to the g4/o2 arms.
27011        static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27012        let p1 = !floor
27013            && !f32_stage
27014            && *P1_ON.get_or_init(|| {
27015                std::env::var("MEMRA_FAW_P1")
27016                    .map(|v| v != "0")
27017                    .unwrap_or(true)
27018            });
27019        let hp = p1
27020            && fa_f16pv_on()
27021            && faw_hp_on()
27022            && n_head.is_multiple_of(2)
27023            && (n_head / n_head_kv).is_multiple_of(2);
27024        if hp {
27025            const BLOCK_QH: usize = 32;
27026            let f = self.func("fa_prefill_w_bf16_p1h2");
27027            let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
27028            use cudarc::driver::sys::CUfunction_attribute_enum as A;
27029            f.set_attribute(
27030                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27031                shmem as i32,
27032            )?;
27033            let cfg = LaunchConfig {
27034                grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
27035                block_dim: (32, 4, 1),
27036                shared_mem_bytes: shmem,
27037            };
27038            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
27039                head_dim as i32,
27040                n_head as i32,
27041                n_head_kv as i32,
27042                t as i32,
27043                t_kv as i32,
27044                causal as i32,
27045                window as i32,
27046            );
27047            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27048            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27049            let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
27050            let __s_b = self.gpu.stream();
27051            let mut b = __s_b.launch_builder(&f);
27052            b.arg(&qb)
27053                .arg(&kb)
27054                .arg(&vh)
27055                .arg(o)
27056                .arg(&hd)
27057                .arg(&nh)
27058                .arg(&nhkv)
27059                .arg(&ti)
27060                .arg(&tkvi)
27061                .arg(&scale)
27062                .arg(&cz)
27063                .arg(&wi);
27064            unsafe {
27065                b.launch(cfg)?;
27066            }
27067            return Ok(());
27068        }
27069        if p1 {
27070            let f = self.func("fa_prefill_w_bf16_p1");
27071            let shmem =
27072                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
27073            use cudarc::driver::sys::CUfunction_attribute_enum as A;
27074            f.set_attribute(
27075                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27076                shmem as i32,
27077            )?;
27078            let cfg = LaunchConfig {
27079                grid_dim: (
27080                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27081                    n_head as u32,
27082                    1,
27083                ),
27084                block_dim: (32, 4, 1),
27085                shared_mem_bytes: shmem,
27086            };
27087            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
27088                head_dim as i32,
27089                n_head as i32,
27090                n_head_kv as i32,
27091                t as i32,
27092                t_kv as i32,
27093                causal as i32,
27094                window as i32,
27095            );
27096            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27097            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27098            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27099            let __s_b = self.gpu.stream();
27100            let mut b = __s_b.launch_builder(&f);
27101            b.arg(&qb)
27102                .arg(&kb)
27103                .arg(&vb)
27104                .arg(o)
27105                .arg(&hd)
27106                .arg(&nh)
27107                .arg(&nhkv)
27108                .arg(&ti)
27109                .arg(&tkvi)
27110                .arg(&scale)
27111                .arg(&cz)
27112                .arg(&wi);
27113            unsafe {
27114                b.launch(cfg)?;
27115            }
27116            return Ok(());
27117        }
27118        // MQA head-grouping (MEMRA_FAW_G4=0 reverts): 4 heads/CTA share the staged K/V —
27119        // per-(head,row) FP chain identical to the per-head stamp -> bit-identical (gated).
27120        static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27121        let g4 = !floor
27122            && !f32_stage
27123            && n_head_kv == 1
27124            && n_head.is_multiple_of(4)
27125            && *G4_ON.get_or_init(|| {
27126                std::env::var("MEMRA_FAW_G4")
27127                    .map(|v| v != "0")
27128                    .unwrap_or(true)
27129            });
27130        if g4 {
27131            const SP_M: usize = 16;
27132            // Occupancy-2 twin (MEMRA_FAW_O2=0 reverts): one shared K/V buffer inside the dead
27133            // Q-stage region -> ~36.5KB smem, 2 CTA/SM (the llama hd256 mechanism). Bit-identical.
27134            static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27135            let o2 = *O2_ON.get_or_init(|| {
27136                std::env::var("MEMRA_FAW_O2")
27137                    .map(|v| v != "0")
27138                    .unwrap_or(true)
27139            });
27140            let f = self.func(if o2 {
27141                "fa_prefill_w_bf16_g4o2"
27142            } else {
27143                "fa_prefill_w_bf16_g4"
27144            });
27145            let shmem = if o2 {
27146                (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
27147            } else {
27148                (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
27149                    as u32
27150            };
27151            use cudarc::driver::sys::CUfunction_attribute_enum as A;
27152            f.set_attribute(
27153                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27154                shmem as i32,
27155            )?;
27156            let cfg = LaunchConfig {
27157                grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
27158                block_dim: (32, 4, 1),
27159                shared_mem_bytes: shmem,
27160            };
27161            let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
27162                head_dim as i32,
27163                n_head as i32,
27164                n_head_kv as i32,
27165                t as i32,
27166                t_kv as i32,
27167                causal as i32,
27168                window as i32,
27169            );
27170            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27171            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27172            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27173            let __s_b = self.gpu.stream();
27174            let mut b = __s_b.launch_builder(&f);
27175            b.arg(&qb)
27176                .arg(&kb)
27177                .arg(&vb)
27178                .arg(o)
27179                .arg(&hd)
27180                .arg(&nh)
27181                .arg(&nhkv)
27182                .arg(&ti)
27183                .arg(&tkvi)
27184                .arg(&scale)
27185                .arg(&cz)
27186                .arg(&wi);
27187            unsafe {
27188                b.launch(cfg)?;
27189            }
27190            return Ok(());
27191        }
27192        let f = self.func(if floor {
27193            "fa_prefill_w_f32"
27194        } else if f32_stage {
27195            "fa_prefill_w_f32_pp"
27196        } else {
27197            "fa_prefill_w_bf16_pp"
27198        });
27199        let shmem =
27200            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
27201        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27202        f.set_attribute(
27203            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27204            shmem as i32,
27205        )?;
27206        let cfg = LaunchConfig {
27207            grid_dim: (
27208                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27209                n_head as u32,
27210                1,
27211            ),
27212            block_dim: (32, 4, 1),
27213            shared_mem_bytes: shmem,
27214        };
27215        let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
27216            head_dim as i32,
27217            n_head as i32,
27218            n_head_kv as i32,
27219            t as i32,
27220            t_kv as i32,
27221            causal as i32,
27222            window as i32,
27223        );
27224        if f32_stage {
27225            let __s_b = self.gpu.stream();
27226            let mut b = __s_b.launch_builder(&f);
27227            b.arg(q)
27228                .arg(k)
27229                .arg(v)
27230                .arg(o)
27231                .arg(&hd)
27232                .arg(&nh)
27233                .arg(&nhkv)
27234                .arg(&ti)
27235                .arg(&tkvi)
27236                .arg(&scale)
27237                .arg(&cz)
27238                .arg(&wi);
27239            unsafe {
27240                b.launch(cfg)?;
27241            }
27242        } else {
27243            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27244            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27245            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27246            let __s_b = self.gpu.stream();
27247            let mut b = __s_b.launch_builder(&f);
27248            b.arg(&qb)
27249                .arg(&kb)
27250                .arg(&vb)
27251                .arg(o)
27252                .arg(&hd)
27253                .arg(&nh)
27254                .arg(&nhkv)
27255                .arg(&ti)
27256                .arg(&tkvi)
27257                .arg(&scale)
27258                .arg(&cz)
27259                .arg(&wi);
27260            unsafe {
27261                b.launch(cfg)?;
27262            }
27263        }
27264        Ok(())
27265    }
27266
27267    /// hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem,
27268    /// grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the
27269    /// V dims). Replaces the scalar sdpa_naive on the prime path's globals.
27270    #[allow(clippy::too_many_arguments)]
27271    pub fn fa_prefill_hd512(
27272        &self,
27273        q: &CudaSlice<f32>,
27274        k: &CudaSlice<f32>,
27275        v: &CudaSlice<f32>,
27276        o: &mut CudaSlice<f32>,
27277        head_dim: usize,
27278        n_head: usize,
27279        n_head_kv: usize,
27280        t: usize,
27281        t_kv: usize,
27282        scale: f32,
27283        causal: bool,
27284    ) -> Result<(), Box<dyn std::error::Error>> {
27285        // sm_90a rides the mma twins (portable_mma_gated, 2026-07-31 — same flip as _w).
27286        if portable_mma_gated() {
27287            return self.sdpa_naive(
27288                q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
27289            );
27290        }
27291        // Default: pre-convert Q/K/V to bf16 once and stage int4 (8 bf16/copy) — at 1 CTA/SM the
27292        // synchronous stage serializes with compute and MQA re-stages the same K/V per head CTA;
27293        // pre-converting halves staged bytes and cuts stage instructions 8x. BIT-IDENTICAL to the
27294        // f32-staged kernel (the converter applies the same __float2bfloat16 the stage applied;
27295        // kernel_check gates the identity). MEMRA_FA512_STAGE=f32 = rollback to the f32 kernel.
27296        static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27297        let f32_stage =
27298            *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
27299        // Single-pass arm (MEMRA_FA512_SP=0 reverts to the z=2 bf16 kernel): GEMM0 split-K across
27300        // the 2 warps instead of recomputed per O-half CTA — the 2026-07-22 kernel-diff excess.
27301        // Own numeric config (partial-sum order) — battery-gated.
27302        static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27303        let sp = !f32_stage
27304            && *SP_ON.get_or_init(|| {
27305                std::env::var("MEMRA_FA512_SP")
27306                    .map(|v| v != "0")
27307                    .unwrap_or(true)
27308            });
27309        self.fa_prefill_hd512_arm(
27310            q,
27311            k,
27312            v,
27313            o,
27314            head_dim,
27315            n_head,
27316            n_head_kv,
27317            t,
27318            t_kv,
27319            scale,
27320            causal,
27321            f32_stage,
27322            sp,
27323            sp && fa_f16pv_on(),
27324        )
27325    }
27326
27327    /// hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).
27328    #[allow(clippy::too_many_arguments)]
27329    pub fn fa_prefill_hd512_pre(
27330        &self,
27331        qb: &CudaSlice<u8>,
27332        kb: &CudaSlice<u8>,
27333        vb: &CudaSlice<u8>,
27334        o: &mut CudaSlice<f32>,
27335        head_dim: usize,
27336        n_head: usize,
27337        n_head_kv: usize,
27338        t: usize,
27339        t_kv: usize,
27340        scale: f32,
27341        causal: bool,
27342        v_f16: bool,
27343    ) -> Result<(), Box<dyn std::error::Error>> {
27344        debug_assert_eq!(head_dim, 512);
27345        const SP_M: usize = 16;
27346        const BKS: usize = 32;
27347        // f16-P/V door (MEMRA_FA_F16PV=1): P and the P@V accumulation in f16 (llama's fa=1 VKQ
27348        // class); KQ/softmax/rescale-band/final-normalize stay f32. Own numeric config —
27349        // battery-gated. V bytes must be f16 for the sp16 kernel (stage/ldmatrix are typeless).
27350        let f16pv = fa_f16pv_on();
27351        let nw = if f16pv { fa512_wide_warps() } else { 2 };
27352        let hp = f16pv
27353            && fa512_hp_on()
27354            && n_head.is_multiple_of(2)
27355            && (n_head / n_head_kv).is_multiple_of(2);
27356        debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
27357        let mut vguard = self.fa_vf16_scratch.lock().unwrap();
27358        let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
27359            // Fallback re-encode (producer emitted bf16); the emit lane normally hands f16.
27360            let n = t_kv * n_head_kv * head_dim;
27361            let need = n * 2;
27362            if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
27363                *vguard = Some(self.alloc_uninit::<u8>(need)?);
27364            }
27365            let dst = vguard.as_mut().unwrap();
27366            self.bf16_to_f16_into(vb, n, dst)?;
27367            vguard.as_ref().unwrap()
27368        } else {
27369            vb
27370        };
27371        let f = self.func(if hp {
27372            "fa_prefill_bf16_hd512_sp16h2"
27373        } else {
27374            match (f16pv, nw) {
27375                (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
27376                (true, _) => "fa_prefill_bf16_hd512_sp16",
27377                _ => "fa_prefill_bf16_hd512_sp",
27378            }
27379        });
27380        let (nwarp, npart) = if hp {
27381            (4usize, 4usize)
27382        } else if nw > 2 {
27383            (nw, nw)
27384        } else {
27385            (2, 1)
27386        };
27387        // h2 drops sQ (Q register-resident) and doubles sP/sS/sL for the head pair.
27388        let shmem = if hp {
27389            (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
27390                as u32
27391        } else {
27392            (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
27393                + 4 * (npart * SP_M * BKS + SP_M)) as u32
27394        };
27395        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27396        f.set_attribute(
27397            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27398            shmem as i32,
27399        )?;
27400        let grid_y = if hp {
27401            (n_head / 2) as u32
27402        } else {
27403            n_head as u32
27404        };
27405        let cfg = LaunchConfig {
27406            grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
27407            block_dim: (32, nwarp as u32, 1),
27408            shared_mem_bytes: shmem,
27409        };
27410        let (hd, nh, nhkv, ti, tkvi, cz) = (
27411            head_dim as i32,
27412            n_head as i32,
27413            n_head_kv as i32,
27414            t as i32,
27415            t_kv as i32,
27416            causal as i32,
27417        );
27418        let __s_b = self.gpu.stream();
27419        let mut b = __s_b.launch_builder(&f);
27420        b.arg(qb)
27421            .arg(kb)
27422            .arg(vref)
27423            .arg(o)
27424            .arg(&hd)
27425            .arg(&nh)
27426            .arg(&nhkv)
27427            .arg(&ti)
27428            .arg(&tkvi)
27429            .arg(&scale)
27430            .arg(&cz);
27431        unsafe {
27432            b.launch(cfg)?;
27433        }
27434        Ok(())
27435    }
27436
27437    /// Absorbed-form MLA prefill attention over a DSA-GATHERED index list, on tensor cores —
27438    /// the MEMRA_MLA_TC_PREFILL kernel (`fa_mla_gathered_bf16`, cu/flash_attn.cu). One CTA per
27439    /// (query, 16-head band); the query's index list is shared across heads (the DSA indexer
27440    /// mixes heads BEFORE top-k), which is exactly what gives the MMA its m axis. V is K
27441    /// (NoPE latent rows), so the kernel is `kv_rank == 512, d_rope == 0` ONLY and this
27442    /// launcher refuses anything else rather than approximate.
27443    #[allow(clippy::too_many_arguments)]
27444    pub fn mla_attn_gathered_tc(
27445        &self,
27446        q_lat_bf: &CudaSlice<u8>,   // [t_q, n_head, 512] bf16
27447        cache_bf: &CudaSlice<u8>,   // [t_kv, 512] bf16 latent rows
27448        idx: &CudaSlice<i32>,       // [t_q, width], ascending, -1 trailing
27449        o_lat: &mut CudaSlice<f32>, // [t_q, n_head, 512] f32
27450        n_head: usize,
27451        kv_rank: usize,
27452        t_q: usize,
27453        width: usize,
27454        scale: f32,
27455    ) -> Result<(), Box<dyn std::error::Error>> {
27456        if kv_rank != 512 {
27457            return Err(format!(
27458                "mla_attn_gathered_tc is stamped at kv_rank 512 (the glm5_next latent width); \
27459                 got {kv_rank} — the caller's door must fall back to the f32 gathered kernel"
27460            )
27461            .into());
27462        }
27463        if t_q == 0 || n_head == 0 {
27464            return Ok(());
27465        }
27466        const SP_M: usize = 16;
27467        const BKS: usize = 32;
27468        const HD: usize = 512;
27469        let f = self.func("fa_mla_gathered_bf16");
27470        // sQ + sK (V aliases K) + sP bf16, sS + sL f32, sIdx i32.
27471        let shmem =
27472            (2 * (SP_M * HD + BKS * HD + SP_M * BKS) + 4 * (SP_M * BKS + SP_M) + 4 * BKS) as u32;
27473        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27474        f.set_attribute(
27475            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27476            shmem as i32,
27477        )?;
27478        let cfg = LaunchConfig {
27479            grid_dim: (t_q as u32, (n_head as u32).div_ceil(SP_M as u32), 1),
27480            block_dim: (32, 2, 1),
27481            shared_mem_bytes: shmem,
27482        };
27483        let (nh, tq, w) = (n_head as i32, t_q as i32, width as i32);
27484        let __s_b = self.gpu.stream();
27485        let mut b = __s_b.launch_builder(&f);
27486        b.arg(q_lat_bf)
27487            .arg(cache_bf)
27488            .arg(idx)
27489            .arg(o_lat)
27490            .arg(&nh)
27491            .arg(&tq)
27492            .arg(&w)
27493            .arg(&scale);
27494        unsafe {
27495            b.launch(cfg)?;
27496        }
27497        Ok(())
27498    }
27499
27500    /// hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry
27501    /// (`fa_prefill_hd512` picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).
27502    #[allow(clippy::too_many_arguments)]
27503    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
27504    pub fn fa_prefill_hd512_arm(
27505        &self,
27506        q: &CudaSlice<f32>,
27507        k: &CudaSlice<f32>,
27508        v: &CudaSlice<f32>,
27509        o: &mut CudaSlice<f32>,
27510        head_dim: usize,
27511        n_head: usize,
27512        n_head_kv: usize,
27513        t: usize,
27514        t_kv: usize,
27515        scale: f32,
27516        causal: bool,
27517        f32_stage: bool,
27518        sp: bool,
27519        f16pv: bool,
27520    ) -> Result<(), Box<dyn std::error::Error>> {
27521        debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
27522        if sp && !f32_stage {
27523            // Single-pass: 16 q-rows/CTA, 2 warps, grid (ceil(T/16), n_head, 1).
27524            // smem: sQ[16][512] + sK[32][512] + sV[32][512] + sP[16][32] (bf16) + sS[16][32]+sL f32.
27525            // f16pv: sp16 kernel — f16 P + f16 P@V accum, V operand encoded f16.
27526            const SP_M: usize = 16;
27527            const BKS: usize = 32;
27528            let nw = if f16pv { fa512_wide_warps() } else { 2 };
27529            let hp = f16pv
27530                && fa512_hp_on()
27531                && n_head.is_multiple_of(2)
27532                && (n_head / n_head_kv).is_multiple_of(2);
27533            let f = self.func(if hp {
27534                "fa_prefill_bf16_hd512_sp16h2"
27535            } else {
27536                match (f16pv, nw) {
27537                    (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
27538                    (true, _) => "fa_prefill_bf16_hd512_sp16",
27539                    _ => "fa_prefill_bf16_hd512_sp",
27540                }
27541            });
27542            let (nwarp, npart) = if hp {
27543                (4usize, 4usize)
27544            } else if nw > 2 {
27545                (nw, nw)
27546            } else {
27547                (2, 1)
27548            };
27549            let shmem = if hp {
27550                (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
27551                    + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
27552            } else {
27553                (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
27554                    + 4 * (npart * SP_M * BKS + SP_M)) as u32
27555            };
27556            use cudarc::driver::sys::CUfunction_attribute_enum as A;
27557            f.set_attribute(
27558                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27559                shmem as i32,
27560            )?;
27561            let grid_y = if hp {
27562                (n_head / 2) as u32
27563            } else {
27564                n_head as u32
27565            };
27566            let cfg = LaunchConfig {
27567                grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
27568                block_dim: (32, nwarp as u32, 1),
27569                shared_mem_bytes: shmem,
27570            };
27571            let (hd, nh, nhkv, ti, tkvi, cz) = (
27572                head_dim as i32,
27573                n_head as i32,
27574                n_head_kv as i32,
27575                t as i32,
27576                t_kv as i32,
27577                causal as i32,
27578            );
27579            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27580            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27581            let vb = if f16pv {
27582                self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
27583            } else {
27584                self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
27585            };
27586            let __s_b = self.gpu.stream();
27587            let mut b = __s_b.launch_builder(&f);
27588            b.arg(&qb)
27589                .arg(&kb)
27590                .arg(&vb)
27591                .arg(o)
27592                .arg(&hd)
27593                .arg(&nh)
27594                .arg(&nhkv)
27595                .arg(&ti)
27596                .arg(&tkvi)
27597                .arg(&scale)
27598                .arg(&cz);
27599            unsafe {
27600                b.launch(cfg)?;
27601            }
27602            return Ok(());
27603        }
27604        const BLOCK_Q: usize = 32;
27605        const BK: usize = 32;
27606        const HALF: usize = 256;
27607        let f = self.func(if f32_stage {
27608            "fa_prefill_f32_hd512"
27609        } else {
27610            "fa_prefill_bf16_hd512"
27611        });
27612        // sQ[32][512] + sK[BK][512] + sV[BK][256] + sP[32][BK] (bf16) + sL[32] f32
27613        let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
27614            + 4 * BLOCK_Q) as u32;
27615        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27616        f.set_attribute(
27617            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27618            shmem as i32,
27619        )?;
27620        let cfg = LaunchConfig {
27621            grid_dim: (
27622                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27623                n_head as u32,
27624                2,
27625            ),
27626            block_dim: (32, 2, 1),
27627            shared_mem_bytes: shmem,
27628        };
27629        let (hd, nh, nhkv, ti, tkvi, cz) = (
27630            head_dim as i32,
27631            n_head as i32,
27632            n_head_kv as i32,
27633            t as i32,
27634            t_kv as i32,
27635            causal as i32,
27636        );
27637        if f32_stage {
27638            let __s_b = self.gpu.stream();
27639            let mut b = __s_b.launch_builder(&f);
27640            b.arg(q)
27641                .arg(k)
27642                .arg(v)
27643                .arg(o)
27644                .arg(&hd)
27645                .arg(&nh)
27646                .arg(&nhkv)
27647                .arg(&ti)
27648                .arg(&tkvi)
27649                .arg(&scale)
27650                .arg(&cz);
27651            unsafe {
27652                b.launch(cfg)?;
27653            }
27654        } else {
27655            let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27656            let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27657            let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27658            let __s_b = self.gpu.stream();
27659            let mut b = __s_b.launch_builder(&f);
27660            b.arg(&qb)
27661                .arg(&kb)
27662                .arg(&vb)
27663                .arg(o)
27664                .arg(&hd)
27665                .arg(&nh)
27666                .arg(&nhkv)
27667                .arg(&ti)
27668                .arg(&tkvi)
27669                .arg(&scale)
27670                .arg(&cz);
27671            unsafe {
27672                b.launch(cfg)?;
27673            }
27674        }
27675        Ok(())
27676    }
27677
27678    /// rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope
27679    /// values written as bf16 — the FA q/k operands come from this launch (bit-identical to the
27680    /// separate f32_to_bf16 the FA entries would run).
27681    #[allow(clippy::too_many_arguments)]
27682    pub fn rope_neox2_bf16e(
27683        &self,
27684        q: &mut CudaSlice<f32>,
27685        k: &mut CudaSlice<f32>,
27686        qb: &mut CudaSlice<u8>,
27687        kb: &mut CudaSlice<u8>,
27688        pos: &CudaSlice<i32>,
27689        head_dim: usize,
27690        n_dims: usize,
27691        nh_q: usize,
27692        nh_k: usize,
27693        n_tokens: usize,
27694        base: f32,
27695        freq_scale: f32,
27696        ff: Option<&CudaSlice<f32>>,
27697    ) -> Result<(), Box<dyn std::error::Error>> {
27698        let f = self.func("rope_neox2_bf16e_f32");
27699        let rows = ((nh_q + nh_k) * n_tokens) as u32;
27700        let cfg = LaunchConfig {
27701            grid_dim: (rows, 1, 1),
27702            block_dim: ((head_dim / 2) as u32, 1, 1),
27703            shared_mem_bytes: 0,
27704        };
27705        let theta_scale = base.powf(-2.0 / n_dims as f32);
27706        let (hd, nd, nhq, nhk, nt) = (
27707            head_dim as i32,
27708            n_dims as i32,
27709            nh_q as i32,
27710            nh_k as i32,
27711            n_tokens as i32,
27712        );
27713        let __s_b = self.gpu.stream();
27714        let mut b = __s_b.launch_builder(&f);
27715        match ff {
27716            Some(t) => {
27717                b.arg(&mut *q)
27718                    .arg(&mut *k)
27719                    .arg(&mut *qb)
27720                    .arg(&mut *kb)
27721                    .arg(pos)
27722                    .arg(&hd)
27723                    .arg(&nd)
27724                    .arg(&nhq)
27725                    .arg(&nhk)
27726                    .arg(&nt)
27727                    .arg(&theta_scale)
27728                    .arg(&freq_scale)
27729                    .arg(t);
27730                unsafe {
27731                    b.launch(cfg)?;
27732                }
27733            }
27734            None => {
27735                let null: u64 = 0;
27736                b.arg(&mut *q)
27737                    .arg(&mut *k)
27738                    .arg(&mut *qb)
27739                    .arg(&mut *kb)
27740                    .arg(pos)
27741                    .arg(&hd)
27742                    .arg(&nd)
27743                    .arg(&nhq)
27744                    .arg(&nhk)
27745                    .arg(&nt)
27746                    .arg(&theta_scale)
27747                    .arg(&freq_scale)
27748                    .arg(&null);
27749                unsafe {
27750                    b.launch(cfg)?;
27751                }
27752            }
27753        }
27754        Ok(())
27755    }
27756
27757    /// Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). `n % 4 == 0`
27758    /// (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.
27759    pub fn f32_to_bf16(
27760        &self,
27761        x: &CudaSlice<f32>,
27762        n: usize,
27763    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27764        assert!(
27765            n.is_multiple_of(4),
27766            "f32_to_bf16 requires n % 4 == 0, got {n}"
27767        );
27768        let mut y = self.alloc_uninit::<u8>(n * 2)?;
27769        let f = self.func("f32_to_bf16_flat");
27770        let n_i = n as i64;
27771        let cfg = LaunchConfig {
27772            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
27773            block_dim: (256, 1, 1),
27774            shared_mem_bytes: 0,
27775        };
27776        let __s_b = self.gpu.stream();
27777        let mut b = __s_b.launch_builder(&f);
27778        b.arg(x).arg(&mut y).arg(&n_i);
27779        unsafe {
27780            b.launch(cfg)?;
27781        }
27782        Ok(y)
27783    }
27784
27785    pub fn f32_to_f16(
27786        &self,
27787        x: &CudaSlice<f32>,
27788        n: usize,
27789    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27790        assert!(
27791            n.is_multiple_of(4),
27792            "f32_to_f16 requires n % 4 == 0, got {n}"
27793        );
27794        let mut y = self.alloc_uninit::<u8>(n * 2)?;
27795        let f = self.func("f32_to_f16_flat");
27796        let n_i = n as i64;
27797        let cfg = LaunchConfig {
27798            grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
27799            block_dim: (256, 1, 1),
27800            shared_mem_bytes: 0,
27801        };
27802        let __s_b = self.gpu.stream();
27803        let mut b = __s_b.launch_builder(&f);
27804        b.arg(x).arg(&mut y).arg(&n_i);
27805        unsafe {
27806            b.launch(cfg)?;
27807        }
27808        Ok(y)
27809    }
27810
27811    /// bf16 bytes -> f16 bytes, n elements (the f16-P/V door's V re-encode on the emit lane).
27812    pub fn bf16_to_f16(
27813        &self,
27814        xb: &CudaSlice<u8>,
27815        n: usize,
27816    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27817        let mut y = self.alloc_uninit::<u8>(n * 2)?;
27818        self.bf16_to_f16_into(xb, n, &mut y)?;
27819        Ok(y)
27820    }
27821
27822    /// Same conversion into a caller-owned (pooled) buffer; `y.len() >= n*2`.
27823    pub fn bf16_to_f16_into(
27824        &self,
27825        xb: &CudaSlice<u8>,
27826        n: usize,
27827        y: &mut CudaSlice<u8>,
27828    ) -> Result<(), Box<dyn std::error::Error>> {
27829        assert!(
27830            n.is_multiple_of(2),
27831            "bf16_to_f16 requires n % 2 == 0, got {n}"
27832        );
27833        assert!(y.len() >= n * 2);
27834        let f = self.func("bf16_to_f16_flat");
27835        let n2 = (n / 2) as i64;
27836        let cfg = LaunchConfig {
27837            grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
27838            block_dim: (256, 1, 1),
27839            shared_mem_bytes: 0,
27840        };
27841        let __s_b = self.gpu.stream();
27842        let mut b = __s_b.launch_builder(&f);
27843        b.arg(xb).arg(y).arg(&n2);
27844        unsafe {
27845            b.launch(cfg)?;
27846        }
27847        Ok(())
27848    }
27849
27850    /// task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE
27851    /// fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the
27852    /// per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T),
27853    /// head_dim in {256, 128}, bf16kv lane on.
27854    #[allow(clippy::too_many_arguments)]
27855    pub fn fa_prefill_vl8(
27856        &self,
27857        seqs: &[FaSeqVl],
27858        head_dim: usize,
27859        n_head: usize,
27860        n_head_kv: usize,
27861        scale: f32,
27862    ) -> Result<(), Box<dyn std::error::Error>> {
27863        const BK: usize = 32;
27864        let b = seqs.len();
27865        assert!((1..=8).contains(&b));
27866        let mut packed = [FaSeqVl::default(); 8];
27867        packed[..b].copy_from_slice(seqs);
27868        let v = FaVl8(packed);
27869        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
27870        let ept = (n_head_kv * head_dim) as i32;
27871        {
27872            let f = self.func("fa_mirror_vl");
27873            let max_n = (max_t as i64) * ept as i64;
27874            let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
27875            for which in 0..2i32 {
27876                let cfg = LaunchConfig {
27877                    grid_dim: (blocks, 1, b as u32),
27878                    block_dim: (256, 1, 1),
27879                    shared_mem_bytes: 0,
27880                };
27881                let __s_lb = self.gpu.stream();
27882                let mut lb = __s_lb.launch_builder(&f);
27883                lb.arg(&v).arg(&ept).arg(&which);
27884                unsafe {
27885                    lb.launch(cfg)?;
27886                }
27887            }
27888        }
27889        let hd_sfx = fa_hd_suffix(head_dim)?;
27890        let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
27891        let block_q = 64usize;
27892        let kv_stages = 2usize;
27893        let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
27894            + 4 * (block_q * BK + 2 * block_q)) as u32;
27895        use cudarc::driver::sys::CUfunction_attribute_enum as A;
27896        f.set_attribute(
27897            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27898            shmem as i32,
27899        )?;
27900        let cfg = LaunchConfig {
27901            grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
27902            block_dim: (32, 4, 1),
27903            shared_mem_bytes: shmem,
27904        };
27905        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27906        let __s_lb = self.gpu.stream();
27907        let mut lb = __s_lb.launch_builder(&f);
27908        lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
27909        unsafe {
27910            lb.launch(cfg)?;
27911        }
27912        Ok(())
27913    }
27914
27915    /// task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches
27916    /// for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view
27917    /// inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).
27918    #[allow(clippy::too_many_arguments)]
27919    pub fn attn_pre_vl8(
27920        &self,
27921        seqs: &[AttnPreVl],
27922        wq: &CudaSlice<f32>,
27923        wk: &CudaSlice<f32>,
27924        head_dim: usize,
27925        rope_dims: usize,
27926        n_head: usize,
27927        n_head_kv: usize,
27928        eps: f32,
27929        freq_base: f32,
27930        freq_scale: f32,
27931        kv_dim_k: usize,
27932        kv_dim_v: usize,
27933        k_tok_bytes: usize,
27934        v_tok_bytes: usize,
27935    ) -> Result<(), Box<dyn std::error::Error>> {
27936        let b = seqs.len();
27937        assert!((1..=8).contains(&b));
27938        let mut packed = [AttnPreVl::default(); 8];
27939        packed[..b].copy_from_slice(seqs);
27940        let v = AttnPreVl8(packed);
27941        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
27942        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27943        {
27944            let f = self.func("q_gate_split_vl");
27945            let n = max_t * (n_head * head_dim) as u32;
27946            let cfg = LaunchConfig {
27947                grid_dim: (n.div_ceil(256), 1, b as u32),
27948                block_dim: (256, 1, 1),
27949                shared_mem_bytes: 0,
27950            };
27951            let __s_lb = self.gpu.stream();
27952            let mut lb = __s_lb.launch_builder(&f);
27953            lb.arg(&v).arg(&hd).arg(&nh);
27954            unsafe {
27955                lb.launch(cfg)?;
27956            }
27957        }
27958        {
27959            let f = self.func("attn_rms_vl");
27960            let cfg = LaunchConfig {
27961                grid_dim: (max_t * n_head as u32, 2, b as u32),
27962                block_dim: (rms_block(), 1, 1),
27963                shared_mem_bytes: 0,
27964            };
27965            let __s_lb = self.gpu.stream();
27966            let mut lb = __s_lb.launch_builder(&f);
27967            lb.arg(&v)
27968                .arg(wq)
27969                .arg(wk)
27970                .arg(&hd)
27971                .arg(&nh)
27972                .arg(&nhkv)
27973                .arg(&eps);
27974            unsafe {
27975                lb.launch(cfg)?;
27976            }
27977        }
27978        {
27979            let f = self.func("attn_rope_vl");
27980            let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
27981            let nd = rope_dims as i32;
27982            let cfg = LaunchConfig {
27983                grid_dim: (max_t * n_head as u32, 2, b as u32),
27984                block_dim: ((head_dim / 2) as u32, 1, 1),
27985                shared_mem_bytes: 0,
27986            };
27987            let __s_lb = self.gpu.stream();
27988            let mut lb = __s_lb.launch_builder(&f);
27989            lb.arg(&v)
27990                .arg(&hd)
27991                .arg(&nd)
27992                .arg(&nh)
27993                .arg(&nhkv)
27994                .arg(&theta_scale)
27995                .arg(&freq_scale);
27996            unsafe {
27997                lb.launch(cfg)?;
27998            }
27999        }
28000        {
28001            let f = self.func("append_kv_vl");
28002            let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
28003            let cfg = LaunchConfig {
28004                grid_dim: (nblk, max_t, b as u32),
28005                block_dim: (32, 1, 1),
28006                shared_mem_bytes: 0,
28007            };
28008            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
28009            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28010            let __s_lb = self.gpu.stream();
28011            let mut lb = __s_lb.launch_builder(&f);
28012            lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
28013            unsafe {
28014                lb.launch(cfg)?;
28015            }
28016        }
28017        Ok(())
28018    }
28019
28020    /// FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify
28021    /// path, MTP-PLAN §D.3). Uses `fa_prefill_q` (inline-dequant during stage-to-smem). The view's
28022    /// base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query
28023    /// rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.
28024    #[allow(clippy::too_many_arguments)]
28025    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
28026    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28027    pub fn fa_prefill_view(
28028        &self,
28029        q: &CudaSlice<f32>,
28030        k: &cudarc::driver::CudaView<u8>,
28031        v: &cudarc::driver::CudaView<u8>,
28032        o: &mut CudaSlice<f32>,
28033        head_dim: usize,
28034        n_head: usize,
28035        n_head_kv: usize,
28036        t: usize,
28037        t_kv: usize,
28038        scale: f32,
28039        causal: bool,
28040        k_tok_bytes: usize,
28041        v_tok_bytes: usize,
28042        g: bool,
28043    ) -> Result<(), Box<dyn std::error::Error>> {
28044        if portable_mma_gated() {
28045            return self.sdpa_naive_quantized_view(
28046                q,
28047                k,
28048                v,
28049                o,
28050                head_dim,
28051                n_head,
28052                n_head_kv,
28053                t,
28054                t_kv,
28055                scale,
28056                causal,
28057                k_tok_bytes,
28058                v_tok_bytes,
28059            );
28060        }
28061        const BLOCK_Q: usize = 64;
28062        const BK: usize = 32;
28063        // g = e4m3 cache: the kernel parses via DQ_K_ELEM/DQ_V_ELEM (format macros) — the
28064        // kf8vf8-module stamp reads fp8 with the identical MMA/softmax/PV body.
28065        let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
28066        let f = if g {
28067            self.func_g(&name)
28068        } else {
28069            self.func(&name)
28070        };
28071        let shmem =
28072            (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
28073        use cudarc::driver::sys::CUfunction_attribute_enum as A;
28074        f.set_attribute(
28075            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28076            shmem as i32,
28077        )?;
28078        let cfg = LaunchConfig {
28079            grid_dim: (
28080                (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
28081                n_head as u32,
28082                1,
28083            ),
28084            block_dim: (32, 4, 1),
28085            shared_mem_bytes: shmem,
28086        };
28087        let (hd, nh, nhkv, ti, tkvi, cz) = (
28088            head_dim as i32,
28089            n_head as i32,
28090            n_head_kv as i32,
28091            t as i32,
28092            t_kv as i32,
28093            causal as i32,
28094        );
28095        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28096        let __s_b = self.gpu.stream();
28097        let mut b = __s_b.launch_builder(&f);
28098        b.arg(q)
28099            .arg(k)
28100            .arg(v)
28101            .arg(o)
28102            .arg(&hd)
28103            .arg(&nh)
28104            .arg(&nhkv)
28105            .arg(&ti)
28106            .arg(&tkvi)
28107            .arg(&scale)
28108            .arg(&cz)
28109            .arg(&ktb)
28110            .arg(&vtb);
28111        unsafe {
28112            b.launch(cfg)?;
28113        }
28114        Ok(())
28115    }
28116
28117    /// ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as `fa_prefill_view`, but
28118    /// instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline
28119    /// (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full
28120    /// [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then
28121    /// run `fa_prefill_qw` (the bf16-workspace twin) over it. EXACT: the workspace holds the same
28122    /// __float2bfloat16(dq_*_elem(...)) values fa_prefill_q stages to smem, and the twin's MMA/
28123    /// softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0).
28124    /// The workspace allocation is REUSED across layers/chunks (grown to the largest shape);
28125    /// contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).
28126    #[allow(clippy::too_many_arguments)]
28127    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28128    pub fn fa_prefill_view_ws(
28129        &self,
28130        q: &CudaSlice<f32>,
28131        k: &cudarc::driver::CudaView<u8>,
28132        v: &cudarc::driver::CudaView<u8>,
28133        o: &mut CudaSlice<f32>,
28134        head_dim: usize,
28135        n_head: usize,
28136        n_head_kv: usize,
28137        t: usize,
28138        t_kv: usize,
28139        scale: f32,
28140        causal: bool,
28141        k_tok_bytes: usize,
28142        v_tok_bytes: usize,
28143        g: bool,
28144    ) -> Result<(), Box<dyn std::error::Error>> {
28145        if portable_mma_gated() {
28146            return self.sdpa_naive_quantized_view(
28147                q,
28148                k,
28149                v,
28150                o,
28151                head_dim,
28152                n_head,
28153                n_head_kv,
28154                t,
28155                t_kv,
28156                scale,
28157                causal,
28158                k_tok_bytes,
28159                v_tok_bytes,
28160            );
28161        }
28162        const BLOCK_Q: usize = 64;
28163        const BK: usize = 32;
28164        let kv_dim_k = n_head_kv * head_dim;
28165        let kv_dim_v = n_head_kv * head_dim;
28166        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
28167        let v_ws_bytes = t_kv * kv_dim_v * 2;
28168        // Lock held across BOTH launches: enqueue-only (µs), all compute serializes on gpu.stream.
28169        let mut guard = self.prime_deqw_ws.lock().unwrap();
28170        let need_grow = match guard.as_ref() {
28171            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
28172            None => true,
28173        };
28174        if need_grow {
28175            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
28176            let (ck, cv) = guard
28177                .as_ref()
28178                .map(|(a, b)| (a.len(), b.len()))
28179                .unwrap_or((0, 0));
28180            *guard = Some((
28181                self.alloc_u8(grow(ck, k_ws_bytes))?,
28182                self.alloc_u8(grow(cv, v_ws_bytes))?,
28183            ));
28184        }
28185        let (kw, vw) = guard.as_mut().unwrap();
28186        // pass 1: dequant K+V once into the bf16 workspace (grid-stride, 1 thread/elem)
28187        {
28188            // only THIS pass parses KV bytes — pass 2 reads the bf16 workspace (format-free).
28189            let f = if g {
28190                self.func_g("fa_dequant_kv_ws_bf16")
28191            } else {
28192                self.func("fa_dequant_kv_ws_bf16")
28193            };
28194            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
28195            #[allow(clippy::manual_div_ceil)]
28196            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28197            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
28198            let cfg = LaunchConfig {
28199                grid_dim: (nblk.max(1), 1, 1),
28200                block_dim: (256, 1, 1),
28201                shared_mem_bytes: 0,
28202            };
28203            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
28204            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28205            let __s_b = self.gpu.stream();
28206            let mut b = __s_b.launch_builder(&f);
28207            b.arg(k)
28208                .arg(v)
28209                .arg(&mut *kw)
28210                .arg(&mut *vw)
28211                .arg(&kdk)
28212                .arg(&kdv)
28213                .arg(&tkvi)
28214                .arg(&ktb)
28215                .arg(&vtb);
28216            unsafe {
28217                b.launch(cfg)?;
28218            }
28219        }
28220        // pass 2: the bf16-workspace prefill twin (same tile sizes/loop structure as fa_prefill_q).
28221        // DEFAULT: cp.async double-buffered staging twin (fa_prefill_qw_db, +32KB smem for the
28222        // second K/V tile pair, 1 CTA/SM): overlaps tile n+1's L2->smem copy with tile n's MMA.
28223        // Bit-identical output (staging is a pure byte copy; kernel_check pins bitdiff=0 under
28224        // both twins). A/B (27B rtx6000, N=3): 32k prime 17.10->16.51s, 16k 9.09->8.65s — the copy
28225        // latency hides behind the MMA pipe and beats the 2-CTA/SM occupancy of the sync twin.
28226        // MEMRA_PRIME_DEQW_DB=0 falls back to the single-buffer twin.
28227        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
28228            .map(|v| v != "0")
28229            .unwrap_or(true);
28230        {
28231            let hd_sfx = fa_hd_suffix(head_dim)?;
28232            let f = self.func(&format!(
28233                "fa_prefill_qw{}{hd_sfx}",
28234                if db { "_db" } else { "" }
28235            ));
28236            let shmem = if db {
28237                // 4x KV tile buffers (bf16) + sP (bf16) + sL (f32)
28238                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
28239            } else {
28240                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
28241            };
28242            use cudarc::driver::sys::CUfunction_attribute_enum as A;
28243            f.set_attribute(
28244                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28245                shmem as i32,
28246            )?;
28247            let cfg = LaunchConfig {
28248                grid_dim: (
28249                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
28250                    n_head as u32,
28251                    1,
28252                ),
28253                block_dim: (32, 4, 1),
28254                shared_mem_bytes: shmem,
28255            };
28256            let (hd, nh, nhkv, ti, tkvi, cz) = (
28257                head_dim as i32,
28258                n_head as i32,
28259                n_head_kv as i32,
28260                t as i32,
28261                t_kv as i32,
28262                causal as i32,
28263            );
28264            let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
28265            let __s_b = self.gpu.stream();
28266            let mut b = __s_b.launch_builder(&f);
28267            b.arg(q)
28268                .arg(&*kw)
28269                .arg(&*vw)
28270                .arg(o)
28271                .arg(&hd)
28272                .arg(&nh)
28273                .arg(&nhkv)
28274                .arg(&ti)
28275                .arg(&tkvi)
28276                .arg(&scale)
28277                .arg(&cz)
28278                .arg(&kdk)
28279                .arg(&kdv);
28280            unsafe {
28281                b.launch(cfg)?;
28282            }
28283        }
28284        Ok(())
28285    }
28286
28287    /// WINDOWED `fa_prefill_view_ws` twin at head_dim 128 (lane/pp-prefill 2026-08-07):
28288    /// step35's SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
28289    /// stamp — every windowed twin was hd256-only — and took `sdpa_naive_w_quantized_view`,
28290    /// the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
28291    /// causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
28292    /// Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
28293    /// bf16 workspace, then the windowed qw kernel (`fa_prefill_qw_db_w_hd128`, cp.async
28294    /// double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
28295    /// mask is `fa_prefill_f32_body`'s exact predicate; `window == 0` is bit-identical to
28296    /// `fa_prefill_view_ws` by construction (default-arg body). NEW NUMERIC CLASS vs the
28297    /// f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
28298    /// gated by the full battery, and the class must change UNIFORMLY for a whole request
28299    /// (kernel selection keys on seq_end, never per chunk — the chunkfix law).
28300    /// hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
28301    /// gemma4 (hd256), which already has `fa_prefill_w_f32`.
28302    #[allow(clippy::too_many_arguments)]
28303    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28304    pub fn fa_prefill_view_ws_w_hd128(
28305        &self,
28306        q: &CudaSlice<f32>,
28307        k: &cudarc::driver::CudaView<u8>,
28308        v: &cudarc::driver::CudaView<u8>,
28309        o: &mut CudaSlice<f32>,
28310        head_dim: usize,
28311        n_head: usize,
28312        n_head_kv: usize,
28313        t: usize,
28314        t_kv: usize,
28315        scale: f32,
28316        causal: bool,
28317        window: usize,
28318        k_tok_bytes: usize,
28319        v_tok_bytes: usize,
28320    ) -> Result<(), Box<dyn std::error::Error>> {
28321        assert_eq!(
28322            head_dim, 128,
28323            "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
28324        );
28325        if portable_mma_gated() {
28326            return self.sdpa_naive_w_quantized_view(
28327                q,
28328                k,
28329                v,
28330                o,
28331                head_dim,
28332                n_head,
28333                n_head_kv,
28334                t,
28335                t_kv,
28336                scale,
28337                causal,
28338                window,
28339                k_tok_bytes,
28340                v_tok_bytes,
28341            );
28342        }
28343        const BLOCK_Q: usize = 64;
28344        const BK: usize = 32;
28345        let kv_dim_k = n_head_kv * head_dim;
28346        let kv_dim_v = n_head_kv * head_dim;
28347        let k_ws_bytes = t_kv * kv_dim_k * 2; // bf16
28348        let v_ws_bytes = t_kv * kv_dim_v * 2;
28349        let mut guard = self.prime_deqw_ws.lock().unwrap();
28350        let need_grow = match guard.as_ref() {
28351            Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
28352            None => true,
28353        };
28354        if need_grow {
28355            let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
28356            let (ck, cv) = guard
28357                .as_ref()
28358                .map(|(a, b)| (a.len(), b.len()))
28359                .unwrap_or((0, 0));
28360            *guard = Some((
28361                self.alloc_u8(grow(ck, k_ws_bytes))?,
28362                self.alloc_u8(grow(cv, v_ws_bytes))?,
28363            ));
28364        }
28365        let (kw, vw) = guard.as_mut().unwrap();
28366        // pass 1: dequant K+V once into the bf16 workspace (identical to fa_prefill_view_ws —
28367        // the workspace bytes are the SAME __float2bfloat16(dq(...)) values either way).
28368        {
28369            let f = self.func("fa_dequant_kv_ws_bf16");
28370            let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
28371            #[allow(clippy::manual_div_ceil)]
28372            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28373            let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
28374            let cfg = LaunchConfig {
28375                grid_dim: (nblk.max(1), 1, 1),
28376                block_dim: (256, 1, 1),
28377                shared_mem_bytes: 0,
28378            };
28379            let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
28380            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28381            let __s_b = self.gpu.stream();
28382            let mut b = __s_b.launch_builder(&f);
28383            b.arg(k)
28384                .arg(v)
28385                .arg(&mut *kw)
28386                .arg(&mut *vw)
28387                .arg(&kdk)
28388                .arg(&kdv)
28389                .arg(&tkvi)
28390                .arg(&ktb)
28391                .arg(&vtb);
28392            unsafe {
28393                b.launch(cfg)?;
28394            }
28395        }
28396        // pass 2: the WINDOWED qw twin (db default, same as the unwindowed wrapper).
28397        let db = std::env::var("MEMRA_PRIME_DEQW_DB")
28398            .map(|v| v != "0")
28399            .unwrap_or(true);
28400        {
28401            let f = self.func(if db {
28402                "fa_prefill_qw_db_w_hd128"
28403            } else {
28404                "fa_prefill_qw_w_hd128"
28405            });
28406            let shmem = if db {
28407                (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
28408            } else {
28409                (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
28410            };
28411            use cudarc::driver::sys::CUfunction_attribute_enum as A;
28412            f.set_attribute(
28413                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28414                shmem as i32,
28415            )?;
28416            let cfg = LaunchConfig {
28417                grid_dim: (
28418                    (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
28419                    n_head as u32,
28420                    1,
28421                ),
28422                block_dim: (32, 4, 1),
28423                shared_mem_bytes: shmem,
28424            };
28425            let (hd, nh, nhkv, ti, tkvi, cz) = (
28426                head_dim as i32,
28427                n_head as i32,
28428                n_head_kv as i32,
28429                t as i32,
28430                t_kv as i32,
28431                causal as i32,
28432            );
28433            let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
28434            let __s_b = self.gpu.stream();
28435            let mut b = __s_b.launch_builder(&f);
28436            b.arg(q)
28437                .arg(&*kw)
28438                .arg(&*vw)
28439                .arg(o)
28440                .arg(&hd)
28441                .arg(&nh)
28442                .arg(&nhkv)
28443                .arg(&ti)
28444                .arg(&tkvi)
28445                .arg(&scale)
28446                .arg(&cz)
28447                .arg(&kdk)
28448                .arg(&kdv)
28449                .arg(&wnd);
28450            unsafe {
28451                b.launch(cfg)?;
28452            }
28453        }
28454        Ok(())
28455    }
28456
28457    /// FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views.
28458    /// Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are
28459    /// the per-token byte strides (differ: q8_0=34*nblk, q5_1=24*nblk per token).
28460    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
28461    pub fn fa_decode(
28462        &self,
28463        q: &CudaSlice<f32>,
28464        k: &cudarc::driver::CudaView<u8>,
28465        v: &cudarc::driver::CudaView<u8>,
28466        o: &mut CudaSlice<f32>,
28467        head_dim: usize,
28468        n_head: usize,
28469        n_head_kv: usize,
28470        t_kv: usize,
28471        scale: f32,
28472        k_tok_bytes: usize,
28473        v_tok_bytes: usize,
28474    ) -> Result<(), Box<dyn std::error::Error>> {
28475        self.fa_decode_kvmod(
28476            q,
28477            k,
28478            v,
28479            o,
28480            head_dim,
28481            n_head,
28482            n_head_kv,
28483            t_kv,
28484            scale,
28485            k_tok_bytes,
28486            v_tok_bytes,
28487            false,
28488        )
28489    }
28490
28491    /// `fa_decode` with an explicit fp8-module flag (`g`): gemma windowed layers under
28492    /// MEMRA_GEMMA_WKV read an e4m3 cache — every kernel must come from the kf8vf8 module
28493    /// and the v4 lane (q8_0-hardcoded staging) is excluded.
28494    #[allow(clippy::too_many_arguments)]
28495    /// UNIFIED scalar decode launch (fa_decode_f32, nullable-ctr): ONE symbol for host-len
28496    /// (kvmod eager) and device-len (graph/stream) callers — the textually-identical f32_dc
28497    /// twin compiled apart and its ULP drift flipped 31B verify argmaxes (2026-07-12).
28498    #[allow(clippy::too_many_arguments)]
28499    #[allow(clippy::too_many_arguments)]
28500    fn fa_decode_scalar_unified(
28501        &self,
28502        q: &cudarc::driver::CudaView<f32>,
28503        k: &cudarc::driver::CudaView<u8>,
28504        v: &cudarc::driver::CudaView<u8>,
28505        o: &mut cudarc::driver::CudaViewMut<f32>,
28506        head_dim: usize,
28507        n_head: usize,
28508        n_head_kv: usize,
28509        t_kv_host: usize,
28510        t_kv_dev: Option<&CudaSlice<i32>>,
28511        scale: f32,
28512        n_splits: usize,
28513        split_keys: usize,
28514        k_tok_bytes: usize,
28515        v_tok_bytes: usize,
28516        g: bool,
28517        part_o: &mut CudaSlice<f32>,
28518        part_m: &mut CudaSlice<f32>,
28519        part_l: &mut CudaSlice<f32>,
28520        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
28521    ) -> Result<(), Box<dyn std::error::Error>> {
28522        let f = if g {
28523            self.func_g("fa_decode_f32")
28524        } else {
28525            self.fa_func("fa_decode_f32", head_dim)
28526        };
28527        let cfg = LaunchConfig {
28528            grid_dim: (n_head as u32, n_splits as u32, 1),
28529            block_dim: (head_dim as u32, 1, 1),
28530            shared_mem_bytes: (4 * (head_dim + 32)) as u32,
28531        };
28532        let (hd, nh, nhkv, nsp) = (
28533            head_dim as i32,
28534            n_head as i32,
28535            n_head_kv as i32,
28536            n_splits as i32,
28537        );
28538        let (ktb, vtb, tkvi, ski) = (
28539            k_tok_bytes as i64,
28540            v_tok_bytes as i64,
28541            t_kv_host as i32,
28542            split_keys as i32,
28543        );
28544        let __s_b = self.gpu.stream();
28545        let mut b = __s_b.launch_builder(&f);
28546        match t_kv_dev {
28547            Some(d) => {
28548                b.arg(q)
28549                    .arg(k)
28550                    .arg(v)
28551                    .arg(&mut *part_o)
28552                    .arg(&mut *part_m)
28553                    .arg(&mut *part_l)
28554                    .arg(&hd)
28555                    .arg(&nh)
28556                    .arg(&nhkv)
28557                    .arg(&tkvi)
28558                    .arg(d)
28559                    .arg(&scale)
28560                    .arg(&nsp)
28561                    .arg(&ski)
28562                    .arg(&ktb)
28563                    .arg(&vtb);
28564                unsafe {
28565                    b.launch(cfg)?;
28566                }
28567            }
28568            None => {
28569                let null: u64 = 0;
28570                b.arg(q)
28571                    .arg(k)
28572                    .arg(v)
28573                    .arg(&mut *part_o)
28574                    .arg(&mut *part_m)
28575                    .arg(&mut *part_l)
28576                    .arg(&hd)
28577                    .arg(&nh)
28578                    .arg(&nhkv)
28579                    .arg(&tkvi)
28580                    .arg(&null)
28581                    .arg(&scale)
28582                    .arg(&nsp)
28583                    .arg(&ski)
28584                    .arg(&ktb)
28585                    .arg(&vtb);
28586                unsafe {
28587                    b.launch(cfg)?;
28588                }
28589            }
28590        }
28591        let cfg2 = LaunchConfig {
28592            grid_dim: (n_head as u32, 1, 1),
28593            block_dim: (head_dim as u32, 1, 1),
28594            shared_mem_bytes: 0,
28595        };
28596        if let Some((oq, od)) = q8_out {
28597            // wave-5b: q8-emitting combine — the wo matmul_pre consumes the pair directly.
28598            let fc = if g {
28599                self.func_g("fa_decode_combine_q8_1")
28600            } else {
28601                self.fa_func("fa_decode_combine_q8_1", head_dim)
28602            };
28603            let __s_b2 = self.gpu.stream();
28604            let mut b2 = __s_b2.launch_builder(&fc);
28605            b2.arg(&*part_o)
28606                .arg(&*part_m)
28607                .arg(&*part_l)
28608                .arg(oq)
28609                .arg(od)
28610                .arg(&hd)
28611                .arg(&nh)
28612                .arg(&nsp);
28613            unsafe {
28614                b2.launch(cfg2)?;
28615            }
28616            return Ok(());
28617        }
28618        let fc = if g {
28619            self.func_g("fa_decode_combine_f32")
28620        } else {
28621            self.fa_func("fa_decode_combine_f32", head_dim)
28622        };
28623        let __s_b2 = self.gpu.stream();
28624        let mut b2 = __s_b2.launch_builder(&fc);
28625        b2.arg(&*part_o)
28626            .arg(&*part_m)
28627            .arg(&*part_l)
28628            .arg(o)
28629            .arg(&hd)
28630            .arg(&nh)
28631            .arg(&nsp);
28632        unsafe {
28633            b2.launch(cfg2)?;
28634        }
28635        Ok(())
28636    }
28637
28638    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
28639    pub fn fa_decode_kvmod(
28640        &self,
28641        q: &CudaSlice<f32>,
28642        k: &cudarc::driver::CudaView<u8>,
28643        v: &cudarc::driver::CudaView<u8>,
28644        o: &mut CudaSlice<f32>,
28645        head_dim: usize,
28646        n_head: usize,
28647        n_head_kv: usize,
28648        t_kv: usize,
28649        scale: f32,
28650        k_tok_bytes: usize,
28651        v_tok_bytes: usize,
28652        g: bool,
28653    ) -> Result<(), Box<dyn std::error::Error>> {
28654        let q_view = q.as_view();
28655        let mut o_view = o.as_view_mut();
28656        self.fa_decode_kvmod_view(
28657            &q_view,
28658            k,
28659            v,
28660            &mut o_view,
28661            head_dim,
28662            n_head,
28663            n_head_kv,
28664            t_kv,
28665            scale,
28666            k_tok_bytes,
28667            v_tok_bytes,
28668            g,
28669        )
28670    }
28671
28672    /// Row-view entry into `fa_decode_kvmod`. The kernel sees the selected Q/output rows as its
28673    /// base pointers, so the launch geometry and arithmetic are identical to the owned-slice entry.
28674    /// Batched fallback callers use this to avoid materializing rows around an otherwise unchanged
28675    /// per-session KV view and FA launch.
28676    #[allow(clippy::too_many_arguments)]
28677    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
28678    pub fn fa_decode_kvmod_view(
28679        &self,
28680        q: &cudarc::driver::CudaView<f32>,
28681        k: &cudarc::driver::CudaView<u8>,
28682        v: &cudarc::driver::CudaView<u8>,
28683        o: &mut cudarc::driver::CudaViewMut<f32>,
28684        head_dim: usize,
28685        n_head: usize,
28686        n_head_kv: usize,
28687        t_kv: usize,
28688        scale: f32,
28689        k_tok_bytes: usize,
28690        v_tok_bytes: usize,
28691        g: bool,
28692    ) -> Result<(), Box<dyn std::error::Error>> {
28693        // PERF-4: the warp-per-token vec path replaces the scalar element-per-thread fa_decode_f32 —
28694        // warp-per-token fa_decode_vec_q (grid=(n_head_kv,n_splits), block=(32,gqa_ratio)).
28695        // The block dequants each KV tile ONCE into smem (bf16) and broadcasts to all gqa Q-head
28696        // warps -> each KV byte leaves HBM/L2 ~1x/group (vs 4x). ARGS identical; func/grid/block/
28697        // smem/n_splits differ. fa_decode_f32 stays the bit-reference fallback. Combine is shared.
28698        //
28699        // SPLIT-K: the scalar path has grid.x=n_head (32) blocks; the vec path only has
28700        // grid.x=n_head_kv (8). To avoid starving the GPU at mid ctx, the vec path splits MORE
28701        // aggressively (64 keys/split vs 256) so grid.y rises and 8*n_splits fills the SMs.
28702        // At VERY short ctx (t_kv<96) even 1 split can't fill the GPU from 8 KV heads, so the
28703        // broadcast can't beat the scalar path's 4x-more-blocks latency hiding — fall back to
28704        // scalar there (measured crossover: vec 0.68x at t_kv=64, 1.23x at t_kv=96, 2.2x at 256).
28705        // DEFAULT-ON (2026-06-28): clean clock-locked sweep proved vec beats scalar at every
28706        // t_kv>=96 and the gain WIDENS with ctx (graph decode: +9.5% @128, +11.6% @512, +11.8%
28707        // @2048) — the KV-byte-broadcast (4x fewer HBM reads/group) compounds as attention grows.
28708        // MEMRA_NO_FA_VEC forces the scalar bit-reference. Below FA_VEC_MIN_TKV the scalar path's
28709        // 4x-more-blocks (grid.x=n_head=32 vs n_head_kv=8) hides latency better, so keep scalar there.
28710        // g + no-v4: the g-module REGISTER twin mis-decodes the gemma windowed shape
28711        // (root-cause open, jsonl) — only reachable by forcing v4 off (MEMRA_FA_V4_MAX);
28712        // fall to the exact scalar there instead of the broken register arm.
28713        let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
28714        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
28715        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
28716        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
28717        if g && head_dim == 256 && !fa_v4_at(t_kv) {
28718            fa_vec = false;
28719        }
28720        let sp = fa_split_keys(t_kv, n_head_kv);
28721        let n_splits = if fa_vec {
28722            ((t_kv + sp - 1) / sp).max(1)
28723        } else {
28724            ((t_kv + 255) / 256).max(1)
28725        };
28726        let o_len = n_head * n_splits * head_dim;
28727        let ml_len = n_head * n_splits;
28728        let mut part_guard = self.fa_part_pool.lock().unwrap();
28729        if part_guard
28730            .as_ref()
28731            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
28732            .unwrap_or(true)
28733        {
28734            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
28735            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
28736            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
28737            // later live allocations land at those addresses, and the next graph REPLAY writes
28738            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
28739            // output corruption began the burst after the trunk's t_kv growth first realloc'd
28740            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
28741            // the baked addresses alive (single-stream: eager writes the new buffers, replays
28742            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
28743            // (total retired < final size).
28744            let old = part_guard.take();
28745            let (co, cm) = old
28746                .as_ref()
28747                .map(|pp| (pp.0.len(), pp.1.len()))
28748                .unwrap_or((0, 0));
28749            if let Some(old) = old {
28750                self.fa_part_retired.lock().unwrap().push(old);
28751            }
28752            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
28753                eprintln!(
28754                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
28755                    co, o_len, cm, ml_len
28756                );
28757            }
28758            *part_guard =
28759                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
28760        }
28761        let pg = part_guard.as_mut().unwrap();
28762        self.gpu
28763            .stream()
28764            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
28765        self.gpu
28766            .stream()
28767            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
28768        self.gpu
28769            .stream()
28770            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
28771        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28772        let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
28773        let (hd, nh, nhkv, tkvi, nsp) = (
28774            head_dim as i32,
28775            n_head as i32,
28776            n_head_kv as i32,
28777            t_kv as i32,
28778            n_splits as i32,
28779        );
28780        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28781        // The vec kernel holds head_dim/32 register accumulators (FA_DEC_MAX_DPL=8 -> head_dim<=256).
28782        // All shipped models use head_dim=256; fall back to scalar for anything wider rather than
28783        // silently truncating the accumulator.
28784        let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
28785        // hd-512 vec crossover (MEMRA_FA512_MIN, default 512): the DPL16 twin wins at depth
28786        // (82.5 -> vec at 1736) but the scalar's more-blocks latency hiding wins at tiny t_kv
28787        // (the same scalar-floor physics as hd256's old 96 floor; short-ctx plain regressed
28788        // 178.4 -> 173.7 when 512 rode vec unconditionally).
28789        let fa512_min = fa512_min_tkv();
28790        // FA-DEEP pick (bit-identical twins, see fa_deep_at): default module only — the
28791        // g-module keeps the v4 pick (its class is not the depth-decay class).
28792        let deep = fa_vec
28793            && head_dim == 256
28794            && fa_v4_at(t_kv)
28795            && !g
28796            && fa_deep_at(t_kv)
28797            && !matches!(fa_v4_mode(), "noB3" | "stage");
28798        let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
28799            // gemma4 globals (hd 512): the DPL16 register twin (fa_decode_vec_q body with a
28800            // 16-slot accumulator ceiling). Scalar fallback measured 82.5us/layer at 1736 ctx.
28801            let gqa = (n_head / n_head_kv).max(1) as u32;
28802            let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
28803            (
28804                fv,
28805                LaunchConfig {
28806                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28807                    block_dim: (32, gqa, 1),
28808                    shared_mem_bytes: 0,
28809                },
28810            )
28811        } else if fa_vec && head_dim <= 256 {
28812            let gqa = (n_head / n_head_kv).max(1) as u32;
28813            // DEEP-CTX smem twin (2026-07-05): the register-dequant path's GQA reuse rides L2,
28814            // which holds to ~8k ctx but dies at 40k (layer KV ~37MB) — the 4 GQA warps then
28815            // re-read every KV byte from DRAM (4x traffic). Above MEMRA_FA_SMEM_TKV (default
28816            // 1024 — the 2026-07-05 crossover re-sweep on real prompts: p3 spec 73.8->79.2 at
28817            // 2048, flat down to 512, p2 +5%, p1/9B unchanged; the ARC-A probe's synthetic
28818            // 2.1x smem-at-all-depths pointed here; 0=never) dispatch the smem-broadcast twin:
28819            // dequant each tile ONCE per block.
28820            // Bit-identical per (token,split): same bf16 round-trip, same accumulation order,
28821            // same partial layout -> same combine. Short/mid ctx keeps the register path (it won
28822            // there by 12x — latency, not bandwidth, rules small KV).
28823            static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
28824            let smem_tkv = *SMEM_TKV.get_or_init(|| {
28825                std::env::var("MEMRA_FA_SMEM_TKV")
28826                    .ok()
28827                    .and_then(|v| v.parse().ok())
28828                    .unwrap_or_else(|| {
28829                        FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
28830                    })
28831            });
28832            if fa_v4_at(t_kv) && head_dim == 256 {
28833                // FA v4 lane (2026-07-10): key-per-lane score phase, zero shuffles per key.
28834                // NEW NUMERIC CONFIG (chunk-serial per-key dot) — battery-arbitrated.
28835                // g (fp8-windowed): the v4 staging is format-aware (2026-07-12) — kf8vf8 module.
28836                let v4name = match fa_v4_mode() {
28837                    "noB3" => "fa_decode_vec_q_v4_noB3", // phase probe (WRONG OUTPUT)
28838                    "stage" => "fa_decode_vec_q_v4_stage", // phase probe (WRONG OUTPUT)
28839                    _ if deep => "fa_decode_vec_q_v4_deep",
28840                    _ => "fa_decode_vec_q_v4",
28841                };
28842                let fv = if g {
28843                    self.func_g(v4name)
28844                } else {
28845                    self.func(v4name)
28846                };
28847                // fa_v4_smem (deep: fa_v4_deep_smem, +640B row pads) + sV (g: raw e4m3 sV
28848                // tile = 1B/elem — half the smem, 3->5 blocks/SM)
28849                let shmem = (if deep { 12160 } else { 11520 }
28850                    + 32 * head_dim * if g { 1 } else { 2 }) as u32;
28851                use cudarc::driver::sys::CUfunction_attribute_enum as A;
28852                fv.set_attribute(
28853                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28854                    shmem as i32,
28855                )?;
28856                (
28857                    fv,
28858                    LaunchConfig {
28859                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28860                        block_dim: (32, gqa, 1),
28861                        shared_mem_bytes: shmem,
28862                    },
28863                )
28864            } else if fa_v3_active(head_dim) {
28865                // FA v3 lane: dp4a-K hybrid (register-quantized Q, raw q8_0 K, staged-V kept).
28866                // smem = sV only (half of v2's).
28867                let fv = if g {
28868                    self.func_g("fa_decode_vec_q_v3")
28869                } else {
28870                    self.func("fa_decode_vec_q_v3")
28871                };
28872                let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
28873                (
28874                    fv,
28875                    LaunchConfig {
28876                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28877                        block_dim: (32, gqa, 1),
28878                        shared_mem_bytes: shmem,
28879                    },
28880                )
28881            } else if fa_v2_on() {
28882                // FAVENDOR lane: llama fattn-vec tile-batched softmax + wide-load staging on
28883                // OUR smem KV broadcast. Replaces BOTH per-key twins when on; same grid/block/
28884                // partials; same 32KB sK+sV tile as the smem twin.
28885                let fv = if g {
28886                    self.func_g("fa_decode_vec_q_v2")
28887                } else {
28888                    self.func("fa_decode_vec_q_v2")
28889                };
28890                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
28891                (
28892                    fv,
28893                    LaunchConfig {
28894                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28895                        block_dim: (32, gqa, 1),
28896                        shared_mem_bytes: shmem,
28897                    },
28898                )
28899            } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
28900            {
28901                // (fp8 exclusions: the smem twin's V-stage is q5_1-hardcoded — neither the wkv
28902                // windowed layers (g) nor the gkv globals (hd512) may be forced onto it via
28903                // MEMRA_FA_SMEM_TKV; they fall through to the format-clean register/scalar arms.)
28904                let fv = if g {
28905                    self.func_g("fa_decode_vec_q_smem")
28906                } else {
28907                    self.func("fa_decode_vec_q_smem")
28908                };
28909                let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
28910                use cudarc::driver::sys::CUfunction_attribute_enum as A;
28911                fv.set_attribute(
28912                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28913                    shmem as i32,
28914                )?;
28915                (
28916                    fv,
28917                    LaunchConfig {
28918                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28919                        block_dim: (32, gqa, 1),
28920                        shared_mem_bytes: shmem,
28921                    },
28922                )
28923            } else {
28924                // REGISTER-DEQUANT kernel (2026-07-03): per-warp direct q8_0/q5_1 register
28925                // dequant, zero dynamic shared memory.
28926                let fv = if g {
28927                    self.func_g("fa_decode_vec_q")
28928                } else {
28929                    self.func("fa_decode_vec_q")
28930                };
28931                (
28932                    fv,
28933                    LaunchConfig {
28934                        grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28935                        block_dim: (32, gqa, 1),
28936                        shared_mem_bytes: 0,
28937                    },
28938                )
28939            }
28940        } else {
28941            // UNIFIED scalar (nullable-ctr symbol shared with graph/stream callers). The
28942            // split ladder value rides along so ns_eff reproduces THIS n_splits in-kernel.
28943            return self.fa_decode_scalar_unified(
28944                q,
28945                k,
28946                v,
28947                o,
28948                head_dim,
28949                n_head,
28950                n_head_kv,
28951                t_kv,
28952                None,
28953                scale,
28954                n_splits,
28955                if fa_vec { sp } else { 256 },
28956                k_tok_bytes,
28957                v_tok_bytes,
28958                g,
28959                part_o,
28960                part_m,
28961                part_l,
28962                None,
28963            );
28964        };
28965        let __s_b = self.gpu.stream();
28966        let mut b = __s_b.launch_builder(&f);
28967        b.arg(q)
28968            .arg(k)
28969            .arg(v)
28970            .arg(&mut *part_o)
28971            .arg(&mut *part_m)
28972            .arg(&mut *part_l)
28973            .arg(&hd)
28974            .arg(&nh)
28975            .arg(&nhkv)
28976            .arg(&tkvi)
28977            .arg(&scale)
28978            .arg(&nsp)
28979            .arg(&ktb)
28980            .arg(&vtb);
28981        unsafe {
28982            b.launch(cfg)?;
28983        }
28984        // (combine re-tile refuted in the fa-deep lane — flat/worse both shapes; the v4
28985        // combine stays for all arms. Receipts research/fa-decode-deep-20260802/.)
28986        let (fc, cfg2) = (
28987            if g {
28988                self.func_g("fa_decode_combine_f32")
28989            } else {
28990                self.fa_func("fa_decode_combine_f32", head_dim)
28991            },
28992            LaunchConfig {
28993                grid_dim: (n_head as u32, 1, 1),
28994                block_dim: (head_dim as u32, 1, 1),
28995                shared_mem_bytes: 0,
28996            },
28997        );
28998        let __s_b2 = self.gpu.stream();
28999        let mut b2 = __s_b2.launch_builder(&fc);
29000        b2.arg(&*part_o)
29001            .arg(&*part_m)
29002            .arg(&*part_l)
29003            .arg(o)
29004            .arg(&hd)
29005            .arg(&nh)
29006            .arg(&nsp);
29007        unsafe {
29008            b2.launch(cfg2)?;
29009        }
29010        Ok(())
29011    }
29012
29013    /// BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the
29014    /// batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device
29015    /// pointer table (`kv_ptrs`, [2B] interleaved k0,v0,...); per-seq key bounds ride the
29016    /// tick's position table (`pos_seq`, T_kv = pos+1). v4-lane only: the CALLER
29017    /// (decode_batch) gates every row through `fa_seqs_eligible` AND one `fa_split_keys`
29018    /// rung (`split_keys`), so each sequence's split partition, key walk and combine order
29019    /// reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit
29020    /// identity; decode-batch-gate strict pins the whole tick vs decode_step_h).
29021    /// q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q
29022    /// copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).
29023    #[allow(clippy::too_many_arguments)]
29024    pub fn fa_decode_batch_seqs_v4(
29025        &self,
29026        q: &CudaSlice<f32>,
29027        kv_ptrs: &cudarc::driver::CudaView<u64>,
29028        pos_seq: &CudaSlice<i32>,
29029        o: &mut CudaSlice<f32>,
29030        head_dim: usize,
29031        n_head: usize,
29032        n_head_kv: usize,
29033        b_n: usize,
29034        t_kv_max: usize,
29035        scale: f32,
29036        split_keys: usize,
29037        k_tok_bytes: usize,
29038        v_tok_bytes: usize,
29039    ) -> Result<(), Box<dyn std::error::Error>> {
29040        debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
29041        #[allow(clippy::manual_div_ceil)]
29042        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29043        let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
29044        let o_len = b_n * n_head * n_splits_max * head_dim;
29045        let ml_len = b_n * n_head * n_splits_max;
29046        let mut part_guard = self.fa_part_pool.lock().unwrap();
29047        if part_guard
29048            .as_ref()
29049            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29050            .unwrap_or(true)
29051        {
29052            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
29053            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
29054            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
29055            // later live allocations land at those addresses, and the next graph REPLAY writes
29056            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
29057            // output corruption began the burst after the trunk's t_kv growth first realloc'd
29058            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
29059            // the baked addresses alive (single-stream: eager writes the new buffers, replays
29060            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
29061            // (total retired < final size).
29062            let old = part_guard.take();
29063            let (co, cm) = old
29064                .as_ref()
29065                .map(|pp| (pp.0.len(), pp.1.len()))
29066                .unwrap_or((0, 0));
29067            if let Some(old) = old {
29068                self.fa_part_retired.lock().unwrap().push(old);
29069            }
29070            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29071                eprintln!(
29072                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29073                    co, o_len, cm, ml_len
29074                );
29075            }
29076            *part_guard =
29077                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29078        }
29079        let pg = part_guard.as_mut().unwrap();
29080        self.gpu
29081            .stream()
29082            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29083        self.gpu
29084            .stream()
29085            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29086        self.gpu
29087            .stream()
29088            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29089        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29090        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29091        let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
29092        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29093        let gqa = (n_head / n_head_kv).max(1) as u32;
29094        let f = self.func("fa_decode_vec_q_seqs_v4");
29095        // fa_v4_smem (11520B) + sV bf16 tile — the v4 eager arm's sizing on the default module.
29096        let shmem = (11520 + 32 * head_dim * 2) as u32;
29097        use cudarc::driver::sys::CUfunction_attribute_enum as A;
29098        f.set_attribute(
29099            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29100            shmem as i32,
29101        )?;
29102        let cfg = LaunchConfig {
29103            grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
29104            block_dim: (32, gqa, 1),
29105            shared_mem_bytes: shmem,
29106        };
29107        {
29108            let __s_b = self.gpu.stream();
29109            let mut b = __s_b.launch_builder(&f);
29110            b.arg(q)
29111                .arg(kv_ptrs)
29112                .arg(pos_seq)
29113                .arg(&mut *part_o)
29114                .arg(&mut *part_m)
29115                .arg(&mut *part_l)
29116                .arg(&hd)
29117                .arg(&nh)
29118                .arg(&nhkv)
29119                .arg(&scale)
29120                .arg(&nspm)
29121                .arg(&spk)
29122                .arg(&ktb)
29123                .arg(&vtb);
29124            unsafe {
29125                b.launch(cfg)?;
29126            }
29127        }
29128        let fc = self.func("fa_decode_combine_seqs");
29129        let cfg2 = LaunchConfig {
29130            grid_dim: (n_head as u32, b_n as u32, 1),
29131            block_dim: (head_dim as u32, 1, 1),
29132            shared_mem_bytes: 0,
29133        };
29134        let __s_b2 = self.gpu.stream();
29135        let mut b2 = __s_b2.launch_builder(&fc);
29136        b2.arg(&*part_o)
29137            .arg(&*part_m)
29138            .arg(&*part_l)
29139            .arg(o)
29140            .arg(&hd)
29141            .arg(&nh)
29142            .arg(pos_seq)
29143            .arg(&nspm)
29144            .arg(&spk);
29145        unsafe {
29146            b2.launch(cfg2)?;
29147        }
29148        Ok(())
29149    }
29150
29151    /// BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this
29152    /// step's B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same
29153    /// [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the
29154    /// per-token appender's exact warp program on row z of the stacked [B, kv_dim] k/v —
29155    /// written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces
29156    /// (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).
29157    #[allow(clippy::too_many_arguments)]
29158    pub fn append_kv_quantized_seqs(
29159        &self,
29160        k_rows: &CudaSlice<f32>,
29161        v_rows: &CudaSlice<f32>,
29162        kv_ptrs: &cudarc::driver::CudaView<u64>,
29163        pos_seq: &CudaSlice<i32>,
29164        b_n: usize,
29165        kv_dim_k: usize,
29166        kv_dim_v: usize,
29167        k_tok_bytes: usize,
29168        v_tok_bytes: usize,
29169    ) -> Result<(), Box<dyn std::error::Error>> {
29170        let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
29171        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
29172        let cfg = LaunchConfig {
29173            grid_dim: (nblk, b_n as u32, 1),
29174            block_dim: (32, 1, 1),
29175            shared_mem_bytes: 0,
29176        };
29177        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
29178        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29179        let __s_b = self.gpu.stream();
29180        let mut b = __s_b.launch_builder(&f);
29181        b.arg(k_rows)
29182            .arg(v_rows)
29183            .arg(kv_ptrs)
29184            .arg(pos_seq)
29185            .arg(&kdk)
29186            .arg(&kdv)
29187            .arg(&ktb)
29188            .arg(&vtb);
29189        unsafe {
29190            b.launch(cfg)?;
29191        }
29192        Ok(())
29193    }
29194
29195    /// True iff the MULTI-ROW verify FA (`fa_decode_rows`) is usable for a verify batch whose
29196    /// FIRST row attends `base_len + 1` keys: every row must take the SAME kernel eager decode
29197    /// would (the vec path) — mirrors fa_decode's gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV +
29198    /// head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar.
29199    /// MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).
29200    pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
29201        std::env::var("MEMRA_NO_FA_VEC").is_err()
29202            && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
29203            && base_len + 1 >= fa_vec_min_tkv()
29204            && head_dim <= 256
29205            && head_dim.is_multiple_of(32)
29206    }
29207
29208    /// MULTI-ROW verify FA: run fa_decode_vec_q's EXACT per-row program for T causal query rows
29209    /// (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE
29210    /// row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec
29211    /// verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the
29212    /// fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs
29213    /// per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate).
29214    /// Caller must have checked `fa_rows_eligible(base_len, head_dim)`.
29215    /// q is the verify's token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].
29216    #[allow(clippy::too_many_arguments)]
29217    pub fn fa_decode_rows(
29218        &self,
29219        q: &CudaSlice<f32>,
29220        k: &cudarc::driver::CudaView<u8>,
29221        v: &cudarc::driver::CudaView<u8>,
29222        o: &mut CudaSlice<f32>,
29223        head_dim: usize,
29224        n_head: usize,
29225        n_head_kv: usize,
29226        base_len: usize,
29227        t: usize,
29228        scale: f32,
29229        k_tok_bytes: usize,
29230        v_tok_bytes: usize,
29231        // hd512 dpl16 twin is DEVICE-LEN (graph arc): base_dev/plus feed the
29232        // kernel; host base_len keeps sizing the splits/partials. hd256 twins
29233        // keep the host arg. None is a bug for hd512 (asserted below).
29234        base_dev: Option<(&CudaSlice<i32>, i32)>,
29235        // K and V planes hold the same values (gemma globals, wv:=wk): pick
29236        // the _kv twin — V plane never read, value rides the q8_0 key dq.
29237        kv_shared: bool,
29238        // this layer's cache is e4m3 (gemma windowed under wkv): resolve the
29239        // hd256 rows kernel from the kf8vf8 module. PER-CALL — a global env
29240        // check here hijacked qwen/kernel-check hd256 rows (8 FAILs, 230ebbe).
29241        g: bool,
29242        // t=1 decode arm only: emit (int8, per-32 scales) from the dc combine
29243        // (hd512 path) — the standalone quantize launch folds away.
29244        mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
29245    ) -> Result<(), Box<dyn std::error::Error>> {
29246        debug_assert!(
29247            base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim.is_multiple_of(32)
29248        );
29249        let t_kv_max = base_len + t; // LAST row's key bound
29250        let mut sp = fa_split_keys(t_kv_max, n_head_kv); // env/default — same value every row
29251        // hd512 split override (MEMRA_FA_SP512, 2026-07-11): gemma globals have n_head_kv=2 so
29252        // the grid is (2 x n_splits) — at depth ~29 splits = 58 blocks on 82 SMs (half idle,
29253        // rows_dpl16 8x off its byte floor). EVERY gemma hd512 caller shares THIS wrapper
29254        // (parity law), so the partition is freely tunable — verify and decode move together.
29255        if head_dim == 512 {
29256            static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29257            // default 16 (2026-07-11 depth sweep, N=2: plain 155.4->156.5, depth spec
29258            // 236.9->250.4; 12/24/32 all worse). hd512 exists only on gemma globals.
29259            let v = *SP512.get_or_init(|| {
29260                std::env::var("MEMRA_FA_SP512")
29261                    .ok()
29262                    .and_then(|x| x.parse().ok())
29263                    .unwrap_or(0)
29264            });
29265            sp = if v >= 8 {
29266                v
29267            } else {
29268                FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
29269            };
29270        }
29271        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29272        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29273        let gqa = (n_head / n_head_kv).max(1) as u32;
29274        // LADDER-RUNG STRADDLE FIX (issue #10, 2026-07-13, rtx6000-proven): one sp for every row
29275        // diverges from eager decode when a split-ladder rung falls INSIDE the batch — row r's
29276        // eager twin used fa_split_keys(t_kv_r), the batch used fa_split_keys(t_kv_max), and
29277        // the different partition changes the combine's FP order (greedy tie flips at depth;
29278        // MEMRA_FA_SPLIT=64 pin -> PASS on the exact rtx6000 failing config). Fix: group
29279        // consecutive rows by their OWN ladder value and launch once per group — each row then
29280        // executes the exact per-row program eager ran. Rungs land once per doubling, so this
29281        // is 1 launch in the common case and 2 on a crossing round. hd512 keeps one group (its
29282        // sp override is t_kv-independent by construction).
29283        let mut groups: Vec<(usize, usize, usize)> = Vec::new(); // (row0, t_g, sp_g)
29284        if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
29285            groups.push((0, t, sp));
29286        } else {
29287            let mut r0 = 0usize;
29288            while r0 < t {
29289                let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
29290                let mut r1 = r0 + 1;
29291                while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
29292                    r1 += 1;
29293                }
29294                groups.push((r0, r1 - r0, sp_g));
29295                r0 = r1;
29296            }
29297        }
29298        // Deep-ctx smem twin for the VERIFY rows (2026-07-05): same threshold + rationale as
29299        // fa_decode's dispatch — at 40k the register path's GQA L2-reuse premise is dead and the
29300        // verify multiplies the 4x DRAM re-read by T rows. Bit-identical per (row,token,split).
29301        static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29302        let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
29303            std::env::var("MEMRA_FA_SMEM_TKV")
29304                .ok()
29305                .and_then(|v| v.parse().ok())
29306                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
29307        });
29308        let v4 = fa_v4_at(base_len + t) && head_dim == 256;
29309        let v3 = fa_v3_active(head_dim);
29310        let smem_rows =
29311            head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
29312        // kv_shared twin RETIRED (2026-07-11 depth run-gen gate): the wv:=wk premise fails
29313        // POST-cache — cached K is k-normed+roped, cached V is not; the twin fed roped keys
29314        // in as values. Verify/decode/stream gates were blind (both sides shared the wrong
29315        // symbol — the parity law's blind spot); only prefill-vs-decode at depth caught it.
29316        let _ = kv_shared;
29317        // i2 twin: 2-key interleaved walk (MEMRA_FA_I2=0 reverts). i4 probed NEGATIVE
29318        // (157.3 vs 161.2 depth plain — register pressure past i2's sweet spot; jsonl).
29319        let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
29320        // v4-hd512 (MEMRA_FA_V512=1 opt-in, 2026-07-14): the v4 key-per-lane recipe on the
29321        // globals lane (depth profile: i2 ~4.6x off its byte floor — the v3-class
29322        // reduce-per-key latency signature). NEW NUMERIC CONFIG shared by every hd512
29323        // caller (decode+verify flip together); run-gen argmax + acceptance arbitrate.
29324        // T-BATCHED hd512 (DEFAULT ON 2026-07-14, MEMRA_FA_TB512=0 seam): one block per
29325        // (kv_head, split) stages its tile once and loops the rows over it — kills the
29326        // x t DRAM re-read of the full-ctx globals (depth cell +1.4%, plain flat, N=3
29327        // interleaved). FIXED absolute partition = NEW NUMERIC for the combine order,
29328        // shared by every hd512 caller through this wrapper (decode+verify flip together;
29329        // depth stream identical, acceptance unshifted, spec 256/256 x3 models).
29330        // Requires sp <= 32 (single staged tile; acc reused per row). The z-form v4_512
29331        // sibling (in-kernel dp4a port alone) probed FLAT — hd512 was DRAM-re-read-bound,
29332        // not unpack-bound; jsonl 2026-07-14.
29333        static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
29334        // gqa <= 16 = fa_v4_smem_512's q-array capacity; past it fall to the register twins.
29335        let tb512 = head_dim == 512
29336            && sp <= 32
29337            && n_head / n_head_kv.max(1) <= 16
29338            && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
29339        let fname = if tb512 {
29340            "fa_decode_vec_q_rows_v4_512_tb"
29341        } else if i2 {
29342            "fa_decode_vec_q_rows_dpl16_i2"
29343        } else if head_dim == 512 {
29344            "fa_decode_vec_q_rows_dpl16"
29345        }
29346        // gemma globals (parity law)
29347        else if v4 {
29348            "fa_decode_vec_q_rows_v4"
29349        } else if v3 {
29350            "fa_decode_vec_q_rows_v3"
29351        } else if fa_v2_on() {
29352            "fa_decode_vec_q_rows_v2"
29353        } else if smem_rows {
29354            "fa_decode_vec_q_rows_smem"
29355        } else {
29356            "fa_decode_vec_q_rows"
29357        };
29358        let f = if head_dim == 512 {
29359            self.fa_func(fname, head_dim)
29360        } else if g {
29361            // FP8-WINDOWED: hd256 rows over an e4m3 cache — kf8vf8 module, SAME symbol
29362            // choice as decode's kvmod dispatch (parity law: excluding v4 here paired
29363            // g-module rows against decode's g-module v4 — different programs, short-VG
29364            // maxdiff 2.0 / spec stream 0/128, 2026-07-12). rows_v4 is format-aware
29365            // since fda9790; only the smem twin stays excluded (V-stage q5_1-only).
29366            // hd128 (qwen fp8-KV) lands on the base/register rows via fname — the
29367            // dq macros are format-aware.
29368            self.func_g(if smem_rows {
29369                "fa_decode_vec_q_rows"
29370            } else {
29371                fname
29372            })
29373        } else {
29374            self.func(fname)
29375        };
29376        let shmem = if tb512 {
29377            // fa_v4_smem_512 (q 9KB gqa<=16 + k tile 18KB) + sV 32*512 (e4m3 module halves it)
29378            let gk = Self::gkv_on();
29379            let sh =
29380                (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
29381            use cudarc::driver::sys::CUfunction_attribute_enum as A;
29382            f.set_attribute(
29383                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29384                sh as i32,
29385            )?;
29386            sh
29387        } else if v4 || v3 || smem_rows || fa_v2_on() {
29388            // v4: fa_v4_smem (11.5KB) + sV; v3 stages sV only; v2/smem twins stage sK+sV.
29389            let sh = (if v4 {
29390                11520 + 32 * head_dim * if g { 1 } else { 2 }
29391            } else if v3 {
29392                32 * head_dim * 2
29393            } else {
29394                2 * 32 * head_dim * 2
29395            }) as u32;
29396            use cudarc::driver::sys::CUfunction_attribute_enum as A;
29397            f.set_attribute(
29398                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29399                sh as i32,
29400            )?;
29401            sh
29402        } else {
29403            0
29404        };
29405        // Per-GROUP launches (single group in the common case — identical to the pre-fix
29406        // single launch there): each group gets its own partials (the rows kernel indexes
29407        // partials by its LOCAL grid.z row) and q/o row-offset views.
29408        for &(r0, t_g, sp_g) in &groups {
29409            let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
29410            let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
29411            let base_i = (base_len + r0) as i32;
29412            let o_len = t_g * n_head * n_splits_g * head_dim;
29413            let ml_len = t_g * n_head * n_splits_g;
29414            let mut part_guard = self.fa_part_pool.lock().unwrap();
29415            if part_guard
29416                .as_ref()
29417                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29418                .unwrap_or(true)
29419            {
29420                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
29421                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
29422                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
29423                // later live allocations land at those addresses, and the next graph REPLAY writes
29424                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
29425                // output corruption began the burst after the trunk's t_kv growth first realloc'd
29426                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
29427                // the baked addresses alive (single-stream: eager writes the new buffers, replays
29428                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
29429                // (total retired < final size).
29430                let old = part_guard.take();
29431                let (co, cm) = old
29432                    .as_ref()
29433                    .map(|pp| (pp.0.len(), pp.1.len()))
29434                    .unwrap_or((0, 0));
29435                if let Some(old) = old {
29436                    self.fa_part_retired.lock().unwrap().push(old);
29437                }
29438                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29439                    eprintln!(
29440                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29441                        co, o_len, cm, ml_len
29442                    );
29443                }
29444                *part_guard =
29445                    Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29446            }
29447            let pg = part_guard.as_mut().unwrap();
29448            self.gpu
29449                .stream()
29450                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29451            self.gpu
29452                .stream()
29453                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29454            self.gpu
29455                .stream()
29456                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29457            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29458            let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
29459            let qv = self.view(q, t * n_head * head_dim);
29460            let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
29461            let cfg = LaunchConfig {
29462                grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
29463                block_dim: (32, gqa, 1),
29464                shared_mem_bytes: shmem,
29465            };
29466            {
29467                let __s_b = self.gpu.stream();
29468                let mut b = __s_b.launch_builder(&f);
29469                if tb512 {
29470                    // rows-inner launch: grid.z dropped, the kernel loops n_rows itself.
29471                    let (bd, plus) =
29472                        base_dev.expect("hd512 rows twin requires a device base counter");
29473                    let plus_g = plus + r0 as i32;
29474                    let nr = t_g as i32;
29475                    if Self::pdl_on() && Self::pdl_wb_on() {
29476                        // wave-B2b: flavor mirrors fa_func(fname, 512) = gkv.
29477                        use cudarc::driver::{DevicePtr, DevicePtrMut};
29478                        let s = &self.gpu.stream();
29479                        let (pq, _b0) = q_g.device_ptr(s);
29480                        let (pk, _b1) = k.device_ptr(s);
29481                        let (pv, _b2) = v.device_ptr(s);
29482                        let (po, _b3) = part_o.device_ptr_mut(s);
29483                        let (pm, _b4) = part_m.device_ptr_mut(s);
29484                        let (pl, _b5) = part_l.device_ptr_mut(s);
29485                        let (pb, _b6) = bd.device_ptr(s);
29486                        let mut ps = [
29487                            &pq as *const _ as *mut std::ffi::c_void,
29488                            &pk as *const _ as *mut _,
29489                            &pv as *const _ as *mut _,
29490                            &po as *const _ as *mut _,
29491                            &pm as *const _ as *mut _,
29492                            &pl as *const _ as *mut _,
29493                            &hd as *const _ as *mut _,
29494                            &nh as *const _ as *mut _,
29495                            &nhkv as *const _ as *mut _,
29496                            &pb as *const _ as *mut _,
29497                            &plus_g as *const _ as *mut _,
29498                            &scale as *const _ as *mut _,
29499                            &nspm as *const _ as *mut _,
29500                            &spk as *const _ as *mut _,
29501                            &ktb as *const _ as *mut _,
29502                            &vtb as *const _ as *mut _,
29503                            &nr as *const _ as *mut _,
29504                        ];
29505                        unsafe {
29506                            self.launch_pdl_flash(
29507                                Self::gkv_on(),
29508                                "fa_decode_vec_q_rows_v4_512_tb",
29509                                (n_head_kv as u32, n_splits_g as u32, 1),
29510                                (32, gqa, 1),
29511                                shmem,
29512                                &mut ps,
29513                            )?;
29514                        }
29515                    } else {
29516                        let cfg_tb = LaunchConfig {
29517                            grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
29518                            block_dim: (32, gqa, 1),
29519                            shared_mem_bytes: shmem,
29520                        };
29521                        b.arg(&q_g)
29522                            .arg(k)
29523                            .arg(v)
29524                            .arg(&mut *part_o)
29525                            .arg(&mut *part_m)
29526                            .arg(&mut *part_l)
29527                            .arg(&hd)
29528                            .arg(&nh)
29529                            .arg(&nhkv)
29530                            .arg(bd)
29531                            .arg(&plus_g)
29532                            .arg(&scale)
29533                            .arg(&nspm)
29534                            .arg(&spk)
29535                            .arg(&ktb)
29536                            .arg(&vtb)
29537                            .arg(&nr);
29538                        unsafe {
29539                            b.launch(cfg_tb)?;
29540                        }
29541                    }
29542                } else if head_dim == 512 {
29543                    let (bd, plus) =
29544                        base_dev.expect("hd512 rows twin requires a device base counter");
29545                    let plus_g = plus + r0 as i32;
29546                    b.arg(&q_g)
29547                        .arg(k)
29548                        .arg(v)
29549                        .arg(&mut *part_o)
29550                        .arg(&mut *part_m)
29551                        .arg(&mut *part_l)
29552                        .arg(&hd)
29553                        .arg(&nh)
29554                        .arg(&nhkv)
29555                        .arg(bd)
29556                        .arg(&plus_g)
29557                        .arg(&scale)
29558                        .arg(&nspm)
29559                        .arg(&spk)
29560                        .arg(&ktb)
29561                        .arg(&vtb);
29562                    unsafe {
29563                        b.launch(cfg)?;
29564                    }
29565                } else {
29566                    b.arg(&q_g)
29567                        .arg(k)
29568                        .arg(v)
29569                        .arg(&mut *part_o)
29570                        .arg(&mut *part_m)
29571                        .arg(&mut *part_l)
29572                        .arg(&hd)
29573                        .arg(&nh)
29574                        .arg(&nhkv)
29575                        .arg(&base_i)
29576                        .arg(&scale)
29577                        .arg(&nspm)
29578                        .arg(&spk)
29579                        .arg(&ktb)
29580                        .arg(&vtb);
29581                    unsafe {
29582                        b.launch(cfg)?;
29583                    }
29584                }
29585            }
29586            let cfg2 = LaunchConfig {
29587                grid_dim: (n_head as u32, t_g as u32, 1),
29588                block_dim: (head_dim as u32, 1, 1),
29589                shared_mem_bytes: 0,
29590            };
29591            let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
29592            if head_dim == 512 {
29593                // device-len combine (shared by verify/eager/graph — parity by symbol): the
29594                // per-row n_splits derives from the SAME counter the rows kernel read.
29595                let (bd, plus) = base_dev.unwrap();
29596                let plus_g = plus + r0 as i32;
29597                if let Some((oq, od)) = q8_out.as_mut() {
29598                    // wave-5b port (2026-07-23, t=1 decode only): q8-emitting dc combine.
29599                    debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
29600                    if Self::pdl_on() && Self::pdl_wb_on() {
29601                        // wave-B2: flavor mirrors fa_func (hd512 + gkv → kf8vf8).
29602                        use cudarc::driver::{DevicePtr, DevicePtrMut};
29603                        let s = &self.gpu.stream();
29604                        let (po, _g0) = part_o.device_ptr(s);
29605                        let (pm, _g1) = part_m.device_ptr(s);
29606                        let (pl, _g2) = part_l.device_ptr(s);
29607                        let (pq, _g3) = oq.device_ptr_mut(s);
29608                        let (pd, _g4) = od.device_ptr_mut(s);
29609                        let (pb, _g5) = bd.device_ptr(s);
29610                        let mut ps = [
29611                            &po as *const _ as *mut std::ffi::c_void,
29612                            &pm as *const _ as *mut _,
29613                            &pl as *const _ as *mut _,
29614                            &pq as *const _ as *mut _,
29615                            &pd as *const _ as *mut _,
29616                            &hd as *const _ as *mut _,
29617                            &nh as *const _ as *mut _,
29618                            &pb as *const _ as *mut _,
29619                            &plus_g as *const _ as *mut _,
29620                            &nspm as *const _ as *mut _,
29621                            &spk as *const _ as *mut _,
29622                        ];
29623                        unsafe {
29624                            self.launch_pdl_flash(
29625                                Self::gkv_on(),
29626                                "fa_decode_combine_rows_dc_q8_1",
29627                                cfg2.grid_dim,
29628                                cfg2.block_dim,
29629                                0,
29630                                &mut ps,
29631                            )?;
29632                        }
29633                        continue;
29634                    }
29635                    let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
29636                    let __s_b2 = self.gpu.stream();
29637                    let mut b2 = __s_b2.launch_builder(&fc);
29638                    b2.arg(&*part_o)
29639                        .arg(&*part_m)
29640                        .arg(&*part_l)
29641                        .arg(&mut **oq)
29642                        .arg(&mut **od)
29643                        .arg(&hd)
29644                        .arg(&nh)
29645                        .arg(bd)
29646                        .arg(&plus_g)
29647                        .arg(&nspm)
29648                        .arg(&spk);
29649                    unsafe {
29650                        b2.launch(cfg2)?;
29651                    }
29652                    continue;
29653                }
29654                let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
29655                let __s_b2 = self.gpu.stream();
29656                let mut b2 = __s_b2.launch_builder(&fc);
29657                b2.arg(&*part_o)
29658                    .arg(&*part_m)
29659                    .arg(&*part_l)
29660                    .arg(&mut o_g)
29661                    .arg(&hd)
29662                    .arg(&nh)
29663                    .arg(bd)
29664                    .arg(&plus_g)
29665                    .arg(&nspm)
29666                    .arg(&spk);
29667                unsafe {
29668                    b2.launch(cfg2)?;
29669                }
29670            } else {
29671                // q8 emit is wired for the hd512 dc-combine arm only — a Some here would
29672                // leave the caller's pair unwritten (consumer would read garbage).
29673                assert!(
29674                    q8_out.is_none(),
29675                    "rows q8 emit requires the hd512 dc combine"
29676                );
29677                let fc = self.func("fa_decode_combine_rows");
29678                let __s_b2 = self.gpu.stream();
29679                let mut b2 = __s_b2.launch_builder(&fc);
29680                b2.arg(&*part_o)
29681                    .arg(&*part_m)
29682                    .arg(&*part_l)
29683                    .arg(&mut o_g)
29684                    .arg(&hd)
29685                    .arg(&nh)
29686                    .arg(&base_i)
29687                    .arg(&nspm)
29688                    .arg(&spk);
29689                unsafe {
29690                    b2.launch(cfg2)?;
29691                }
29692            }
29693        }
29694        Ok(())
29695    }
29696
29697    /// WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly `window` keys —
29698    /// bit-identical per row to the T=1 decode's fa_decode over the window VIEW. Caller gates
29699    /// base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).
29700    #[allow(clippy::too_many_arguments)]
29701    pub fn fa_decode_rows_w(
29702        &self,
29703        q: &CudaSlice<f32>,
29704        k: &cudarc::driver::CudaView<u8>,
29705        v: &cudarc::driver::CudaView<u8>,
29706        o: &mut CudaSlice<f32>,
29707        head_dim: usize,
29708        n_head: usize,
29709        n_head_kv: usize,
29710        base_dev: &CudaSlice<i32>,
29711        base_plus: i32,
29712        t: usize,
29713        scale: f32,
29714        window: usize,
29715        k_tok_bytes: usize,
29716        v_tok_bytes: usize,
29717        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
29718    ) -> Result<(), Box<dyn std::error::Error>> {
29719        // DEVICE-LEN (graph arc step 1, 2026-07-11): the causal base rides an i32 counter
29720        // (kernel T_kv = dev[0] + base_plus + r + 1) so depth graphs can replay with len
29721        // advancing on-device. dc paths pass kvl.len_d with plus=-1; verify/eager sync the
29722        // counter with one async set_i32_one first. Partials/splits size from `window` (host).
29723        debug_assert!(head_dim == 256);
29724        // windowed split (MEMRA_FA_SPW, default 32 — re-swept 2026-07-12 under the raw-e4m3 sV
29725        // occupancy ceiling (4 blocks/SM): t=1 decode is GRID-limited (win/sp splits x nkv
29726        // blocks), so smaller splits fill the ceiling — 1.7k 174.4/174.0 vs 48's 170.7/170.3,
29727        // 4.9k 159.8 vs 157.4 (N=2 interleaved, stable window). Spec serving prefers 64
29728        // (verify t=K+1 fills the grid via grid.z=t; depth K=7 281.3 vs 249.3 at 32) — set
29729        // MEMRA_FA_SPW=64 there, same config law as MEMRA_GEMMA_GKV=0. MUST be one value for
29730        // ALL widths: a t-keyed probe broke decode-vs-verify combine order (stream 9/128).
29731        let sp = {
29732            static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29733            let v = *SPW.get_or_init(|| {
29734                std::env::var("MEMRA_FA_SPW")
29735                    .ok()
29736                    .and_then(|x| x.parse().ok())
29737                    .unwrap_or(0)
29738            });
29739            if v >= 8 {
29740                v
29741            } else {
29742                FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
29743            }
29744        };
29745        #[allow(clippy::manual_div_ceil)]
29746        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
29747        let n_splits_max = (window + sp - 1) / sp;
29748        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29749        let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
29750        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29751        let gqa = (n_head / n_head_kv).max(1) as u32;
29752        let o_len = t * n_head * n_splits_max * head_dim;
29753        let ml_len = t * n_head * n_splits_max;
29754        let mut part_guard = self.fa_part_pool.lock().unwrap();
29755        if part_guard
29756            .as_ref()
29757            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29758            .unwrap_or(true)
29759        {
29760            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
29761            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
29762            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
29763            // later live allocations land at those addresses, and the next graph REPLAY writes
29764            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
29765            // output corruption began the burst after the trunk's t_kv growth first realloc'd
29766            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
29767            // the baked addresses alive (single-stream: eager writes the new buffers, replays
29768            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
29769            // (total retired < final size).
29770            let old = part_guard.take();
29771            let (co, cm) = old
29772                .as_ref()
29773                .map(|pp| (pp.0.len(), pp.1.len()))
29774                .unwrap_or((0, 0));
29775            if let Some(old) = old {
29776                self.fa_part_retired.lock().unwrap().push(old);
29777            }
29778            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29779                eprintln!(
29780                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29781                    co, o_len, cm, ml_len
29782                );
29783            }
29784            *part_guard =
29785                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29786        }
29787        let pg = part_guard.as_mut().unwrap();
29788        self.gpu
29789            .stream()
29790            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29791        self.gpu
29792            .stream()
29793            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29794        self.gpu
29795            .stream()
29796            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29797        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29798        // Lane pick: decode AND verify both land here in the windowed regime (parity law —
29799        // hybrid_forward verify_attn), so the pick only needs internal consistency, not
29800        // clone-of-decode bit fidelity (SASS-proven impossible for textually identical
29801        // kernels, jsonl 2026-07-10). v4 under the threshold; smem twin at/above the smem
29802        // floor (deep-ctx broadcast win); register twin between.
29803        static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29804        let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
29805            std::env::var("MEMRA_FA_SMEM_TKV")
29806                .ok()
29807                .and_then(|v| v.parse().ok())
29808                .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
29809        });
29810        // MULTI-ROW v4: resurrected 2026-07-14 (the '33 tok/s collapse' was a paired-map
29811        // partial-write bug, not the mechanism) and falsified HONESTLY at gqa 2: bit-exact
29812        // but −1.7% on the 31B depth cell — the sp helper warp already hides staging
29813        // in-block, and mr trades L2-cheap redundant bytes for serialized per-warp gqa
29814        // score/B3 chains. Arm deleted; jsonl row 2026-07-14 is the record.
29815        use cudarc::driver::sys::CUfunction_attribute_enum as A;
29816        // FP8-WINDOWED (wkv): the v4 family is format-aware (2026-07-12 KFMT/VFMT staging
29817        // arms) — wkv rides the SAME lane logic, resolved from the kf8vf8 module. One symbol
29818        // per (lane, format-module) keeps parity structural; the old register-i2 detour
29819        // (-33%) is retired.
29820        let wg = Self::wkv_on();
29821        // STAGING-PARALLEL v4 (MEMRA_FA_SPW2, default ON at gqa==1): warp 1 = staging helper
29822        // (v4 is 61% staging); score phases identical to v4_w. Same symbol all t.
29823        let sp2 =
29824            gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
29825        if sp2 {
29826            let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
29827            if Self::pdl_on() && Self::pdl_wb_on() {
29828                // wave-B2b: flavor mirrors wg.
29829                use cudarc::driver::{DevicePtr, DevicePtrMut};
29830                let s = &self.gpu.stream();
29831                let (pq, _b0) = q.device_ptr(s);
29832                let (pk, _b1) = k.device_ptr(s);
29833                let (pv, _b2) = v.device_ptr(s);
29834                let (po, _b3) = part_o.device_ptr_mut(s);
29835                let (pm, _b4) = part_m.device_ptr_mut(s);
29836                let (pl, _b5) = part_l.device_ptr_mut(s);
29837                let (pb, _b6) = base_dev.device_ptr(s);
29838                let mut ps = [
29839                    &pq as *const _ as *mut std::ffi::c_void,
29840                    &pk as *const _ as *mut _,
29841                    &pv as *const _ as *mut _,
29842                    &po as *const _ as *mut _,
29843                    &pm as *const _ as *mut _,
29844                    &pl as *const _ as *mut _,
29845                    &hd as *const _ as *mut _,
29846                    &nh as *const _ as *mut _,
29847                    &nhkv as *const _ as *mut _,
29848                    &pb as *const _ as *mut _,
29849                    &base_plus as *const _ as *mut _,
29850                    &scale as *const _ as *mut _,
29851                    &nspm as *const _ as *mut _,
29852                    &spk as *const _ as *mut _,
29853                    &ktb as *const _ as *mut _,
29854                    &vtb as *const _ as *mut _,
29855                    &wini as *const _ as *mut _,
29856                ];
29857                unsafe {
29858                    self.launch_pdl_flash(
29859                        wg,
29860                        "fa_decode_vec_q_rows_v4_w_sp",
29861                        (n_head_kv as u32, n_splits_max as u32, t as u32),
29862                        (32, gqa + 1, 1),
29863                        sh,
29864                        &mut ps,
29865                    )?;
29866                }
29867            } else {
29868                let f = if wg {
29869                    self.func_g("fa_decode_vec_q_rows_v4_w_sp")
29870                } else {
29871                    self.func("fa_decode_vec_q_rows_v4_w_sp")
29872                };
29873                f.set_attribute(
29874                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29875                    sh as i32,
29876                )?;
29877                let cfg = LaunchConfig {
29878                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
29879                    block_dim: (32, gqa + 1, 1),
29880                    shared_mem_bytes: sh,
29881                };
29882                let __s_b = self.gpu.stream();
29883                let mut b = __s_b.launch_builder(&f);
29884                b.arg(q)
29885                    .arg(k)
29886                    .arg(v)
29887                    .arg(&mut *part_o)
29888                    .arg(&mut *part_m)
29889                    .arg(&mut *part_l)
29890                    .arg(&hd)
29891                    .arg(&nh)
29892                    .arg(&nhkv)
29893                    .arg(base_dev)
29894                    .arg(&base_plus)
29895                    .arg(&scale)
29896                    .arg(&nspm)
29897                    .arg(&spk)
29898                    .arg(&ktb)
29899                    .arg(&vtb)
29900                    .arg(&wini);
29901                unsafe {
29902                    b.launch(cfg)?;
29903                }
29904            }
29905        } else {
29906            if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
29907                // wave-B2b: the v4_w pick only (smem/reg twins stay builder-launched).
29908                let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
29909                use cudarc::driver::{DevicePtr, DevicePtrMut};
29910                let s = &self.gpu.stream();
29911                let (pq, _b0) = q.device_ptr(s);
29912                let (pk, _b1) = k.device_ptr(s);
29913                let (pv, _b2) = v.device_ptr(s);
29914                let (po, _b3) = part_o.device_ptr_mut(s);
29915                let (pm, _b4) = part_m.device_ptr_mut(s);
29916                let (pl, _b5) = part_l.device_ptr_mut(s);
29917                let (pb, _b6) = base_dev.device_ptr(s);
29918                let mut ps = [
29919                    &pq as *const _ as *mut std::ffi::c_void,
29920                    &pk as *const _ as *mut _,
29921                    &pv as *const _ as *mut _,
29922                    &po as *const _ as *mut _,
29923                    &pm as *const _ as *mut _,
29924                    &pl as *const _ as *mut _,
29925                    &hd as *const _ as *mut _,
29926                    &nh as *const _ as *mut _,
29927                    &nhkv as *const _ as *mut _,
29928                    &pb as *const _ as *mut _,
29929                    &base_plus as *const _ as *mut _,
29930                    &scale as *const _ as *mut _,
29931                    &nspm as *const _ as *mut _,
29932                    &spk as *const _ as *mut _,
29933                    &ktb as *const _ as *mut _,
29934                    &vtb as *const _ as *mut _,
29935                    &wini as *const _ as *mut _,
29936                ];
29937                unsafe {
29938                    self.launch_pdl_flash(
29939                        wg,
29940                        "fa_decode_vec_q_rows_v4_w",
29941                        (n_head_kv as u32, n_splits_max as u32, t as u32),
29942                        (32, gqa, 1),
29943                        sh,
29944                        &mut ps,
29945                    )?;
29946                }
29947            } else {
29948                let pick = |name: &str| {
29949                    if wg {
29950                        self.func_g(name)
29951                    } else {
29952                        self.func(name)
29953                    }
29954                };
29955                let (f, sh) = if fa_v4_at(window) {
29956                    let f = pick("fa_decode_vec_q_rows_v4_w");
29957                    (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
29958                } else if smem_tkv > 0 && window >= smem_tkv {
29959                    // NOTE: the smem twin's V-stage is still q5_1-hardcoded — unreachable under wkv
29960                    // at the gemma window (v4 covers it); revisit if the smem floor ever drops.
29961                    (
29962                        pick("fa_decode_vec_q_rows_smem_w"),
29963                        (2 * 32 * head_dim * 2) as u32,
29964                    )
29965                } else {
29966                    (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
29967                };
29968                f.set_attribute(
29969                    A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29970                    sh as i32,
29971                )?;
29972                let cfg = LaunchConfig {
29973                    grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
29974                    block_dim: (32, gqa, 1),
29975                    shared_mem_bytes: sh,
29976                };
29977                let __s_b = self.gpu.stream();
29978                let mut b = __s_b.launch_builder(&f);
29979                b.arg(q)
29980                    .arg(k)
29981                    .arg(v)
29982                    .arg(&mut *part_o)
29983                    .arg(&mut *part_m)
29984                    .arg(&mut *part_l)
29985                    .arg(&hd)
29986                    .arg(&nh)
29987                    .arg(&nhkv)
29988                    .arg(base_dev)
29989                    .arg(&base_plus)
29990                    .arg(&scale)
29991                    .arg(&nspm)
29992                    .arg(&spk)
29993                    .arg(&ktb)
29994                    .arg(&vtb)
29995                    .arg(&wini);
29996                unsafe {
29997                    b.launch(cfg)?;
29998                }
29999            }
30000        }
30001        let cfg2 = LaunchConfig {
30002            grid_dim: (n_head as u32, t as u32, 1),
30003            block_dim: (head_dim as u32, 1, 1),
30004            shared_mem_bytes: 0,
30005        };
30006        if let Some((oq, od)) = q8_out {
30007            // wave-5b port (2026-07-23): q8-emitting combine — the t=1 decode's wo matvec
30008            // consumes the pair directly; the standalone quantize launch folds away.
30009            if Self::pdl_on() && Self::pdl_wb_on() {
30010                // wave-B2: flavor mirrors the builder's wg choice.
30011                use cudarc::driver::{DevicePtr, DevicePtrMut};
30012                let s = &self.gpu.stream();
30013                let (po, _g0) = part_o.device_ptr(s);
30014                let (pm, _g1) = part_m.device_ptr(s);
30015                let (pl, _g2) = part_l.device_ptr(s);
30016                let (pq, _g3) = oq.device_ptr_mut(s);
30017                let (pd, _g4) = od.device_ptr_mut(s);
30018                let mut ps = [
30019                    &po as *const _ as *mut std::ffi::c_void,
30020                    &pm as *const _ as *mut _,
30021                    &pl as *const _ as *mut _,
30022                    &pq as *const _ as *mut _,
30023                    &pd as *const _ as *mut _,
30024                    &hd as *const _ as *mut _,
30025                    &nh as *const _ as *mut _,
30026                    &nspm as *const _ as *mut _,
30027                    &spk as *const _ as *mut _,
30028                    &wini as *const _ as *mut _,
30029                ];
30030                unsafe {
30031                    self.launch_pdl_flash(
30032                        wg,
30033                        "fa_decode_combine_rows_w_q8_1",
30034                        cfg2.grid_dim,
30035                        cfg2.block_dim,
30036                        0,
30037                        &mut ps,
30038                    )?;
30039                }
30040                return Ok(());
30041            }
30042            let fc = if wg {
30043                self.func_g("fa_decode_combine_rows_w_q8_1")
30044            } else {
30045                self.func("fa_decode_combine_rows_w_q8_1")
30046            };
30047            let __s_b2 = self.gpu.stream();
30048            let mut b2 = __s_b2.launch_builder(&fc);
30049            b2.arg(&*part_o)
30050                .arg(&*part_m)
30051                .arg(&*part_l)
30052                .arg(oq)
30053                .arg(od)
30054                .arg(&hd)
30055                .arg(&nh)
30056                .arg(&nspm)
30057                .arg(&spk)
30058                .arg(&wini);
30059            unsafe {
30060                b2.launch(cfg2)?;
30061            }
30062            return Ok(());
30063        }
30064        let fc = if wg {
30065            self.func_g("fa_decode_combine_rows_w")
30066        } else {
30067            self.func("fa_decode_combine_rows_w")
30068        };
30069        let __s_b2 = self.gpu.stream();
30070        let mut b2 = __s_b2.launch_builder(&fc);
30071        b2.arg(&*part_o)
30072            .arg(&*part_m)
30073            .arg(&*part_l)
30074            .arg(o)
30075            .arg(&hd)
30076            .arg(&nh)
30077            .arg(&nspm)
30078            .arg(&spk)
30079            .arg(&wini);
30080        unsafe {
30081            b2.launch(cfg2)?;
30082        }
30083        Ok(())
30084    }
30085
30086    /// ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes:
30087    /// v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware);
30088    /// `t_kv_upper` sizes splits/partials — the same one-sp-for-all-rows approximation class
30089    /// the host rows path already uses (battery-arbitrated); actual per-row bounds derive
30090    /// in-kernel from the counter (+ base_plus, v4 lane only — v3's kernel has no plus arg).
30091    #[allow(clippy::too_many_arguments)]
30092    pub fn fa_decode_rows_dc(
30093        &self,
30094        q: &CudaSlice<f32>,
30095        k: &cudarc::driver::CudaView<u8>,
30096        v: &cudarc::driver::CudaView<u8>,
30097        o: &mut CudaSlice<f32>,
30098        head_dim: usize,
30099        n_head: usize,
30100        n_head_kv: usize,
30101        base_dev: &CudaSlice<i32>,
30102        t_kv_upper: usize,
30103        t: usize,
30104        scale: f32,
30105        k_tok_bytes: usize,
30106        v_tok_bytes: usize,
30107        base_plus: i32,
30108        g: bool,
30109    ) -> Result<(), Box<dyn std::error::Error>> {
30110        let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
30111        assert!(
30112            v4 || fa_v3_active(head_dim),
30113            "stream fa rows requires the v3 or v4 lane"
30114        );
30115        assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
30116        if v4 {
30117            let sp = fa_split_keys(t_kv_upper, n_head_kv);
30118            #[allow(clippy::manual_div_ceil)]
30119            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30120            let n_splits_max = (t_kv_upper + sp - 1) / sp;
30121            let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
30122            let (nspm, spk) = (n_splits_max as i32, sp as i32);
30123            let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30124            let gqa = (n_head / n_head_kv).max(1) as u32;
30125            let o_len = t * n_head * n_splits_max * head_dim;
30126            let ml_len = t * n_head * n_splits_max;
30127            let mut part_guard = self.fa_part_pool.lock().unwrap();
30128            if part_guard
30129                .as_ref()
30130                .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30131                .unwrap_or(true)
30132            {
30133                // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
30134                // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
30135                // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
30136                // later live allocations land at those addresses, and the next graph REPLAY writes
30137                // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
30138                // output corruption began the burst after the trunk's t_kv growth first realloc'd
30139                // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
30140                // the baked addresses alive (single-stream: eager writes the new buffers, replays
30141                // touch only the old — never concurrently). Doubling growth bounds retired VRAM
30142                // (total retired < final size).
30143                let old = part_guard.take();
30144                let (co, cm) = old
30145                    .as_ref()
30146                    .map(|pp| (pp.0.len(), pp.1.len()))
30147                    .unwrap_or((0, 0));
30148                if let Some(old) = old {
30149                    self.fa_part_retired.lock().unwrap().push(old);
30150                }
30151                if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
30152                    eprintln!(
30153                        "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
30154                        co, o_len, cm, ml_len
30155                    );
30156                }
30157                *part_guard =
30158                    Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30159            }
30160            let pg = part_guard.as_mut().unwrap();
30161            self.gpu
30162                .stream()
30163                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
30164            self.gpu
30165                .stream()
30166                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
30167            self.gpu
30168                .stream()
30169                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
30170            let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30171            let f = if g {
30172                self.func_g("fa_decode_vec_q_rows_v4_dc")
30173            } else {
30174                self.func("fa_decode_vec_q_rows_v4_dc")
30175            };
30176            let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
30177            use cudarc::driver::sys::CUfunction_attribute_enum as A;
30178            f.set_attribute(
30179                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30180                sh as i32,
30181            )?;
30182            let cfg = LaunchConfig {
30183                grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
30184                block_dim: (32, gqa, 1),
30185                shared_mem_bytes: sh,
30186            };
30187            let __s_b = self.gpu.stream();
30188            let mut b = __s_b.launch_builder(&f);
30189            b.arg(q)
30190                .arg(k)
30191                .arg(v)
30192                .arg(&mut *part_o)
30193                .arg(&mut *part_m)
30194                .arg(&mut *part_l)
30195                .arg(&hd)
30196                .arg(&nh)
30197                .arg(&nhkv)
30198                .arg(base_dev)
30199                .arg(&base_plus)
30200                .arg(&scale)
30201                .arg(&nspm)
30202                .arg(&spk)
30203                .arg(&ktb)
30204                .arg(&vtb);
30205            unsafe {
30206                b.launch(cfg)?;
30207            }
30208            let fc = self.func("fa_decode_combine_rows_dc");
30209            let cfg2 = LaunchConfig {
30210                grid_dim: (n_head as u32, t as u32, 1),
30211                block_dim: (head_dim as u32, 1, 1),
30212                shared_mem_bytes: 0,
30213            };
30214            let __s_b2 = self.gpu.stream();
30215            let mut b2 = __s_b2.launch_builder(&fc);
30216            b2.arg(&*part_o)
30217                .arg(&*part_m)
30218                .arg(&*part_l)
30219                .arg(o)
30220                .arg(&hd)
30221                .arg(&nh)
30222                .arg(base_dev)
30223                .arg(&base_plus)
30224                .arg(&nspm)
30225                .arg(&spk);
30226            unsafe {
30227                b2.launch(cfg2)?;
30228            }
30229            return Ok(());
30230        }
30231        let sp = fa_split_keys(t_kv_upper, n_head_kv);
30232        #[allow(clippy::manual_div_ceil)]
30233        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30234        let n_splits_max = (t_kv_upper + sp - 1) / sp;
30235        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
30236        let (nspm, spk) = (n_splits_max as i32, sp as i32);
30237        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30238        let gqa = (n_head / n_head_kv).max(1) as u32;
30239        let o_len = t * n_head * n_splits_max * head_dim;
30240        let ml_len = t * n_head * n_splits_max;
30241        let mut part_guard = self.fa_part_pool.lock().unwrap();
30242        if part_guard
30243            .as_ref()
30244            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30245            .unwrap_or(true)
30246        {
30247            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
30248            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
30249            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
30250            // later live allocations land at those addresses, and the next graph REPLAY writes
30251            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
30252            // output corruption began the burst after the trunk's t_kv growth first realloc'd
30253            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
30254            // the baked addresses alive (single-stream: eager writes the new buffers, replays
30255            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
30256            // (total retired < final size).
30257            let old = part_guard.take();
30258            let (co, cm) = old
30259                .as_ref()
30260                .map(|pp| (pp.0.len(), pp.1.len()))
30261                .unwrap_or((0, 0));
30262            if let Some(old) = old {
30263                self.fa_part_retired.lock().unwrap().push(old);
30264            }
30265            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
30266                eprintln!(
30267                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
30268                    co, o_len, cm, ml_len
30269                );
30270            }
30271            *part_guard =
30272                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30273        }
30274        let pg = part_guard.as_mut().unwrap();
30275        self.gpu
30276            .stream()
30277            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
30278        self.gpu
30279            .stream()
30280            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
30281        self.gpu
30282            .stream()
30283            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
30284        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30285        let f = self.func("fa_decode_vec_q_rows_v3_dc");
30286        let sh = (32 * head_dim * 2) as u32;
30287        use cudarc::driver::sys::CUfunction_attribute_enum as A;
30288        f.set_attribute(
30289            A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30290            sh as i32,
30291        )?;
30292        let cfg = LaunchConfig {
30293            grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
30294            block_dim: (32, gqa, 1),
30295            shared_mem_bytes: sh,
30296        };
30297        let __s_b = self.gpu.stream();
30298        let mut b = __s_b.launch_builder(&f);
30299        b.arg(q)
30300            .arg(k)
30301            .arg(v)
30302            .arg(&mut *part_o)
30303            .arg(&mut *part_m)
30304            .arg(&mut *part_l)
30305            .arg(&hd)
30306            .arg(&nh)
30307            .arg(&nhkv)
30308            .arg(base_dev)
30309            .arg(&scale)
30310            .arg(&nspm)
30311            .arg(&spk)
30312            .arg(&ktb)
30313            .arg(&vtb);
30314        unsafe {
30315            b.launch(cfg)?;
30316        }
30317        let fc = self.func("fa_decode_combine_rows_dc");
30318        let cfg2 = LaunchConfig {
30319            grid_dim: (n_head as u32, t as u32, 1),
30320            block_dim: (head_dim as u32, 1, 1),
30321            shared_mem_bytes: 0,
30322        };
30323        let plus0 = 0i32;
30324        let __s_b2 = self.gpu.stream();
30325        let mut b2 = __s_b2.launch_builder(&fc);
30326        b2.arg(&*part_o)
30327            .arg(&*part_m)
30328            .arg(&*part_l)
30329            .arg(o)
30330            .arg(&hd)
30331            .arg(&nh)
30332            .arg(base_dev)
30333            .arg(&plus0)
30334            .arg(&nspm)
30335            .arg(&spk);
30336        unsafe {
30337            b2.launch(cfg2)?;
30338        }
30339        Ok(())
30340    }
30341
30342    /// Device-counter variant of `fa_decode` (CUDA-GRAPH-PLAN Phase 2). The sequence length is read
30343    /// from `t_kv_dev[0]` (resident device i32[1]) for the attention loop bound + per-split key range;
30344    /// the GRID `n_splits` is sized for `bucket_max` (the bucket's max t_kv — baked at capture time).
30345    /// Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the
30346    /// shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.
30347    ///
30348    /// BIT-IDENTITY (the gate): pass `bucket_max == actual_t_kv` and this reproduces `fa_decode`
30349    /// EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from
30350    /// device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split
30351    /// grouping (different but mathematically-equal log-sum-exp merge).
30352    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
30353    pub fn fa_decode_dc(
30354        &self,
30355        q: &CudaSlice<f32>,
30356        k: &cudarc::driver::CudaView<u8>,
30357        v: &cudarc::driver::CudaView<u8>,
30358        o: &mut CudaSlice<f32>,
30359        head_dim: usize,
30360        n_head: usize,
30361        n_head_kv: usize,
30362        t_kv_dev: &CudaSlice<i32>,
30363        bucket_max: usize,
30364        scale: f32,
30365        k_tok_bytes: usize,
30366        v_tok_bytes: usize,
30367        g: bool,
30368    ) -> Result<(), Box<dyn std::error::Error>> {
30369        self.fa_decode_dc_q8(
30370            q,
30371            k,
30372            v,
30373            o,
30374            head_dim,
30375            n_head,
30376            n_head_kv,
30377            t_kv_dev,
30378            bucket_max,
30379            scale,
30380            k_tok_bytes,
30381            v_tok_bytes,
30382            g,
30383            None,
30384        )
30385    }
30386
30387    /// `fa_decode_dc` with an optional q8_1 sink (wave 5b): when `q8_out` is given the
30388    /// combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.
30389    #[allow(clippy::too_many_arguments)]
30390    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30391    pub fn fa_decode_dc_q8(
30392        &self,
30393        q: &CudaSlice<f32>,
30394        k: &cudarc::driver::CudaView<u8>,
30395        v: &cudarc::driver::CudaView<u8>,
30396        o: &mut CudaSlice<f32>,
30397        head_dim: usize,
30398        n_head: usize,
30399        n_head_kv: usize,
30400        t_kv_dev: &CudaSlice<i32>,
30401        bucket_max: usize,
30402        scale: f32,
30403        k_tok_bytes: usize,
30404        v_tok_bytes: usize,
30405        g: bool,
30406        q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
30407    ) -> Result<(), Box<dyn std::error::Error>> {
30408        // The fa_vec gate + n_splits are sized from bucket_max (host, fixed at capture). The kernel
30409        // reads the ACTUAL t_kv from t_kv_dev for the per-split bound. DEFAULT-ON to MATCH the eager
30410        // `fa_decode` gate above — graph capture must mirror eager's kernel choice or the graph-vs-eager
30411        // bit-identity gate breaks. MEMRA_NO_FA_VEC forces scalar on BOTH paths in lockstep.
30412        // `g` = this layer's cache is e4m3 (gemma windowed under wkv) — every pick below must
30413        // mirror fa_decode_kvmod's g-routing or the graph diverges from eager (short/mid 1/96,
30414        // 2026-07-12).
30415        let mut fa_vec =
30416            std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
30417        if g && head_dim == 256 && !fa_v4_at(bucket_max) {
30418            fa_vec = false;
30419        } // mirror kvmod/geom
30420        let sp = fa_split_keys(bucket_max, n_head_kv);
30421        let n_splits = if fa_vec {
30422            ((bucket_max + sp - 1) / sp).max(1)
30423        } else {
30424            ((bucket_max + 255) / 256).max(1)
30425        };
30426        let o_len = n_head * n_splits * head_dim;
30427        let ml_len = n_head * n_splits;
30428        let mut part_guard = self.fa_part_pool.lock().unwrap();
30429        if part_guard
30430            .as_ref()
30431            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30432            .unwrap_or(true)
30433        {
30434            // RETIRE-ON-GROW, never free (#68 root cause, 2026-08-04): captured graphs (the
30435            // per-session persistent draft graph, the decode/prime graph doors) BAKE these pool
30436            // buffer addresses. Dropping the old buffers on grow returns them to the async pool,
30437            // later live allocations land at those addresses, and the next graph REPLAY writes
30438            // its fa partials over them — the ST serve-spec corruption (acceptance collapse +
30439            // output corruption began the burst after the trunk's t_kv growth first realloc'd
30440            // this pool past the draft-capture size; research/fp8ship-20260804). Retiring keeps
30441            // the baked addresses alive (single-stream: eager writes the new buffers, replays
30442            // touch only the old — never concurrently). Doubling growth bounds retired VRAM
30443            // (total retired < final size).
30444            let old = part_guard.take();
30445            let (co, cm) = old
30446                .as_ref()
30447                .map(|pp| (pp.0.len(), pp.1.len()))
30448                .unwrap_or((0, 0));
30449            if let Some(old) = old {
30450                self.fa_part_retired.lock().unwrap().push(old);
30451            }
30452            if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
30453                eprintln!(
30454                    "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
30455                    co, o_len, cm, ml_len
30456                );
30457            }
30458            *part_guard =
30459                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30460        }
30461        let pg = part_guard.as_mut().unwrap();
30462        self.gpu
30463            .stream()
30464            .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
30465        self.gpu
30466            .stream()
30467            .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
30468        self.gpu
30469            .stream()
30470            .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
30471        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30472        let (hd, nh, nhkv, nsp) = (
30473            head_dim as i32,
30474            n_head as i32,
30475            n_head_kv as i32,
30476            n_splits as i32,
30477        );
30478        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30479        let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
30480        // FA-DEEP pick keyed on bucket_max (the fa_v4_at precedent) — bit-identical twins,
30481        // so a threshold falling between t_kv and bucket_max cannot diverge eager-vs-graph.
30482        let deep = fa_vec
30483            && head_dim == 256
30484            && fa_v4_at(bucket_max)
30485            && !g
30486            && fa_deep_at(bucket_max)
30487            && !matches!(fa_v4_mode(), "noB3" | "stage");
30488        let (f, cfg) = if fa_vec
30489            && head_dim == 512
30490            && bucket_max >= {
30491                static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
30492                *FA512_MIN_DC.get_or_init(|| {
30493                    std::env::var("MEMRA_FA512_MIN")
30494                        .ok()
30495                        .and_then(|v| v.parse().ok())
30496                        .unwrap_or(512)
30497                })
30498            } {
30499            // gemma globals dc twin (mirror the eager dpl16 pick incl the crossover floor).
30500            let gqa = (n_head / n_head_kv).max(1) as u32;
30501            (
30502                self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
30503                LaunchConfig {
30504                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30505                    block_dim: (32, gqa, 1),
30506                    shared_mem_bytes: 0,
30507                },
30508            )
30509        } else if fa_vec && head_dim == 512 {
30510            // under the 512 floor eager runs scalar — the SAME unified symbol, ctr non-null;
30511            // ns_eff in-kernel reproduces eager's ceil(t_kv/sp) partition for the LIVE len.
30512            let q_view = q.as_view();
30513            let mut o_view = o.as_view_mut();
30514            return self.fa_decode_scalar_unified(
30515                &q_view,
30516                k,
30517                v,
30518                &mut o_view,
30519                head_dim,
30520                n_head,
30521                n_head_kv,
30522                0,
30523                Some(t_kv_dev),
30524                scale,
30525                n_splits,
30526                sp,
30527                k_tok_bytes,
30528                v_tok_bytes,
30529                g,
30530                &mut *part_o,
30531                &mut *part_m,
30532                &mut *part_l,
30533                q8_out,
30534            );
30535        } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
30536            // gemma/qwen v4 dc twin (eager default lane) — capture must mirror eager's pick,
30537            // incl the g-module route + raw-e4m3 sV sizing.
30538            let gqa = (n_head / n_head_kv).max(1) as u32;
30539            let fv = if g {
30540                self.func_g("fa_decode_vec_q_v4_dc")
30541            } else if deep {
30542                self.func("fa_decode_vec_q_v4_deep_dc")
30543            } else {
30544                self.func("fa_decode_vec_q_v4_dc")
30545            };
30546            let shmem =
30547                (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
30548            use cudarc::driver::sys::CUfunction_attribute_enum as A;
30549            fv.set_attribute(
30550                A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30551                shmem as i32,
30552            )?;
30553            (
30554                fv,
30555                LaunchConfig {
30556                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30557                    block_dim: (32, gqa, 1),
30558                    shared_mem_bytes: shmem,
30559                },
30560            )
30561        } else if fa_vec && fa_v3_active(head_dim) {
30562            // FA v3 lane _dc twin: the captured graph must run the SAME walk body as eager
30563            // under MEMRA_FA_V3=1 (eager, rows-verify and graph switch together).
30564            let gqa = (n_head / n_head_kv).max(1) as u32;
30565            let fv = if g {
30566                self.func_g("fa_decode_vec_q_v3_dc")
30567            } else {
30568                self.func("fa_decode_vec_q_v3_dc")
30569            };
30570            let shmem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd]
30571            (
30572                fv,
30573                LaunchConfig {
30574                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30575                    block_dim: (32, gqa, 1),
30576                    shared_mem_bytes: shmem,
30577                },
30578            )
30579        } else if fa_vec && fa_v2_on() {
30580            // FAVENDOR lane: v2 _dc twin — the captured graph must run the SAME walk body as
30581            // eager under MEMRA_FA_V2=1 or graph_decode_gate's bit-identity breaks (the flag is
30582            // a numeric config; eager, rows-verify and graph all switch together).
30583            let gqa = (n_head / n_head_kv).max(1) as u32;
30584            let fv = if g {
30585                self.func_g("fa_decode_vec_q_v2_dc")
30586            } else {
30587                self.func("fa_decode_vec_q_v2_dc")
30588            };
30589            let shmem = (2 * 32 * head_dim * 2) as u32; // sK+sV bf16 [FA_DEC_TILE=32][hd]
30590            (
30591                fv,
30592                LaunchConfig {
30593                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30594                    block_dim: (32, gqa, 1),
30595                    shared_mem_bytes: shmem,
30596                },
30597            )
30598        } else if fa_vec {
30599            let gqa = (n_head / n_head_kv).max(1) as u32;
30600            // REGISTER-DEQUANT twin: zero dynamic smem (see fa_decode above).
30601            let fv = if g {
30602                self.func_g("fa_decode_vec_q_dc")
30603            } else {
30604                self.func("fa_decode_vec_q_dc")
30605            };
30606            (
30607                fv,
30608                LaunchConfig {
30609                    grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30610                    block_dim: (32, gqa, 1),
30611                    shared_mem_bytes: 0,
30612                },
30613            )
30614        } else {
30615            let q_view = q.as_view();
30616            let mut o_view = o.as_view_mut();
30617            return self.fa_decode_scalar_unified(
30618                &q_view,
30619                k,
30620                v,
30621                &mut o_view,
30622                head_dim,
30623                n_head,
30624                n_head_kv,
30625                0,
30626                Some(t_kv_dev),
30627                scale,
30628                n_splits,
30629                if fa_vec { sp } else { 256 },
30630                k_tok_bytes,
30631                v_tok_bytes,
30632                g,
30633                &mut *part_o,
30634                &mut *part_m,
30635                &mut *part_l,
30636                q8_out,
30637            );
30638        };
30639        let ski = sp as i32; // one-partition law: the twins derive ns_eff from (T_kv, ski)
30640        let __s_b = self.gpu.stream();
30641        let mut b = __s_b.launch_builder(&f);
30642        b.arg(q)
30643            .arg(k)
30644            .arg(v)
30645            .arg(&mut *part_o)
30646            .arg(&mut *part_m)
30647            .arg(&mut *part_l)
30648            .arg(&hd)
30649            .arg(&nh)
30650            .arg(&nhkv)
30651            .arg(t_kv_dev)
30652            .arg(&scale)
30653            .arg(&nsp)
30654            .arg(&ski)
30655            .arg(&ktb)
30656            .arg(&vtb);
30657        unsafe {
30658            b.launch(cfg)?;
30659        }
30660        let cfg2 = LaunchConfig {
30661            grid_dim: (n_head as u32, 1, 1),
30662            block_dim: (head_dim as u32, 1, 1),
30663            shared_mem_bytes: 0,
30664        };
30665        if let Some((oq, od)) = q8_out {
30666            let fc = if g {
30667                self.func_g("fa_decode_combine_q8_1")
30668            } else {
30669                self.fa_func("fa_decode_combine_q8_1", head_dim)
30670            };
30671            let __s_b2 = self.gpu.stream();
30672            let mut b2 = __s_b2.launch_builder(&fc);
30673            b2.arg(&*part_o)
30674                .arg(&*part_m)
30675                .arg(&*part_l)
30676                .arg(oq)
30677                .arg(od)
30678                .arg(&hd)
30679                .arg(&nh)
30680                .arg(&nsp);
30681            unsafe {
30682                b2.launch(cfg2)?;
30683            }
30684            return Ok(());
30685        }
30686        let fc = if g {
30687            self.func_g("fa_decode_combine_f32")
30688        } else {
30689            self.fa_func("fa_decode_combine_f32", head_dim)
30690        };
30691        let __s_b2 = self.gpu.stream();
30692        let mut b2 = __s_b2.launch_builder(&fc);
30693        b2.arg(&*part_o)
30694            .arg(&*part_m)
30695            .arg(&*part_l)
30696            .arg(o)
30697            .arg(&hd)
30698            .arg(&nh)
30699            .arg(&nsp);
30700        unsafe {
30701            b2.launch(cfg2)?;
30702        }
30703        Ok(())
30704    }
30705
30706    /// _dcw append (t=1): physical write row = len_dev[0] - base_dev[0] in-kernel; follow
30707    /// with `inc_i32(len_dev)` on the same stream. Bit-identical bytes to the host-row append
30708    /// at equal rows.
30709    #[allow(clippy::too_many_arguments)]
30710    pub fn append_kv_quantized_dcw(
30711        &self,
30712        k_row: &CudaSlice<f32>,
30713        v_row: &CudaSlice<f32>,
30714        kc: &mut CudaSlice<u8>,
30715        vc: &mut CudaSlice<u8>,
30716        len_dev: &CudaSlice<i32>,
30717        base_dev: Option<&CudaSlice<i32>>,
30718        kv_dim_k: usize,
30719        kv_dim_v: usize,
30720        k_tok_bytes: usize,
30721        v_tok_bytes: usize,
30722    ) -> Result<(), Box<dyn std::error::Error>> {
30723        let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
30724        let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
30725        let cfg = LaunchConfig {
30726            grid_dim: (nblk, 1, 1),
30727            block_dim: (32, 1, 1),
30728            shared_mem_bytes: 0,
30729        };
30730        let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
30731        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30732        let null: u64 = 0;
30733        let __s_b = self.gpu.stream();
30734        let mut b = __s_b.launch_builder(&f);
30735        b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
30736        match base_dev {
30737            Some(base) => {
30738                b.arg(base);
30739            }
30740            None => {
30741                b.arg(&null);
30742            }
30743        }
30744        b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
30745        unsafe {
30746            b.launch(cfg)?;
30747        }
30748        Ok(())
30749    }
30750
30751    /// Increment a device i32 counter (graph-capturable; the `inc_i32` kernel).
30752    pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
30753        let f = self.func("inc_i32");
30754        let cfg = LaunchConfig {
30755            grid_dim: (1, 1, 1),
30756            block_dim: (1, 1, 1),
30757            shared_mem_bytes: 0,
30758        };
30759        let __s_b = self.gpu.stream();
30760        let mut b = __s_b.launch_builder(&f);
30761        b.arg(counter);
30762        unsafe {
30763            b.launch(cfg)?;
30764        }
30765        Ok(())
30766    }
30767
30768    /// Windowed device-counter fa decode (step TP graph increment A): the KV view derives
30769    /// entirely from device state — `len_dev` (staged length), `base_dev` (physical row of
30770    /// logical 0 after the last ring rebase; None reads as 0), and `window` (0 = global) — so
30771    /// a captured child replays with ZERO per-token node updates. v3-vec only (the default
30772    /// kernel class on this lane); callers keep eager below the vec floor and for any other
30773    /// class. Scratch comes from the engine's fa partial pool sized at `bucket_max` (for SWA
30774    /// layers pass min(bucket, window)); the pool's retire-on-grow keeps captured addresses
30775    /// alive across bucket growth.
30776    #[allow(clippy::too_many_arguments)]
30777    /// Retire-on-grow ensure for the fa partial pool (see the #68 comment on the eager
30778    /// twin). Split out so graph capture can pre-run it OUTSIDE the capture region — an
30779    /// alloc inside a captured section becomes a mem node, and child graphs reject those.
30780    /// THE ONE PLACE THE FA PARTIAL POOL IS ALLOCATED.
30781    ///
30782    /// Eight call sites grow this pool and all eight retire-on-grow correctly, but only ONE
30783    /// of them carried the `[fa-pool] grow` receipt, so that receipt under-reported grows by
30784    /// seven eighths and no grow could honestly be dated against a request. Routing every
30785    /// grower through here makes the count real. The receipt names the site so a ladder can
30786    /// be attributed, and stays bounded so a pathological ladder cannot flood a serving log.
30787    ///
30788    /// `MEMRA_FA_PART_ZERO=1` (DEFAULT OFF, diagnostic only) zeroes the fresh buffers. A grow
30789    /// hands every subsequent launch three UNINITIALIZED banks; if the poison is a combine
30790    /// reading a partial bank its producer never wrote, that makes every row and every head
30791    /// non-finite at once, which is the shape the level-2 bad-row bitmap reports at the
30792    /// global-attention join.
30793    ///
30794    /// READ IT IN ONE DIRECTION ONLY. Zeroed banks carry m = 0.0, not NEG_INF, so the
30795    /// empty-split no-op guard never engages: a bank that is entirely unwritten still
30796    /// combines to L = 0 and O/L = 0/0 = NaN. So **silence under this arm convicts the pool;
30797    /// continued trapping acquits nothing**, because only the PARTIALLY unwritten class (real
30798    /// splits beside stale zeroed ones) goes quiet. Discriminator, never a fix, and never a
30799    /// serving arm: where it does go quiet the output is still wrong, it just looks plausible.
30800    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
30801    fn fa_part_alloc(
30802        &self,
30803        o_len: usize,
30804        ml_len: usize,
30805        co: usize,
30806        cm: usize,
30807    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
30808        static GROWS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
30809        let n = GROWS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
30810        if n < 64 {
30811            eprintln!(
30812                "[fa-pool] grow #{n} dev={} o_len {co} -> {o_len} ml_len {cm} -> {ml_len} (retired kept, zero={})",
30813                self.ctx().ordinal(),
30814                fa_part_zero_on()
30815            );
30816        }
30817        let mut po = self.alloc_uninit::<f32>(o_len)?;
30818        let mut pm = self.alloc_uninit::<f32>(ml_len)?;
30819        let mut pl = self.alloc_uninit::<f32>(ml_len)?;
30820        if fa_part_zero_on() {
30821            self.gpu.stream().memset_zeros(&mut po)?;
30822            self.gpu.stream().memset_zeros(&mut pm)?;
30823            self.gpu.stream().memset_zeros(&mut pl)?;
30824        }
30825        Ok((po, pm, pl))
30826    }
30827
30828    fn fa_part_pool_grow(
30829        &self,
30830        part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
30831        o_len: usize,
30832        ml_len: usize,
30833    ) -> Result<(), Box<dyn std::error::Error>> {
30834        if part_guard
30835            .as_ref()
30836            .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30837            .unwrap_or(true)
30838        {
30839            let old = part_guard.take();
30840            let (co, cm) = old
30841                .as_ref()
30842                .map(|pp| (pp.0.len(), pp.1.len()))
30843                .unwrap_or((0, 0));
30844            if let Some(old) = old {
30845                self.fa_part_retired.lock().unwrap().push(old);
30846            }
30847            // GROW RECEIPT. This pool is grow-only, retires-on-grow and never frees, and every
30848            // FA decode/verify launch in the process reads and writes it. A grow is therefore a
30849            // process-lifetime EVENT — new addresses, a retired buffer kept alive forever, and
30850            // a different partial layout — and it is invisible in every log we have. The step37
30851            // spec fault is clean for the first two or three requests of a process and then
30852            // poisons trunk layer 20 (research: MEMRA_SPEC_NAN_SCAN), which is exactly the
30853            // shape a mid-life pool grow would produce, so the grows have to be datable
30854            // against the requests. Cap raised from 8 after the first run measured FOUR
30855            // grows per device (380928 -> 761856 -> 1523712 -> 3047424): with two devices the
30856            // 8 slots were spent before any grow could be dated against a request, which was
30857            // the entire point of the receipt. Still bounded so a pathological ladder cannot
30858            // flood a serving log.
30859            *part_guard =
30860                Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30861        }
30862        Ok(())
30863    }
30864
30865    /// Pre-grow the fa partial pool for a dcw call at (n_head, bucket_max) geometry, from
30866    /// OUTSIDE any capture region. Idempotent and cheap when already big enough.
30867    pub fn fa_dcw_pool_ensure(
30868        &self,
30869        head_dim: usize,
30870        n_head: usize,
30871        n_head_kv: usize,
30872        bucket_max: usize,
30873    ) -> Result<(), Box<dyn std::error::Error>> {
30874        let sp = fa_split_keys(bucket_max, n_head_kv);
30875        #[allow(clippy::manual_div_ceil)]
30876        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30877        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
30878        let o_len = n_head * n_splits * head_dim;
30879        let ml_len = n_head * n_splits;
30880        let mut part_guard = self.fa_part_pool.lock().unwrap();
30881        self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
30882    }
30883
30884    /// T=2 dcw decode attention (MEMRA_SPEC_FA2): both verify columns' rows are ALREADY
30885    /// appended; one launch walks the KV stream once with two query rows (per-row causal
30886    /// bounds len-1 / len) and the per-row combine consumes each half of the partials.
30887    /// BIT-IDENTICAL per row to that row's own per-column launch under the equal-partition
30888    /// guard the caller enforces (ns_eff/per equal for both bounds; boundary rounds fall
30889    /// back per column). `q2` = [2, n_head, head_dim]; `o2` = [2, n_head*head_dim] gated
30890    /// outputs (the head gate fuses into the combine as in the t=1 path).
30891    #[allow(clippy::too_many_arguments)]
30892    pub fn fa_decode_dcw2(
30893        &self,
30894        q2: &CudaSlice<f32>,
30895        k_ring: &cudarc::driver::CudaView<u8>,
30896        v_ring: &cudarc::driver::CudaView<u8>,
30897        o2: &mut CudaSlice<f32>,
30898        head_dim: usize,
30899        n_head: usize,
30900        n_head_kv: usize,
30901        len_dev: &CudaSlice<i32>,
30902        base_dev: Option<&CudaSlice<i32>>,
30903        window: usize,
30904        bucket_max: usize,
30905        scale: f32,
30906        k_tok_bytes: usize,
30907        v_tok_bytes: usize,
30908        gate2: &CudaSlice<f32>,
30909    ) -> Result<(), Box<dyn std::error::Error>> {
30910        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
30911        if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) || !fa_v3_on() {
30912            return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
30913        }
30914        let sp = fa_split_keys(bucket_max, n_head_kv);
30915        #[allow(clippy::manual_div_ceil)]
30916        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
30917        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
30918        // Partials for BOTH rows: row-major halves.
30919        let o_len = 2 * n_head * n_splits * head_dim;
30920        let ml_len = 2 * n_head * n_splits;
30921        let mut part_guard = self.fa_part_pool.lock().unwrap();
30922        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
30923        let pg = part_guard.as_mut().unwrap();
30924        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30925        let (hd, nh, nhkv, nsp) = (
30926            head_dim as i32,
30927            n_head as i32,
30928            n_head_kv as i32,
30929            n_splits as i32,
30930        );
30931        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30932        let (ski, win) = (sp as i32, window as i32);
30933        let gqa = (n_head / n_head_kv).max(1) as u32;
30934        let smem = (32 * head_dim * 2) as u32;
30935        let f = self.func("fa_decode_vec_q_v3_dcw2");
30936        let cfg = LaunchConfig {
30937            grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30938            block_dim: (32, gqa, 1),
30939            shared_mem_bytes: smem,
30940        };
30941        let null: u64 = 0;
30942        {
30943            let __s_b = self.gpu.stream();
30944            let mut b = __s_b.launch_builder(&f);
30945            b.arg(q2)
30946                .arg(k_ring)
30947                .arg(v_ring)
30948                .arg(&mut *part_o)
30949                .arg(&mut *part_m)
30950                .arg(&mut *part_l)
30951                .arg(&hd)
30952                .arg(&nh)
30953                .arg(&nhkv)
30954                .arg(len_dev);
30955            match base_dev {
30956                Some(base) => {
30957                    b.arg(base);
30958                }
30959                None => {
30960                    b.arg(&null);
30961                }
30962            }
30963            b.arg(&win)
30964                .arg(&scale)
30965                .arg(&nsp)
30966                .arg(&ski)
30967                .arg(&ktb)
30968                .arg(&vtb);
30969            unsafe {
30970                b.launch(cfg)?;
30971            }
30972        }
30973        // Per-row combine+gate: the t=1 combine kernel over each half (its `head` axis spans
30974        // 2*n_head rows laid out row-major, and the gate rows are stacked the same way), so
30975        // one launch covers both rows with the exact t=1 program per (row, head).
30976        let fc = {
30977            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
30978            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
30979                self.func("fa_decode_combine_gate_f32_s")
30980            } else {
30981                self.func("fa_decode_combine_gate_f32")
30982            }
30983        };
30984        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
30985        let nh2 = (2 * n_head) as i32;
30986        let cfg2 = LaunchConfig {
30987            grid_dim: ((2 * n_head) as u32, 1, 1),
30988            block_dim: (head_dim as u32, 1, 1),
30989            shared_mem_bytes: if combine_shared {
30990                (2 * n_splits * 4) as u32
30991            } else {
30992                0
30993            },
30994        };
30995        let __s_b2 = self.gpu.stream();
30996        let mut b2 = __s_b2.launch_builder(&fc);
30997        b2.arg(&*part_o)
30998            .arg(&*part_m)
30999            .arg(&*part_l)
31000            .arg(gate2)
31001            .arg(o2)
31002            .arg(&hd)
31003            .arg(&nh2)
31004            .arg(&nsp);
31005        unsafe {
31006            b2.launch(cfg2)?;
31007        }
31008        Ok(())
31009    }
31010
31011    /// T-ROW dcw decode attention over a per-row session table (the per-session
31012    /// distributed-KV primitive). `tab` = t entries of five u64 words {k_ring, v_ring,
31013    /// len_ptr, base_ptr, len_back}; every (row, head, split) block runs the t=1 dcw
31014    /// program verbatim with that row's ring/len/base and its own split geometry, so each
31015    /// row is bit-identical to its own per-row launch. The kernel embeds the big-rig
31016    /// split ladder, so this refuses when the ladder env overrides are armed or the rig
31017    /// is not the >=128-SM class. `q_rows` = [t, n_head, head_dim]; `o_rows` = [t,
31018    /// n_head*head_dim] gated; `gate_rows` = [t, n_head].
31019    #[allow(clippy::too_many_arguments)]
31020    pub fn fa_decode_dcw_rows(
31021        &self,
31022        q_rows: &CudaSlice<f32>,
31023        tab: &CudaSlice<u64>,
31024        o_rows: &mut CudaSlice<f32>,
31025        t: usize,
31026        head_dim: usize,
31027        n_head: usize,
31028        n_head_kv: usize,
31029        window: usize,
31030        max_ns: usize,
31031        scale: f32,
31032        k_tok_bytes: usize,
31033        v_tok_bytes: usize,
31034        gate_rows: &CudaSlice<f32>,
31035    ) -> Result<(), Box<dyn std::error::Error>> {
31036        if std::env::var("MEMRA_NO_FA_VEC").is_ok()
31037            || head_dim > 256
31038            || !head_dim.is_multiple_of(32)
31039            || !fa_v3_on()
31040        {
31041            return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
31042        }
31043        if fa_sm_count() < 128
31044            || std::env::var("MEMRA_FA_SPLIT").is_ok()
31045            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
31046            || std::env::var("MEMRA_FA_SP16").is_ok()
31047        {
31048            return Err(
31049                "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
31050                 (or a <128-SM rig) keep the per-row path"
31051                    .into(),
31052            );
31053        }
31054        if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
31055            return Err("fa_decode_dcw_rows geometry".into());
31056        }
31057        let o_len = t * n_head * max_ns * head_dim;
31058        let ml_len = t * n_head * max_ns;
31059        let mut part_guard = self.fa_part_pool.lock().unwrap();
31060        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
31061        let pg = part_guard.as_mut().unwrap();
31062        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
31063        let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
31064        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
31065        let (win, mns) = (window as i32, max_ns as i32);
31066        let gqa = (n_head / n_head_kv).max(1) as u32;
31067        let smem = (32 * head_dim * 2) as u32;
31068        let f = self.func("fa_decode_vec_q_v3_dcw_rows");
31069        let cfg = LaunchConfig {
31070            grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
31071            block_dim: (32, gqa, 1),
31072            shared_mem_bytes: smem,
31073        };
31074        {
31075            let __s_b = self.gpu.stream();
31076            let mut b = __s_b.launch_builder(&f);
31077            b.arg(q_rows)
31078                .arg(tab)
31079                .arg(&mut *part_o)
31080                .arg(&mut *part_m)
31081                .arg(&mut *part_l)
31082                .arg(&hd)
31083                .arg(&nh)
31084                .arg(&nhkv)
31085                .arg(&win)
31086                .arg(&scale)
31087                .arg(&mns)
31088                .arg(&ktb)
31089                .arg(&vtb);
31090            unsafe {
31091                b.launch(cfg)?;
31092            }
31093        }
31094        // Per-(row, head) combine+gate: the t=1 combine over t*n_head stacked heads —
31095        // row r head h reads its own partial bank; splits past a row's ns_eff carry
31096        // (-inf, 0) partials the NEG_INF guard no-ops bit-exactly.
31097        let fc = {
31098            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31099            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
31100                self.func("fa_decode_combine_gate_f32_s")
31101            } else {
31102                self.func("fa_decode_combine_gate_f32")
31103            }
31104        };
31105        let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
31106        let nht = (t * n_head) as i32;
31107        let cfg2 = LaunchConfig {
31108            grid_dim: ((t * n_head) as u32, 1, 1),
31109            block_dim: (head_dim as u32, 1, 1),
31110            shared_mem_bytes: if combine_shared {
31111                (2 * max_ns * 4) as u32
31112            } else {
31113                0
31114            },
31115        };
31116        let __s_b2 = self.gpu.stream();
31117        let mut b2 = __s_b2.launch_builder(&fc);
31118        b2.arg(&*part_o)
31119            .arg(&*part_m)
31120            .arg(&*part_l)
31121            .arg(gate_rows)
31122            .arg(o_rows)
31123            .arg(&hd)
31124            .arg(&nht)
31125            .arg(&mns);
31126        unsafe {
31127            b2.launch(cfg2)?;
31128        }
31129        Ok(())
31130    }
31131
31132    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31133    pub fn fa_decode_dcw(
31134        &self,
31135        q: &CudaSlice<f32>,
31136        k_ring: &cudarc::driver::CudaView<u8>,
31137        v_ring: &cudarc::driver::CudaView<u8>,
31138        o: &mut CudaSlice<f32>,
31139        head_dim: usize,
31140        n_head: usize,
31141        n_head_kv: usize,
31142        len_dev: &CudaSlice<i32>,
31143        base_dev: Option<&CudaSlice<i32>>,
31144        window: usize,
31145        bucket_max: usize,
31146        scale: f32,
31147        k_tok_bytes: usize,
31148        v_tok_bytes: usize,
31149        // FUSION #2d: Some(gate_row) fuses the head gate into the combine (bit-identical,
31150        // one launch saved); `o` then receives the GATED output and the caller skips its
31151        // attn_head_gate call.
31152        fused_gate: Option<&CudaSlice<f32>>,
31153    ) -> Result<(), Box<dyn std::error::Error>> {
31154        let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
31155        if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) || !fa_v3_on() {
31156            return Err("fa_decode_dcw supports the default v3-vec class only                         (bucket >= vec floor, head_dim <= 256, MEMRA_FA_V3 on);                         keep eager outside it"
31157                .into());
31158        }
31159        let sp = fa_split_keys(bucket_max, n_head_kv);
31160        #[allow(clippy::manual_div_ceil)]
31161        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
31162        let n_splits = ((bucket_max + sp - 1) / sp).max(1);
31163        let o_len = n_head * n_splits * head_dim;
31164        let ml_len = n_head * n_splits;
31165        let mut part_guard = self.fa_part_pool.lock().unwrap();
31166        Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
31167        let pg = part_guard.as_mut().unwrap();
31168        // MEMRA_FA_DCW_MEMSET=0: skip the partial-pool zeroing — every (head, split) in
31169        // [0, nsp) writes its partial before the combine reads it (per = ceil(len/nsp), so
31170        // split s starts at s*per < len for all s < nsp), making the zeros dead stores.
31171        // Door-gated pending the identity battery; =0 saves 3 memset launches/rank/layer.
31172        static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31173        // Token-graph capture ALWAYS keeps the memsets: the retarget path (increment C)
31174        // finds the attention children BY their three-memset signature and updates the
31175        // memset widths per bucket — capturing without them silently kills retargeting
31176        // (battery-v8 token drift, 2026-08-21).
31177        let memset_on = *MEMSET_ON
31178            .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
31179            || crate::tp::token_graph_building();
31180        if memset_on {
31181            self.gpu
31182                .stream()
31183                .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
31184            self.gpu
31185                .stream()
31186                .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
31187            self.gpu
31188                .stream()
31189                .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
31190        }
31191        let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
31192        let (hd, nh, nhkv, nsp) = (
31193            head_dim as i32,
31194            n_head as i32,
31195            n_head_kv as i32,
31196            n_splits as i32,
31197        );
31198        let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
31199        let (ski, win) = (sp as i32, window as i32);
31200        let gqa = (n_head / n_head_kv).max(1) as u32;
31201        let smem = (32 * head_dim * 2) as u32; // sV bf16 [FA_DEC_TILE=32][hd] (v3 uses sV only)
31202        // MEMRA_FA_UNROLL=8: the B1-unroll-8 twin (deeper K load pipeline, bit-identical —
31203        // see fa_dec_v3_walk_u). Same launch geometry.
31204        static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31205        static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
31206        let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
31207            Ok("2") => 2,
31208            Ok("1") => 1,
31209            _ => 0,
31210        });
31211        // MEMRA_FA_PROF=1: clock64() phase profile of the decode-attention walk. ncu is
31212        // permission-blocked in this container and the module params are not exposed, so this
31213        // is how the ~1.18us/key gets localised. Diagnostic only (extra atomics per block);
31214        // prints cumulative cycle shares every 430 launches.
31215        static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31216        let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
31217        static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
31218            std::sync::Mutex::new(None);
31219        // MEMRA_FA_HSPLIT=2: split each kv_head's gqa warp group across TWO blocks (2x grid,
31220        // duplicated Phase A staging) — bit-identical per (head, split). Tests whether B1's
31221        // 59-63% cycle share is occupancy-starved latency (grid is only n_head_kv x n_splits).
31222        static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31223        let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
31224            && (n_head / n_head_kv).is_multiple_of(2)
31225            && (n_head / n_head_kv) >= 2;
31226        let f = if fprof {
31227            self.func("fa_decode_vec_q_v3_dcw_prof")
31228        } else if hs2 {
31229            self.func("fa_decode_vec_q_v3_dcw_hs2")
31230        } else if hoist == 2 {
31231            // + typed 4-byte K loads (memcpy from uint8_t* can lower to byte loads).
31232            self.func("fa_decode_vec_q_v3_dcw_hc")
31233        } else if hoist == 1 {
31234            // Loop-invariant K alignment class hoisted out of B1 (bit-identical).
31235            self.func("fa_decode_vec_q_v3_dcw_h")
31236        } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
31237            self.func("fa_decode_vec_q_v3_dcw_u8")
31238        } else {
31239            self.func("fa_decode_vec_q_v3_dcw")
31240        };
31241        let cfg = LaunchConfig {
31242            grid_dim: if hs2 {
31243                ((2 * n_head_kv) as u32, n_splits as u32, 1)
31244            } else {
31245                (n_head_kv as u32, n_splits as u32, 1)
31246            },
31247            block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
31248            shared_mem_bytes: smem,
31249        };
31250        let null: u64 = 0;
31251        let __s_b = self.gpu.stream();
31252        let mut b = __s_b.launch_builder(&f);
31253        b.arg(q)
31254            .arg(k_ring)
31255            .arg(v_ring)
31256            .arg(&mut *part_o)
31257            .arg(&mut *part_m)
31258            .arg(&mut *part_l)
31259            .arg(&hd)
31260            .arg(&nh)
31261            .arg(&nhkv)
31262            .arg(len_dev);
31263        match base_dev {
31264            Some(base) => {
31265                b.arg(base);
31266            }
31267            None => {
31268                b.arg(&null);
31269            }
31270        }
31271        b.arg(&win)
31272            .arg(&scale)
31273            .arg(&nsp)
31274            .arg(&ski)
31275            .arg(&ktb)
31276            .arg(&vtb);
31277        if fprof {
31278            let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
31279            if guard
31280                .as_ref()
31281                .is_none_or(|(d, _)| *d != self.ctx().ordinal())
31282            {
31283                *guard = Some((self.ctx().ordinal(), self.htod_u64(&[0u64; 8])?));
31284            }
31285            let (_, buf) = guard.as_mut().expect("armed above");
31286            b.arg(&*buf);
31287            unsafe {
31288                b.launch(cfg)?;
31289            }
31290            static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
31291            let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
31292            if n.is_multiple_of(430) {
31293                self.stream().synchronize()?;
31294                let h = self.dtoh_u64(buf)?;
31295                let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
31296                let tot: u64 = h[..6].iter().sum();
31297                let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
31298                for (i, name) in phases.iter().enumerate() {
31299                    let pct = if tot > 0 {
31300                        h[i] as f64 / tot as f64 * 100.0
31301                    } else {
31302                        0.0
31303                    };
31304                    line.push_str(&format!(" {name}={pct:.1}%"));
31305                }
31306                if h[6] > 0 {
31307                    line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
31308                }
31309                eprintln!("{line}");
31310            }
31311        } else {
31312            unsafe {
31313                b.launch(cfg)?;
31314            }
31315        }
31316        let mut combine_shared = false;
31317        let fc = if fused_gate.is_some() {
31318            // MEMRA_FA_COMBINE_S=1: shared-staged split metadata (bit-identical; kills the
31319            // n_splits-deep dependent global load chain every thread used to walk twice).
31320            static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31321            if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
31322                combine_shared = true;
31323                self.func("fa_decode_combine_gate_f32_s")
31324            } else {
31325                self.func("fa_decode_combine_gate_f32")
31326            }
31327        } else {
31328            self.fa_func("fa_decode_combine_f32", head_dim)
31329        };
31330        let cfg2 = LaunchConfig {
31331            grid_dim: (n_head as u32, 1, 1),
31332            block_dim: (head_dim as u32, 1, 1),
31333            shared_mem_bytes: if combine_shared {
31334                (2 * n_splits * 4) as u32
31335            } else {
31336                0
31337            },
31338        };
31339        let __s_b2 = self.gpu.stream();
31340        let mut b2 = __s_b2.launch_builder(&fc);
31341        b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
31342        if let Some(gate_row) = fused_gate {
31343            b2.arg(gate_row);
31344        }
31345        b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
31346        unsafe {
31347            b2.launch(cfg2)?;
31348        }
31349        Ok(())
31350    }
31351
31352    /// EAGER fa_decode geometry for a given actual `t_kv` (CUDA-GRAPH-PLAN §3.3 bucketing). Returns
31353    /// `(fa_vec, n_splits)` EXACTLY as `fa_decode` computes them so the graph-capture path can key its
31354    /// bucket on the same `(kernel, n_splits)` pair and pass a `bucket_max` that reproduces eager's
31355    /// n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside
31356    /// the kernel and matches eager when n_splits matches — the bit-identity contract.)
31357    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
31358    pub fn fa_geom_eager(
31359        &self,
31360        t_kv: usize,
31361        head_dim: usize,
31362        n_head_kv: usize,
31363        g: bool,
31364    ) -> (bool, usize) {
31365        // MUST mirror `fa_decode` / `fa_decode_dc` (default-ON 2026-06-28). This is the bucket-key
31366        // source: if it disagrees with the actual kernel pick, the graph captures the wrong path and
31367        // replay diverges from eager. All three sites read MEMRA_NO_FA_VEC in lockstep.
31368        let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
31369        // hd512 dpl16 vec lane (gemma globals, 2026-07-11 graph-arc fix): the original key
31370        // hardcoded vec = hd<=256, so for hd512 it bucketed by the SCALAR 256-key splits while
31371        // the dpl16/rows_dpl16 kernels split by the ladder — n_splits changed WITHIN a bucket
31372        // (mid-ctx graph mismatch at pos 19 + partials OOB at longer runs). Mirror the real
31373        // fa_decode dispatch: vec512 above the fa512 floor, vec256 as before.
31374        let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
31375        let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim.is_multiple_of(32));
31376        // g (fp8-windowed): mirror kvmod's clamp — only the v4 lane parses e4m3 in the vec
31377        // family; everything else falls to the g-module scalar.
31378        // hd256 under g: v4-or-scalar. hd128/other under g ride the REGISTER g-lane (the
31379        // dq_K_lane/dq_V_lane macros are format-aware; v3 is format-gated off, v2/smem
31380        // arms are excluded under g). Mirrored in kvmod / fa_decode_dc / fa_geom_eager.
31381        if g && head_dim == 256 && !fa_v4_at(t_kv) {
31382            fa_vec = false;
31383        }
31384        let sp = fa_split_keys(t_kv, n_head_kv);
31385        let n_splits = if fa_vec {
31386            ((t_kv + sp - 1) / sp).max(1)
31387        } else {
31388            ((t_kv + 255) / 256).max(1)
31389        };
31390        (fa_vec, n_splits)
31391    }
31392
31393    /// `bucket_max` (host t_kv to feed `fa_decode_dc` / `full_attn_decode_dc`) that makes the _dc
31394    /// kernel pick the SAME (fa_vec, n_splits) as eager would for actual `t_kv`. Because the dc
31395    /// launcher derives both from `bucket_max` via the same formulas, we just hand it `t_kv` itself:
31396    /// the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in
31397    /// the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is `(fa_vec, n_splits)`.
31398    pub fn fa_bucket_key(
31399        &self,
31400        t_kv: usize,
31401        head_dim: usize,
31402        n_head_kv: usize,
31403        g: bool,
31404    ) -> (bool, usize) {
31405        self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
31406    }
31407
31408    /// CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs `step`
31409    /// inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time
31410    /// kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine's
31411    /// decode stream (RELAXED mode) and instantiates it into a replayable `CudaGraph`. The closure
31412    /// must enqueue ONLY device work on `e.stream()` (no dtoh / no synchronize / no host branch on
31413    /// device data) — every per-step varying scalar must come from a device counter. Returns the
31414    /// instantiated graph; `CudaGraph::launch()` replays the whole step in one dispatch.
31415    /// `capture_graph` with CAPTURE-RETAIN: every Engine allocation made during the warmups
31416    /// and the capture is kept alive in the returned keeper — hold it as long as the graph
31417    /// replays (transients returning to the pool get reused by unrelated work and corrupt
31418    /// replays; the draft-graph root cause). Model-generic, next capture reuses it.
31419    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
31420    pub fn capture_graph_retained<F>(
31421        &self,
31422        step: F,
31423    ) -> Result<
31424        (
31425            cudarc::driver::CudaGraph,
31426            Vec<Box<dyn std::any::Any + Send>>,
31427        ),
31428        Box<dyn std::error::Error>,
31429    >
31430    where
31431        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
31432    {
31433        use cudarc::driver::sys::CUgraphInstantiate_flags;
31434        self.capture_graph_retained_flags(
31435            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
31436            step,
31437        )
31438    }
31439
31440    /// Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs
31441    /// (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of
31442    /// AUTO_FREE_ON_LAUNCH: the auto-free flag's launch-time mem-pool scan was measured at
31443    /// ~0.25us/node (205us on the 826-node step) even with nothing to free.
31444    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
31445    pub fn capture_graph_retained_flags<F>(
31446        &self,
31447        flags: cudarc::driver::sys::CUgraphInstantiate_flags,
31448        mut step: F,
31449    ) -> Result<
31450        (
31451            cudarc::driver::CudaGraph,
31452            Vec<Box<dyn std::any::Any + Send>>,
31453        ),
31454        Box<dyn std::error::Error>,
31455    >
31456    where
31457        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
31458    {
31459        use cudarc::driver::sys::CUstreamCaptureMode;
31460        // KEEP scope = WARMUPS ONLY (2026-07-13): keep_if_capturing retains via
31461        // CudaSlice::clone, which is a device ALLOC + D2D COPY on the stream — clones made
31462        // while the capture region is open become dead copy NODES replayed every launch
31463        // (E4B: 1440 copies = 0.74ms/token, the whole graph-vs-eager regression). The
31464        // warmup runs allocate the same transient sequence at the same pool addresses, so
31465        // retaining the warmup clones preserves the draft-graph fix without polluting the
31466        // captured graph.
31467        self.capture_keep.lock().unwrap().clear();
31468        let was_tracking = self.gpu.ctx.is_event_tracking();
31469        if was_tracking {
31470            unsafe {
31471                self.gpu.ctx.disable_event_tracking();
31472            }
31473        }
31474        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
31475            self.capture_keep_on
31476                .store(true, std::sync::atomic::Ordering::Relaxed);
31477            let w = (|| {
31478                step(self)?;
31479                step(self)
31480            })();
31481            self.capture_keep_on
31482                .store(false, std::sync::atomic::Ordering::Relaxed);
31483            w?;
31484            self.gpu.stream().synchronize()?;
31485            self.gpu
31486                .stream()
31487                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
31488            let r = step(self);
31489            let g = self.gpu.stream().end_capture(flags);
31490            r?;
31491            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
31492            graph.upload()?;
31493            Ok(graph)
31494        };
31495        let result = run();
31496        self.capture_keep_on
31497            .store(false, std::sync::atomic::Ordering::Relaxed);
31498        if was_tracking {
31499            unsafe {
31500                self.gpu.ctx.enable_event_tracking();
31501            }
31502        }
31503        let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
31504        Ok((result?, keeper))
31505    }
31506
31507    /// Retained capture WITHOUT the two warmup executions. The warmups exist for transient
31508    /// pool-address stability (draft-graph lanes); the step TP token-graph sections are
31509    /// alloc-free with persistent operands, and their bodies carry device side effects
31510    /// (dcw KV appends + counter incs) that a warmup would REALLY EXECUTE — measured as a
31511    /// +2/rank len_d drift per bucket build that marched appends past the ring planes.
31512    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
31513    pub fn capture_graph_retained_nowarm<F>(
31514        &self,
31515        mut step: F,
31516    ) -> Result<
31517        (
31518            cudarc::driver::CudaGraph,
31519            Vec<Box<dyn std::any::Any + Send>>,
31520        ),
31521        Box<dyn std::error::Error>,
31522    >
31523    where
31524        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
31525    {
31526        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
31527        let was_tracking = self.gpu.ctx.is_event_tracking();
31528        if was_tracking {
31529            unsafe {
31530                self.gpu.ctx.disable_event_tracking();
31531            }
31532        }
31533        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
31534            self.gpu.stream().synchronize()?;
31535            self.gpu
31536                .stream()
31537                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
31538            let r = step(self);
31539            let g = self.gpu.stream().end_capture(
31540                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
31541            );
31542            r?;
31543            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
31544            graph.upload()?;
31545            Ok(graph)
31546        };
31547        let result = run();
31548        if was_tracking {
31549            unsafe {
31550                self.gpu.ctx.enable_event_tracking();
31551            }
31552        }
31553        Ok((result?, Vec::new()))
31554    }
31555
31556    pub fn capture_graph<F>(
31557        &self,
31558        mut step: F,
31559    ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
31560    where
31561        F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
31562    {
31563        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
31564        // EVENT TRACKING OFF for capture. The Engine creates a 2nd stream (copy_stream) so cudarc is in
31565        // multi-stream mode and, by default, records a CudaEvent per CudaSlice alloc/use to serialize
31566        // cross-stream access. Those per-buffer event waits issue stream ops that are NOT permitted
31567        // inside a capture region (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). The captured decode step is
31568        // strictly SINGLE-STREAM (every kernel on gpu.stream), so this synchronization is unnecessary
31569        // here — disable it for the whole warmup+capture, re-enable after. SAFETY: the decode-dc path
31570        // touches only gpu.stream; no buffer crosses to copy_stream during capture.
31571        let was_tracking = self.gpu.ctx.is_event_tracking();
31572        if was_tracking {
31573            unsafe {
31574                self.gpu.ctx.disable_event_tracking();
31575            }
31576        }
31577        // Q1 PROBE (MEMRA_GRAPH_IFLAG): the generic capture body's cuMemAllocAsync nodes are
31578        // EXACTLY BALANCED by in-graph free nodes (measured census q27: 1589 ALLOC / 1589
31579        // FREE), so AUTO_FREE_ON_LAUNCH has nothing to reclaim at launch — it only pays its
31580        // per-node launch-time mem-pool scan. `upload` / `none` select the alternatives to
31581        // measure that scan's real cost on the generic path. Diagnostic door only; the
31582        // default stays AUTO_FREE until a measured A/B justifies moving it.
31583        let iflag = {
31584            static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
31585            *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
31586                // UPLOAD = the gemma slotted door's zero-mem-node choice; PRIORITY = the flag
31587                // hybrid_forward.rs:5935 actually ships (both drop the auto-free launch scan).
31588                Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
31589                Ok("priority") => {
31590                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
31591                }
31592                _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
31593            })
31594        };
31595        // MEMRA_GRAPH_CAPTIME=1 (Q1 lane): phase-resolved capture cost. Recapture is paid at
31596        // every kernel-class crossing, so it — not steady-state decode — is the quantity a
31597        // mem-node reduction could plausibly shrink. Only `instantiate` (cuStreamEndCapture +
31598        // cuGraphInstantiateWithFlags) and `upload` scale with node count; the warmups are
31599        // eager step executions and are node-count-invariant. Printing the split bounds the
31600        // refactor's ceiling instead of assuming it.
31601        let ct = {
31602            static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31603            *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
31604        };
31605        // MEMRA_GRAPH_WARMUPS (Q1 lane; DEFAULT 1 since lane/graph-warmups 2026-08-05): the
31606        // phase split showed the eager warmups are 80% of recapture cost (q27 27.4 of 34.4 ms
31607        // pod / 42% of 52.6 ms 5090) — 3x larger than the ENTIRE mem-node ceiling the audit
31608        // chased, and node-count-invariant, so no capture-body refactor could touch it.
31609        // Warmup 2's theorized job was async-pool ADDRESS STABILITY: warmup 1's allocs may
31610        // grow/map the pool, warmup 2 re-walks the same sequence over the freed blocks so the
31611        // captured third run bakes settled addresses. That hazard is the #68 stale-baked-
31612        // address class — which the engine now guards STRUCTURALLY rather than by re-walking:
31613        // in-body transients are captured as BALANCED in-graph alloc/free node pairs (census
31614        // 1589/1589 — replays allocate for themselves; no baked transient pointers), every
31615        // externally-referenced buffer is stable-pointer by design (fa_part_pool retires-on-
31616        // grow and never frees, resident counters/scratch, cache set in place), and the
31617        // draft-graph path additionally rides capture_graph_retained (capture_keep holds all
31618        // warmup+capture allocs alive). One warmup therefore suffices for kernel-attr
31619        // settling and pool mapping. Arbitrated adversarially, not by taste:
31620        // graph-warmup-stress (pool-growth cycles large<->small x10, overlap arm, forced
31621        // recaptures over freed blocks — bit-identity vs eager + canary teeth) is GREEN at
31622        // warmups=1 on the deployment rig, plus graph-decode-gate 256-step bit-identity,
31623        // graph-session-gate, run-spec K=1..8 (receipts research/graph-warmups-5090-20260805/
31624        // + the pod's research/graph-allocfree-20260805/). Measured: recapture -38..-42% q27 /
31625        // -41% q9, decode +~1%, capture+prime -13ms. MEMRA_GRAPH_WARMUPS=2 = the rollback
31626        // seam; tools/graph-warmup-stress-gate.sh = the gate any regression re-runs.
31627        let warmups = {
31628            static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
31629            *W.get_or_init(|| {
31630                std::env::var("MEMRA_GRAPH_WARMUPS")
31631                    .ok()
31632                    .and_then(|v| v.parse().ok())
31633                    .filter(|n| *n >= 1)
31634                    .unwrap_or(1)
31635            })
31636        };
31637        let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
31638            let t_w = std::time::Instant::now();
31639            // warmup: inline runs (no capture) so allocator pointers + kernel attrs are stable.
31640            for _ in 0..warmups {
31641                step(self)?;
31642            }
31643            self.gpu.stream().synchronize()?;
31644            let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
31645            // capture the third run.
31646            let t_c = std::time::Instant::now();
31647            self.gpu
31648                .stream()
31649                .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
31650            // If the body errors mid-capture, end the capture before propagating so the stream isn't
31651            // left in a capturing state.
31652            let r = step(self);
31653            let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
31654            let t_i = std::time::Instant::now();
31655            let g = self.gpu.stream().end_capture(iflag);
31656            let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
31657            r?;
31658            let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
31659            let t_u = std::time::Instant::now();
31660            graph.upload()?;
31661            if ct {
31662                println!(
31663                    "[graph-captime] warmup2x {ms_warm:.2} ms  capture-body {ms_body:.2} ms  \
31664                          instantiate {ms_inst:.2} ms  upload {:.2} ms",
31665                    t_u.elapsed().as_secs_f64() * 1e3
31666                );
31667            }
31668            Ok(graph)
31669        };
31670        let result = run();
31671        if was_tracking {
31672            unsafe {
31673                self.gpu.ctx.enable_event_tracking();
31674            }
31675        }
31676        result
31677    }
31678
31679    /// gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).
31680    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31681    pub fn gdn_scan_s128_view(
31682        &self,
31683        q: &CudaSlice<f32>,
31684        k: &CudaSlice<f32>,
31685        v: &CudaSlice<f32>,
31686        g: &CudaSlice<f32>,
31687        beta: &CudaSlice<f32>,
31688        state_in: &cudarc::driver::CudaView<f32>,
31689        state_out: &mut cudarc::driver::CudaViewMut<f32>,
31690        o: &mut CudaSlice<f32>,
31691        n_head: usize,
31692        t: usize,
31693        scale: f32,
31694    ) -> Result<(), Box<dyn std::error::Error>> {
31695        let f = self.func("gdn_scan_s128");
31696        const S_V: u32 = 128;
31697        const WARP: u32 = 32;
31698        const COLS: u32 = 4;
31699        let cfg = LaunchConfig {
31700            grid_dim: (n_head as u32, 1, S_V / COLS),
31701            block_dim: (WARP, COLS, 1),
31702            shared_mem_bytes: 0,
31703        };
31704        let (h, ti) = (n_head as i32, t as i32);
31705        let __s_b = self.gpu.stream();
31706        let mut b = __s_b.launch_builder(&f);
31707        b.arg(q)
31708            .arg(k)
31709            .arg(v)
31710            .arg(g)
31711            .arg(beta)
31712            .arg(state_in)
31713            .arg(state_out)
31714            .arg(o)
31715            .arg(&h)
31716            .arg(&ti)
31717            .arg(&scale);
31718        unsafe {
31719            b.launch(cfg)?;
31720        }
31721        Ok(())
31722    }
31723
31724    /// conv1d where the input is a CudaView (resident conv state assembled in place).
31725    #[allow(clippy::too_many_arguments)]
31726    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31727    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
31728    pub fn ssm_conv1d_view(
31729        &self,
31730        x: &cudarc::driver::CudaView<f32>,
31731        w: &CudaSlice<f32>,
31732        y: &mut CudaSlice<f32>,
31733        conv_dim: usize,
31734        t: usize,
31735        d_conv: usize,
31736        silu: bool,
31737    ) -> Result<(), Box<dyn std::error::Error>> {
31738        let f = self.func("ssm_conv1d_silu_f32");
31739        // grid.x = channel, grid.y = T-tiles (block 256 strides over T) — parallel over both axes.
31740        let cfg = LaunchConfig {
31741            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
31742            block_dim: (256, 1, 1),
31743            shared_mem_bytes: 0,
31744        };
31745        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
31746        let __s_b = self.gpu.stream();
31747        let mut b = __s_b.launch_builder(&f);
31748        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
31749        unsafe {
31750            b.launch(cfg)?;
31751        }
31752        Ok(())
31753    }
31754
31755    /// Depthwise causal conv1d + optional SiLU.
31756    /// x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state),
31757    /// w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major.
31758    /// FUSED prefill conv (token-major input, zero left-state): replaces
31759    /// transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output
31760    /// directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.
31761    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
31762    pub fn ssm_conv1d_tm(
31763        &self,
31764        qkv_tm: &CudaSlice<f32>,
31765        w: &CudaSlice<f32>,
31766        y: &mut CudaSlice<f32>,
31767        conv_dim: usize,
31768        t: usize,
31769        d_conv: usize,
31770    ) -> Result<(), Box<dyn std::error::Error>> {
31771        let f = self.func("ssm_conv1d_tm_f32");
31772        let cfg = LaunchConfig {
31773            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31774            block_dim: (256, 1, 1),
31775            shared_mem_bytes: 0,
31776        };
31777        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31778        let __s_b = self.gpu.stream();
31779        let mut b = __s_b.launch_builder(&f);
31780        b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
31781        unsafe {
31782            b.launch(cfg)?;
31783        }
31784        Ok(())
31785    }
31786
31787    /// BATCHED verify conv (T>1, carried state): window reads the resident conv ring for
31788    /// negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the
31789    /// T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path);
31790    /// T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the
31791    /// update kernel would race reading the ring it rewrites, so that arm clones the ring
31792    /// (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input
31793    /// columns; the final ring == what T sequential decode ring rolls leave).
31794    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31795    pub fn ssm_conv1d_tm_state(
31796        &self,
31797        qkv_tm: &CudaSlice<f32>,
31798        conv_state: &mut CudaSlice<f32>,
31799        w: &CudaSlice<f32>,
31800        y: &mut CudaSlice<f32>,
31801        conv_dim: usize,
31802        t: usize,
31803        d_conv: usize,
31804    ) -> Result<(), Box<dyn std::error::Error>> {
31805        self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
31806    }
31807
31808    /// task #14: `pad_len` = device true length for PADDED prime graphs — the ring update
31809    /// reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.
31810    #[allow(clippy::too_many_arguments)]
31811    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
31812    pub fn ssm_conv1d_tm_state_pad(
31813        &self,
31814        qkv_tm: &CudaSlice<f32>,
31815        conv_state: &mut CudaSlice<f32>,
31816        w: &CudaSlice<f32>,
31817        y: &mut CudaSlice<f32>,
31818        conv_dim: usize,
31819        t: usize,
31820        d_conv: usize,
31821        pad_len: Option<&CudaSlice<i32>>,
31822    ) -> Result<(), Box<dyn std::error::Error>> {
31823        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
31824        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
31825        // the window kernel both read the pre-roll ring; the roll launches after both) — but
31826        // cloning first keeps the ordering trivially correct under any future stream split.
31827        let ring_old = if t < d_conv - 1 {
31828            Some(self.clone_dtod(conv_state)?)
31829        } else {
31830            None
31831        };
31832        {
31833            let f = self.func("ssm_conv1d_tm_state_f32");
31834            let cfg = LaunchConfig {
31835                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31836                block_dim: (256, 1, 1),
31837                shared_mem_bytes: 0,
31838            };
31839            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31840            let __s_b = self.gpu.stream();
31841            let mut b = __s_b.launch_builder(&f);
31842            b.arg(qkv_tm)
31843                .arg(&*conv_state)
31844                .arg(w)
31845                .arg(y)
31846                .arg(&cd)
31847                .arg(&ti)
31848                .arg(&dc);
31849            unsafe {
31850                b.launch(cfg)?;
31851            }
31852        }
31853        match (ring_old, pad_len) {
31854            (None, Some(len_d)) => {
31855                let f = self.func("ssm_conv_ring_update_dev_f32");
31856                let n = conv_dim * (d_conv - 1);
31857                let cfg = LaunchConfig::for_num_elems(n as u32);
31858                let (cd, dc) = (conv_dim as i32, d_conv as i32);
31859                let __s_b = self.gpu.stream();
31860                let mut b = __s_b.launch_builder(&f);
31861                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
31862                unsafe {
31863                    b.launch(cfg)?;
31864                }
31865            }
31866            (None, None) => {
31867                let f = self.func("ssm_conv_ring_update_f32");
31868                let n = conv_dim * (d_conv - 1);
31869                let cfg = LaunchConfig::for_num_elems(n as u32);
31870                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31871                let __s_b = self.gpu.stream();
31872                let mut b = __s_b.launch_builder(&f);
31873                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
31874                unsafe {
31875                    b.launch(cfg)?;
31876                }
31877            }
31878            (Some(old), _) => {
31879                self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
31880            }
31881        }
31882        Ok(())
31883    }
31884
31885    /// qkv-view twin (task #16): batched prime reads the concat GEMM output directly.
31886    #[allow(clippy::too_many_arguments)]
31887    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
31888    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
31889    pub fn ssm_conv1d_tm_state_pad_v(
31890        &self,
31891        qkv_tm: &cudarc::driver::CudaView<f32>,
31892        conv_state: &mut CudaSlice<f32>,
31893        w: &CudaSlice<f32>,
31894        y: &mut CudaSlice<f32>,
31895        conv_dim: usize,
31896        t: usize,
31897        d_conv: usize,
31898        pad_len: Option<&CudaSlice<i32>>,
31899    ) -> Result<(), Box<dyn std::error::Error>> {
31900        assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
31901        // clone BEFORE the window kernel is issued is not required (stream-ordered: the dtod and
31902        // the window kernel both read the pre-roll ring; the roll launches after both) — but
31903        // cloning first keeps the ordering trivially correct under any future stream split.
31904        let ring_old = if t < d_conv - 1 {
31905            Some(self.clone_dtod(conv_state)?)
31906        } else {
31907            None
31908        };
31909        {
31910            let f = self.func("ssm_conv1d_tm_state_f32");
31911            let cfg = LaunchConfig {
31912                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31913                block_dim: (256, 1, 1),
31914                shared_mem_bytes: 0,
31915            };
31916            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31917            let __s_b = self.gpu.stream();
31918            let mut b = __s_b.launch_builder(&f);
31919            b.arg(qkv_tm)
31920                .arg(&*conv_state)
31921                .arg(w)
31922                .arg(y)
31923                .arg(&cd)
31924                .arg(&ti)
31925                .arg(&dc);
31926            unsafe {
31927                b.launch(cfg)?;
31928            }
31929        }
31930        match (ring_old, pad_len) {
31931            (None, Some(len_d)) => {
31932                let f = self.func("ssm_conv_ring_update_dev_f32");
31933                let n = conv_dim * (d_conv - 1);
31934                let cfg = LaunchConfig::for_num_elems(n as u32);
31935                let (cd, dc) = (conv_dim as i32, d_conv as i32);
31936                let __s_b = self.gpu.stream();
31937                let mut b = __s_b.launch_builder(&f);
31938                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
31939                unsafe {
31940                    b.launch(cfg)?;
31941                }
31942            }
31943            (None, None) => {
31944                let f = self.func("ssm_conv_ring_update_f32");
31945                let n = conv_dim * (d_conv - 1);
31946                let cfg = LaunchConfig::for_num_elems(n as u32);
31947                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31948                let __s_b = self.gpu.stream();
31949                let mut b = __s_b.launch_builder(&f);
31950                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
31951                unsafe {
31952                    b.launch(cfg)?;
31953                }
31954            }
31955            (Some(_), _) => unreachable!(
31956                "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
31957            ),
31958        }
31959        Ok(())
31960    }
31961
31962    /// PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring
31963    /// with the state a T=1 chain holds after only the FIRST `tc` columns of `qkv_tm` — the last
31964    /// `pad` entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no
31965    /// arithmetic, cannot perturb FP order). `ring_old` = the pre-round snapshot ring.
31966    pub fn ssm_conv_ring_rebuild(
31967        &self,
31968        qkv_tm: &CudaSlice<f32>,
31969        ring_old: &CudaSlice<f32>,
31970        conv_state: &mut CudaSlice<f32>,
31971        conv_dim: usize,
31972        tc: usize,
31973        d_conv: usize,
31974    ) -> Result<(), Box<dyn std::error::Error>> {
31975        let f = self.func("ssm_conv_ring_rebuild_f32");
31976        let n = conv_dim * (d_conv - 1);
31977        let cfg = LaunchConfig::for_num_elems(n as u32);
31978        let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
31979        let __s_b = self.gpu.stream();
31980        let mut b = __s_b.launch_builder(&f);
31981        b.arg(qkv_tm)
31982            .arg(ring_old)
31983            .arg(conv_state)
31984            .arg(&cd)
31985            .arg(&ti)
31986            .arg(&dc);
31987        unsafe {
31988            b.launch(cfg)?;
31989        }
31990        Ok(())
31991    }
31992
31993    /// FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch.
31994    /// Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane
31995    /// warp tree (vs l2_norm_f32's 256-thread two-level tree) — same math, different FP sum order;
31996    /// the argmax + run-spec gates are the authority.
31997    #[allow(clippy::too_many_arguments)]
31998    pub fn gdn_prep_decode(
31999        &self,
32000        conv_out: &CudaSlice<f32>,
32001        beta_raw: &CudaSlice<f32>,
32002        alpha: &CudaSlice<f32>,
32003        dt_bias: &CudaSlice<f32>,
32004        a: &CudaSlice<f32>,
32005        q_l2: &mut CudaSlice<f32>,
32006        k_l2: &mut CudaSlice<f32>,
32007        v_g: &mut CudaSlice<f32>,
32008        beta: &mut CudaSlice<f32>,
32009        g_log: &mut CudaSlice<f32>,
32010        d_state: usize,
32011        num_v: usize,
32012        num_k: usize,
32013        key_dim: usize,
32014        eps: f32,
32015    ) -> Result<(), Box<dyn std::error::Error>> {
32016        let f = self.func("gdn_prep_decode_f32");
32017        let cfg = LaunchConfig {
32018            grid_dim: (num_v as u32, 1, 1),
32019            block_dim: (32, 4, 1),
32020            shared_mem_bytes: 0,
32021        };
32022        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
32023        let __s_b = self.gpu.stream();
32024        let mut b = __s_b.launch_builder(&f);
32025        b.arg(conv_out)
32026            .arg(beta_raw)
32027            .arg(alpha)
32028            .arg(dt_bias)
32029            .arg(a)
32030            .arg(q_l2)
32031            .arg(k_l2)
32032            .arg(v_g)
32033            .arg(beta)
32034            .arg(g_log)
32035            .arg(&ds)
32036            .arg(&nv)
32037            .arg(&nk)
32038            .arg(&kd)
32039            .arg(&eps);
32040        unsafe {
32041            b.launch(cfg)?;
32042        }
32043        Ok(())
32044    }
32045
32046    /// FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out
32047    /// materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches
32048    /// qkv_to_gdn_repack's modulo head-repeat mapping exactly.
32049    #[allow(clippy::too_many_arguments)]
32050    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
32051    pub fn ssm_conv1d_gdn(
32052        &self,
32053        qkv_tm: &CudaSlice<f32>,
32054        w: &CudaSlice<f32>,
32055        q_g: &mut CudaSlice<f32>,
32056        k_g: &mut CudaSlice<f32>,
32057        v_g: &mut CudaSlice<f32>,
32058        conv_dim: usize,
32059        t: usize,
32060        d_conv: usize,
32061        d_state: usize,
32062        num_v: usize,
32063        num_k: usize,
32064        key_dim: usize,
32065    ) -> Result<(), Box<dyn std::error::Error>> {
32066        let f = self.func("ssm_conv1d_gdn_f32");
32067        let cfg = LaunchConfig {
32068            grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
32069            block_dim: (256, 1, 1),
32070            shared_mem_bytes: 0,
32071        };
32072        let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
32073        let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
32074        let __s_b = self.gpu.stream();
32075        let mut b = __s_b.launch_builder(&f);
32076        b.arg(qkv_tm)
32077            .arg(w)
32078            .arg(q_g)
32079            .arg(k_g)
32080            .arg(v_g)
32081            .arg(&cd)
32082            .arg(&ti)
32083            .arg(&dc)
32084            .arg(&ds)
32085            .arg(&nv)
32086            .arg(&nk)
32087            .arg(&kd);
32088        unsafe {
32089            b.launch(cfg)?;
32090        }
32091        Ok(())
32092    }
32093
32094    #[allow(clippy::too_many_arguments)]
32095    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
32096    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
32097    pub fn ssm_conv1d(
32098        &self,
32099        x: &CudaSlice<f32>,
32100        w: &CudaSlice<f32>,
32101        y: &mut CudaSlice<f32>,
32102        conv_dim: usize,
32103        t: usize,
32104        d_conv: usize,
32105        silu: bool,
32106    ) -> Result<(), Box<dyn std::error::Error>> {
32107        let f = self.func("ssm_conv1d_silu_f32");
32108        let cfg = LaunchConfig {
32109            grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
32110            block_dim: (256, 1, 1),
32111            shared_mem_bytes: 0,
32112        };
32113        let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
32114        let __s_b = self.gpu.stream();
32115        let mut b = __s_b.launch_builder(&f);
32116        b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
32117        unsafe {
32118            b.launch(cfg)?;
32119        }
32120        Ok(())
32121    }
32122
32123    /// Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed;
32124    /// o:[128,H,T]. Single sequence.
32125    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
32126    pub fn gdn_scan_s128(
32127        &self,
32128        q: &CudaSlice<f32>,
32129        k: &CudaSlice<f32>,
32130        v: &CudaSlice<f32>,
32131        g: &CudaSlice<f32>,
32132        beta: &CudaSlice<f32>,
32133        state_in: &CudaSlice<f32>,
32134        state_out: &mut CudaSlice<f32>,
32135        o: &mut CudaSlice<f32>,
32136        n_head: usize,
32137        t: usize,
32138        scale: f32,
32139    ) -> Result<(), Box<dyn std::error::Error>> {
32140        let f = self.func("gdn_scan_s128");
32141        const S_V: u32 = 128;
32142        const WARP: u32 = 32;
32143        const COLS_PER_BLOCK: u32 = 4;
32144        let cfg = LaunchConfig {
32145            grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
32146            block_dim: (WARP, COLS_PER_BLOCK, 1),
32147            shared_mem_bytes: 0,
32148        };
32149        let (h, ti) = (n_head as i32, t as i32);
32150        let __s_b = self.gpu.stream();
32151        let mut b = __s_b.launch_builder(&f);
32152        b.arg(q)
32153            .arg(k)
32154            .arg(v)
32155            .arg(g)
32156            .arg(beta)
32157            .arg(state_in)
32158            .arg(state_out)
32159            .arg(o)
32160            .arg(&h)
32161            .arg(&ti)
32162            .arg(&scale);
32163        unsafe {
32164            b.launch(cfg)?;
32165        }
32166        Ok(())
32167    }
32168
32169    // ==== B2' batched decode state ops (decode_batch.rs) ====
32170    // Per-seq state pointers ride device u64 arrays (views into the per-step pointer table).
32171    // Bodies are the single-seq kernels per sequence — bit-identical per row.
32172
32173    #[allow(clippy::too_many_arguments)]
32174    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
32175    pub fn ssm_conv1d_fused_decode_b(
32176        &self,
32177        qkv_cols: &CudaSlice<f32>,
32178        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
32179        w: &CudaSlice<f32>,
32180        conv_outs: &mut CudaSlice<f32>,
32181        conv_dim: usize,
32182        d_conv: usize,
32183        b_n: usize,
32184    ) -> Result<(), Box<dyn std::error::Error>> {
32185        let f = self.func("ssm_conv1d_fused_decode_b_f32");
32186        let cfg = LaunchConfig {
32187            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
32188            block_dim: (256, 1, 1),
32189            shared_mem_bytes: 0,
32190        };
32191        let (cd, dc) = (conv_dim as i32, d_conv as i32);
32192        let __s_b = self.gpu.stream();
32193        let mut b = __s_b.launch_builder(&f);
32194        b.arg(qkv_cols)
32195            .arg(conv_state_ptrs)
32196            .arg(w)
32197            .arg(conv_outs)
32198            .arg(&cd)
32199            .arg(&dc);
32200        unsafe {
32201            b.launch(cfg)?;
32202        }
32203        Ok(())
32204    }
32205
32206    #[allow(clippy::too_many_arguments)]
32207    pub fn gdn_prep_decode_b(
32208        &self,
32209        conv_outs: &CudaSlice<f32>,
32210        beta_raws: &CudaSlice<f32>,
32211        alphas: &CudaSlice<f32>,
32212        dt_bias: &CudaSlice<f32>,
32213        a: &CudaSlice<f32>,
32214        q_l2: &mut CudaSlice<f32>,
32215        k_l2: &mut CudaSlice<f32>,
32216        v_g: &mut CudaSlice<f32>,
32217        beta: &mut CudaSlice<f32>,
32218        g_log: &mut CudaSlice<f32>,
32219        d_state: usize,
32220        num_v: usize,
32221        num_k: usize,
32222        key_dim: usize,
32223        eps: f32,
32224        conv_dim: usize,
32225        b_n: usize,
32226    ) -> Result<(), Box<dyn std::error::Error>> {
32227        let f = self.func("gdn_prep_decode_b_f32");
32228        let cfg = LaunchConfig {
32229            grid_dim: (num_v as u32, 1, b_n as u32),
32230            block_dim: (32, 4, 1),
32231            shared_mem_bytes: 0,
32232        };
32233        let (ds, nv, nk, kd, cd) = (
32234            d_state as i32,
32235            num_v as i32,
32236            num_k as i32,
32237            key_dim as i32,
32238            conv_dim as i32,
32239        );
32240        let __s_b = self.gpu.stream();
32241        let mut b = __s_b.launch_builder(&f);
32242        b.arg(conv_outs)
32243            .arg(beta_raws)
32244            .arg(alphas)
32245            .arg(dt_bias)
32246            .arg(a)
32247            .arg(q_l2)
32248            .arg(k_l2)
32249            .arg(v_g)
32250            .arg(beta)
32251            .arg(g_log)
32252            .arg(&ds)
32253            .arg(&nv)
32254            .arg(&nk)
32255            .arg(&kd)
32256            .arg(&eps)
32257            .arg(&cd);
32258        unsafe {
32259            b.launch(cfg)?;
32260        }
32261        Ok(())
32262    }
32263
32264    #[allow(clippy::too_many_arguments)]
32265    pub fn gdn_scan_s128_batched(
32266        &self,
32267        q: &CudaSlice<f32>,
32268        k: &CudaSlice<f32>,
32269        v: &CudaSlice<f32>,
32270        g: &CudaSlice<f32>,
32271        beta: &CudaSlice<f32>,
32272        state_in_ptrs: &cudarc::driver::CudaView<u64>,
32273        state_out_ptrs: &cudarc::driver::CudaView<u64>,
32274        o: &mut CudaSlice<f32>,
32275        n_head: usize,
32276        b_n: usize,
32277        scale: f32,
32278    ) -> Result<(), Box<dyn std::error::Error>> {
32279        let f = self.func("gdn_scan_s128_b");
32280        const S_V: u32 = 128;
32281        const WARP: u32 = 32;
32282        const COLS_PER_BLOCK: u32 = 4;
32283        let cfg = LaunchConfig {
32284            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
32285            block_dim: (WARP, COLS_PER_BLOCK, 1),
32286            shared_mem_bytes: 0,
32287        };
32288        let h = n_head as i32;
32289        let __s_b = self.gpu.stream();
32290        let mut b = __s_b.launch_builder(&f);
32291        b.arg(q)
32292            .arg(k)
32293            .arg(v)
32294            .arg(g)
32295            .arg(beta)
32296            .arg(state_in_ptrs)
32297            .arg(state_out_ptrs)
32298            .arg(o)
32299            .arg(&h)
32300            .arg(&scale);
32301        unsafe {
32302            b.launch(cfg)?;
32303        }
32304        Ok(())
32305    }
32306
32307    /// VIEW twins of the three GDN decode state wrappers (lane/verify-launchslim): identical
32308    /// launches, row args as views into the caller's packed [T, ...] buffers — the t-parallel
32309    /// verify's per-row loop passes slices instead of paying an arithmetic-free dtod per row
32310    /// per kernel (48 layers x T rows x 4 copies/round on the money path). Same kernels, same
32311    /// numeric class; only the pointer arithmetic moved host-side.
32312    #[allow(clippy::too_many_arguments)]
32313    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
32314    pub fn ssm_conv1d_fused_decode_b_view(
32315        &self,
32316        qkv_cols: &cudarc::driver::CudaView<f32>,
32317        conv_state_ptrs: &cudarc::driver::CudaView<u64>,
32318        w: &CudaSlice<f32>,
32319        conv_outs: &mut CudaSlice<f32>,
32320        conv_dim: usize,
32321        d_conv: usize,
32322        b_n: usize,
32323    ) -> Result<(), Box<dyn std::error::Error>> {
32324        let f = self.func("ssm_conv1d_fused_decode_b_f32");
32325        let cfg = LaunchConfig {
32326            grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
32327            block_dim: (256, 1, 1),
32328            shared_mem_bytes: 0,
32329        };
32330        let (cd, dc) = (conv_dim as i32, d_conv as i32);
32331        let __s_b = self.gpu.stream();
32332        let mut b = __s_b.launch_builder(&f);
32333        b.arg(qkv_cols)
32334            .arg(conv_state_ptrs)
32335            .arg(w)
32336            .arg(conv_outs)
32337            .arg(&cd)
32338            .arg(&dc);
32339        unsafe {
32340            b.launch(cfg)?;
32341        }
32342        Ok(())
32343    }
32344
32345    #[allow(clippy::too_many_arguments)]
32346    pub fn gdn_prep_decode_b_view(
32347        &self,
32348        conv_outs: &CudaSlice<f32>,
32349        beta_raws: &cudarc::driver::CudaView<f32>,
32350        alphas: &cudarc::driver::CudaView<f32>,
32351        dt_bias: &CudaSlice<f32>,
32352        a: &CudaSlice<f32>,
32353        q_l2: &mut CudaSlice<f32>,
32354        k_l2: &mut CudaSlice<f32>,
32355        v_g: &mut CudaSlice<f32>,
32356        beta: &mut CudaSlice<f32>,
32357        g_log: &mut CudaSlice<f32>,
32358        d_state: usize,
32359        num_v: usize,
32360        num_k: usize,
32361        key_dim: usize,
32362        eps: f32,
32363        conv_dim: usize,
32364        b_n: usize,
32365    ) -> Result<(), Box<dyn std::error::Error>> {
32366        let f = self.func("gdn_prep_decode_b_f32");
32367        let cfg = LaunchConfig {
32368            grid_dim: (num_v as u32, 1, b_n as u32),
32369            block_dim: (32, 4, 1),
32370            shared_mem_bytes: 0,
32371        };
32372        let (ds, nv, nk, kd, cd) = (
32373            d_state as i32,
32374            num_v as i32,
32375            num_k as i32,
32376            key_dim as i32,
32377            conv_dim as i32,
32378        );
32379        let __s_b = self.gpu.stream();
32380        let mut b = __s_b.launch_builder(&f);
32381        b.arg(conv_outs)
32382            .arg(beta_raws)
32383            .arg(alphas)
32384            .arg(dt_bias)
32385            .arg(a)
32386            .arg(q_l2)
32387            .arg(k_l2)
32388            .arg(v_g)
32389            .arg(beta)
32390            .arg(g_log)
32391            .arg(&ds)
32392            .arg(&nv)
32393            .arg(&nk)
32394            .arg(&kd)
32395            .arg(&eps)
32396            .arg(&cd);
32397        unsafe {
32398            b.launch(cfg)?;
32399        }
32400        Ok(())
32401    }
32402
32403    #[allow(clippy::too_many_arguments)]
32404    pub fn gdn_scan_s128_batched_view(
32405        &self,
32406        q: &CudaSlice<f32>,
32407        k: &CudaSlice<f32>,
32408        v: &CudaSlice<f32>,
32409        g: &CudaSlice<f32>,
32410        beta: &CudaSlice<f32>,
32411        state_in_ptrs: &cudarc::driver::CudaView<u64>,
32412        state_out_ptrs: &cudarc::driver::CudaView<u64>,
32413        o: &mut cudarc::driver::CudaViewMut<f32>,
32414        n_head: usize,
32415        b_n: usize,
32416        scale: f32,
32417    ) -> Result<(), Box<dyn std::error::Error>> {
32418        let f = self.func("gdn_scan_s128_b");
32419        const S_V: u32 = 128;
32420        const WARP: u32 = 32;
32421        const COLS_PER_BLOCK: u32 = 4;
32422        let cfg = LaunchConfig {
32423            grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
32424            block_dim: (WARP, COLS_PER_BLOCK, 1),
32425            shared_mem_bytes: 0,
32426        };
32427        let h = n_head as i32;
32428        let __s_b = self.gpu.stream();
32429        let mut b = __s_b.launch_builder(&f);
32430        b.arg(q)
32431            .arg(k)
32432            .arg(v)
32433            .arg(g)
32434            .arg(beta)
32435            .arg(state_in_ptrs)
32436            .arg(state_out_ptrs)
32437            .arg(o)
32438            .arg(&h)
32439            .arg(&scale);
32440        unsafe {
32441            b.launch(cfg)?;
32442        }
32443        Ok(())
32444    }
32445
32446    /// A4 seam: chunked WY GDN prefill. DEFAULT ON (`MEMRA_GDN_CHUNKED=0` = rollback to the
32447    /// sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL
32448    /// GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models
32449    /// on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text,
32450    /// 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125
32451    /// on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent).
32452    /// PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch
32453    /// identity law); prime_cache/forward/forward_last are the only callers.
32454    pub fn gdn_chunked_enabled() -> bool {
32455        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
32456        *E.get_or_init(|| {
32457            std::env::var("MEMRA_GDN_CHUNKED")
32458                .map(|v| v != "0")
32459                .unwrap_or(true)
32460        })
32461    }
32462
32463    /// A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk
32464    /// matrices grow with C while the sequential state pass is C-flat, so smaller chunks
32465    /// win; C=32/64 also get the register-history solve template). Clamped to multiples
32466    /// of 32 in [32, 128] (kernel row mappings require it).
32467    pub fn gdn_chunk_size() -> usize {
32468        static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
32469        *C.get_or_init(|| {
32470            let c: usize = std::env::var("MEMRA_GDN_CHUNK")
32471                .ok()
32472                .and_then(|v| v.parse().ok())
32473                .unwrap_or(32);
32474            c.clamp(32, 128) / 32 * 32
32475        })
32476    }
32477
32478    /// A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the
32479    /// math). Same contract as `gdn_scan_s128` (layouts, state ping-pong) but chunk-parallel:
32480    /// NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen
32481    /// argmax + run-spec batteries are the accuracy authority. PREFILL callers only.
32482    #[allow(clippy::too_many_arguments)]
32483    /// task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the
32484    /// batched-prime varlen path). Returns (gcum, P, U, W); `A` is K3-internal.
32485    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
32486    #[allow(clippy::too_many_arguments)]
32487    pub fn gdn_chunk_k123(
32488        &self,
32489        q: &CudaSlice<f32>,
32490        k: &CudaSlice<f32>,
32491        v: &CudaSlice<f32>,
32492        g: &CudaSlice<f32>,
32493        beta: &CudaSlice<f32>,
32494        wb16: Option<&mut CudaSlice<u8>>,
32495        n_head: usize,
32496        t: usize,
32497        c: usize,
32498        hk: usize,
32499        k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
32500    ) -> Result<
32501        (
32502            CudaSlice<f32>,
32503            CudaSlice<f32>,
32504            CudaSlice<f32>,
32505            CudaSlice<f32>,
32506        ),
32507        Box<dyn std::error::Error>,
32508    > {
32509        const D: usize = 128;
32510        let h = n_head;
32511        #[allow(clippy::manual_div_ceil)]
32512        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
32513        let nc = (t + c - 1) / c;
32514        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
32515        let mut gcum = self.uninit(t * h)?;
32516        let mut a = self.uninit(nc * h * c * c)?;
32517        let mut p = self.uninit(nc * h * c * c)?;
32518        let mut u = self.uninit(nc * h * c * D)?;
32519        let mut w = self.uninit(nc * h * c * D)?;
32520        {
32521            // K1
32522            let f = self.func("gdn_chunk_cumgate_f32");
32523            let cfg = LaunchConfig {
32524                grid_dim: (nc as u32, h as u32, 1),
32525                block_dim: (32, 1, 1),
32526                shared_mem_bytes: 0,
32527            };
32528            let __s_b = self.gpu.stream();
32529            let mut b = __s_b.launch_builder(&f);
32530            b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
32531            unsafe {
32532                b.launch(cfg)?;
32533            }
32534        }
32535        if let Some((qb, kb, pb)) = k2w {
32536            // K2-wgmma (MEMRA_GDN_WGMMA path, c==32): A + pre-masked Pb16 in one kernel;
32537            // the P f32 buffer stays UNWRITTEN (its only wgmma-path consumer is Pb16).
32538            assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
32539            let f = self.func("gdn_k2_wgmma");
32540            let cfg = LaunchConfig {
32541                grid_dim: (nc as u32, h as u32, 1),
32542                block_dim: (128, 1, 1),
32543                shared_mem_bytes: 0,
32544            };
32545            let hki = hk as i32;
32546            let __s_b = self.gpu.stream();
32547            let mut b = __s_b.launch_builder(&f);
32548            b.arg(qb)
32549                .arg(kb)
32550                .arg(&gcum)
32551                .arg(beta)
32552                .arg(&mut a)
32553                .arg(&mut *pb)
32554                .arg(&hi)
32555                .arg(&ti)
32556                .arg(&ci)
32557                .arg(&hki);
32558            unsafe {
32559                b.launch(cfg)?;
32560            }
32561        } else if c <= 64 && !portable_mma_gated() {
32562            // K2 register-tiled (2x2 outputs/thread, whole-chunk smem k tile)
32563            let f = self.func("gdn_chunk_attn_f32");
32564            f.set_attribute(
32565                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
32566                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
32567            )?;
32568            #[allow(clippy::manual_div_ceil)]
32569            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
32570            let jt = ((c + 31) / 32) as u32;
32571            let cfg = LaunchConfig {
32572                grid_dim: (nc as u32, h as u32, jt),
32573                block_dim: (256, 1, 1),
32574                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
32575            };
32576            let hki = hk as i32;
32577            let __s_b = self.gpu.stream();
32578            let mut b = __s_b.launch_builder(&f);
32579            b.arg(q)
32580                .arg(k)
32581                .arg(&gcum)
32582                .arg(beta)
32583                .arg(&mut a)
32584                .arg(&mut p)
32585                .arg(&hi)
32586                .arg(&ti)
32587                .arg(&ci)
32588                .arg(&hki);
32589            unsafe {
32590                b.launch(cfg)?;
32591            }
32592        } else {
32593            // K2 generic (C = 128, or the portable target's low-smem fallback)
32594            assert!(
32595                hk == h,
32596                "generic K2 is broadcast-only (de-broadcast rides C==32)"
32597            );
32598            let f = self.func("gdn_chunk_attn_g_f32");
32599            let cfg = LaunchConfig {
32600                grid_dim: (nc as u32, h as u32, 1),
32601                block_dim: (32, 8, 1),
32602                shared_mem_bytes: 0,
32603            };
32604            let __s_b = self.gpu.stream();
32605            let mut b = __s_b.launch_builder(&f);
32606            b.arg(q)
32607                .arg(k)
32608                .arg(&gcum)
32609                .arg(beta)
32610                .arg(&mut a)
32611                .arg(&mut p)
32612                .arg(&hi)
32613                .arg(&ti)
32614                .arg(&ci);
32615            unsafe {
32616                b.launch(cfg)?;
32617            }
32618        }
32619        {
32620            // K3 (register-history templates for C=32/64; local-memory generic otherwise)
32621            let cfg = LaunchConfig {
32622                grid_dim: (nc as u32, h as u32, 1),
32623                block_dim: (256, 1, 1),
32624                shared_mem_bytes: 0,
32625            };
32626            match c {
32627                32 | 64 => {
32628                    let f = self.func(if c == 32 {
32629                        "gdn_chunk_solve32_f32"
32630                    } else {
32631                        "gdn_chunk_solve64_f32"
32632                    });
32633                    // mirror-fold: W's bf16 twin emitted on store (0 = skip)
32634                    let wb: u64 = match wb16 {
32635                        Some(d) => self.addr_u8(d),
32636                        None => 0,
32637                    };
32638                    let hki = hk as i32;
32639                    let __s_b = self.gpu.stream();
32640                    let mut b = __s_b.launch_builder(&f);
32641                    b.arg(v)
32642                        .arg(k)
32643                        .arg(&a)
32644                        .arg(&gcum)
32645                        .arg(&mut u)
32646                        .arg(&mut w)
32647                        .arg(&wb)
32648                        .arg(&hi)
32649                        .arg(&ti)
32650                        .arg(&hki);
32651                    unsafe {
32652                        b.launch(cfg)?;
32653                    }
32654                }
32655                _ => {
32656                    assert!(hk == h, "generic K3 is broadcast-only");
32657                    let f = self.func("gdn_chunk_solve_f32");
32658                    let __s_b = self.gpu.stream();
32659                    let mut b = __s_b.launch_builder(&f);
32660                    b.arg(v)
32661                        .arg(k)
32662                        .arg(&a)
32663                        .arg(&gcum)
32664                        .arg(&mut u)
32665                        .arg(&mut w)
32666                        .arg(&hi)
32667                        .arg(&ti)
32668                        .arg(&ci);
32669                    unsafe {
32670                        b.launch(cfg)?;
32671                    }
32672                }
32673            }
32674        }
32675        Ok((gcum, p, u, w))
32676    }
32677
32678    /// task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of
32679    /// the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path
32680    /// consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).
32681    pub fn gdn_db_on() -> bool {
32682        std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
32683    }
32684
32685    /// Whether the K4/K5 mma pair serves at chunk size `c` (mirrors gdn_scan_chunked's
32686    /// seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).
32687    /// DEFAULT ON for sm_120a builds too (lane/moeprime-nvfp4-direct, 2026-08-21): the pair
32688    /// was qualified on 90a only and left env-opt-in elsewhere; measured on Blackwell it
32689    /// wins on BOTH rigs — one RTX PRO 6000 (ornith15 pp14715 12,036 -> 12,751/12,957,
32690    /// +6-8%, both orders) and the local 5090 (q38-27b pp6435 1,397/1,429 -> 1,427/1,446,
32691    /// both orders) — with kernel-check/run-gen/margin-gate/run-spec green under the flag.
32692    /// bf16 HMMA (m16n8k16) is sm_80-class PTX; only the wgmma nest stays Hopper-gated.
32693    pub fn gdn_mma_enabled(&self, c: usize) -> bool {
32694        !portable_mma_gated()
32695            && c == 32
32696            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
32697                Ok("1") => true,
32698                Ok("0") => false,
32699                _ => gdn_mma_default_on(),
32700            }
32701    }
32702
32703    /// task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the
32704    /// mma config; same per-call env read discipline). HARD-gated to the Hopper build:
32705    /// the wgmma asm bodies exist only at __CUDA_ARCH__ == 900 (MEMRA_K45_REAL,
32706    /// wgmma_common.cuh) — on every other arch the kernel compiles EMPTY, so an env
32707    /// force would silently produce garbage. Required since the sm_120a mma default
32708    /// flip made MEMRA_GDN_WGMMA=1 alone reach this branch there.
32709    pub fn gdn_wgmma_on(&self, c: usize) -> bool {
32710        cfg!(memra_hopper_mma)
32711            && self.gdn_mma_enabled(c)
32712            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
32713    }
32714
32715    /// task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the
32716    /// conv_out intermediate and its transposed re-read disappear — 11.8ms of the
32717    /// T=2048 prime). Ring update stays the separate follow-up launch (pad-aware).
32718    /// BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.
32719    #[allow(clippy::too_many_arguments)]
32720    #[allow(clippy::manual_div_ceil)] // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
32721    pub fn ssm_conv1d_gdn_state_pad(
32722        &self,
32723        qkv_tm: &cudarc::driver::CudaView<f32>,
32724        conv_state: &mut CudaSlice<f32>,
32725        w: &CudaSlice<f32>,
32726        q_g: &mut CudaSlice<f32>,
32727        k_g: &mut CudaSlice<f32>,
32728        v_g: &mut CudaSlice<f32>,
32729        conv_dim: usize,
32730        t: usize,
32731        d_conv: usize,
32732        d_state: usize,
32733        num_v: usize,
32734        num_k: usize,
32735        key_dim: usize,
32736        hk: usize,
32737        pad_len: Option<&CudaSlice<i32>>,
32738    ) -> Result<(), Box<dyn std::error::Error>> {
32739        assert!(
32740            t >= d_conv - 1,
32741            "fused state conv requires T >= pad (PRIME_MIN_T gates)"
32742        );
32743        {
32744            let f = self.func("ssm_conv1d_gdn_state_f32");
32745            let cfg = LaunchConfig {
32746                grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
32747                block_dim: (256, 1, 1),
32748                shared_mem_bytes: 0,
32749            };
32750            let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
32751            let (ds, nv, nk, kd, hki) = (
32752                d_state as i32,
32753                num_v as i32,
32754                num_k as i32,
32755                key_dim as i32,
32756                hk as i32,
32757            );
32758            let __s_b = self.gpu.stream();
32759            let mut b = __s_b.launch_builder(&f);
32760            b.arg(qkv_tm)
32761                .arg(&*conv_state)
32762                .arg(w)
32763                .arg(q_g)
32764                .arg(k_g)
32765                .arg(v_g)
32766                .arg(&cd)
32767                .arg(&ti)
32768                .arg(&dc)
32769                .arg(&ds)
32770                .arg(&nv)
32771                .arg(&nk)
32772                .arg(&kd)
32773                .arg(&hki);
32774            unsafe {
32775                b.launch(cfg)?;
32776            }
32777        }
32778        match pad_len {
32779            Some(len_d) => {
32780                let f = self.func("ssm_conv_ring_update_dev_f32");
32781                let n = conv_dim * (d_conv - 1);
32782                let cfg = LaunchConfig::for_num_elems(n as u32);
32783                let (cd, dc) = (conv_dim as i32, d_conv as i32);
32784                let __s_b = self.gpu.stream();
32785                let mut b = __s_b.launch_builder(&f);
32786                b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
32787                unsafe {
32788                    b.launch(cfg)?;
32789                }
32790            }
32791            None => {
32792                let f = self.func("ssm_conv_ring_update_f32");
32793                let n = conv_dim * (d_conv - 1);
32794                let cfg = LaunchConfig::for_num_elems(n as u32);
32795                let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
32796                let __s_b = self.gpu.stream();
32797                let mut b = __s_b.launch_builder(&f);
32798                b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
32799                unsafe {
32800                    b.launch(cfg)?;
32801                }
32802            }
32803        }
32804        Ok(())
32805    }
32806
32807    /// task #18 increment 2: allocate ONE sequence's chunk buffers (no launches) —
32808    /// K1-K5 all run varlen afterwards. `a`/`w` become struct members so the varlen
32809    /// K2/K3 can write them.
32810    pub fn gdn_chunk_alloc(
32811        &self,
32812        n_head: usize,
32813        t: usize,
32814        c: usize,
32815        hk: usize,
32816    ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
32817        const D: usize = 128;
32818        assert!(
32819            c == 32,
32820            "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
32821        );
32822        let h = n_head;
32823        #[allow(clippy::manual_div_ceil)]
32824        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
32825        let nc = (t + c - 1) / c;
32826        Ok(GdnChunkBufs {
32827            gcum: self.uninit(t * h)?,
32828            a: self.uninit(nc * h * c * c)?,
32829            p: self.uninit(nc * h * c * c)?,
32830            u: self.uninit(nc * h * c * D)?,
32831            w: self.uninit(nc * h * c * D)?,
32832            kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
32833            wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
32834            y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
32835            ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
32836            qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
32837            pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
32838            o: self.uninit(D * h * t)?,
32839            t,
32840            nc,
32841        })
32842    }
32843
32844    /// view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).
32845    pub fn f32_to_bf16_v(
32846        &self,
32847        x: &cudarc::driver::CudaView<f32>,
32848        dst: &mut CudaSlice<u8>,
32849        n: usize,
32850    ) -> Result<(), Box<dyn std::error::Error>> {
32851        let f = self.func("f32_to_bf16_bulk");
32852        let ni = n as i64;
32853        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
32854        let __s_b = self.gpu.stream();
32855        let mut b = __s_b.launch_builder(&f);
32856        b.arg(x).arg(dst).arg(&ni);
32857        unsafe {
32858            b.launch(cfg)?;
32859        }
32860        Ok(())
32861    }
32862
32863    /// f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).
32864    pub fn f32_to_bf16_into(
32865        &self,
32866        x: &CudaSlice<f32>,
32867        dst: &mut CudaSlice<u8>,
32868        n: usize,
32869    ) -> Result<(), Box<dyn std::error::Error>> {
32870        let f = self.func("f32_to_bf16_bulk");
32871        let ni = n as i64;
32872        let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
32873        let __s_b = self.gpu.stream();
32874        let mut b = __s_b.launch_builder(&f);
32875        b.arg(x).arg(dst).arg(&ni);
32876        unsafe {
32877            b.launch(cfg)?;
32878        }
32879        Ok(())
32880    }
32881
32882    /// task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence's
32883    /// cumgate/attn/solve (per-block math identical to the per-seq kernels).
32884    pub fn gdn_chunk_k123_vl8(
32885        &self,
32886        seqs: &[GdnSeqVl],
32887        n_head: usize,
32888        hk: usize,
32889        wq: Option<&GdnWVl8>,
32890    ) -> Result<(), Box<dyn std::error::Error>> {
32891        let b = seqs.len();
32892        assert!((1..=8).contains(&b), "gdn_chunk_k123_vl8: 1..=8 sequences");
32893        let mut packed = [GdnSeqVl::default(); 8];
32894        packed[..b].copy_from_slice(seqs);
32895        let v = GdnVl8(packed);
32896        let (hi, ci) = (n_head as i32, 32i32);
32897        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
32898        {
32899            let f = self.func("gdn_chunk_cumgate_vl");
32900            let cfg = LaunchConfig {
32901                grid_dim: (max_nc, n_head as u32, b as u32),
32902                block_dim: (32, 1, 1),
32903                shared_mem_bytes: 0,
32904            };
32905            let __s_lb = self.gpu.stream();
32906            let mut lb = __s_lb.launch_builder(&f);
32907            lb.arg(&v).arg(&hi).arg(&ci);
32908            unsafe {
32909                lb.launch(cfg)?;
32910            }
32911        }
32912        let hki = hk as i32;
32913        if let Some(w) = wq {
32914            // K2-wgmma vl twin (writes A + pre-masked Pb16)
32915            let f = self.func("gdn_k2_wgmma_vl");
32916            let cfg = LaunchConfig {
32917                grid_dim: (max_nc, n_head as u32, b as u32),
32918                block_dim: (128, 1, 1),
32919                shared_mem_bytes: 0,
32920            };
32921            let __s_lb = self.gpu.stream();
32922            let mut lb = __s_lb.launch_builder(&f);
32923            lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
32924            unsafe {
32925                lb.launch(cfg)?;
32926            }
32927        } else {
32928            let f = self.func("gdn_chunk_attn_vl");
32929            f.set_attribute(
32930                CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
32931                GDN_K2_DYNAMIC_SHARED_BYTES as i32,
32932            )?;
32933            let cfg = LaunchConfig {
32934                grid_dim: (max_nc, n_head as u32, b as u32),
32935                block_dim: (256, 1, 1),
32936                shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
32937            };
32938            let __s_lb = self.gpu.stream();
32939            let mut lb = __s_lb.launch_builder(&f);
32940            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
32941            unsafe {
32942                lb.launch(cfg)?;
32943            }
32944        }
32945        {
32946            let f = self.func("gdn_chunk_solve32_vl");
32947            let cfg = LaunchConfig {
32948                grid_dim: (max_nc, n_head as u32, b as u32),
32949                block_dim: (256, 1, 1),
32950                shared_mem_bytes: 0,
32951            };
32952            let __s_lb = self.gpu.stream();
32953            let mut lb = __s_lb.launch_builder(&f);
32954            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
32955            unsafe {
32956                lb.launch(cfg)?;
32957            }
32958        }
32959        Ok(())
32960    }
32961
32962    /// task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 /
32963    /// fused gate-prep, 5 launches for every sequence (per-element math identical
32964    /// to the per-seq kernels; l2/gate fusions write disjoint outputs).
32965    #[allow(clippy::too_many_arguments)]
32966    pub fn gdn_prep_vl8(
32967        &self,
32968        seqs: &[GdnPrepVl],
32969        conv_w: &CudaSlice<f32>,
32970        dt_bias: &CudaSlice<f32>,
32971        a: &CudaSlice<f32>,
32972        conv_dim: usize,
32973        d_conv: usize,
32974        d_state: usize,
32975        num_v: usize,
32976        num_k: usize,
32977        key_dim: usize,
32978        hk: usize,
32979        eps: f32,
32980    ) -> Result<(), Box<dyn std::error::Error>> {
32981        let b = seqs.len();
32982        assert!((1..=8).contains(&b));
32983        let mut packed = [GdnPrepVl::default(); 8];
32984        packed[..b].copy_from_slice(seqs);
32985        let v = GdnPrepVl8(packed);
32986        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
32987        let (cdi, dci) = (conv_dim as i32, d_conv as i32);
32988        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
32989        assert!(
32990            conv_fuse || hk == num_v,
32991            "de-broadcast requires the fused conv"
32992        );
32993        if conv_fuse {
32994            let f = self.func("ssm_conv1d_gdn_state_vl");
32995            let cfg = LaunchConfig {
32996                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
32997                block_dim: (256, 1, 1),
32998                shared_mem_bytes: 0,
32999            };
33000            let (dsi, nvi, nki, kdi, hki) = (
33001                d_state as i32,
33002                num_v as i32,
33003                num_k as i32,
33004                key_dim as i32,
33005                hk as i32,
33006            );
33007            let __s_lb = self.gpu.stream();
33008            let mut lb = __s_lb.launch_builder(&f);
33009            lb.arg(&v)
33010                .arg(conv_w)
33011                .arg(&cdi)
33012                .arg(&dci)
33013                .arg(&dsi)
33014                .arg(&nvi)
33015                .arg(&nki)
33016                .arg(&kdi)
33017                .arg(&hki);
33018            unsafe {
33019                lb.launch(cfg)?;
33020            }
33021        } else {
33022            let f = self.func("ssm_conv1d_tm_state_vl");
33023            let cfg = LaunchConfig {
33024                grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
33025                block_dim: (256, 1, 1),
33026                shared_mem_bytes: 0,
33027            };
33028            let __s_lb = self.gpu.stream();
33029            let mut lb = __s_lb.launch_builder(&f);
33030            lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
33031            unsafe {
33032                lb.launch(cfg)?;
33033            }
33034        }
33035        {
33036            let f = self.func("ssm_conv_ring_update_vl");
33037            let n = (conv_dim * (d_conv - 1)) as u32;
33038            let cfg = LaunchConfig {
33039                grid_dim: (n.div_ceil(256), 1, b as u32),
33040                block_dim: (256, 1, 1),
33041                shared_mem_bytes: 0,
33042            };
33043            let __s_lb = self.gpu.stream();
33044            let mut lb = __s_lb.launch_builder(&f);
33045            lb.arg(&v).arg(&cdi).arg(&dci);
33046            unsafe {
33047                lb.launch(cfg)?;
33048            }
33049        }
33050        if !conv_fuse {
33051            let f = self.func("qkv_to_gdn_repack_vl");
33052            let n = max_t * (num_v * d_state) as u32;
33053            let cfg = LaunchConfig {
33054                grid_dim: (n.div_ceil(256), 1, b as u32),
33055                block_dim: (256, 1, 1),
33056                shared_mem_bytes: 0,
33057            };
33058            let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
33059            let __s_lb = self.gpu.stream();
33060            let mut lb = __s_lb.launch_builder(&f);
33061            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
33062            unsafe {
33063                lb.launch(cfg)?;
33064            }
33065        }
33066        if Self::l2_v2_on(d_state) {
33067            let f = self.func("gdn_l2_v2_vl");
33068            let cfg = LaunchConfig {
33069                grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
33070                block_dim: (256, 1, 1),
33071                shared_mem_bytes: 0,
33072            };
33073            let (dsi, nvi) = (d_state as i32, hk as i32);
33074            let __s_lb = self.gpu.stream();
33075            let mut lb = __s_lb.launch_builder(&f);
33076            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
33077            unsafe {
33078                lb.launch(cfg)?;
33079            }
33080        } else {
33081            let f = self.func("gdn_l2_vl");
33082            let cfg = LaunchConfig {
33083                grid_dim: (max_t * hk as u32, 2, b as u32),
33084                block_dim: (256, 1, 1),
33085                shared_mem_bytes: 0,
33086            };
33087            let (dsi, nvi) = (d_state as i32, hk as i32);
33088            let __s_lb = self.gpu.stream();
33089            let mut lb = __s_lb.launch_builder(&f);
33090            lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
33091            unsafe {
33092                lb.launch(cfg)?;
33093            }
33094        }
33095        {
33096            let f = self.func("gdn_gate_prep_vl");
33097            let n = max_t * num_v as u32;
33098            let cfg = LaunchConfig {
33099                grid_dim: (n.div_ceil(256), 1, b as u32),
33100                block_dim: (256, 1, 1),
33101                shared_mem_bytes: 0,
33102            };
33103            let nvi = num_v as i32;
33104            let __s_lb = self.gpu.stream();
33105            let mut lb = __s_lb.launch_builder(&f);
33106            lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
33107            unsafe {
33108                lb.launch(cfg)?;
33109            }
33110        }
33111        Ok(())
33112    }
33113
33114    /// varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).
33115    pub fn gdn_mirror_vl8(
33116        &self,
33117        seqs: &[GdnSeqVl],
33118        n_head: usize,
33119        which: i32,
33120        hk: usize,
33121    ) -> Result<(), Box<dyn std::error::Error>> {
33122        let b = seqs.len();
33123        assert!((1..=8).contains(&b));
33124        let mut packed = [GdnSeqVl::default(); 8];
33125        packed[..b].copy_from_slice(seqs);
33126        let v = GdnVl8(packed);
33127        let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
33128        let max_n = seqs
33129            .iter()
33130            .map(|s| {
33131                if which == 0 {
33132                    s.t as i64 * ept as i64
33133                } else {
33134                    s.nc as i64 * ept as i64 * 32
33135                }
33136            })
33137            .max()
33138            .unwrap();
33139        let f = self.func("gdn_mirror_vl");
33140        let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
33141        let cfg = LaunchConfig {
33142            grid_dim: (blocks, 1, b as u32),
33143            block_dim: (256, 1, 1),
33144            shared_mem_bytes: 0,
33145        };
33146        let __s_lb = self.gpu.stream();
33147        let mut lb = __s_lb.launch_builder(&f);
33148        lb.arg(&v).arg(&ept).arg(&which);
33149        unsafe {
33150            lb.launch(cfg)?;
33151        }
33152        Ok(())
33153    }
33154
33155    /// varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.
33156    pub fn gdn_tail_vl8(
33157        &self,
33158        seqs: &[GdnPrepVl],
33159        norm_w: &CudaSlice<f32>,
33160        d_state: usize,
33161        num_v: usize,
33162        eps: f32,
33163    ) -> Result<(), Box<dyn std::error::Error>> {
33164        let b = seqs.len();
33165        assert!((1..=8).contains(&b));
33166        let mut packed = [GdnPrepVl::default(); 8];
33167        packed[..b].copy_from_slice(seqs);
33168        let v = GdnPrepVl8(packed);
33169        let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
33170        let f = self.func("gated_rmsnorm_f16out_vl");
33171        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
33172        let cfg = LaunchConfig {
33173            grid_dim: (max_t * num_v as u32, 1, b as u32),
33174            block_dim: (128, 1, 1),
33175            shared_mem_bytes: 0,
33176        };
33177        let (dsi, nvi) = (d_state as i32, num_v as i32);
33178        let __s_lb = self.gpu.stream();
33179        let mut lb = __s_lb.launch_builder(&f);
33180        lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
33181        unsafe {
33182            lb.launch(cfg)?;
33183        }
33184        Ok(())
33185    }
33186
33187    /// Raw device address helpers for the varlen by-value arg struct (single-stream
33188    /// launches; every buffer outlives the call — the f16 FFI discipline).
33189    pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
33190        use cudarc::driver::DevicePtr;
33191        let s = self.gpu.stream();
33192        let (p, _g) = x.device_ptr(&s);
33193        p
33194    }
33195    pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
33196        use cudarc::driver::DevicePtrMut;
33197        let s = self.gpu.stream();
33198        let (p, _g) = x.device_ptr_mut(&s);
33199        p
33200    }
33201    pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
33202        use cudarc::driver::DevicePtr;
33203        let s = self.gpu.stream();
33204        let (p, _g) = x.device_ptr(&s);
33205        p
33206    }
33207    pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
33208        use cudarc::driver::DevicePtr;
33209        let s = self.gpu.stream();
33210        let (p, _g) = x.device_ptr(&s);
33211        p
33212    }
33213
33214    /// task #18: the varlen K4+K5 pair — TWO launches run every sequence's state pass
33215    /// and output pass (grid gains a seq dim; per-block math identical to the per-seq
33216    /// launches, so this is strictly bit-gateable against them).
33217    pub fn gdn_chunk_vl8(
33218        &self,
33219        seqs: &[GdnSeqVl],
33220        n_head: usize,
33221        scale: f32,
33222        hk: usize,
33223        wq: Option<&GdnWVl8>,
33224    ) -> Result<(), Box<dyn std::error::Error>> {
33225        const NSPLIT: u32 = 4;
33226        let b = seqs.len();
33227        assert!((1..=8).contains(&b), "gdn_chunk_vl8: 1..=8 sequences");
33228        let mut packed = [GdnSeqVl::default(); 8];
33229        packed[..b].copy_from_slice(seqs);
33230        let v = GdnVl8(packed);
33231        let (hi, ci) = (n_head as i32, 32i32);
33232        let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
33233        let hki = hk as i32;
33234        if let Some(w) = wq {
33235            // K4+K5 fused wgmma vl twin: one launch, Y/Ssnap never materialized.
33236            let f = self.func("gdn_k45_wgmma_vl");
33237            let cfg = LaunchConfig {
33238                grid_dim: (n_head as u32, NSPLIT, b as u32),
33239                block_dim: (256, 1, 1),
33240                shared_mem_bytes: 0,
33241            };
33242            let __s_lb = self.gpu.stream();
33243            let mut lb = __s_lb.launch_builder(&f);
33244            lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
33245            unsafe {
33246                lb.launch(cfg)?;
33247            }
33248            let _ = max_nc;
33249            return Ok(());
33250        }
33251        {
33252            let f = self.func("gdn_chunk_state_mma_vl");
33253            let cfg = LaunchConfig {
33254                grid_dim: (n_head as u32, NSPLIT, b as u32),
33255                block_dim: (256, 1, 1),
33256                shared_mem_bytes: 0,
33257            };
33258            let __s_lb = self.gpu.stream();
33259            let mut lb = __s_lb.launch_builder(&f);
33260            lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
33261            unsafe {
33262                lb.launch(cfg)?;
33263            }
33264        }
33265        {
33266            let f = self.func("gdn_chunk_output_mma_vl");
33267            let cfg = LaunchConfig {
33268                grid_dim: (max_nc, n_head as u32, b as u32),
33269                block_dim: (256, 1, 1),
33270                shared_mem_bytes: 0,
33271            };
33272            let __s_lb = self.gpu.stream();
33273            let mut lb = __s_lb.launch_builder(&f);
33274            lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
33275            unsafe {
33276                lb.launch(cfg)?;
33277            }
33278        }
33279        Ok(())
33280    }
33281    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
33282    pub fn gdn_scan_chunked(
33283        &self,
33284        q: &CudaSlice<f32>,
33285        k: &CudaSlice<f32>,
33286        v: &CudaSlice<f32>,
33287        g: &CudaSlice<f32>,
33288        beta: &CudaSlice<f32>,
33289        kb16_pre: Option<&CudaSlice<u8>>,
33290        qb16_pre: Option<&CudaSlice<u8>>,
33291        state_in: &CudaSlice<f32>,
33292        state_out: &mut CudaSlice<f32>,
33293        o: &mut CudaSlice<f32>,
33294        n_head: usize,
33295        t: usize,
33296        scale: f32,
33297        c: usize,
33298        hk: usize,
33299    ) -> Result<(), Box<dyn std::error::Error>> {
33300        const D: usize = 128;
33301        const NSPLIT: u32 = 4;
33302        assert!(
33303            (1..=128).contains(&c),
33304            "gdn_scan_chunked: C must be in 1..=128"
33305        );
33306        let h = n_head;
33307        #[allow(clippy::manual_div_ceil)]
33308        // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
33309        let nc = (t + c - 1) / c;
33310        let (hi, ti, ci) = (h as i32, t as i32, c as i32);
33311        // mirror-fold (round 27): on the mma path W's bf16 twin is emitted by K3's store
33312        // (wb16 pre-allocated and threaded through k123) and k's by the producer l2 when
33313        // the caller hands `kb16_pre` — both standalone mirror passes disappear.
33314        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON, lane/moeprime-nvfp4-direct)
33315        let gdn_mma_pre = !portable_mma_gated()
33316            && c == 32
33317            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
33318                Ok("1") => true,
33319                Ok("0") => false,
33320                _ => gdn_mma_default_on(),
33321            };
33322        let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
33323            Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
33324        } else {
33325            None
33326        };
33327        // K2-wgmma pre-work (MEMRA_GDN_WGMMA): the kb16/qb16 mirrors hoist ABOVE K123 so
33328        // K2 rides them via cp.async; K2 writes the pre-masked Pb16 directly (the
33329        // gdn_p_bf16_masked pass and the in-branch mirror builds disappear).
33330        // mirrors gdn_wgmma_on: hard Hopper-build gate (empty wgmma bodies elsewhere)
33331        let gdn_wgmma_pre = cfg!(memra_hopper_mma)
33332            && gdn_mma_pre
33333            && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
33334        let nk = t * hk * D;
33335        let mut kb16_local: Option<CudaSlice<u8>> = None;
33336        if gdn_mma_pre && kb16_pre.is_none() {
33337            let mut kb = self.alloc_u8_uninit(nk * 2)?;
33338            let f = self.func("f32_to_bf16_bulk");
33339            let n2 = nk as i64;
33340            let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
33341            let __s_b = self.gpu.stream();
33342            let mut b = __s_b.launch_builder(&f);
33343            b.arg(k).arg(&mut kb).arg(&n2);
33344            unsafe {
33345                b.launch(cfg2)?;
33346            }
33347            kb16_local = Some(kb);
33348        }
33349        let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
33350        if let Some(kb) = kb16_pre {
33351            assert!(kb.len() >= nk * 2, "kb16_pre too small");
33352        }
33353        let mut qb16: Option<CudaSlice<u8>> = None;
33354        let mut pb16: Option<CudaSlice<u8>> = None;
33355        if gdn_wgmma_pre {
33356            // mirror-fold (round 35): prep's l2 v2 emits qb16 in-epilogue (kb16 pattern);
33357            // the standalone bulk cvt only serves callers without the prep mirror.
33358            if qb16_pre.is_none() {
33359                let mut qb = self.alloc_u8_uninit(nk * 2)?;
33360                let f = self.func("f32_to_bf16_bulk");
33361                let n2 = nk as i64;
33362                let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
33363                let __s_b = self.gpu.stream();
33364                let mut b = __s_b.launch_builder(&f);
33365                b.arg(q).arg(&mut qb).arg(&n2);
33366                unsafe {
33367                    b.launch(cfg2)?;
33368                }
33369                qb16 = Some(qb);
33370            } else if let Some(qb) = qb16_pre {
33371                assert!(qb.len() >= nk * 2, "qb16_pre too small");
33372            }
33373            pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
33374        }
33375        let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
33376        let k2w = if gdn_wgmma_pre {
33377            Some((
33378                *qb16_ref0.as_ref().unwrap(),
33379                *kb16_ref0.as_ref().unwrap(),
33380                pb16.as_mut().unwrap(),
33381            ))
33382        } else {
33383            None
33384        };
33385        let (gcum, p, u, w) =
33386            self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
33387        let _ = &w;
33388        let mut y = self.uninit(nc * h * c * D)?;
33389        let mut ssnap = self.uninit(nc * h * D * D)?; // chunk-start state snapshots (K5 phase 1)
33390        // K4-MMA seam (MEMRA_GDN_MMA; harness verdict 1.75x — tools/bench_gdn_k4.cu, ledger
33391        // 2026-07-26): M in mma accumulator fragments, bf16 W/k mirrors through a cp.async
33392        // ring. C==32 only (the kernel's tile). PROMOTED default-ON on the Hopper lane
33393        // after the STATE-CARRY battery (2026-07-26): 2048-token prime (64 in-kernel state
33394        // carries) -> 256 greedy decode tokens IDENTICAL to f32 on 3 seeds, AND chunked-
33395        // continuation prime (MEMRA_PRIME_CHUNK=512, 4 cross-call carries via cache.recur)
33396        // IDENTICAL on 2 seeds; plus argmax MATCH, pp512 +3.5% (17286), oracle out
33397        // mean_rel ~1e-4. kernel-check pins BOTH configs (f32 tight band forced =0; mma
33398        // band 8e-2/8e-1 vs f64 truth). =0 reverts; portable stays f32. NOT read via
33399        // OnceLock ON PURPOSE: kernel-check toggles the env per call to pin both forms.
33400        // default mirrors gdn_mma_enabled (incl. the sm_120a-build ON) — all three read
33401        // sites must agree or the pre-work arms while the scan takes the scalar route.
33402        let gdn_mma = !portable_mma_gated()
33403            && c == 32
33404            && match std::env::var("MEMRA_GDN_MMA").as_deref() {
33405                Ok("1") => true,
33406                Ok("0") => false,
33407                _ => gdn_mma_default_on(),
33408            };
33409        if gdn_mma {
33410            let wb16 = wb16_pre
33411                .take()
33412                .expect("mma path pre-allocates wb16 (K3 store fold)");
33413            let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
33414            // K4+K5 FUSED wgmma seam (MEMRA_GDN_WGMMA, task #22; harness verdict
33415            // tools/bench_gdn_wgmma.cu v5, ledger 1f08b997: in-band Y 1.07e-2 / state
33416            // 1.03e-2 / O 1.08e-2, 91.3us vs 70.4 K4-only at H=32 T=512). K5's output
33417            // pass runs inside the persistent-M kernel; Y and Ssnap are never
33418            // materialized. New numeric class (gk folds into k^T instead of ys) —
33419            // explicit opt-in until the state-carry battery promotes it. Env read per
33420            // call (kernel-check pins configs by toggling env, GDN_MMA precedent).
33421            // PROMOTED default-ON hopper (2026-07-27): full battery green — harness
33422            // in-band, argmax gate PASS, 3-seed greedy IDENTICAL after ~2k prime,
33423            // chunked-continuation IDENTICAL, kernel-check + decode-batch gates green,
33424            // official prefill lane +0.74% interleaved x5 (5/5 rounds). =0 reverts.
33425            if gdn_wgmma_pre {
33426                // qb16/pb16 pre-built above K123 (K2-wgmma wrote the masked Pb16).
33427                let qb16 = qb16_ref0.unwrap();
33428                let pb16 = pb16.as_ref().unwrap();
33429                {
33430                    let f = self.func("gdn_k45_wgmma");
33431                    let cfg = LaunchConfig {
33432                        grid_dim: (h as u32, 4, 1),
33433                        block_dim: (256, 1, 1),
33434                        shared_mem_bytes: 0,
33435                    };
33436                    let hki = hk as i32;
33437                    let __s_b = self.gpu.stream();
33438                    let mut b = __s_b.launch_builder(&f);
33439                    b.arg(kb16_ref)
33440                        .arg(&gcum)
33441                        .arg(beta)
33442                        .arg(&u)
33443                        .arg(&wb16)
33444                        .arg(qb16)
33445                        .arg(pb16)
33446                        .arg(o)
33447                        .arg(&scale)
33448                        .arg(state_in)
33449                        .arg(&mut *state_out)
33450                        .arg(&hi)
33451                        .arg(&ti)
33452                        .arg(&ci)
33453                        .arg(&hki);
33454                    unsafe {
33455                        b.launch(cfg)?;
33456                    }
33457                }
33458                return Ok(());
33459            }
33460            // COUPLED PAIR: K4-mma writes Y and Ssnap as bf16 (their only consumer is
33461            // K5-mma, which rounds to bf16 regardless — identical numerics, half the
33462            // traffic; harness K5 63.0 -> 35.3us). Fresh bf16 buffers replace the f32 ones.
33463            let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
33464            let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
33465            {
33466                let f = self.func("gdn_chunk_state_mma");
33467                let cfg = LaunchConfig {
33468                    grid_dim: (h as u32, NSPLIT, 1),
33469                    block_dim: (256, 1, 1),
33470                    shared_mem_bytes: 0,
33471                };
33472                let hki = hk as i32;
33473                let __s_b = self.gpu.stream();
33474                let mut b = __s_b.launch_builder(&f);
33475                b.arg(kb16_ref)
33476                    .arg(&gcum)
33477                    .arg(beta)
33478                    .arg(&u)
33479                    .arg(&wb16)
33480                    .arg(&mut y16)
33481                    .arg(&mut ssnap16)
33482                    .arg(state_in)
33483                    .arg(&mut *state_out)
33484                    .arg(&hi)
33485                    .arg(&ti)
33486                    .arg(&ci)
33487                    .arg(&hki);
33488                unsafe {
33489                    b.launch(cfg)?;
33490                }
33491            }
33492            {
33493                // K5-mma (bf16 St/Y consumers)
33494                let f = self.func("gdn_chunk_output_mma");
33495                #[allow(clippy::manual_div_ceil)]
33496                // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
33497                let jt = ((c + 31) / 32) as u32;
33498                let cfg = LaunchConfig {
33499                    grid_dim: (nc as u32, h as u32, jt),
33500                    block_dim: (256, 1, 1),
33501                    shared_mem_bytes: 0,
33502                };
33503                let hki = hk as i32;
33504                let __s_b = self.gpu.stream();
33505                let mut b = __s_b.launch_builder(&f);
33506                b.arg(q)
33507                    .arg(&gcum)
33508                    .arg(&p)
33509                    .arg(&y16)
33510                    .arg(&ssnap16)
33511                    .arg(o)
33512                    .arg(&hi)
33513                    .arg(&ti)
33514                    .arg(&ci)
33515                    .arg(&scale)
33516                    .arg(&hki);
33517                unsafe {
33518                    b.launch(cfg)?;
33519                }
33520            }
33521            return Ok(());
33522        }
33523        {
33524            // K4 (sequential over chunks inside; blocks col-partition the state)
33525            let f = self.func("gdn_chunk_state_f32");
33526            let cfg = LaunchConfig {
33527                grid_dim: (h as u32, NSPLIT, 1),
33528                block_dim: (256, 1, 1),
33529                shared_mem_bytes: 0,
33530            };
33531            let __s_b = self.gpu.stream();
33532            let mut b = __s_b.launch_builder(&f);
33533            b.arg(k)
33534                .arg(&gcum)
33535                .arg(beta)
33536                .arg(&u)
33537                .arg(&w)
33538                .arg(&mut y)
33539                .arg(&mut ssnap)
33540                .arg(state_in)
33541                .arg(&mut *state_out)
33542                .arg(&hi)
33543                .arg(&ti)
33544                .arg(&ci);
33545            unsafe {
33546                b.launch(cfg)?;
33547            }
33548        }
33549        {
33550            // K5 (j-blocked: grid.z = 32-row output blocks per chunk; writes o fully)
33551            let f = self.func("gdn_chunk_output_f32");
33552            #[allow(clippy::manual_div_ceil)]
33553            // allow: explicit (n + k - 1) / k is the load-bearing sizing form, kept textually identical to the kernel-side math
33554            let jt = ((c + 31) / 32) as u32;
33555            let cfg = LaunchConfig {
33556                grid_dim: (nc as u32, h as u32, jt),
33557                block_dim: (256, 1, 1),
33558                shared_mem_bytes: 0,
33559            };
33560            let __s_b = self.gpu.stream();
33561            let mut b = __s_b.launch_builder(&f);
33562            b.arg(q)
33563                .arg(&gcum)
33564                .arg(&p)
33565                .arg(&y)
33566                .arg(&ssnap)
33567                .arg(o)
33568                .arg(&hi)
33569                .arg(&ti)
33570                .arg(&ci)
33571                .arg(&scale);
33572            unsafe {
33573                b.launch(cfg)?;
33574            }
33575        }
33576        Ok(())
33577    }
33578
33579    /// PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the
33580    /// batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn
33581    /// (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec
33582    /// verify call `gdn_scan_s128` DIRECTLY — the decode==verify dispatch identity is untouched.
33583    ///
33584    /// MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the
33585    /// per-call (== per-layer, in call order) output/state error distribution, and keeps the
33586    /// SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).
33587    #[allow(clippy::too_many_arguments)]
33588    #[allow(clippy::too_many_arguments)]
33589    pub fn gdn_scan_prefill(
33590        &self,
33591        q: &CudaSlice<f32>,
33592        k: &CudaSlice<f32>,
33593        v: &CudaSlice<f32>,
33594        g: &CudaSlice<f32>,
33595        beta: &CudaSlice<f32>,
33596        kb16_pre: Option<&CudaSlice<u8>>,
33597        qb16_pre: Option<&CudaSlice<u8>>,
33598        state_in: &CudaSlice<f32>,
33599        state_out: &mut CudaSlice<f32>,
33600        o: &mut CudaSlice<f32>,
33601        n_head: usize,
33602        t: usize,
33603        scale: f32,
33604        hk: usize,
33605    ) -> Result<(), Box<dyn std::error::Error>> {
33606        if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
33607            assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
33608            return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
33609        }
33610        if Self::gdn_chunked_enabled() && t >= 16 {
33611            self.gdn_scan_chunked(
33612                q,
33613                k,
33614                v,
33615                g,
33616                beta,
33617                kb16_pre,
33618                qb16_pre,
33619                state_in,
33620                state_out,
33621                o,
33622                n_head,
33623                t,
33624                scale,
33625                Self::gdn_chunk_size(),
33626                hk,
33627            )
33628        } else {
33629            assert!(
33630                hk == n_head,
33631                "s128 scan is broadcast-only (prep guarantees by predicate)"
33632            );
33633            self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
33634        }
33635    }
33636
33637    /// Stage-1 oracle: run sequential AND chunked, report per-call error stats, keep sequential.
33638    #[allow(clippy::too_many_arguments)]
33639    fn gdn_scan_diff(
33640        &self,
33641        q: &CudaSlice<f32>,
33642        k: &CudaSlice<f32>,
33643        v: &CudaSlice<f32>,
33644        g: &CudaSlice<f32>,
33645        beta: &CudaSlice<f32>,
33646        state_in: &CudaSlice<f32>,
33647        state_out: &mut CudaSlice<f32>,
33648        o: &mut CudaSlice<f32>,
33649        n_head: usize,
33650        t: usize,
33651        scale: f32,
33652    ) -> Result<(), Box<dyn std::error::Error>> {
33653        static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
33654        let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
33655        let mut o_c = self.uninit(o.len())?;
33656        let mut st_c = self.uninit(state_out.len())?;
33657        self.gdn_scan_chunked(
33658            q,
33659            k,
33660            v,
33661            g,
33662            beta,
33663            None,
33664            None,
33665            state_in,
33666            &mut st_c,
33667            &mut o_c,
33668            n_head,
33669            t,
33670            scale,
33671            Self::gdn_chunk_size(),
33672            n_head,
33673        )?;
33674        self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
33675        let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
33676        let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
33677        let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
33678            let mut max_abs = 0f32;
33679            let mut max_rel = 0f32;
33680            let mut sum_rel = 0f64;
33681            for (x, y) in a.iter().zip(b) {
33682                let ad = (x - y).abs();
33683                let rel = ad / x.abs().max(y.abs()).max(1e-3);
33684                if ad > max_abs {
33685                    max_abs = ad;
33686                }
33687                if rel > max_rel {
33688                    max_rel = rel;
33689                }
33690                sum_rel += rel as f64;
33691            }
33692            (max_abs, max_rel, sum_rel / a.len() as f64)
33693        };
33694        let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
33695        let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
33696        println!(
33697            "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
33698                  state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
33699            Self::gdn_chunk_size()
33700        );
33701        Ok(())
33702    }
33703
33704    /// softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.
33705    pub fn gdn_glog(
33706        &self,
33707        alpha: &CudaSlice<f32>,
33708        dt_bias: &CudaSlice<f32>,
33709        a: &CudaSlice<f32>,
33710        g_log: &mut CudaSlice<f32>,
33711        n_head: usize,
33712        t: usize,
33713    ) -> Result<(), Box<dyn std::error::Error>> {
33714        let f = self.func("gdn_glog_f32");
33715        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
33716        let (h, ti) = (n_head as i32, t as i32);
33717        let __s_b = self.gpu.stream();
33718        let mut b = __s_b.launch_builder(&f);
33719        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
33720        unsafe {
33721            b.launch(cfg)?;
33722        }
33723        Ok(())
33724    }
33725
33726    /// view twins (task #16): the batched prime's GDN core reads the CONCAT projection
33727    /// buffers at row offsets (CudaView) — same kernels, same values, no split copies.
33728    pub fn sigmoid_v(
33729        &self,
33730        x: &cudarc::driver::CudaView<f32>,
33731        y: &mut CudaSlice<f32>,
33732        n: usize,
33733    ) -> Result<(), Box<dyn std::error::Error>> {
33734        let f = self.func("sigmoid_f32");
33735        let cfg = LaunchConfig::for_num_elems(n as u32);
33736        let ni = n as i32;
33737        let __s_b = self.gpu.stream();
33738        let mut b = __s_b.launch_builder(&f);
33739        b.arg(x).arg(y).arg(&ni);
33740        unsafe {
33741            b.launch(cfg)?;
33742        }
33743        Ok(())
33744    }
33745
33746    pub fn gdn_glog_v(
33747        &self,
33748        alpha: &cudarc::driver::CudaView<f32>,
33749        dt_bias: &CudaSlice<f32>,
33750        a: &CudaSlice<f32>,
33751        g_log: &mut CudaSlice<f32>,
33752        n_head: usize,
33753        t: usize,
33754    ) -> Result<(), Box<dyn std::error::Error>> {
33755        let f = self.func("gdn_glog_f32");
33756        let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
33757        let (h, ti) = (n_head as i32, t as i32);
33758        let __s_b = self.gpu.stream();
33759        let mut b = __s_b.launch_builder(&f);
33760        b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
33761        unsafe {
33762            b.launch(cfg)?;
33763        }
33764        Ok(())
33765    }
33766
33767    pub fn sigmoid(
33768        &self,
33769        x: &CudaSlice<f32>,
33770        y: &mut CudaSlice<f32>,
33771        n: usize,
33772    ) -> Result<(), Box<dyn std::error::Error>> {
33773        let f = self.func("sigmoid_f32");
33774        let cfg = LaunchConfig::for_num_elems(n as u32);
33775        let ni = n as i32;
33776        let __s_b = self.gpu.stream();
33777        let mut b = __s_b.launch_builder(&f);
33778        b.arg(x).arg(y).arg(&ni);
33779        unsafe {
33780            b.launch(cfg)?;
33781        }
33782        Ok(())
33783    }
33784
33785    /// attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch
33786    /// (replaces sigmoid + mul + convert). Bit-identical class.
33787    pub fn sig_mul_f16out(
33788        &self,
33789        a: &CudaSlice<f32>,
33790        g: &CudaSlice<f32>,
33791        dst: &mut CudaSlice<f32>,
33792        dst16: &mut CudaSlice<u8>,
33793        n: usize,
33794    ) -> Result<(), Box<dyn std::error::Error>> {
33795        let f = self.func("sig_mul_f16out_f32");
33796        let cfg = LaunchConfig::for_num_elems(n as u32);
33797        let ni = n as i32;
33798        let __s_b = self.gpu.stream();
33799        let mut b = __s_b.launch_builder(&f);
33800        b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
33801        unsafe {
33802            b.launch(cfg)?;
33803        }
33804        Ok(())
33805    }
33806
33807    /// step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
33808    /// broadcast over head_dim. `dst = a * sigmoid(g)` where `a`/`dst` are `[head_dim, n_head, T]`
33809    /// (the `q_gate_split` layout) and `g` is the PRE-sigmoid `attn_gate` projection output in
33810    /// token-major `[T, n_head]`. `dst16` is the optional fp16 operand for wo (None -> skipped).
33811    ///
33812    /// NOT interchangeable with `sig_mul_f16out`, which gates FULL WIDTH (qwen35 packs one gate
33813    /// value per (head, dim) element inside wq). Using this for that, or that for this, silently
33814    /// applies the wrong number of distinct gate values.
33815    #[allow(clippy::too_many_arguments)]
33816    pub fn attn_head_gate(
33817        &self,
33818        a: &CudaSlice<f32>,
33819        g: &CudaSlice<f32>,
33820        dst: &mut CudaSlice<f32>,
33821        dst16: Option<&mut CudaSlice<u8>>,
33822        head_dim: usize,
33823        n_head: usize,
33824        t: usize,
33825    ) -> Result<(), Box<dyn std::error::Error>> {
33826        let f = self.func("attn_head_gate_f32");
33827        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
33828        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
33829        // nullable device pointer by value (0 = skip), same convention as `l2_norm_pp`.
33830        let d16: u64 = match dst16 {
33831            Some(d) => self.addr_u8(d),
33832            None => 0,
33833        };
33834        let __s_b = self.gpu.stream();
33835        let mut b = __s_b.launch_builder(&f);
33836        b.arg(a)
33837            .arg(g)
33838            .arg(dst)
33839            .arg(&d16)
33840            .arg(&hd)
33841            .arg(&nh)
33842            .arg(&ti);
33843        unsafe {
33844            b.launch(cfg)?;
33845        }
33846        Ok(())
33847    }
33848
33849    /// step35 CLAMPED SwiGLU: `dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit)`.
33850    /// Verbatim from llama.cpp `llama-graph.cpp:2146-2165` (routed, `swiglu_clamp_exp`) and
33851    /// `:1751-1770` (shared, `swiglu_clamp_shexp`), non-DEEPSEEK4 branch.
33852    ///
33853    /// This is NOT `swigluoai_mul_scaled`: that one clamps the gate BEFORE swish and multiplies by
33854    /// `(1 + clamp(up))`. Caller MUST check `limit > 1e-6` (upstream's eps gate) and use the plain
33855    /// `silu_mul_scaled` path otherwise — at limit=0 this kernel would clamp every positive
33856    /// activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
33857    #[allow(clippy::too_many_arguments)]
33858    pub fn swiglu_clamped_mul_scaled(
33859        &self,
33860        gate: &CudaSlice<f32>,
33861        up: &CudaSlice<f32>,
33862        gs: f32,
33863        us: f32,
33864        limit: f32,
33865        dst: &mut CudaSlice<f32>,
33866        n: usize,
33867    ) -> Result<(), Box<dyn std::error::Error>> {
33868        debug_assert!(
33869            limit > 1e-6,
33870            "swiglu_clamped needs a live limit; use silu_mul_scaled"
33871        );
33872        let f = self.func("swiglu_clamped_mul_scaled_f32");
33873        let cfg = LaunchConfig::for_num_elems(n as u32);
33874        let ni = n as i32;
33875        let __s_b = self.gpu.stream();
33876        let mut b = __s_b.launch_builder(&f);
33877        b.arg(gate)
33878            .arg(up)
33879            .arg(&gs)
33880            .arg(&us)
33881            .arg(&limit)
33882            .arg(dst)
33883            .arg(&ni);
33884        unsafe {
33885            b.launch(cfg)?;
33886        }
33887        Ok(())
33888    }
33889
33890    /// glm5_next PRE-clamped SwiGLU: `dst = silu(min(gate*gs, limit)) * clamp(up*us, +-limit)`.
33891    /// The gate clamp is BEFORE silu and one-sided — vendor `Glm5NextTextMLP.forward` /
33892    /// `Glm5NextTextExperts._apply_gate`, one `swiglu_limit` shared by the dense MLP, the routed
33893    /// experts and the shared expert on every layer.
33894    ///
33895    /// This is NOT `swiglu_clamped_mul_scaled` (step35 clamps the silu OUTPUT) and NOT
33896    /// `swigluoai_mul_scaled` (alpha-swish plus a `1 +` linear term). Same caller contract as the
33897    /// post-clamp sibling: `limit > 1e-6`, else the plain `silu_mul_scaled` path.
33898    #[allow(clippy::too_many_arguments)]
33899    pub fn swiglu_preclamped_mul_scaled(
33900        &self,
33901        gate: &CudaSlice<f32>,
33902        up: &CudaSlice<f32>,
33903        gs: f32,
33904        us: f32,
33905        limit: f32,
33906        dst: &mut CudaSlice<f32>,
33907        n: usize,
33908    ) -> Result<(), Box<dyn std::error::Error>> {
33909        debug_assert!(
33910            limit > 1e-6,
33911            "swiglu_preclamped needs a live limit; use silu_mul_scaled"
33912        );
33913        let f = self.func("swiglu_preclamped_mul_scaled_f32");
33914        let cfg = LaunchConfig::for_num_elems(n as u32);
33915        let ni = n as i32;
33916        let __s_b = self.gpu.stream();
33917        let mut b = __s_b.launch_builder(&f);
33918        b.arg(gate)
33919            .arg(up)
33920            .arg(&gs)
33921            .arg(&us)
33922            .arg(&limit)
33923            .arg(dst)
33924            .arg(&ni);
33925        unsafe {
33926            b.launch(cfg)?;
33927        }
33928        Ok(())
33929    }
33930
33931    /// gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.
33932    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
33933    pub fn gated_rmsnorm(
33934        &self,
33935        o: &CudaSlice<f32>,
33936        w: &CudaSlice<f32>,
33937        z: &CudaSlice<f32>,
33938        dst: &mut CudaSlice<f32>,
33939        ncols: usize,
33940        nrows: usize,
33941        eps: f32,
33942    ) -> Result<(), Box<dyn std::error::Error>> {
33943        let f = self.func("gated_rmsnorm_f32");
33944        let cfg = LaunchConfig {
33945            grid_dim: (nrows as u32, 1, 1),
33946            block_dim: (128, 1, 1),
33947            shared_mem_bytes: 0,
33948        };
33949        let (nc, e) = (ncols as i32, eps);
33950        let __s_b = self.gpu.stream();
33951        let mut b = __s_b.launch_builder(&f);
33952        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
33953        unsafe {
33954            b.launch(cfg)?;
33955        }
33956        Ok(())
33957    }
33958
33959    /// f16out twin of `gated_rmsnorm` (task #17): epilogue also emits the fp16 operand for
33960    /// the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel's __float2half).
33961    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
33962    pub fn gated_rmsnorm_f16out(
33963        &self,
33964        o: &CudaSlice<f32>,
33965        w: &CudaSlice<f32>,
33966        z: &CudaSlice<f32>,
33967        dst: &mut CudaSlice<f32>,
33968        dst16: &mut CudaSlice<u8>,
33969        ncols: usize,
33970        nrows: usize,
33971        eps: f32,
33972    ) -> Result<(), Box<dyn std::error::Error>> {
33973        let f = self.func("gated_rmsnorm_f16out_f32");
33974        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
33975        let cfg = LaunchConfig {
33976            grid_dim: (nrows as u32, 1, 1),
33977            block_dim: (128, 1, 1),
33978            shared_mem_bytes: 0,
33979        };
33980        let (nc, e) = (ncols as i32, eps);
33981        let __s_b = self.gpu.stream();
33982        let mut b = __s_b.launch_builder(&f);
33983        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
33984        unsafe {
33985            b.launch(cfg)?;
33986        }
33987        Ok(())
33988    }
33989
33990    /// add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE
33991    /// layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to
33992    /// add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.
33993    #[allow(clippy::too_many_arguments)]
33994    pub fn add_rms_norm_zq8(
33995        &self,
33996        a: &CudaSlice<f32>,
33997        b_in: &CudaSlice<f32>,
33998        w: &CudaSlice<f32>,
33999        res: &mut CudaSlice<f32>,
34000        z: &mut CudaSlice<f32>,
34001        ncols: usize,
34002        nrows: usize,
34003        eps: f32,
34004    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
34005        assert!(ncols.is_multiple_of(32));
34006        let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
34007        let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
34008        let f = self.func("add_rms_norm_zq8");
34009        let cfg = LaunchConfig {
34010            grid_dim: (nrows as u32, 1, 1),
34011            block_dim: (1024, 1, 1),
34012            shared_mem_bytes: 0,
34013        };
34014        let (nc, ep) = (ncols as i32, eps);
34015        let __s_b = self.gpu.stream();
34016        let mut b = __s_b.launch_builder(&f);
34017        b.arg(a)
34018            .arg(b_in)
34019            .arg(w)
34020            .arg(res)
34021            .arg(z)
34022            .arg(&mut q)
34023            .arg(&mut d)
34024            .arg(&nc)
34025            .arg(&ep);
34026        unsafe {
34027            b.launch(cfg)?;
34028        }
34029        Ok((q, d))
34030    }
34031
34032    /// gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input.
34033    /// BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle
34034    /// rows). Saves one launch per linear-attn layer (36/token on the 9B).
34035    /// z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.
34036    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
34037    pub fn gated_rmsnorm_zv(
34038        &self,
34039        o: &CudaSlice<f32>,
34040        w: &CudaSlice<f32>,
34041        z: &cudarc::driver::CudaView<f32>,
34042        dst: &mut CudaSlice<f32>,
34043        ncols: usize,
34044        nrows: usize,
34045        eps: f32,
34046    ) -> Result<(), Box<dyn std::error::Error>> {
34047        let f = self.func("gated_rmsnorm_f32");
34048        let cfg = LaunchConfig {
34049            grid_dim: (nrows as u32, 1, 1),
34050            block_dim: (128, 1, 1),
34051            shared_mem_bytes: 0,
34052        };
34053        let (nc, e) = (ncols as i32, eps);
34054        let __s_b = self.gpu.stream();
34055        let mut b = __s_b.launch_builder(&f);
34056        b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
34057        unsafe {
34058            b.launch(cfg)?;
34059        }
34060        Ok(())
34061    }
34062
34063    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
34064    pub fn gated_rmsnorm_f16out_zv(
34065        &self,
34066        o: &CudaSlice<f32>,
34067        w: &CudaSlice<f32>,
34068        z: &cudarc::driver::CudaView<f32>,
34069        dst: &mut CudaSlice<f32>,
34070        dst16: &mut CudaSlice<u8>,
34071        ncols: usize,
34072        nrows: usize,
34073        eps: f32,
34074    ) -> Result<(), Box<dyn std::error::Error>> {
34075        let f = self.func("gated_rmsnorm_f16out_f32");
34076        // block_dim MUST match gated_rmsnorm's (128): the reduction tree order pins the scale
34077        let cfg = LaunchConfig {
34078            grid_dim: (nrows as u32, 1, 1),
34079            block_dim: (128, 1, 1),
34080            shared_mem_bytes: 0,
34081        };
34082        let (nc, e) = (ncols as i32, eps);
34083        let __s_b = self.gpu.stream();
34084        let mut b = __s_b.launch_builder(&f);
34085        b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
34086        unsafe {
34087            b.launch(cfg)?;
34088        }
34089        Ok(())
34090    }
34091
34092    pub fn gated_rmsnorm_q8_1(
34093        &self,
34094        o: &CudaSlice<f32>,
34095        w: &CudaSlice<f32>,
34096        z: &CudaSlice<f32>,
34097        ncols: usize,
34098        nrows: usize,
34099        eps: f32,
34100    ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
34101        assert!(ncols.is_multiple_of(32));
34102        let f = self.func("gated_rmsnorm_q8_1");
34103        let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
34104        let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
34105        let cfg = LaunchConfig {
34106            grid_dim: (nrows as u32, 1, 1),
34107            block_dim: (128, 1, 1),
34108            shared_mem_bytes: 0,
34109        };
34110        let (nc, ep) = (ncols as i32, eps);
34111        let __s_b = self.gpu.stream();
34112        let mut b = __s_b.launch_builder(&f);
34113        b.arg(o)
34114            .arg(w)
34115            .arg(z)
34116            .arg(&mut out_q)
34117            .arg(&mut out_d)
34118            .arg(&nc)
34119            .arg(&ep);
34120        unsafe {
34121            b.launch(cfg)?;
34122        }
34123        Ok((out_q, out_d))
34124    }
34125
34126    /// transpose [rows,cols] row-major -> [cols,rows] row-major.
34127    pub fn transpose(
34128        &self,
34129        inp: &CudaSlice<f32>,
34130        rows: usize,
34131        cols: usize,
34132    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34133        let f = self.func("transpose_f32");
34134        let mut out = self.zeros(rows * cols)?;
34135        let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
34136        let (r, c) = (rows as i32, cols as i32);
34137        let __s_b = self.gpu.stream();
34138        let mut b = __s_b.launch_builder(&f);
34139        b.arg(inp).arg(&mut out).arg(&r).arg(&c);
34140        unsafe {
34141            b.launch(cfg)?;
34142        }
34143        Ok(out)
34144    }
34145
34146    /// repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].
34147    pub fn repeat_heads(
34148        &self,
34149        inp: &CudaSlice<f32>,
34150        out: &mut CudaSlice<f32>,
34151        head_dim: usize,
34152        n_in: usize,
34153        n_out: usize,
34154        t: usize,
34155    ) -> Result<(), Box<dyn std::error::Error>> {
34156        let f = self.func("repeat_heads_f32");
34157        let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
34158        let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
34159        let __s_b = self.gpu.stream();
34160        let mut b = __s_b.launch_builder(&f);
34161        b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
34162        unsafe {
34163            b.launch(cfg)?;
34164        }
34165        Ok(())
34166    }
34167
34168    /// q|gate split (on-device). qf:[T, n_head*2*head_dim] -> q_out,gate_out:[head_dim,n_head,T].
34169    /// Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.
34170    ///
34171    /// Bounds-checked against `qf` before the launch: the kernel reads `2*head_dim*n_head*t`
34172    /// floats, and running it on a `wq` output that carries no fused gate reads 2x off the end
34173    /// (silently, on the device). A layout mismatch is a typed `FusedQGateExtent` here instead.
34174    pub fn q_gate_split(
34175        &self,
34176        qf: &CudaSlice<f32>,
34177        q_out: &mut CudaSlice<f32>,
34178        gate_out: &mut CudaSlice<f32>,
34179        head_dim: usize,
34180        n_head: usize,
34181        t: usize,
34182    ) -> Result<(), Box<dyn std::error::Error>> {
34183        memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
34184        let out_need = head_dim * n_head * t;
34185        if q_out.len() < out_need || gate_out.len() < out_need {
34186            return Err(format!(
34187                "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
34188                q_out.len(),
34189                gate_out.len()
34190            )
34191            .into());
34192        }
34193        let f = self.func("q_gate_split_f32");
34194        let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
34195        let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
34196        let __s_b = self.gpu.stream();
34197        let mut b = __s_b.launch_builder(&f);
34198        b.arg(qf)
34199            .arg(q_out)
34200            .arg(gate_out)
34201            .arg(&hd)
34202            .arg(&nh)
34203            .arg(&ti);
34204        unsafe {
34205            b.launch(cfg)?;
34206        }
34207        Ok(())
34208    }
34209
34210    /// qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major ->
34211    /// q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping).
34212    /// Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.
34213    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
34214    pub fn qkv_to_gdn_repack(
34215        &self,
34216        conv_out: &CudaSlice<f32>,
34217        q_g: &mut CudaSlice<f32>,
34218        k_g: &mut CudaSlice<f32>,
34219        v_g: &mut CudaSlice<f32>,
34220        d_state: usize,
34221        num_v: usize,
34222        num_k: usize,
34223        key_dim: usize,
34224        t: usize,
34225    ) -> Result<(), Box<dyn std::error::Error>> {
34226        let f = self.func("qkv_to_gdn_repack_f32");
34227        let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
34228        let (ds, nv, nk, kd, ti) = (
34229            d_state as i32,
34230            num_v as i32,
34231            num_k as i32,
34232            key_dim as i32,
34233            t as i32,
34234        );
34235        let __s_b = self.gpu.stream();
34236        let mut b = __s_b.launch_builder(&f);
34237        b.arg(conv_out)
34238            .arg(q_g)
34239            .arg(k_g)
34240            .arg(v_g)
34241            .arg(&ds)
34242            .arg(&nv)
34243            .arg(&nk)
34244            .arg(&kd)
34245            .arg(&ti);
34246        unsafe {
34247            b.launch(cfg)?;
34248        }
34249        Ok(())
34250    }
34251
34252    /// conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad],
34253    /// cols 0..pad = 0, cols pad..pad+T = src. `dst` MUST be pre-zeroed. No dtoh/host-loop/htod.
34254    pub fn conv_left_pad(
34255        &self,
34256        src: &CudaSlice<f32>,
34257        dst: &mut CudaSlice<f32>,
34258        conv_dim: usize,
34259        t: usize,
34260        pad: usize,
34261    ) -> Result<(), Box<dyn std::error::Error>> {
34262        let f = self.func("conv_left_pad_f32");
34263        let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
34264        let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
34265        let __s_b = self.gpu.stream();
34266        let mut b = __s_b.launch_builder(&f);
34267        b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
34268        unsafe {
34269            b.launch(cfg)?;
34270        }
34271        Ok(())
34272    }
34273
34274    /// conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident),
34275    /// qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols).
34276    /// Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.
34277    pub fn conv_assemble_and_roll(
34278        &self,
34279        qkv_col: &CudaSlice<f32>,
34280        conv_state: &mut CudaSlice<f32>,
34281        conv_in: &mut CudaSlice<f32>,
34282        conv_dim: usize,
34283        pad: usize,
34284    ) -> Result<(), Box<dyn std::error::Error>> {
34285        let f = self.func("conv_assemble_and_roll_f32");
34286        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
34287        let (cd, p) = (conv_dim as i32, pad as i32);
34288        let __s_b = self.gpu.stream();
34289        let mut b = __s_b.launch_builder(&f);
34290        b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
34291        unsafe {
34292            b.launch(cfg)?;
34293        }
34294        Ok(())
34295    }
34296
34297    /// RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE
34298    /// launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise
34299    /// causal conv + SiLU into `conv_out`, and rolls the ring — never materializing conv_in to HBM.
34300    /// Replaces e.conv_assemble_and_roll(...) + e.ssm_conv1d(...). BIT-IDENTICAL to that two-kernel
34301    /// sequence (same 8-wide accumulation order, same SiLU). `conv_out` is [conv_dim] (T=1).
34302    pub fn ssm_conv1d_fused_decode(
34303        &self,
34304        qkv_col: &CudaSlice<f32>,
34305        conv_state: &mut CudaSlice<f32>,
34306        w: &CudaSlice<f32>,
34307        conv_out: &mut CudaSlice<f32>,
34308        conv_dim: usize,
34309        d_conv: usize,
34310    ) -> Result<(), Box<dyn std::error::Error>> {
34311        let f = self.func("ssm_conv1d_fused_decode_f32");
34312        let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
34313        let (cd, dc) = (conv_dim as i32, d_conv as i32);
34314        let __s_b = self.gpu.stream();
34315        let mut b = __s_b.launch_builder(&f);
34316        b.arg(qkv_col)
34317            .arg(conv_state)
34318            .arg(w)
34319            .arg(conv_out)
34320            .arg(&cd)
34321            .arg(&dc);
34322        unsafe {
34323            b.launch(cfg)?;
34324        }
34325        Ok(())
34326    }
34327
34328    /// Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host).
34329    /// Used for qkv split views. Small/rare; not perf-critical in Stage 1.
34330    pub fn slice_range(
34331        &self,
34332        src: &CudaSlice<f32>,
34333        start: usize,
34334        len: usize,
34335    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34336        let host = self.gpu.stream().clone_dtoh(src)?;
34337        self.gpu.stream().synchronize()?;
34338        self.htod(&host[start..start + len])
34339    }
34340}
34341
34342#[cfg(test)]
34343mod target_dispatch_tests {
34344    use super::legacy_quant_gemm_allowed;
34345
34346    #[test]
34347    fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
34348        // sm_120a native lane
34349        assert!(legacy_quant_gemm_allowed(false, false, false));
34350        assert!(!legacy_quant_gemm_allowed(false, false, true));
34351        // pure portable lane (sm_89): gated
34352        assert!(!legacy_quant_gemm_allowed(true, false, false));
34353        assert!(!legacy_quant_gemm_allowed(true, false, true));
34354        // Hopper-MMA lane (sm_90a): portable build, int8-MMA GEMM re-admitted
34355        assert!(legacy_quant_gemm_allowed(true, true, false));
34356        assert!(!legacy_quant_gemm_allowed(true, true, true));
34357    }
34358
34359    #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
34360    #[test]
34361    fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
34362        assert!(!legacy_quant_gemm_allowed(
34363            cfg!(memra_portable_cuda),
34364            cfg!(memra_hopper_mma),
34365            false
34366        ));
34367    }
34368
34369    #[cfg(memra_hopper_mma)]
34370    #[test]
34371    fn hopper_mma_build_re_admits_legacy_quant_gemm() {
34372        assert!(legacy_quant_gemm_allowed(
34373            cfg!(memra_portable_cuda),
34374            cfg!(memra_hopper_mma),
34375            false
34376        ));
34377        assert!(super::portable_mma_gated() == false);
34378    }
34379}
34380
34381/// The memra-kv device seam (Phase D): the cache's 7 ops delegate to the engine's
34382/// inherent methods (inherent methods win name resolution, so no recursion).
34383impl memra_kv::KvDev for Engine {
34384    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34385        Engine::zeros(self, n)
34386    }
34387    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34388        Engine::uninit(self, n)
34389    }
34390    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
34391        Engine::alloc_u8(self, n)
34392    }
34393    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
34394        Engine::htod_i32(self, v)
34395    }
34396    fn clone_dtod(
34397        &self,
34398        src: &CudaSlice<f32>,
34399    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34400        Engine::clone_dtod(self, src)
34401    }
34402    fn copy_into(
34403        &self,
34404        dst: &mut CudaSlice<f32>,
34405        off: usize,
34406        src: &CudaSlice<f32>,
34407        len: usize,
34408    ) -> Result<(), Box<dyn std::error::Error>> {
34409        Engine::copy_into(self, dst, off, src, len)
34410    }
34411    fn copy_range_into(
34412        &self,
34413        dst: &mut CudaSlice<f32>,
34414        dst_off: usize,
34415        src: &CudaSlice<f32>,
34416        src_off: usize,
34417        len: usize,
34418    ) -> Result<(), Box<dyn std::error::Error>> {
34419        Engine::copy_range_into(self, dst, dst_off, src, src_off, len)
34420    }
34421    fn set_i32_one(
34422        &self,
34423        d: &mut CudaSlice<i32>,
34424        v: i32,
34425    ) -> Result<(), Box<dyn std::error::Error>> {
34426        Engine::set_i32_one(self, d, v)
34427    }
34428}
34429
34430#[cfg(test)]
34431mod fused_gate_bounds_tests {
34432    use super::*;
34433
34434    /// The fused `[q|gate]` split's read-site guard, on the device.
34435    ///
34436    /// `q_gate_split_f32` reads `2*head_dim*n_head*T` floats out of `qf`. A checkpoint whose gate
34437    /// is a SEPARATE tensor produces a `wq` output of exactly half that, so before 2026-08-19 the
34438    /// kernel launched and read 2x past the end of the allocation — an out-of-bounds DEVICE read:
34439    /// no panic, no error, just whatever memory follows. The guard turns it into a typed
34440    /// `FusedQGateExtent` before the launch.
34441    ///
34442    /// Catch demonstration for this test (guard temporarily removed, then restored):
34443    /// `compute-sanitizer --tool memcheck` on the half-width case reported invalid `__global__`
34444    /// reads of size 4 in `q_gate_split_f32`; with the guard in place the same run is clean and
34445    /// the call returns `Err`. Receipt in the lane report.
34446    #[test]
34447    #[ignore = "requires a CUDA GPU"]
34448    fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
34449        let e = Engine::new(0).unwrap();
34450        let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
34451        let fused = 2 * head_dim * n_head * t;
34452        let out_n = head_dim * n_head * t;
34453
34454        // half-width `qf` = the separate-gate / ungated layout. MUST be refused.
34455        let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
34456        let mut q = e.uninit(out_n).unwrap();
34457        let mut gate = e.uninit(out_n).unwrap();
34458        let err = e
34459            .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
34460            .expect_err("half-width wq must be refused, not read past")
34461            .to_string();
34462        assert!(err.contains("NO fused gate"), "{err}");
34463        assert!(err.contains(&format!("{fused}")), "{err}");
34464
34465        // full-width `qf` = a real qwen3.5 fused layout. MUST still run, and split correctly:
34466        // per head hh the block is [q(head_dim) | gate(head_dim)] at stride 2*head_dim.
34467        let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
34468        let wide = e.htod(&host).unwrap();
34469        e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
34470            .expect("full-width wq splits");
34471        let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
34472        for tok in 0..t {
34473            for hh in 0..n_head {
34474                for d in 0..head_dim {
34475                    let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
34476                    let idx = tok * (n_head * head_dim) + hh * head_dim + d;
34477                    assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
34478                    assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
34479                }
34480            }
34481        }
34482
34483        // undersized destinations are refused too (the other half of the extent contract)
34484        let mut small = e.uninit(out_n - 1).unwrap();
34485        assert!(
34486            e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
34487                .is_err()
34488        );
34489    }
34490}
34491
34492/// FULL-WIDTH-ROPE CONTRACT on the fused rms_norm+qkv+rope kernels
34493/// (lane/graph-s-key-exactness-20260819, probe O-10 follow-up). CPU-only: the guard runs before
34494/// any launch, so the refusal is testable without a device.
34495#[cfg(test)]
34496mod fused_rope_width_tests {
34497    use super::Engine;
34498
34499    /// gemma-4: rotary width == head width on both classes (GGUF 256/256 and 512/512, and the
34500    /// safetensors route derives the same), which is why the fusion is legal there today.
34501    #[test]
34502    fn full_width_is_accepted() {
34503        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
34504        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
34505        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
34506    }
34507
34508    /// The widths the gemma-4 31B OFFICIAL artifact declares, read from its own GGUF header
34509    /// (`gemma-4-31B-it-official-Q8_0-MTP.gguf`, box3, 2026-08-19):
34510    ///
34511    /// ```text
34512    /// attention.key_length     512   rope.dimension_count     512   (global class)
34513    /// attention.key_length_swa 256   rope.dimension_count_swa 256   (SWA class)
34514    /// ```
34515    ///
34516    /// Both classes satisfy `n_rot == head_dim`, which is why the fusion is legal for gemma and
34517    /// why `HybridModel::gemma4_rope_dims` can feed this guard without refusing what we serve.
34518    /// An artifact that ever declares otherwise gets a loud refusal at the first fused launch
34519    /// instead of a silently over-rotated head.
34520    #[test]
34521    fn gemma4_official_artifact_widths_pass() {
34522        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
34523        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
34524    }
34525
34526    /// The bug this guard exists to make impossible: a partial-rotary arch fused onto a kernel
34527    /// with no `n_dims`, silently rotating the pass-through band.
34528    #[test]
34529    fn partial_rotary_is_refused_with_the_geometry_named() {
34530        // qwen3.5: n_rot 64 of head_dim 256 (the shape probe O-10 pinned in the split path).
34531        let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
34532            .expect_err("partial rotary must refuse");
34533        let msg = err.to_string();
34534        assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
34535        assert!(msg.contains("n_rot 64"), "{msg}");
34536        assert!(msg.contains("head_dim 256"), "{msg}");
34537        assert!(
34538            msg.contains("64..256"),
34539            "names the band it would corrupt: {msg}"
34540        );
34541        // step35 full-attn: 64 of 128 (upstream halves n_rot_full).
34542        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
34543        // and the reverse mismatch (a wider rope than the head) is not "close enough" either.
34544        assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
34545    }
34546}