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

memra_engine/
f16_ffi.rs
