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

memra_engine/
f16_ffi.rs
