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        // W8A8 PILOT act half (MEMRA_W8A8_SIM=2): per-TOKEN int8 fake-quant of the
233        // activation rows before the f16 GEMM — with the =1 weight half this models
234        // the full w8a8 numeric class through the unchanged lane. Slow host roundtrip,
235        // pilot only.
236        static SIM_ACT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
237        let sim_act =
238            *SIM_ACT.get_or_init(|| std::env::var("MEMRA_W8A8_SIM").as_deref() == Ok("2"));
239        let mut y = if sim_act {
240            let mut hx = self.dtoh(x)?;
241            hx.truncate(m * in_f);
242            for row in hx.chunks_mut(in_f) {
243                let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
244                if amax > 0.0 {
245                    let d = amax / 127.0;
246                    for v in row.iter_mut() {
247                        *v = (*v / d).round().clamp(-127.0, 127.0) * d;
248                    }
249                }
250            }
251            let xq = self.htod(&hx)?;
252            self.qmatvec_gemm_f16_raw(w16, &xq, m, in_f, out_f)?
253        } else {
254            self.qmatvec_gemm_f16_raw(w16, x, m, in_f, out_f)?
255        };
256        if scale != 1.0 {
257            self.scale_inplace(&mut y, scale, m * out_f)?;
258        }
259        Ok(Some(y))
260    }
261
262    /// BF16 tensor-core prefill GEMM on RESIDENT checkpoint bytes: y[m,out] = x[m,in] @ W^T,
263    /// f32 accumulate. `data` is the untouched row-major [out_f, in_f] bf16 weight — no mirror,
264    /// no dequant, no extra VRAM. Shares the fp16 scratch (a bf16 activation is the same 2 B/elem).
265    pub fn bf16_tc_gemm(
266        &self,
267        data: &CudaSlice<u8>,
268        x: &CudaSlice<f32>,
269        m: usize,
270        in_f: usize,
271        out_f: usize,
272    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
273        let need_xh = m * in_f * 2;
274        let mut guard = self.f16_scratch.lock().unwrap();
275        if guard.is_none() {
276            *guard = Some(F16Scratch {
277                xh: self.alloc_u8_uninit(need_xh)?,
278                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
279                cap_xh: need_xh,
280            });
281        }
282        let s = guard.as_mut().unwrap();
283        if need_xh > s.cap_xh {
284            s.xh = self.alloc_u8_uninit(need_xh)?;
285            s.cap_xh = need_xh;
286        }
287        let mut y = self.uninit(m * out_f)?; // full-overwrite GEMM output: skip memset
288        let rc = {
289            let stream = self.gpu.stream();
290            let (w_p, _gw) = data.device_ptr(&stream);
291            let (x_p, _gx) = x.device_ptr(&stream);
292            let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
293            let (y_p, _gy) = y.device_ptr_mut(&stream);
294            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
295            // cuBLASLt wants 16B-aligned operands and its heuristic does not inspect the
296            // pointer, so an unaligned resident slice reaches cublasLtMatmul and comes back
297            // NOT_SUPPORTED. Decline it here instead.
298            if !(w_p as usize).is_multiple_of(16) {
299                -1
300            } else {
301                unsafe {
302                    memra_bf16_pp_gemm(
303                        w_p as *const core::ffi::c_void,
304                        x_p as *const f32,
305                        h_p as *mut core::ffi::c_void,
306                        y_p as *mut f32,
307                        m as i32,
308                        out_f as i32,
309                        in_f as i32,
310                        ws_p as *mut core::ffi::c_void,
311                        F16_WS_BYTES,
312                        stream.cu_stream() as *mut core::ffi::c_void,
313                    )
314                }
315            }
316        };
317        // NOT a hard error: cuBLASLt refuses some (m,n,k)/alignment combinations that its own
318        // heuristic accepted (measured 2026-08-28: rc=30014 = CUBLAS_STATUS_NOT_SUPPORTED at
319        // m=43 n=4096 k=1024, after the same door served every 4096-token prime shape). The f32
320        // dequant path below this door is always correct, so a refusal DECLINES the shape rather
321        // than failing the request. Announced once per shape so the decline can never be silent —
322        // a door that quietly stops engaging reads exactly like a door that never helped.
323        if rc != 0 {
324            static SAID: std::sync::Mutex<
325                Option<std::collections::HashSet<(usize, usize, usize)>>,
326            > = std::sync::Mutex::new(None);
327            let mut g = SAID.lock().unwrap();
328            let seen = g.get_or_insert_with(std::collections::HashSet::new);
329            if seen.insert((m, out_f, in_f)) {
330                eprintln!(
331                    "[bf16-tc] DECLINED m={m} n={out_f} k={in_f} rc={rc} \
332                     (1xxxx=cudaError convert, 2xxxx=no cublasLt algo, 3xxxx=matmul status, \
333                     4xxxx=cublasLtCreate, -1=weight not 16B-aligned) — this shape falls back to \
334                     the f32 dequant GEMM; every other shape keeps the tensor-core path"
335                );
336            }
337            return Ok(None);
338        }
339        // ENGAGEMENT RECEIPT, and it is not optional. Only the DECLINE was announced, so an arm
340        // with MEMRA_PP_BF16=1 that never actually took this path was indistinguishable from one
341        // that did -- and a correctness gate whose two arms ran the SAME code reports a
342        // byte-identical MATCH for the wrong reason. Announced once per shape, same as the
343        // decline, so it costs one line per distinct GEMM and nothing per token.
344        BF16_TC_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
345        {
346            static ACCEPTED: std::sync::Mutex<
347                Option<std::collections::HashSet<(usize, usize, usize)>>,
348            > = std::sync::Mutex::new(None);
349            let mut g = ACCEPTED.lock().unwrap();
350            let seen = g.get_or_insert_with(std::collections::HashSet::new);
351            if seen.insert((m, out_f, in_f)) {
352                eprintln!(
353                    "[bf16-tc] ENGAGED m={m} n={out_f} k={in_f} (bf16 tensor-core GEMM on resident checkpoint bytes)"
354                );
355            }
356        }
357        Ok(Some(y))
358    }
359
360    /// Bare FP16 GEMM launch on an fp16 mirror — also the kernel_check gate entry.
361    pub fn qmatvec_gemm_f16_raw(
362        &self,
363        w16: &CudaSlice<u8>,
364        x: &CudaSlice<f32>,
365        m: usize,
366        in_f: usize,
367        out_f: usize,
368    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
369        let need_xh = m * in_f * 2;
370        let mut guard = self.f16_scratch.lock().unwrap();
371        if guard.is_none() {
372            *guard = Some(F16Scratch {
373                xh: self.alloc_u8_uninit(need_xh)?,
374                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
375                cap_xh: need_xh,
376            });
377        }
378        let s = guard.as_mut().unwrap();
379        if need_xh > s.cap_xh {
380            s.xh = self.alloc_u8_uninit(need_xh)?;
381            s.cap_xh = need_xh;
382        }
383        let mut y = self.uninit(m * out_f)?; // full-overwrite GEMM output: skip memset
384        let rc = {
385            let stream = self.gpu.stream();
386            let (w_p, _gw) = w16.device_ptr(&stream);
387            let (x_p, _gx) = x.device_ptr(&stream);
388            let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
389            let (y_p, _gy) = y.device_ptr_mut(&stream);
390            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
391            unsafe {
392                memra_f16_pp_gemm(
393                    w_p as *const core::ffi::c_void,
394                    x_p as *const f32,
395                    h_p as *mut core::ffi::c_void,
396                    y_p as *mut f32,
397                    m as i32,
398                    out_f as i32,
399                    in_f as i32,
400                    ws_p as *mut core::ffi::c_void,
401                    F16_WS_BYTES,
402                    stream.cu_stream() as *mut core::ffi::c_void,
403                )
404            }
405        };
406        if rc != 0 {
407            return Err(format!(
408                "memra_f16_pp_gemm rc={rc} (m={m} n={out_f} k={in_f}; 1xxxx=cudaError convert, \
409                 2xxxx=no cublasLt algo, 3xxxx=matmul status)"
410            )
411            .into());
412        }
413        Ok(y)
414    }
415
416    /// f32 -> fp16 activation convert into a fresh buffer (matmul_group: ONE convert feeds
417    /// every mirror-carrying weight in the group; the standalone per-GEMM converts were ~250
418    /// launches/prime of gap-cluster fuel, nsys 2026-07-26).
419    #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
420    pub fn f16_act(
421        &self,
422        x: &CudaSlice<f32>,
423        nelem: usize,
424        in_f: usize,
425    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
426        // W8A8 PILOT act half (MEMRA_W8A8_SIM=2): per-TOKEN int8 fake-quant of the
427        // activation rows before the fp16 convert — every pre-converted GEMM in the
428        // group inherits it. Slow host roundtrip, pilot only; default path unchanged.
429        static SIM_ACT2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
430        if *SIM_ACT2.get_or_init(|| std::env::var("MEMRA_W8A8_SIM").as_deref() == Ok("2"))
431            && in_f > 0
432            && nelem % in_f == 0
433        {
434            static ONCE: std::sync::Once = std::sync::Once::new();
435            ONCE.call_once(|| {
436                eprintln!("[w8a8-sim] act per-token int8 fake-quant ACTIVE (f16_act)")
437            });
438            let mut hx = self.dtoh(x)?;
439            hx.truncate(nelem);
440            for row in hx.chunks_mut(in_f) {
441                let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
442                if amax > 0.0 {
443                    let d = amax / 127.0;
444                    for v in row.iter_mut() {
445                        *v = (*v / d).round().clamp(-127.0, 127.0) * d;
446                    }
447                }
448            }
449            let xq = self.htod(&hx)?;
450            let mut xh = self.alloc_u8_uninit(nelem * 2)?;
451            let rc = {
452                let stream = self.gpu.stream();
453                let (x_p, _gx) = xq.device_ptr(&stream);
454                let (h_p, _gh) = xh.device_ptr_mut(&stream);
455                unsafe {
456                    memra_f16_cvt(
457                        x_p as *const f32,
458                        h_p as *mut core::ffi::c_void,
459                        nelem,
460                        stream.cu_stream() as *mut core::ffi::c_void,
461                    )
462                }
463            };
464            if rc != 0 {
465                return Err(format!("memra_f16_cvt rc={rc}").into());
466            }
467            return Ok(xh);
468        }
469        let mut xh = self.alloc_u8_uninit(nelem * 2)?;
470        let rc = {
471            let stream = self.gpu.stream();
472            let (x_p, _gx) = x.device_ptr(&stream);
473            let (h_p, _gh) = xh.device_ptr_mut(&stream);
474            unsafe {
475                memra_f16_cvt(
476                    x_p as *const f32,
477                    h_p as *mut core::ffi::c_void,
478                    nelem,
479                    stream.cu_stream() as *mut core::ffi::c_void,
480                )
481            }
482        };
483        if rc != 0 {
484            return Err(format!("memra_f16_cvt rc={rc}").into());
485        }
486        Ok(xh)
487    }
488
489    /// `_into` twin of `try_f16_gemm_pre` (piecewise-slab plumbing): the GEMM writes into
490    /// a caller-provided buffer (a resident slab view) instead of a fresh allocation —
491    /// the FFI has always taken the y pointer; only the wrapper allocated. Returns
492    /// Ok(false) when the weight has no mirror (caller falls back and copies).
493    pub fn try_f16_gemm_pre_into(
494        &self,
495        w: &crate::model::GpuTensor,
496        xh: &CudaSlice<u8>,
497        m: usize,
498        y: &mut CudaSlice<f32>,
499    ) -> Result<bool, Box<dyn std::error::Error>> {
500        use crate::model::GpuTensor;
501        let (w16, ne, scale) = match w {
502            GpuTensor::Quant {
503                f16: Some(w16),
504                ne,
505                scale,
506                ..
507            } => (w16, ne, *scale),
508            _ => return Ok(false),
509        };
510        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
511        assert!(
512            y.len() >= m * out_f,
513            "try_f16_gemm_pre_into: output slab too small"
514        );
515        let mut guard = self.f16_scratch.lock().unwrap();
516        if guard.is_none() {
517            *guard = Some(F16Scratch {
518                xh: self.alloc_u8_uninit(2)?,
519                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
520                cap_xh: 2,
521            });
522        }
523        let s = guard.as_mut().unwrap();
524        let rc = {
525            let stream = self.gpu.stream();
526            let (w_p, _gw) = w16.device_ptr(&stream);
527            let (h_p, _gh) = xh.device_ptr(&stream);
528            let (y_p, _gy) = y.device_ptr_mut(&stream);
529            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
530            unsafe {
531                memra_f16_pp_gemm_pre(
532                    w_p as *const core::ffi::c_void,
533                    h_p as *const core::ffi::c_void,
534                    y_p as *mut f32,
535                    m as i32,
536                    out_f as i32,
537                    in_f as i32,
538                    ws_p as *mut core::ffi::c_void,
539                    F16_WS_BYTES,
540                    stream.cu_stream() as *mut core::ffi::c_void,
541                )
542            }
543        };
544        if rc != 0 {
545            return Err(
546                format!("memra_f16_pp_gemm_pre(into) rc={rc} (m={m} n={out_f} k={in_f})").into(),
547            );
548        }
549        if scale != 1.0 {
550            self.scale_inplace(y, scale, m * out_f)?;
551        }
552        Ok(true)
553    }
554
555    /// `_into` at a ROW OFFSET (task #16): the batched prime's per-seq out-GEMMs write
556    /// straight into the concat `mixed` trunk at offs[s] — removing the per-seq gather
557    /// copy. off_elems must keep the pointer's alignment class (n_embd rows do).
558    pub fn try_f16_gemm_pre_into_off(
559        &self,
560        w: &crate::model::GpuTensor,
561        xh: &CudaSlice<u8>,
562        m: usize,
563        y: &mut CudaSlice<f32>,
564        off_elems: usize,
565    ) -> Result<bool, Box<dyn std::error::Error>> {
566        use crate::model::GpuTensor;
567        let (w16, ne, scale) = match w {
568            GpuTensor::Quant {
569                f16: Some(w16),
570                ne,
571                scale,
572                ..
573            } => (w16, ne, *scale),
574            _ => return Ok(false),
575        };
576        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
577        assert!(
578            y.len() >= off_elems + m * out_f,
579            "try_f16_gemm_pre_into_off: output slab too small"
580        );
581        if scale != 1.0 {
582            return Ok(false); // post-scale would need a strided view; caller falls back
583        }
584        let mut guard = self.f16_scratch.lock().unwrap();
585        if guard.is_none() {
586            *guard = Some(F16Scratch {
587                xh: self.alloc_u8_uninit(2)?,
588                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
589                cap_xh: 2,
590            });
591        }
592        let s = guard.as_mut().unwrap();
593        let rc = {
594            let stream = self.gpu.stream();
595            let (w_p, _gw) = w16.device_ptr(&stream);
596            let (h_p, _gh) = xh.device_ptr(&stream);
597            let (y_p, _gy) = y.device_ptr_mut(&stream);
598            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
599            unsafe {
600                memra_f16_pp_gemm_pre(
601                    w_p as *const core::ffi::c_void,
602                    h_p as *const core::ffi::c_void,
603                    (y_p as *mut f32).add(off_elems),
604                    m as i32,
605                    out_f as i32,
606                    in_f as i32,
607                    ws_p as *mut core::ffi::c_void,
608                    F16_WS_BYTES,
609                    stream.cu_stream() as *mut core::ffi::c_void,
610                )
611            }
612        };
613        if rc != 0 {
614            return Err(format!(
615                "memra_f16_pp_gemm_pre(into_off) rc={rc} (m={m} n={out_f} k={in_f})"
616            )
617            .into());
618        }
619        Ok(true)
620    }
621
622    /// FP16 GEMM on a pre-converted activation — the matmul_group arm. Same contract as
623    /// `try_f16_gemm` minus the convert.
624    pub fn try_f16_gemm_pre(
625        &self,
626        w: &crate::model::GpuTensor,
627        xh: &CudaSlice<u8>,
628        m: usize,
629    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
630        use crate::model::GpuTensor;
631        let (w16, ne, scale) = match w {
632            GpuTensor::Quant {
633                f16: Some(w16),
634                ne,
635                scale,
636                ..
637            } => (w16, ne, *scale),
638            _ => return Ok(None),
639        };
640        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
641        // workspace from the shared scratch (xh is caller-owned here)
642        let mut guard = self.f16_scratch.lock().unwrap();
643        if guard.is_none() {
644            *guard = Some(F16Scratch {
645                xh: self.alloc_u8_uninit(2)?,
646                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
647                cap_xh: 2,
648            });
649        }
650        let s = guard.as_mut().unwrap();
651        let mut y = self.uninit(m * out_f)?;
652        let rc = {
653            let stream = self.gpu.stream();
654            let (w_p, _gw) = w16.device_ptr(&stream);
655            let (h_p, _gh) = xh.device_ptr(&stream);
656            let (y_p, _gy) = y.device_ptr_mut(&stream);
657            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
658            unsafe {
659                memra_f16_pp_gemm_pre(
660                    w_p as *const core::ffi::c_void,
661                    h_p as *const core::ffi::c_void,
662                    y_p as *mut f32,
663                    m as i32,
664                    out_f as i32,
665                    in_f as i32,
666                    ws_p as *mut core::ffi::c_void,
667                    F16_WS_BYTES,
668                    stream.cu_stream() as *mut core::ffi::c_void,
669                )
670            }
671        };
672        if rc != 0 {
673            return Err(format!("memra_f16_pp_gemm_pre rc={rc} (m={m} n={out_f} k={in_f})").into());
674        }
675        if scale != 1.0 {
676            self.scale_inplace(&mut y, scale, m * out_f)?;
677        }
678        Ok(Some(y))
679    }
680
681    /// Raw fp16 mirror build from GGUF Q8_0 device bytes (gates/benches; also the loader's
682    /// worker via `build_q8_f16`).
683    pub fn build_q8_f16_raw(
684        &self,
685        bytes: &CudaSlice<u8>,
686        in_f: usize,
687        out_f: usize,
688    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
689        assert!(in_f.is_multiple_of(32));
690        let nblk = in_f / 32;
691        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
692        let rc = {
693            let stream = self.gpu.stream();
694            let (s_p, _gs) = bytes.device_ptr(&stream);
695            let (d_p, _gd) = dst.device_ptr_mut(&stream);
696            unsafe {
697                memra_q8_0_dequant_f16(
698                    s_p as *const core::ffi::c_void,
699                    d_p as *mut core::ffi::c_void,
700                    out_f as i64,
701                    nblk as i64,
702                    stream.cu_stream() as *mut core::ffi::c_void,
703                )
704            }
705        };
706        if rc != 0 {
707            return Err(format!("memra_q8_0_dequant_f16 rc={rc}").into());
708        }
709        Ok(dst)
710    }
711
712    /// Load-time fp16 mirror pass for one tensor (hybrid.rs calls this under MEMRA_PP_F16=1,
713    /// next to `build_q8_rp4`). No-op unless 2D Q8_0 with integral rows and budget headroom.
714    pub fn build_q8_f16(
715        &self,
716        t: &mut crate::model::GpuTensor,
717    ) -> Result<(), Box<dyn std::error::Error>> {
718        use crate::model::GpuTensor;
719        let GpuTensor::Quant {
720            bytes,
721            qtype,
722            row_bytes,
723            ne,
724            f16,
725            ..
726        } = t
727        else {
728            return Ok(());
729        };
730        // Q4_0 admitted 2026-07-31 (campaign A): the gemma QAT trunk rides the same Lt
731        // f16 lane — int4 magnitudes exact in fp16, same rounding class as Q8_0.
732        // Q6_K admitted round 47: the q27 Q4_K_M mix packs attn_v/ffn_down/head as Q6_K
733        // with NO MMQ arm — its 6.7ms/call dequant-GEMMs were the prefill wall.
734        // Q4_K admitted round 49: the q27 trunk bulk (294 tensors) rides mul_mat_q_q45k
735        // int8-MMA; the Lt f16 lane beats that class at large m (campaign-A precedent).
736        // Q5_K admitted round 49b: q27's 48 ssm_out projections — same MMQ class.
737        let q4 = *qtype == crate::QT_Q4_0;
738        let q6k = *qtype == crate::QT_Q6_K;
739        let q4k = *qtype == crate::QT_Q4_K;
740        let q5k = *qtype == crate::QT_Q5_K;
741        if (*qtype != crate::QT_Q8_0 && !q4 && !q6k && !q4k && !q5k)
742            || f16.is_some()
743            || ne.len() != 2
744        {
745            return Ok(());
746        }
747        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
748        if q6k || q4k || q5k {
749            let sb = if q6k {
750                210
751            } else if q5k {
752                176
753            } else {
754                144
755            };
756            if in_f % 256 != 0 || *row_bytes != (in_f / 256) * sb {
757                return Ok(());
758            }
759        } else if in_f % 32 != 0 || *row_bytes != (in_f / 32) * (if q4 { 18 } else { 34 }) {
760            return Ok(());
761        }
762        // Budget (layer-order prefix, MEMRA_PP_FP8_BUDGET_MB pattern): default 32GB — the whole
763        // 9B mirror on an 80GB box; smaller rigs set MEMRA_PP_F16_BUDGET_MB down.
764        use std::sync::atomic::{AtomicUsize, Ordering};
765        static SPENT: AtomicUsize = AtomicUsize::new(0);
766        static BUDGET: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
767        let budget = *BUDGET.get_or_init(|| {
768            std::env::var("MEMRA_PP_F16_BUDGET_MB")
769                .ok()
770                .and_then(|v| v.parse::<usize>().ok())
771                .unwrap_or(32768)
772                << 20
773        });
774        let sz = out_f * in_f * 2;
775        if SPENT.fetch_add(sz, Ordering::Relaxed) + sz > budget {
776            SPENT.fetch_sub(sz, Ordering::Relaxed);
777            return Ok(());
778        }
779        let mut mirror = if q6k {
780            self.build_q6k_f16_raw(bytes, in_f, out_f)?
781        } else if q4k {
782            self.build_q4k_f16_raw(bytes, in_f, out_f)?
783        } else if q5k {
784            self.build_q5k_f16_raw(bytes, in_f, out_f)?
785        } else if q4 {
786            self.build_q4_f16_raw(bytes, in_f, out_f)?
787        } else {
788            self.build_q8_f16_raw(bytes, in_f, out_f)?
789        };
790        // W8A8 ACCURACY PILOT (MEMRA_W8A8_SIM=1, 2026-07-31, round-41 arc): fake-quant
791        // the mirror per ROW to int8 (absmax) and round-trip back to f16 — the exact
792        // weight-precision class of the proposed w8a8 crossing (per-row scales replacing
793        // per-32-block), run through the UNCHANGED f16 GEMM lane. Slow host pass, sim
794        // only; the pilot compares greedy streams vs the default config to price the
795        // accuracy relaxation with receipts. Activations stay f16 in this step (the
796        // act-int8 half is additive and strictly smaller — per-token absmax on smooth
797        // activations).
798        static SIM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
799        if *SIM.get_or_init(|| {
800            matches!(
801                std::env::var("MEMRA_W8A8_SIM").as_deref(),
802                Ok("1") | Ok("2")
803            )
804        }) {
805            fn f16_bits_to_f32(b: u16) -> f32 {
806                let (s, e, m) = (
807                    (b >> 15) as u32,
808                    ((b >> 10) & 0x1f) as u32,
809                    (b & 0x3ff) as u32,
810                );
811                let bits = if e == 0 {
812                    if m == 0 {
813                        s << 31
814                    } else {
815                        // subnormal: normalize
816                        let mut e2 = 127 - 15 + 1;
817                        let mut m2 = m;
818                        while m2 & 0x400 == 0 {
819                            m2 <<= 1;
820                            e2 -= 1;
821                        }
822                        (s << 31) | ((e2 as u32) << 23) | ((m2 & 0x3ff) << 13)
823                    }
824                } else if e == 0x1f {
825                    (s << 31) | (0xff << 23) | (m << 13)
826                } else {
827                    (s << 31) | ((e + 127 - 15) << 23) | (m << 13)
828                };
829                f32::from_bits(bits)
830            }
831            fn f32_to_f16_bits(v: f32) -> u16 {
832                let b = v.to_bits();
833                let (s, e, m) = ((b >> 31) as u16, ((b >> 23) & 0xff) as i32, b & 0x7fffff);
834                if e == 0xff {
835                    return (s << 15) | 0x7c00 | ((m >> 13) as u16 & 0x3ff);
836                }
837                let e2 = e - 127 + 15;
838                if e2 >= 0x1f {
839                    return (s << 15) | 0x7c00;
840                }
841                if e2 <= 0 {
842                    if e2 < -10 {
843                        return s << 15;
844                    }
845                    let m2 = (m | 0x800000) >> (1 - e2);
846                    // round-to-nearest-even on the shifted mantissa
847                    let r = (m2 >> 13) as u16 + ((m2 >> 12) & 1) as u16;
848                    return (s << 15) | r;
849                }
850                let mut r = ((e2 as u32) << 10) as u16 | (m >> 13) as u16;
851                if m & 0x1000 != 0 {
852                    r += 1;
853                }
854                (s << 15) | r
855            }
856            let host: Vec<u8> = self.dtoh_u8(&mirror)?;
857            let mut vals: Vec<f32> = host
858                .chunks_exact(2)
859                .map(|c| f16_bits_to_f32(u16::from_le_bytes([c[0], c[1]])))
860                .collect();
861            for row in vals.chunks_mut(in_f) {
862                let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
863                if amax > 0.0 {
864                    let d = amax / 127.0;
865                    for v in row.iter_mut() {
866                        *v = (*v / d).round().clamp(-127.0, 127.0) * d;
867                    }
868                }
869            }
870            let out: Vec<u8> = vals
871                .iter()
872                .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes())
873                .collect();
874            mirror = self.htod_bytes(&out)?;
875        }
876        *f16 = Some(mirror);
877        Ok(())
878    }
879
880    /// Q4_0 twin of `build_q8_f16_raw` (18B blocks, campaign A 2026-07-31).
881    pub fn build_q4_f16_raw(
882        &self,
883        bytes: &CudaSlice<u8>,
884        in_f: usize,
885        out_f: usize,
886    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
887        assert!(in_f.is_multiple_of(32));
888        let nblk = in_f / 32;
889        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
890        let rc = {
891            let stream = self.gpu.stream();
892            let (s_p, _gs) = bytes.device_ptr(&stream);
893            let (d_p, _gd) = dst.device_ptr_mut(&stream);
894            unsafe {
895                memra_q4_0_dequant_f16(
896                    s_p as *const core::ffi::c_void,
897                    d_p as *mut core::ffi::c_void,
898                    out_f as i64,
899                    nblk as i64,
900                    stream.cu_stream() as *mut core::ffi::c_void,
901                )
902            }
903        };
904        if rc != 0 {
905            return Err(format!("memra_q4_0_dequant_f16 rc={rc}").into());
906        }
907        Ok(dst)
908    }
909
910    /// Q5_K twin (176B superblocks, round 49b). Also the kernel_check gate entry for the
911    /// Q5_K f16-mirror class.
912    pub fn build_q5k_f16_raw(
913        &self,
914        bytes: &CudaSlice<u8>,
915        in_f: usize,
916        out_f: usize,
917    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
918        assert!(in_f.is_multiple_of(256));
919        let nsb = in_f / 256;
920        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
921        let rc = {
922            let stream = self.gpu.stream();
923            let (s_p, _gs) = bytes.device_ptr(&stream);
924            let (d_p, _gd) = dst.device_ptr_mut(&stream);
925            unsafe {
926                memra_q5_K_dequant_f16(
927                    s_p as *const core::ffi::c_void,
928                    d_p as *mut core::ffi::c_void,
929                    out_f as i64,
930                    nsb as i64,
931                    stream.cu_stream() as *mut core::ffi::c_void,
932                )
933            }
934        };
935        if rc != 0 {
936            return Err(format!("memra_q5_K_dequant_f16 rc={rc}").into());
937        }
938        Ok(dst)
939    }
940
941    /// Q4_K twin of `build_q6k_f16_raw` (144B superblocks, round 49). Also the kernel_check
942    /// gate entry for the Q4_K f16-mirror class.
943    pub fn build_q4k_f16_raw(
944        &self,
945        bytes: &CudaSlice<u8>,
946        in_f: usize,
947        out_f: usize,
948    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
949        assert!(in_f.is_multiple_of(256));
950        let nsb = in_f / 256;
951        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
952        let rc = {
953            let stream = self.gpu.stream();
954            let (s_p, _gs) = bytes.device_ptr(&stream);
955            let (d_p, _gd) = dst.device_ptr_mut(&stream);
956            unsafe {
957                memra_q4_K_dequant_f16(
958                    s_p as *const core::ffi::c_void,
959                    d_p as *mut core::ffi::c_void,
960                    out_f as i64,
961                    nsb as i64,
962                    stream.cu_stream() as *mut core::ffi::c_void,
963                )
964            }
965        };
966        if rc != 0 {
967            return Err(format!("memra_q4_K_dequant_f16 rc={rc}").into());
968        }
969        Ok(dst)
970    }
971
972    pub fn build_q6k_f16_raw(
973        &self,
974        bytes: &CudaSlice<u8>,
975        in_f: usize,
976        out_f: usize,
977    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
978        assert!(in_f.is_multiple_of(256));
979        let nsb = in_f / 256;
980        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
981        let rc = {
982            let stream = self.gpu.stream();
983            let (s_p, _gs) = bytes.device_ptr(&stream);
984            let (d_p, _gd) = dst.device_ptr_mut(&stream);
985            unsafe {
986                memra_q6_K_dequant_f16(
987                    s_p as *const core::ffi::c_void,
988                    d_p as *mut core::ffi::c_void,
989                    out_f as i64,
990                    nsb as i64,
991                    stream.cu_stream() as *mut core::ffi::c_void,
992                )
993            }
994        };
995        if rc != 0 {
996            return Err(format!("memra_q6_K_dequant_f16 rc={rc}").into());
997        }
998        Ok(dst)
999    }
1000}
memra_engine/f16_ffi.rs

memra_engine/
f16_ffi.rs
