memra_engine/
f16_ffi.rs

1//! FP16-MIRROR PREFILL (MEMRA_PP_F16=1): cuBLASLt FP16 TN GEMM on a resident fp16 dequant
2//! mirror of the Q8_0 trunk weights.
3//!
4//! Probe verdict 2026-07-26 (tools/bench_lt_f16.cu on the H100 box): 611-687 TF at the 9B
5//! m=512 prefill shapes vs the vendored MMQ per-shape medians = **3.2-3.7x per launch**
6//! (MMQ = 60% of prime). Why fp16 and not faster int8: the exact wgmma arc proved Q8_0's
7//! per-32-block scale fold serializes Hopper's warpgroup MMA pipe (ptxas C7514, ledger'd);
8//! fp16 f32-accumulate has no mid-loop accumulator reads and streams at tensor-core rate.
9//!
10//! NUMERIC CONFIG (new, explicit, opt-in — MEMRA_PP_FP8/GDN-chunked precedent): the int8 part
11//! of the dequant is exact in fp16 (7 mantissa bits into 11); rounding enters at d*q products
12//! and the activation f32->fp16 cast. run-gen argmax battery + kernel-check tolerance gate
13//! arbitrate. Decode (m<16) keeps the Q8_0 dp4a/MMVQ chain untouched — decode==verify law holds.
14//!
15//! VRAM: the mirror duplicates every 2D Q8_0 projection at 2 B/w (9B model ~+17GB) — an 80GB
16//! H100 lane feature. MEMRA_PP_F16_BUDGET_MB (default 32768) caps the spend, layer-order prefix.
17
18use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut};
19
20unsafe extern "C" {
21    /// One FP16 prefill GEMM: f32->fp16 activation convert + cublasLtMatmul TN on one stream.
22    fn memra_f16_pp_gemm(
23        w_f16: *const core::ffi::c_void,
24        x_f32: *const f32,
25        xh_f16: *mut core::ffi::c_void,
26        y_f32: *mut f32,
27        m: i32,
28        n: i32,
29        k: i32,
30        ws: *mut core::ffi::c_void,
31        ws_bytes: usize,
32        stream: *mut core::ffi::c_void,
33    ) -> i32;
34    /// Standalone f32->fp16 convert (grouped dispatch: one convert feeds N GEMMs).
35    fn memra_f16_cvt(
36        x_f32: *const f32,
37        xh_f16: *mut core::ffi::c_void,
38        nelem: usize,
39        stream: *mut core::ffi::c_void,
40    ) -> i32;
41    /// GEMM on a PRE-CONVERTED fp16 activation (see memra_f16_cvt).
42    fn memra_f16_pp_gemm_pre(
43        w_f16: *const core::ffi::c_void,
44        xh_f16: *const core::ffi::c_void,
45        y_f32: *mut f32,
46        m: i32,
47        n: i32,
48        k: i32,
49        ws: *mut core::ffi::c_void,
50        ws_bytes: usize,
51        stream: *mut core::ffi::c_void,
52    ) -> i32;
53    /// One BF16 prefill GEMM on a RESIDENT bf16 weight (no mirror): f32->bf16 activation
54    /// convert + cublasLtMatmul TN (CUDA_R_16BF) on one stream.
55    fn memra_bf16_pp_gemm(
56        w_bf16: *const core::ffi::c_void,
57        x_f32: *const f32,
58        xb_bf16: *mut core::ffi::c_void,
59        y_f32: *mut f32,
60        m: i32,
61        n: i32,
62        k: i32,
63        ws: *mut core::ffi::c_void,
64        ws_bytes: usize,
65        stream: *mut core::ffi::c_void,
66    ) -> i32;
67    /// GGUF Q8_0 34B blocks -> row-major fp16 mirror (load-time).
68    fn memra_q8_0_dequant_f16(
69        w_q8: *const core::ffi::c_void,
70        w_f16: *mut core::ffi::c_void,
71        out_f: i64,
72        nblk_row: i64,
73        stream: *mut core::ffi::c_void,
74    ) -> i32;
75    /// GGUF Q4_0 18B blocks -> row-major fp16 mirror (campaign A, 2026-07-31).
76    fn memra_q4_0_dequant_f16(
77        w_q4: *const core::ffi::c_void,
78        w_f16: *mut core::ffi::c_void,
79        out_f: i64,
80        nblk_row: i64,
81        stream: *mut core::ffi::c_void,
82    ) -> i32;
83    /// GGUF Q6_K 210B superblocks -> row-major fp16 mirror (round 47: the q27 prefill wall).
84    fn memra_q6_K_dequant_f16(
85        w_q6: *const core::ffi::c_void,
86        w_f16: *mut core::ffi::c_void,
87        out_f: i64,
88        nsb_row: i64,
89        stream: *mut core::ffi::c_void,
90    ) -> i32;
91    /// GGUF Q4_K 144B superblocks -> row-major fp16 mirror (round 49: the q27 trunk bulk).
92    fn memra_q4_K_dequant_f16(
93        w_q4k: *const core::ffi::c_void,
94        w_f16: *mut core::ffi::c_void,
95        out_f: i64,
96        nsb_row: i64,
97        stream: *mut core::ffi::c_void,
98    ) -> i32;
99    /// GGUF Q5_K 176B superblocks -> row-major fp16 mirror (round 49b: q27 ssm_out).
100    fn memra_q5_K_dequant_f16(
101        w_q5k: *const core::ffi::c_void,
102        w_f16: *mut core::ffi::c_void,
103        out_f: i64,
104        nsb_row: i64,
105        stream: *mut core::ffi::c_void,
106    ) -> i32;
107}
108
109/// MEMRA_PP_F16 gate, read once. DEFAULT ON on the Hopper lane (80GB — the mirror costs
110/// 2 B/w, ~17GB on the 9B; the box carries it), opt-in elsewhere; =1/=0 overrides either way.
111/// Promotion battery (2026-07-26, H100): kernel-check ALL GREEN (f16 rel <= 6.5e-3, band 1e-2);
112/// run-gen argmax MATCH on p1/p2/p3 long prompts; greedy streams IDENTICAL to the MMQ config
113/// on all three; pp512 8674 -> 15626 tok/s (+80%, N=5 medians). Decode untouched (m>=16 arm).
114pub fn pp_f16_enabled() -> bool {
115    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
116    *ON.get_or_init(|| match std::env::var("MEMRA_PP_F16").as_deref() {
117        Ok("1") => true,
118        Ok("0") => false,
119        _ => cfg!(memra_hopper_mma),
120    })
121}
122
123/// Capacity-keyed f16-mirror admission (zoo-fusion arc, 2026-08-17): when MEMRA_PP_F16
124/// is UNSET, the mirror walk may still turn on iff free VRAM covers `need` (the
125/// admissible mirror mass the caller computed) plus serving headroom. The Q6_K prefill
126/// dequant-GEMM wall costs 3.46ms/layer-call on the 31B downQ6K trunk (30% of c8 GPU
127/// time, ttft 1.38s); the f16 lane removes it (measured c8 agg +37%, ttft -70%). The
128/// env keeps absolute priority in pp_f16_enabled(); this fn only decides the UNSET
129/// case, per-boot, from the measured free after weights. 24GB rigs refuse by
130/// construction (need + 8GiB never fits).
131pub fn pp_f16_capacity_ok(free: usize, need: usize) -> bool {
132    if std::env::var("MEMRA_PP_F16").is_ok() {
133        return false; // explicit env decided already via pp_f16_enabled()
134    }
135    need > 0 && free >= need + (8usize << 30)
136}
137
138/// MEMRA_PP_BF16 gate, read once — the resident-BF16 tensor-core prefill GEMM.
139///
140/// WHY IT EXISTS (2026-08-28, step37 prime): a BF16 checkpoint has no Q8_0 fp16 mirror, so every
141/// prefill projection fell to `linear_bf16_chunked_inner`, which dequants the FULL weight to f32
142/// and runs an f32 GEMM — no tensor cores, and a fresh 2x-weight f32 buffer per call. cuBLASLt
143/// consumes CUDA_R_16BF directly and the checkpoint bytes are already the operand layout the TN
144/// form wants, so this path costs no mirror and no VRAM.
145///
146/// DEFAULT (explicit, per the new-flag law): OFF until this lane's A/B + argmax receipts land;
147/// the flip to family-default-ON carries its FLAGS.md row and receipts in the same PR.
148/// Decode never reaches it — callers gate on m >= 16.
149pub fn pp_bf16_enabled() -> bool {
150    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
151    *ON.get_or_init(|| {
152        let on = matches!(std::env::var("MEMRA_PP_BF16").as_deref(), Ok("1"));
153        // ANNOUNCE PER FLAG VALUE, printed in BOTH arms (the moe-grouped-prefill pattern):
154        // an A/B grep must distinguish "flag off" from "flag on but never consulted" without
155        // the line being an arm-local cost. First consult happens on the first bf16-resident
156        // prefill GEMM (the door's own condition chain), so a boot that never primes a
157        // bf16-resident weight prints nothing, which is itself the honest reading.
158        eprintln!(
159            "[bf16-tc] flag={} (MEMRA_PP_BF16; engagement is the per-shape ENGAGED line + \
160             the dispatch counter)",
161            if on { "on" } else { "off" }
162        );
163        on
164    })
165}
166
167/// Engagement counter for the resident-BF16 tensor-core prefill GEMM (`MEMRA_PP_BF16`),
168/// incremented once per ACCEPTED `bf16_tc_gemm` launch at the invocation itself, after the
169/// cuBLASLt decline check, so a declined shape does not count. Same reason the
170/// moe-grouped-prefill counter exists: the per-shape ENGAGED eprintln dedups by shape, so a
171/// gate that must assert "the door ran N times for this workload" needs a counter at the
172/// invocation, not a log grep (LAW:wiring-assertions-match-prose).
173pub static BF16_TC_DISPATCHES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
174
175/// Snapshot of [`BF16_TC_DISPATCHES`]. Gates take a before/after pair around a workload and
176/// assert on the delta.
177pub fn bf16_tc_dispatches() -> u64 {
178    BF16_TC_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
179}
180
181/// Resident scratch (fp8_ffi::Fp8Scratch pattern): fp16 activation (grown to the largest m*k
182/// seen) + the cuBLASLt workspace. Single GPU worker; the Mutex guards lazy build/grow only.
183pub struct F16Scratch {
184    pub xh: CudaSlice<u8>,
185    pub ws: CudaSlice<u8>,
186    cap_xh: usize,
187}
188
189impl F16Scratch {
190    /// Pre-sized scratch (task #14: the captured prime gets a PRIVATE scratch so the
191    /// graph's baked cvt/Lt pointers are never mutated by eager GEMMs between replays).
192    pub fn with_capacity(
193        e: &crate::Engine,
194        xh_bytes: usize,
195    ) -> Result<Self, Box<dyn std::error::Error>> {
196        Ok(F16Scratch {
197            xh: e.alloc_u8_uninit(xh_bytes)?,
198            ws: e.alloc_u8_uninit(F16_WS_BYTES)?,
199            cap_xh: xh_bytes,
200        })
201    }
202}
203
204pub(crate) const F16_WS_BYTES: usize = 64 << 20;
205
206impl crate::Engine {
207    /// Swap the resident f16 scratch (task #14 capture isolation). Returns the previous
208    /// contents; pass them back to restore.
209    pub fn f16_scratch_swap(&self, new: Option<F16Scratch>) -> Option<F16Scratch> {
210        std::mem::replace(&mut *self.f16_scratch.lock().unwrap(), new)
211    }
212
213    /// FP16 prefill GEMM for a weight carrying the f16 mirror: y[m,out] = x[m,in] @ (fp16 W)^T,
214    /// f32 accumulate. Returns None when the weight has no mirror (caller falls through to MMQ).
215    pub fn try_f16_gemm(
216        &self,
217        w: &crate::model::GpuTensor,
218        x: &CudaSlice<f32>,
219        m: usize,
220    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
221        use crate::model::GpuTensor;
222        let (w16, ne, scale) = match w {
223            GpuTensor::Quant {
224                f16: Some(w16),
225                ne,
226                scale,
227                ..
228            } => (w16, ne, *scale),
229            _ => return Ok(None),
230        };
231        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
232        let mut y = self.qmatvec_gemm_f16_raw(w16, x, m, in_f, out_f)?;
233        if scale != 1.0 {
234            self.scale_inplace(&mut y, scale, m * out_f)?;
235        }
236        Ok(Some(y))
237    }
238
239    /// BF16 tensor-core prefill GEMM on RESIDENT checkpoint bytes: y[m,out] = x[m,in] @ W^T,
240    /// f32 accumulate. `data` is the untouched row-major [out_f, in_f] bf16 weight — no mirror,
241    /// no dequant, no extra VRAM. Shares the fp16 scratch (a bf16 activation is the same 2 B/elem).
242    pub fn bf16_tc_gemm(
243        &self,
244        data: &CudaSlice<u8>,
245        x: &CudaSlice<f32>,
246        m: usize,
247        in_f: usize,
248        out_f: usize,
249    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
250        let need_xh = m * in_f * 2;
251        let mut guard = self.f16_scratch.lock().unwrap();
252        if guard.is_none() {
253            *guard = Some(F16Scratch {
254                xh: self.alloc_u8_uninit(need_xh)?,
255                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
256                cap_xh: need_xh,
257            });
258        }
259        let s = guard.as_mut().unwrap();
260        if need_xh > s.cap_xh {
261            s.xh = self.alloc_u8_uninit(need_xh)?;
262            s.cap_xh = need_xh;
263        }
264        let mut y = self.uninit(m * out_f)?; // full-overwrite GEMM output: skip memset
265        let rc = {
266            let stream = self.gpu.stream();
267            let (w_p, _gw) = data.device_ptr(&stream);
268            let (x_p, _gx) = x.device_ptr(&stream);
269            let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
270            let (y_p, _gy) = y.device_ptr_mut(&stream);
271            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
272            // cuBLASLt wants 16B-aligned operands and its heuristic does not inspect the
273            // pointer, so an unaligned resident slice reaches cublasLtMatmul and comes back
274            // NOT_SUPPORTED. Decline it here instead.
275            if !(w_p as usize).is_multiple_of(16) {
276                -1
277            } else {
278                unsafe {
279                    memra_bf16_pp_gemm(
280                        w_p as *const core::ffi::c_void,
281                        x_p as *const f32,
282                        h_p as *mut core::ffi::c_void,
283                        y_p as *mut f32,
284                        m as i32,
285                        out_f as i32,
286                        in_f as i32,
287                        ws_p as *mut core::ffi::c_void,
288                        F16_WS_BYTES,
289                        stream.cu_stream() as *mut core::ffi::c_void,
290                    )
291                }
292            }
293        };
294        // NOT a hard error: cuBLASLt refuses some (m,n,k)/alignment combinations that its own
295        // heuristic accepted (measured 2026-08-28: rc=30014 = CUBLAS_STATUS_NOT_SUPPORTED at
296        // m=43 n=4096 k=1024, after the same door served every 4096-token prime shape). The f32
297        // dequant path below this door is always correct, so a refusal DECLINES the shape rather
298        // than failing the request. Announced once per shape so the decline can never be silent —
299        // a door that quietly stops engaging reads exactly like a door that never helped.
300        if rc != 0 {
301            static SAID: std::sync::Mutex<
302                Option<std::collections::HashSet<(usize, usize, usize)>>,
303            > = std::sync::Mutex::new(None);
304            let mut g = SAID.lock().unwrap();
305            let seen = g.get_or_insert_with(std::collections::HashSet::new);
306            if seen.insert((m, out_f, in_f)) {
307                eprintln!(
308                    "[bf16-tc] DECLINED m={m} n={out_f} k={in_f} rc={rc} \
309                     (1xxxx=cudaError convert, 2xxxx=no cublasLt algo, 3xxxx=matmul status, \
310                     4xxxx=cublasLtCreate, -1=weight not 16B-aligned) — this shape falls back to \
311                     the f32 dequant GEMM; every other shape keeps the tensor-core path"
312                );
313            }
314            return Ok(None);
315        }
316        // ENGAGEMENT RECEIPT, and it is not optional. Only the DECLINE was announced, so an arm
317        // with MEMRA_PP_BF16=1 that never actually took this path was indistinguishable from one
318        // that did -- and a correctness gate whose two arms ran the SAME code reports a
319        // byte-identical MATCH for the wrong reason. Announced once per shape, same as the
320        // decline, so it costs one line per distinct GEMM and nothing per token.
321        BF16_TC_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
322        {
323            static ACCEPTED: std::sync::Mutex<
324                Option<std::collections::HashSet<(usize, usize, usize)>>,
325            > = std::sync::Mutex::new(None);
326            let mut g = ACCEPTED.lock().unwrap();
327            let seen = g.get_or_insert_with(std::collections::HashSet::new);
328            if seen.insert((m, out_f, in_f)) {
329                eprintln!(
330                    "[bf16-tc] ENGAGED m={m} n={out_f} k={in_f} (bf16 tensor-core GEMM on resident checkpoint bytes)"
331                );
332            }
333        }
334        Ok(Some(y))
335    }
336
337    /// cuBLASLt REFERENCE GEMV for the bf16 decode rows (`MEMRA_B200_BF16_GEMV_LT`,
338    /// lane/b200-gemv-hbm-20260902). Same plan `bf16_tc_gemm` builds (TN, CUDA_R_16BF operands,
339    /// f32 accumulate/output, per-device handle) but WRITES INTO the caller's `y` instead of
340    /// allocating one, so the bench times the library GEMV and nothing else. `t` is the token
341    /// count (1 for plain decode); `y` must hold `t * out_f` floats.
342    ///
343    /// NUMERIC CLASS `bf16_gemv_lt`, not a bit-identical twin: the activation is cast f32 ->
344    /// bf16 before the multiply and the K summation order is the library's. Default OFF,
345    /// reference only; see `crate::b200_bf16_gemv_lt_on`.
346    ///
347    /// Returns `Ok(false)` when cuBLASLt declines the shape (unaligned weight, no algo) so the
348    /// caller falls through to the shipped kernel rather than failing the request — the same
349    /// contract `bf16_tc_gemm` uses, and announced once per shape for the same reason.
350    pub fn bf16_gemv_lt_into(
351        &self,
352        data: &CudaSlice<u8>,
353        x: &CudaSlice<f32>,
354        y: &mut CudaSlice<f32>,
355        in_f: usize,
356        out_f: usize,
357        t: usize,
358    ) -> Result<bool, Box<dyn std::error::Error>> {
359        if t == 0 || x.len() < t * in_f || y.len() < t * out_f {
360            return Err("bf16_gemv_lt geometry".into());
361        }
362        if data.len() < out_f * in_f * 2 {
363            return Err("bf16_gemv_lt weight too small".into());
364        }
365        let need_xh = t * in_f * 2;
366        let mut guard = self.f16_scratch.lock().unwrap();
367        if guard.is_none() {
368            *guard = Some(F16Scratch {
369                xh: self.alloc_u8_uninit(need_xh)?,
370                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
371                cap_xh: need_xh,
372            });
373        }
374        let s = guard.as_mut().unwrap();
375        if need_xh > s.cap_xh {
376            s.xh = self.alloc_u8_uninit(need_xh)?;
377            s.cap_xh = need_xh;
378        }
379        let rc = {
380            let stream = self.gpu.stream();
381            let (w_p, _gw) = data.device_ptr(&stream);
382            let (x_p, _gx) = x.device_ptr(&stream);
383            let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
384            let (y_p, _gy) = y.device_ptr_mut(&stream);
385            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
386            if !(w_p as usize).is_multiple_of(16) {
387                -1
388            } else {
389                unsafe {
390                    memra_bf16_pp_gemm(
391                        w_p as *const core::ffi::c_void,
392                        x_p as *const f32,
393                        h_p as *mut core::ffi::c_void,
394                        y_p as *mut f32,
395                        t as i32,
396                        out_f as i32,
397                        in_f as i32,
398                        ws_p as *mut core::ffi::c_void,
399                        F16_WS_BYTES,
400                        stream.cu_stream() as *mut core::ffi::c_void,
401                    )
402                }
403            }
404        };
405        if rc != 0 {
406            static SAID: std::sync::Mutex<
407                Option<std::collections::HashSet<(usize, usize, usize)>>,
408            > = std::sync::Mutex::new(None);
409            let mut g = SAID.lock().unwrap();
410            let seen = g.get_or_insert_with(std::collections::HashSet::new);
411            if seen.insert((t, out_f, in_f)) {
412                eprintln!(
413                    "[b200-bf16-gemv-lt] DECLINED t={t} n={out_f} k={in_f} rc={rc} \
414                     (1xxxx=cudaError convert, 2xxxx=no cublasLt algo, 3xxxx=matmul status, \
415                     4xxxx=cublasLtCreate, -1=weight not 16B-aligned) — this shape keeps the \
416                     shipped matvec kernel"
417                );
418            }
419            return Ok(false);
420        }
421        {
422            static ACCEPTED: std::sync::Mutex<
423                Option<std::collections::HashSet<(usize, usize, usize)>>,
424            > = std::sync::Mutex::new(None);
425            let mut g = ACCEPTED.lock().unwrap();
426            let seen = g.get_or_insert_with(std::collections::HashSet::new);
427            if seen.insert((t, out_f, in_f)) {
428                eprintln!(
429                    "[b200-bf16-gemv-lt] ENGAGED t={t} n={out_f} k={in_f} (cuBLASLt reference \
430                     GEMV, numeric class bf16_gemv_lt, MEMRA_B200_BF16_GEMV_LT=1)"
431                );
432            }
433        }
434        Ok(true)
435    }
436
437    /// Bare FP16 GEMM launch on an fp16 mirror — also the kernel_check gate entry.
438    pub fn qmatvec_gemm_f16_raw(
439        &self,
440        w16: &CudaSlice<u8>,
441        x: &CudaSlice<f32>,
442        m: usize,
443        in_f: usize,
444        out_f: usize,
445    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
446        let need_xh = m * in_f * 2;
447        let mut guard = self.f16_scratch.lock().unwrap();
448        if guard.is_none() {
449            *guard = Some(F16Scratch {
450                xh: self.alloc_u8_uninit(need_xh)?,
451                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
452                cap_xh: need_xh,
453            });
454        }
455        let s = guard.as_mut().unwrap();
456        if need_xh > s.cap_xh {
457            s.xh = self.alloc_u8_uninit(need_xh)?;
458            s.cap_xh = need_xh;
459        }
460        let mut y = self.uninit(m * out_f)?; // full-overwrite GEMM output: skip memset
461        let rc = {
462            let stream = self.gpu.stream();
463            let (w_p, _gw) = w16.device_ptr(&stream);
464            let (x_p, _gx) = x.device_ptr(&stream);
465            let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
466            let (y_p, _gy) = y.device_ptr_mut(&stream);
467            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
468            unsafe {
469                memra_f16_pp_gemm(
470                    w_p as *const core::ffi::c_void,
471                    x_p as *const f32,
472                    h_p as *mut core::ffi::c_void,
473                    y_p as *mut f32,
474                    m as i32,
475                    out_f as i32,
476                    in_f as i32,
477                    ws_p as *mut core::ffi::c_void,
478                    F16_WS_BYTES,
479                    stream.cu_stream() as *mut core::ffi::c_void,
480                )
481            }
482        };
483        if rc != 0 {
484            return Err(format!(
485                "memra_f16_pp_gemm rc={rc} (m={m} n={out_f} k={in_f}; 1xxxx=cudaError convert, \
486                 2xxxx=no cublasLt algo, 3xxxx=matmul status)"
487            )
488            .into());
489        }
490        Ok(y)
491    }
492
493    /// f32 -> fp16 activation convert into a fresh buffer (matmul_group: ONE convert feeds
494    /// every mirror-carrying weight in the group; the standalone per-GEMM converts were ~250
495    /// launches/prime of gap-cluster fuel, nsys 2026-07-26).
496    pub fn f16_act(
497        &self,
498        x: &CudaSlice<f32>,
499        nelem: usize,
500    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
501        let mut xh = self.alloc_u8_uninit(nelem * 2)?;
502        let rc = {
503            let stream = self.gpu.stream();
504            let (x_p, _gx) = x.device_ptr(&stream);
505            let (h_p, _gh) = xh.device_ptr_mut(&stream);
506            unsafe {
507                memra_f16_cvt(
508                    x_p as *const f32,
509                    h_p as *mut core::ffi::c_void,
510                    nelem,
511                    stream.cu_stream() as *mut core::ffi::c_void,
512                )
513            }
514        };
515        if rc != 0 {
516            return Err(format!("memra_f16_cvt rc={rc}").into());
517        }
518        Ok(xh)
519    }
520
521    /// `_into` twin of `try_f16_gemm_pre` (piecewise-slab plumbing): the GEMM writes into
522    /// a caller-provided buffer (a resident slab view) instead of a fresh allocation —
523    /// the FFI has always taken the y pointer; only the wrapper allocated. Returns
524    /// Ok(false) when the weight has no mirror (caller falls back and copies).
525    pub fn try_f16_gemm_pre_into(
526        &self,
527        w: &crate::model::GpuTensor,
528        xh: &CudaSlice<u8>,
529        m: usize,
530        y: &mut CudaSlice<f32>,
531    ) -> Result<bool, Box<dyn std::error::Error>> {
532        use crate::model::GpuTensor;
533        let (w16, ne, scale) = match w {
534            GpuTensor::Quant {
535                f16: Some(w16),
536                ne,
537                scale,
538                ..
539            } => (w16, ne, *scale),
540            _ => return Ok(false),
541        };
542        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
543        assert!(
544            y.len() >= m * out_f,
545            "try_f16_gemm_pre_into: output slab too small"
546        );
547        let mut guard = self.f16_scratch.lock().unwrap();
548        if guard.is_none() {
549            *guard = Some(F16Scratch {
550                xh: self.alloc_u8_uninit(2)?,
551                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
552                cap_xh: 2,
553            });
554        }
555        let s = guard.as_mut().unwrap();
556        let rc = {
557            let stream = self.gpu.stream();
558            let (w_p, _gw) = w16.device_ptr(&stream);
559            let (h_p, _gh) = xh.device_ptr(&stream);
560            let (y_p, _gy) = y.device_ptr_mut(&stream);
561            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
562            unsafe {
563                memra_f16_pp_gemm_pre(
564                    w_p as *const core::ffi::c_void,
565                    h_p as *const core::ffi::c_void,
566                    y_p as *mut f32,
567                    m as i32,
568                    out_f as i32,
569                    in_f as i32,
570                    ws_p as *mut core::ffi::c_void,
571                    F16_WS_BYTES,
572                    stream.cu_stream() as *mut core::ffi::c_void,
573                )
574            }
575        };
576        if rc != 0 {
577            return Err(
578                format!("memra_f16_pp_gemm_pre(into) rc={rc} (m={m} n={out_f} k={in_f})").into(),
579            );
580        }
581        if scale != 1.0 {
582            self.scale_inplace(y, scale, m * out_f)?;
583        }
584        Ok(true)
585    }
586
587    /// `_into` at a ROW OFFSET (task #16): the batched prime's per-seq out-GEMMs write
588    /// straight into the concat `mixed` trunk at offs[s] — removing the per-seq gather
589    /// copy. off_elems must keep the pointer's alignment class (n_embd rows do).
590    pub fn try_f16_gemm_pre_into_off(
591        &self,
592        w: &crate::model::GpuTensor,
593        xh: &CudaSlice<u8>,
594        m: usize,
595        y: &mut CudaSlice<f32>,
596        off_elems: usize,
597    ) -> Result<bool, Box<dyn std::error::Error>> {
598        use crate::model::GpuTensor;
599        let (w16, ne, scale) = match w {
600            GpuTensor::Quant {
601                f16: Some(w16),
602                ne,
603                scale,
604                ..
605            } => (w16, ne, *scale),
606            _ => return Ok(false),
607        };
608        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
609        assert!(
610            y.len() >= off_elems + m * out_f,
611            "try_f16_gemm_pre_into_off: output slab too small"
612        );
613        if scale != 1.0 {
614            return Ok(false); // post-scale would need a strided view; caller falls back
615        }
616        let mut guard = self.f16_scratch.lock().unwrap();
617        if guard.is_none() {
618            *guard = Some(F16Scratch {
619                xh: self.alloc_u8_uninit(2)?,
620                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
621                cap_xh: 2,
622            });
623        }
624        let s = guard.as_mut().unwrap();
625        let rc = {
626            let stream = self.gpu.stream();
627            let (w_p, _gw) = w16.device_ptr(&stream);
628            let (h_p, _gh) = xh.device_ptr(&stream);
629            let (y_p, _gy) = y.device_ptr_mut(&stream);
630            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
631            unsafe {
632                memra_f16_pp_gemm_pre(
633                    w_p as *const core::ffi::c_void,
634                    h_p as *const core::ffi::c_void,
635                    (y_p as *mut f32).add(off_elems),
636                    m as i32,
637                    out_f as i32,
638                    in_f as i32,
639                    ws_p as *mut core::ffi::c_void,
640                    F16_WS_BYTES,
641                    stream.cu_stream() as *mut core::ffi::c_void,
642                )
643            }
644        };
645        if rc != 0 {
646            return Err(format!(
647                "memra_f16_pp_gemm_pre(into_off) rc={rc} (m={m} n={out_f} k={in_f})"
648            )
649            .into());
650        }
651        Ok(true)
652    }
653
654    /// FP16 GEMM on a pre-converted activation — the matmul_group arm. Same contract as
655    /// `try_f16_gemm` minus the convert.
656    pub fn try_f16_gemm_pre(
657        &self,
658        w: &crate::model::GpuTensor,
659        xh: &CudaSlice<u8>,
660        m: usize,
661    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
662        use crate::model::GpuTensor;
663        let (w16, ne, scale) = match w {
664            GpuTensor::Quant {
665                f16: Some(w16),
666                ne,
667                scale,
668                ..
669            } => (w16, ne, *scale),
670            _ => return Ok(None),
671        };
672        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
673        // workspace from the shared scratch (xh is caller-owned here)
674        let mut guard = self.f16_scratch.lock().unwrap();
675        if guard.is_none() {
676            *guard = Some(F16Scratch {
677                xh: self.alloc_u8_uninit(2)?,
678                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
679                cap_xh: 2,
680            });
681        }
682        let s = guard.as_mut().unwrap();
683        let mut y = self.uninit(m * out_f)?;
684        let rc = {
685            let stream = self.gpu.stream();
686            let (w_p, _gw) = w16.device_ptr(&stream);
687            let (h_p, _gh) = xh.device_ptr(&stream);
688            let (y_p, _gy) = y.device_ptr_mut(&stream);
689            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
690            unsafe {
691                memra_f16_pp_gemm_pre(
692                    w_p as *const core::ffi::c_void,
693                    h_p as *const core::ffi::c_void,
694                    y_p as *mut f32,
695                    m as i32,
696                    out_f as i32,
697                    in_f as i32,
698                    ws_p as *mut core::ffi::c_void,
699                    F16_WS_BYTES,
700                    stream.cu_stream() as *mut core::ffi::c_void,
701                )
702            }
703        };
704        if rc != 0 {
705            return Err(format!("memra_f16_pp_gemm_pre rc={rc} (m={m} n={out_f} k={in_f})").into());
706        }
707        if scale != 1.0 {
708            self.scale_inplace(&mut y, scale, m * out_f)?;
709        }
710        Ok(Some(y))
711    }
712
713    /// Raw fp16 mirror build from GGUF Q8_0 device bytes (gates/benches; also the loader's
714    /// worker via `build_q8_f16`).
715    pub fn build_q8_f16_raw(
716        &self,
717        bytes: &CudaSlice<u8>,
718        in_f: usize,
719        out_f: usize,
720    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
721        assert!(in_f.is_multiple_of(32));
722        let nblk = in_f / 32;
723        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
724        let rc = {
725            let stream = self.gpu.stream();
726            let (s_p, _gs) = bytes.device_ptr(&stream);
727            let (d_p, _gd) = dst.device_ptr_mut(&stream);
728            unsafe {
729                memra_q8_0_dequant_f16(
730                    s_p as *const core::ffi::c_void,
731                    d_p as *mut core::ffi::c_void,
732                    out_f as i64,
733                    nblk as i64,
734                    stream.cu_stream() as *mut core::ffi::c_void,
735                )
736            }
737        };
738        if rc != 0 {
739            return Err(format!("memra_q8_0_dequant_f16 rc={rc}").into());
740        }
741        Ok(dst)
742    }
743
744    /// Load-time fp16 mirror pass for one tensor (hybrid.rs calls this under MEMRA_PP_F16=1,
745    /// next to `build_q8_rp4`). No-op unless 2D Q8_0 with integral rows and budget headroom.
746    pub fn build_q8_f16(
747        &self,
748        t: &mut crate::model::GpuTensor,
749    ) -> Result<(), Box<dyn std::error::Error>> {
750        use crate::model::GpuTensor;
751        let GpuTensor::Quant {
752            bytes,
753            qtype,
754            row_bytes,
755            ne,
756            f16,
757            ..
758        } = t
759        else {
760            return Ok(());
761        };
762        // Q4_0 admitted 2026-07-31 (campaign A): the gemma QAT trunk rides the same Lt
763        // f16 lane — int4 magnitudes exact in fp16, same rounding class as Q8_0.
764        // Q6_K admitted round 47: the q27 Q4_K_M mix packs attn_v/ffn_down/head as Q6_K
765        // with NO MMQ arm — its 6.7ms/call dequant-GEMMs were the prefill wall.
766        // Q4_K admitted round 49: the q27 trunk bulk (294 tensors) rides mul_mat_q_q45k
767        // int8-MMA; the Lt f16 lane beats that class at large m (campaign-A precedent).
768        // Q5_K admitted round 49b: q27's 48 ssm_out projections — same MMQ class.
769        let q4 = *qtype == crate::QT_Q4_0;
770        let q6k = *qtype == crate::QT_Q6_K;
771        let q4k = *qtype == crate::QT_Q4_K;
772        let q5k = *qtype == crate::QT_Q5_K;
773        if (*qtype != crate::QT_Q8_0 && !q4 && !q6k && !q4k && !q5k)
774            || f16.is_some()
775            || ne.len() != 2
776        {
777            return Ok(());
778        }
779        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
780        if q6k || q4k || q5k {
781            let sb = if q6k {
782                210
783            } else if q5k {
784                176
785            } else {
786                144
787            };
788            if in_f % 256 != 0 || *row_bytes != (in_f / 256) * sb {
789                return Ok(());
790            }
791        } else if in_f % 32 != 0 || *row_bytes != (in_f / 32) * (if q4 { 18 } else { 34 }) {
792            return Ok(());
793        }
794        // Budget (layer-order prefix, MEMRA_PP_FP8_BUDGET_MB pattern): default 32GB — the whole
795        // 9B mirror on an 80GB box; smaller rigs set MEMRA_PP_F16_BUDGET_MB down.
796        use std::sync::atomic::{AtomicUsize, Ordering};
797        static SPENT: AtomicUsize = AtomicUsize::new(0);
798        static BUDGET: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
799        let budget = *BUDGET.get_or_init(|| {
800            std::env::var("MEMRA_PP_F16_BUDGET_MB")
801                .ok()
802                .and_then(|v| v.parse::<usize>().ok())
803                .unwrap_or(32768)
804                << 20
805        });
806        let sz = out_f * in_f * 2;
807        if SPENT.fetch_add(sz, Ordering::Relaxed) + sz > budget {
808            SPENT.fetch_sub(sz, Ordering::Relaxed);
809            return Ok(());
810        }
811        let mirror = if q6k {
812            self.build_q6k_f16_raw(bytes, in_f, out_f)?
813        } else if q4k {
814            self.build_q4k_f16_raw(bytes, in_f, out_f)?
815        } else if q5k {
816            self.build_q5k_f16_raw(bytes, in_f, out_f)?
817        } else if q4 {
818            self.build_q4_f16_raw(bytes, in_f, out_f)?
819        } else {
820            self.build_q8_f16_raw(bytes, in_f, out_f)?
821        };
822        *f16 = Some(mirror);
823        Ok(())
824    }
825
826    /// Q4_0 twin of `build_q8_f16_raw` (18B blocks, campaign A 2026-07-31).
827    pub fn build_q4_f16_raw(
828        &self,
829        bytes: &CudaSlice<u8>,
830        in_f: usize,
831        out_f: usize,
832    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
833        assert!(in_f.is_multiple_of(32));
834        let nblk = in_f / 32;
835        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
836        let rc = {
837            let stream = self.gpu.stream();
838            let (s_p, _gs) = bytes.device_ptr(&stream);
839            let (d_p, _gd) = dst.device_ptr_mut(&stream);
840            unsafe {
841                memra_q4_0_dequant_f16(
842                    s_p as *const core::ffi::c_void,
843                    d_p as *mut core::ffi::c_void,
844                    out_f as i64,
845                    nblk as i64,
846                    stream.cu_stream() as *mut core::ffi::c_void,
847                )
848            }
849        };
850        if rc != 0 {
851            return Err(format!("memra_q4_0_dequant_f16 rc={rc}").into());
852        }
853        Ok(dst)
854    }
855
856    /// Q5_K twin (176B superblocks, round 49b). Also the kernel_check gate entry for the
857    /// Q5_K f16-mirror class.
858    pub fn build_q5k_f16_raw(
859        &self,
860        bytes: &CudaSlice<u8>,
861        in_f: usize,
862        out_f: usize,
863    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
864        assert!(in_f.is_multiple_of(256));
865        let nsb = in_f / 256;
866        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
867        let rc = {
868            let stream = self.gpu.stream();
869            let (s_p, _gs) = bytes.device_ptr(&stream);
870            let (d_p, _gd) = dst.device_ptr_mut(&stream);
871            unsafe {
872                memra_q5_K_dequant_f16(
873                    s_p as *const core::ffi::c_void,
874                    d_p as *mut core::ffi::c_void,
875                    out_f as i64,
876                    nsb as i64,
877                    stream.cu_stream() as *mut core::ffi::c_void,
878                )
879            }
880        };
881        if rc != 0 {
882            return Err(format!("memra_q5_K_dequant_f16 rc={rc}").into());
883        }
884        Ok(dst)
885    }
886
887    /// Q4_K twin of `build_q6k_f16_raw` (144B superblocks, round 49). Also the kernel_check
888    /// gate entry for the Q4_K f16-mirror class.
889    pub fn build_q4k_f16_raw(
890        &self,
891        bytes: &CudaSlice<u8>,
892        in_f: usize,
893        out_f: usize,
894    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
895        assert!(in_f.is_multiple_of(256));
896        let nsb = in_f / 256;
897        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
898        let rc = {
899            let stream = self.gpu.stream();
900            let (s_p, _gs) = bytes.device_ptr(&stream);
901            let (d_p, _gd) = dst.device_ptr_mut(&stream);
902            unsafe {
903                memra_q4_K_dequant_f16(
904                    s_p as *const core::ffi::c_void,
905                    d_p as *mut core::ffi::c_void,
906                    out_f as i64,
907                    nsb as i64,
908                    stream.cu_stream() as *mut core::ffi::c_void,
909                )
910            }
911        };
912        if rc != 0 {
913            return Err(format!("memra_q4_K_dequant_f16 rc={rc}").into());
914        }
915        Ok(dst)
916    }
917
918    pub fn build_q6k_f16_raw(
919        &self,
920        bytes: &CudaSlice<u8>,
921        in_f: usize,
922        out_f: usize,
923    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
924        assert!(in_f.is_multiple_of(256));
925        let nsb = in_f / 256;
926        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
927        let rc = {
928            let stream = self.gpu.stream();
929            let (s_p, _gs) = bytes.device_ptr(&stream);
930            let (d_p, _gd) = dst.device_ptr_mut(&stream);
931            unsafe {
932                memra_q6_K_dequant_f16(
933                    s_p as *const core::ffi::c_void,
934                    d_p as *mut core::ffi::c_void,
935                    out_f as i64,
936                    nsb as i64,
937                    stream.cu_stream() as *mut core::ffi::c_void,
938                )
939            }
940        };
941        if rc != 0 {
942            return Err(format!("memra_q6_K_dequant_f16 rc={rc}").into());
943        }
944        Ok(dst)
945    }
946}
memra_engine/f16_ffi.rs

memra_engine/
f16_ffi.rs
