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    /// GGUF Q8_0 34B blocks -> row-major fp16 mirror (load-time).
54    fn memra_q8_0_dequant_f16(
55        w_q8: *const core::ffi::c_void,
56        w_f16: *mut core::ffi::c_void,
57        out_f: i64,
58        nblk_row: i64,
59        stream: *mut core::ffi::c_void,
60    ) -> i32;
61    /// GGUF Q4_0 18B blocks -> row-major fp16 mirror (campaign A, 2026-07-31).
62    fn memra_q4_0_dequant_f16(
63        w_q4: *const core::ffi::c_void,
64        w_f16: *mut core::ffi::c_void,
65        out_f: i64,
66        nblk_row: i64,
67        stream: *mut core::ffi::c_void,
68    ) -> i32;
69    /// GGUF Q6_K 210B superblocks -> row-major fp16 mirror (round 47: the q27 prefill wall).
70    fn memra_q6_K_dequant_f16(
71        w_q6: *const core::ffi::c_void,
72        w_f16: *mut core::ffi::c_void,
73        out_f: i64,
74        nsb_row: i64,
75        stream: *mut core::ffi::c_void,
76    ) -> i32;
77    /// GGUF Q4_K 144B superblocks -> row-major fp16 mirror (round 49: the q27 trunk bulk).
78    fn memra_q4_K_dequant_f16(
79        w_q4k: *const core::ffi::c_void,
80        w_f16: *mut core::ffi::c_void,
81        out_f: i64,
82        nsb_row: i64,
83        stream: *mut core::ffi::c_void,
84    ) -> i32;
85    /// GGUF Q5_K 176B superblocks -> row-major fp16 mirror (round 49b: q27 ssm_out).
86    fn memra_q5_K_dequant_f16(
87        w_q5k: *const core::ffi::c_void,
88        w_f16: *mut core::ffi::c_void,
89        out_f: i64,
90        nsb_row: i64,
91        stream: *mut core::ffi::c_void,
92    ) -> i32;
93}
94
95/// MEMRA_PP_F16 gate, read once. DEFAULT ON on the Hopper lane (80GB — the mirror costs
96/// 2 B/w, ~17GB on the 9B; the box carries it), opt-in elsewhere; =1/=0 overrides either way.
97/// Promotion battery (2026-07-26, H100): kernel-check ALL GREEN (f16 rel <= 6.5e-3, band 1e-2);
98/// run-gen argmax MATCH on p1/p2/p3 long prompts; greedy streams IDENTICAL to the MMQ config
99/// on all three; pp512 8674 -> 15626 tok/s (+80%, N=5 medians). Decode untouched (m>=16 arm).
100pub fn pp_f16_enabled() -> bool {
101    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
102    *ON.get_or_init(|| match std::env::var("MEMRA_PP_F16").as_deref() {
103        Ok("1") => true,
104        Ok("0") => false,
105        _ => cfg!(memra_hopper_mma),
106    })
107}
108
109/// Resident scratch (fp8_ffi::Fp8Scratch pattern): fp16 activation (grown to the largest m*k
110/// seen) + the cuBLASLt workspace. Single GPU worker; the Mutex guards lazy build/grow only.
111pub struct F16Scratch {
112    pub xh: CudaSlice<u8>,
113    pub ws: CudaSlice<u8>,
114    cap_xh: usize,
115}
116
117impl F16Scratch {
118    /// Pre-sized scratch (task #14: the captured prime gets a PRIVATE scratch so the
119    /// graph's baked cvt/Lt pointers are never mutated by eager GEMMs between replays).
120    pub fn with_capacity(e: &crate::Engine, xh_bytes: usize)
121                         -> Result<Self, Box<dyn std::error::Error>> {
122        Ok(F16Scratch {
123            xh: e.alloc_u8_uninit(xh_bytes)?,
124            ws: e.alloc_u8_uninit(F16_WS_BYTES)?,
125            cap_xh: xh_bytes,
126        })
127    }
128}
129
130const F16_WS_BYTES: usize = 64 << 20;
131
132impl crate::Engine {
133    /// Swap the resident f16 scratch (task #14 capture isolation). Returns the previous
134    /// contents; pass them back to restore.
135    pub fn f16_scratch_swap(&self, new: Option<F16Scratch>) -> Option<F16Scratch> {
136        std::mem::replace(&mut *self.f16_scratch.lock().unwrap(), new)
137    }
138
139    /// FP16 prefill GEMM for a weight carrying the f16 mirror: y[m,out] = x[m,in] @ (fp16 W)^T,
140    /// f32 accumulate. Returns None when the weight has no mirror (caller falls through to MMQ).
141    pub fn try_f16_gemm(
142        &self,
143        w: &crate::model::GpuTensor,
144        x: &CudaSlice<f32>,
145        m: usize,
146    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
147        use crate::model::GpuTensor;
148        let (w16, ne, scale) = match w {
149            GpuTensor::Quant {
150                f16: Some(w16),
151                ne,
152                scale,
153                ..
154            } => (w16, ne, *scale),
155            _ => return Ok(None),
156        };
157        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
158        // W8A8 PILOT act half (MEMRA_W8A8_SIM=2): per-TOKEN int8 fake-quant of the
159        // activation rows before the f16 GEMM — with the =1 weight half this models
160        // the full w8a8 numeric class through the unchanged lane. Slow host roundtrip,
161        // pilot only.
162        static SIM_ACT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163        let sim_act = *SIM_ACT.get_or_init(||
164            std::env::var("MEMRA_W8A8_SIM").as_deref() == Ok("2"));
165        let mut y = if sim_act {
166            let mut hx = self.dtoh(x)?;
167            hx.truncate(m * in_f);
168            for row in hx.chunks_mut(in_f) {
169                let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
170                if amax > 0.0 {
171                    let d = amax / 127.0;
172                    for v in row.iter_mut() {
173                        *v = (*v / d).round().clamp(-127.0, 127.0) * d;
174                    }
175                }
176            }
177            let xq = self.htod(&hx)?;
178            self.qmatvec_gemm_f16_raw(w16, &xq, m, in_f, out_f)?
179        } else {
180            self.qmatvec_gemm_f16_raw(w16, x, m, in_f, out_f)?
181        };
182        if scale != 1.0 {
183            self.scale_inplace(&mut y, scale, m * out_f)?;
184        }
185        Ok(Some(y))
186    }
187
188    /// Bare FP16 GEMM launch on an fp16 mirror — also the kernel_check gate entry.
189    pub fn qmatvec_gemm_f16_raw(
190        &self,
191        w16: &CudaSlice<u8>,
192        x: &CudaSlice<f32>,
193        m: usize,
194        in_f: usize,
195        out_f: usize,
196    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
197        let need_xh = m * in_f * 2;
198        let mut guard = self.f16_scratch.lock().unwrap();
199        if guard.is_none() {
200            *guard = Some(F16Scratch {
201                xh: self.alloc_u8_uninit(need_xh)?,
202                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
203                cap_xh: need_xh,
204            });
205        }
206        let s = guard.as_mut().unwrap();
207        if need_xh > s.cap_xh {
208            s.xh = self.alloc_u8_uninit(need_xh)?;
209            s.cap_xh = need_xh;
210        }
211        let mut y = self.uninit(m * out_f)?; // full-overwrite GEMM output: skip memset
212        let rc = {
213            let stream = self.gpu.stream();
214            let (w_p, _gw) = w16.device_ptr(&stream);
215            let (x_p, _gx) = x.device_ptr(&stream);
216            let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
217            let (y_p, _gy) = y.device_ptr_mut(&stream);
218            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
219            unsafe {
220                memra_f16_pp_gemm(
221                    w_p as *const core::ffi::c_void,
222                    x_p as *const f32,
223                    h_p as *mut core::ffi::c_void,
224                    y_p as *mut f32,
225                    m as i32,
226                    out_f as i32,
227                    in_f as i32,
228                    ws_p as *mut core::ffi::c_void,
229                    F16_WS_BYTES,
230                    stream.cu_stream() as *mut core::ffi::c_void,
231                )
232            }
233        };
234        if rc != 0 {
235            return Err(format!(
236                "memra_f16_pp_gemm rc={rc} (m={m} n={out_f} k={in_f}; 1xxxx=cudaError convert, \
237                 2xxxx=no cublasLt algo, 3xxxx=matmul status)"
238            )
239            .into());
240        }
241        Ok(y)
242    }
243
244    /// f32 -> fp16 activation convert into a fresh buffer (matmul_group: ONE convert feeds
245    /// every mirror-carrying weight in the group; the standalone per-GEMM converts were ~250
246    /// launches/prime of gap-cluster fuel, nsys 2026-07-26).
247    pub fn f16_act(
248        &self,
249        x: &CudaSlice<f32>,
250        nelem: usize,
251        in_f: usize,
252    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
253        // W8A8 PILOT act half (MEMRA_W8A8_SIM=2): per-TOKEN int8 fake-quant of the
254        // activation rows before the fp16 convert — every pre-converted GEMM in the
255        // group inherits it. Slow host roundtrip, pilot only; default path unchanged.
256        static SIM_ACT2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
257        if *SIM_ACT2.get_or_init(|| std::env::var("MEMRA_W8A8_SIM").as_deref() == Ok("2"))
258            && in_f > 0 && nelem % in_f == 0 {
259            static ONCE: std::sync::Once = std::sync::Once::new();
260            ONCE.call_once(|| eprintln!("[w8a8-sim] act per-token int8 fake-quant ACTIVE (f16_act)"));
261            let mut hx = self.dtoh(x)?;
262            hx.truncate(nelem);
263            for row in hx.chunks_mut(in_f) {
264                let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
265                if amax > 0.0 {
266                    let d = amax / 127.0;
267                    for v in row.iter_mut() {
268                        *v = (*v / d).round().clamp(-127.0, 127.0) * d;
269                    }
270                }
271            }
272            let xq = self.htod(&hx)?;
273            let mut xh = self.alloc_u8_uninit(nelem * 2)?;
274            let rc = {
275                let stream = self.gpu.stream();
276                let (x_p, _gx) = xq.device_ptr(&stream);
277                let (h_p, _gh) = xh.device_ptr_mut(&stream);
278                unsafe {
279                    memra_f16_cvt(x_p as *const f32, h_p as *mut core::ffi::c_void,
280                                  nelem, stream.cu_stream() as *mut core::ffi::c_void)
281                }
282            };
283            if rc != 0 { return Err(format!("memra_f16_cvt rc={rc}").into()); }
284            return Ok(xh);
285        }
286        let mut xh = self.alloc_u8_uninit(nelem * 2)?;
287        let rc = {
288            let stream = self.gpu.stream();
289            let (x_p, _gx) = x.device_ptr(&stream);
290            let (h_p, _gh) = xh.device_ptr_mut(&stream);
291            unsafe {
292                memra_f16_cvt(
293                    x_p as *const f32,
294                    h_p as *mut core::ffi::c_void,
295                    nelem,
296                    stream.cu_stream() as *mut core::ffi::c_void,
297                )
298            }
299        };
300        if rc != 0 {
301            return Err(format!("memra_f16_cvt rc={rc}").into());
302        }
303        Ok(xh)
304    }
305
306    /// `_into` twin of `try_f16_gemm_pre` (piecewise-slab plumbing): the GEMM writes into
307    /// a caller-provided buffer (a resident slab view) instead of a fresh allocation —
308    /// the FFI has always taken the y pointer; only the wrapper allocated. Returns
309    /// Ok(false) when the weight has no mirror (caller falls back and copies).
310    pub fn try_f16_gemm_pre_into(
311        &self,
312        w: &crate::model::GpuTensor,
313        xh: &CudaSlice<u8>,
314        m: usize,
315        y: &mut CudaSlice<f32>,
316    ) -> Result<bool, Box<dyn std::error::Error>> {
317        use crate::model::GpuTensor;
318        let (w16, ne, scale) = match w {
319            GpuTensor::Quant { f16: Some(w16), ne, scale, .. } => (w16, ne, *scale),
320            _ => return Ok(false),
321        };
322        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
323        assert!(y.len() >= m * out_f, "try_f16_gemm_pre_into: output slab too small");
324        let mut guard = self.f16_scratch.lock().unwrap();
325        if guard.is_none() {
326            *guard = Some(F16Scratch {
327                xh: self.alloc_u8_uninit(2)?,
328                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
329                cap_xh: 2,
330            });
331        }
332        let s = guard.as_mut().unwrap();
333        let rc = {
334            let stream = self.gpu.stream();
335            let (w_p, _gw) = w16.device_ptr(&stream);
336            let (h_p, _gh) = xh.device_ptr(&stream);
337            let (y_p, _gy) = y.device_ptr_mut(&stream);
338            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
339            unsafe {
340                memra_f16_pp_gemm_pre(
341                    w_p as *const core::ffi::c_void,
342                    h_p as *const core::ffi::c_void,
343                    y_p as *mut f32,
344                    m as i32,
345                    out_f as i32,
346                    in_f as i32,
347                    ws_p as *mut core::ffi::c_void,
348                    F16_WS_BYTES,
349                    stream.cu_stream() as *mut core::ffi::c_void,
350                )
351            }
352        };
353        if rc != 0 {
354            return Err(format!("memra_f16_pp_gemm_pre(into) rc={rc} (m={m} n={out_f} k={in_f})").into());
355        }
356        if scale != 1.0 {
357            self.scale_inplace(y, scale, m * out_f)?;
358        }
359        Ok(true)
360    }
361
362    /// `_into` at a ROW OFFSET (task #16): the batched prime's per-seq out-GEMMs write
363    /// straight into the concat `mixed` trunk at offs[s] — removing the per-seq gather
364    /// copy. off_elems must keep the pointer's alignment class (n_embd rows do).
365    pub fn try_f16_gemm_pre_into_off(
366        &self,
367        w: &crate::model::GpuTensor,
368        xh: &CudaSlice<u8>,
369        m: usize,
370        y: &mut CudaSlice<f32>,
371        off_elems: usize,
372    ) -> Result<bool, Box<dyn std::error::Error>> {
373        use crate::model::GpuTensor;
374        let (w16, ne, scale) = match w {
375            GpuTensor::Quant { f16: Some(w16), ne, scale, .. } => (w16, ne, *scale),
376            _ => return Ok(false),
377        };
378        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
379        assert!(y.len() >= off_elems + m * out_f, "try_f16_gemm_pre_into_off: output slab too small");
380        if scale != 1.0 {
381            return Ok(false); // post-scale would need a strided view; caller falls back
382        }
383        let mut guard = self.f16_scratch.lock().unwrap();
384        if guard.is_none() {
385            *guard = Some(F16Scratch {
386                xh: self.alloc_u8_uninit(2)?,
387                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
388                cap_xh: 2,
389            });
390        }
391        let s = guard.as_mut().unwrap();
392        let rc = {
393            let stream = self.gpu.stream();
394            let (w_p, _gw) = w16.device_ptr(&stream);
395            let (h_p, _gh) = xh.device_ptr(&stream);
396            let (y_p, _gy) = y.device_ptr_mut(&stream);
397            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
398            unsafe {
399                memra_f16_pp_gemm_pre(
400                    w_p as *const core::ffi::c_void,
401                    h_p as *const core::ffi::c_void,
402                    (y_p as *mut f32).add(off_elems),
403                    m as i32,
404                    out_f as i32,
405                    in_f as i32,
406                    ws_p as *mut core::ffi::c_void,
407                    F16_WS_BYTES,
408                    stream.cu_stream() as *mut core::ffi::c_void,
409                )
410            }
411        };
412        if rc != 0 {
413            return Err(format!("memra_f16_pp_gemm_pre(into_off) rc={rc} (m={m} n={out_f} k={in_f})").into());
414        }
415        Ok(true)
416    }
417
418    /// FP16 GEMM on a pre-converted activation — the matmul_group arm. Same contract as
419    /// `try_f16_gemm` minus the convert.
420    pub fn try_f16_gemm_pre(
421        &self,
422        w: &crate::model::GpuTensor,
423        xh: &CudaSlice<u8>,
424        m: usize,
425    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
426        use crate::model::GpuTensor;
427        let (w16, ne, scale) = match w {
428            GpuTensor::Quant {
429                f16: Some(w16),
430                ne,
431                scale,
432                ..
433            } => (w16, ne, *scale),
434            _ => return Ok(None),
435        };
436        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
437        // workspace from the shared scratch (xh is caller-owned here)
438        let mut guard = self.f16_scratch.lock().unwrap();
439        if guard.is_none() {
440            *guard = Some(F16Scratch {
441                xh: self.alloc_u8_uninit(2)?,
442                ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
443                cap_xh: 2,
444            });
445        }
446        let s = guard.as_mut().unwrap();
447        let mut y = self.uninit(m * out_f)?;
448        let rc = {
449            let stream = self.gpu.stream();
450            let (w_p, _gw) = w16.device_ptr(&stream);
451            let (h_p, _gh) = xh.device_ptr(&stream);
452            let (y_p, _gy) = y.device_ptr_mut(&stream);
453            let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
454            unsafe {
455                memra_f16_pp_gemm_pre(
456                    w_p as *const core::ffi::c_void,
457                    h_p as *const core::ffi::c_void,
458                    y_p as *mut f32,
459                    m as i32,
460                    out_f as i32,
461                    in_f as i32,
462                    ws_p as *mut core::ffi::c_void,
463                    F16_WS_BYTES,
464                    stream.cu_stream() as *mut core::ffi::c_void,
465                )
466            }
467        };
468        if rc != 0 {
469            return Err(format!("memra_f16_pp_gemm_pre rc={rc} (m={m} n={out_f} k={in_f})").into());
470        }
471        if scale != 1.0 {
472            self.scale_inplace(&mut y, scale, m * out_f)?;
473        }
474        Ok(Some(y))
475    }
476
477    /// Raw fp16 mirror build from GGUF Q8_0 device bytes (gates/benches; also the loader's
478    /// worker via `build_q8_f16`).
479    pub fn build_q8_f16_raw(
480        &self,
481        bytes: &CudaSlice<u8>,
482        in_f: usize,
483        out_f: usize,
484    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
485        assert!(in_f % 32 == 0);
486        let nblk = in_f / 32;
487        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
488        let rc = {
489            let stream = self.gpu.stream();
490            let (s_p, _gs) = bytes.device_ptr(&stream);
491            let (d_p, _gd) = dst.device_ptr_mut(&stream);
492            unsafe {
493                memra_q8_0_dequant_f16(
494                    s_p as *const core::ffi::c_void,
495                    d_p as *mut core::ffi::c_void,
496                    out_f as i64,
497                    nblk as i64,
498                    stream.cu_stream() as *mut core::ffi::c_void,
499                )
500            }
501        };
502        if rc != 0 {
503            return Err(format!("memra_q8_0_dequant_f16 rc={rc}").into());
504        }
505        Ok(dst)
506    }
507
508    /// Load-time fp16 mirror pass for one tensor (hybrid.rs calls this under MEMRA_PP_F16=1,
509    /// next to `build_q8_rp4`). No-op unless 2D Q8_0 with integral rows and budget headroom.
510    pub fn build_q8_f16(
511        &self,
512        t: &mut crate::model::GpuTensor,
513    ) -> Result<(), Box<dyn std::error::Error>> {
514        use crate::model::GpuTensor;
515        let GpuTensor::Quant {
516            bytes,
517            qtype,
518            row_bytes,
519            ne,
520            f16,
521            ..
522        } = t
523        else {
524            return Ok(());
525        };
526        // Q4_0 admitted 2026-07-31 (campaign A): the gemma QAT trunk rides the same Lt
527        // f16 lane — int4 magnitudes exact in fp16, same rounding class as Q8_0.
528        // Q6_K admitted round 47: the q27 Q4_K_M mix packs attn_v/ffn_down/head as Q6_K
529        // with NO MMQ arm — its 6.7ms/call dequant-GEMMs were the prefill wall.
530        // Q4_K admitted round 49: the q27 trunk bulk (294 tensors) rides mul_mat_q_q45k
531        // int8-MMA; the Lt f16 lane beats that class at large m (campaign-A precedent).
532        // Q5_K admitted round 49b: q27's 48 ssm_out projections — same MMQ class.
533        let q4 = *qtype == crate::QT_Q4_0;
534        let q6k = *qtype == crate::QT_Q6_K;
535        let q4k = *qtype == crate::QT_Q4_K;
536        let q5k = *qtype == crate::QT_Q5_K;
537        if (*qtype != crate::QT_Q8_0 && !q4 && !q6k && !q4k && !q5k)
538            || f16.is_some() || ne.len() != 2 {
539            return Ok(());
540        }
541        let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
542        if q6k || q4k || q5k {
543            let sb = if q6k { 210 } else if q5k { 176 } else { 144 };
544            if in_f % 256 != 0 || *row_bytes != (in_f / 256) * sb {
545                return Ok(());
546            }
547        } else if in_f % 32 != 0 || *row_bytes != (in_f / 32) * (if q4 { 18 } else { 34 }) {
548            return Ok(());
549        }
550        // Budget (layer-order prefix, MEMRA_PP_FP8_BUDGET_MB pattern): default 32GB — the whole
551        // 9B mirror on an 80GB box; smaller rigs set MEMRA_PP_F16_BUDGET_MB down.
552        use std::sync::atomic::{AtomicUsize, Ordering};
553        static SPENT: AtomicUsize = AtomicUsize::new(0);
554        static BUDGET: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
555        let budget = *BUDGET.get_or_init(|| {
556            std::env::var("MEMRA_PP_F16_BUDGET_MB")
557                .ok()
558                .and_then(|v| v.parse::<usize>().ok())
559                .unwrap_or(32768)
560                << 20
561        });
562        let sz = out_f * in_f * 2;
563        if SPENT.fetch_add(sz, Ordering::Relaxed) + sz > budget {
564            SPENT.fetch_sub(sz, Ordering::Relaxed);
565            return Ok(());
566        }
567        let mut mirror = if q6k { self.build_q6k_f16_raw(bytes, in_f, out_f)? }
568                         else if q4k { self.build_q4k_f16_raw(bytes, in_f, out_f)? }
569                         else if q5k { self.build_q5k_f16_raw(bytes, in_f, out_f)? }
570                         else if q4 { self.build_q4_f16_raw(bytes, in_f, out_f)? }
571                         else { self.build_q8_f16_raw(bytes, in_f, out_f)? };
572        // W8A8 ACCURACY PILOT (MEMRA_W8A8_SIM=1, 2026-07-31, round-41 arc): fake-quant
573        // the mirror per ROW to int8 (absmax) and round-trip back to f16 — the exact
574        // weight-precision class of the proposed w8a8 crossing (per-row scales replacing
575        // per-32-block), run through the UNCHANGED f16 GEMM lane. Slow host pass, sim
576        // only; the pilot compares greedy streams vs the default config to price the
577        // accuracy relaxation with receipts. Activations stay f16 in this step (the
578        // act-int8 half is additive and strictly smaller — per-token absmax on smooth
579        // activations).
580        static SIM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
581        if *SIM.get_or_init(|| matches!(std::env::var("MEMRA_W8A8_SIM").as_deref(), Ok("1") | Ok("2"))) {
582            fn f16_bits_to_f32(b: u16) -> f32 {
583                let (s, e, m) = ((b >> 15) as u32, ((b >> 10) & 0x1f) as u32, (b & 0x3ff) as u32);
584                let bits = if e == 0 {
585                    if m == 0 { s << 31 } else {
586                        // subnormal: normalize
587                        let mut e2 = 127 - 15 + 1;
588                        let mut m2 = m;
589                        while m2 & 0x400 == 0 { m2 <<= 1; e2 -= 1; }
590                        (s << 31) | ((e2 as u32) << 23) | ((m2 & 0x3ff) << 13)
591                    }
592                } else if e == 0x1f {
593                    (s << 31) | (0xff << 23) | (m << 13)
594                } else {
595                    (s << 31) | ((e + 127 - 15) << 23) | (m << 13)
596                };
597                f32::from_bits(bits)
598            }
599            fn f32_to_f16_bits(v: f32) -> u16 {
600                let b = v.to_bits();
601                let (s, e, m) = ((b >> 31) as u16, ((b >> 23) & 0xff) as i32, b & 0x7fffff);
602                if e == 0xff { return (s << 15) | 0x7c00 | ((m >> 13) as u16 & 0x3ff); }
603                let e2 = e - 127 + 15;
604                if e2 >= 0x1f { return (s << 15) | 0x7c00; }
605                if e2 <= 0 {
606                    if e2 < -10 { return s << 15; }
607                    let m2 = (m | 0x800000) >> (1 - e2);
608                    // round-to-nearest-even on the shifted mantissa
609                    let r = (m2 >> 13) as u16 + ((m2 >> 12) & 1) as u16;
610                    return (s << 15) | r;
611                }
612                let mut r = ((e2 as u32) << 10) as u16 | (m >> 13) as u16;
613                if m & 0x1000 != 0 { r += 1; }
614                (s << 15) | r
615            }
616            let host: Vec<u8> = self.dtoh_u8(&mirror)?;
617            let mut vals: Vec<f32> = host.chunks_exact(2)
618                .map(|c| f16_bits_to_f32(u16::from_le_bytes([c[0], c[1]]))).collect();
619            for row in vals.chunks_mut(in_f) {
620                let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
621                if amax > 0.0 {
622                    let d = amax / 127.0;
623                    for v in row.iter_mut() {
624                        *v = (*v / d).round().clamp(-127.0, 127.0) * d;
625                    }
626                }
627            }
628            let out: Vec<u8> = vals.iter()
629                .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes()).collect();
630            mirror = self.htod_bytes(&out)?;
631        }
632        *f16 = Some(mirror);
633        Ok(())
634    }
635
636    /// Q4_0 twin of `build_q8_f16_raw` (18B blocks, campaign A 2026-07-31).
637    pub fn build_q4_f16_raw(
638        &self,
639        bytes: &CudaSlice<u8>,
640        in_f: usize,
641        out_f: usize,
642    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
643        assert!(in_f % 32 == 0);
644        let nblk = in_f / 32;
645        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
646        let rc = {
647            let stream = self.gpu.stream();
648            let (s_p, _gs) = bytes.device_ptr(&stream);
649            let (d_p, _gd) = dst.device_ptr_mut(&stream);
650            unsafe {
651                memra_q4_0_dequant_f16(
652                    s_p as *const core::ffi::c_void,
653                    d_p as *mut core::ffi::c_void,
654                    out_f as i64,
655                    nblk as i64,
656                    stream.cu_stream() as *mut core::ffi::c_void,
657                )
658            }
659        };
660        if rc != 0 {
661            return Err(format!("memra_q4_0_dequant_f16 rc={rc}").into());
662        }
663        Ok(dst)
664    }
665
666    /// Q5_K twin (176B superblocks, round 49b). Also the kernel_check gate entry for the
667    /// Q5_K f16-mirror class.
668    pub fn build_q5k_f16_raw(
669        &self,
670        bytes: &CudaSlice<u8>,
671        in_f: usize,
672        out_f: usize,
673    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
674        assert!(in_f % 256 == 0);
675        let nsb = in_f / 256;
676        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
677        let rc = {
678            let stream = self.gpu.stream();
679            let (s_p, _gs) = bytes.device_ptr(&stream);
680            let (d_p, _gd) = dst.device_ptr_mut(&stream);
681            unsafe {
682                memra_q5_K_dequant_f16(
683                    s_p as *const core::ffi::c_void,
684                    d_p as *mut core::ffi::c_void,
685                    out_f as i64,
686                    nsb as i64,
687                    stream.cu_stream() as *mut core::ffi::c_void,
688                )
689            }
690        };
691        if rc != 0 {
692            return Err(format!("memra_q5_K_dequant_f16 rc={rc}").into());
693        }
694        Ok(dst)
695    }
696
697    /// Q4_K twin of `build_q6k_f16_raw` (144B superblocks, round 49). Also the kernel_check
698    /// gate entry for the Q4_K f16-mirror class.
699    pub fn build_q4k_f16_raw(
700        &self,
701        bytes: &CudaSlice<u8>,
702        in_f: usize,
703        out_f: usize,
704    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
705        assert!(in_f % 256 == 0);
706        let nsb = in_f / 256;
707        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
708        let rc = {
709            let stream = self.gpu.stream();
710            let (s_p, _gs) = bytes.device_ptr(&stream);
711            let (d_p, _gd) = dst.device_ptr_mut(&stream);
712            unsafe {
713                memra_q4_K_dequant_f16(
714                    s_p as *const core::ffi::c_void,
715                    d_p as *mut core::ffi::c_void,
716                    out_f as i64,
717                    nsb as i64,
718                    stream.cu_stream() as *mut core::ffi::c_void,
719                )
720            }
721        };
722        if rc != 0 {
723            return Err(format!("memra_q4_K_dequant_f16 rc={rc}").into());
724        }
725        Ok(dst)
726    }
727
728    pub fn build_q6k_f16_raw(
729        &self,
730        bytes: &CudaSlice<u8>,
731        in_f: usize,
732        out_f: usize,
733    ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
734        assert!(in_f % 256 == 0);
735        let nsb = in_f / 256;
736        let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
737        let rc = {
738            let stream = self.gpu.stream();
739            let (s_p, _gs) = bytes.device_ptr(&stream);
740            let (d_p, _gd) = dst.device_ptr_mut(&stream);
741            unsafe {
742                memra_q6_K_dequant_f16(
743                    s_p as *const core::ffi::c_void,
744                    d_p as *mut core::ffi::c_void,
745                    out_f as i64,
746                    nsb as i64,
747                    stream.cu_stream() as *mut core::ffi::c_void,
748                )
749            }
750        };
751        if rc != 0 {
752            return Err(format!("memra_q6_K_dequant_f16 rc={rc}").into());
753        }
754        Ok(dst)
755    }
756}
memra_engine/f16_ffi.rs

memra_engine/
f16_ffi.rs
