Skip to main content

memra_engine/
mmq_ffi.rs

1//! FFI to the MMQ prefill GEMMs (cu/mmq_fp4.cu + cu/mmq_q45k.cu) — vendored floor kernels.
2//!
3//! NVFP4: the 5150-pp512 kernel from llama.cpp, ggml-decoupled into a static lib with a C-ABI host
4//! launcher. The launcher quantizes the f32 activation to block_fp4_mmq internally (llama's 2-level
5//! FP8-e8m0/UE4M3 scale = the accurate W4A8-via-FP8 path that fixes memra's W4A4 maxdiff 1.46), then
6//! launches the native mxf4nvf4 block-scale tensor-core mma.
7//!
8//! Q4_K/Q5_K: llama's k-quant int8-MMA MMQ (dequant to int8 at tile-load, q8_1 DS4 activation with
9//! the (d, sum) pair that feeds the k-quant min-offset term, shared m16n8k32 s8 mma inner loop).
10//! Replaces the hand-rolled qmatvec_gemm k-quant GEMMs that dominate prefill (32% + 28% busy).
11//!
12//! All dispatched behind MEMRA_MMQ=1. Always built (no external deps) — unlike cutlass_ffi which is
13//! MEMRA_CUTLASS-gated.
14
15use crate::Engine;
16use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut};
17
18/// Quantize-once seam state (see `Engine::mmq_act_begin`): window epoch + one cached
19/// (epoch, act_ptr, m, in_f, D4 scratch) slot. Slot drops (freeing the scratch) on each new window.
20static MMQ_ACT_EPOCH: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21#[allow(clippy::type_complexity)]
22static MMQ_ACT_SLOT: std::sync::Mutex<Option<(u64, u64, usize, usize, CudaSlice<u8>)>> =
23    std::sync::Mutex::new(None);
24/// Stream-k fixup scratch (lazy; sized once per process — one slot per SM).
25static MMQ_FIXUP_SLOT: std::sync::Mutex<Option<cudarc::driver::CudaSlice<u8>>> =
26    std::sync::Mutex::new(None);
27
28unsafe extern "C" {
29    /// Bytes needed for the block_fp4_mmq activation scratch for (in_f, n_tokens).
30    pub fn memra_mmq_nvfp4_act_bytes(in_f: i32, n_tokens: i32) -> usize;
31    /// Run the NVFP4 W4A4 MMQ prefill GEMM. y[n_tokens, out_f] = act[n_tokens, in_f] @ W[out_f, in_f]^T.
32    ///   W_nvfp4_blocks : raw memra NVFP4 weight rows (block_nvfp4 36B blocks, in_f/64 per row).
33    ///   act_f32        : f32 activation [n_tokens, in_f] (contiguous).
34    ///   y              : f32 output [n_tokens, out_f].
35    ///   act_scratch    : pre-alloc'd quant buffer >= memra_mmq_nvfp4_act_bytes(in_f, n_tokens).
36    /// Returns 0 on success, else (1000 + cudaError).
37    pub fn memra_mmq_nvfp4(
38        w_nvfp4_blocks: *const core::ffi::c_void,
39        act_f32: *const f32,
40        y: *mut f32,
41        in_f: i32,
42        out_f: i32,
43        n_tokens: i32,
44        act_scratch: *mut core::ffi::c_void,
45        stream: *mut core::ffi::c_void,
46        out_scale: f32,
47    ) -> i32;
48    /// Same as `memra_mmq_nvfp4`, plus the activation-quantizer selector.
49    ///   per_token_scale = 1: two-level scaling (per-token row amax folded into the GEMM epilogue
50    ///     + per-sub-block UE4M3). This is what `memra_mmq_nvfp4` does.
51    ///   per_token_scale = 0: the v1 sub-block-only quantizer, retained as the numeric oracle so
52    ///     kernel-check can measure what the row scale bought, and as the rollback seam.
53    pub fn memra_mmq_nvfp4_ex(
54        w_nvfp4_blocks: *const core::ffi::c_void,
55        act_f32: *const f32,
56        y: *mut f32,
57        in_f: i32,
58        out_f: i32,
59        n_tokens: i32,
60        act_scratch: *mut core::ffi::c_void,
61        stream: *mut core::ffi::c_void,
62        out_scale: f32,
63        per_token_scale: i32,
64    ) -> i32;
65    /// Same as `memra_mmq_nvfp4_ex`, plus the residual high-precision channel count.
66    ///   residual_k = 0: off.
67    ///   residual_k > 0: the k largest-magnitude activation channels (ranked across the batch) are
68    ///     zeroed before quantization and their exact f32 contribution is added back as a rank-k
69    ///     correction. Requires per_token_scale = 1. Clamped to MMQ_MAX_RESIDUAL_K (64).
70    pub fn memra_mmq_nvfp4_ex2(
71        w_nvfp4_blocks: *const core::ffi::c_void,
72        act_f32: *const f32,
73        y: *mut f32,
74        in_f: i32,
75        out_f: i32,
76        n_tokens: i32,
77        act_scratch: *mut core::ffi::c_void,
78        stream: *mut core::ffi::c_void,
79        out_scale: f32,
80        per_token_scale: i32,
81        residual_k: i32,
82    ) -> i32;
83    /// Bytes needed for the block_q8_1_mmq activation scratch for the NVFP4 W4A8 path.
84    pub fn memra_mmq_nvfp4_w4a8_act_bytes(in_f: i32, n_tokens: i32) -> usize;
85    /// Run the NVFP4 W4A8 MMQ prefill GEMM (STAGE 2 accuracy-safe rung). Same fast MMQ tile as
86    /// memra_mmq_nvfp4 (W4A4) but the non-Blackwell int8 pair: weight FP4 LUT-dequantized to int8 at
87    /// tile-load, activation stays q8_1 int8 (D4, the same quant class as the default int8 GEMM).
88    /// `rp`: 0 = GGUF 36B-block weight layout, 1 = A6 split-plane repack (the resident decode
89    /// layout). The rp tile loader is a pure address remap of the GGUF loader (same dequant math,
90    /// same FP op order) — output is bit-identical either way.
91    /// Same contract as memra_mmq_nvfp4 otherwise. Returns 0 or (1000 + cudaError).
92    pub fn memra_mmq_nvfp4_w4a8(
93        w_nvfp4_blocks: *const core::ffi::c_void,
94        act_f32: *const f32,
95        y: *mut f32,
96        in_f: i32,
97        out_f: i32,
98        n_tokens: i32,
99        act_scratch: *mut core::ffi::c_void,
100        stream: *mut core::ffi::c_void,
101        out_scale: f32,
102        rp: i32,
103    ) -> i32;
104    /// Bytes for the block_e4m3_mmq activation scratch (footprint-identical to block_q8_1_mmq).
105    pub fn memra_mmq_nvfp4_f8f4_act_bytes(in_f: i32, n_tokens: i32) -> usize;
106    /// R-B W4A8-FP8 MMQ prefill GEMM (research/prefill-mxf8f6f4-design.md): NVFP4 per-16 scales
107    /// fold into e4m3 weight VALUES at tile load; e4m3 activations; ONE kind::f8f6f4 m16n8k32
108    /// MMA (381-TF class) where the int8 path issues two imma k16. NEW NUMERIC CONFIG — own
109    /// battery. Same contract/rp semantics as memra_mmq_nvfp4_w4a8. Returns 0 / 1000+cudaError /
110    /// 2000+cudaError.
111    pub fn memra_mmq_nvfp4_f8f4(
112        w_nvfp4_blocks: *const core::ffi::c_void,
113        act_f32: *const f32,
114        y: *mut f32,
115        in_f: i32,
116        out_f: i32,
117        n_tokens: i32,
118        act_scratch: *mut core::ffi::c_void,
119        stream: *mut core::ffi::c_void,
120        out_scale: f32,
121        rp: i32,
122    ) -> i32;
123    /// Bytes for the per-block FP8 MMQ activation scratch (delegates to the F8F4 sizing — the
124    /// two arms deliberately share ONE activation format, `block_e4m3_mmq`).
125    pub fn memra_mmq_fp8_blk_act_bytes(in_f: i32, n_tokens: i32) -> usize;
126    /// Scale-grid dims for an [out_f x in_f] block-128 FP8 tensor (ceil-div by 128).
127    pub fn memra_mmq_fp8_blk_scale_rows(out_f: i32) -> i32;
128    pub fn memra_mmq_fp8_blk_scale_cols(in_f: i32) -> i32;
129    /// PER-BLOCK FP8 MMQ prefill GEMM (cu/mmq_fp8_blk.cu, P1 option (b)): consumes the
130    /// Qwen-official e4m3 weight bytes + the per-[128x128] f32 scale grid DIRECTLY. The weight
131    /// side is never re-quantized (the checkpoint bytes are the MMA A operand), so unlike ARM A's
132    /// per-tensor fold there is no precision loss; unlike ARM B' it does not land on the Q8_0
133    /// floor. `blk_scales` is device f32 [ceil(out_f/128) x ceil(in_f/128)], row-major.
134    /// Requires in_f % 16 == 0. Returns 0 / 1 (bad dims) / 1000+cudaError / 2000+cudaError.
135    pub fn memra_mmq_fp8_blk(
136        w_e4m3: *const core::ffi::c_void,
137        blk_scales: *const f32,
138        act_f32: *const f32,
139        y: *mut f32,
140        in_f: i32,
141        out_f: i32,
142        n_tokens: i32,
143        act_scratch: *mut core::ffi::c_void,
144        stream: *mut core::ffi::c_void,
145        out_scale: f32,
146    ) -> i32;
147    /// Count e4m3 NaN codes (magnitude 0x7F) in a device weight buffer. Those decode to NaN in
148    /// hardware but to 0.0 in the host/ARM B' convention, so a tensor containing any must NOT
149    /// ride `memra_mmq_fp8_blk`. `out_count` is a device u32 (zeroed by the call).
150    pub fn memra_fp8_blk_count_nan(
151        w_e4m3: *const core::ffi::c_void,
152        nbytes: usize,
153        out_count: *mut u32,
154        stream: *mut core::ffi::c_void,
155    ) -> i32;
156    /// Bytes needed for the block_q8_1_mmq activation scratch (shared by Q4_K and Q5_K).
157    pub fn memra_mmq_q45k_act_bytes(in_f: i32, n_tokens: i32) -> usize;
158    /// Run the Q4_K W4A8 MMQ prefill GEMM. Same contract as memra_mmq_nvfp4 (raw ggml block_q4_K
159    /// weight rows, in_f/256 144B superblocks per row). Returns 0 or (1000 + cudaError).
160    pub fn memra_mmq_q4_K(
161        w_q4k_blocks: *const core::ffi::c_void,
162        act_f32: *const f32,
163        y: *mut f32,
164        in_f: i32,
165        out_f: i32,
166        n_tokens: i32,
167        act_scratch: *mut core::ffi::c_void,
168        stream: *mut core::ffi::c_void,
169    ) -> i32;
170    /// Run the Q5_K W4A8 MMQ prefill GEMM (176B superblocks). Same contract as memra_mmq_q4_K.
171    pub fn memra_mmq_q5_K(
172        w_q5k_blocks: *const core::ffi::c_void,
173        act_f32: *const f32,
174        y: *mut f32,
175        in_f: i32,
176        out_f: i32,
177        n_tokens: i32,
178        act_scratch: *mut core::ffi::c_void,
179        stream: *mut core::ffi::c_void,
180    ) -> i32;
181
182    /// Bytes needed for the block_q8_1_mmq (D4) activation scratch for the Q8_0 MMQ path.
183    pub fn memra_mmq_q8_0_act_bytes(in_f: i32, n_tokens: i32) -> usize;
184    /// Run the Q8_0 int8-MMA MMQ prefill GEMM (MEMRA_PP_Q8MMQ). Conventional xy-tiling only (no fixup
185    /// scratch). Weight = raw ggml block_q8_0 rows (34B blocks, in_f/32 per row); activation is
186    /// quantized internally to q8_1 D4. Requires in_f % 32 == 0. Returns 0 or (1000 + cudaError).
187    pub fn memra_mmq_q8_0(
188        w_q8_0_blocks: *const core::ffi::c_void,
189        act_f32: *const f32,
190        y: *mut f32,
191        in_f: i32,
192        out_f: i32,
193        n_tokens: i32,
194        act_scratch: *mut core::ffi::c_void,
195        stream: *mut core::ffi::c_void,
196    ) -> i32;
197
198    // ---- Q1 accumulator instrument (cu/mmq_q8_0_f32acc.cu, lane/fp8-v3-gate) ----
199    // The Q8_0 MMQ floor's GEMM with the accumulator as its ONE free variable: arm S32 is the
200    // floor's `mma...s32.s8.s8.s32`, arm F32 is the same m16n8k32 shape and the same A/B/D fragment
201    // ABI with `mma...kind::f8f6f4...f32.e4m3.e4m3.f32` — the op cu/mmq_fp8_blk.cu accumulates in.
202    // Both take a PRE-QUANTIZED block_q8_1_mmq activation buffer, so the measurement is GEMM-only
203    // and cannot differ by a quantizer. Research instrument only: no dispatch seam, and neither arm's
204    // output is a numeric claim (see the TU header).
205    /// Activation-scratch bytes for the accumulator instrument (same padding rule as the floor).
206    pub fn memra_accprobe_act_bytes(in_f: i32, n_tokens: i32) -> usize;
207    /// ARM S32 — the floor's GEMM verbatim, s32 accumulate. Returns 0, 1, or 1000+cudaError.
208    pub fn memra_accprobe_gemm_s32(
209        w_q8_0_blocks: *const core::ffi::c_void,
210        act_q: *const core::ffi::c_void,
211        y: *mut f32,
212        in_f: i32,
213        out_f: i32,
214        n_tokens: i32,
215        stream: *mut core::ffi::c_void,
216    ) -> i32;
217    /// ARM F32 — byte-identical kernel, f32 accumulate over the e4m3 reading of the same bytes.
218    pub fn memra_accprobe_gemm_f32(
219        w_q8_0_blocks: *const core::ffi::c_void,
220        act_q: *const core::ffi::c_void,
221        y: *mut f32,
222        in_f: i32,
223        out_f: i32,
224        n_tokens: i32,
225        stream: *mut core::ffi::c_void,
226    ) -> i32;
227
228    /// Bytes needed for the block_q8_1_mmq (D4) activation scratch for the Q4_0 MMQ path.
229    pub fn memra_mmq_q4_0_act_bytes(in_f: i32, n_tokens: i32) -> usize;
230    /// Run the Q4_0 int8-MMA MMQ prefill GEMM (MEMRA_PP_Q4MMQ). Nibbles dequant to int8 at
231    /// tile-load (the -8 zero-point folds into the quants, D4 epilogue — same accuracy class as
232    /// the Q8_0 MMQ). `rp`: 0 = raw ggml 18B blocks, 1 = MEMRA_Q4RP split-plane repack (qs plane +
233    /// fp16 d plane) — pure address remap, bit-identical output either way. Requires
234    /// in_f % 32 == 0. Returns 0 or (1000 + cudaError).
235    pub fn memra_mmq_q4_0(
236        w_q4_0: *const core::ffi::c_void,
237        act_f32: *const f32,
238        y: *mut f32,
239        in_f: i32,
240        out_f: i32,
241        n_tokens: i32,
242        act_scratch: *mut core::ffi::c_void,
243        stream: *mut core::ffi::c_void,
244        rp: i32,
245    ) -> i32;
246    /// Quantize-only entry (quantize-once seam): f32 activation -> block_q8_1_mmq scratch.
247    pub fn memra_mmq_q4_0_quant_act(
248        act_f32: *const f32,
249        act_scratch: *mut core::ffi::c_void,
250        in_f: i32,
251        n_tokens: i32,
252        stream: *mut core::ffi::c_void,
253    ) -> i32;
254    /// GEMM-only entry: consumes a pre-quantized scratch (from memra_mmq_q4_0_quant_act).
255    pub fn memra_mmq_q4_0_gemm(
256        w_q4_0: *const core::ffi::c_void,
257        act_scratch: *const core::ffi::c_void,
258        y: *mut f32,
259        in_f: i32,
260        out_f: i32,
261        n_tokens: i32,
262        stream: *mut core::ffi::c_void,
263        rp: i32,
264    ) -> i32;
265    /// Stream-k fixup scratch bytes (one [MMQ_X x MMQ_Y] f32 slot per SM).
266    pub fn memra_mmq_q4_0_fixup_bytes() -> usize;
267    /// Force the CLC work-stealing arm: 1 = on, 0 = off (static grid), -1 = MEMRA_MMQ_CLC env
268    /// default. Schedule-only swap of the xy-tiling kernel — bit-identical output by
269    /// construction (perf-frontier lever #1). Returns 1 when the CLC kernel is compiled in
270    /// (SM_100+ gencode), 0 on sm_89/90a builds (force is a no-op there; static grid always).
271    pub fn memra_mmq_q4_0_set_clc(force: i32) -> i32;
272    /// Stream-k GEMM entry: tiling when wave efficiency >= 90%, else stream-k + fixup.
273    pub fn memra_mmq_q4_0_gemm_sk(
274        w_q4_0: *const core::ffi::c_void,
275        act_scratch: *const core::ffi::c_void,
276        y: *mut f32,
277        fixup_scratch: *mut core::ffi::c_void,
278        in_f: i32,
279        out_f: i32,
280        n_tokens: i32,
281        stream: *mut core::ffi::c_void,
282        rp: i32,
283    ) -> i32;
284
285    // ---- IQ3_S / IQ4_XS expert-segmented int8-MMA MMQ (cu/mmq_iq_experts.cu, MEMRA_MOE_MMA) ----
286    /// Bytes for the token-major block_q8_1_mmq activation scratch (in_f, n_tokens).
287    pub fn memra_mmq_iq_experts_act_bytes(in_f: i32, n_tokens: i32) -> usize;
288    /// Quantize token-major f32 activation [n_tokens, in_f] -> block_q8_1_mmq (D4). Returns 0 or 1000+err.
289    pub fn memra_mmq_iq_quantize_act(
290        act_f32: *const f32,
291        act_scratch: *mut core::ffi::c_void,
292        in_f: i32,
293        n_tokens: i32,
294        stream: *mut core::ffi::c_void,
295    ) -> i32;
296    /// Fused act-epilogue: silu/gelu(gate)*up + q8_1_mmq (D4) quantize in ONE launch — no f32 act
297    /// buffer. gate/up pair-major [n_tokens, in_f]; scratch identical to memra_mmq_iq_quantize_act.
298    /// act_kind: 0=silu*mul, 1=gelu_tanh*mul. Byte-identical to the two-pass path (kernel-check gated).
299    pub fn memra_mmq_iq_fused_act_quant(
300        gate: *const f32,
301        up: *const f32,
302        act_scratch: *mut core::ffi::c_void,
303        in_f: i32,
304        n_tokens: i32,
305        act_kind: i32,
306        stream: *mut core::ffi::c_void,
307    ) -> i32;
308    /// Expert-segmented IQ MMA MMQ. Same CSR shape as moe_pairs_matvec_q8_dec: `table` = [3,n_expert]
309    /// device slab ptrs, CSR ex_ids/ex_off/ex_pairs group pairs by expert, pair_tok gathers the
310    /// activation row. y = [n_pairs, out_f] pair-major. `act_scratch` pre-quantized over n_tokens.
311    /// qtype: 5=IQ4_XS, 6=IQ3_S. Returns 0 or 1000+cudaError.
312    /// Dense-trunk IQ4_XS MMQ (lane/kquant-tile-loaders): the dense analog of the expert
313    /// kernel for non-expert IQ4_XS 2-D matmuls (the KAT-Coder trunk class). Quantizes the
314    /// f32 activation to D4 q8_1_mmq internally; `act_scratch` sized by
315    /// `memra_mmq_iq_experts_act_bytes`. Requires in_f % 256 == 0.
316    pub fn memra_mmq_iq4xs_dense(
317        w_blocks: *const core::ffi::c_void,
318        act_f32: *const f32,
319        y: *mut f32,
320        in_f: i32,
321        out_f: i32,
322        n_tokens: i32,
323        row_bytes: i64,
324        act_scratch: *mut core::ffi::c_void,
325        stream: *mut core::ffi::c_void,
326    ) -> i32;
327    pub fn memra_mmq_iq_experts(
328        table: *const u64,
329        proj: i32,
330        n_expert: i32,
331        ex_ids: *const i32,
332        ex_off: *const i32,
333        ex_pairs: *const i32,
334        pair_tok: *const i32,
335        act_scratch: *const core::ffi::c_void,
336        y: *mut f32,
337        in_f: i32,
338        out_f: i32,
339        n_active: i32,
340        n_tokens: i32,
341        qtype: i32,
342        row_bytes: i64,
343        stream: *mut core::ffi::c_void,
344    ) -> i32;
345
346    // ---- MoE grouped f16 GEMM (cu/moe_f16_grouped.cu, round 46 arc 2) ----
347    pub fn memra_moe_f16g_dequant(
348        table: *const u64, proj: i32, n_expert: i32, ex_ids: *const i32,
349        w_f16: *mut core::ffi::c_void, in_f: i32, out_f: i32, n_active: i32,
350        qtype: i32, row_bytes: i64, stream: *mut core::ffi::c_void,
351    ) -> i32;
352    pub fn memra_moe_f16g_gather_act(
353        x: *const f32, pair_tok_or_null: *const i32, act_f16: *mut core::ffi::c_void,
354        row_scale: *mut f32, in_f: i32, n_pairs: i32, stream: *mut core::ffi::c_void,
355    ) -> i32;
356    pub fn memra_moe_f16g_h2f_scaled(
357        src_f16: *const core::ffi::c_void, dst: *mut f32, row_scale: *const f32,
358        ncols: i32, nrows: i32, stream: *mut core::ffi::c_void,
359    ) -> i32;
360    pub fn memra_moe_f16g_gemm(
361        w_f16: *const core::ffi::c_void, act_f16: *const core::ffi::c_void,
362        y_f16: *mut core::ffi::c_void,
363        ex_off_host: *const i32, n_active: i32, in_f: i32, out_f: i32,
364        stream: *mut core::ffi::c_void,
365    ) -> i32;
366    pub fn memra_moe_f16g_h2f(
367        src_f16: *const core::ffi::c_void, dst: *mut f32, n: usize,
368        stream: *mut core::ffi::c_void,
369    ) -> i32;
370    // Single-kernel grouped GEMM (MEMRA_MOE_F16G=2, rounds 49+51): on OUR stream, f32 C with
371    // the act row-scale folded in — no cublas internal-stream race, no sync. Round 51 runs it
372    // as a persistent problem-visitor over the real tiles with two tile forms (32x64 tail
373    // / 128x64x64 3-stage): shape_sel < 0 = the round-49 grid-scan kernel (rollback
374    // arm); else groups with m_e >= cross ride the 128 form. ex_off_host sizes the visitor
375    // grids host-side (the offsets are already there at the call site — no extra transfer).
376    // tail != 0 (lane/sk-tail-form): sub-cross groups ride the DEEP tail (32x64x64 3-stage);
377    // 0 = the round-51 2-stage 32x64x32 (MEMRA_F16G_TAIL=0 rollback). Byte-identical arms.
378    pub fn memra_moe_f16g_gemm_sk(
379        w_f16: *const core::ffi::c_void, act_f16: *const core::ffi::c_void,
380        y_f32: *mut f32, row_scale: *const f32, ex_off_dev: *const i32,
381        ex_off_host: *const i32,
382        n_active: i32, max_m: i32, in_f: i32, out_f: i32, shape_sel: i32, cross: i32,
383        tail: i32, stream: *mut core::ffi::c_void,
384    ) -> i32;
385    // DIRECT-FROM-QUANT sk visitor grouped GEMM (lane/kquant-tile-loaders + iq-direct-loaders):
386    // the visitor forms with the B (weight) tiles dequanted in-register from the expert
387    // superblocks — no f16 dequant workspace pass. Bit-identical to the workspace path by
388    // construction (kernel-check "f16g-kq-direct"). qtype: QT_Q4_K | QT_Q6_K | QT_IQ4_XS |
389    // QT_IQ3_S; rc=2 = not admitted here (caller keeps the dequant-workspace path).
390    // tail: as memra_moe_f16g_gemm_sk.
391    pub fn memra_moe_kq_gemm_sk(
392        table: *const u64, proj: i32, n_expert: i32, ex_ids: *const i32,
393        act_f16: *const core::ffi::c_void, y_f32: *mut f32,
394        row_scale: *const f32, ex_off_dev: *const i32, ex_off_host: *const i32,
395        n_active: i32, max_m: i32, in_f: i32, out_f: i32, qtype: i32, cross: i32,
396        tail: i32, row_bytes: i64, stream: *mut core::ffi::c_void,
397    ) -> i32;
398}
399
400/// W4A8-MMQ DEFAULT-FLIP seam (2026-07-05): the vendored MMQ prefill suite is DEFAULT-ON — NVFP4
401/// takes the W4A8 MMQ tile (same int8 accuracy class as the int8 GEMM it replaces, all exactness
402/// gates hold, ~1.9x pp512; the rp tile-loader arm coexists with the A6 split-plane repack) and
403/// Q4_K/Q5_K take the vendored k-quant int8-MMA MMQ (also int8-class; gated with W4A8 in the same
404/// battery — the predecessor's `MEMRA_MMQ_W4A8=1` arm engaged BOTH, this flip preserves exactly
405/// that measured config). `MEMRA_MMQ_W4A8=0` = escape hatch back to the int8 GEMM prefill
406/// everywhere. `MEMRA_MMQ=1` additionally switches GGUF-layout NVFP4 to the W4A4 mxf4nvf4 tile
407/// (speed/accuracy tradeoff opt-in, unchanged).
408pub fn mmq_w4a8_enabled() -> bool {
409    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
410    *ON.get_or_init(|| {
411        std::env::var("MEMRA_MMQ_W4A8")
412            .map(|v| v != "0")
413            .unwrap_or(true)
414    })
415}
416
417/// Residual high-precision activation channels for the W4A4 MMQ prefill path.
418/// `MEMRA_MMQ_RESIDUAL_K=<k>` keeps the k largest-magnitude activation channels out of the e2m1
419/// quantized path and adds their exact f32 contribution back as a rank-k correction. k=0 (default)
420/// is off; the kernel clamps to MMQ_MAX_RESIDUAL_K (64).
421///
422/// Read LIVE per call, not OnceLock'd, for the same reason `MEMRA_MMQ` is: the W4A4 exactness gate
423/// sweeps arms inside ONE process against ONE set of loaded weights, and a cached first read would
424/// pin every later arm to whatever the first one saw.
425pub fn mmq_residual_k() -> i32 {
426    std::env::var("MEMRA_MMQ_RESIDUAL_K")
427        .ok()
428        .and_then(|v| v.parse::<i32>().ok())
429        .unwrap_or(0)
430        .clamp(0, 64)
431}
432
433/// Q8_0 MMQ prefill seam (lane/ppmmq lever 2, DEFAULT ON since 2026-07-09 — `MEMRA_PP_Q8MMQ=0`
434/// reverts): routes Q8_0 dense
435/// projections (m>=16) through the vendored int8-MMA MMQ (cu/mmq_q8_0.cu) instead of the hand-rolled
436/// `qmatvec_gemm_q8_0` tiling GEMM. Its own numeric config (MMA f32 reduction order != the tiling
437/// GEMM's) — gated with the full exactness battery. Default OFF until the battery is green.
438pub fn mmq_q8_enabled() -> bool {
439    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
440    // Promotion battery (2026-07-09): argmax MATCH on 35B p1/p2/p3 + 9B p2/p3 (p4-16k OOMs
441    // identically with and without the flag — pre-existing gate capacity limit, not this seam);
442    // kernel-check ALL GREEN; run-spec K=1..8 PASS on 9B+35B. 35B pp 2456->3069 free-clock.
443    *ON.get_or_init(|| {
444        std::env::var("MEMRA_PP_Q8MMQ")
445            .map(|v| v != "0")
446            .unwrap_or(true)
447    })
448}
449
450/// IQ4_XS dense-trunk MMQ prefill seam (lane/kquant-tile-loaders, 2026-08-02): routes
451/// NON-expert IQ4_XS 2-D projections (m>=16) through the vendored-machinery int8-MMA dense
452/// MMQ (cu/mmq_iq_experts.cu `mmq_iq4xs_dense_kernel`) instead of the per-column dp4a grid
453/// — the KAT-Coder prefill wall (0.169x vs llama; zero weight reuse across tokens,
454/// research/kat-anomaly-20260802 §6). Its own numeric config (MMA reduction order) — gated
455/// with the full exactness battery. m=1..15 decode/verify keep dp4a (dispatch parity).
456/// `MEMRA_PP_IQMMQ=0` reverts.
457pub fn mmq_iq4xs_enabled() -> bool {
458    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
459    *ON.get_or_init(|| {
460        std::env::var("MEMRA_PP_IQMMQ")
461            .map(|v| v != "0")
462            .unwrap_or(true)
463    })
464}
465
466/// Q4_0 MMQ prefill seam (gemma-4-12B lane, 2026-07-22): routes Q4_0 dense projections (m>=16)
467/// through the vendored int8-MMA MMQ (cu/mmq_q4_0.cu) instead of the hand-rolled
468/// `qmatvec_gemm_q4_0[_rp]` tiling GEMM (measured 77% of the 12B prime pass). Its own numeric
469/// config (MMA f32 reduction order != the tiling GEMM's) — gated with the full exactness battery
470/// before default-flip; `MEMRA_PP_Q4MMQ=0` reverts.
471pub fn mmq_q4_enabled() -> bool {
472    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
473    *ON.get_or_init(|| {
474        std::env::var("MEMRA_PP_Q4MMQ")
475            .map(|v| v != "0")
476            .unwrap_or(true)
477    })
478}
479
480impl Engine {
481    /// True if `w` should take a vendored MMQ GEMM under the current env policy (see
482    /// `mmq_w4a8_enabled`): NVFP4 needs in_f % 64 == 0, Q4_K/Q5_K need in_f % 256 == 0.
483    pub fn mmq_supports(&self, w: &crate::model::GpuTensor) -> bool {
484        use crate::model::GpuTensor;
485        if crate::portable_mma_gated() {
486            return false;
487        }
488        let mmq_opt_in = std::env::var("MEMRA_MMQ").is_ok();
489        match w {
490            // A6 split-plane repacked NVFP4: ONLY the W4A8 loader has an rp arm (pure address
491            // remap, bit-identical output — mmq_nvfp4_w4a8.cu load_tiles_nvfp4_w4a8<is_rp>).
492            // The W4A4 loader (mmq_fp4.cu load_tiles_nvfp4_nvfp4) reads 36B GGUF blocks only,
493            // so an rp weight with W4A8 disabled falls through to the rp-ported int8 GEMM.
494            // NVFP4 W4A8/W4A4 launchers use .kind::f8f6f4 / mxf4nvf4 tile MMA — sm_100a+/
495            // sm_120a-only. On every portable build (incl. the 90a Hopper-MMA lane) they are
496            // fail-closed link stubs (build.rs), so never offer them here.
497            GpuTensor::Quant { qtype, rp, .. } if *qtype == crate::QT_NVFP4 && *rp => {
498                !cfg!(memra_portable_cuda) && mmq_w4a8_enabled() && w.in_features() % 64 == 0
499            }
500            // GGUF-layout NVFP4 (MEMRA_RP=0): W4A8 (default-on) or the explicit W4A4 opt-in.
501            GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_NVFP4 => {
502                !cfg!(memra_portable_cuda) && (mmq_w4a8_enabled() || mmq_opt_in)
503                    && w.in_features() % 64 == 0
504            }
505            GpuTensor::Quant { qtype, .. }
506                if *qtype == crate::QT_Q4_K || *qtype == crate::QT_Q5_K =>
507            {
508                (mmq_w4a8_enabled() || mmq_opt_in) && w.in_features() % 256 == 0
509            }
510            // Q8_0 dense projections (35B attn/ssm/shexp): opt-in only (MEMRA_PP_Q8MMQ=1), its own
511            // numeric config vs qmatvec_gemm_q8_0. in_f % 256 == 0: MMQ_ITER_K=256 loads 8-block
512            // groups, so a non-multiple row would read a garbage weight tail (fp16 d bytes can be
513            // NaN-pattern, and NaN * 0-padded-activation = NaN — the 26B ffn_down lesson).
514            GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_Q8_0 => {
515                mmq_q8_enabled() && w.in_features() % 256 == 0
516            }
517            // Q4_0 dense projections (gemma QAT ggufs): MEMRA_PP_Q4MMQ seam. Both weight layouts
518            // (raw 18B blocks and the MEMRA_Q4RP split-plane repack) have loader arms. Same
519            // in_f % 256 == 0 tail rule as Q8_0 (26B ffn_down in_f=2112 NaN'd on the %32 gate);
520            // non-multiples fall back to the hand-rolled qmatvec_gemm_q4_0[_rp].
521            GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_Q4_0 => {
522                mmq_q4_enabled() && w.in_features() % 256 == 0
523            }
524            // IQ4_XS dense projections (KAT-Coder trunk): m>=16 prefill only — decode and
525            // spec-verify (m<16) keep the qmatvec_iq4_XS_dp4a per-column program (the
526            // kat-anomaly dispatch-parity law). Requires the dp4a fast path itself enabled:
527            // MEMRA_IQ_FAST=0 (the Stage-A oracle rollback) must also kill this arm so the
528            // rollback stays a full-path seam. in_f % 256: MMQ_ITER_K walks whole superblocks.
529            GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_IQ4_XS => {
530                mmq_iq4xs_enabled() && Self::iq_fast_enabled() && w.in_features() % 256 == 0
531            }
532            _ => false,
533        }
534    }
535
536    /// Unified vendored-MMQ dispatch: routes to the NVFP4 or Q4_K/Q5_K launcher by qtype.
537    /// Caller MUST have checked `mmq_supports(w)`. `x` is the RAW f32 activation.
538    pub fn qmatvec_mmq(
539        &self,
540        w: &crate::model::GpuTensor,
541        x: &CudaSlice<f32>,
542        m: usize,
543    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
544        use crate::model::GpuTensor;
545        let (in_f, out_f) = (w.in_features(), w.out_features());
546        let GpuTensor::Quant {
547            bytes,
548            scale,
549            qtype,
550            rp,
551            ..
552        } = w
553        else {
554            return Err("qmatvec_mmq: not a Quant tensor".into());
555        };
556        // NVFP4 tile choice: W4A8 (accuracy-safe int8 pair, DEFAULT since the flip) vs W4A4
557        // (mxf4nvf4 mma, explicit MEMRA_MMQ=1 speed/accuracy tradeoff). An rp weight ALWAYS takes
558        // W4A8 — only its loader has the split-plane arm (pure address remap, bit-identical).
559        // Explicit MEMRA_MMQ_W4A8=1 still overrides a simultaneous MEMRA_MMQ=1 (predecessor rule).
560        let w4a8_explicit = std::env::var("MEMRA_MMQ_W4A8")
561            .map(|v| v != "0")
562            .unwrap_or(false);
563        let use_w4a8 =
564            *rp || w4a8_explicit || (mmq_w4a8_enabled() && std::env::var("MEMRA_MMQ").is_err());
565        match *qtype {
566            // STAGE 2: the accuracy-safe int8 W4A8 MMQ tile (weight FP4->int8 dequant + q8_1
567            // activation) — handles BOTH weight layouts (rp = A6 split-plane vs GGUF blocks).
568            q if q == crate::QT_NVFP4 && use_w4a8 => {
569                self.qmatvec_mmq_nvfp4_w4a8(bytes, x, m, in_f, out_f, *scale, *rp)
570            }
571            q if q == crate::QT_NVFP4 => self.qmatvec_mmq_nvfp4(bytes, x, m, in_f, out_f, *scale),
572            q if q == crate::QT_Q4_K || q == crate::QT_Q5_K => {
573                let mut y = self.qmatvec_mmq_q45k_raw(bytes, x, m, in_f, out_f, q)?;
574                if *scale != 1.0 {
575                    self.scale_inplace(&mut y, *scale, m * out_f)?;
576                }
577                Ok(y)
578            }
579            q if q == crate::QT_Q8_0 => {
580                // wgmma arm (sm_90a, task 8): OPT-IN via MEMRA_WGMMA=1 — v0 measured 3845
581                // vs MMQ 8692 tok/s pp512 (2026-07-26 N=5), so MMQ stays the default until
582                // the pipelined wgmma wins. Reads the rp4 split-plane mirror + the engine's
583                // q8_1 activation planes. Same numeric class as MMQ (exact s32 per 32-block,
584                // one f32 fold per block, ascending K) — kernel-check tolerance-gated.
585                if cfg!(memra_hopper_mma) && out_f % 64 == 0 && crate::wgmma_gemm_enabled() {
586                    if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
587                        let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
588                        let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, &aq, &ad, m, in_f, out_f)?;
589                        if *scale != 1.0 {
590                            self.scale_inplace(&mut y, *scale, m * out_f)?;
591                        }
592                        return Ok(y);
593                    }
594                }
595                let mut y = self.qmatvec_mmq_q8_0_raw(bytes, x, m, in_f, out_f)?;
596                if *scale != 1.0 {
597                    self.scale_inplace(&mut y, *scale, m * out_f)?;
598                }
599                Ok(y)
600            }
601            q if q == crate::QT_Q4_0 => {
602                let mut y = self.qmatvec_mmq_q4_0_raw(bytes, x, m, in_f, out_f, *rp)?;
603                if *scale != 1.0 {
604                    self.scale_inplace(&mut y, *scale, m * out_f)?;
605                }
606                Ok(y)
607            }
608            q if q == crate::QT_IQ4_XS => {
609                let GpuTensor::Quant { row_bytes, .. } = w else { unreachable!() };
610                let mut y = self.qmatvec_mmq_iq4xs_raw(bytes, x, m, in_f, out_f, *row_bytes)?;
611                if *scale != 1.0 {
612                    self.scale_inplace(&mut y, *scale, m * out_f)?;
613                }
614                Ok(y)
615            }
616            q => Err(format!("qmatvec_mmq: unsupported qtype {q}").into()),
617        }
618    }
619
620    /// Bare IQ4_XS dense MMQ launch (no macro-scale) — also the kernel_check gate entry.
621    pub fn qmatvec_mmq_iq4xs_raw(
622        &self,
623        bytes: &CudaSlice<u8>,
624        x: &CudaSlice<f32>,
625        m: usize,
626        in_f: usize,
627        out_f: usize,
628        row_bytes: usize,
629    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
630        assert!(
631            in_f % 256 == 0,
632            "MMQ IQ4_XS requires in_f % 256 == 0, got {in_f}"
633        );
634        let act_bytes = unsafe { memra_mmq_iq_experts_act_bytes(in_f as i32, m as i32) };
635        let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
636        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
637        {
638            let stream = self.gpu.stream();
639            let (w_p, _gw) = bytes.device_ptr(&stream);
640            let (x_p, _gx) = x.device_ptr(&stream);
641            let (y_p, _gy) = y.device_ptr_mut(&stream);
642            let (s_p, _gs) = scratch.device_ptr_mut(&stream);
643            let rc = unsafe {
644                memra_mmq_iq4xs_dense(
645                    w_p as *const core::ffi::c_void,
646                    x_p as *const f32,
647                    y_p as *mut f32,
648                    in_f as i32,
649                    out_f as i32,
650                    m as i32,
651                    row_bytes as i64,
652                    s_p as *mut core::ffi::c_void,
653                    stream.cu_stream() as *mut core::ffi::c_void,
654                )
655            };
656            if rc != 0 {
657                return Err(format!("memra_mmq_iq4xs_dense rc={rc}").into());
658            }
659        }
660        Ok(y)
661    }
662
663    /// Bare Q4_K/Q5_K MMQ launch (no macro-scale) — also the kernel_check accuracy-gate entry.
664    /// Conventional xy-tiling only (the vendored stream-K arm — MEMRA_MMQ_STREAMK — was removed
665    /// 2026-07-08: 1.11x per-GEMM but its k-split f32 reorder flipped the model argmax gate;
666    /// rig5090.jsonl 2026-07-03 has the record).
667    pub fn qmatvec_mmq_q45k_raw(
668        &self,
669        bytes: &CudaSlice<u8>,
670        x: &CudaSlice<f32>,
671        m: usize,
672        in_f: usize,
673        out_f: usize,
674        qtype: i32,
675    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
676        assert!(
677            in_f % 256 == 0,
678            "MMQ Q4_K/Q5_K requires in_f % 256 == 0, got {in_f}"
679        );
680        let act_bytes = unsafe { memra_mmq_q45k_act_bytes(in_f as i32, m as i32) };
681        let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
682        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
683        {
684            let stream = self.gpu.stream();
685            let (w_p, _gw) = bytes.device_ptr(&stream);
686            let (x_p, _gx) = x.device_ptr(&stream);
687            let (y_p, _gy) = y.device_ptr_mut(&stream);
688            let (s_p, _gs) = scratch.device_ptr_mut(&stream);
689            let launcher = if qtype == crate::QT_Q4_K {
690                memra_mmq_q4_K
691            } else {
692                memra_mmq_q5_K
693            };
694            let rc = unsafe {
695                launcher(
696                    w_p as *const core::ffi::c_void,
697                    x_p as *const f32,
698                    y_p as *mut f32,
699                    in_f as i32,
700                    out_f as i32,
701                    m as i32,
702                    s_p as *mut core::ffi::c_void,
703                    stream.cu_stream() as *mut core::ffi::c_void,
704                )
705            };
706            if rc != 0 {
707                return Err(format!("memra_mmq_q45k(qtype={qtype}) rc={rc}").into());
708            }
709        }
710        Ok(y)
711    }
712
713    /// Bare Q8_0 int8-MMA MMQ launch (no macro-scale) — the kernel_check accuracy-gate entry and
714    /// the `qmatvec_mmq` dispatch body. Conventional xy-tiling only (no stream-K / fixup scratch).
715    pub fn qmatvec_mmq_q8_0_raw(
716        &self,
717        bytes: &CudaSlice<u8>,
718        x: &CudaSlice<f32>,
719        m: usize,
720        in_f: usize,
721        out_f: usize,
722    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
723        assert!(
724            in_f % 32 == 0,
725            "MMQ Q8_0 requires in_f % 32 == 0, got {in_f}"
726        );
727        let act_bytes = unsafe { memra_mmq_q8_0_act_bytes(in_f as i32, m as i32) };
728        let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
729        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
730        {
731            let stream = self.gpu.stream();
732            let (w_p, _gw) = bytes.device_ptr(&stream);
733            let (x_p, _gx) = x.device_ptr(&stream);
734            let (y_p, _gy) = y.device_ptr_mut(&stream);
735            let (s_p, _gs) = scratch.device_ptr_mut(&stream);
736            let rc = unsafe {
737                memra_mmq_q8_0(
738                    w_p as *const core::ffi::c_void,
739                    x_p as *const f32,
740                    y_p as *mut f32,
741                    in_f as i32,
742                    out_f as i32,
743                    m as i32,
744                    s_p as *mut core::ffi::c_void,
745                    stream.cu_stream() as *mut core::ffi::c_void,
746                )
747            };
748            if rc != 0 {
749                return Err(format!("memra_mmq_q8_0 rc={rc}").into());
750            }
751        }
752        Ok(y)
753    }
754
755    /// Accumulator-instrument bytes for a pre-quantized block_q8_1_mmq activation buffer
756    /// (cu/mmq_q8_0_f32acc.cu). The caller synthesizes that buffer itself — see `accprobe_gemm`.
757    pub fn accprobe_act_bytes(&self, in_f: usize, m: usize) -> usize {
758        unsafe { memra_accprobe_act_bytes(in_f as i32, m as i32) }
759    }
760
761    /// Run one arm of the Q1 accumulator instrument. `f32acc=false` is the Q8_0 MMQ floor's GEMM
762    /// verbatim (s32 accumulate); `f32acc=true` is the byte-identical kernel with the f8f6f4 f32
763    /// accumulate. `act_q` is a PRE-QUANTIZED block_q8_1_mmq buffer of at least
764    /// `accprobe_act_bytes(in_f, m)` bytes — keeping the quantizer out of the timed region is the
765    /// point, so this wrapper does not build it. Research instrument: the output is not a numeric
766    /// claim.
767    pub fn accprobe_gemm(
768        &self,
769        w_q8_0: &CudaSlice<u8>,
770        act_q: &CudaSlice<u8>,
771        m: usize,
772        in_f: usize,
773        out_f: usize,
774        f32acc: bool,
775    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
776        assert!(in_f % 32 == 0, "accprobe requires in_f % 32 == 0, got {in_f}");
777        assert!(
778            act_q.len() >= self.accprobe_act_bytes(in_f, m),
779            "accprobe act_q too small: {} < {}",
780            act_q.len(),
781            self.accprobe_act_bytes(in_f, m)
782        );
783        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
784        {
785            let stream = self.gpu.stream();
786            let (w_p, _gw) = w_q8_0.device_ptr(&stream);
787            let (a_p, _ga) = act_q.device_ptr(&stream);
788            let (y_p, _gy) = y.device_ptr_mut(&stream);
789            let f = if f32acc {
790                memra_accprobe_gemm_f32
791            } else {
792                memra_accprobe_gemm_s32
793            };
794            let rc = unsafe {
795                f(
796                    w_p as *const core::ffi::c_void,
797                    a_p as *const core::ffi::c_void,
798                    y_p as *mut f32,
799                    in_f as i32,
800                    out_f as i32,
801                    m as i32,
802                    stream.cu_stream() as *mut core::ffi::c_void,
803                )
804            };
805            if rc != 0 {
806                let arm = if f32acc { "f32" } else { "s32" };
807                return Err(format!("memra_accprobe_gemm_{arm} rc={rc}").into());
808            }
809        }
810        Ok(y)
811    }
812
813    /// Open a quantize-once sharing window for the NEXT activation (quantize-once seam): sibling
814    /// Q4_0 MMQ matmuls on the SAME input (q/k/v; gate/up) quantize its D4 scratch once. Safe by
815    /// construction: a hit requires the same window epoch AND the same (ptr, m, in_f) — the caller
816    /// opens a window while it holds the shared input alive, so its address can neither change nor
817    /// be recycled inside the window. Paths that never call this never hit the cache.
818    pub fn mmq_act_begin(&self) {
819        use std::sync::atomic::Ordering;
820        MMQ_ACT_EPOCH.fetch_add(1, Ordering::Relaxed);
821        *MMQ_ACT_SLOT.lock().unwrap() = None;
822    }
823
824    /// Bare Q4_0 int8-MMA MMQ launch (no macro-scale) — the kernel_check accuracy-gate entry and
825    /// the `qmatvec_mmq` dispatch body. `rp` selects the weight layout (MEMRA_Q4RP split-plane vs
826    /// raw ggml 18B blocks) — pure address remap, bit-identical output.
827    pub fn qmatvec_mmq_q4_0_raw(
828        &self,
829        bytes: &CudaSlice<u8>,
830        x: &CudaSlice<f32>,
831        m: usize,
832        in_f: usize,
833        out_f: usize,
834        rp: bool,
835    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
836        use std::sync::atomic::Ordering;
837        assert!(
838            in_f % 32 == 0,
839            "MMQ Q4_0 requires in_f % 32 == 0, got {in_f}"
840        );
841        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
842        let stream = self.gpu.stream();
843        let (x_p, _gx) = x.device_ptr(&stream);
844        let epoch = MMQ_ACT_EPOCH.load(Ordering::Relaxed);
845        // quantize-once: reuse the window's scratch when the SAME activation comes back.
846        let mut slot = MMQ_ACT_SLOT.lock().unwrap();
847        let hit = matches!(&*slot,
848            Some((e, p, mm, inf, _)) if *e == epoch && *p == x_p as u64 && *mm == m && *inf == in_f);
849        if !hit {
850            let act_bytes = unsafe { memra_mmq_q4_0_act_bytes(in_f as i32, m as i32) };
851            let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
852            {
853                let (s_p, _gs) = scratch.device_ptr_mut(&stream);
854                let rc = unsafe {
855                    memra_mmq_q4_0_quant_act(
856                        x_p as *const f32,
857                        s_p as *mut core::ffi::c_void,
858                        in_f as i32,
859                        m as i32,
860                        stream.cu_stream() as *mut core::ffi::c_void,
861                    )
862                };
863                if rc != 0 {
864                    return Err(
865                        format!("memra_mmq_q4_0_quant_act(in_f={in_f}, m={m}) rc={rc}").into()
866                    );
867                }
868            }
869            *slot = Some((epoch, x_p as u64, m, in_f, scratch));
870        }
871        let scratch = &slot.as_ref().unwrap().4;
872        {
873            let (w_p, _gw) = bytes.device_ptr(&stream);
874            let (y_p, _gy) = y.device_ptr_mut(&stream);
875            let (s_p, _gs) = scratch.device_ptr(&stream);
876            // Stream-k arm (DEFAULT since 2026-07-23; MEMRA_MMQ_SK=0 reverts to xy-tiling):
877            // small-batch tail-wave fix — the sk entry itself falls back to (bit-identical)
878            // tiling at >=90% wave efficiency. Band-class fold order below that. Gate: 12B
879            // pp512 +3.3% (1.005x vs llama), pp1736 +1.0%; 31B +0.5%; D512 sentinel MATCH.
880            //
881            // SPEC-SERVING FLIP (2026-07-27, the f16pv/wkv acceptance-law pattern): with
882            // MEMRA_DRAFT set the default is OFF. Two mechanisms, both measured: (1) sk's
883            // fold order shifts PRIME numerics in the drafter's sub-argmax logit space
884            // (26B d1736 accept 0.846 w/o sk vs 0.826 — bisect row 0.861 vs 0.405-era);
885            // (2) the shape-keyed AUTOTUNE re-times both forms per process, so knife-edge
886            // shapes pick DIFFERENT kernels run-to-run — the 12B depth cell was BIMODAL
887            // (205 @ 0.756 / 260 @ 0.943 across identical invocations; sk-off x6 = stable
888            // 263-269 @ 0.953). Explicit MEMRA_MMQ_SK always wins; plain serving keeps sk.
889            // HOPPER DEFAULT OFF (2026-07-31, #23): on sm_90a the SK arm computes WRONG
890            // values for the 26B a4b's non-rp Q4_0 shapes once the autotune m-bucket
891            // crosses 256 (prefill argmax garbage, maxdiff ~10; MEMRA_MMQ_SK=0 -> MATCH,
892            // one-variable kill x confirmed on-box). The winner is timing-picked and the
893            // SK split/fixup are SM-count dependent (132 vs 170) — until the kernel is
894            // fixed for that class, Hopper fails CLOSED to the bit-identical xy-tiling
895            // (cost on the healthy models: g12 -1.4%, g31 -0.6% prefill, N=3 on-box).
896            // sm_120a keeps SK on (rig-divergence law). MEMRA_MMQ_SK=1 still forces.
897            static SK_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
898            let sk = match crate::MMQ_SK_FORCE.load(std::sync::atomic::Ordering::Relaxed) {
899                0 => false,
900                1 => true,
901                _ => *SK_ON.get_or_init(|| std::env::var("MEMRA_MMQ_SK")
902                        .map(|v| v != "0").unwrap_or(!cfg!(memra_hopper_mma))),
903            };
904            let rc = if sk {
905                let mut fx = MMQ_FIXUP_SLOT.lock().unwrap();
906                if fx.is_none() {
907                    let nb = unsafe { memra_mmq_q4_0_fixup_bytes() };
908                    *fx = Some(self.alloc_uninit::<u8>(nb)?);
909                }
910                let (f_p, _gf) = fx.as_mut().unwrap().device_ptr_mut(&stream);
911                unsafe {
912                    memra_mmq_q4_0_gemm_sk(
913                        w_p as *const core::ffi::c_void,
914                        s_p as *const core::ffi::c_void,
915                        y_p as *mut f32,
916                        f_p as *mut core::ffi::c_void,
917                        in_f as i32,
918                        out_f as i32,
919                        m as i32,
920                        stream.cu_stream() as *mut core::ffi::c_void,
921                        rp as i32,
922                    )
923                }
924            } else { unsafe {
925                memra_mmq_q4_0_gemm(
926                    w_p as *const core::ffi::c_void,
927                    s_p as *const core::ffi::c_void,
928                    y_p as *mut f32,
929                    in_f as i32,
930                    out_f as i32,
931                    m as i32,
932                    stream.cu_stream() as *mut core::ffi::c_void,
933                    rp as i32,
934                )
935            } };
936            if rc != 0 {
937                return Err(format!(
938                    "memra_mmq_q4_0_gemm(rp={rp}, in_f={in_f}, out_f={out_f}, m={m}, wbytes={}) rc={rc}",
939                    bytes.len()
940                )
941                .into());
942            }
943        }
944        Ok(y)
945    }
946
947    /// Run the vendored NVFP4 MMQ prefill GEMM from raw weight bytes + f32 activation.
948    /// y[m, out_f] = x[m, in_f] @ W^T. The per-tensor NVFP4 macro-scale is FOLDED into the MMQ
949    /// write-back epilogue (was a separate scale_inplace launch + full y round-trip per matmul).
950    /// Same elementwise multiply -> bit-identical to the two-launch form.
951    /// `x` is the RAW f32 activation (the launcher quantizes it to block_fp4_mmq internally).
952    pub fn qmatvec_mmq_nvfp4(
953        &self,
954        bytes: &CudaSlice<u8>,
955        x: &CudaSlice<f32>,
956        m: usize,
957        in_f: usize,
958        out_f: usize,
959        scale: f32,
960    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
961        self.qmatvec_mmq_nvfp4_scaled(bytes, x, m, in_f, out_f, scale)
962    }
963
964    /// Bare MMQ launch (no macro-scale) — for the kernel_check accuracy gate.
965    pub fn qmatvec_mmq_nvfp4_raw(
966        &self,
967        bytes: &CudaSlice<u8>,
968        x: &CudaSlice<f32>,
969        m: usize,
970        in_f: usize,
971        out_f: usize,
972    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
973        self.qmatvec_mmq_nvfp4_scaled(bytes, x, m, in_f, out_f, 1.0)
974    }
975
976    /// Bare MMQ launch on the PRE-PORT activation quantizer (per-sub-block UE4M3 scale only, no
977    /// per-token row amax). The numeric oracle for the two-level quantizer: kernel-check runs both
978    /// and reports the accuracy delta, so the port's value is measured rather than asserted.
979    pub fn qmatvec_mmq_nvfp4_raw_v1(
980        &self,
981        bytes: &CudaSlice<u8>,
982        x: &CudaSlice<f32>,
983        m: usize,
984        in_f: usize,
985        out_f: usize,
986    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
987        self.qmatvec_mmq_nvfp4_inner(bytes, x, m, in_f, out_f, 1.0, false, 0)
988    }
989
990    /// Bare MMQ launch with an explicit residual-channel count — for the kernel-check k sweep.
991    pub fn qmatvec_mmq_nvfp4_raw_res(
992        &self,
993        bytes: &CudaSlice<u8>,
994        x: &CudaSlice<f32>,
995        m: usize,
996        in_f: usize,
997        out_f: usize,
998        residual_k: i32,
999    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1000        self.qmatvec_mmq_nvfp4_inner(bytes, x, m, in_f, out_f, 1.0, true, residual_k)
1001    }
1002
1003    fn qmatvec_mmq_nvfp4_scaled(
1004        &self,
1005        bytes: &CudaSlice<u8>,
1006        x: &CudaSlice<f32>,
1007        m: usize,
1008        in_f: usize,
1009        out_f: usize,
1010        scale: f32,
1011    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1012        self.qmatvec_mmq_nvfp4_inner(bytes, x, m, in_f, out_f, scale, true, mmq_residual_k())
1013    }
1014
1015    fn qmatvec_mmq_nvfp4_inner(
1016        &self,
1017        bytes: &CudaSlice<u8>,
1018        x: &CudaSlice<f32>,
1019        m: usize,
1020        in_f: usize,
1021        out_f: usize,
1022        scale: f32,
1023        per_token_scale: bool,
1024        residual_k: i32,
1025    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1026        assert!(
1027            in_f % 64 == 0,
1028            "MMQ NVFP4 requires in_f % 64 == 0, got {in_f}"
1029        );
1030        let act_bytes = unsafe { memra_mmq_nvfp4_act_bytes(in_f as i32, m as i32) };
1031        let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1032        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
1033        {
1034            let stream = self.gpu.stream();
1035            let (w_p, _gw) = bytes.device_ptr(&stream);
1036            let (x_p, _gx) = x.device_ptr(&stream);
1037            let (y_p, _gy) = y.device_ptr_mut(&stream);
1038            let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1039            let rc = unsafe {
1040                memra_mmq_nvfp4_ex2(
1041                    w_p as *const core::ffi::c_void,
1042                    x_p as *const f32,
1043                    y_p as *mut f32,
1044                    in_f as i32,
1045                    out_f as i32,
1046                    m as i32,
1047                    s_p as *mut core::ffi::c_void,
1048                    stream.cu_stream() as *mut core::ffi::c_void,
1049                    scale,
1050                    per_token_scale as i32,
1051                    residual_k,
1052                )
1053            };
1054            if rc != 0 {
1055                return Err(format!("memra_mmq_nvfp4_ex2 rc={rc}").into());
1056            }
1057        }
1058        Ok(y)
1059    }
1060
1061    /// STAGE 2 W4A8 MMQ NVFP4: same tile as the W4A4 path, but weight FP4 is LUT-dequantized to
1062    /// int8 at tile-load and the activation stays q8_1 int8 — the accuracy-safe rung. Macro-scale
1063    /// folded into the write-back epilogue (bit-identical to a post-matmul scale_inplace).
1064    /// `rp` selects the weight layout (A6 split-plane vs GGUF blocks) — bit-identical output.
1065    pub fn qmatvec_mmq_nvfp4_w4a8(
1066        &self,
1067        bytes: &CudaSlice<u8>,
1068        x: &CudaSlice<f32>,
1069        m: usize,
1070        in_f: usize,
1071        out_f: usize,
1072        scale: f32,
1073        rp: bool,
1074    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1075        self.qmatvec_mmq_nvfp4_w4a8_scaled(bytes, x, m, in_f, out_f, scale, rp)
1076    }
1077
1078    /// Bare W4A8 MMQ launch (no macro-scale, GGUF layout) — for the kernel_check accuracy gate.
1079    pub fn qmatvec_mmq_nvfp4_w4a8_raw(
1080        &self,
1081        bytes: &CudaSlice<u8>,
1082        x: &CudaSlice<f32>,
1083        m: usize,
1084        in_f: usize,
1085        out_f: usize,
1086    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1087        self.qmatvec_mmq_nvfp4_w4a8_scaled(bytes, x, m, in_f, out_f, 1.0, false)
1088    }
1089
1090    /// Bare W4A8 MMQ launch on an A6 split-plane repacked weight — the rp-loader bit-identity gate
1091    /// compares this against `qmatvec_mmq_nvfp4_w4a8_raw` on the same weight.
1092    pub fn qmatvec_mmq_nvfp4_w4a8_raw_rp(
1093        &self,
1094        bytes: &CudaSlice<u8>,
1095        x: &CudaSlice<f32>,
1096        m: usize,
1097        in_f: usize,
1098        out_f: usize,
1099    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1100        self.qmatvec_mmq_nvfp4_w4a8_scaled(bytes, x, m, in_f, out_f, 1.0, true)
1101    }
1102
1103    fn qmatvec_mmq_nvfp4_w4a8_scaled(
1104        &self,
1105        bytes: &CudaSlice<u8>,
1106        x: &CudaSlice<f32>,
1107        m: usize,
1108        in_f: usize,
1109        out_f: usize,
1110        scale: f32,
1111        rp: bool,
1112    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1113        assert!(
1114            in_f % 64 == 0,
1115            "MMQ NVFP4 W4A8 requires in_f % 64 == 0, got {in_f}"
1116        );
1117        let act_bytes = unsafe { memra_mmq_nvfp4_w4a8_act_bytes(in_f as i32, m as i32) };
1118        let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1119        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
1120        {
1121            let stream = self.gpu.stream();
1122            let (w_p, _gw) = bytes.device_ptr(&stream);
1123            let (x_p, _gx) = x.device_ptr(&stream);
1124            let (y_p, _gy) = y.device_ptr_mut(&stream);
1125            let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1126            // MEMRA_MMQ_F8F4=1: the R-B W4A8-FP8 tile (own numeric config; battery-gated seam).
1127            // Scratch layouts are footprint-identical, so only the entry point swaps.
1128            static F8F4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1129            let f8f4 = *F8F4.get_or_init(|| std::env::var("MEMRA_MMQ_F8F4").as_deref() == Ok("1"));
1130            let rc = unsafe {
1131                if f8f4 {
1132                    memra_mmq_nvfp4_f8f4(
1133                        w_p as *const core::ffi::c_void,
1134                        x_p as *const f32,
1135                        y_p as *mut f32,
1136                        in_f as i32,
1137                        out_f as i32,
1138                        m as i32,
1139                        s_p as *mut core::ffi::c_void,
1140                        stream.cu_stream() as *mut core::ffi::c_void,
1141                        scale,
1142                        rp as i32,
1143                    )
1144                } else {
1145                    memra_mmq_nvfp4_w4a8(
1146                        w_p as *const core::ffi::c_void,
1147                        x_p as *const f32,
1148                        y_p as *mut f32,
1149                        in_f as i32,
1150                        out_f as i32,
1151                        m as i32,
1152                        s_p as *mut core::ffi::c_void,
1153                        stream.cu_stream() as *mut core::ffi::c_void,
1154                        scale,
1155                        rp as i32,
1156                    )
1157                }
1158            };
1159            if rc != 0 {
1160                return Err(format!("memra_mmq_nvfp4_w4a8(f8f4={f8f4}) rc={rc}").into());
1161            }
1162        }
1163        Ok(y)
1164    }
1165
1166    /// PER-BLOCK FP8 MMQ prefill GEMM (cu/mmq_fp8_blk.cu). `w_e4m3` is the raw checkpoint e4m3
1167    /// plane [out_f x in_f] and `blk_scales` the device f32 grid [ceil(out_f/128) x
1168    /// ceil(in_f/128)] — no re-quantization of either.
1169    pub fn qmatvec_mmq_fp8_blk(
1170        &self,
1171        w_e4m3: &CudaSlice<u8>,
1172        blk_scales: &CudaSlice<f32>,
1173        x: &CudaSlice<f32>,
1174        m: usize,
1175        in_f: usize,
1176        out_f: usize,
1177    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1178        self.qmatvec_mmq_fp8_blk_scaled(w_e4m3, blk_scales, x, m, in_f, out_f, 1.0)
1179    }
1180
1181    pub fn qmatvec_mmq_fp8_blk_scaled(
1182        &self,
1183        w_e4m3: &CudaSlice<u8>,
1184        blk_scales: &CudaSlice<f32>,
1185        x: &CudaSlice<f32>,
1186        m: usize,
1187        in_f: usize,
1188        out_f: usize,
1189        scale: f32,
1190    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1191        assert!(
1192            in_f % 16 == 0,
1193            "per-block FP8 MMQ requires in_f % 16 == 0, got {in_f}"
1194        );
1195        let want_scales = ((out_f + 127) / 128) * ((in_f + 127) / 128);
1196        assert!(
1197            blk_scales.len() >= want_scales,
1198            "blk_scales too small: {} < {want_scales}",
1199            blk_scales.len()
1200        );
1201        assert!(
1202            w_e4m3.len() >= out_f * in_f,
1203            "e4m3 plane too small: {} < {}",
1204            w_e4m3.len(),
1205            out_f * in_f
1206        );
1207        let act_bytes = unsafe { memra_mmq_fp8_blk_act_bytes(in_f as i32, m as i32) };
1208        let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1209        let mut y = self.alloc_uninit::<f32>(m * out_f)?;
1210        {
1211            let stream = self.gpu.stream();
1212            let (w_p, _gw) = w_e4m3.device_ptr(&stream);
1213            let (sc_p, _gsc) = blk_scales.device_ptr(&stream);
1214            let (x_p, _gx) = x.device_ptr(&stream);
1215            let (y_p, _gy) = y.device_ptr_mut(&stream);
1216            let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1217            let rc = unsafe {
1218                memra_mmq_fp8_blk(
1219                    w_p as *const core::ffi::c_void,
1220                    sc_p as *const f32,
1221                    x_p as *const f32,
1222                    y_p as *mut f32,
1223                    in_f as i32,
1224                    out_f as i32,
1225                    m as i32,
1226                    s_p as *mut core::ffi::c_void,
1227                    stream.cu_stream() as *mut core::ffi::c_void,
1228                    scale,
1229                )
1230            };
1231            if rc != 0 {
1232                return Err(format!("memra_mmq_fp8_blk rc={rc}").into());
1233            }
1234        }
1235        Ok(y)
1236    }
1237
1238    /// Count e4m3 NaN codes (magnitude 0x7F) in a device e4m3 plane. 0 is the precondition for
1239    /// routing that tensor through `qmatvec_mmq_fp8_blk` (hardware decodes them to NaN, the
1240    /// host/ARM B' reference to 0.0).
1241    pub fn fp8_blk_nan_count(
1242        &self,
1243        w_e4m3: &CudaSlice<u8>,
1244    ) -> Result<u32, Box<dyn std::error::Error>> {
1245        let mut cnt = self.htod_u32_v(&[0u32])?;
1246        let n = w_e4m3.len();
1247        {
1248            let stream = self.gpu.stream();
1249            let (w_p, _gw) = w_e4m3.device_ptr(&stream);
1250            let (c_p, _gc) = cnt.device_ptr_mut(&stream);
1251            let rc = unsafe {
1252                memra_fp8_blk_count_nan(
1253                    w_p as *const core::ffi::c_void,
1254                    n,
1255                    c_p as *mut u32,
1256                    stream.cu_stream() as *mut core::ffi::c_void,
1257                )
1258            };
1259            if rc != 0 {
1260                return Err(format!("memra_fp8_blk_count_nan rc={rc}").into());
1261            }
1262        }
1263        Ok(self.dtoh_u32(&cnt)?[0])
1264    }
1265
1266    /// Quantize token-major f32 activation [n_tokens, in_f] to the block_q8_1_mmq (D4) scratch the
1267    /// IQ expert-MMA kernel consumes. Returns the scratch buffer (one per proj input per layer).
1268    pub fn mmq_iq_quantize_act(
1269        &self,
1270        x: &CudaSlice<f32>,
1271        in_f: usize,
1272        n_tokens: usize,
1273    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1274        let act_bytes = unsafe { memra_mmq_iq_experts_act_bytes(in_f as i32, n_tokens as i32) };
1275        let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1276        {
1277            let stream = self.gpu.stream();
1278            let (x_p, _gx) = x.device_ptr(&stream);
1279            let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1280            let rc = unsafe {
1281                memra_mmq_iq_quantize_act(
1282                    x_p as *const f32,
1283                    s_p as *mut core::ffi::c_void,
1284                    in_f as i32,
1285                    n_tokens as i32,
1286                    stream.cu_stream() as *mut core::ffi::c_void,
1287                )
1288            };
1289            if rc != 0 {
1290                return Err(format!("memra_mmq_iq_quantize_act rc={rc}").into());
1291            }
1292        }
1293        Ok(scratch)
1294    }
1295
1296    /// Fused act-epilogue (research lever #3): silu/gelu(gate)*up + D4 quantize in one launch —
1297    /// replaces moe_pairs_{silu,gelu}_mul + mmq_iq_quantize_act without materializing the f32 act
1298    /// buffer (saves one full write + one full read pass over [n_pairs x n_ff]). Scratch bytes are
1299    /// BYTE-IDENTICAL to the two-pass path (kernel-check `iq fused act+quant` gates it).
1300    /// `act_kind`: 0 = silu*mul (qwen35moe), 1 = gelu_tanh*mul (gemma4).
1301    pub fn mmq_iq_fused_act_quant(
1302        &self,
1303        gate: &CudaSlice<f32>,
1304        up: &CudaSlice<f32>,
1305        in_f: usize,
1306        n_tokens: usize,
1307        act_kind: i32,
1308    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1309        let act_bytes = unsafe { memra_mmq_iq_experts_act_bytes(in_f as i32, n_tokens as i32) };
1310        let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1311        {
1312            let stream = self.gpu.stream();
1313            let (g_p, _gg) = gate.device_ptr(&stream);
1314            let (u_p, _gu) = up.device_ptr(&stream);
1315            let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1316            let rc = unsafe {
1317                memra_mmq_iq_fused_act_quant(
1318                    g_p as *const f32,
1319                    u_p as *const f32,
1320                    s_p as *mut core::ffi::c_void,
1321                    in_f as i32,
1322                    n_tokens as i32,
1323                    act_kind,
1324                    stream.cu_stream() as *mut core::ffi::c_void,
1325                )
1326            };
1327            if rc != 0 {
1328                return Err(format!("memra_mmq_iq_fused_act_quant rc={rc}").into());
1329            }
1330        }
1331        Ok(scratch)
1332    }
1333
1334    /// Expert-segmented IQ3_S/IQ4_XS int8-MMA MMQ (the m16n8k16.s8 analog of moe_pairs_matvec_q8_dec).
1335    /// Same CSR inputs (table/ex_ids/ex_off/ex_pairs/pair_tok) + a pre-quantized q8_1_mmq activation
1336    /// scratch (from `mmq_iq_quantize_act` over n_tokens). y = [n_pairs, out_f] pair-major.
1337    #[allow(clippy::too_many_arguments)]
1338    pub fn mmq_iq_experts(
1339        &self,
1340        table: &CudaSlice<u64>,
1341        proj: i32,
1342        n_expert: usize,
1343        ex_ids: &CudaSlice<i32>,
1344        ex_off: &CudaSlice<i32>,
1345        ex_pairs: &CudaSlice<i32>,
1346        pair_tok: &CudaSlice<i32>,
1347        act_scratch: &CudaSlice<u8>,
1348        in_f: usize,
1349        out_f: usize,
1350        n_active: usize,
1351        n_pairs: usize,
1352        n_tokens: usize,
1353        qtype: i32,
1354        row_bytes: usize,
1355    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1356        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1357        {
1358            let stream = self.gpu.stream();
1359            let (tab_p, _g0) = table.device_ptr(&stream);
1360            let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1361            let (eo_p, _g2) = ex_off.device_ptr(&stream);
1362            let (ep_p, _g3) = ex_pairs.device_ptr(&stream);
1363            let (pt_p, _g4) = pair_tok.device_ptr(&stream);
1364            let (as_p, _g5) = act_scratch.device_ptr(&stream);
1365            let (y_p, _g6) = y.device_ptr_mut(&stream);
1366            let rc = unsafe {
1367                memra_mmq_iq_experts(
1368                    tab_p as *const u64,
1369                    proj,
1370                    n_expert as i32,
1371                    ei_p as *const i32,
1372                    eo_p as *const i32,
1373                    ep_p as *const i32,
1374                    pt_p as *const i32,
1375                    as_p as *const core::ffi::c_void,
1376                    y_p as *mut f32,
1377                    in_f as i32,
1378                    out_f as i32,
1379                    n_active as i32,
1380                    n_tokens as i32,
1381                    qtype,
1382                    row_bytes as i64,
1383                    stream.cu_stream() as *mut core::ffi::c_void,
1384                )
1385            };
1386            if rc != 0 {
1387                return Err(format!("memra_mmq_iq_experts rc={rc}").into());
1388            }
1389        }
1390        Ok(y)
1391    }
1392
1393    /// Gather+convert the activation to f16 pair-major [n_pairs, in_f] for the grouped
1394    /// GEMM, normalized per row by its amax (raw f16 overflows on gemma's activation
1395    /// spikes — round 46 NaN find). Returns (act_f16, row_scales) — the scales fold back
1396    /// into the GEMM output. `pair_tok` = None when the input is already pair-major.
1397    pub fn moe_f16g_act(
1398        &self,
1399        x: &CudaSlice<f32>,
1400        pair_tok: Option<&CudaSlice<i32>>,
1401        in_f: usize,
1402        n_pairs: usize,
1403    ) -> Result<(CudaSlice<u8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1404        let mut act = self.alloc_uninit::<u8>(n_pairs * in_f * 2)?;
1405        let mut scales = self.alloc_uninit::<f32>(n_pairs)?;
1406        {
1407            let stream = self.gpu.stream();
1408            let (x_p, _gx) = x.device_ptr(&stream);
1409            let pt_p = match pair_tok {
1410                Some(pt) => { let (p, _g) = pt.device_ptr(&stream); p as *const i32 }
1411                None => std::ptr::null(),
1412            };
1413            let (a_p, _ga) = act.device_ptr_mut(&stream);
1414            let (s_p, _gs) = scales.device_ptr_mut(&stream);
1415            let rc = unsafe {
1416                memra_moe_f16g_gather_act(x_p as *const f32, pt_p,
1417                    a_p as *mut core::ffi::c_void, s_p as *mut f32,
1418                    in_f as i32, n_pairs as i32,
1419                    stream.cu_stream() as *mut core::ffi::c_void)
1420            };
1421            if rc != 0 { return Err(format!("memra_moe_f16g_gather_act rc={rc}").into()); }
1422        }
1423        Ok((act, scales))
1424    }
1425
1426    /// One projection through the grouped f16 lane: dequant the active experts' rows to an
1427    /// f16 workspace, then ONE grouped GEMM over the CSR groups (variable m per expert).
1428    /// y = f32 [n_pairs, out_f] pair-major — same layout as mmq_iq_experts.
1429    /// MEMRA_MOE_F16G=1: cublasGemmGroupedBatchedEx (+ h2f pass + per-projection sync — the
1430    /// grouped API runs on internal streams unordered with ours, round-47 ledger).
1431    /// MEMRA_MOE_F16G=2: single-kernel grouped GEMM on the engine stream (round 49) — the
1432    /// row scale folds into the kernel epilogue; no f16 C, no h2f, NO sync (ordered by
1433    /// construction). f16-MIRROR numeric class either way (argmax/spec gated, not
1434    /// byte-identity). Errors on unsupported qtype (caller keeps the MMQ arm as fallback).
1435    #[allow(clippy::too_many_arguments)]
1436    pub fn moe_f16_grouped(
1437        &self,
1438        table: &CudaSlice<u64>,
1439        proj: i32,
1440        n_expert: usize,
1441        ex_ids: &CudaSlice<i32>,
1442        ex_off_host: &[i32],
1443        ex_off_dev: &CudaSlice<i32>,
1444        act_f16: &CudaSlice<u8>,
1445        act_scale: &CudaSlice<f32>,
1446        in_f: usize,
1447        out_f: usize,
1448        n_active: usize,
1449        n_pairs: usize,
1450        qtype: i32,
1451        row_bytes: usize,
1452    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1453        let sk = crate::moe_f16g_mode() >= 2 && in_f % 32 == 0;
1454        // DIRECT-FROM-QUANT lane (lane/kquant-tile-loaders + lane/iq-direct-loaders, default
1455        // ON — MEMRA_F16G_DIRECT=0 is the rollback seam): Q4_K/Q6_K/IQ4_XS/IQ3_S expert
1456        // projections skip the dequant-workspace pass entirely; the sk visitor forms dequant
1457        // B tiles in-register from the superblocks. Bit-identical to the workspace path by
1458        // construction (kernel-check "f16g-kq-direct") — this is a pure data-movement change,
1459        // not a numeric-class change. Admission mirrors the C-side guards; the grid-scan
1460        // rollback arm (MEMRA_F16G_SK=0) keeps the workspace.
1461        let (shape_sel, cross) = crate::moe_f16g_sk_params();
1462        if sk && shape_sel >= 0 && crate::moe_f16g_direct_on(qtype)
1463            && (qtype == crate::QT_Q4_K || qtype == crate::QT_Q6_K
1464                || qtype == crate::QT_IQ4_XS || qtype == crate::QT_IQ3_S)
1465            && in_f % 256 == 0 && n_active <= 512 && n_active > 0
1466        {
1467            let max_m = ex_off_host.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
1468            let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1469            {
1470                let stream = self.gpu.stream();
1471                let (tab_p, _g0) = table.device_ptr(&stream);
1472                let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1473                let (a_p, _g2) = act_f16.device_ptr(&stream);
1474                let (s_p, _g3) = act_scale.device_ptr(&stream);
1475                let (off_p, _g4) = ex_off_dev.device_ptr(&stream);
1476                let (y_p, _g5) = y.device_ptr_mut(&stream);
1477                let rc = unsafe {
1478                    memra_moe_kq_gemm_sk(tab_p as *const u64, proj, n_expert as i32,
1479                        ei_p as *const i32, a_p as *const core::ffi::c_void, y_p as *mut f32,
1480                        s_p as *const f32, off_p as *const i32, ex_off_host.as_ptr(),
1481                        n_active as i32, max_m, in_f as i32, out_f as i32, qtype, cross,
1482                        crate::moe_f16g_tail_on() as i32, row_bytes as i64,
1483                        stream.cu_stream() as *mut core::ffi::c_void)
1484                };
1485                if rc != 0 { return Err(format!("memra_moe_kq_gemm_sk rc={rc}").into()); }
1486            }
1487            return Ok(y);
1488        }
1489        // one-time cublas grouped init (algo heuristics + module load cost ~10% of a cold
1490        // g26 prime when paid inside the first projection): a tiny dummy grouped GEMM at
1491        // first use, synced, so the real prime runs warm. The =2 path never touches cublas.
1492        if !sk {
1493            static WARM: std::sync::Once = std::sync::Once::new();
1494            let mut warm_err = None;
1495            WARM.call_once(|| {
1496                let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1497                    let w = self.alloc_uninit::<u8>(2 * 32 * 64 * 2)?;
1498                    let a = self.alloc_uninit::<u8>(4 * 64 * 2)?;
1499                    let mut yw = self.alloc_uninit::<u8>(4 * 32 * 2)?;
1500                    let off = [0i32, 2, 4];
1501                    let stream = self.gpu.stream();
1502                    let (w_p, _a1) = w.device_ptr(&stream);
1503                    let (a_p, _a2) = a.device_ptr(&stream);
1504                    let (y_p, _a3) = yw.device_ptr_mut(&stream);
1505                    let rc = unsafe {
1506                        memra_moe_f16g_gemm(w_p as *const core::ffi::c_void,
1507                            a_p as *const core::ffi::c_void, y_p as *mut core::ffi::c_void,
1508                            off.as_ptr(), 2, 64, 32,
1509                            stream.cu_stream() as *mut core::ffi::c_void)
1510                    };
1511                    if rc != 0 { return Err(format!("f16g warmup rc={rc}").into()); }
1512                    self.gpu.stream().synchronize()?;
1513                    Ok(())
1514                })();
1515                if let Err(e) = r { warm_err = Some(e.to_string()); }
1516            });
1517            if let Some(we) = warm_err { return Err(we.into()); }
1518        }
1519        let w_bytes = n_active * out_f * in_f * 2;
1520        let mut w_f16 = self.alloc_uninit::<u8>(w_bytes)?;
1521        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1522        {
1523            let stream = self.gpu.stream();
1524            let (tab_p, _g0) = table.device_ptr(&stream);
1525            let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1526            let (w_p, _g2) = w_f16.device_ptr_mut(&stream);
1527            let rc = unsafe {
1528                memra_moe_f16g_dequant(tab_p as *const u64, proj, n_expert as i32,
1529                    ei_p as *const i32, w_p as *mut core::ffi::c_void,
1530                    in_f as i32, out_f as i32, n_active as i32, qtype, row_bytes as i64,
1531                    stream.cu_stream() as *mut core::ffi::c_void)
1532            };
1533            if rc != 0 { return Err(format!("memra_moe_f16g_dequant rc={rc}").into()); }
1534            let (a_p, _g3) = act_f16.device_ptr(&stream);
1535            let (s_p, _g6) = act_scale.device_ptr(&stream);
1536            let (y_p, _g5) = y.device_ptr_mut(&stream);
1537            if sk {
1538                let max_m = ex_off_host.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
1539                let (off_p, _g7) = ex_off_dev.device_ptr(&stream);
1540                let (shape_sel, cross) = crate::moe_f16g_sk_params();
1541                let rc = unsafe {
1542                    memra_moe_f16g_gemm_sk(w_p as *const core::ffi::c_void,
1543                        a_p as *const core::ffi::c_void, y_p as *mut f32,
1544                        s_p as *const f32, off_p as *const i32, ex_off_host.as_ptr(),
1545                        n_active as i32, max_m, in_f as i32, out_f as i32, shape_sel, cross,
1546                        crate::moe_f16g_tail_on() as i32,
1547                        stream.cu_stream() as *mut core::ffi::c_void)
1548                };
1549                if rc != 0 { return Err(format!("memra_moe_f16g_gemm_sk rc={rc}").into()); }
1550            } else {
1551                let mut y16 = self.alloc_uninit::<u8>(n_pairs * out_f * 2)?;
1552                let (y16_p, _g4) = y16.device_ptr_mut(&stream);
1553                let rc = unsafe {
1554                    memra_moe_f16g_gemm(w_p as *const core::ffi::c_void,
1555                        a_p as *const core::ffi::c_void, y16_p as *mut core::ffi::c_void,
1556                        ex_off_host.as_ptr(), n_active as i32, in_f as i32, out_f as i32,
1557                        stream.cu_stream() as *mut core::ffi::c_void)
1558                };
1559                if rc != 0 { return Err(format!("memra_moe_f16g_gemm rc={rc}").into()); }
1560                let rc = unsafe {
1561                    memra_moe_f16g_h2f_scaled(y16_p as *const core::ffi::c_void, y_p as *mut f32,
1562                        s_p as *const f32, out_f as i32, n_pairs as i32,
1563                        stream.cu_stream() as *mut core::ffi::c_void)
1564                };
1565                if rc != 0 { return Err(format!("memra_moe_f16g_h2f_scaled rc={rc}").into()); }
1566            }
1567        }
1568        // MODE 1 ONLY: cublasGemmGroupedBatchedEx issues through internal streams NOT ordered
1569        // with ours (round 46: NaN race, clean under sync — 205=205 MATCH). Full sync per
1570        // projection. Mode 2 (single kernel, our stream) is ordered by construction — no sync,
1571        // that is the point of this arc.
1572        if !sk {
1573            self.gpu.stream().synchronize()?;
1574        }
1575        if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
1576            // FULL NaN/Inf scan of w, act (through h2f) and y — localizes the corrupt stage.
1577            let wn = n_active * out_f * in_f;
1578            let an = n_pairs * in_f;
1579            let mut wf = self.alloc_uninit::<f32>(wn)?;
1580            let mut af = self.alloc_uninit::<f32>(an)?;
1581            {
1582                let stream = self.gpu.stream();
1583                let (w_p, _a) = w_f16.device_ptr(&stream);
1584                let (a_p, _b) = act_f16.device_ptr(&stream);
1585                let (wf_p, _c) = wf.device_ptr_mut(&stream);
1586                let (af_p, _d) = af.device_ptr_mut(&stream);
1587                unsafe {
1588                    memra_moe_f16g_h2f(w_p as *const core::ffi::c_void, wf_p as *mut f32, wn,
1589                        stream.cu_stream() as *mut core::ffi::c_void);
1590                    memra_moe_f16g_h2f(a_p as *const core::ffi::c_void, af_p as *mut f32, an,
1591                        stream.cu_stream() as *mut core::ffi::c_void);
1592                }
1593            }
1594            let (wh, ah, yh) = (self.dtoh(&wf)?, self.dtoh(&af)?, self.dtoh(&y)?);
1595            let scan = |v: &[f32]| -> (usize, f32) {
1596                let bad = v.iter().filter(|x| !x.is_finite()).count();
1597                let mx = v.iter().filter(|x| x.is_finite()).fold(0.0f32, |m, x| m.max(x.abs()));
1598                (bad, mx)
1599            };
1600            let (wb, wm) = scan(&wh); let (ab, am) = scan(&ah); let (yb, ym) = scan(&yh);
1601            eprintln!("[f16g-debug] proj={proj} w: bad={wb} max={wm:.3e} | act: bad={ab} \
1602                       max={am:.3e} | y: bad={yb} max={ym:.3e} (na={n_active} np={n_pairs} \
1603                       in={in_f} out={out_f})");
1604        }
1605        Ok(y)
1606    }
1607
1608    /// Raw sk grouped-GEMM entry for kernel-check ("f16g-sk" section): explicit shape/cross
1609    /// instead of the env policy. shape_sel < 0 = the round-49 grid-scan rollback arm; else
1610    /// the round-51 problem-visitor split at `cross` (1 forces all-128, i32::MAX all-32).
1611    /// tail: 1 = the deep tail (32x64x64 3-stage, lane/sk-tail-form) on sub-cross groups,
1612    /// 0 = the round-51 2-stage 32x64x32 tail.
1613    /// w_f16 = [n_active][out_f][in_f] f16 bytes, act_f16 = [n_pairs][in_f] f16 bytes.
1614    #[allow(clippy::too_many_arguments)]
1615    pub fn moe_f16g_gemm_sk_raw(&self, w_f16: &CudaSlice<u8>, act_f16: &CudaSlice<u8>,
1616        row_scale: &CudaSlice<f32>, ex_off_host: &[i32], ex_off_dev: &CudaSlice<i32>,
1617        in_f: usize, out_f: usize, n_pairs: usize, shape_sel: i32, cross: i32, tail: i32)
1618        -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1619        let n_active = ex_off_host.len() - 1;
1620        let max_m = ex_off_host.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
1621        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1622        {
1623            let stream = self.gpu.stream();
1624            let (w_p, _g0) = w_f16.device_ptr(&stream);
1625            let (a_p, _g1) = act_f16.device_ptr(&stream);
1626            let (s_p, _g2) = row_scale.device_ptr(&stream);
1627            let (off_p, _g3) = ex_off_dev.device_ptr(&stream);
1628            let (y_p, _g4) = y.device_ptr_mut(&stream);
1629            let rc = unsafe {
1630                memra_moe_f16g_gemm_sk(w_p as *const core::ffi::c_void,
1631                    a_p as *const core::ffi::c_void, y_p as *mut f32,
1632                    s_p as *const f32, off_p as *const i32, ex_off_host.as_ptr(),
1633                    n_active as i32, max_m, in_f as i32, out_f as i32, shape_sel, cross,
1634                    tail, stream.cu_stream() as *mut core::ffi::c_void)
1635            };
1636            if rc != 0 { return Err(format!("memra_moe_f16g_gemm_sk rc={rc}").into()); }
1637        }
1638        Ok(y)
1639    }
1640
1641    /// Raw direct-from-quant sk grouped-GEMM entry for kernel-check ("f16g-kq-direct"):
1642    /// explicit cross/tail instead of the env policy. `table` = device u64 pointer table
1643    /// (proj-major, [n_proj][n_expert] — same contract as moe_f16_grouped), `ex_ids` =
1644    /// active-expert ids (device). Visitor forms only (the C side rejects anything else).
1645    #[allow(clippy::too_many_arguments)]
1646    pub fn moe_kq_gemm_sk_raw(&self, table: &CudaSlice<u64>, proj: i32, n_expert: usize,
1647        ex_ids: &CudaSlice<i32>, act_f16: &CudaSlice<u8>, row_scale: &CudaSlice<f32>,
1648        ex_off_host: &[i32], ex_off_dev: &CudaSlice<i32>,
1649        in_f: usize, out_f: usize, n_pairs: usize, qtype: i32, row_bytes: usize, cross: i32,
1650        tail: i32)
1651        -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1652        let n_active = ex_off_host.len() - 1;
1653        let max_m = ex_off_host.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
1654        let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1655        {
1656            let stream = self.gpu.stream();
1657            let (tab_p, _g0) = table.device_ptr(&stream);
1658            let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1659            let (a_p, _g2) = act_f16.device_ptr(&stream);
1660            let (s_p, _g3) = row_scale.device_ptr(&stream);
1661            let (off_p, _g4) = ex_off_dev.device_ptr(&stream);
1662            let (y_p, _g5) = y.device_ptr_mut(&stream);
1663            let rc = unsafe {
1664                memra_moe_kq_gemm_sk(tab_p as *const u64, proj, n_expert as i32,
1665                    ei_p as *const i32, a_p as *const core::ffi::c_void, y_p as *mut f32,
1666                    s_p as *const f32, off_p as *const i32, ex_off_host.as_ptr(),
1667                    n_active as i32, max_m, in_f as i32, out_f as i32, qtype, cross,
1668                    tail, row_bytes as i64,
1669                    stream.cu_stream() as *mut core::ffi::c_void)
1670            };
1671            if rc != 0 { return Err(format!("memra_moe_kq_gemm_sk rc={rc}").into()); }
1672        }
1673        Ok(y)
1674    }
1675
1676    /// Raw dequant-workspace entry for kernel-check: dequant the active experts' rows to a
1677    /// fresh f16 workspace via the same kernel `moe_f16_grouped` uses (the direct loaders'
1678    /// bitwise reference).
1679    pub fn moe_f16g_dequant_raw(&self, table: &CudaSlice<u64>, proj: i32, n_expert: usize,
1680        ex_ids: &CudaSlice<i32>, in_f: usize, out_f: usize, n_active: usize, qtype: i32,
1681        row_bytes: usize)
1682        -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1683        let mut w_f16 = self.alloc_uninit::<u8>(n_active * out_f * in_f * 2)?;
1684        {
1685            let stream = self.gpu.stream();
1686            let (tab_p, _g0) = table.device_ptr(&stream);
1687            let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1688            let (w_p, _g2) = w_f16.device_ptr_mut(&stream);
1689            let rc = unsafe {
1690                memra_moe_f16g_dequant(tab_p as *const u64, proj, n_expert as i32,
1691                    ei_p as *const i32, w_p as *mut core::ffi::c_void,
1692                    in_f as i32, out_f as i32, n_active as i32, qtype, row_bytes as i64,
1693                    stream.cu_stream() as *mut core::ffi::c_void)
1694            };
1695            if rc != 0 { return Err(format!("memra_moe_f16g_dequant rc={rc}").into()); }
1696        }
1697        Ok(w_f16)
1698    }
1699}