1use crate::Engine;
16use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut};
17
18static MMQ_ACT_EPOCH: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21#[allow(clippy::type_complexity)]
22static MMQ_ACT_SLOT: std::sync::Mutex<Option<(u64, u64, usize, usize, CudaSlice<u8>)>> =
23 std::sync::Mutex::new(None);
24static MMQ_FIXUP_SLOT: std::sync::Mutex<Option<cudarc::driver::CudaSlice<u8>>> =
26 std::sync::Mutex::new(None);
27
28unsafe extern "C" {
29 pub fn memra_mmq_nvfp4_act_bytes(in_f: i32, n_tokens: i32) -> usize;
31 pub fn memra_mmq_nvfp4(
38 w_nvfp4_blocks: *const core::ffi::c_void,
39 act_f32: *const f32,
40 y: *mut f32,
41 in_f: i32,
42 out_f: i32,
43 n_tokens: i32,
44 act_scratch: *mut core::ffi::c_void,
45 stream: *mut core::ffi::c_void,
46 out_scale: f32,
47 ) -> i32;
48 pub fn memra_mmq_nvfp4_ex(
54 w_nvfp4_blocks: *const core::ffi::c_void,
55 act_f32: *const f32,
56 y: *mut f32,
57 in_f: i32,
58 out_f: i32,
59 n_tokens: i32,
60 act_scratch: *mut core::ffi::c_void,
61 stream: *mut core::ffi::c_void,
62 out_scale: f32,
63 per_token_scale: i32,
64 ) -> i32;
65 pub fn memra_mmq_nvfp4_ex2(
71 w_nvfp4_blocks: *const core::ffi::c_void,
72 act_f32: *const f32,
73 y: *mut f32,
74 in_f: i32,
75 out_f: i32,
76 n_tokens: i32,
77 act_scratch: *mut core::ffi::c_void,
78 stream: *mut core::ffi::c_void,
79 out_scale: f32,
80 per_token_scale: i32,
81 residual_k: i32,
82 ) -> i32;
83 pub fn memra_mmq_nvfp4_w4a8_act_bytes(in_f: i32, n_tokens: i32) -> usize;
85 pub fn memra_mmq_nvfp4_w4a8(
93 w_nvfp4_blocks: *const core::ffi::c_void,
94 act_f32: *const f32,
95 y: *mut f32,
96 in_f: i32,
97 out_f: i32,
98 n_tokens: i32,
99 act_scratch: *mut core::ffi::c_void,
100 stream: *mut core::ffi::c_void,
101 out_scale: f32,
102 rp: i32,
103 ) -> i32;
104 pub fn memra_mmq_nvfp4_f8f4_act_bytes(in_f: i32, n_tokens: i32) -> usize;
106 pub fn memra_mmq_nvfp4_f8f4(
112 w_nvfp4_blocks: *const core::ffi::c_void,
113 act_f32: *const f32,
114 y: *mut f32,
115 in_f: i32,
116 out_f: i32,
117 n_tokens: i32,
118 act_scratch: *mut core::ffi::c_void,
119 stream: *mut core::ffi::c_void,
120 out_scale: f32,
121 rp: i32,
122 ) -> i32;
123 pub fn memra_mmq_fp8_blk_act_bytes(in_f: i32, n_tokens: i32) -> usize;
126 pub fn memra_mmq_fp8_blk_scale_rows(out_f: i32) -> i32;
128 pub fn memra_mmq_fp8_blk_scale_cols(in_f: i32) -> i32;
129 pub fn memra_mmq_fp8_blk(
136 w_e4m3: *const core::ffi::c_void,
137 blk_scales: *const f32,
138 act_f32: *const f32,
139 y: *mut f32,
140 in_f: i32,
141 out_f: i32,
142 n_tokens: i32,
143 act_scratch: *mut core::ffi::c_void,
144 stream: *mut core::ffi::c_void,
145 out_scale: f32,
146 ) -> i32;
147 pub fn memra_fp8_blk_count_nan(
151 w_e4m3: *const core::ffi::c_void,
152 nbytes: usize,
153 out_count: *mut u32,
154 stream: *mut core::ffi::c_void,
155 ) -> i32;
156 pub fn memra_mmq_q45k_act_bytes(in_f: i32, n_tokens: i32) -> usize;
158 pub fn memra_mmq_q4_K(
161 w_q4k_blocks: *const core::ffi::c_void,
162 act_f32: *const f32,
163 y: *mut f32,
164 in_f: i32,
165 out_f: i32,
166 n_tokens: i32,
167 act_scratch: *mut core::ffi::c_void,
168 stream: *mut core::ffi::c_void,
169 ) -> i32;
170 pub fn memra_mmq_q5_K(
172 w_q5k_blocks: *const core::ffi::c_void,
173 act_f32: *const f32,
174 y: *mut f32,
175 in_f: i32,
176 out_f: i32,
177 n_tokens: i32,
178 act_scratch: *mut core::ffi::c_void,
179 stream: *mut core::ffi::c_void,
180 ) -> i32;
181
182 pub fn memra_mmq_q8_0_act_bytes(in_f: i32, n_tokens: i32) -> usize;
184 pub fn memra_mmq_q8_0(
188 w_q8_0_blocks: *const core::ffi::c_void,
189 act_f32: *const f32,
190 y: *mut f32,
191 in_f: i32,
192 out_f: i32,
193 n_tokens: i32,
194 act_scratch: *mut core::ffi::c_void,
195 stream: *mut core::ffi::c_void,
196 ) -> i32;
197
198 pub fn memra_accprobe_act_bytes(in_f: i32, n_tokens: i32) -> usize;
207 pub fn memra_accprobe_gemm_s32(
209 w_q8_0_blocks: *const core::ffi::c_void,
210 act_q: *const core::ffi::c_void,
211 y: *mut f32,
212 in_f: i32,
213 out_f: i32,
214 n_tokens: i32,
215 stream: *mut core::ffi::c_void,
216 ) -> i32;
217 pub fn memra_accprobe_gemm_f32(
219 w_q8_0_blocks: *const core::ffi::c_void,
220 act_q: *const core::ffi::c_void,
221 y: *mut f32,
222 in_f: i32,
223 out_f: i32,
224 n_tokens: i32,
225 stream: *mut core::ffi::c_void,
226 ) -> i32;
227
228 pub fn memra_mmq_q4_0_act_bytes(in_f: i32, n_tokens: i32) -> usize;
230 pub fn memra_mmq_q4_0(
236 w_q4_0: *const core::ffi::c_void,
237 act_f32: *const f32,
238 y: *mut f32,
239 in_f: i32,
240 out_f: i32,
241 n_tokens: i32,
242 act_scratch: *mut core::ffi::c_void,
243 stream: *mut core::ffi::c_void,
244 rp: i32,
245 ) -> i32;
246 pub fn memra_mmq_q4_0_quant_act(
248 act_f32: *const f32,
249 act_scratch: *mut core::ffi::c_void,
250 in_f: i32,
251 n_tokens: i32,
252 stream: *mut core::ffi::c_void,
253 ) -> i32;
254 pub fn memra_mmq_q4_0_gemm(
256 w_q4_0: *const core::ffi::c_void,
257 act_scratch: *const core::ffi::c_void,
258 y: *mut f32,
259 in_f: i32,
260 out_f: i32,
261 n_tokens: i32,
262 stream: *mut core::ffi::c_void,
263 rp: i32,
264 ) -> i32;
265 pub fn memra_mmq_q4_0_fixup_bytes() -> usize;
267 pub fn memra_mmq_q4_0_set_clc(force: i32) -> i32;
272 pub fn memra_mmq_q4_0_gemm_sk(
274 w_q4_0: *const core::ffi::c_void,
275 act_scratch: *const core::ffi::c_void,
276 y: *mut f32,
277 fixup_scratch: *mut core::ffi::c_void,
278 in_f: i32,
279 out_f: i32,
280 n_tokens: i32,
281 stream: *mut core::ffi::c_void,
282 rp: i32,
283 ) -> i32;
284
285 pub fn memra_mmq_iq_experts_act_bytes(in_f: i32, n_tokens: i32) -> usize;
288 pub fn memra_mmq_iq_quantize_act(
290 act_f32: *const f32,
291 act_scratch: *mut core::ffi::c_void,
292 in_f: i32,
293 n_tokens: i32,
294 stream: *mut core::ffi::c_void,
295 ) -> i32;
296 pub fn memra_mmq_iq_fused_act_quant(
300 gate: *const f32,
301 up: *const f32,
302 act_scratch: *mut core::ffi::c_void,
303 in_f: i32,
304 n_tokens: i32,
305 act_kind: i32,
306 stream: *mut core::ffi::c_void,
307 ) -> i32;
308 pub fn memra_mmq_iq4xs_dense(
317 w_blocks: *const core::ffi::c_void,
318 act_f32: *const f32,
319 y: *mut f32,
320 in_f: i32,
321 out_f: i32,
322 n_tokens: i32,
323 row_bytes: i64,
324 act_scratch: *mut core::ffi::c_void,
325 stream: *mut core::ffi::c_void,
326 ) -> i32;
327 pub fn memra_mmq_iq_experts(
328 table: *const u64,
329 proj: i32,
330 n_expert: i32,
331 ex_ids: *const i32,
332 ex_off: *const i32,
333 ex_pairs: *const i32,
334 pair_tok: *const i32,
335 act_scratch: *const core::ffi::c_void,
336 y: *mut f32,
337 in_f: i32,
338 out_f: i32,
339 n_active: i32,
340 n_tokens: i32,
341 qtype: i32,
342 row_bytes: i64,
343 stream: *mut core::ffi::c_void,
344 ) -> i32;
345
346 pub fn memra_moe_f16g_dequant(
348 table: *const u64, proj: i32, n_expert: i32, ex_ids: *const i32,
349 w_f16: *mut core::ffi::c_void, in_f: i32, out_f: i32, n_active: i32,
350 qtype: i32, row_bytes: i64, stream: *mut core::ffi::c_void,
351 ) -> i32;
352 pub fn memra_moe_f16g_gather_act(
353 x: *const f32, pair_tok_or_null: *const i32, act_f16: *mut core::ffi::c_void,
354 row_scale: *mut f32, in_f: i32, n_pairs: i32, stream: *mut core::ffi::c_void,
355 ) -> i32;
356 pub fn memra_moe_f16g_h2f_scaled(
357 src_f16: *const core::ffi::c_void, dst: *mut f32, row_scale: *const f32,
358 ncols: i32, nrows: i32, stream: *mut core::ffi::c_void,
359 ) -> i32;
360 pub fn memra_moe_f16g_gemm(
361 w_f16: *const core::ffi::c_void, act_f16: *const core::ffi::c_void,
362 y_f16: *mut core::ffi::c_void,
363 ex_off_host: *const i32, n_active: i32, in_f: i32, out_f: i32,
364 stream: *mut core::ffi::c_void,
365 ) -> i32;
366 pub fn memra_moe_f16g_h2f(
367 src_f16: *const core::ffi::c_void, dst: *mut f32, n: usize,
368 stream: *mut core::ffi::c_void,
369 ) -> i32;
370 pub fn memra_moe_f16g_gemm_sk(
379 w_f16: *const core::ffi::c_void, act_f16: *const core::ffi::c_void,
380 y_f32: *mut f32, row_scale: *const f32, ex_off_dev: *const i32,
381 ex_off_host: *const i32,
382 n_active: i32, max_m: i32, in_f: i32, out_f: i32, shape_sel: i32, cross: i32,
383 tail: i32, stream: *mut core::ffi::c_void,
384 ) -> i32;
385 pub fn memra_moe_kq_gemm_sk(
392 table: *const u64, proj: i32, n_expert: i32, ex_ids: *const i32,
393 act_f16: *const core::ffi::c_void, y_f32: *mut f32,
394 row_scale: *const f32, ex_off_dev: *const i32, ex_off_host: *const i32,
395 n_active: i32, max_m: i32, in_f: i32, out_f: i32, qtype: i32, cross: i32,
396 tail: i32, row_bytes: i64, stream: *mut core::ffi::c_void,
397 ) -> i32;
398}
399
400pub fn mmq_w4a8_enabled() -> bool {
409 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
410 *ON.get_or_init(|| {
411 std::env::var("MEMRA_MMQ_W4A8")
412 .map(|v| v != "0")
413 .unwrap_or(true)
414 })
415}
416
417pub fn mmq_residual_k() -> i32 {
426 std::env::var("MEMRA_MMQ_RESIDUAL_K")
427 .ok()
428 .and_then(|v| v.parse::<i32>().ok())
429 .unwrap_or(0)
430 .clamp(0, 64)
431}
432
433pub fn mmq_q8_enabled() -> bool {
439 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
440 *ON.get_or_init(|| {
444 std::env::var("MEMRA_PP_Q8MMQ")
445 .map(|v| v != "0")
446 .unwrap_or(true)
447 })
448}
449
450pub fn mmq_iq4xs_enabled() -> bool {
458 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
459 *ON.get_or_init(|| {
460 std::env::var("MEMRA_PP_IQMMQ")
461 .map(|v| v != "0")
462 .unwrap_or(true)
463 })
464}
465
466pub fn mmq_q4_enabled() -> bool {
472 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
473 *ON.get_or_init(|| {
474 std::env::var("MEMRA_PP_Q4MMQ")
475 .map(|v| v != "0")
476 .unwrap_or(true)
477 })
478}
479
480impl Engine {
481 pub fn mmq_supports(&self, w: &crate::model::GpuTensor) -> bool {
484 use crate::model::GpuTensor;
485 if crate::portable_mma_gated() {
486 return false;
487 }
488 let mmq_opt_in = std::env::var("MEMRA_MMQ").is_ok();
489 match w {
490 GpuTensor::Quant { qtype, rp, .. } if *qtype == crate::QT_NVFP4 && *rp => {
498 !cfg!(memra_portable_cuda) && mmq_w4a8_enabled() && w.in_features() % 64 == 0
499 }
500 GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_NVFP4 => {
502 !cfg!(memra_portable_cuda) && (mmq_w4a8_enabled() || mmq_opt_in)
503 && w.in_features() % 64 == 0
504 }
505 GpuTensor::Quant { qtype, .. }
506 if *qtype == crate::QT_Q4_K || *qtype == crate::QT_Q5_K =>
507 {
508 (mmq_w4a8_enabled() || mmq_opt_in) && w.in_features() % 256 == 0
509 }
510 GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_Q8_0 => {
515 mmq_q8_enabled() && w.in_features() % 256 == 0
516 }
517 GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_Q4_0 => {
522 mmq_q4_enabled() && w.in_features() % 256 == 0
523 }
524 GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_IQ4_XS => {
530 mmq_iq4xs_enabled() && Self::iq_fast_enabled() && w.in_features() % 256 == 0
531 }
532 _ => false,
533 }
534 }
535
536 pub fn qmatvec_mmq(
539 &self,
540 w: &crate::model::GpuTensor,
541 x: &CudaSlice<f32>,
542 m: usize,
543 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
544 use crate::model::GpuTensor;
545 let (in_f, out_f) = (w.in_features(), w.out_features());
546 let GpuTensor::Quant {
547 bytes,
548 scale,
549 qtype,
550 rp,
551 ..
552 } = w
553 else {
554 return Err("qmatvec_mmq: not a Quant tensor".into());
555 };
556 let w4a8_explicit = std::env::var("MEMRA_MMQ_W4A8")
561 .map(|v| v != "0")
562 .unwrap_or(false);
563 let use_w4a8 =
564 *rp || w4a8_explicit || (mmq_w4a8_enabled() && std::env::var("MEMRA_MMQ").is_err());
565 match *qtype {
566 q if q == crate::QT_NVFP4 && use_w4a8 => {
569 self.qmatvec_mmq_nvfp4_w4a8(bytes, x, m, in_f, out_f, *scale, *rp)
570 }
571 q if q == crate::QT_NVFP4 => self.qmatvec_mmq_nvfp4(bytes, x, m, in_f, out_f, *scale),
572 q if q == crate::QT_Q4_K || q == crate::QT_Q5_K => {
573 let mut y = self.qmatvec_mmq_q45k_raw(bytes, x, m, in_f, out_f, q)?;
574 if *scale != 1.0 {
575 self.scale_inplace(&mut y, *scale, m * out_f)?;
576 }
577 Ok(y)
578 }
579 q if q == crate::QT_Q8_0 => {
580 if cfg!(memra_hopper_mma) && out_f % 64 == 0 && crate::wgmma_gemm_enabled() {
586 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
587 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
588 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, &aq, &ad, m, in_f, out_f)?;
589 if *scale != 1.0 {
590 self.scale_inplace(&mut y, *scale, m * out_f)?;
591 }
592 return Ok(y);
593 }
594 }
595 let mut y = self.qmatvec_mmq_q8_0_raw(bytes, x, m, in_f, out_f)?;
596 if *scale != 1.0 {
597 self.scale_inplace(&mut y, *scale, m * out_f)?;
598 }
599 Ok(y)
600 }
601 q if q == crate::QT_Q4_0 => {
602 let mut y = self.qmatvec_mmq_q4_0_raw(bytes, x, m, in_f, out_f, *rp)?;
603 if *scale != 1.0 {
604 self.scale_inplace(&mut y, *scale, m * out_f)?;
605 }
606 Ok(y)
607 }
608 q if q == crate::QT_IQ4_XS => {
609 let GpuTensor::Quant { row_bytes, .. } = w else { unreachable!() };
610 let mut y = self.qmatvec_mmq_iq4xs_raw(bytes, x, m, in_f, out_f, *row_bytes)?;
611 if *scale != 1.0 {
612 self.scale_inplace(&mut y, *scale, m * out_f)?;
613 }
614 Ok(y)
615 }
616 q => Err(format!("qmatvec_mmq: unsupported qtype {q}").into()),
617 }
618 }
619
620 pub fn qmatvec_mmq_iq4xs_raw(
622 &self,
623 bytes: &CudaSlice<u8>,
624 x: &CudaSlice<f32>,
625 m: usize,
626 in_f: usize,
627 out_f: usize,
628 row_bytes: usize,
629 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
630 assert!(
631 in_f % 256 == 0,
632 "MMQ IQ4_XS requires in_f % 256 == 0, got {in_f}"
633 );
634 let act_bytes = unsafe { memra_mmq_iq_experts_act_bytes(in_f as i32, m as i32) };
635 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
636 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
637 {
638 let stream = self.gpu.stream();
639 let (w_p, _gw) = bytes.device_ptr(&stream);
640 let (x_p, _gx) = x.device_ptr(&stream);
641 let (y_p, _gy) = y.device_ptr_mut(&stream);
642 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
643 let rc = unsafe {
644 memra_mmq_iq4xs_dense(
645 w_p as *const core::ffi::c_void,
646 x_p as *const f32,
647 y_p as *mut f32,
648 in_f as i32,
649 out_f as i32,
650 m as i32,
651 row_bytes as i64,
652 s_p as *mut core::ffi::c_void,
653 stream.cu_stream() as *mut core::ffi::c_void,
654 )
655 };
656 if rc != 0 {
657 return Err(format!("memra_mmq_iq4xs_dense rc={rc}").into());
658 }
659 }
660 Ok(y)
661 }
662
663 pub fn qmatvec_mmq_q45k_raw(
668 &self,
669 bytes: &CudaSlice<u8>,
670 x: &CudaSlice<f32>,
671 m: usize,
672 in_f: usize,
673 out_f: usize,
674 qtype: i32,
675 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
676 assert!(
677 in_f % 256 == 0,
678 "MMQ Q4_K/Q5_K requires in_f % 256 == 0, got {in_f}"
679 );
680 let act_bytes = unsafe { memra_mmq_q45k_act_bytes(in_f as i32, m as i32) };
681 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
682 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
683 {
684 let stream = self.gpu.stream();
685 let (w_p, _gw) = bytes.device_ptr(&stream);
686 let (x_p, _gx) = x.device_ptr(&stream);
687 let (y_p, _gy) = y.device_ptr_mut(&stream);
688 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
689 let launcher = if qtype == crate::QT_Q4_K {
690 memra_mmq_q4_K
691 } else {
692 memra_mmq_q5_K
693 };
694 let rc = unsafe {
695 launcher(
696 w_p as *const core::ffi::c_void,
697 x_p as *const f32,
698 y_p as *mut f32,
699 in_f as i32,
700 out_f as i32,
701 m as i32,
702 s_p as *mut core::ffi::c_void,
703 stream.cu_stream() as *mut core::ffi::c_void,
704 )
705 };
706 if rc != 0 {
707 return Err(format!("memra_mmq_q45k(qtype={qtype}) rc={rc}").into());
708 }
709 }
710 Ok(y)
711 }
712
713 pub fn qmatvec_mmq_q8_0_raw(
716 &self,
717 bytes: &CudaSlice<u8>,
718 x: &CudaSlice<f32>,
719 m: usize,
720 in_f: usize,
721 out_f: usize,
722 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
723 assert!(
724 in_f % 32 == 0,
725 "MMQ Q8_0 requires in_f % 32 == 0, got {in_f}"
726 );
727 let act_bytes = unsafe { memra_mmq_q8_0_act_bytes(in_f as i32, m as i32) };
728 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
729 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
730 {
731 let stream = self.gpu.stream();
732 let (w_p, _gw) = bytes.device_ptr(&stream);
733 let (x_p, _gx) = x.device_ptr(&stream);
734 let (y_p, _gy) = y.device_ptr_mut(&stream);
735 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
736 let rc = unsafe {
737 memra_mmq_q8_0(
738 w_p as *const core::ffi::c_void,
739 x_p as *const f32,
740 y_p as *mut f32,
741 in_f as i32,
742 out_f as i32,
743 m as i32,
744 s_p as *mut core::ffi::c_void,
745 stream.cu_stream() as *mut core::ffi::c_void,
746 )
747 };
748 if rc != 0 {
749 return Err(format!("memra_mmq_q8_0 rc={rc}").into());
750 }
751 }
752 Ok(y)
753 }
754
755 pub fn accprobe_act_bytes(&self, in_f: usize, m: usize) -> usize {
758 unsafe { memra_accprobe_act_bytes(in_f as i32, m as i32) }
759 }
760
761 pub fn accprobe_gemm(
768 &self,
769 w_q8_0: &CudaSlice<u8>,
770 act_q: &CudaSlice<u8>,
771 m: usize,
772 in_f: usize,
773 out_f: usize,
774 f32acc: bool,
775 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
776 assert!(in_f % 32 == 0, "accprobe requires in_f % 32 == 0, got {in_f}");
777 assert!(
778 act_q.len() >= self.accprobe_act_bytes(in_f, m),
779 "accprobe act_q too small: {} < {}",
780 act_q.len(),
781 self.accprobe_act_bytes(in_f, m)
782 );
783 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
784 {
785 let stream = self.gpu.stream();
786 let (w_p, _gw) = w_q8_0.device_ptr(&stream);
787 let (a_p, _ga) = act_q.device_ptr(&stream);
788 let (y_p, _gy) = y.device_ptr_mut(&stream);
789 let f = if f32acc {
790 memra_accprobe_gemm_f32
791 } else {
792 memra_accprobe_gemm_s32
793 };
794 let rc = unsafe {
795 f(
796 w_p as *const core::ffi::c_void,
797 a_p as *const core::ffi::c_void,
798 y_p as *mut f32,
799 in_f as i32,
800 out_f as i32,
801 m as i32,
802 stream.cu_stream() as *mut core::ffi::c_void,
803 )
804 };
805 if rc != 0 {
806 let arm = if f32acc { "f32" } else { "s32" };
807 return Err(format!("memra_accprobe_gemm_{arm} rc={rc}").into());
808 }
809 }
810 Ok(y)
811 }
812
813 pub fn mmq_act_begin(&self) {
819 use std::sync::atomic::Ordering;
820 MMQ_ACT_EPOCH.fetch_add(1, Ordering::Relaxed);
821 *MMQ_ACT_SLOT.lock().unwrap() = None;
822 }
823
824 pub fn qmatvec_mmq_q4_0_raw(
828 &self,
829 bytes: &CudaSlice<u8>,
830 x: &CudaSlice<f32>,
831 m: usize,
832 in_f: usize,
833 out_f: usize,
834 rp: bool,
835 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
836 use std::sync::atomic::Ordering;
837 assert!(
838 in_f % 32 == 0,
839 "MMQ Q4_0 requires in_f % 32 == 0, got {in_f}"
840 );
841 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
842 let stream = self.gpu.stream();
843 let (x_p, _gx) = x.device_ptr(&stream);
844 let epoch = MMQ_ACT_EPOCH.load(Ordering::Relaxed);
845 let mut slot = MMQ_ACT_SLOT.lock().unwrap();
847 let hit = matches!(&*slot,
848 Some((e, p, mm, inf, _)) if *e == epoch && *p == x_p as u64 && *mm == m && *inf == in_f);
849 if !hit {
850 let act_bytes = unsafe { memra_mmq_q4_0_act_bytes(in_f as i32, m as i32) };
851 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
852 {
853 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
854 let rc = unsafe {
855 memra_mmq_q4_0_quant_act(
856 x_p as *const f32,
857 s_p as *mut core::ffi::c_void,
858 in_f as i32,
859 m as i32,
860 stream.cu_stream() as *mut core::ffi::c_void,
861 )
862 };
863 if rc != 0 {
864 return Err(
865 format!("memra_mmq_q4_0_quant_act(in_f={in_f}, m={m}) rc={rc}").into()
866 );
867 }
868 }
869 *slot = Some((epoch, x_p as u64, m, in_f, scratch));
870 }
871 let scratch = &slot.as_ref().unwrap().4;
872 {
873 let (w_p, _gw) = bytes.device_ptr(&stream);
874 let (y_p, _gy) = y.device_ptr_mut(&stream);
875 let (s_p, _gs) = scratch.device_ptr(&stream);
876 static SK_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
898 let sk = match crate::MMQ_SK_FORCE.load(std::sync::atomic::Ordering::Relaxed) {
899 0 => false,
900 1 => true,
901 _ => *SK_ON.get_or_init(|| std::env::var("MEMRA_MMQ_SK")
902 .map(|v| v != "0").unwrap_or(!cfg!(memra_hopper_mma))),
903 };
904 let rc = if sk {
905 let mut fx = MMQ_FIXUP_SLOT.lock().unwrap();
906 if fx.is_none() {
907 let nb = unsafe { memra_mmq_q4_0_fixup_bytes() };
908 *fx = Some(self.alloc_uninit::<u8>(nb)?);
909 }
910 let (f_p, _gf) = fx.as_mut().unwrap().device_ptr_mut(&stream);
911 unsafe {
912 memra_mmq_q4_0_gemm_sk(
913 w_p as *const core::ffi::c_void,
914 s_p as *const core::ffi::c_void,
915 y_p as *mut f32,
916 f_p as *mut core::ffi::c_void,
917 in_f as i32,
918 out_f as i32,
919 m as i32,
920 stream.cu_stream() as *mut core::ffi::c_void,
921 rp as i32,
922 )
923 }
924 } else { unsafe {
925 memra_mmq_q4_0_gemm(
926 w_p as *const core::ffi::c_void,
927 s_p as *const core::ffi::c_void,
928 y_p as *mut f32,
929 in_f as i32,
930 out_f as i32,
931 m as i32,
932 stream.cu_stream() as *mut core::ffi::c_void,
933 rp as i32,
934 )
935 } };
936 if rc != 0 {
937 return Err(format!(
938 "memra_mmq_q4_0_gemm(rp={rp}, in_f={in_f}, out_f={out_f}, m={m}, wbytes={}) rc={rc}",
939 bytes.len()
940 )
941 .into());
942 }
943 }
944 Ok(y)
945 }
946
947 pub fn qmatvec_mmq_nvfp4(
953 &self,
954 bytes: &CudaSlice<u8>,
955 x: &CudaSlice<f32>,
956 m: usize,
957 in_f: usize,
958 out_f: usize,
959 scale: f32,
960 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
961 self.qmatvec_mmq_nvfp4_scaled(bytes, x, m, in_f, out_f, scale)
962 }
963
964 pub fn qmatvec_mmq_nvfp4_raw(
966 &self,
967 bytes: &CudaSlice<u8>,
968 x: &CudaSlice<f32>,
969 m: usize,
970 in_f: usize,
971 out_f: usize,
972 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
973 self.qmatvec_mmq_nvfp4_scaled(bytes, x, m, in_f, out_f, 1.0)
974 }
975
976 pub fn qmatvec_mmq_nvfp4_raw_v1(
980 &self,
981 bytes: &CudaSlice<u8>,
982 x: &CudaSlice<f32>,
983 m: usize,
984 in_f: usize,
985 out_f: usize,
986 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
987 self.qmatvec_mmq_nvfp4_inner(bytes, x, m, in_f, out_f, 1.0, false, 0)
988 }
989
990 pub fn qmatvec_mmq_nvfp4_raw_res(
992 &self,
993 bytes: &CudaSlice<u8>,
994 x: &CudaSlice<f32>,
995 m: usize,
996 in_f: usize,
997 out_f: usize,
998 residual_k: i32,
999 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1000 self.qmatvec_mmq_nvfp4_inner(bytes, x, m, in_f, out_f, 1.0, true, residual_k)
1001 }
1002
1003 fn qmatvec_mmq_nvfp4_scaled(
1004 &self,
1005 bytes: &CudaSlice<u8>,
1006 x: &CudaSlice<f32>,
1007 m: usize,
1008 in_f: usize,
1009 out_f: usize,
1010 scale: f32,
1011 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1012 self.qmatvec_mmq_nvfp4_inner(bytes, x, m, in_f, out_f, scale, true, mmq_residual_k())
1013 }
1014
1015 fn qmatvec_mmq_nvfp4_inner(
1016 &self,
1017 bytes: &CudaSlice<u8>,
1018 x: &CudaSlice<f32>,
1019 m: usize,
1020 in_f: usize,
1021 out_f: usize,
1022 scale: f32,
1023 per_token_scale: bool,
1024 residual_k: i32,
1025 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1026 assert!(
1027 in_f % 64 == 0,
1028 "MMQ NVFP4 requires in_f % 64 == 0, got {in_f}"
1029 );
1030 let act_bytes = unsafe { memra_mmq_nvfp4_act_bytes(in_f as i32, m as i32) };
1031 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1032 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
1033 {
1034 let stream = self.gpu.stream();
1035 let (w_p, _gw) = bytes.device_ptr(&stream);
1036 let (x_p, _gx) = x.device_ptr(&stream);
1037 let (y_p, _gy) = y.device_ptr_mut(&stream);
1038 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1039 let rc = unsafe {
1040 memra_mmq_nvfp4_ex2(
1041 w_p as *const core::ffi::c_void,
1042 x_p as *const f32,
1043 y_p as *mut f32,
1044 in_f as i32,
1045 out_f as i32,
1046 m as i32,
1047 s_p as *mut core::ffi::c_void,
1048 stream.cu_stream() as *mut core::ffi::c_void,
1049 scale,
1050 per_token_scale as i32,
1051 residual_k,
1052 )
1053 };
1054 if rc != 0 {
1055 return Err(format!("memra_mmq_nvfp4_ex2 rc={rc}").into());
1056 }
1057 }
1058 Ok(y)
1059 }
1060
1061 pub fn qmatvec_mmq_nvfp4_w4a8(
1066 &self,
1067 bytes: &CudaSlice<u8>,
1068 x: &CudaSlice<f32>,
1069 m: usize,
1070 in_f: usize,
1071 out_f: usize,
1072 scale: f32,
1073 rp: bool,
1074 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1075 self.qmatvec_mmq_nvfp4_w4a8_scaled(bytes, x, m, in_f, out_f, scale, rp)
1076 }
1077
1078 pub fn qmatvec_mmq_nvfp4_w4a8_raw(
1080 &self,
1081 bytes: &CudaSlice<u8>,
1082 x: &CudaSlice<f32>,
1083 m: usize,
1084 in_f: usize,
1085 out_f: usize,
1086 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1087 self.qmatvec_mmq_nvfp4_w4a8_scaled(bytes, x, m, in_f, out_f, 1.0, false)
1088 }
1089
1090 pub fn qmatvec_mmq_nvfp4_w4a8_raw_rp(
1093 &self,
1094 bytes: &CudaSlice<u8>,
1095 x: &CudaSlice<f32>,
1096 m: usize,
1097 in_f: usize,
1098 out_f: usize,
1099 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1100 self.qmatvec_mmq_nvfp4_w4a8_scaled(bytes, x, m, in_f, out_f, 1.0, true)
1101 }
1102
1103 fn qmatvec_mmq_nvfp4_w4a8_scaled(
1104 &self,
1105 bytes: &CudaSlice<u8>,
1106 x: &CudaSlice<f32>,
1107 m: usize,
1108 in_f: usize,
1109 out_f: usize,
1110 scale: f32,
1111 rp: bool,
1112 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1113 assert!(
1114 in_f % 64 == 0,
1115 "MMQ NVFP4 W4A8 requires in_f % 64 == 0, got {in_f}"
1116 );
1117 let act_bytes = unsafe { memra_mmq_nvfp4_w4a8_act_bytes(in_f as i32, m as i32) };
1118 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1119 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
1120 {
1121 let stream = self.gpu.stream();
1122 let (w_p, _gw) = bytes.device_ptr(&stream);
1123 let (x_p, _gx) = x.device_ptr(&stream);
1124 let (y_p, _gy) = y.device_ptr_mut(&stream);
1125 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1126 static F8F4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1129 let f8f4 = *F8F4.get_or_init(|| std::env::var("MEMRA_MMQ_F8F4").as_deref() == Ok("1"));
1130 let rc = unsafe {
1131 if f8f4 {
1132 memra_mmq_nvfp4_f8f4(
1133 w_p as *const core::ffi::c_void,
1134 x_p as *const f32,
1135 y_p as *mut f32,
1136 in_f as i32,
1137 out_f as i32,
1138 m as i32,
1139 s_p as *mut core::ffi::c_void,
1140 stream.cu_stream() as *mut core::ffi::c_void,
1141 scale,
1142 rp as i32,
1143 )
1144 } else {
1145 memra_mmq_nvfp4_w4a8(
1146 w_p as *const core::ffi::c_void,
1147 x_p as *const f32,
1148 y_p as *mut f32,
1149 in_f as i32,
1150 out_f as i32,
1151 m as i32,
1152 s_p as *mut core::ffi::c_void,
1153 stream.cu_stream() as *mut core::ffi::c_void,
1154 scale,
1155 rp as i32,
1156 )
1157 }
1158 };
1159 if rc != 0 {
1160 return Err(format!("memra_mmq_nvfp4_w4a8(f8f4={f8f4}) rc={rc}").into());
1161 }
1162 }
1163 Ok(y)
1164 }
1165
1166 pub fn qmatvec_mmq_fp8_blk(
1170 &self,
1171 w_e4m3: &CudaSlice<u8>,
1172 blk_scales: &CudaSlice<f32>,
1173 x: &CudaSlice<f32>,
1174 m: usize,
1175 in_f: usize,
1176 out_f: usize,
1177 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1178 self.qmatvec_mmq_fp8_blk_scaled(w_e4m3, blk_scales, x, m, in_f, out_f, 1.0)
1179 }
1180
1181 pub fn qmatvec_mmq_fp8_blk_scaled(
1182 &self,
1183 w_e4m3: &CudaSlice<u8>,
1184 blk_scales: &CudaSlice<f32>,
1185 x: &CudaSlice<f32>,
1186 m: usize,
1187 in_f: usize,
1188 out_f: usize,
1189 scale: f32,
1190 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1191 assert!(
1192 in_f % 16 == 0,
1193 "per-block FP8 MMQ requires in_f % 16 == 0, got {in_f}"
1194 );
1195 let want_scales = ((out_f + 127) / 128) * ((in_f + 127) / 128);
1196 assert!(
1197 blk_scales.len() >= want_scales,
1198 "blk_scales too small: {} < {want_scales}",
1199 blk_scales.len()
1200 );
1201 assert!(
1202 w_e4m3.len() >= out_f * in_f,
1203 "e4m3 plane too small: {} < {}",
1204 w_e4m3.len(),
1205 out_f * in_f
1206 );
1207 let act_bytes = unsafe { memra_mmq_fp8_blk_act_bytes(in_f as i32, m as i32) };
1208 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1209 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
1210 {
1211 let stream = self.gpu.stream();
1212 let (w_p, _gw) = w_e4m3.device_ptr(&stream);
1213 let (sc_p, _gsc) = blk_scales.device_ptr(&stream);
1214 let (x_p, _gx) = x.device_ptr(&stream);
1215 let (y_p, _gy) = y.device_ptr_mut(&stream);
1216 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1217 let rc = unsafe {
1218 memra_mmq_fp8_blk(
1219 w_p as *const core::ffi::c_void,
1220 sc_p as *const f32,
1221 x_p as *const f32,
1222 y_p as *mut f32,
1223 in_f as i32,
1224 out_f as i32,
1225 m as i32,
1226 s_p as *mut core::ffi::c_void,
1227 stream.cu_stream() as *mut core::ffi::c_void,
1228 scale,
1229 )
1230 };
1231 if rc != 0 {
1232 return Err(format!("memra_mmq_fp8_blk rc={rc}").into());
1233 }
1234 }
1235 Ok(y)
1236 }
1237
1238 pub fn fp8_blk_nan_count(
1242 &self,
1243 w_e4m3: &CudaSlice<u8>,
1244 ) -> Result<u32, Box<dyn std::error::Error>> {
1245 let mut cnt = self.htod_u32_v(&[0u32])?;
1246 let n = w_e4m3.len();
1247 {
1248 let stream = self.gpu.stream();
1249 let (w_p, _gw) = w_e4m3.device_ptr(&stream);
1250 let (c_p, _gc) = cnt.device_ptr_mut(&stream);
1251 let rc = unsafe {
1252 memra_fp8_blk_count_nan(
1253 w_p as *const core::ffi::c_void,
1254 n,
1255 c_p as *mut u32,
1256 stream.cu_stream() as *mut core::ffi::c_void,
1257 )
1258 };
1259 if rc != 0 {
1260 return Err(format!("memra_fp8_blk_count_nan rc={rc}").into());
1261 }
1262 }
1263 Ok(self.dtoh_u32(&cnt)?[0])
1264 }
1265
1266 pub fn mmq_iq_quantize_act(
1269 &self,
1270 x: &CudaSlice<f32>,
1271 in_f: usize,
1272 n_tokens: usize,
1273 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1274 let act_bytes = unsafe { memra_mmq_iq_experts_act_bytes(in_f as i32, n_tokens as i32) };
1275 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1276 {
1277 let stream = self.gpu.stream();
1278 let (x_p, _gx) = x.device_ptr(&stream);
1279 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1280 let rc = unsafe {
1281 memra_mmq_iq_quantize_act(
1282 x_p as *const f32,
1283 s_p as *mut core::ffi::c_void,
1284 in_f as i32,
1285 n_tokens as i32,
1286 stream.cu_stream() as *mut core::ffi::c_void,
1287 )
1288 };
1289 if rc != 0 {
1290 return Err(format!("memra_mmq_iq_quantize_act rc={rc}").into());
1291 }
1292 }
1293 Ok(scratch)
1294 }
1295
1296 pub fn mmq_iq_fused_act_quant(
1302 &self,
1303 gate: &CudaSlice<f32>,
1304 up: &CudaSlice<f32>,
1305 in_f: usize,
1306 n_tokens: usize,
1307 act_kind: i32,
1308 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1309 let act_bytes = unsafe { memra_mmq_iq_experts_act_bytes(in_f as i32, n_tokens as i32) };
1310 let mut scratch = self.alloc_uninit::<u8>(act_bytes)?;
1311 {
1312 let stream = self.gpu.stream();
1313 let (g_p, _gg) = gate.device_ptr(&stream);
1314 let (u_p, _gu) = up.device_ptr(&stream);
1315 let (s_p, _gs) = scratch.device_ptr_mut(&stream);
1316 let rc = unsafe {
1317 memra_mmq_iq_fused_act_quant(
1318 g_p as *const f32,
1319 u_p as *const f32,
1320 s_p as *mut core::ffi::c_void,
1321 in_f as i32,
1322 n_tokens as i32,
1323 act_kind,
1324 stream.cu_stream() as *mut core::ffi::c_void,
1325 )
1326 };
1327 if rc != 0 {
1328 return Err(format!("memra_mmq_iq_fused_act_quant rc={rc}").into());
1329 }
1330 }
1331 Ok(scratch)
1332 }
1333
1334 #[allow(clippy::too_many_arguments)]
1338 pub fn mmq_iq_experts(
1339 &self,
1340 table: &CudaSlice<u64>,
1341 proj: i32,
1342 n_expert: usize,
1343 ex_ids: &CudaSlice<i32>,
1344 ex_off: &CudaSlice<i32>,
1345 ex_pairs: &CudaSlice<i32>,
1346 pair_tok: &CudaSlice<i32>,
1347 act_scratch: &CudaSlice<u8>,
1348 in_f: usize,
1349 out_f: usize,
1350 n_active: usize,
1351 n_pairs: usize,
1352 n_tokens: usize,
1353 qtype: i32,
1354 row_bytes: usize,
1355 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1356 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1357 {
1358 let stream = self.gpu.stream();
1359 let (tab_p, _g0) = table.device_ptr(&stream);
1360 let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1361 let (eo_p, _g2) = ex_off.device_ptr(&stream);
1362 let (ep_p, _g3) = ex_pairs.device_ptr(&stream);
1363 let (pt_p, _g4) = pair_tok.device_ptr(&stream);
1364 let (as_p, _g5) = act_scratch.device_ptr(&stream);
1365 let (y_p, _g6) = y.device_ptr_mut(&stream);
1366 let rc = unsafe {
1367 memra_mmq_iq_experts(
1368 tab_p as *const u64,
1369 proj,
1370 n_expert as i32,
1371 ei_p as *const i32,
1372 eo_p as *const i32,
1373 ep_p as *const i32,
1374 pt_p as *const i32,
1375 as_p as *const core::ffi::c_void,
1376 y_p as *mut f32,
1377 in_f as i32,
1378 out_f as i32,
1379 n_active as i32,
1380 n_tokens as i32,
1381 qtype,
1382 row_bytes as i64,
1383 stream.cu_stream() as *mut core::ffi::c_void,
1384 )
1385 };
1386 if rc != 0 {
1387 return Err(format!("memra_mmq_iq_experts rc={rc}").into());
1388 }
1389 }
1390 Ok(y)
1391 }
1392
1393 pub fn moe_f16g_act(
1398 &self,
1399 x: &CudaSlice<f32>,
1400 pair_tok: Option<&CudaSlice<i32>>,
1401 in_f: usize,
1402 n_pairs: usize,
1403 ) -> Result<(CudaSlice<u8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1404 let mut act = self.alloc_uninit::<u8>(n_pairs * in_f * 2)?;
1405 let mut scales = self.alloc_uninit::<f32>(n_pairs)?;
1406 {
1407 let stream = self.gpu.stream();
1408 let (x_p, _gx) = x.device_ptr(&stream);
1409 let pt_p = match pair_tok {
1410 Some(pt) => { let (p, _g) = pt.device_ptr(&stream); p as *const i32 }
1411 None => std::ptr::null(),
1412 };
1413 let (a_p, _ga) = act.device_ptr_mut(&stream);
1414 let (s_p, _gs) = scales.device_ptr_mut(&stream);
1415 let rc = unsafe {
1416 memra_moe_f16g_gather_act(x_p as *const f32, pt_p,
1417 a_p as *mut core::ffi::c_void, s_p as *mut f32,
1418 in_f as i32, n_pairs as i32,
1419 stream.cu_stream() as *mut core::ffi::c_void)
1420 };
1421 if rc != 0 { return Err(format!("memra_moe_f16g_gather_act rc={rc}").into()); }
1422 }
1423 Ok((act, scales))
1424 }
1425
1426 #[allow(clippy::too_many_arguments)]
1436 pub fn moe_f16_grouped(
1437 &self,
1438 table: &CudaSlice<u64>,
1439 proj: i32,
1440 n_expert: usize,
1441 ex_ids: &CudaSlice<i32>,
1442 ex_off_host: &[i32],
1443 ex_off_dev: &CudaSlice<i32>,
1444 act_f16: &CudaSlice<u8>,
1445 act_scale: &CudaSlice<f32>,
1446 in_f: usize,
1447 out_f: usize,
1448 n_active: usize,
1449 n_pairs: usize,
1450 qtype: i32,
1451 row_bytes: usize,
1452 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1453 let sk = crate::moe_f16g_mode() >= 2 && in_f % 32 == 0;
1454 let (shape_sel, cross) = crate::moe_f16g_sk_params();
1462 if sk && shape_sel >= 0 && crate::moe_f16g_direct_on(qtype)
1463 && (qtype == crate::QT_Q4_K || qtype == crate::QT_Q6_K
1464 || qtype == crate::QT_IQ4_XS || qtype == crate::QT_IQ3_S)
1465 && in_f % 256 == 0 && n_active <= 512 && n_active > 0
1466 {
1467 let max_m = ex_off_host.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
1468 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1469 {
1470 let stream = self.gpu.stream();
1471 let (tab_p, _g0) = table.device_ptr(&stream);
1472 let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1473 let (a_p, _g2) = act_f16.device_ptr(&stream);
1474 let (s_p, _g3) = act_scale.device_ptr(&stream);
1475 let (off_p, _g4) = ex_off_dev.device_ptr(&stream);
1476 let (y_p, _g5) = y.device_ptr_mut(&stream);
1477 let rc = unsafe {
1478 memra_moe_kq_gemm_sk(tab_p as *const u64, proj, n_expert as i32,
1479 ei_p as *const i32, a_p as *const core::ffi::c_void, y_p as *mut f32,
1480 s_p as *const f32, off_p as *const i32, ex_off_host.as_ptr(),
1481 n_active as i32, max_m, in_f as i32, out_f as i32, qtype, cross,
1482 crate::moe_f16g_tail_on() as i32, row_bytes as i64,
1483 stream.cu_stream() as *mut core::ffi::c_void)
1484 };
1485 if rc != 0 { return Err(format!("memra_moe_kq_gemm_sk rc={rc}").into()); }
1486 }
1487 return Ok(y);
1488 }
1489 if !sk {
1493 static WARM: std::sync::Once = std::sync::Once::new();
1494 let mut warm_err = None;
1495 WARM.call_once(|| {
1496 let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1497 let w = self.alloc_uninit::<u8>(2 * 32 * 64 * 2)?;
1498 let a = self.alloc_uninit::<u8>(4 * 64 * 2)?;
1499 let mut yw = self.alloc_uninit::<u8>(4 * 32 * 2)?;
1500 let off = [0i32, 2, 4];
1501 let stream = self.gpu.stream();
1502 let (w_p, _a1) = w.device_ptr(&stream);
1503 let (a_p, _a2) = a.device_ptr(&stream);
1504 let (y_p, _a3) = yw.device_ptr_mut(&stream);
1505 let rc = unsafe {
1506 memra_moe_f16g_gemm(w_p as *const core::ffi::c_void,
1507 a_p as *const core::ffi::c_void, y_p as *mut core::ffi::c_void,
1508 off.as_ptr(), 2, 64, 32,
1509 stream.cu_stream() as *mut core::ffi::c_void)
1510 };
1511 if rc != 0 { return Err(format!("f16g warmup rc={rc}").into()); }
1512 self.gpu.stream().synchronize()?;
1513 Ok(())
1514 })();
1515 if let Err(e) = r { warm_err = Some(e.to_string()); }
1516 });
1517 if let Some(we) = warm_err { return Err(we.into()); }
1518 }
1519 let w_bytes = n_active * out_f * in_f * 2;
1520 let mut w_f16 = self.alloc_uninit::<u8>(w_bytes)?;
1521 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1522 {
1523 let stream = self.gpu.stream();
1524 let (tab_p, _g0) = table.device_ptr(&stream);
1525 let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1526 let (w_p, _g2) = w_f16.device_ptr_mut(&stream);
1527 let rc = unsafe {
1528 memra_moe_f16g_dequant(tab_p as *const u64, proj, n_expert as i32,
1529 ei_p as *const i32, w_p as *mut core::ffi::c_void,
1530 in_f as i32, out_f as i32, n_active as i32, qtype, row_bytes as i64,
1531 stream.cu_stream() as *mut core::ffi::c_void)
1532 };
1533 if rc != 0 { return Err(format!("memra_moe_f16g_dequant rc={rc}").into()); }
1534 let (a_p, _g3) = act_f16.device_ptr(&stream);
1535 let (s_p, _g6) = act_scale.device_ptr(&stream);
1536 let (y_p, _g5) = y.device_ptr_mut(&stream);
1537 if sk {
1538 let max_m = ex_off_host.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
1539 let (off_p, _g7) = ex_off_dev.device_ptr(&stream);
1540 let (shape_sel, cross) = crate::moe_f16g_sk_params();
1541 let rc = unsafe {
1542 memra_moe_f16g_gemm_sk(w_p as *const core::ffi::c_void,
1543 a_p as *const core::ffi::c_void, y_p as *mut f32,
1544 s_p as *const f32, off_p as *const i32, ex_off_host.as_ptr(),
1545 n_active as i32, max_m, in_f as i32, out_f as i32, shape_sel, cross,
1546 crate::moe_f16g_tail_on() as i32,
1547 stream.cu_stream() as *mut core::ffi::c_void)
1548 };
1549 if rc != 0 { return Err(format!("memra_moe_f16g_gemm_sk rc={rc}").into()); }
1550 } else {
1551 let mut y16 = self.alloc_uninit::<u8>(n_pairs * out_f * 2)?;
1552 let (y16_p, _g4) = y16.device_ptr_mut(&stream);
1553 let rc = unsafe {
1554 memra_moe_f16g_gemm(w_p as *const core::ffi::c_void,
1555 a_p as *const core::ffi::c_void, y16_p as *mut core::ffi::c_void,
1556 ex_off_host.as_ptr(), n_active as i32, in_f as i32, out_f as i32,
1557 stream.cu_stream() as *mut core::ffi::c_void)
1558 };
1559 if rc != 0 { return Err(format!("memra_moe_f16g_gemm rc={rc}").into()); }
1560 let rc = unsafe {
1561 memra_moe_f16g_h2f_scaled(y16_p as *const core::ffi::c_void, y_p as *mut f32,
1562 s_p as *const f32, out_f as i32, n_pairs as i32,
1563 stream.cu_stream() as *mut core::ffi::c_void)
1564 };
1565 if rc != 0 { return Err(format!("memra_moe_f16g_h2f_scaled rc={rc}").into()); }
1566 }
1567 }
1568 if !sk {
1573 self.gpu.stream().synchronize()?;
1574 }
1575 if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
1576 let wn = n_active * out_f * in_f;
1578 let an = n_pairs * in_f;
1579 let mut wf = self.alloc_uninit::<f32>(wn)?;
1580 let mut af = self.alloc_uninit::<f32>(an)?;
1581 {
1582 let stream = self.gpu.stream();
1583 let (w_p, _a) = w_f16.device_ptr(&stream);
1584 let (a_p, _b) = act_f16.device_ptr(&stream);
1585 let (wf_p, _c) = wf.device_ptr_mut(&stream);
1586 let (af_p, _d) = af.device_ptr_mut(&stream);
1587 unsafe {
1588 memra_moe_f16g_h2f(w_p as *const core::ffi::c_void, wf_p as *mut f32, wn,
1589 stream.cu_stream() as *mut core::ffi::c_void);
1590 memra_moe_f16g_h2f(a_p as *const core::ffi::c_void, af_p as *mut f32, an,
1591 stream.cu_stream() as *mut core::ffi::c_void);
1592 }
1593 }
1594 let (wh, ah, yh) = (self.dtoh(&wf)?, self.dtoh(&af)?, self.dtoh(&y)?);
1595 let scan = |v: &[f32]| -> (usize, f32) {
1596 let bad = v.iter().filter(|x| !x.is_finite()).count();
1597 let mx = v.iter().filter(|x| x.is_finite()).fold(0.0f32, |m, x| m.max(x.abs()));
1598 (bad, mx)
1599 };
1600 let (wb, wm) = scan(&wh); let (ab, am) = scan(&ah); let (yb, ym) = scan(&yh);
1601 eprintln!("[f16g-debug] proj={proj} w: bad={wb} max={wm:.3e} | act: bad={ab} \
1602 max={am:.3e} | y: bad={yb} max={ym:.3e} (na={n_active} np={n_pairs} \
1603 in={in_f} out={out_f})");
1604 }
1605 Ok(y)
1606 }
1607
1608 #[allow(clippy::too_many_arguments)]
1615 pub fn moe_f16g_gemm_sk_raw(&self, w_f16: &CudaSlice<u8>, act_f16: &CudaSlice<u8>,
1616 row_scale: &CudaSlice<f32>, ex_off_host: &[i32], ex_off_dev: &CudaSlice<i32>,
1617 in_f: usize, out_f: usize, n_pairs: usize, shape_sel: i32, cross: i32, tail: i32)
1618 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1619 let n_active = ex_off_host.len() - 1;
1620 let max_m = ex_off_host.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
1621 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1622 {
1623 let stream = self.gpu.stream();
1624 let (w_p, _g0) = w_f16.device_ptr(&stream);
1625 let (a_p, _g1) = act_f16.device_ptr(&stream);
1626 let (s_p, _g2) = row_scale.device_ptr(&stream);
1627 let (off_p, _g3) = ex_off_dev.device_ptr(&stream);
1628 let (y_p, _g4) = y.device_ptr_mut(&stream);
1629 let rc = unsafe {
1630 memra_moe_f16g_gemm_sk(w_p as *const core::ffi::c_void,
1631 a_p as *const core::ffi::c_void, y_p as *mut f32,
1632 s_p as *const f32, off_p as *const i32, ex_off_host.as_ptr(),
1633 n_active as i32, max_m, in_f as i32, out_f as i32, shape_sel, cross,
1634 tail, stream.cu_stream() as *mut core::ffi::c_void)
1635 };
1636 if rc != 0 { return Err(format!("memra_moe_f16g_gemm_sk rc={rc}").into()); }
1637 }
1638 Ok(y)
1639 }
1640
1641 #[allow(clippy::too_many_arguments)]
1646 pub fn moe_kq_gemm_sk_raw(&self, table: &CudaSlice<u64>, proj: i32, n_expert: usize,
1647 ex_ids: &CudaSlice<i32>, act_f16: &CudaSlice<u8>, row_scale: &CudaSlice<f32>,
1648 ex_off_host: &[i32], ex_off_dev: &CudaSlice<i32>,
1649 in_f: usize, out_f: usize, n_pairs: usize, qtype: i32, row_bytes: usize, cross: i32,
1650 tail: i32)
1651 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1652 let n_active = ex_off_host.len() - 1;
1653 let max_m = ex_off_host.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
1654 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
1655 {
1656 let stream = self.gpu.stream();
1657 let (tab_p, _g0) = table.device_ptr(&stream);
1658 let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1659 let (a_p, _g2) = act_f16.device_ptr(&stream);
1660 let (s_p, _g3) = row_scale.device_ptr(&stream);
1661 let (off_p, _g4) = ex_off_dev.device_ptr(&stream);
1662 let (y_p, _g5) = y.device_ptr_mut(&stream);
1663 let rc = unsafe {
1664 memra_moe_kq_gemm_sk(tab_p as *const u64, proj, n_expert as i32,
1665 ei_p as *const i32, a_p as *const core::ffi::c_void, y_p as *mut f32,
1666 s_p as *const f32, off_p as *const i32, ex_off_host.as_ptr(),
1667 n_active as i32, max_m, in_f as i32, out_f as i32, qtype, cross,
1668 tail, row_bytes as i64,
1669 stream.cu_stream() as *mut core::ffi::c_void)
1670 };
1671 if rc != 0 { return Err(format!("memra_moe_kq_gemm_sk rc={rc}").into()); }
1672 }
1673 Ok(y)
1674 }
1675
1676 pub fn moe_f16g_dequant_raw(&self, table: &CudaSlice<u64>, proj: i32, n_expert: usize,
1680 ex_ids: &CudaSlice<i32>, in_f: usize, out_f: usize, n_active: usize, qtype: i32,
1681 row_bytes: usize)
1682 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
1683 let mut w_f16 = self.alloc_uninit::<u8>(n_active * out_f * in_f * 2)?;
1684 {
1685 let stream = self.gpu.stream();
1686 let (tab_p, _g0) = table.device_ptr(&stream);
1687 let (ei_p, _g1) = ex_ids.device_ptr(&stream);
1688 let (w_p, _g2) = w_f16.device_ptr_mut(&stream);
1689 let rc = unsafe {
1690 memra_moe_f16g_dequant(tab_p as *const u64, proj, n_expert as i32,
1691 ei_p as *const i32, w_p as *mut core::ffi::c_void,
1692 in_f as i32, out_f as i32, n_active as i32, qtype, row_bytes as i64,
1693 stream.cu_stream() as *mut core::ffi::c_void)
1694 };
1695 if rc != 0 { return Err(format!("memra_moe_f16g_dequant rc={rc}").into()); }
1696 }
1697 Ok(w_f16)
1698 }
1699}