1use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut};
19
20unsafe extern "C" {
21 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 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 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 fn memra_bf16_pp_gemm(
56 w_bf16: *const core::ffi::c_void,
57 x_f32: *const f32,
58 xb_bf16: *mut core::ffi::c_void,
59 y_f32: *mut f32,
60 m: i32,
61 n: i32,
62 k: i32,
63 ws: *mut core::ffi::c_void,
64 ws_bytes: usize,
65 stream: *mut core::ffi::c_void,
66 ) -> i32;
67 fn memra_q8_0_dequant_f16(
69 w_q8: *const core::ffi::c_void,
70 w_f16: *mut core::ffi::c_void,
71 out_f: i64,
72 nblk_row: i64,
73 stream: *mut core::ffi::c_void,
74 ) -> i32;
75 fn memra_q4_0_dequant_f16(
77 w_q4: *const core::ffi::c_void,
78 w_f16: *mut core::ffi::c_void,
79 out_f: i64,
80 nblk_row: i64,
81 stream: *mut core::ffi::c_void,
82 ) -> i32;
83 fn memra_q6_K_dequant_f16(
85 w_q6: *const core::ffi::c_void,
86 w_f16: *mut core::ffi::c_void,
87 out_f: i64,
88 nsb_row: i64,
89 stream: *mut core::ffi::c_void,
90 ) -> i32;
91 fn memra_q4_K_dequant_f16(
93 w_q4k: *const core::ffi::c_void,
94 w_f16: *mut core::ffi::c_void,
95 out_f: i64,
96 nsb_row: i64,
97 stream: *mut core::ffi::c_void,
98 ) -> i32;
99 fn memra_q5_K_dequant_f16(
101 w_q5k: *const core::ffi::c_void,
102 w_f16: *mut core::ffi::c_void,
103 out_f: i64,
104 nsb_row: i64,
105 stream: *mut core::ffi::c_void,
106 ) -> i32;
107}
108
109pub fn pp_f16_enabled() -> bool {
115 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
116 *ON.get_or_init(|| match std::env::var("MEMRA_PP_F16").as_deref() {
117 Ok("1") => true,
118 Ok("0") => false,
119 _ => cfg!(memra_hopper_mma),
120 })
121}
122
123pub fn pp_f16_capacity_ok(free: usize, need: usize) -> bool {
132 if std::env::var("MEMRA_PP_F16").is_ok() {
133 return false; }
135 need > 0 && free >= need + (8usize << 30)
136}
137
138pub fn pp_bf16_enabled() -> bool {
150 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
151 *ON.get_or_init(|| {
152 let on = matches!(std::env::var("MEMRA_PP_BF16").as_deref(), Ok("1"));
153 eprintln!(
159 "[bf16-tc] flag={} (MEMRA_PP_BF16; engagement is the per-shape ENGAGED line + \
160 the dispatch counter)",
161 if on { "on" } else { "off" }
162 );
163 on
164 })
165}
166
167pub static BF16_TC_DISPATCHES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
174
175pub fn bf16_tc_dispatches() -> u64 {
178 BF16_TC_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
179}
180
181pub struct F16Scratch {
184 pub xh: CudaSlice<u8>,
185 pub ws: CudaSlice<u8>,
186 cap_xh: usize,
187}
188
189impl F16Scratch {
190 pub fn with_capacity(
193 e: &crate::Engine,
194 xh_bytes: usize,
195 ) -> Result<Self, Box<dyn std::error::Error>> {
196 Ok(F16Scratch {
197 xh: e.alloc_u8_uninit(xh_bytes)?,
198 ws: e.alloc_u8_uninit(F16_WS_BYTES)?,
199 cap_xh: xh_bytes,
200 })
201 }
202}
203
204pub(crate) const F16_WS_BYTES: usize = 64 << 20;
205
206impl crate::Engine {
207 pub fn f16_scratch_swap(&self, new: Option<F16Scratch>) -> Option<F16Scratch> {
210 std::mem::replace(&mut *self.f16_scratch.lock().unwrap(), new)
211 }
212
213 pub fn try_f16_gemm(
216 &self,
217 w: &crate::model::GpuTensor,
218 x: &CudaSlice<f32>,
219 m: usize,
220 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
221 use crate::model::GpuTensor;
222 let (w16, ne, scale) = match w {
223 GpuTensor::Quant {
224 f16: Some(w16),
225 ne,
226 scale,
227 ..
228 } => (w16, ne, *scale),
229 _ => return Ok(None),
230 };
231 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
232 static SIM_ACT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
237 let sim_act =
238 *SIM_ACT.get_or_init(|| std::env::var("MEMRA_W8A8_SIM").as_deref() == Ok("2"));
239 let mut y = if sim_act {
240 let mut hx = self.dtoh(x)?;
241 hx.truncate(m * in_f);
242 for row in hx.chunks_mut(in_f) {
243 let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
244 if amax > 0.0 {
245 let d = amax / 127.0;
246 for v in row.iter_mut() {
247 *v = (*v / d).round().clamp(-127.0, 127.0) * d;
248 }
249 }
250 }
251 let xq = self.htod(&hx)?;
252 self.qmatvec_gemm_f16_raw(w16, &xq, m, in_f, out_f)?
253 } else {
254 self.qmatvec_gemm_f16_raw(w16, x, m, in_f, out_f)?
255 };
256 if scale != 1.0 {
257 self.scale_inplace(&mut y, scale, m * out_f)?;
258 }
259 Ok(Some(y))
260 }
261
262 pub fn bf16_tc_gemm(
266 &self,
267 data: &CudaSlice<u8>,
268 x: &CudaSlice<f32>,
269 m: usize,
270 in_f: usize,
271 out_f: usize,
272 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
273 let need_xh = m * in_f * 2;
274 let mut guard = self.f16_scratch.lock().unwrap();
275 if guard.is_none() {
276 *guard = Some(F16Scratch {
277 xh: self.alloc_u8_uninit(need_xh)?,
278 ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
279 cap_xh: need_xh,
280 });
281 }
282 let s = guard.as_mut().unwrap();
283 if need_xh > s.cap_xh {
284 s.xh = self.alloc_u8_uninit(need_xh)?;
285 s.cap_xh = need_xh;
286 }
287 let mut y = self.uninit(m * out_f)?; let rc = {
289 let stream = self.gpu.stream();
290 let (w_p, _gw) = data.device_ptr(&stream);
291 let (x_p, _gx) = x.device_ptr(&stream);
292 let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
293 let (y_p, _gy) = y.device_ptr_mut(&stream);
294 let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
295 if !(w_p as usize).is_multiple_of(16) {
299 -1
300 } else {
301 unsafe {
302 memra_bf16_pp_gemm(
303 w_p as *const core::ffi::c_void,
304 x_p as *const f32,
305 h_p as *mut core::ffi::c_void,
306 y_p as *mut f32,
307 m as i32,
308 out_f as i32,
309 in_f as i32,
310 ws_p as *mut core::ffi::c_void,
311 F16_WS_BYTES,
312 stream.cu_stream() as *mut core::ffi::c_void,
313 )
314 }
315 }
316 };
317 if rc != 0 {
324 static SAID: std::sync::Mutex<
325 Option<std::collections::HashSet<(usize, usize, usize)>>,
326 > = std::sync::Mutex::new(None);
327 let mut g = SAID.lock().unwrap();
328 let seen = g.get_or_insert_with(std::collections::HashSet::new);
329 if seen.insert((m, out_f, in_f)) {
330 eprintln!(
331 "[bf16-tc] DECLINED m={m} n={out_f} k={in_f} rc={rc} \
332 (1xxxx=cudaError convert, 2xxxx=no cublasLt algo, 3xxxx=matmul status, \
333 4xxxx=cublasLtCreate, -1=weight not 16B-aligned) — this shape falls back to \
334 the f32 dequant GEMM; every other shape keeps the tensor-core path"
335 );
336 }
337 return Ok(None);
338 }
339 BF16_TC_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
345 {
346 static ACCEPTED: std::sync::Mutex<
347 Option<std::collections::HashSet<(usize, usize, usize)>>,
348 > = std::sync::Mutex::new(None);
349 let mut g = ACCEPTED.lock().unwrap();
350 let seen = g.get_or_insert_with(std::collections::HashSet::new);
351 if seen.insert((m, out_f, in_f)) {
352 eprintln!(
353 "[bf16-tc] ENGAGED m={m} n={out_f} k={in_f} (bf16 tensor-core GEMM on resident checkpoint bytes)"
354 );
355 }
356 }
357 Ok(Some(y))
358 }
359
360 pub fn bf16_gemv_lt_into(
374 &self,
375 data: &CudaSlice<u8>,
376 x: &CudaSlice<f32>,
377 y: &mut CudaSlice<f32>,
378 in_f: usize,
379 out_f: usize,
380 t: usize,
381 ) -> Result<bool, Box<dyn std::error::Error>> {
382 if t == 0 || x.len() < t * in_f || y.len() < t * out_f {
383 return Err("bf16_gemv_lt geometry".into());
384 }
385 if data.len() < out_f * in_f * 2 {
386 return Err("bf16_gemv_lt weight too small".into());
387 }
388 let need_xh = t * in_f * 2;
389 let mut guard = self.f16_scratch.lock().unwrap();
390 if guard.is_none() {
391 *guard = Some(F16Scratch {
392 xh: self.alloc_u8_uninit(need_xh)?,
393 ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
394 cap_xh: need_xh,
395 });
396 }
397 let s = guard.as_mut().unwrap();
398 if need_xh > s.cap_xh {
399 s.xh = self.alloc_u8_uninit(need_xh)?;
400 s.cap_xh = need_xh;
401 }
402 let rc = {
403 let stream = self.gpu.stream();
404 let (w_p, _gw) = data.device_ptr(&stream);
405 let (x_p, _gx) = x.device_ptr(&stream);
406 let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
407 let (y_p, _gy) = y.device_ptr_mut(&stream);
408 let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
409 if !(w_p as usize).is_multiple_of(16) {
410 -1
411 } else {
412 unsafe {
413 memra_bf16_pp_gemm(
414 w_p as *const core::ffi::c_void,
415 x_p as *const f32,
416 h_p as *mut core::ffi::c_void,
417 y_p as *mut f32,
418 t as i32,
419 out_f as i32,
420 in_f as i32,
421 ws_p as *mut core::ffi::c_void,
422 F16_WS_BYTES,
423 stream.cu_stream() as *mut core::ffi::c_void,
424 )
425 }
426 }
427 };
428 if rc != 0 {
429 static SAID: std::sync::Mutex<
430 Option<std::collections::HashSet<(usize, usize, usize)>>,
431 > = std::sync::Mutex::new(None);
432 let mut g = SAID.lock().unwrap();
433 let seen = g.get_or_insert_with(std::collections::HashSet::new);
434 if seen.insert((t, out_f, in_f)) {
435 eprintln!(
436 "[b200-bf16-gemv-lt] DECLINED t={t} n={out_f} k={in_f} rc={rc} \
437 (1xxxx=cudaError convert, 2xxxx=no cublasLt algo, 3xxxx=matmul status, \
438 4xxxx=cublasLtCreate, -1=weight not 16B-aligned) — this shape keeps the \
439 shipped matvec kernel"
440 );
441 }
442 return Ok(false);
443 }
444 {
445 static ACCEPTED: std::sync::Mutex<
446 Option<std::collections::HashSet<(usize, usize, usize)>>,
447 > = std::sync::Mutex::new(None);
448 let mut g = ACCEPTED.lock().unwrap();
449 let seen = g.get_or_insert_with(std::collections::HashSet::new);
450 if seen.insert((t, out_f, in_f)) {
451 eprintln!(
452 "[b200-bf16-gemv-lt] ENGAGED t={t} n={out_f} k={in_f} (cuBLASLt reference \
453 GEMV, numeric class bf16_gemv_lt, MEMRA_B200_BF16_GEMV_LT=1)"
454 );
455 }
456 }
457 Ok(true)
458 }
459
460 pub fn qmatvec_gemm_f16_raw(
462 &self,
463 w16: &CudaSlice<u8>,
464 x: &CudaSlice<f32>,
465 m: usize,
466 in_f: usize,
467 out_f: usize,
468 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
469 let need_xh = m * in_f * 2;
470 let mut guard = self.f16_scratch.lock().unwrap();
471 if guard.is_none() {
472 *guard = Some(F16Scratch {
473 xh: self.alloc_u8_uninit(need_xh)?,
474 ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
475 cap_xh: need_xh,
476 });
477 }
478 let s = guard.as_mut().unwrap();
479 if need_xh > s.cap_xh {
480 s.xh = self.alloc_u8_uninit(need_xh)?;
481 s.cap_xh = need_xh;
482 }
483 let mut y = self.uninit(m * out_f)?; let rc = {
485 let stream = self.gpu.stream();
486 let (w_p, _gw) = w16.device_ptr(&stream);
487 let (x_p, _gx) = x.device_ptr(&stream);
488 let (h_p, _gh) = s.xh.device_ptr_mut(&stream);
489 let (y_p, _gy) = y.device_ptr_mut(&stream);
490 let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
491 unsafe {
492 memra_f16_pp_gemm(
493 w_p as *const core::ffi::c_void,
494 x_p as *const f32,
495 h_p as *mut core::ffi::c_void,
496 y_p as *mut f32,
497 m as i32,
498 out_f as i32,
499 in_f as i32,
500 ws_p as *mut core::ffi::c_void,
501 F16_WS_BYTES,
502 stream.cu_stream() as *mut core::ffi::c_void,
503 )
504 }
505 };
506 if rc != 0 {
507 return Err(format!(
508 "memra_f16_pp_gemm rc={rc} (m={m} n={out_f} k={in_f}; 1xxxx=cudaError convert, \
509 2xxxx=no cublasLt algo, 3xxxx=matmul status)"
510 )
511 .into());
512 }
513 Ok(y)
514 }
515
516 #[allow(clippy::manual_is_multiple_of)] pub fn f16_act(
521 &self,
522 x: &CudaSlice<f32>,
523 nelem: usize,
524 in_f: usize,
525 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
526 static SIM_ACT2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
530 if *SIM_ACT2.get_or_init(|| std::env::var("MEMRA_W8A8_SIM").as_deref() == Ok("2"))
531 && in_f > 0
532 && nelem % in_f == 0
533 {
534 static ONCE: std::sync::Once = std::sync::Once::new();
535 ONCE.call_once(|| {
536 eprintln!("[w8a8-sim] act per-token int8 fake-quant ACTIVE (f16_act)")
537 });
538 let mut hx = self.dtoh(x)?;
539 hx.truncate(nelem);
540 for row in hx.chunks_mut(in_f) {
541 let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
542 if amax > 0.0 {
543 let d = amax / 127.0;
544 for v in row.iter_mut() {
545 *v = (*v / d).round().clamp(-127.0, 127.0) * d;
546 }
547 }
548 }
549 let xq = self.htod(&hx)?;
550 let mut xh = self.alloc_u8_uninit(nelem * 2)?;
551 let rc = {
552 let stream = self.gpu.stream();
553 let (x_p, _gx) = xq.device_ptr(&stream);
554 let (h_p, _gh) = xh.device_ptr_mut(&stream);
555 unsafe {
556 memra_f16_cvt(
557 x_p as *const f32,
558 h_p as *mut core::ffi::c_void,
559 nelem,
560 stream.cu_stream() as *mut core::ffi::c_void,
561 )
562 }
563 };
564 if rc != 0 {
565 return Err(format!("memra_f16_cvt rc={rc}").into());
566 }
567 return Ok(xh);
568 }
569 let mut xh = self.alloc_u8_uninit(nelem * 2)?;
570 let rc = {
571 let stream = self.gpu.stream();
572 let (x_p, _gx) = x.device_ptr(&stream);
573 let (h_p, _gh) = xh.device_ptr_mut(&stream);
574 unsafe {
575 memra_f16_cvt(
576 x_p as *const f32,
577 h_p as *mut core::ffi::c_void,
578 nelem,
579 stream.cu_stream() as *mut core::ffi::c_void,
580 )
581 }
582 };
583 if rc != 0 {
584 return Err(format!("memra_f16_cvt rc={rc}").into());
585 }
586 Ok(xh)
587 }
588
589 pub fn try_f16_gemm_pre_into(
594 &self,
595 w: &crate::model::GpuTensor,
596 xh: &CudaSlice<u8>,
597 m: usize,
598 y: &mut CudaSlice<f32>,
599 ) -> Result<bool, Box<dyn std::error::Error>> {
600 use crate::model::GpuTensor;
601 let (w16, ne, scale) = match w {
602 GpuTensor::Quant {
603 f16: Some(w16),
604 ne,
605 scale,
606 ..
607 } => (w16, ne, *scale),
608 _ => return Ok(false),
609 };
610 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
611 assert!(
612 y.len() >= m * out_f,
613 "try_f16_gemm_pre_into: output slab too small"
614 );
615 let mut guard = self.f16_scratch.lock().unwrap();
616 if guard.is_none() {
617 *guard = Some(F16Scratch {
618 xh: self.alloc_u8_uninit(2)?,
619 ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
620 cap_xh: 2,
621 });
622 }
623 let s = guard.as_mut().unwrap();
624 let rc = {
625 let stream = self.gpu.stream();
626 let (w_p, _gw) = w16.device_ptr(&stream);
627 let (h_p, _gh) = xh.device_ptr(&stream);
628 let (y_p, _gy) = y.device_ptr_mut(&stream);
629 let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
630 unsafe {
631 memra_f16_pp_gemm_pre(
632 w_p as *const core::ffi::c_void,
633 h_p as *const core::ffi::c_void,
634 y_p as *mut f32,
635 m as i32,
636 out_f as i32,
637 in_f as i32,
638 ws_p as *mut core::ffi::c_void,
639 F16_WS_BYTES,
640 stream.cu_stream() as *mut core::ffi::c_void,
641 )
642 }
643 };
644 if rc != 0 {
645 return Err(
646 format!("memra_f16_pp_gemm_pre(into) rc={rc} (m={m} n={out_f} k={in_f})").into(),
647 );
648 }
649 if scale != 1.0 {
650 self.scale_inplace(y, scale, m * out_f)?;
651 }
652 Ok(true)
653 }
654
655 pub fn try_f16_gemm_pre_into_off(
659 &self,
660 w: &crate::model::GpuTensor,
661 xh: &CudaSlice<u8>,
662 m: usize,
663 y: &mut CudaSlice<f32>,
664 off_elems: usize,
665 ) -> Result<bool, Box<dyn std::error::Error>> {
666 use crate::model::GpuTensor;
667 let (w16, ne, scale) = match w {
668 GpuTensor::Quant {
669 f16: Some(w16),
670 ne,
671 scale,
672 ..
673 } => (w16, ne, *scale),
674 _ => return Ok(false),
675 };
676 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
677 assert!(
678 y.len() >= off_elems + m * out_f,
679 "try_f16_gemm_pre_into_off: output slab too small"
680 );
681 if scale != 1.0 {
682 return Ok(false); }
684 let mut guard = self.f16_scratch.lock().unwrap();
685 if guard.is_none() {
686 *guard = Some(F16Scratch {
687 xh: self.alloc_u8_uninit(2)?,
688 ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
689 cap_xh: 2,
690 });
691 }
692 let s = guard.as_mut().unwrap();
693 let rc = {
694 let stream = self.gpu.stream();
695 let (w_p, _gw) = w16.device_ptr(&stream);
696 let (h_p, _gh) = xh.device_ptr(&stream);
697 let (y_p, _gy) = y.device_ptr_mut(&stream);
698 let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
699 unsafe {
700 memra_f16_pp_gemm_pre(
701 w_p as *const core::ffi::c_void,
702 h_p as *const core::ffi::c_void,
703 (y_p as *mut f32).add(off_elems),
704 m as i32,
705 out_f as i32,
706 in_f as i32,
707 ws_p as *mut core::ffi::c_void,
708 F16_WS_BYTES,
709 stream.cu_stream() as *mut core::ffi::c_void,
710 )
711 }
712 };
713 if rc != 0 {
714 return Err(format!(
715 "memra_f16_pp_gemm_pre(into_off) rc={rc} (m={m} n={out_f} k={in_f})"
716 )
717 .into());
718 }
719 Ok(true)
720 }
721
722 pub fn try_f16_gemm_pre(
725 &self,
726 w: &crate::model::GpuTensor,
727 xh: &CudaSlice<u8>,
728 m: usize,
729 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
730 use crate::model::GpuTensor;
731 let (w16, ne, scale) = match w {
732 GpuTensor::Quant {
733 f16: Some(w16),
734 ne,
735 scale,
736 ..
737 } => (w16, ne, *scale),
738 _ => return Ok(None),
739 };
740 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
741 let mut guard = self.f16_scratch.lock().unwrap();
743 if guard.is_none() {
744 *guard = Some(F16Scratch {
745 xh: self.alloc_u8_uninit(2)?,
746 ws: self.alloc_u8_uninit(F16_WS_BYTES)?,
747 cap_xh: 2,
748 });
749 }
750 let s = guard.as_mut().unwrap();
751 let mut y = self.uninit(m * out_f)?;
752 let rc = {
753 let stream = self.gpu.stream();
754 let (w_p, _gw) = w16.device_ptr(&stream);
755 let (h_p, _gh) = xh.device_ptr(&stream);
756 let (y_p, _gy) = y.device_ptr_mut(&stream);
757 let (ws_p, _gws) = s.ws.device_ptr_mut(&stream);
758 unsafe {
759 memra_f16_pp_gemm_pre(
760 w_p as *const core::ffi::c_void,
761 h_p as *const core::ffi::c_void,
762 y_p as *mut f32,
763 m as i32,
764 out_f as i32,
765 in_f as i32,
766 ws_p as *mut core::ffi::c_void,
767 F16_WS_BYTES,
768 stream.cu_stream() as *mut core::ffi::c_void,
769 )
770 }
771 };
772 if rc != 0 {
773 return Err(format!("memra_f16_pp_gemm_pre rc={rc} (m={m} n={out_f} k={in_f})").into());
774 }
775 if scale != 1.0 {
776 self.scale_inplace(&mut y, scale, m * out_f)?;
777 }
778 Ok(Some(y))
779 }
780
781 pub fn build_q8_f16_raw(
784 &self,
785 bytes: &CudaSlice<u8>,
786 in_f: usize,
787 out_f: usize,
788 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
789 assert!(in_f.is_multiple_of(32));
790 let nblk = in_f / 32;
791 let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
792 let rc = {
793 let stream = self.gpu.stream();
794 let (s_p, _gs) = bytes.device_ptr(&stream);
795 let (d_p, _gd) = dst.device_ptr_mut(&stream);
796 unsafe {
797 memra_q8_0_dequant_f16(
798 s_p as *const core::ffi::c_void,
799 d_p as *mut core::ffi::c_void,
800 out_f as i64,
801 nblk as i64,
802 stream.cu_stream() as *mut core::ffi::c_void,
803 )
804 }
805 };
806 if rc != 0 {
807 return Err(format!("memra_q8_0_dequant_f16 rc={rc}").into());
808 }
809 Ok(dst)
810 }
811
812 pub fn build_q8_f16(
815 &self,
816 t: &mut crate::model::GpuTensor,
817 ) -> Result<(), Box<dyn std::error::Error>> {
818 use crate::model::GpuTensor;
819 let GpuTensor::Quant {
820 bytes,
821 qtype,
822 row_bytes,
823 ne,
824 f16,
825 ..
826 } = t
827 else {
828 return Ok(());
829 };
830 let q4 = *qtype == crate::QT_Q4_0;
838 let q6k = *qtype == crate::QT_Q6_K;
839 let q4k = *qtype == crate::QT_Q4_K;
840 let q5k = *qtype == crate::QT_Q5_K;
841 if (*qtype != crate::QT_Q8_0 && !q4 && !q6k && !q4k && !q5k)
842 || f16.is_some()
843 || ne.len() != 2
844 {
845 return Ok(());
846 }
847 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
848 if q6k || q4k || q5k {
849 let sb = if q6k {
850 210
851 } else if q5k {
852 176
853 } else {
854 144
855 };
856 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * sb {
857 return Ok(());
858 }
859 } else if in_f % 32 != 0 || *row_bytes != (in_f / 32) * (if q4 { 18 } else { 34 }) {
860 return Ok(());
861 }
862 use std::sync::atomic::{AtomicUsize, Ordering};
865 static SPENT: AtomicUsize = AtomicUsize::new(0);
866 static BUDGET: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
867 let budget = *BUDGET.get_or_init(|| {
868 std::env::var("MEMRA_PP_F16_BUDGET_MB")
869 .ok()
870 .and_then(|v| v.parse::<usize>().ok())
871 .unwrap_or(32768)
872 << 20
873 });
874 let sz = out_f * in_f * 2;
875 if SPENT.fetch_add(sz, Ordering::Relaxed) + sz > budget {
876 SPENT.fetch_sub(sz, Ordering::Relaxed);
877 return Ok(());
878 }
879 let mut mirror = if q6k {
880 self.build_q6k_f16_raw(bytes, in_f, out_f)?
881 } else if q4k {
882 self.build_q4k_f16_raw(bytes, in_f, out_f)?
883 } else if q5k {
884 self.build_q5k_f16_raw(bytes, in_f, out_f)?
885 } else if q4 {
886 self.build_q4_f16_raw(bytes, in_f, out_f)?
887 } else {
888 self.build_q8_f16_raw(bytes, in_f, out_f)?
889 };
890 static SIM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
899 if *SIM.get_or_init(|| {
900 matches!(
901 std::env::var("MEMRA_W8A8_SIM").as_deref(),
902 Ok("1") | Ok("2")
903 )
904 }) {
905 fn f16_bits_to_f32(b: u16) -> f32 {
906 let (s, e, m) = (
907 (b >> 15) as u32,
908 ((b >> 10) & 0x1f) as u32,
909 (b & 0x3ff) as u32,
910 );
911 let bits = if e == 0 {
912 if m == 0 {
913 s << 31
914 } else {
915 let mut e2 = 127 - 15 + 1;
917 let mut m2 = m;
918 while m2 & 0x400 == 0 {
919 m2 <<= 1;
920 e2 -= 1;
921 }
922 (s << 31) | ((e2 as u32) << 23) | ((m2 & 0x3ff) << 13)
923 }
924 } else if e == 0x1f {
925 (s << 31) | (0xff << 23) | (m << 13)
926 } else {
927 (s << 31) | ((e + 127 - 15) << 23) | (m << 13)
928 };
929 f32::from_bits(bits)
930 }
931 fn f32_to_f16_bits(v: f32) -> u16 {
932 let b = v.to_bits();
933 let (s, e, m) = ((b >> 31) as u16, ((b >> 23) & 0xff) as i32, b & 0x7fffff);
934 if e == 0xff {
935 return (s << 15) | 0x7c00 | ((m >> 13) as u16 & 0x3ff);
936 }
937 let e2 = e - 127 + 15;
938 if e2 >= 0x1f {
939 return (s << 15) | 0x7c00;
940 }
941 if e2 <= 0 {
942 if e2 < -10 {
943 return s << 15;
944 }
945 let m2 = (m | 0x800000) >> (1 - e2);
946 let r = (m2 >> 13) as u16 + ((m2 >> 12) & 1) as u16;
948 return (s << 15) | r;
949 }
950 let mut r = ((e2 as u32) << 10) as u16 | (m >> 13) as u16;
951 if m & 0x1000 != 0 {
952 r += 1;
953 }
954 (s << 15) | r
955 }
956 let host: Vec<u8> = self.dtoh_u8(&mirror)?;
957 let mut vals: Vec<f32> = host
958 .chunks_exact(2)
959 .map(|c| f16_bits_to_f32(u16::from_le_bytes([c[0], c[1]])))
960 .collect();
961 for row in vals.chunks_mut(in_f) {
962 let amax = row.iter().fold(0f32, |a, &v| a.max(v.abs()));
963 if amax > 0.0 {
964 let d = amax / 127.0;
965 for v in row.iter_mut() {
966 *v = (*v / d).round().clamp(-127.0, 127.0) * d;
967 }
968 }
969 }
970 let out: Vec<u8> = vals
971 .iter()
972 .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes())
973 .collect();
974 mirror = self.htod_bytes(&out)?;
975 }
976 *f16 = Some(mirror);
977 Ok(())
978 }
979
980 pub fn build_q4_f16_raw(
982 &self,
983 bytes: &CudaSlice<u8>,
984 in_f: usize,
985 out_f: usize,
986 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
987 assert!(in_f.is_multiple_of(32));
988 let nblk = in_f / 32;
989 let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
990 let rc = {
991 let stream = self.gpu.stream();
992 let (s_p, _gs) = bytes.device_ptr(&stream);
993 let (d_p, _gd) = dst.device_ptr_mut(&stream);
994 unsafe {
995 memra_q4_0_dequant_f16(
996 s_p as *const core::ffi::c_void,
997 d_p as *mut core::ffi::c_void,
998 out_f as i64,
999 nblk as i64,
1000 stream.cu_stream() as *mut core::ffi::c_void,
1001 )
1002 }
1003 };
1004 if rc != 0 {
1005 return Err(format!("memra_q4_0_dequant_f16 rc={rc}").into());
1006 }
1007 Ok(dst)
1008 }
1009
1010 pub fn build_q5k_f16_raw(
1013 &self,
1014 bytes: &CudaSlice<u8>,
1015 in_f: usize,
1016 out_f: usize,
1017 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1018 assert!(in_f.is_multiple_of(256));
1019 let nsb = in_f / 256;
1020 let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
1021 let rc = {
1022 let stream = self.gpu.stream();
1023 let (s_p, _gs) = bytes.device_ptr(&stream);
1024 let (d_p, _gd) = dst.device_ptr_mut(&stream);
1025 unsafe {
1026 memra_q5_K_dequant_f16(
1027 s_p as *const core::ffi::c_void,
1028 d_p as *mut core::ffi::c_void,
1029 out_f as i64,
1030 nsb as i64,
1031 stream.cu_stream() as *mut core::ffi::c_void,
1032 )
1033 }
1034 };
1035 if rc != 0 {
1036 return Err(format!("memra_q5_K_dequant_f16 rc={rc}").into());
1037 }
1038 Ok(dst)
1039 }
1040
1041 pub fn build_q4k_f16_raw(
1044 &self,
1045 bytes: &CudaSlice<u8>,
1046 in_f: usize,
1047 out_f: usize,
1048 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1049 assert!(in_f.is_multiple_of(256));
1050 let nsb = in_f / 256;
1051 let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
1052 let rc = {
1053 let stream = self.gpu.stream();
1054 let (s_p, _gs) = bytes.device_ptr(&stream);
1055 let (d_p, _gd) = dst.device_ptr_mut(&stream);
1056 unsafe {
1057 memra_q4_K_dequant_f16(
1058 s_p as *const core::ffi::c_void,
1059 d_p as *mut core::ffi::c_void,
1060 out_f as i64,
1061 nsb as i64,
1062 stream.cu_stream() as *mut core::ffi::c_void,
1063 )
1064 }
1065 };
1066 if rc != 0 {
1067 return Err(format!("memra_q4_K_dequant_f16 rc={rc}").into());
1068 }
1069 Ok(dst)
1070 }
1071
1072 pub fn build_q6k_f16_raw(
1073 &self,
1074 bytes: &CudaSlice<u8>,
1075 in_f: usize,
1076 out_f: usize,
1077 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1078 assert!(in_f.is_multiple_of(256));
1079 let nsb = in_f / 256;
1080 let mut dst = self.alloc_u8_uninit(out_f * in_f * 2)?;
1081 let rc = {
1082 let stream = self.gpu.stream();
1083 let (s_p, _gs) = bytes.device_ptr(&stream);
1084 let (d_p, _gd) = dst.device_ptr_mut(&stream);
1085 unsafe {
1086 memra_q6_K_dequant_f16(
1087 s_p as *const core::ffi::c_void,
1088 d_p as *mut core::ffi::c_void,
1089 out_f as i64,
1090 nsb as i64,
1091 stream.cu_stream() as *mut core::ffi::c_void,
1092 )
1093 }
1094 };
1095 if rc != 0 {
1096 return Err(format!("memra_q6_K_dequant_f16 rc={rc}").into());
1097 }
1098 Ok(dst)
1099 }
1100}