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

memra_engine/
f16_ffi.rs
