1use cudarc::driver::{
4 CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
5};
6use cudarc::nvrtc::Ptx;
7use std::sync::{Arc, Mutex};
8
9#[cfg(debug_assertions)]
10pub(crate) fn debug_assert_tensor_stream_device<T>(
11 tensor: &CudaSlice<T>,
12 stream: &CudaStream,
13 site: &str,
14) {
15 let tensor_dev = tensor.ordinal();
16 let stream_dev = stream.context().ordinal();
17 assert_eq!(
18 tensor_dev, stream_dev,
19 "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
20 );
21}
22
23pub use memra_gguf;
24pub use memra_runtime;
25
26pub mod forward;
27pub mod hybrid;
28pub mod hybrid_forward;
29pub mod model;
30pub mod sigrouter_contract;
31pub mod vision;
32pub mod vision_gemma;
33pub mod vision_pre;
34pub mod cache {
37 pub use memra_kv::*;
38}
39pub mod decode;
40pub mod decode_batch;
41pub mod dflash;
42pub mod eagle;
43pub mod gemma_spec;
44pub mod graph_update;
45pub mod mla;
49pub mod moesd;
50pub mod parallel;
51pub mod pp;
52pub mod round_stream;
53pub mod spec;
54pub use memra_sampling as sampler;
55
56pub fn moe_f16g_mode() -> u8 {
100 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
101 *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
102 Ok("0") => 0,
103 Ok("2") => 2,
104 Ok("3") => 3,
105 Ok(_) => 1,
106 Err(_) => 2,
109 })
110}
111pub fn moe_f16g_sk_params() -> (i32, i32) {
125 static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
126 *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
127 Ok("0") => (-1, 0),
128 Ok("32") => (0, i32::MAX),
129 Ok("128") => (0, 1),
130 _ => {
131 let cross = std::env::var("MEMRA_F16G_SK_CROSS")
132 .ok()
133 .and_then(|v| v.parse().ok())
134 .unwrap_or(64);
135 (0, cross)
136 }
137 })
138}
139pub fn moe_f16g_direct_on(qtype: i32) -> bool {
150 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
151 let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
152 Ok("0") => 0,
153 Ok("kq") => 1,
154 _ => 2,
155 });
156 match m {
157 0 => false,
158 1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
159 _ => true,
160 }
161}
162pub fn moe_f16g_tail_on() -> bool {
171 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
172 *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
173}
174
175pub fn moe_f16g_gemma_on() -> bool {
182 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
183 *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
184}
185
186pub fn moe_fuse_actq_on() -> bool {
190 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
191 *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
192}
193
194pub fn router_prefill_exact_on() -> bool {
204 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
205 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
206}
207
208pub fn router_kernel_on() -> bool {
209 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
210 *ON.get_or_init(|| {
211 let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
212 if !on {
213 eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
214 }
215 on
216 })
217}
218
219pub const ROUTER_BATCH_MIN_T: usize = 8;
234pub fn router_batch_on() -> bool {
235 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
236 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
237}
238mod cpu_experts;
239#[cfg(memra_cutlass)]
240pub mod cutlass_ffi;
241pub mod f16_ffi;
242pub mod fp8_ffi;
243pub mod mmq_ffi;
244pub mod moe_cache;
245pub mod prime_graph;
246pub mod spill;
247mod spill_pread;
248
249const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
256const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
257const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
258const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
259const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
260const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
261const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
263
264fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
270 assert!(
271 !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
272 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
273 );
274 match std::env::var("MEMRA_GEMM_FATBIN") {
275 Ok(path) => std::borrow::Cow::Owned(
276 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
277 ),
278 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
279 }
280}
281
282pub(crate) const fn portable_mma_gated() -> bool {
289 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
290}
291
292const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
297 (!portable_cuda || hopper_mma) && !no_gemm
298}
299
300const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
308const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
309const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
310const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
311const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
312
313pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
316
317fn flash_fatbin_bytes() -> &'static [u8] {
319 match kv_cache_formats() {
320 ("q8_0", "q5_1") => FLASH_FATBIN,
321 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
322 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
323 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
324 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
325 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
326 other => unreachable!("kv_cache_formats returned {other:?}"),
327 }
328}
329
330fn k1_launch_override() -> Option<(u32, u32, u32)> {
337 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
338 *K1.get_or_init(|| {
339 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
340 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
341 match p.as_slice() {
342 [bm, bn, w] => Some((*bm, *bn, *w)),
343 _ => None,
344 }
345 })
346}
347
348pub(crate) fn wgmma_gemm_enabled() -> bool {
355 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
356 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
357}
358
359pub const FA_VEC_MIN_TKV: usize = 96;
374pub fn fa_vec_min_tkv() -> usize {
378 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
379 *V.get_or_init(|| {
380 std::env::var("MEMRA_FA_VEC_MIN")
381 .ok()
382 .and_then(|v| v.parse().ok())
383 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
384 })
385}
386
387pub fn fa_f16pv_on() -> bool {
398 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
399 *ON.get_or_init(|| {
400 std::env::var("MEMRA_FA_F16PV")
401 .map(|v| v != "0")
402 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
403 })
404}
405
406pub fn fa512_hp_on() -> bool {
410 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
411 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
412}
413
414pub fn faw_hp_on() -> bool {
418 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
419 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
420}
421
422pub fn fa512_wide_warps() -> usize {
426 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
427 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
428 Ok("1") => 4,
429 _ => 2,
430 })
431}
432
433pub fn fa512_min_tkv() -> usize {
436 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
437 *FA512_MIN.get_or_init(|| {
438 std::env::var("MEMRA_FA512_MIN")
439 .ok()
440 .and_then(|v| v.parse().ok())
441 .unwrap_or(512)
442 })
443}
444pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
448 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
449pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
453pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
459 std::sync::atomic::AtomicBool::new(false);
460pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
467 std::sync::atomic::AtomicBool::new(true);
468pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
469 std::sync::atomic::AtomicUsize::new(16);
470pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
475pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
477pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
482pub use memra_kv::KV_FP8_FORCE;
485pub(crate) fn rms_block() -> u32 {
486 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
487 *V.get_or_init(|| {
488 std::env::var("MEMRA_RMS_BLOCK")
489 .ok()
490 .and_then(|v| v.parse().ok())
491 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
492 })
493}
494
495pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
496 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
497 if let Some(forced) = *S.get_or_init(|| {
498 std::env::var("MEMRA_FA_SPLIT")
499 .ok()
500 .and_then(|v| v.parse().ok())
501 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
502 }) {
503 return forced;
504 }
505 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
523 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
524 {
525 return if t_kv <= 8192 {
526 16
527 } else if t_kv <= 16384 {
528 64
529 } else {
530 128
531 };
532 }
533 let big_rig = fa_sm_count() >= 128;
534 if big_rig {
535 let _ = n_head_kv;
536 if t_kv <= 2048 {
537 16
538 } else if t_kv <= 16384 {
539 64
540 } else {
541 128
542 }
543 } else if n_head_kv <= 4 {
544 if t_kv <= 512 {
565 8
566 } else if t_kv <= 16384 {
567 64
568 } else {
569 128
570 }
571 } else {
572 if t_kv <= 8192 {
573 32
574 } else if t_kv <= 16384 {
575 64
576 } else {
577 128
578 }
579 }
580}
581
582fn fa_sm_count() -> i32 {
585 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
586 *N.get_or_init(|| {
587 cudarc::driver::result::init().ok();
588 cudarc::driver::result::device::get(0)
589 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
590 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
591 .unwrap_or(82)
592 })
593}
594
595fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
599 match head_dim {
600 256 => Ok(""),
601 128 => Ok("_hd128"),
602 d => Err(format!(
603 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
604 callers must gate to sdpa_naive"
605 )
606 .into()),
607 }
608}
609
610pub const QT_Q8_0: i32 = 0;
612pub const QT_Q4_K: i32 = 1;
613pub const QT_Q6_K: i32 = 2;
614pub const QT_Q5_K: i32 = 3;
615pub const QT_Q3_K: i32 = 4;
616pub const QT_IQ4_XS: i32 = 5;
617pub const QT_IQ3_S: i32 = 6;
618pub const QT_NVFP4: i32 = 7;
619pub const QT_F8_E4M3: i32 = 10;
625pub const QT_NVFP4_RP: i32 = 9;
628pub const QT_F32: i32 = 8;
630pub const QT_BF16: i32 = 11;
631pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
636pub const QT_F8_E4M3_BLK: i32 = 14;
652
653pub struct Engine {
655 pub gpu: memra_runtime::Gpu,
656 module: Arc<CudaModule>,
657 hybrid: Arc<CudaModule>,
658 qmatvec: Arc<CudaModule>,
659 flash: Arc<CudaModule>,
660 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
664 gemm: Arc<CudaModule>,
665 router: Arc<CudaModule>,
666 sample: Arc<CudaModule>,
668 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
672 moe_cache_layout: Mutex<Option<Vec<usize>>>,
676 capture_keep_on: std::sync::atomic::AtomicBool,
682 verify_exact: std::sync::atomic::AtomicBool,
687 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
688 pub copy_stream: Arc<CudaStream>,
690 #[cfg(memra_cutlass)]
697 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
698 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
702 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
705 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
709 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
713 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
715 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
716 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
721 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
726 router_stage: Mutex<Option<PinnedStage>>,
730}
731
732fn fa_v2_on() -> bool {
742 std::env::var("MEMRA_FA_V2")
748 .map(|v| v != "0")
749 .unwrap_or(true)
750}
751
752fn fa_v3_on() -> bool {
760 std::env::var("MEMRA_FA_V3")
764 .map(|v| v != "0")
765 .unwrap_or(true)
766}
767
768fn fa_v4_mode() -> &'static str {
773 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
774 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
775}
776fn fa_v4_on() -> bool {
777 fa_v4_mode() != "0"
778} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
787 std::sync::atomic::AtomicUsize::new(1024);
788pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
789 std::sync::atomic::AtomicUsize::new(usize::MAX);
790pub fn fa_v4_at_pub(t_kv: usize) -> bool {
791 fa_v4_at(t_kv)
792}
793fn fa_v4_at(t_kv: usize) -> bool {
794 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
795 let mx = *M.get_or_init(|| {
796 std::env::var("MEMRA_FA_V4_MAX")
797 .ok()
798 .and_then(|v| v.parse().ok())
799 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
800 });
801 fa_v4_on() && t_kv < mx
802}
803pub const FA_DEEP_MIN_DEFAULT: usize = 0;
817fn fa_deep_at(t_kv: usize) -> bool {
818 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
819 return false;
820 }
821 let min = std::env::var("MEMRA_FA_DEEP_MIN")
822 .ok()
823 .and_then(|v| v.parse().ok())
824 .unwrap_or(FA_DEEP_MIN_DEFAULT);
825 t_kv >= min
826}
827pub fn fa_deep_at_pub(t_kv: usize) -> bool {
829 fa_deep_at(t_kv)
830}
831
832fn fa_v3_active(head_dim: usize) -> bool {
833 fa_v3_on()
836 && head_dim % 128 == 0
837 && kv_cache_formats() == ("q8_0", "q5_1")
838 && !Engine::kv_fp8_on()
839}
840
841pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
849 std::env::var("MEMRA_NO_FA_VEC").is_err()
850 && t_kv >= fa_vec_min_tkv()
851 && head_dim == 256
852 && fa_v4_at(t_kv)
853 && !matches!(fa_v4_mode(), "noB3" | "stage")
854 && !Engine::kv_fp8_on()
855}
856pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
858 fa_split_keys(t_kv, n_head_kv)
859}
860
861struct PinnedStage {
866 ptr: *mut u8,
867 cap: usize,
868}
869unsafe impl Send for PinnedStage {}
870impl PinnedStage {
871 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
872 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
873 Ok(PinnedStage { ptr, cap })
874 }
875}
876impl Drop for PinnedStage {
877 fn drop(&mut self) {
878 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
879 }
880}
881
882pub const ARGMAX_NB: usize = 256;
885
886pub(crate) use memra_fa3_vl as fa3_vl_raw;
888
889unsafe extern "C" {
890 fn memra_fa3_prefill(
892 q16: *const core::ffi::c_void,
893 k16: *const core::ffi::c_void,
894 v16: *const core::ffi::c_void,
895 o: *mut f32,
896 t: i32,
897 h: i32,
898 hkv: i32,
899 d: i32,
900 scale: f32,
901 stream: *mut core::ffi::c_void,
902 ) -> i32;
903 pub(crate) fn memra_fa3_vl(
905 q16s: *const *const core::ffi::c_void,
906 k16s: *const *const core::ffi::c_void,
907 v16s: *const *const core::ffi::c_void,
908 os: *const *mut f32,
909 ts: *const i32,
910 b: i32,
911 h: i32,
912 hkv: i32,
913 d: i32,
914 scale: f32,
915 stream: *mut core::ffi::c_void,
916 ) -> i32;
917}
918
919#[repr(C)]
924#[derive(Clone, Copy)]
925pub struct WPtr8(pub [u64; 8]);
926unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
927
928#[repr(C)]
933#[derive(Clone, Copy, Default)]
934pub struct GdnSeqVl {
935 pub kb16: u64,
936 pub gcum: u64,
937 pub beta: u64,
938 pub u: u64,
939 pub wb16: u64,
940 pub y: u64,
941 pub ssnap: u64,
942 pub state_in: u64,
943 pub state_out: u64,
944 pub q: u64,
945 pub p: u64,
946 pub o: u64,
947 pub k: u64,
948 pub v: u64,
949 pub g: u64,
950 pub a: u64,
951 pub w: u64,
952 pub t: i32,
953 pub nc: i32,
954}
955unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
956#[repr(C)]
957#[derive(Clone, Copy)]
958pub struct GdnVl8(pub [GdnSeqVl; 8]);
959unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
960
961#[repr(C)]
964#[derive(Clone, Copy, Default)]
965pub struct GdnWVl {
966 pub qb16: u64,
967 pub pb16: u64,
968}
969unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
970#[repr(C)]
971#[derive(Clone, Copy)]
972pub struct GdnWVl8(pub [GdnWVl; 8]);
973unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
974
975#[repr(C)]
977#[derive(Clone, Copy, Default)]
978pub struct GdnPrepVl {
979 pub qkv: u64,
980 pub conv_state: u64,
981 pub conv_out: u64,
982 pub q_g: u64,
983 pub k_g: u64,
984 pub v_g: u64,
985 pub q_l2: u64,
986 pub k_l2: u64,
987 pub beta_raw: u64,
988 pub alpha: u64,
989 pub beta: u64,
990 pub g_log: u64,
991 pub o: u64,
992 pub z: u64,
993 pub gn: u64,
994 pub gn16: u64,
995 pub kb16: u64,
996 pub qb16: u64,
997 pub t: i32,
998 pub pad: i32,
999}
1000unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1001#[repr(C)]
1002#[derive(Clone, Copy)]
1003pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1004unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1005
1006#[repr(C)]
1008#[derive(Clone, Copy, Default)]
1009pub struct FaSeqVl {
1010 pub q: u64,
1011 pub k16: u64,
1012 pub v16: u64,
1013 pub o: u64,
1014 pub kf: u64,
1015 pub vf: u64,
1016 pub t: i32,
1017 pub pad: i32,
1018}
1019unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1020#[repr(C)]
1021#[derive(Clone, Copy)]
1022pub struct FaVl8(pub [FaSeqVl; 8]);
1023unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1024
1025#[repr(C)]
1027#[derive(Clone, Copy, Default)]
1028pub struct AttnPreVl {
1029 pub qf: u64,
1030 pub kf: u64,
1031 pub vf: u64,
1032 pub q: u64,
1033 pub gate: u64,
1034 pub qn: u64,
1035 pub kn: u64,
1036 pub kc: u64,
1037 pub vc: u64,
1038 pub t: i32,
1039 pub pad: i32,
1040}
1041unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1042#[repr(C)]
1043#[derive(Clone, Copy)]
1044pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1045unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1046
1047pub struct GdnChunkBufs {
1050 pub gcum: CudaSlice<f32>,
1051 pub a: CudaSlice<f32>,
1052 pub p: CudaSlice<f32>,
1053 pub u: CudaSlice<f32>,
1054 pub w: CudaSlice<f32>,
1055 pub kb16: CudaSlice<u8>,
1056 pub wb16: CudaSlice<u8>,
1057 pub y16: CudaSlice<u8>,
1058 pub ssnap16: CudaSlice<u8>,
1059 pub qb16: CudaSlice<u8>,
1060 pub pb16: CudaSlice<u8>,
1061 pub o: CudaSlice<f32>,
1062 pub t: usize,
1063 pub nc: usize,
1064}
1065
1066#[repr(C)]
1068#[derive(Clone, Copy)]
1069pub struct F32x8(pub [f32; 8]);
1070unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1071
1072pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1076
1077impl Engine {
1078 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1079 let gpu = memra_runtime::Gpu::new(ordinal)?;
1080 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1084 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1085 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1086 .and_then(|d| unsafe {
1087 Ok((
1088 cudarc::driver::result::device::get_attribute(
1089 d,
1090 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1091 )?,
1092 cudarc::driver::result::device::get_attribute(
1093 d,
1094 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1095 )?,
1096 ))
1097 })
1098 .unwrap_or((0, 0));
1099 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1100 let ok = matches!(
1101 (built, maj, min),
1102 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1103 );
1104 if !ok {
1105 return Err(format!(
1106 "memra was built for sm_{built} but device {ordinal} reports compute \
1107 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1108 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1109 )
1110 .into());
1111 }
1112 }
1113 unsafe {
1118 use cudarc::driver::sys;
1119 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1120 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1121 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1122 let mut thresh: u64 = u64::MAX;
1123 let _ = sys::cuMemPoolSetAttribute(
1124 pool,
1125 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1126 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1127 );
1128 }
1129 }
1130 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1131 let hybrid = gpu
1132 .ctx
1133 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1134 let qmatvec = gpu
1135 .ctx
1136 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1137 let flash = gpu
1138 .ctx
1139 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1140 let gemm = gpu
1141 .ctx
1142 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1143 let router = gpu
1144 .ctx
1145 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1146 let sample = gpu
1147 .ctx
1148 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1149 let copy_stream = gpu.ctx.new_stream()?;
1150 if std::env::var("MEMRA_EVT")
1166 .map(|v| v == "1")
1167 .unwrap_or(false)
1168 {
1169 } else {
1171 unsafe {
1172 gpu.ctx.disable_event_tracking();
1173 }
1174 }
1175 Ok(Self {
1176 gpu,
1177 module,
1178 hybrid,
1179 qmatvec,
1180 flash,
1181 flash_g: std::sync::OnceLock::new(),
1182 gemm,
1183 router,
1184 sample,
1185 moe_cache: Mutex::new(None),
1186 moe_cache_layout: Mutex::new(None),
1187 copy_stream,
1188 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1189 verify_exact: std::sync::atomic::AtomicBool::new(false),
1190 capture_keep: Mutex::new(Vec::new()),
1191 argmax_partials: Mutex::new(None),
1192 prime_deqw_ws: Mutex::new(None),
1193 router_stage: Mutex::new(None),
1194 fp8_scratch: Mutex::new(None),
1195 fa_vf16_scratch: Mutex::new(None),
1196 fa_part_pool: Mutex::new(None),
1197 fa_part_retired: Mutex::new(Vec::new()),
1198 fn_cache: Mutex::new(Default::default()),
1199 f16_scratch: Mutex::new(None),
1200 #[cfg(memra_cutlass)]
1201 cutlass_scratch: Mutex::new(None),
1202 })
1203 }
1204
1205 pub fn ctx(&self) -> &Arc<CudaContext> {
1206 &self.gpu.ctx
1207 }
1208
1209 pub fn pool_cached_bytes(&self) -> usize {
1227 let (reserved, used) = self.pool_reserved_used();
1228 reserved.saturating_sub(used)
1229 }
1230
1231 pub fn pool_reserved_used(&self) -> (usize, usize) {
1238 use cudarc::driver::sys;
1239 unsafe {
1240 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1241 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1242 != sys::CUresult::CUDA_SUCCESS
1243 {
1244 return (0, 0);
1245 }
1246 let (mut reserved, mut used) = (0u64, 0u64);
1247 if sys::cuMemPoolGetAttribute(
1248 pool,
1249 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1250 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1251 ) != sys::CUresult::CUDA_SUCCESS
1252 {
1253 return (0, 0);
1254 }
1255 if sys::cuMemPoolGetAttribute(
1256 pool,
1257 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1258 &mut used as *mut u64 as *mut core::ffi::c_void,
1259 ) != sys::CUresult::CUDA_SUCCESS
1260 {
1261 return (0, 0);
1262 }
1263 (reserved as usize, used as usize)
1264 }
1265 }
1266
1267 pub fn stream(&self) -> Arc<CudaStream> {
1270 self.gpu.stream()
1271 }
1272 pub fn gkv_on() -> bool {
1275 memra_kv::gkv_on()
1276 }
1277
1278 pub fn wkv_on() -> bool {
1290 memra_kv::wkv_on()
1291 }
1292
1293 pub fn kv_fp8_on() -> bool {
1299 memra_kv::kv_fp8_on()
1300 }
1301
1302 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1305 if head_dim == 512 && Self::gkv_on() {
1306 self.func_g(name)
1307 } else {
1308 self.func(name)
1309 }
1310 }
1311
1312 fn func_g(&self, name: &str) -> CudaFunction {
1316 let m = self.flash_g.get_or_init(|| {
1317 self.gpu
1318 .ctx
1319 .load_module(cudarc::nvrtc::Ptx::from_binary(
1320 FLASH_FATBIN_KF8VF8.to_vec(),
1321 ))
1322 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1323 });
1324 let key = format!("g:{name}");
1325 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1326 return f.clone();
1327 }
1328 let f = match m.load_function(name) {
1329 Ok(f) => f,
1330 Err(_) => self.func(name),
1331 };
1332 self.fn_cache.lock().unwrap().insert(key, f.clone());
1333 f
1334 }
1335
1336 fn func(&self, name: &str) -> CudaFunction {
1337 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1340 return f.clone();
1341 }
1342 let f = self
1343 .module
1344 .load_function(name)
1345 .or_else(|_| self.hybrid.load_function(name))
1346 .or_else(|_| self.qmatvec.load_function(name))
1347 .or_else(|_| self.flash.load_function(name))
1348 .or_else(|_| self.gemm.load_function(name))
1349 .or_else(|_| self.router.load_function(name))
1350 .or_else(|_| self.sample.load_function(name))
1351 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1352 self.fn_cache
1353 .lock()
1354 .unwrap()
1355 .insert(name.to_string(), f.clone());
1356 f
1357 }
1358
1359 pub fn scatter_trim_logits(
1362 &self,
1363 src: &CudaSlice<f32>,
1364 d2t: &CudaSlice<u32>,
1365 dst: &mut CudaSlice<f32>,
1366 d_vocab: usize,
1367 n_vocab: usize,
1368 ) -> Result<(), Box<dyn std::error::Error>> {
1369 let f1 = self.func("scatter_trim_logits_f32");
1370 let f2 = self.func("scatter_trim_logits_pass2_f32");
1371 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1372 let cfg1 = LaunchConfig {
1373 grid_dim: (256, 1, 1),
1374 block_dim: (256, 1, 1),
1375 shared_mem_bytes: 0,
1376 };
1377 let __s_b1 = self.gpu.stream();
1378 let mut b1 = __s_b1.launch_builder(&f1);
1379 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1380 unsafe {
1381 b1.launch(cfg1)?;
1382 }
1383 let cfg2 = LaunchConfig {
1384 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1385 block_dim: (256, 1, 1),
1386 shared_mem_bytes: 0,
1387 };
1388 let __s_b2 = self.gpu.stream();
1389 let mut b2 = __s_b2.launch_builder(&f2);
1390 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1391 unsafe {
1392 b2.launch(cfg2)?;
1393 }
1394 Ok(())
1395 }
1396
1397 #[allow(clippy::too_many_arguments)]
1403 pub fn filter_stats(
1404 &self,
1405 x: &CudaSlice<f32>,
1406 row_stride: usize,
1407 rows: &CudaSlice<i32>,
1408 out_th: &mut CudaSlice<f32>,
1409 out_z: &mut CudaSlice<f32>,
1410 out_max: &mut CudaSlice<f32>,
1411 n: usize,
1412 nrow: usize,
1413 temp: f32,
1414 top_k: i32,
1415 top_p: f32,
1416 min_p: f32,
1417 ) -> Result<(), Box<dyn std::error::Error>> {
1418 let f = self.func("filter_stats_f32");
1419 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1420 let cfg = LaunchConfig {
1421 grid_dim: (nrow as u32, 1, 1),
1422 block_dim: (1024, 1, 1),
1423 shared_mem_bytes: 0,
1424 };
1425 let __s_b = self.gpu.stream();
1426 let mut b = __s_b.launch_builder(&f);
1427 b.arg(x)
1428 .arg(&rs)
1429 .arg(rows)
1430 .arg(&mut *out_th)
1431 .arg(&mut *out_z)
1432 .arg(&mut *out_max)
1433 .arg(&ni)
1434 .arg(&nr)
1435 .arg(&temp)
1436 .arg(&top_k)
1437 .arg(&top_p)
1438 .arg(&min_p);
1439 unsafe {
1440 b.launch(cfg)?;
1441 }
1442 Ok(())
1443 }
1444
1445 #[allow(clippy::too_many_arguments)]
1447 pub fn softmax_gather_filtered(
1448 &self,
1449 x: &CudaSlice<f32>,
1450 row_stride: usize,
1451 ids: &CudaSlice<u32>,
1452 rows: &CudaSlice<i32>,
1453 th: &CudaSlice<f32>,
1454 z: &CudaSlice<f32>,
1455 out: &mut CudaSlice<f32>,
1456 n: usize,
1457 npair: usize,
1458 temp: f32,
1459 ) -> Result<(), Box<dyn std::error::Error>> {
1460 let f = self.func("softmax_gather_filtered_f32");
1461 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1462 let cfg = LaunchConfig {
1463 grid_dim: (npair as u32, 1, 1),
1464 block_dim: (256, 1, 1),
1465 shared_mem_bytes: 0,
1466 };
1467 let __s_b = self.gpu.stream();
1468 let mut b = __s_b.launch_builder(&f);
1469 b.arg(x)
1470 .arg(&rs)
1471 .arg(ids)
1472 .arg(rows)
1473 .arg(th)
1474 .arg(z)
1475 .arg(&mut *out)
1476 .arg(&ni)
1477 .arg(&np)
1478 .arg(&temp);
1479 unsafe {
1480 b.launch(cfg)?;
1481 }
1482 Ok(())
1483 }
1484
1485 #[allow(clippy::too_many_arguments)]
1487 pub fn residual_sample_filtered(
1488 &self,
1489 p: &CudaSlice<f32>,
1490 q: Option<&CudaSlice<f32>>,
1491 n: usize,
1492 temp: f32,
1493 seed: u64,
1494 stream_pos: u32,
1495 p_stats: (f32, f32, f32),
1496 q_stats: (f32, f32, f32),
1497 out_tok: &mut CudaSlice<u32>,
1498 ) -> Result<(), Box<dyn std::error::Error>> {
1499 let f = self.func("residual_sample_filtered_f32");
1500 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1501 let has_q: i32 = q.is_some() as i32;
1502 let qbuf = q.unwrap_or(p);
1503 let (pm, pth, pz) = p_stats;
1504 let (qm, qth, qz) = q_stats;
1505 let cfg = LaunchConfig {
1506 grid_dim: (1, 1, 1),
1507 block_dim: (1024, 1, 1),
1508 shared_mem_bytes: 0,
1509 };
1510 let __s_b = self.gpu.stream();
1511 let mut b = __s_b.launch_builder(&f);
1512 b.arg(p)
1513 .arg(qbuf)
1514 .arg(&has_q)
1515 .arg(&ni)
1516 .arg(&temp)
1517 .arg(&slo)
1518 .arg(&shi)
1519 .arg(&stream_pos)
1520 .arg(&pm)
1521 .arg(&pth)
1522 .arg(&pz)
1523 .arg(&qm)
1524 .arg(&qth)
1525 .arg(&qz)
1526 .arg(&mut *out_tok);
1527 unsafe {
1528 b.launch(cfg)?;
1529 }
1530 Ok(())
1531 }
1532
1533 #[allow(clippy::too_many_arguments)]
1535 pub fn gumbel_perturb_filtered(
1536 &self,
1537 x: &CudaSlice<f32>,
1538 y: &mut CudaSlice<f32>,
1539 n: usize,
1540 seed: u64,
1541 stream_pos: u32,
1542 temp: f32,
1543 row_max: f32,
1544 th: f32,
1545 ) -> Result<(), Box<dyn std::error::Error>> {
1546 let f = self.func("gumbel_perturb_filtered_f32");
1547 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1548 let cfg = LaunchConfig {
1549 grid_dim: (n.div_ceil(256) as u32, 1, 1),
1550 block_dim: (256, 1, 1),
1551 shared_mem_bytes: 0,
1552 };
1553 let __s_b = self.gpu.stream();
1554 let mut b = __s_b.launch_builder(&f);
1555 b.arg(x)
1556 .arg(&mut *y)
1557 .arg(&ni)
1558 .arg(&slo)
1559 .arg(&shi)
1560 .arg(&stream_pos)
1561 .arg(&temp)
1562 .arg(&row_max)
1563 .arg(&th);
1564 unsafe {
1565 b.launch(cfg)?;
1566 }
1567 Ok(())
1568 }
1569
1570 #[allow(clippy::too_many_arguments)]
1574 pub fn penalize_logits(
1575 &self,
1576 x: &mut CudaSlice<f32>,
1577 hist: &CudaSlice<u32>,
1578 n_hist: usize,
1579 rep: f32,
1580 freq: f32,
1581 present: f32,
1582 n: usize,
1583 ) -> Result<(), Box<dyn std::error::Error>> {
1584 if n_hist == 0 {
1585 return Ok(());
1586 }
1587 let f = self.func("penalize_logits_f32");
1588 let (nh, ni) = (n_hist as i32, n as i32);
1589 let cfg = LaunchConfig {
1590 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1591 block_dim: (128, 1, 1),
1592 shared_mem_bytes: 0,
1593 };
1594 let __s_b = self.gpu.stream();
1595 let mut b = __s_b.launch_builder(&f);
1596 b.arg(&mut *x)
1597 .arg(hist)
1598 .arg(&nh)
1599 .arg(&rep)
1600 .arg(&freq)
1601 .arg(&present)
1602 .arg(&ni);
1603 unsafe {
1604 b.launch(cfg)?;
1605 }
1606 Ok(())
1607 }
1608
1609 #[allow(clippy::too_many_arguments)]
1611 pub fn penalize_logits_rows(
1612 &self,
1613 x: &mut CudaSlice<f32>,
1614 hist: &CudaSlice<u32>,
1615 n_hist: usize,
1616 rep: f32,
1617 freq: f32,
1618 present: f32,
1619 n: usize,
1620 nrow: usize,
1621 ) -> Result<(), Box<dyn std::error::Error>> {
1622 if n_hist == 0 || nrow == 0 {
1623 return Ok(());
1624 }
1625 let f = self.func("penalize_logits_rows_f32");
1626 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1627 let cfg = LaunchConfig {
1628 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1629 block_dim: (128, 1, 1),
1630 shared_mem_bytes: 0,
1631 };
1632 let __s_b = self.gpu.stream();
1633 let mut b = __s_b.launch_builder(&f);
1634 b.arg(&mut *x)
1635 .arg(hist)
1636 .arg(&nh)
1637 .arg(&rep)
1638 .arg(&freq)
1639 .arg(&present)
1640 .arg(&ni)
1641 .arg(&nr);
1642 unsafe {
1643 b.launch(cfg)?;
1644 }
1645 Ok(())
1646 }
1647
1648 pub fn wpf_level() -> u32 {
1656 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1657 *ON.get_or_init(|| {
1658 std::env::var("MEMRA_WPF")
1659 .ok()
1660 .and_then(|v| v.parse().ok())
1661 .unwrap_or(1)
1662 })
1663 }
1664
1665 pub fn set_verify_exact(&self, on: bool) {
1677 self.verify_exact
1678 .store(on, std::sync::atomic::Ordering::Relaxed);
1679 }
1680 pub(crate) fn verify_exact_on(&self) -> bool {
1681 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1682 }
1683
1684 pub fn qkv_append_on() -> bool {
1687 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1688 *ON.get_or_init(|| {
1689 std::env::var("MEMRA_QKV_APPEND")
1690 .map(|v| v != "0")
1691 .unwrap_or(true)
1692 })
1693 }
1694
1695 pub fn pdl_wb_on() -> bool {
1698 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1699 *ON.get_or_init(|| {
1700 std::env::var("MEMRA_PDL_WB")
1701 .map(|v| v != "0")
1702 .unwrap_or(true)
1703 })
1704 }
1705
1706 pub fn pdl_mmvq_on() -> bool {
1710 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1711 *ON.get_or_init(|| {
1712 std::env::var("MEMRA_PDL_MMVQ")
1713 .map(|v| v != "0")
1714 .unwrap_or(true)
1715 })
1716 }
1717
1718 pub fn pdl_on() -> bool {
1719 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1720 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1721 }
1722
1723 pub fn pdl_nvfp4q8_on() -> bool {
1729 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1730 *ON.get_or_init(|| {
1731 std::env::var("MEMRA_PDL_NVFP4")
1732 .map(|v| v != "0")
1733 .unwrap_or(true)
1734 })
1735 }
1736
1737 fn q40_mr1_on() -> bool {
1743 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1744 match *Q40MR.get_or_init(|| {
1745 std::env::var("MEMRA_Q40_MR")
1746 .ok()
1747 .and_then(|v| v.parse().ok())
1748 }) {
1749 Some(v) => v == 1,
1750 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1751 }
1752 }
1753
1754 fn pdl_func_flash(
1759 &self,
1760 g: bool,
1761 name: &'static str,
1762 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1763 use cudarc::driver::sys as cu;
1764 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1771 std::sync::Mutex::new(None);
1772 static FNS: std::sync::Mutex<
1773 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1774 > = std::sync::Mutex::new(None);
1775 let ctx_key = self.ctx().cu_ctx() as usize;
1776 if let Some(&f) = FNS
1777 .lock()
1778 .unwrap()
1779 .get_or_insert_with(Default::default)
1780 .get(&(ctx_key, g, name))
1781 {
1782 return Ok(f as cu::CUfunction);
1783 }
1784 let module = {
1785 let mut mods = MODS.lock().unwrap();
1786 let map = mods.get_or_insert_with(Default::default);
1787 match map.get(&(ctx_key, g)) {
1788 Some(&m) => m,
1789 None => {
1790 let m = self.pdl_load_module_in_ctx(if g {
1791 FLASH_FATBIN_KF8VF8
1792 } else {
1793 FLASH_FATBIN
1794 })?;
1795 map.insert((ctx_key, g), m);
1796 m
1797 }
1798 }
1799 };
1800 let cname = std::ffi::CString::new(name)?;
1801 let mut f: cu::CUfunction = std::ptr::null_mut();
1802 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1803 if r != cu::CUresult::CUDA_SUCCESS {
1804 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
1805 }
1806 FNS.lock()
1807 .unwrap()
1808 .get_or_insert_with(Default::default)
1809 .insert((ctx_key, g, name), f as usize);
1810 Ok(f)
1811 }
1812
1813 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1818 use cudarc::driver::sys as cu;
1819 let mut prev: cu::CUcontext = std::ptr::null_mut();
1820 unsafe {
1821 cu::cuCtxGetCurrent(&mut prev).result()?;
1822 }
1823 self.ctx().bind_to_thread()?;
1824 let mut m: cu::CUmodule = std::ptr::null_mut();
1825 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1826 let restore = if prev.is_null() {
1827 cu::CUresult::CUDA_SUCCESS
1828 } else {
1829 unsafe { cu::cuCtxSetCurrent(prev) }
1830 };
1831 if r != cu::CUresult::CUDA_SUCCESS {
1832 return Err(format!("pdl module load: {r:?}").into());
1833 }
1834 if restore != cu::CUresult::CUDA_SUCCESS {
1835 return Err(format!("pdl module load: ctx restore {restore:?}").into());
1836 }
1837 Ok(m as usize)
1838 }
1839
1840 fn pdl_func(
1841 &self,
1842 name: &'static str,
1843 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1844 use cudarc::driver::sys as cu;
1845 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1848 std::sync::Mutex::new(None);
1849 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1852 std::sync::Mutex::new(None);
1853 static FNS: std::sync::Mutex<
1854 Option<std::collections::HashMap<(usize, &'static str), usize>>,
1855 > = std::sync::Mutex::new(None);
1856 let ctx_key = self.ctx().cu_ctx() as usize;
1857 if let Some(&f) = FNS
1858 .lock()
1859 .unwrap()
1860 .get_or_insert_with(Default::default)
1861 .get(&(ctx_key, name))
1862 {
1863 return Ok(f as cu::CUfunction);
1864 }
1865 let module = {
1866 let mut mods = MODULES.lock().unwrap();
1867 let map = mods.get_or_insert_with(Default::default);
1868 match map.get(&ctx_key) {
1869 Some(&m) => m,
1870 None => {
1871 let m = self.pdl_load_module_in_ctx(FATBIN)?;
1872 map.insert(ctx_key, m);
1873 m
1874 }
1875 }
1876 };
1877 let cname = std::ffi::CString::new(name)?;
1878 let mut f: cu::CUfunction = std::ptr::null_mut();
1879 let mut r =
1880 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1881 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
1882 let qmodule = {
1883 let mut mods = QMODULES.lock().unwrap();
1884 let map = mods.get_or_insert_with(Default::default);
1885 match map.get(&ctx_key) {
1886 Some(&m) => m,
1887 None => {
1888 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
1889 map.insert(ctx_key, m);
1890 m
1891 }
1892 }
1893 };
1894 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
1895 }
1896 if r != cu::CUresult::CUDA_SUCCESS {
1897 return Err(format!("pdl_func {name}: {r:?}").into());
1898 }
1899 FNS.lock()
1900 .unwrap()
1901 .get_or_insert_with(Default::default)
1902 .insert((ctx_key, name), f as usize);
1903 Ok(f)
1904 }
1905
1906 unsafe fn launch_pdl_flash(
1918 &self,
1919 g: bool,
1920 name: &'static str,
1921 grid: (u32, u32, u32),
1922 block: (u32, u32, u32),
1923 smem: u32,
1924 params: &mut [*mut std::ffi::c_void],
1925 ) -> Result<(), Box<dyn std::error::Error>> {
1926 use cudarc::driver::sys as cu;
1927 let f = self.pdl_func_flash(g, name)?;
1928 if smem > 0 {
1929 let r =
1931 unsafe {
1932 cu::cuFuncSetAttribute(f,
1933 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
1934 smem as i32)
1935 };
1936 if r != cu::CUresult::CUDA_SUCCESS {
1937 return Err(format!("pdl smem attr {name}: {r:?}").into());
1938 }
1939 }
1940 let mut attr = cu::CUlaunchAttribute {
1941 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1942 pad: [0; 4],
1943 value: cu::CUlaunchAttributeValue {
1944 programmaticStreamSerializationAllowed: 1,
1945 },
1946 };
1947 let cfg = cu::CUlaunchConfig {
1948 gridDimX: grid.0,
1949 gridDimY: grid.1,
1950 gridDimZ: grid.2,
1951 blockDimX: block.0,
1952 blockDimY: block.1,
1953 blockDimZ: block.2,
1954 sharedMemBytes: smem,
1955 hStream: self.gpu.stream().cu_stream(),
1956 attrs: &mut attr,
1957 numAttrs: 1,
1958 };
1959 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1960 if r != cu::CUresult::CUDA_SUCCESS {
1961 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
1962 }
1963 Ok(())
1964 }
1965
1966 unsafe fn launch_pdl(
1967 &self,
1968 name: &'static str,
1969 grid: (u32, u32, u32),
1970 block: (u32, u32, u32),
1971 params: &mut [*mut std::ffi::c_void],
1972 ) -> Result<(), Box<dyn std::error::Error>> {
1973 use cudarc::driver::sys as cu;
1974 let f = self.pdl_func(name)?;
1975 let mut attr = cu::CUlaunchAttribute {
1976 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1977 pad: [0; 4],
1978 value: cu::CUlaunchAttributeValue {
1979 programmaticStreamSerializationAllowed: 1,
1980 },
1981 };
1982 let cfg = cu::CUlaunchConfig {
1983 gridDimX: grid.0,
1984 gridDimY: grid.1,
1985 gridDimZ: grid.2,
1986 blockDimX: block.0,
1987 blockDimY: block.1,
1988 blockDimZ: block.2,
1989 sharedMemBytes: 0,
1990 hStream: self.gpu.stream().cu_stream(),
1991 attrs: &mut attr,
1992 numAttrs: 1,
1993 };
1994 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1995 if r != cu::CUresult::CUDA_SUCCESS {
1996 return Err(format!("launch_pdl {name}: {r:?}").into());
1997 }
1998 Ok(())
1999 }
2000
2001 pub fn prefetch_weight_l2(
2004 &self,
2005 w: &crate::model::GpuTensor,
2006 ) -> Result<(), Box<dyn std::error::Error>> {
2007 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2008 let p = rp4.as_ref().unwrap_or(bytes);
2009 self.prefetch_l2(p, p.len())?;
2010 }
2011 Ok(())
2012 }
2013
2014 pub fn gather_row_bf16(
2017 &self,
2018 table: &CudaSlice<u8>,
2019 tok: &CudaSlice<u32>,
2020 idx: usize,
2021 dst: &mut CudaSlice<f32>,
2022 ncols: usize,
2023 ) -> Result<(), Box<dyn std::error::Error>> {
2024 let f = self.func("gather_row_bf16_f32");
2025 let cfg = LaunchConfig {
2026 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2027 block_dim: (256, 1, 1),
2028 shared_mem_bytes: 0,
2029 };
2030 let (nc, ix) = (ncols as i32, idx as i32);
2031 let __s_b = self.gpu.stream();
2032 let mut b = __s_b.launch_builder(&f);
2033 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2034 unsafe {
2035 b.launch(cfg)?;
2036 }
2037 Ok(())
2038 }
2039
2040 pub fn add_row_inplace(
2042 &self,
2043 logits: &mut CudaSlice<f32>,
2044 bias: &CudaSlice<f32>,
2045 n: usize,
2046 row_off: usize,
2047 ) -> Result<(), Box<dyn std::error::Error>> {
2048 let f = self.func("add_row_inplace_f32");
2049 let cfg = LaunchConfig {
2050 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2051 block_dim: (256, 1, 1),
2052 shared_mem_bytes: 0,
2053 };
2054 let (ni, off) = (n as i32, row_off as i64);
2055 let __s_b = self.gpu.stream();
2056 let mut b = __s_b.launch_builder(&f);
2057 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2058 unsafe {
2059 b.launch(cfg)?;
2060 }
2061 Ok(())
2062 }
2063
2064 pub fn prefetch_l2(
2066 &self,
2067 p: &CudaSlice<u8>,
2068 n: usize,
2069 ) -> Result<(), Box<dyn std::error::Error>> {
2070 let f = self.func("prefetch_l2_bytes");
2071 let lines = n.div_ceil(128);
2072 let ni = n as i64;
2073 let cfg = LaunchConfig {
2074 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2075 block_dim: (256, 1, 1),
2076 shared_mem_bytes: 0,
2077 };
2078 let __s_b = self.gpu.stream();
2079 let mut b = __s_b.launch_builder(&f);
2080 b.arg(p).arg(&ni);
2081 unsafe {
2082 b.launch(cfg)?;
2083 }
2084 Ok(())
2085 }
2086
2087 pub fn router_gemv(
2090 &self,
2091 w: &CudaSlice<f32>,
2092 x: &CudaSlice<f32>,
2093 n_embd: usize,
2094 n_experts: usize,
2095 t: usize,
2096 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2097 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2103 Ok("0") => false,
2104 Ok(_) => true,
2105 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2106 };
2107 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2117 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2118 }
2119
2120 pub fn router_gemv_form(
2123 &self,
2124 w: &CudaSlice<f32>,
2125 x: &CudaSlice<f32>,
2126 n_embd: usize,
2127 n_experts: usize,
2128 t: usize,
2129 w8: bool,
2130 batch: bool,
2131 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2132 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2133 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2134 let f = if batch {
2135 self.func("router_gemv_f32_w8_batch")
2136 } else if w8 {
2137 self.func("router_gemv_f32_w8")
2138 } else {
2139 self.func("router_gemv_f32")
2140 };
2141 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2142 let cfg = if batch {
2143 LaunchConfig {
2144 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2145 block_dim: (32, 8, 1),
2146 shared_mem_bytes: 0,
2147 }
2148 } else {
2149 LaunchConfig {
2150 grid_dim: (n_experts as u32, t as u32, 1),
2151 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2152 shared_mem_bytes: 0,
2153 }
2154 };
2155 let __s_b = self.gpu.stream();
2156 let mut b = __s_b.launch_builder(&f);
2157 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2158 unsafe {
2159 b.launch(cfg)?;
2160 }
2161 Ok(y)
2162 }
2163
2164 pub fn rows_permute(
2166 &self,
2167 src: &CudaSlice<f32>,
2168 idx: &CudaSlice<i32>,
2169 nrows: usize,
2170 ncols: usize,
2171 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2172 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2173 let f = self.func("rows_permute_f32");
2174 let (nc, nr) = (ncols as i32, nrows as i32);
2175 let cfg = LaunchConfig {
2176 grid_dim: (nrows as u32, 1, 1),
2177 block_dim: (256, 1, 1),
2178 shared_mem_bytes: 0,
2179 };
2180 let __s_b = self.gpu.stream();
2181 let mut b = __s_b.launch_builder(&f);
2182 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2183 unsafe {
2184 b.launch(cfg)?;
2185 }
2186 Ok(dst)
2187 }
2188
2189 pub fn sigmoid_dot_rows(
2194 &self,
2195 x: &CudaSlice<f32>,
2196 w: &CudaSlice<f32>,
2197 n_embd: usize,
2198 t: usize,
2199 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2200 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2203 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2204 let gs = self.linear(x, w, t, n_embd, 1)?;
2205 let mut g = self.uninit(t)?;
2206 self.sigmoid(&gs, &mut g, t)?;
2207 return Ok(g);
2208 }
2209 let mut g = self.alloc_uninit::<f32>(t)?;
2215 let f = self.func("sigmoid_dot_rows_f32");
2216 let (ne, ti) = (n_embd as i32, t as i32);
2217 let cfg = LaunchConfig {
2218 grid_dim: (t as u32, 1, 1),
2219 block_dim: (32, 8, 1),
2220 shared_mem_bytes: 0,
2221 };
2222 let __s_b = self.gpu.stream();
2223 let mut b = __s_b.launch_builder(&f);
2224 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2225 unsafe {
2226 b.launch(cfg)?;
2227 }
2228 Ok(g)
2229 }
2230
2231 pub fn spec_rollback_stream(
2233 &self,
2234 len_ptrs: &CudaSlice<u64>,
2235 pos_start: &CudaSlice<i32>,
2236 acc: &CudaSlice<u32>,
2237 base: usize,
2238 n_rows: usize,
2239 ) -> Result<(), Box<dyn std::error::Error>> {
2240 let f = self.func("spec_rollback_stream");
2241 let (b, nr) = (base as i32, n_rows as i32);
2242 let cfg = LaunchConfig {
2243 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2244 block_dim: (64, 1, 1),
2245 shared_mem_bytes: 0,
2246 };
2247 let __s_bl = self.gpu.stream();
2248 let mut bl = __s_bl.launch_builder(&f);
2249 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2250 unsafe {
2251 bl.launch(cfg)?;
2252 }
2253 Ok(())
2254 }
2255
2256 pub fn plain_tok_ring(
2258 &self,
2259 vam: &CudaSlice<u32>,
2260 pos_start: &CudaSlice<i32>,
2261 base: usize,
2262 ring: &mut CudaSlice<u32>,
2263 ) -> Result<(), Box<dyn std::error::Error>> {
2264 let f = self.func("plain_tok_ring");
2265 let (b, cap) = (base as i32, ring.len() as i32);
2266 let cfg = LaunchConfig {
2267 grid_dim: (1, 1, 1),
2268 block_dim: (32, 1, 1),
2269 shared_mem_bytes: 0,
2270 };
2271 let __s_bl = self.gpu.stream();
2272 let mut bl = __s_bl.launch_builder(&f);
2273 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2274 unsafe {
2275 bl.launch(cfg)?;
2276 }
2277 Ok(())
2278 }
2279
2280 pub fn spec_ring_commit(
2282 &self,
2283 vtok: &CudaSlice<u32>,
2284 acc: &CudaSlice<u32>,
2285 brk: &CudaSlice<u32>,
2286 ring: &mut CudaSlice<u32>,
2287 pend: &mut CudaSlice<u32>,
2288 ) -> Result<(), Box<dyn std::error::Error>> {
2289 let f = self.func("spec_ring_commit");
2290 let cfg = LaunchConfig {
2291 grid_dim: (1, 1, 1),
2292 block_dim: (32, 1, 1),
2293 shared_mem_bytes: 0,
2294 };
2295 let __s_b = self.gpu.stream();
2296 let mut b = __s_b.launch_builder(&f);
2297 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2298 unsafe {
2299 b.launch(cfg)?;
2300 }
2301 Ok(())
2302 }
2303 pub fn i32_copy_add(
2304 &self,
2305 src: &CudaSlice<i32>,
2306 dst: &mut CudaSlice<i32>,
2307 delta: i32,
2308 ) -> Result<(), Box<dyn std::error::Error>> {
2309 let f = self.func("i32_copy_add");
2310 let cfg = LaunchConfig {
2311 grid_dim: (1, 1, 1),
2312 block_dim: (32, 1, 1),
2313 shared_mem_bytes: 0,
2314 };
2315 let __s_b = self.gpu.stream();
2316 let mut b = __s_b.launch_builder(&f);
2317 b.arg(src).arg(dst).arg(&delta);
2318 unsafe {
2319 b.launch(cfg)?;
2320 }
2321 Ok(())
2322 }
2323 pub fn u32_copy(
2324 &self,
2325 src: &CudaSlice<u32>,
2326 dst: &mut CudaSlice<u32>,
2327 ) -> Result<(), Box<dyn std::error::Error>> {
2328 let f = self.func("u32_copy");
2329 let cfg = LaunchConfig {
2330 grid_dim: (1, 1, 1),
2331 block_dim: (32, 1, 1),
2332 shared_mem_bytes: 0,
2333 };
2334 let __s_b = self.gpu.stream();
2335 let mut b = __s_b.launch_builder(&f);
2336 b.arg(src).arg(dst);
2337 unsafe {
2338 b.launch(cfg)?;
2339 }
2340 Ok(())
2341 }
2342
2343 pub fn spec_adapt_k(
2347 &self,
2348 acc: &CudaSlice<u32>,
2349 brk: &mut CudaSlice<u32>,
2350 floor: usize,
2351 cap: usize,
2352 ) -> Result<(), Box<dyn std::error::Error>> {
2353 let f = self.func("spec_adapt_k");
2354 let (fl, cp) = (floor as i32, cap as i32);
2355 let cfg = LaunchConfig {
2356 grid_dim: (1, 1, 1),
2357 block_dim: (32, 1, 1),
2358 shared_mem_bytes: 0,
2359 };
2360 let __s_b = self.gpu.stream();
2361 let mut b = __s_b.launch_builder(&f);
2362 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2363 unsafe {
2364 b.launch(cfg)?;
2365 }
2366 Ok(())
2367 }
2368
2369 pub fn spec_accept_greedy_dc(
2371 &self,
2372 preds: &CudaSlice<u32>,
2373 vtok: &CudaSlice<u32>,
2374 last_pred: &CudaSlice<u32>,
2375 brk: &CudaSlice<u32>,
2376 out: &mut CudaSlice<u32>,
2377 ) -> Result<(), Box<dyn std::error::Error>> {
2378 let f = self.func("spec_accept_greedy_dc");
2379 let cfg = LaunchConfig {
2380 grid_dim: (1, 1, 1),
2381 block_dim: (32, 1, 1),
2382 shared_mem_bytes: 0,
2383 };
2384 let __s_b = self.gpu.stream();
2385 let mut b = __s_b.launch_builder(&f);
2386 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2387 unsafe {
2388 b.launch(cfg)?;
2389 }
2390 Ok(())
2391 }
2392
2393 pub fn pos_iota(
2395 &self,
2396 pos0: &CudaSlice<i32>,
2397 out: &mut CudaSlice<i32>,
2398 t: usize,
2399 ) -> Result<(), Box<dyn std::error::Error>> {
2400 let f = self.func("pos_iota_i32");
2401 let ti = t as i32;
2402 let cfg = LaunchConfig {
2403 grid_dim: (1, 1, 1),
2404 block_dim: (t.max(1) as u32, 1, 1),
2405 shared_mem_bytes: 0,
2406 };
2407 let __s_b = self.gpu.stream();
2408 let mut b = __s_b.launch_builder(&f);
2409 b.arg(pos0).arg(out).arg(&ti);
2410 unsafe {
2411 b.launch(cfg)?;
2412 }
2413 Ok(())
2414 }
2415 #[allow(clippy::too_many_arguments)]
2416 pub fn append_kv_quantized_rows_dc(
2417 &self,
2418 k_rows: &CudaSlice<f32>,
2419 v_rows: &CudaSlice<f32>,
2420 kc: &mut CudaSlice<u8>,
2421 vc: &mut CudaSlice<u8>,
2422 t0_dev: &CudaSlice<i32>,
2423 t: usize,
2424 kv_dim_k: usize,
2425 kv_dim_v: usize,
2426 k_tok_bytes: usize,
2427 v_tok_bytes: usize,
2428 g: bool,
2429 ) -> Result<(), Box<dyn std::error::Error>> {
2430 let f = if g {
2431 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2432 } else {
2433 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2434 };
2435 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2436 let cfg = LaunchConfig {
2437 grid_dim: (nblk, t as u32, 1),
2438 block_dim: (32, 1, 1),
2439 shared_mem_bytes: 0,
2440 };
2441 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2442 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2443 let __s_b = self.gpu.stream();
2444 let mut b = __s_b.launch_builder(&f);
2445 b.arg(k_rows)
2446 .arg(v_rows)
2447 .arg(kc)
2448 .arg(vc)
2449 .arg(t0_dev)
2450 .arg(&kdk)
2451 .arg(&kdv)
2452 .arg(&ktb)
2453 .arg(&vtb);
2454 unsafe {
2455 b.launch(cfg)?;
2456 }
2457 Ok(())
2458 }
2459
2460 #[allow(clippy::too_many_arguments)]
2463 pub fn append_kv_quantized_row_dc_inc(
2464 &self,
2465 k_row: &CudaSlice<f32>,
2466 v_row: &CudaSlice<f32>,
2467 kc: &mut CudaSlice<u8>,
2468 vc: &mut CudaSlice<u8>,
2469 t0_dev: &mut CudaSlice<i32>,
2470 kv_dim_k: usize,
2471 kv_dim_v: usize,
2472 k_tok_bytes: usize,
2473 v_tok_bytes: usize,
2474 g: bool,
2475 ) -> Result<(), Box<dyn std::error::Error>> {
2476 let f = if g {
2477 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2478 } else {
2479 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2480 };
2481 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2482 let cfg = LaunchConfig {
2483 grid_dim: (1, 1, 1),
2484 block_dim: (nthreads, 1, 1),
2485 shared_mem_bytes: 0,
2486 };
2487 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2488 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2489 let __s_b = self.gpu.stream();
2490 let mut b = __s_b.launch_builder(&f);
2491 b.arg(k_row)
2492 .arg(v_row)
2493 .arg(kc)
2494 .arg(vc)
2495 .arg(t0_dev)
2496 .arg(&kdk)
2497 .arg(&kdv)
2498 .arg(&ktb)
2499 .arg(&vtb);
2500 unsafe {
2501 b.launch(cfg)?;
2502 }
2503 Ok(())
2504 }
2505
2506 pub fn pack_tok_p(
2508 &self,
2509 tok: &CudaSlice<u32>,
2510 p: &CudaSlice<f32>,
2511 out: &mut CudaSlice<u32>,
2512 slot: usize,
2513 ) -> Result<(), Box<dyn std::error::Error>> {
2514 let f = self.func("pack_tok_p");
2515 let sl = slot as i32;
2516 let cfg = LaunchConfig {
2517 grid_dim: (1, 1, 1),
2518 block_dim: (32, 1, 1),
2519 shared_mem_bytes: 0,
2520 };
2521 let __s_b = self.gpu.stream();
2522 let mut b = __s_b.launch_builder(&f);
2523 b.arg(tok).arg(p).arg(out).arg(&sl);
2524 unsafe {
2525 b.launch(cfg)?;
2526 }
2527 Ok(())
2528 }
2529 pub fn tok_map_u32(
2530 &self,
2531 tok: &mut CudaSlice<u32>,
2532 map: &CudaSlice<u32>,
2533 ) -> Result<(), Box<dyn std::error::Error>> {
2534 let f = self.func("tok_map_u32");
2535 let cfg = LaunchConfig {
2536 grid_dim: (1, 1, 1),
2537 block_dim: (32, 1, 1),
2538 shared_mem_bytes: 0,
2539 };
2540 let __s_b = self.gpu.stream();
2541 let mut b = __s_b.launch_builder(&f);
2542 b.arg(tok).arg(map);
2543 unsafe {
2544 b.launch(cfg)?;
2545 }
2546 Ok(())
2547 }
2548
2549 #[allow(clippy::too_many_arguments)]
2551 pub fn spec_assemble_verify(
2552 &self,
2553 tokp: &CudaSlice<u32>,
2554 pend: &CudaSlice<u32>,
2555 d2t: Option<&CudaSlice<u32>>,
2556 vtok: &mut CudaSlice<u32>,
2557 brk: &mut CudaSlice<u32>,
2558 p_min: f32,
2559 k: usize,
2560 pmin0: bool,
2561 ) -> Result<(), Box<dyn std::error::Error>> {
2562 let f = self.func("spec_assemble_verify");
2563 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2564 let cfg = LaunchConfig {
2565 grid_dim: (1, 1, 1),
2566 block_dim: (32, 1, 1),
2567 shared_mem_bytes: 0,
2568 };
2569 let __s_b = self.gpu.stream();
2570 let mut b = __s_b.launch_builder(&f);
2571 match d2t {
2572 Some(m) => {
2573 b.arg(tokp)
2574 .arg(pend)
2575 .arg(m)
2576 .arg(vtok)
2577 .arg(brk)
2578 .arg(&p_min)
2579 .arg(&ki)
2580 .arg(&pm);
2581 unsafe {
2582 b.launch(cfg)?;
2583 }
2584 }
2585 None => {
2586 let null: u64 = 0;
2587 b.arg(tokp)
2588 .arg(pend)
2589 .arg(&null)
2590 .arg(vtok)
2591 .arg(brk)
2592 .arg(&p_min)
2593 .arg(&ki)
2594 .arg(&pm);
2595 unsafe {
2596 b.launch(cfg)?;
2597 }
2598 }
2599 }
2600 Ok(())
2601 }
2602
2603 #[allow(clippy::too_many_arguments)]
2605 pub fn ssm_conv_ring_rebuild_dc(
2606 &self,
2607 qkv_tm: &CudaSlice<f32>,
2608 ring_old: &CudaSlice<f32>,
2609 conv_state: &mut CudaSlice<f32>,
2610 conv_dim: usize,
2611 acc: &CudaSlice<u32>,
2612 base: usize,
2613 t_v: usize,
2614 d_conv: usize,
2615 ) -> Result<(), Box<dyn std::error::Error>> {
2616 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2617 let n = conv_dim * (d_conv - 1);
2618 let cfg = LaunchConfig::for_num_elems(n as u32);
2619 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2620 let __s_b = self.gpu.stream();
2621 let mut b = __s_b.launch_builder(&f);
2622 b.arg(qkv_tm)
2623 .arg(ring_old)
2624 .arg(conv_state)
2625 .arg(&cd)
2626 .arg(acc)
2627 .arg(&b0)
2628 .arg(&tv)
2629 .arg(&dc);
2630 unsafe {
2631 b.launch(cfg)?;
2632 }
2633 Ok(())
2634 }
2635 #[allow(clippy::too_many_arguments)]
2636 pub fn gdn_scan_s128_dc(
2637 &self,
2638 q: &CudaSlice<f32>,
2639 k: &CudaSlice<f32>,
2640 v: &CudaSlice<f32>,
2641 g: &CudaSlice<f32>,
2642 beta: &CudaSlice<f32>,
2643 state_in: &CudaSlice<f32>,
2644 state_out: &mut CudaSlice<f32>,
2645 o: &mut CudaSlice<f32>,
2646 n_head: usize,
2647 acc: &CudaSlice<u32>,
2648 base: usize,
2649 t_v: usize,
2650 scale: f32,
2651 ) -> Result<(), Box<dyn std::error::Error>> {
2652 let f = self.func("gdn_scan_s128_dc");
2653 const S_V: u32 = 128;
2654 const WARP: u32 = 32;
2655 const COLS_PER_BLOCK: u32 = 4;
2656 let cfg = LaunchConfig {
2657 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2658 block_dim: (WARP, COLS_PER_BLOCK, 1),
2659 shared_mem_bytes: 0,
2660 };
2661 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2662 let __s_b = self.gpu.stream();
2663 let mut b = __s_b.launch_builder(&f);
2664 b.arg(q)
2665 .arg(k)
2666 .arg(v)
2667 .arg(g)
2668 .arg(beta)
2669 .arg(state_in)
2670 .arg(state_out)
2671 .arg(o)
2672 .arg(&h)
2673 .arg(acc)
2674 .arg(&b0)
2675 .arg(&tv)
2676 .arg(&scale);
2677 unsafe {
2678 b.launch(cfg)?;
2679 }
2680 Ok(())
2681 }
2682
2683 pub fn spec_rollback_kv(
2685 &self,
2686 len_ptrs: &CudaSlice<u64>,
2687 saved: &CudaSlice<i32>,
2688 acc: &CudaSlice<u32>,
2689 base: usize,
2690 n_layer: usize,
2691 ) -> Result<(), Box<dyn std::error::Error>> {
2692 let f = self.func("spec_rollback_kv");
2693 let (b, nl) = (base as i32, n_layer as i32);
2694 let cfg = LaunchConfig {
2695 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2696 block_dim: (64, 1, 1),
2697 shared_mem_bytes: 0,
2698 };
2699 let __s_bl = self.gpu.stream();
2700 let mut bl = __s_bl.launch_builder(&f);
2701 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
2702 unsafe {
2703 bl.launch(cfg)?;
2704 }
2705 Ok(())
2706 }
2707
2708 pub fn spec_fork_valid(
2710 &self,
2711 acc: &CudaSlice<u32>,
2712 optimistic_pending: u32,
2713 valid: &mut CudaSlice<u32>,
2714 ) -> Result<(), Box<dyn std::error::Error>> {
2715 let f = self.func("spec_fork_valid");
2716 let cfg = LaunchConfig {
2717 grid_dim: (1, 1, 1),
2718 block_dim: (1, 1, 1),
2719 shared_mem_bytes: 0,
2720 };
2721 let __s_bl = self.gpu.stream();
2722 let mut bl = __s_bl.launch_builder(&f);
2723 bl.arg(acc).arg(&optimistic_pending).arg(valid);
2724 unsafe {
2725 bl.launch(cfg)?;
2726 }
2727 Ok(())
2728 }
2729
2730 pub fn spec_fork_reconcile_kv(
2732 &self,
2733 len_ptrs: &CudaSlice<u64>,
2734 saved: &CudaSlice<i32>,
2735 acc: &CudaSlice<u32>,
2736 valid: &CudaSlice<u32>,
2737 base: usize,
2738 n_layer: usize,
2739 ) -> Result<(), Box<dyn std::error::Error>> {
2740 let f = self.func("spec_fork_reconcile_kv");
2741 let (b, nl) = (base as i32, n_layer as i32);
2742 let cfg = LaunchConfig {
2743 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2744 block_dim: (64, 1, 1),
2745 shared_mem_bytes: 0,
2746 };
2747 let __s_bl = self.gpu.stream();
2748 let mut bl = __s_bl.launch_builder(&f);
2749 bl.arg(len_ptrs)
2750 .arg(saved)
2751 .arg(acc)
2752 .arg(valid)
2753 .arg(&b)
2754 .arg(&nl);
2755 unsafe {
2756 bl.launch(cfg)?;
2757 }
2758 Ok(())
2759 }
2760
2761 pub fn spec_fork_restore_f32(
2763 &self,
2764 snapshot: &CudaSlice<f32>,
2765 state: &mut CudaSlice<f32>,
2766 valid: &CudaSlice<u32>,
2767 ) -> Result<(), Box<dyn std::error::Error>> {
2768 assert_eq!(
2769 snapshot.len(),
2770 state.len(),
2771 "fork recurrent snapshot shape mismatch"
2772 );
2773 let f = self.func("spec_fork_restore_f32");
2774 let n = state.len() as i32;
2775 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
2776 let cfg = LaunchConfig {
2777 grid_dim: (blocks, 1, 1),
2778 block_dim: (256, 1, 1),
2779 shared_mem_bytes: 0,
2780 };
2781 let __s_bl = self.gpu.stream();
2782 let mut bl = __s_bl.launch_builder(&f);
2783 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
2784 unsafe {
2785 bl.launch(cfg)?;
2786 }
2787 Ok(())
2788 }
2789
2790 pub fn spec_seed_gather(
2793 &self,
2794 vx: &CudaSlice<f32>,
2795 fill_prev: &CudaSlice<f32>,
2796 acc: &CudaSlice<u32>,
2797 h_seed: &mut CudaSlice<f32>,
2798 base: usize,
2799 n_embd: usize,
2800 ) -> Result<(), Box<dyn std::error::Error>> {
2801 let f = self.func("spec_seed_gather");
2802 let (b, ne) = (base as i32, n_embd as i32);
2803 let cfg = LaunchConfig {
2804 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
2805 block_dim: (256, 1, 1),
2806 shared_mem_bytes: 0,
2807 };
2808 let __s_bl = self.gpu.stream();
2809 let mut bl = __s_bl.launch_builder(&f);
2810 bl.arg(vx)
2811 .arg(fill_prev)
2812 .arg(acc)
2813 .arg(h_seed)
2814 .arg(&b)
2815 .arg(&ne);
2816 unsafe {
2817 bl.launch(cfg)?;
2818 }
2819 Ok(())
2820 }
2821
2822 pub fn spec_accept_greedy(
2824 &self,
2825 preds: &CudaSlice<u32>,
2826 draft: &CudaSlice<u32>,
2827 last_pred: u32,
2828 base: usize,
2829 k_round: usize,
2830 out: &mut CudaSlice<u32>,
2831 ) -> Result<(), Box<dyn std::error::Error>> {
2832 let f = self.func("spec_accept_greedy");
2833 let (b, k) = (base as i32, k_round as i32);
2834 let cfg = LaunchConfig {
2835 grid_dim: (1, 1, 1),
2836 block_dim: (32, 1, 1),
2837 shared_mem_bytes: 0,
2838 };
2839 let __s_bl = self.gpu.stream();
2840 let mut bl = __s_bl.launch_builder(&f);
2841 bl.arg(preds)
2842 .arg(draft)
2843 .arg(&last_pred)
2844 .arg(&b)
2845 .arg(&k)
2846 .arg(out);
2847 unsafe {
2848 bl.launch(cfg)?;
2849 }
2850 Ok(())
2851 }
2852
2853 pub fn gumbel_perturb(
2860 &self,
2861 x: &CudaSlice<f32>,
2862 y: &mut CudaSlice<f32>,
2863 n: usize,
2864 seed: u64,
2865 stream_pos: u32,
2866 temp: f32,
2867 ) -> Result<(), Box<dyn std::error::Error>> {
2868 let f = self.func("gumbel_perturb_f32");
2869 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2870 let cfg = LaunchConfig {
2871 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2872 block_dim: (256, 1, 1),
2873 shared_mem_bytes: 0,
2874 };
2875 let __s_b = self.gpu.stream();
2876 let mut b = __s_b.launch_builder(&f);
2877 b.arg(x)
2878 .arg(&mut *y)
2879 .arg(&ni)
2880 .arg(&slo)
2881 .arg(&shi)
2882 .arg(&stream_pos)
2883 .arg(&temp);
2884 unsafe {
2885 b.launch(cfg)?;
2886 }
2887 Ok(())
2888 }
2889
2890 pub fn mask_logits_col(
2898 &self,
2899 logits: &mut CudaSlice<f32>,
2900 mask: &CudaSlice<u32>,
2901 col: usize,
2902 n: usize,
2903 mask_words: usize,
2904 ) -> Result<(), Box<dyn std::error::Error>> {
2905 let f = self.func("mask_logits_f32");
2906 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
2907 let cfg = LaunchConfig {
2908 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
2909 block_dim: (256, 1, 1),
2910 shared_mem_bytes: 0,
2911 };
2912 let __s_b = self.gpu.stream();
2913 let mut b = __s_b.launch_builder(&f);
2914 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
2915 unsafe {
2916 b.launch(cfg)?;
2917 }
2918 Ok(())
2919 }
2920
2921 pub fn gumbel_perturb_col(
2928 &self,
2929 x: &CudaSlice<f32>,
2930 col: usize,
2931 y: &mut CudaSlice<f32>,
2932 n: usize,
2933 seed: u64,
2934 stream_pos: u32,
2935 temp: f32,
2936 ) -> Result<(), Box<dyn std::error::Error>> {
2937 let f = self.func("gumbel_perturb_f32");
2938 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2939 let col_view = x.slice(col * n..(col + 1) * n);
2940 let cfg = LaunchConfig {
2941 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2942 block_dim: (256, 1, 1),
2943 shared_mem_bytes: 0,
2944 };
2945 let __s_b = self.gpu.stream();
2946 let mut b = __s_b.launch_builder(&f);
2947 b.arg(&col_view)
2948 .arg(&mut *y)
2949 .arg(&ni)
2950 .arg(&slo)
2951 .arg(&shi)
2952 .arg(&stream_pos)
2953 .arg(&temp);
2954 unsafe {
2955 b.launch(cfg)?;
2956 }
2957 Ok(())
2958 }
2959
2960 #[allow(clippy::too_many_arguments)]
2966 pub fn gumbel_perturb_filtered_col(
2967 &self,
2968 x: &CudaSlice<f32>,
2969 col: usize,
2970 y: &mut CudaSlice<f32>,
2971 n: usize,
2972 seed: u64,
2973 stream_pos: u32,
2974 temp: f32,
2975 stat_max: &CudaSlice<f32>,
2976 stat_th: &CudaSlice<f32>,
2977 stat_idx: usize,
2978 ) -> Result<(), Box<dyn std::error::Error>> {
2979 let f = self.func("gumbel_perturb_filtered_col_f32");
2980 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2981 let (ci, si) = (col as i32, stat_idx as i32);
2982 let cfg = LaunchConfig {
2983 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2984 block_dim: (256, 1, 1),
2985 shared_mem_bytes: 0,
2986 };
2987 let __s_b = self.gpu.stream();
2988 let mut b = __s_b.launch_builder(&f);
2989 b.arg(x)
2990 .arg(&ci)
2991 .arg(&mut *y)
2992 .arg(&ni)
2993 .arg(&slo)
2994 .arg(&shi)
2995 .arg(&stream_pos)
2996 .arg(&temp)
2997 .arg(stat_max)
2998 .arg(stat_th)
2999 .arg(&si);
3000 unsafe {
3001 b.launch(cfg)?;
3002 }
3003 Ok(())
3004 }
3005
3006 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3011 let f = self.func("memra_sctr_inc");
3012 let cfg = LaunchConfig {
3013 grid_dim: (1, 1, 1),
3014 block_dim: (1, 1, 1),
3015 shared_mem_bytes: 0,
3016 };
3017 let __s_b = self.gpu.stream();
3018 let mut b = __s_b.launch_builder(&f);
3019 b.arg(&mut *ctr);
3020 unsafe {
3021 b.launch(cfg)?;
3022 }
3023 Ok(())
3024 }
3025
3026 pub fn gumbel_perturb_ctr(
3031 &self,
3032 x: &CudaSlice<f32>,
3033 y: &mut CudaSlice<f32>,
3034 n: usize,
3035 seed: u64,
3036 ctr: &CudaSlice<u32>,
3037 temp: f32,
3038 ) -> Result<(), Box<dyn std::error::Error>> {
3039 let f = self.func("gumbel_perturb_ctr_f32");
3040 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3041 let cfg = LaunchConfig {
3042 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3043 block_dim: (256, 1, 1),
3044 shared_mem_bytes: 0,
3045 };
3046 let __s_b = self.gpu.stream();
3047 let mut b = __s_b.launch_builder(&f);
3048 b.arg(x)
3049 .arg(&mut *y)
3050 .arg(&ni)
3051 .arg(&slo)
3052 .arg(&shi)
3053 .arg(ctr)
3054 .arg(&temp);
3055 unsafe {
3056 b.launch(cfg)?;
3057 }
3058 Ok(())
3059 }
3060
3061 pub fn softmax_gather(
3065 &self,
3066 x: &CudaSlice<f32>,
3067 row_stride: usize,
3068 ids: &CudaSlice<u32>,
3069 rows: &CudaSlice<i32>,
3070 out: &mut CudaSlice<f32>,
3071 n: usize,
3072 npair: usize,
3073 temp: f32,
3074 ) -> Result<(), Box<dyn std::error::Error>> {
3075 let f = self.func("softmax_gather_f32");
3076 let (ni, rs) = (n as i32, row_stride as i64);
3077 let np = npair as i32;
3078 let cfg = LaunchConfig {
3079 grid_dim: (npair as u32, 1, 1),
3080 block_dim: (256, 1, 1),
3081 shared_mem_bytes: 0,
3082 };
3083 let __s_b = self.gpu.stream();
3084 let mut b = __s_b.launch_builder(&f);
3085 b.arg(x)
3086 .arg(&rs)
3087 .arg(ids)
3088 .arg(rows)
3089 .arg(&mut *out)
3090 .arg(&ni)
3091 .arg(&np)
3092 .arg(&temp);
3093 unsafe {
3094 b.launch(cfg)?;
3095 }
3096 Ok(())
3097 }
3098
3099 pub fn residual_sample(
3103 &self,
3104 p: &CudaSlice<f32>,
3105 q: Option<&CudaSlice<f32>>,
3106 n: usize,
3107 temp: f32,
3108 seed: u64,
3109 stream_pos: u32,
3110 out_tok: &mut CudaSlice<u32>,
3111 ) -> Result<(), Box<dyn std::error::Error>> {
3112 let f = self.func("residual_sample_f32");
3113 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3114 let nth = 1024u32;
3115 let cfg = LaunchConfig {
3116 grid_dim: (1, 1, 1),
3117 block_dim: (nth, 1, 1),
3118 shared_mem_bytes: 0,
3119 };
3120 let has_q: i32 = q.is_some() as i32;
3121 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3123 let mut b = __s_b.launch_builder(&f);
3124 b.arg(p)
3125 .arg(qbuf)
3126 .arg(&has_q)
3127 .arg(&ni)
3128 .arg(&temp)
3129 .arg(&slo)
3130 .arg(&shi)
3131 .arg(&stream_pos)
3132 .arg(&mut *out_tok);
3133 unsafe {
3134 b.launch(cfg)?;
3135 }
3136 Ok(())
3137 }
3138
3139 pub fn with_moe_cache<R>(
3144 &self,
3145 max_block_bytes: usize,
3146 f: impl FnOnce(
3147 &mut crate::moe_cache::MoeSlotCache,
3148 &Engine,
3149 ) -> Result<R, Box<dyn std::error::Error>>,
3150 ) -> Result<R, Box<dyn std::error::Error>> {
3151 let mut guard = self.moe_cache.lock().unwrap();
3152 if guard.is_none() {
3153 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3154 }
3155 let cache = guard.as_mut().unwrap();
3156 f(cache, self)
3157 }
3158
3159 pub fn freeze_moe_cache(&self) {
3162 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3163 cache.freeze();
3164 }
3165 }
3166
3167 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3170 self.moe_cache
3171 .lock()
3172 .unwrap()
3173 .as_ref()
3174 .map(crate::moe_cache::MoeSlotCache::export_residency)
3175 }
3176
3177 pub(crate) fn moe_cache_frozen(&self) -> bool {
3178 self.moe_cache
3179 .lock()
3180 .unwrap()
3181 .as_ref()
3182 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3183 }
3184
3185 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3192 crate::cpu_experts::configured()
3193 && self.moe_cache_frozen()
3194 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3195 }
3196
3197 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3199 assert!(
3200 self.moe_cache.lock().unwrap().is_none(),
3201 "MoE cache layout configured after cache construction"
3202 );
3203 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3204 }
3205
3206 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3207 self.moe_cache_layout.lock().unwrap().clone()
3208 }
3209
3210 pub fn moe_cache_enabled() -> bool {
3212 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3213 }
3214
3215 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3218 let guard = self.moe_cache.lock().unwrap();
3219 guard
3220 .as_ref()
3221 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3222 }
3223
3224 pub fn cpu_expert_stats(
3228 &self,
3229 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3230 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3231 }
3232
3233 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3236 crate::cpu_experts::predictor_stats()
3237 }
3238
3239 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3240 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3241 }
3242
3243 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3246 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3247 }
3248
3249 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3252 let guard = self.moe_cache.lock().unwrap();
3253 guard
3254 .as_ref()
3255 .and_then(|cache| cache.pread_stats())
3256 .map(|stats| {
3257 (
3258 stats.reads,
3259 stats.bytes,
3260 stats.read_errors,
3261 stats.short_reads,
3262 stats.fallbacks,
3263 stats.buffer_waits,
3264 stats.ring_full,
3265 )
3266 })
3267 }
3268
3269 pub fn spill_config_fallbacks(&self) -> u64 {
3271 crate::spill_pread::config_fallbacks()
3272 }
3273
3274 pub fn moe_cache_reset_counters(&self) {
3276 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3277 c.reset_counters();
3278 }
3279 }
3280
3281 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3282 Ok(self.gpu.stream().clone_htod(v)?)
3283 }
3284
3285 pub fn htod_bytes_padded(
3289 &self,
3290 v: &[u8],
3291 pad: usize,
3292 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3293 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3294 {
3295 let mut view = d.slice_mut(0..v.len());
3296 self.gpu.stream().memcpy_htod(v, &mut view)?;
3297 }
3298 Ok(d)
3299 }
3300
3301 pub fn copy_into(
3303 &self,
3304 dst: &mut CudaSlice<f32>,
3305 off: usize,
3306 src: &CudaSlice<f32>,
3307 len: usize,
3308 ) -> Result<(), Box<dyn std::error::Error>> {
3309 let mut view = dst.slice_mut(off..off + len);
3310 self.gpu
3311 .stream()
3312 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3313 Ok(())
3314 }
3315
3316 pub fn copy_u8_into(
3319 &self,
3320 dst: &mut CudaSlice<u8>,
3321 off: usize,
3322 src: &CudaSlice<u8>,
3323 len: usize,
3324 ) -> Result<(), Box<dyn std::error::Error>> {
3325 let mut view = dst.slice_mut(off..off + len);
3326 self.gpu
3327 .stream()
3328 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3329 Ok(())
3330 }
3331
3332 pub fn copy_u8_range_into(
3334 &self,
3335 dst: &mut CudaSlice<u8>,
3336 dst_off: usize,
3337 src: &CudaSlice<u8>,
3338 src_off: usize,
3339 len: usize,
3340 ) -> Result<(), Box<dyn std::error::Error>> {
3341 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3342 self.gpu
3343 .stream()
3344 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3345 Ok(())
3346 }
3347
3348 pub fn prepare_kv_append(
3352 &self,
3353 kv: &mut crate::cache::KvLayer,
3354 retain_from: usize,
3355 append_rows: usize,
3356 ) -> Result<usize, Box<dyn std::error::Error>> {
3357 let Some(plan) = kv
3358 .ring
3359 .as_ref()
3360 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3361 .transpose()?
3362 else {
3363 return Ok(kv.len);
3364 };
3365 match plan {
3366 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3367 crate::cache::KvRingAppend::Rebase {
3368 src_row,
3369 keep_rows,
3370 new_base,
3371 write_row,
3372 } => {
3373 if keep_rows > 0 {
3374 let k_len = keep_rows * kv.k_tok_bytes;
3375 let v_len = keep_rows * kv.v_tok_bytes;
3376 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3377 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3378 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3379 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3380 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3381 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3382 }
3383 kv.ring.as_mut().unwrap().apply_rebase(new_base);
3384 Ok(write_row)
3385 }
3386 }
3387 }
3388
3389 pub fn htod_u8_into(
3392 &self,
3393 dst: &mut CudaSlice<u8>,
3394 off: usize,
3395 src: &[u8],
3396 ) -> Result<(), Box<dyn std::error::Error>> {
3397 let mut view = dst.slice_mut(off..off + src.len());
3398 self.gpu.stream().memcpy_htod(src, &mut view)?;
3399 Ok(())
3400 }
3401
3402 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3403 b.slice(0..len)
3404 }
3405
3406 pub fn view_u8_range<'a>(
3409 &self,
3410 b: &'a CudaSlice<u8>,
3411 start: usize,
3412 end: usize,
3413 ) -> cudarc::driver::CudaView<'a, u8> {
3414 b.slice(start..end)
3415 }
3416 pub fn view_u8<'a>(
3417 &self,
3418 b: &'a CudaSlice<u8>,
3419 len: usize,
3420 ) -> cudarc::driver::CudaView<'a, u8> {
3421 b.slice(0..len)
3422 }
3423
3424 pub fn append_kv_quantized(
3428 &self,
3429 k_row: &CudaSlice<f32>,
3430 v_row: &CudaSlice<f32>,
3431 kc: &mut CudaSlice<u8>,
3432 vc: &mut CudaSlice<u8>,
3433 t: usize,
3434 kv_dim_k: usize,
3435 kv_dim_v: usize,
3436 k_tok_bytes: usize,
3437 v_tok_bytes: usize,
3438 g: bool,
3439 ) -> Result<(), Box<dyn std::error::Error>> {
3440 let f = if g {
3441 self.func_g("append_quantize_kv_q8_0_q5_1")
3442 } else {
3443 self.func("append_quantize_kv_q8_0_q5_1")
3444 };
3445 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3446 let cfg = LaunchConfig {
3447 grid_dim: (nblk, 1, 1),
3448 block_dim: (32, 1, 1),
3449 shared_mem_bytes: 0,
3450 };
3451 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3452 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3453 let __s_b = self.gpu.stream();
3454 let mut b = __s_b.launch_builder(&f);
3455 b.arg(k_row)
3456 .arg(v_row)
3457 .arg(kc)
3458 .arg(vc)
3459 .arg(&ti)
3460 .arg(&kdk)
3461 .arg(&kdv)
3462 .arg(&ktb)
3463 .arg(&vtb);
3464 unsafe {
3465 b.launch(cfg)?;
3466 }
3467 Ok(())
3468 }
3469
3470 pub fn append_kv_quantized_dc(
3474 &self,
3475 k_row: &CudaSlice<f32>,
3476 v_row: &CudaSlice<f32>,
3477 kc: &mut CudaSlice<u8>,
3478 vc: &mut CudaSlice<u8>,
3479 t_dev: &CudaSlice<i32>,
3480 kv_dim_k: usize,
3481 kv_dim_v: usize,
3482 k_tok_bytes: usize,
3483 v_tok_bytes: usize,
3484 g: bool,
3485 ) -> Result<(), Box<dyn std::error::Error>> {
3486 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3487 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3488 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3489 if Self::pdl_on() && Self::pdl_wb_on() {
3491 use cudarc::driver::{DevicePtr, DevicePtrMut};
3492 let s = &self.gpu.stream();
3493 let (pk, _g0) = k_row.device_ptr(s);
3494 let (pv, _g1) = v_row.device_ptr(s);
3495 let (pkc, _g2) = kc.device_ptr_mut(s);
3496 let (pvc, _g3) = vc.device_ptr_mut(s);
3497 let (pt, _g4) = t_dev.device_ptr(s);
3498 let mut ps = [
3499 &pk as *const _ as *mut std::ffi::c_void,
3500 &pv as *const _ as *mut _,
3501 &pkc as *const _ as *mut _,
3502 &pvc as *const _ as *mut _,
3503 &pt as *const _ as *mut _,
3504 &kdk as *const _ as *mut _,
3505 &kdv as *const _ as *mut _,
3506 &ktb as *const _ as *mut _,
3507 &vtb as *const _ as *mut _,
3508 ];
3509 unsafe {
3510 self.launch_pdl_flash(
3511 g,
3512 "append_quantize_kv_q8_0_q5_1_dc",
3513 (nblk, 1, 1),
3514 (32, 1, 1),
3515 0,
3516 &mut ps,
3517 )?;
3518 }
3519 return Ok(());
3520 }
3521 let f = if g {
3522 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3523 } else {
3524 self.func("append_quantize_kv_q8_0_q5_1_dc")
3525 };
3526 let cfg = LaunchConfig {
3527 grid_dim: (nblk, 1, 1),
3528 block_dim: (32, 1, 1),
3529 shared_mem_bytes: 0,
3530 };
3531 let __s_b = self.gpu.stream();
3532 let mut b = __s_b.launch_builder(&f);
3533 b.arg(k_row)
3534 .arg(v_row)
3535 .arg(kc)
3536 .arg(vc)
3537 .arg(t_dev)
3538 .arg(&kdk)
3539 .arg(&kdv)
3540 .arg(&ktb)
3541 .arg(&vtb);
3542 unsafe {
3543 b.launch(cfg)?;
3544 }
3545 Ok(())
3546 }
3547
3548 #[allow(clippy::too_many_arguments)]
3555 pub fn append_kv_quantized_rows(
3556 &self,
3557 k_rows: &CudaSlice<f32>,
3558 v_rows: &CudaSlice<f32>,
3559 kc: &mut CudaSlice<u8>,
3560 vc: &mut CudaSlice<u8>,
3561 t0: usize,
3562 t: usize,
3563 kv_dim_k: usize,
3564 kv_dim_v: usize,
3565 k_tok_bytes: usize,
3566 v_tok_bytes: usize,
3567 g: bool,
3568 ) -> Result<(), Box<dyn std::error::Error>> {
3569 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3570 for i in 0..t {
3571 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3572 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3573 self.append_kv_quantized_view(
3574 &k_row,
3575 &v_row,
3576 kc,
3577 vc,
3578 t0 + i,
3579 kv_dim_k,
3580 kv_dim_v,
3581 k_tok_bytes,
3582 v_tok_bytes,
3583 g,
3584 )?;
3585 }
3586 return Ok(());
3587 }
3588 let f = if g {
3589 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3590 } else {
3591 self.func("append_quantize_kv_q8_0_q5_1_rows")
3592 };
3593 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3594 let cfg = LaunchConfig {
3595 grid_dim: (nblk, t as u32, 1),
3596 block_dim: (32, 1, 1),
3597 shared_mem_bytes: 0,
3598 };
3599 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3600 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3601 let __s_b = self.gpu.stream();
3602 let mut b = __s_b.launch_builder(&f);
3603 b.arg(k_rows)
3604 .arg(v_rows)
3605 .arg(kc)
3606 .arg(vc)
3607 .arg(&t0i)
3608 .arg(&kdk)
3609 .arg(&kdv)
3610 .arg(&ktb)
3611 .arg(&vtb);
3612 unsafe {
3613 b.launch(cfg)?;
3614 }
3615 Ok(())
3616 }
3617
3618 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3622 let f = self.func("inc_i32");
3623 let cfg = LaunchConfig {
3624 grid_dim: (1, 1, 1),
3625 block_dim: (1, 1, 1),
3626 shared_mem_bytes: 0,
3627 };
3628 let __s_b = self.gpu.stream();
3629 let mut b = __s_b.launch_builder(&f);
3630 b.arg(p);
3631 unsafe {
3632 b.launch(cfg)?;
3633 }
3634 Ok(())
3635 }
3636
3637 pub fn append_kv_quantized_view(
3640 &self,
3641 k_row: &cudarc::driver::CudaView<f32>,
3642 v_row: &cudarc::driver::CudaView<f32>,
3643 kc: &mut CudaSlice<u8>,
3644 vc: &mut CudaSlice<u8>,
3645 t: usize,
3646 kv_dim_k: usize,
3647 kv_dim_v: usize,
3648 k_tok_bytes: usize,
3649 v_tok_bytes: usize,
3650 g: bool,
3651 ) -> Result<(), Box<dyn std::error::Error>> {
3652 let f = if g {
3653 self.func_g("append_quantize_kv_q8_0_q5_1")
3654 } else {
3655 self.func("append_quantize_kv_q8_0_q5_1")
3656 };
3657 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3658 let cfg = LaunchConfig {
3659 grid_dim: (nblk, 1, 1),
3660 block_dim: (32, 1, 1),
3661 shared_mem_bytes: 0,
3662 };
3663 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3664 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3665 let __s_b = self.gpu.stream();
3666 let mut b = __s_b.launch_builder(&f);
3667 b.arg(k_row)
3668 .arg(v_row)
3669 .arg(kc)
3670 .arg(vc)
3671 .arg(&ti)
3672 .arg(&kdk)
3673 .arg(&kdv)
3674 .arg(&ktb)
3675 .arg(&vtb);
3676 unsafe {
3677 b.launch(cfg)?;
3678 }
3679 Ok(())
3680 }
3681
3682 pub fn copy_view_into(
3685 &self,
3686 dst: &mut CudaSlice<f32>,
3687 off: usize,
3688 src: &cudarc::driver::CudaView<f32>,
3689 len: usize,
3690 ) -> Result<(), Box<dyn std::error::Error>> {
3691 let mut view = dst.slice_mut(off..off + len);
3692 self.gpu
3693 .stream()
3694 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3695 Ok(())
3696 }
3697
3698 pub fn clone_dtod(
3702 &self,
3703 src: &CudaSlice<f32>,
3704 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3705 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
3706 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
3707 Ok(dst)
3708 }
3709
3710 pub fn dtod_copy_view(
3713 &self,
3714 src: &cudarc::driver::CudaView<f32>,
3715 dst: &mut CudaSlice<f32>,
3716 ) -> Result<(), Box<dyn std::error::Error>> {
3717 self.gpu.stream().memcpy_dtod(src, dst)?;
3718 Ok(())
3719 }
3720
3721 pub fn dtod_copy_view_i8(
3723 &self,
3724 src: &cudarc::driver::CudaView<i8>,
3725 dst: &mut CudaSlice<i8>,
3726 ) -> Result<(), Box<dyn std::error::Error>> {
3727 self.gpu.stream().memcpy_dtod(src, dst)?;
3728 Ok(())
3729 }
3730
3731 pub fn dtod_copy_into(
3733 &self,
3734 src: &CudaSlice<f32>,
3735 dst: &mut CudaSlice<f32>,
3736 offset: usize,
3737 ) -> Result<(), Box<dyn std::error::Error>> {
3738 let n = src.len();
3739 let mut dv = dst.slice_mut(offset..offset + n);
3740 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
3741 Ok(())
3742 }
3743
3744 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
3746 self.alloc_uninit::<i8>(n)
3747 }
3748
3749 pub fn qmatvec(
3751 &self,
3752 w: &CudaSlice<u8>,
3753 x: &CudaSlice<f32>,
3754 m: usize,
3755 in_f: usize,
3756 out_f: usize,
3757 qtype: i32,
3758 row_bytes: usize,
3759 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3760 let f = self.func("qmatvec_f32");
3761 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
3763 grid_dim: (out_f as u32, m as u32, 1),
3764 block_dim: (256, 1, 1),
3765 shared_mem_bytes: 0,
3766 };
3767 let (inf, outf, mi, qt, rb) =
3768 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
3769 let __s_b = self.gpu.stream();
3770 let mut b = __s_b.launch_builder(&f);
3771 b.arg(w)
3772 .arg(x)
3773 .arg(&mut y)
3774 .arg(&inf)
3775 .arg(&outf)
3776 .arg(&mi)
3777 .arg(&qt)
3778 .arg(&rb);
3779 unsafe {
3780 b.launch(cfg)?;
3781 }
3782 Ok(y)
3783 }
3784
3785 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3787 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
3788 self.keep_if_capturing(&s);
3789 Ok(s)
3790 }
3791
3792 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3796 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
3797 self.keep_if_capturing(&s);
3798 Ok(s)
3799 }
3800
3801 pub fn memset_zeros_view(
3804 &self,
3805 dst: &mut cudarc::driver::CudaViewMut<f32>,
3806 ) -> Result<(), Box<dyn std::error::Error>> {
3807 self.gpu.stream().memset_zeros(dst)?;
3808 Ok(())
3809 }
3810
3811 pub fn stage_expert(
3817 &self,
3818 host_bytes: &[u8],
3819 scratch: &mut CudaSlice<u8>,
3820 off: usize,
3821 ) -> Result<(), Box<dyn std::error::Error>> {
3822 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
3825 }
3826
3827 pub fn moe_router_topk(
3833 &self,
3834 logits: &CudaSlice<f32>,
3835 t: usize,
3836 n_expert: usize,
3837 n_used: usize,
3838 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3839 let f = self.func("moe_router_topk_f32");
3840 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?; let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?; let cfg = LaunchConfig {
3843 grid_dim: (t as u32, 1, 1),
3844 block_dim: (n_expert as u32, 1, 1),
3845 shared_mem_bytes: 0,
3846 };
3847 let (ne, nu) = (n_expert as i32, n_used as i32);
3848 let __s_b = self.gpu.stream();
3849 let mut b = __s_b.launch_builder(&f);
3850 b.arg(logits)
3851 .arg(&mut sel_idx)
3852 .arg(&mut sel_w)
3853 .arg(&ne)
3854 .arg(&nu);
3855 unsafe {
3856 b.launch(cfg)?;
3857 }
3858 Ok((sel_idx, sel_w))
3859 }
3860
3861 pub fn moe_router_topk_scaled(
3864 &self,
3865 logits: &CudaSlice<f32>,
3866 t: usize,
3867 n_expert: usize,
3868 n_used: usize,
3869 ex_scale: &CudaSlice<f32>,
3870 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3871 let f = self.func("moe_router_topk_scaled_f32");
3876 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
3877 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
3878 let cfg = LaunchConfig {
3879 grid_dim: (t as u32, 1, 1),
3880 block_dim: (n_expert as u32, 1, 1),
3881 shared_mem_bytes: 0,
3882 };
3883 let (ne, nu) = (n_expert as i32, n_used as i32);
3884 let __s_b = self.gpu.stream();
3885 let mut b = __s_b.launch_builder(&f);
3886 b.arg(logits)
3887 .arg(&mut sel_idx)
3888 .arg(&mut sel_w)
3889 .arg(&ne)
3890 .arg(&nu)
3891 .arg(ex_scale);
3892 unsafe {
3893 b.launch(cfg)?;
3894 }
3895 Ok((sel_idx, sel_w))
3896 }
3897
3898 pub fn moe_router_topk_host(
3906 &self,
3907 logits: &CudaSlice<f32>,
3908 t: usize,
3909 n_expert: usize,
3910 n_used: usize,
3911 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3912 let f = self.func("moe_router_topk_f32");
3913 let n = t * n_used;
3914 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
3915 let mut sel_w = self.alloc_uninit::<f32>(n)?;
3916 let cfg = LaunchConfig {
3917 grid_dim: (t as u32, 1, 1),
3918 block_dim: (n_expert as u32, 1, 1),
3919 shared_mem_bytes: 0,
3920 };
3921 let (ne, nu) = (n_expert as i32, n_used as i32);
3922 let __s_b = self.gpu.stream();
3923 let mut b = __s_b.launch_builder(&f);
3924 b.arg(logits)
3925 .arg(&mut sel_idx)
3926 .arg(&mut sel_w)
3927 .arg(&ne)
3928 .arg(&nu);
3929 unsafe {
3930 b.launch(cfg)?;
3931 }
3932 let bytes = n * 8;
3934 let mut guard = self.router_stage.lock().unwrap();
3935 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
3936 *guard = Some(PinnedStage::new(bytes.max(4096))?);
3937 }
3938 let stage = guard.as_mut().unwrap();
3939 let (si, sw) = unsafe {
3940 (
3941 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
3942 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
3943 )
3944 };
3945 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?; self.gpu.stream().memcpy_dtoh(&sel_w, sw)?; self.gpu.stream().synchronize()?; Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
3949 }
3950
3951 #[allow(clippy::too_many_arguments)]
3955 pub fn moe_router_sigmoid_topk(
3956 &self,
3957 logits: &CudaSlice<f32>,
3958 t: usize,
3959 n_expert: usize,
3960 n_used: usize,
3961 active_count: usize,
3962 correction_bias: &CudaSlice<f32>,
3963 active: &CudaSlice<u8>,
3964 scaling_factor: f32,
3965 route_norm: bool,
3966 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3967 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
3968 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
3969 return Err(format!(
3970 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
3971 )
3972 .into());
3973 }
3974 if logits.len() < t * n_expert
3975 || correction_bias.len() != n_expert
3976 || active.len() != n_expert
3977 {
3978 return Err(format!(
3979 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
3980 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
3981 ).into());
3982 }
3983 let f = self.func("moe_router_sigmoid_topk_f32");
3984 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
3985 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
3986 let threads = n_expert.div_ceil(32) * 32;
3987 let cfg = LaunchConfig {
3988 grid_dim: (t as u32, 1, 1),
3989 block_dim: (threads as u32, 1, 1),
3990 shared_mem_bytes: 0,
3991 };
3992 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
3993 let __s_b = self.gpu.stream();
3994 let mut b = __s_b.launch_builder(&f);
3995 b.arg(logits)
3996 .arg(correction_bias)
3997 .arg(active)
3998 .arg(&mut sel_idx)
3999 .arg(&mut sel_w)
4000 .arg(&ne)
4001 .arg(&nu)
4002 .arg(&scaling_factor)
4003 .arg(&rn);
4004 unsafe {
4005 b.launch(cfg)?;
4006 }
4007 Ok((sel_idx, sel_w))
4008 }
4009
4010 #[allow(clippy::too_many_arguments)]
4013 pub fn moe_router_sigmoid_topk_host(
4014 &self,
4015 logits: &CudaSlice<f32>,
4016 t: usize,
4017 n_expert: usize,
4018 n_used: usize,
4019 active_count: usize,
4020 correction_bias: &CudaSlice<f32>,
4021 active: &CudaSlice<u8>,
4022 scaling_factor: f32,
4023 route_norm: bool,
4024 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4025 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4026 logits,
4027 t,
4028 n_expert,
4029 n_used,
4030 active_count,
4031 correction_bias,
4032 active,
4033 scaling_factor,
4034 route_norm,
4035 )?;
4036 let n = t * n_used;
4037 let bytes = n * 8;
4038 let mut guard = self.router_stage.lock().unwrap();
4039 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4040 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4041 }
4042 let stage = guard.as_mut().unwrap();
4043 let (si, sw) = unsafe {
4044 (
4045 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4046 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4047 )
4048 };
4049 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4050 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4051 self.gpu.stream().synchronize()?;
4052 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4053 }
4054
4055 pub fn stage_expert_async(
4059 &self,
4060 host_bytes: &[u8],
4061 scratch: &mut CudaSlice<u8>,
4062 off: usize,
4063 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4064 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4065 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4066 Ok(self.copy_stream.record_event(None)?)
4067 }
4068
4069 pub fn compute_wait(
4071 &self,
4072 ev: &cudarc::driver::CudaEvent,
4073 ) -> Result<(), Box<dyn std::error::Error>> {
4074 self.gpu.stream().wait(ev)?;
4075 Ok(())
4076 }
4077
4078 pub fn qmatvec_view(
4083 &self,
4084 w: &CudaSlice<u8>,
4085 range: std::ops::Range<usize>,
4086 x: &cudarc::driver::CudaView<f32>,
4087 m: usize,
4088 in_f: usize,
4089 out_f: usize,
4090 qtype: i32,
4091 row_bytes: usize,
4092 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4093 let f = self.func("qmatvec_f32");
4094 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4097 grid_dim: (out_f as u32, m as u32, 1),
4098 block_dim: (256, 1, 1),
4099 shared_mem_bytes: 0,
4100 };
4101 let (inf, outf, mi, qt, rb) =
4102 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4103 let __s_b = self.gpu.stream();
4104 let mut b = __s_b.launch_builder(&f);
4105 b.arg(&wv)
4106 .arg(x)
4107 .arg(&mut y)
4108 .arg(&inf)
4109 .arg(&outf)
4110 .arg(&mi)
4111 .arg(&qt)
4112 .arg(&rb);
4113 unsafe {
4114 b.launch(cfg)?;
4115 }
4116 Ok(y)
4117 }
4118
4119 #[allow(clippy::too_many_arguments)]
4126 pub fn moe_gate_up_silu8_q8(
4130 &self,
4131 gp: WPtr8,
4132 up: WPtr8,
4133 aq: &CudaSlice<i8>,
4134 ad: &CudaSlice<f32>,
4135 in_f: usize,
4136 n_ff: usize,
4137 n_used: usize,
4138 qt_g: i32,
4139 qt_u: i32,
4140 rb_g: usize,
4141 rb_u: usize,
4142 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4143 let f = self.func("moe_gate_up_silu8_q8");
4144 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4145 let cfg = LaunchConfig {
4146 grid_dim: (n_ff as u32, n_used as u32, 1),
4147 block_dim: (32, 1, 1),
4148 shared_mem_bytes: 0,
4149 };
4150 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4151 let __s_b = self.gpu.stream();
4152 let mut b = __s_b.launch_builder(&f);
4153 b.arg(&gp)
4154 .arg(&up)
4155 .arg(aq)
4156 .arg(ad)
4157 .arg(&mut act)
4158 .arg(&inf)
4159 .arg(&nff)
4160 .arg(&qt_g)
4161 .arg(&qt_u)
4162 .arg(&rbg)
4163 .arg(&rbu);
4164 unsafe {
4165 b.launch(cfg)?;
4166 }
4167 Ok(act)
4168 }
4169
4170 #[allow(clippy::too_many_arguments)]
4171 pub fn moe_down8_fma_q8(
4172 &self,
4173 dp: WPtr8,
4174 w: F32x8,
4175 aq2: &CudaSlice<i8>,
4176 ad2: &CudaSlice<f32>,
4177 dst: &mut cudarc::driver::CudaViewMut<f32>,
4178 in_f: usize,
4179 out_f: usize,
4180 n_used: usize,
4181 qt: i32,
4182 rb: usize,
4183 ) -> Result<(), Box<dyn std::error::Error>> {
4184 let f = self.func("moe_down8_fma_q8");
4185 let cfg = LaunchConfig {
4186 grid_dim: (out_f as u32, 1, 1),
4187 block_dim: (32, 1, 1),
4188 shared_mem_bytes: 0,
4189 };
4190 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4191 let __s_b = self.gpu.stream();
4192 let mut b = __s_b.launch_builder(&f);
4193 b.arg(&dp)
4194 .arg(&w)
4195 .arg(aq2)
4196 .arg(ad2)
4197 .arg(dst)
4198 .arg(&inf)
4199 .arg(&outf)
4200 .arg(&nu)
4201 .arg(&qt)
4202 .arg(&rbi);
4203 unsafe {
4204 b.launch(cfg)?;
4205 }
4206 Ok(())
4207 }
4208
4209 pub fn qmatvec_expert_q8(
4211 &self,
4212 w: &CudaSlice<u8>,
4213 range: std::ops::Range<usize>,
4214 aq: &CudaSlice<i8>,
4215 ad: &CudaSlice<f32>,
4216 m: usize,
4217 in_f: usize,
4218 out_f: usize,
4219 qtype: i32,
4220 row_bytes: usize,
4221 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4222 let f = self.func("qmatvec_expert_q8");
4223 let wv = w.slice(range);
4224 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4225 const ROWS: u32 = 4; let cfg = LaunchConfig {
4227 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4228 block_dim: (32, ROWS, 1),
4229 shared_mem_bytes: 0,
4230 };
4231 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4232 let __s_b = self.gpu.stream();
4233 let mut b = __s_b.launch_builder(&f);
4234 b.arg(&wv)
4235 .arg(aq)
4236 .arg(ad)
4237 .arg(&mut y)
4238 .arg(&inf)
4239 .arg(&outf)
4240 .arg(&mi)
4241 .arg(&qtype)
4242 .arg(&rbi);
4243 unsafe {
4244 b.launch(cfg)?;
4245 }
4246 Ok(y)
4247 }
4248
4249 pub fn moe_gate_up_silu8(
4250 &self,
4251 gp: WPtr8,
4252 up: WPtr8,
4253 x: &cudarc::driver::CudaView<f32>,
4254 in_f: usize,
4255 n_ff: usize,
4256 n_used: usize,
4257 qt_g: i32,
4258 qt_u: i32,
4259 rb_g: usize,
4260 rb_u: usize,
4261 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4262 let f = self.func("moe_gate_up_silu8_f32");
4263 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
4265 grid_dim: (n_ff as u32, n_used as u32, 1),
4266 block_dim: (256, 1, 1),
4267 shared_mem_bytes: 0,
4268 };
4269 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4270 let __s_b = self.gpu.stream();
4271 let mut b = __s_b.launch_builder(&f);
4272 b.arg(&gp)
4273 .arg(&up)
4274 .arg(x)
4275 .arg(&mut act)
4276 .arg(&inf)
4277 .arg(&nff)
4278 .arg(&qt_g)
4279 .arg(&qt_u)
4280 .arg(&rbg)
4281 .arg(&rbu);
4282 unsafe {
4283 b.launch(cfg)?;
4284 }
4285 Ok(act)
4286 }
4287
4288 #[allow(clippy::too_many_arguments)]
4294 pub fn moe_down8_fma_into(
4295 &self,
4296 dp: WPtr8,
4297 w: F32x8,
4298 act: &CudaSlice<f32>,
4299 dst: &mut cudarc::driver::CudaViewMut<f32>,
4300 in_f: usize,
4301 out_f: usize,
4302 n_used: usize,
4303 qt: i32,
4304 rb: usize,
4305 ) -> Result<(), Box<dyn std::error::Error>> {
4306 let f = self.func("moe_down8_fma_f32");
4307 let cfg = LaunchConfig {
4308 grid_dim: (out_f as u32, 1, 1),
4309 block_dim: (256, 1, 1),
4310 shared_mem_bytes: 0,
4311 };
4312 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4313 let __s_b = self.gpu.stream();
4314 let mut b = __s_b.launch_builder(&f);
4315 b.arg(&dp)
4316 .arg(&w)
4317 .arg(act)
4318 .arg(dst)
4319 .arg(&inf)
4320 .arg(&outf)
4321 .arg(&nu)
4322 .arg(&qt)
4323 .arg(&rbv);
4324 unsafe {
4325 b.launch(cfg)?;
4326 }
4327 Ok(())
4328 }
4329
4330 #[allow(clippy::too_many_arguments)]
4335 #[allow(clippy::too_many_arguments)]
4350 #[allow(clippy::too_many_arguments)]
4352 pub fn moe_pairs_matvec_q8(
4353 &self,
4354 table: &CudaSlice<u64>,
4355 proj: i32,
4356 pair_tok: &CudaSlice<i32>,
4357 pair_ex: &CudaSlice<i32>,
4358 aq: &CudaSlice<i8>,
4359 ad: &CudaSlice<f32>,
4360 in_f: usize,
4361 out_f: usize,
4362 n_expert: usize,
4363 n_pairs: usize,
4364 qtype: i32,
4365 row_bytes: usize,
4366 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4367 let f = self.func("moe_pairs_matvec_q8");
4368 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4369 const ROWS: u32 = 4;
4370 let cfg = LaunchConfig {
4371 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4372 block_dim: (32, ROWS, 1),
4373 shared_mem_bytes: 0,
4374 };
4375 let (inf, outf, ne, np, rbi) = (
4376 in_f as i32,
4377 out_f as i32,
4378 n_expert as i32,
4379 n_pairs as i32,
4380 row_bytes as i64,
4381 );
4382 let __s_b = self.gpu.stream();
4383 let mut b = __s_b.launch_builder(&f);
4384 b.arg(table)
4385 .arg(&proj)
4386 .arg(pair_tok)
4387 .arg(pair_ex)
4388 .arg(aq)
4389 .arg(ad)
4390 .arg(&mut y)
4391 .arg(&inf)
4392 .arg(&outf)
4393 .arg(&ne)
4394 .arg(&np)
4395 .arg(&qtype)
4396 .arg(&rbi);
4397 unsafe {
4398 b.launch(cfg)?;
4399 }
4400 Ok(y)
4401 }
4402
4403 #[allow(clippy::too_many_arguments)]
4405 pub fn moe_pairs_matvec_q8_em(
4406 &self,
4407 table: &CudaSlice<u64>,
4408 proj: i32,
4409 ex_ids: &CudaSlice<i32>,
4410 ex_off: &CudaSlice<i32>,
4411 ex_pairs: &CudaSlice<i32>,
4412 pair_tok: &CudaSlice<i32>,
4413 aq: &CudaSlice<i8>,
4414 ad: &CudaSlice<f32>,
4415 in_f: usize,
4416 out_f: usize,
4417 n_expert: usize,
4418 n_active: usize,
4419 n_pairs: usize,
4420 qtype: i32,
4421 row_bytes: usize,
4422 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4423 let f = self.func("moe_pairs_matvec_q8_em");
4424 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4425 const ROWS: u32 = 4;
4426 let cfg = LaunchConfig {
4427 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4428 block_dim: (32, ROWS, 1),
4429 shared_mem_bytes: 0,
4430 };
4431 let (inf, outf, ne, na, rbi) = (
4432 in_f as i32,
4433 out_f as i32,
4434 n_expert as i32,
4435 n_active as i32,
4436 row_bytes as i64,
4437 );
4438 let __s_b = self.gpu.stream();
4439 let mut b = __s_b.launch_builder(&f);
4440 b.arg(table)
4441 .arg(&proj)
4442 .arg(ex_ids)
4443 .arg(ex_off)
4444 .arg(ex_pairs)
4445 .arg(pair_tok)
4446 .arg(aq)
4447 .arg(ad)
4448 .arg(&mut y)
4449 .arg(&inf)
4450 .arg(&outf)
4451 .arg(&ne)
4452 .arg(&na)
4453 .arg(&qtype)
4454 .arg(&rbi);
4455 unsafe {
4456 b.launch(cfg)?;
4457 }
4458 Ok(y)
4459 }
4460
4461 #[allow(clippy::too_many_arguments)]
4464 pub fn moe_pairs_matvec_q8_dec(
4465 &self,
4466 table: &CudaSlice<u64>,
4467 proj: i32,
4468 ex_ids: &CudaSlice<i32>,
4469 ex_off: &CudaSlice<i32>,
4470 ex_pairs: &CudaSlice<i32>,
4471 pair_tok: &CudaSlice<i32>,
4472 aq: &CudaSlice<i8>,
4473 ad: &CudaSlice<f32>,
4474 in_f: usize,
4475 out_f: usize,
4476 n_expert: usize,
4477 n_active: usize,
4478 n_pairs: usize,
4479 qtype: i32,
4480 row_bytes: usize,
4481 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4482 let f = self.func("moe_pairs_matvec_q8_dec");
4483 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4484 const ROWS: u32 = 4;
4485 let cfg = LaunchConfig {
4486 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4487 block_dim: (32, ROWS, 1),
4488 shared_mem_bytes: 0,
4489 };
4490 let (inf, outf, ne, na, rbi) = (
4491 in_f as i32,
4492 out_f as i32,
4493 n_expert as i32,
4494 n_active as i32,
4495 row_bytes as i64,
4496 );
4497 let __s_b = self.gpu.stream();
4498 let mut b = __s_b.launch_builder(&f);
4499 b.arg(table)
4500 .arg(&proj)
4501 .arg(ex_ids)
4502 .arg(ex_off)
4503 .arg(ex_pairs)
4504 .arg(pair_tok)
4505 .arg(aq)
4506 .arg(ad)
4507 .arg(&mut y)
4508 .arg(&inf)
4509 .arg(&outf)
4510 .arg(&ne)
4511 .arg(&na)
4512 .arg(&qtype)
4513 .arg(&rbi);
4514 unsafe {
4515 b.launch(cfg)?;
4516 }
4517 Ok(y)
4518 }
4519
4520 pub fn moe_pairs_gelu_mul(
4521 &self,
4522 gate: &CudaSlice<f32>,
4523 up: &CudaSlice<f32>,
4524 n: usize,
4525 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4526 let f = self.func("moe_pairs_gelu_mul");
4527 let mut act = self.alloc_uninit::<f32>(n)?;
4528 let cfg = LaunchConfig::for_num_elems(n as u32);
4529 let nl = n as i64;
4530 let __s_b = self.gpu.stream();
4531 let mut b = __s_b.launch_builder(&f);
4532 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4533 unsafe {
4534 b.launch(cfg)?;
4535 }
4536 Ok(act)
4537 }
4538
4539 pub fn moe_pairs_silu_mul(
4540 &self,
4541 gate: &CudaSlice<f32>,
4542 up: &CudaSlice<f32>,
4543 n: usize,
4544 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4545 let f = self.func("moe_pairs_silu_mul");
4546 let mut act = self.alloc_uninit::<f32>(n)?;
4547 let cfg = LaunchConfig::for_num_elems(n as u32);
4548 let nl = n as i64;
4549 let __s_b = self.gpu.stream();
4550 let mut b = __s_b.launch_builder(&f);
4551 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4552 unsafe {
4553 b.launch(cfg)?;
4554 }
4555 Ok(act)
4556 }
4557
4558 #[allow(clippy::too_many_arguments)]
4559 pub fn moe_pairs_scatter(
4560 &self,
4561 y_down: &CudaSlice<f32>,
4562 pair_w: &CudaSlice<f32>,
4563 tok_pair_off: &CudaSlice<i32>,
4564 tok_pair_ids: &CudaSlice<i32>,
4565 moe_out: &mut CudaSlice<f32>,
4566 t: usize,
4567 n_embd: usize,
4568 ) -> Result<(), Box<dyn std::error::Error>> {
4569 let f = self.func("moe_pairs_scatter");
4570 let cfg = LaunchConfig {
4571 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4572 block_dim: (256, 1, 1),
4573 shared_mem_bytes: 0,
4574 };
4575 let ne = n_embd as i32;
4576 let __s_b = self.gpu.stream();
4577 let mut b = __s_b.launch_builder(&f);
4578 b.arg(y_down)
4579 .arg(pair_w)
4580 .arg(tok_pair_off)
4581 .arg(tok_pair_ids)
4582 .arg(moe_out)
4583 .arg(&ne);
4584 unsafe {
4585 b.launch(cfg)?;
4586 }
4587 Ok(())
4588 }
4589
4590 #[allow(clippy::too_many_arguments)]
4594 pub fn moe_gate_up_gelu8_dev_q8(
4595 &self,
4596 table: &CudaSlice<u64>,
4597 sel: &cudarc::driver::CudaView<i32>,
4598 aq: &CudaSlice<i8>,
4599 ad: &CudaSlice<f32>,
4600 in_f: usize,
4601 n_ff: usize,
4602 n_used: usize,
4603 n_expert: usize,
4604 qt_g: i32,
4605 qt_u: i32,
4606 rb_g: usize,
4607 rb_u: usize,
4608 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4609 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4610 let (inf, nff, ne, rbg, rbu) = (
4611 in_f as i32,
4612 n_ff as i32,
4613 n_expert as i32,
4614 rb_g as i64,
4615 rb_u as i64,
4616 );
4617 let f = self.func("moe_gate_up_gelu8_dev_q8");
4618 let cfg = LaunchConfig {
4619 grid_dim: (n_ff as u32, n_used as u32, 1),
4620 block_dim: (32, 1, 1),
4621 shared_mem_bytes: 0,
4622 };
4623 let __s_b = self.gpu.stream();
4624 let mut b = __s_b.launch_builder(&f);
4625 b.arg(table)
4626 .arg(sel)
4627 .arg(aq)
4628 .arg(ad)
4629 .arg(&mut act)
4630 .arg(&inf)
4631 .arg(&nff)
4632 .arg(&ne)
4633 .arg(&qt_g)
4634 .arg(&qt_u)
4635 .arg(&rbg)
4636 .arg(&rbu);
4637 unsafe {
4638 b.launch(cfg)?;
4639 }
4640 Ok(act)
4641 }
4642
4643 #[allow(clippy::too_many_arguments)]
4645 pub fn moe_gate_up_gelu8_dev_q8_rows(
4646 &self,
4647 table: &CudaSlice<u64>,
4648 sel: &CudaSlice<i32>,
4649 aq: &CudaSlice<i8>,
4650 ad: &CudaSlice<f32>,
4651 t: usize,
4652 in_f: usize,
4653 n_ff: usize,
4654 n_used: usize,
4655 n_expert: usize,
4656 qt_g: i32,
4657 qt_u: i32,
4658 rb_g: usize,
4659 rb_u: usize,
4660 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4661 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
4662 let (inf, nff, ne, rbg, rbu, nu) = (
4663 in_f as i32,
4664 n_ff as i32,
4665 n_expert as i32,
4666 rb_g as i64,
4667 rb_u as i64,
4668 n_used as i32,
4669 );
4670 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
4671 let cfg = LaunchConfig {
4672 grid_dim: (n_ff as u32, n_used as u32, t as u32),
4673 block_dim: (32, 1, 1),
4674 shared_mem_bytes: 0,
4675 };
4676 let __s_b = self.gpu.stream();
4677 let mut b = __s_b.launch_builder(&f);
4678 b.arg(table)
4679 .arg(sel)
4680 .arg(aq)
4681 .arg(ad)
4682 .arg(&mut act)
4683 .arg(&inf)
4684 .arg(&nff)
4685 .arg(&ne)
4686 .arg(&qt_g)
4687 .arg(&qt_u)
4688 .arg(&rbg)
4689 .arg(&rbu)
4690 .arg(&nu);
4691 unsafe {
4692 b.launch(cfg)?;
4693 }
4694 Ok(act)
4695 }
4696
4697 #[allow(clippy::too_many_arguments)]
4699 pub fn moe_gate_up_gelu8_dev_q8_csr(
4700 &self,
4701 table: &CudaSlice<u64>,
4702 sel: &CudaSlice<i32>,
4703 aq: &CudaSlice<i8>,
4704 ad: &CudaSlice<f32>,
4705 n_pairs: usize,
4706 in_f: usize,
4707 n_ff: usize,
4708 n_used: usize,
4709 n_expert: usize,
4710 qt_g: i32,
4711 qt_u: i32,
4712 rb_g: usize,
4713 rb_u: usize,
4714 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4715 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
4716 let (inf, nff, ne, rbg, rbu, nu, npi) = (
4717 in_f as i32,
4718 n_ff as i32,
4719 n_expert as i32,
4720 rb_g as i64,
4721 rb_u as i64,
4722 n_used as i32,
4723 n_pairs as i32,
4724 );
4725 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
4726 let cfg = LaunchConfig {
4727 grid_dim: (n_ff as u32, n_pairs as u32, 1),
4728 block_dim: (32, 1, 1),
4729 shared_mem_bytes: 0,
4730 };
4731 let __s_b = self.gpu.stream();
4732 let mut b = __s_b.launch_builder(&f);
4733 b.arg(table)
4734 .arg(sel)
4735 .arg(aq)
4736 .arg(ad)
4737 .arg(&mut act)
4738 .arg(&inf)
4739 .arg(&nff)
4740 .arg(&ne)
4741 .arg(&qt_g)
4742 .arg(&qt_u)
4743 .arg(&rbg)
4744 .arg(&rbu)
4745 .arg(&nu)
4746 .arg(&npi);
4747 unsafe {
4748 b.launch(cfg)?;
4749 }
4750 Ok(act)
4751 }
4752
4753 #[allow(clippy::too_many_arguments)]
4755 pub fn moe_down8_fma_dev_q8_rows_g(
4756 &self,
4757 table: &CudaSlice<u64>,
4758 sel: &CudaSlice<i32>,
4759 w: &CudaSlice<f32>,
4760 aq2: &CudaSlice<i8>,
4761 ad2: &CudaSlice<f32>,
4762 dst: &mut CudaSlice<f32>,
4763 t: usize,
4764 in_f: usize,
4765 out_f: usize,
4766 n_used: usize,
4767 n_expert: usize,
4768 qt: i32,
4769 rb: usize,
4770 ) -> Result<(), Box<dyn std::error::Error>> {
4771 let (inf, outf, nu, ne, rbi) = (
4772 in_f as i32,
4773 out_f as i32,
4774 n_used as i32,
4775 n_expert as i32,
4776 rb as i64,
4777 );
4778 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
4782 let f = self.func(if step_b1_w8 {
4783 "moe_down8_fma_dev_q8_rows_w8"
4784 } else {
4785 "moe_down8_fma_dev_q8_rows_g"
4786 });
4787 let cfg = LaunchConfig {
4788 grid_dim: (out_f as u32, 1, t as u32),
4789 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
4790 shared_mem_bytes: 0,
4791 };
4792 let __s_b = self.gpu.stream();
4793 let mut b = __s_b.launch_builder(&f);
4794 b.arg(table)
4795 .arg(sel)
4796 .arg(w)
4797 .arg(aq2)
4798 .arg(ad2)
4799 .arg(dst)
4800 .arg(&inf)
4801 .arg(&outf)
4802 .arg(&nu)
4803 .arg(&ne)
4804 .arg(&qt)
4805 .arg(&rbi);
4806 unsafe {
4807 b.launch(cfg)?;
4808 }
4809 Ok(())
4810 }
4811
4812 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
4816 let (out_f, in_f) = (2048usize, 2816usize);
4817 let nblk = in_f / 32;
4818 let mut seed = 0x9E3779B97F4A7C15u64;
4819 let mut rng = move || {
4820 seed = seed
4821 .wrapping_mul(6364136223846793005)
4822 .wrapping_add(1442695040888963407);
4823 (seed >> 33) as u8
4824 };
4825 let mut w = vec![0u8; out_f * nblk * 18];
4826 for b in w.iter_mut() {
4827 *b = rng();
4828 }
4829 for r in 0..out_f {
4830 for g in 0..nblk {
4831 let off = (r * nblk + g) * 18;
4832 w[off] = 0x00;
4833 w[off + 1] = 0x2C; }
4835 }
4836 let qplane = out_f * nblk * 16;
4837 let mut wrp = vec![0u8; w.len()];
4838 for r in 0..out_f {
4839 for g in 0..nblk {
4840 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
4841 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
4842 .copy_from_slice(&src[0..2]);
4843 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
4844 }
4845 }
4846 let w_d = self.htod_bytes(&w)?;
4847 let wrp_d = self.htod_bytes(&wrp)?;
4848 let mut aq = vec![0i8; m * in_f];
4849 for v in aq.iter_mut() {
4850 *v = rng() as i8;
4851 }
4852 let aq_d = self.htod_i8(&aq)?;
4853 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
4854 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
4855 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
4856 const RPB: u32 = 4;
4857 let cfg = LaunchConfig {
4858 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
4859 block_dim: (32, RPB, 1),
4860 shared_mem_bytes: 0,
4861 };
4862 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
4863 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
4864 let fb = self.func("qmatvec_q4_0_mmvq_b4");
4865 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
4866 {
4867 let __s_b = self.gpu.stream();
4868 let mut b = __s_b.launch_builder(&fb);
4869 b.arg(&w_d)
4870 .arg(&aq_d)
4871 .arg(&ad_d)
4872 .arg(&mut y0)
4873 .arg(&inf)
4874 .arg(&outf)
4875 .arg(&mi)
4876 .arg(&rb);
4877 unsafe {
4878 b.launch(cfg)?;
4879 }
4880 let __s_b = self.gpu.stream();
4881 let mut b = __s_b.launch_builder(&fr);
4882 b.arg(&wrp_d)
4883 .arg(&aq_d)
4884 .arg(&ad_d)
4885 .arg(&mut y1)
4886 .arg(&inf)
4887 .arg(&outf)
4888 .arg(&mi)
4889 .arg(&qp);
4890 unsafe {
4891 b.launch(cfg)?;
4892 }
4893 }
4894 self.gpu.stream().synchronize()?;
4895 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
4896 let nd = h0
4897 .iter()
4898 .zip(&h1)
4899 .filter(|(a, b)| a.to_bits() != b.to_bits())
4900 .count();
4901 if nd != 0 {
4902 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
4903 }
4904 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
4905 self.gpu.stream().synchronize()?;
4906 let t0 = std::time::Instant::now();
4907 for _ in 0..500 {
4908 if rp {
4909 let __s_b = self.gpu.stream();
4910 let mut b = __s_b.launch_builder(&fr);
4911 b.arg(&wrp_d)
4912 .arg(&aq_d)
4913 .arg(&ad_d)
4914 .arg(&mut y1)
4915 .arg(&inf)
4916 .arg(&outf)
4917 .arg(&mi)
4918 .arg(&qp);
4919 unsafe {
4920 b.launch(cfg)?;
4921 }
4922 } else {
4923 let __s_b = self.gpu.stream();
4924 let mut b = __s_b.launch_builder(&fb);
4925 b.arg(&w_d)
4926 .arg(&aq_d)
4927 .arg(&ad_d)
4928 .arg(&mut y0)
4929 .arg(&inf)
4930 .arg(&outf)
4931 .arg(&mi)
4932 .arg(&rb);
4933 unsafe {
4934 b.launch(cfg)?;
4935 }
4936 }
4937 }
4938 self.gpu.stream().synchronize()?;
4939 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
4940 };
4941 let _ = time(false)?;
4942 let _ = time(true)?; Ok((time(false)?, time(true)?))
4944 }
4945
4946 pub fn build_q4_rp4(
4951 &self,
4952 t: &mut crate::model::GpuTensor,
4953 ) -> Result<(), Box<dyn std::error::Error>> {
4954 use crate::model::GpuTensor;
4955 let GpuTensor::Quant {
4956 bytes,
4957 qtype,
4958 row_bytes,
4959 ne,
4960 rp4,
4961 ..
4962 } = t
4963 else {
4964 return Ok(());
4965 };
4966 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
4967 return Ok(());
4968 }
4969 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
4970 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
4971 return Ok(());
4972 }
4973 let nblk = in_f / 32;
4974 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
4975 let f = self.func("q4_0_split_rp_build");
4976 let n = (out_f * nblk) as i32;
4977 let cfg = LaunchConfig {
4978 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
4979 block_dim: (256, 1, 1),
4980 shared_mem_bytes: 0,
4981 };
4982 let (of, nb) = (out_f as i32, nblk as i32);
4983 let _ = n;
4984 let __s_b = self.gpu.stream();
4985 let mut b = __s_b.launch_builder(&f);
4986 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
4987 unsafe {
4988 b.launch(cfg)?;
4989 }
4990 *rp4 = Some(dst);
4991 Ok(())
4992 }
4993
4994 pub fn build_q8_rp4(
4999 &self,
5000 t: &mut crate::model::GpuTensor,
5001 ) -> Result<(), Box<dyn std::error::Error>> {
5002 use crate::model::GpuTensor;
5003 let GpuTensor::Quant {
5004 bytes,
5005 qtype,
5006 row_bytes,
5007 ne,
5008 rp4,
5009 ..
5010 } = t
5011 else {
5012 return Ok(());
5013 };
5014 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5015 return Ok(());
5016 }
5017 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5018 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5019 return Ok(());
5020 }
5021 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5022 Ok(())
5023 }
5024
5025 pub fn build_q8_rp4_raw(
5028 &self,
5029 bytes: &CudaSlice<u8>,
5030 in_f: usize,
5031 out_f: usize,
5032 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5033 assert!(in_f % 32 == 0);
5034 let nblk = in_f / 32;
5035 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5036 let f = self.func("q8_0_split_rp_build");
5037 let cfg = LaunchConfig {
5038 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5039 block_dim: (256, 1, 1),
5040 shared_mem_bytes: 0,
5041 };
5042 let (of, nb) = (out_f as i32, nblk as i32);
5043 let __s_b = self.gpu.stream();
5044 let mut b = __s_b.launch_builder(&f);
5045 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5046 unsafe {
5047 b.launch(cfg)?;
5048 }
5049 Ok(dst)
5050 }
5051
5052 pub fn build_q4k_rp4(
5060 &self,
5061 t: &mut crate::model::GpuTensor,
5062 ) -> Result<(), Box<dyn std::error::Error>> {
5063 use crate::model::GpuTensor;
5064 let GpuTensor::Quant {
5065 bytes,
5066 qtype,
5067 row_bytes,
5068 ne,
5069 rp4,
5070 ..
5071 } = t
5072 else {
5073 return Ok(());
5074 };
5075 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5076 return Ok(());
5077 }
5078 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5079 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5080 return Ok(());
5081 }
5082 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5083 Ok(())
5084 }
5085
5086 pub fn build_q6k_rp4(
5087 &self,
5088 t: &mut crate::model::GpuTensor,
5089 ) -> Result<(), Box<dyn std::error::Error>> {
5090 use crate::model::GpuTensor;
5091 let GpuTensor::Quant {
5092 bytes,
5093 qtype,
5094 row_bytes,
5095 ne,
5096 rp4,
5097 ..
5098 } = t
5099 else {
5100 return Ok(());
5101 };
5102 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5103 return Ok(());
5104 }
5105 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5106 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5107 return Ok(());
5108 }
5109 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5110 Ok(())
5111 }
5112
5113 pub fn build_kq_rp4_raw(
5115 &self,
5116 bytes: &CudaSlice<u8>,
5117 in_f: usize,
5118 out_f: usize,
5119 qtype: i32,
5120 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5121 assert!(in_f % 256 == 0);
5122 let nsbk = in_f / 256;
5123 let (sb_bytes, kname) = match qtype {
5124 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5125 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5126 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5127 };
5128 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5129 let f = self.func(kname);
5130 let cfg = LaunchConfig {
5131 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5132 block_dim: (256, 1, 1),
5133 shared_mem_bytes: 0,
5134 };
5135 let (of, nb) = (out_f as i32, nsbk as i32);
5136 let __s_b = self.gpu.stream();
5137 let mut b = __s_b.launch_builder(&f);
5138 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5139 unsafe {
5140 b.launch(cfg)?;
5141 }
5142 Ok(dst)
5143 }
5144
5145 pub fn kqrp_enabled() -> bool {
5149 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5150 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5151 Ok("0") => false,
5152 Ok(_) => true,
5153 Err(_) => cfg!(memra_hopper_mma),
5154 })
5155 }
5156
5157 pub fn build_q4_rp_swap(
5163 &self,
5164 t: &mut crate::model::GpuTensor,
5165 ) -> Result<bool, Box<dyn std::error::Error>> {
5166 use crate::model::GpuTensor;
5167 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5177 return Ok(false);
5178 }
5179 self.build_q4_rp4(t)?;
5180 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5182 return Ok(false);
5183 };
5184 match rp4.take() {
5185 Some(split) => {
5186 *bytes = split; *rp = true;
5188 Ok(true)
5189 }
5190 None => Ok(false),
5191 }
5192 }
5193
5194 pub fn q4rp_enabled() -> bool {
5196 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5197 *ON.get_or_init(|| {
5198 std::env::var("MEMRA_Q4RP")
5199 .map(|v| v != "0")
5200 .unwrap_or(true)
5201 })
5202 }
5203
5204 pub fn copy_rows_strided(
5207 &self,
5208 src: &CudaSlice<f32>,
5209 dst: &mut CudaSlice<f32>,
5210 row_elems: usize,
5211 n_rows: usize,
5212 src_stride: usize,
5213 src_off: usize,
5214 ) -> Result<(), Box<dyn std::error::Error>> {
5215 let f = self.func("copy_rows_strided_f32");
5216 let cfg = LaunchConfig {
5217 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5218 block_dim: (256, 1, 1),
5219 shared_mem_bytes: 0,
5220 };
5221 let (re, nr) = (row_elems as i32, n_rows as i32);
5222 let (st, off) = (src_stride as i64, src_off as i64);
5223 let __s_b = self.gpu.stream();
5224 let mut b = __s_b.launch_builder(&f);
5225 b.arg(src)
5226 .arg(&mut *dst)
5227 .arg(&re)
5228 .arg(&nr)
5229 .arg(&st)
5230 .arg(&off);
5231 unsafe {
5232 b.launch(cfg)?;
5233 }
5234 Ok(())
5235 }
5236
5237 pub fn u32_set_k(
5239 &self,
5240 dst: &mut CudaSlice<u32>,
5241 v: u32,
5242 idx: usize,
5243 ) -> Result<(), Box<dyn std::error::Error>> {
5244 let f = self.func("u32_set_k");
5245 let cfg = LaunchConfig {
5246 grid_dim: (1, 1, 1),
5247 block_dim: (1, 1, 1),
5248 shared_mem_bytes: 0,
5249 };
5250 let ii = idx as i32;
5251 let __s_b = self.gpu.stream();
5252 let mut b = __s_b.launch_builder(&f);
5253 b.arg(dst).arg(&v).arg(&ii);
5254 unsafe {
5255 b.launch(cfg)?;
5256 }
5257 Ok(())
5258 }
5259
5260 pub fn i32_add_k(
5262 &self,
5263 d: &mut CudaSlice<i32>,
5264 v: i32,
5265 ) -> Result<(), Box<dyn std::error::Error>> {
5266 let f = self.func("i32_add_k");
5267 let cfg = LaunchConfig {
5268 grid_dim: (1, 1, 1),
5269 block_dim: (32, 1, 1),
5270 shared_mem_bytes: 0,
5271 };
5272 let __s_b = self.gpu.stream();
5273 let mut b = __s_b.launch_builder(&f);
5274 b.arg(d).arg(&v);
5275 unsafe {
5276 b.launch(cfg)?;
5277 }
5278 Ok(())
5279 }
5280
5281 pub fn i32_iota_from(
5283 &self,
5284 ctr: &CudaSlice<i32>,
5285 dst: &mut CudaSlice<i32>,
5286 n: usize,
5287 ) -> Result<(), Box<dyn std::error::Error>> {
5288 let f = self.func("i32_iota_from");
5289 let cfg = LaunchConfig::for_num_elems(n as u32);
5290 let ni = n as i32;
5291 let __s_b = self.gpu.stream();
5292 let mut b = __s_b.launch_builder(&f);
5293 b.arg(ctr).arg(dst).arg(&ni);
5294 unsafe {
5295 b.launch(cfg)?;
5296 }
5297 Ok(())
5298 }
5299
5300 pub fn u32_map_k(
5302 &self,
5303 buf: &mut CudaSlice<u32>,
5304 map: &CudaSlice<u32>,
5305 idx: usize,
5306 ) -> Result<(), Box<dyn std::error::Error>> {
5307 let f = self.func("u32_map_k");
5308 let cfg = LaunchConfig {
5309 grid_dim: (1, 1, 1),
5310 block_dim: (1, 1, 1),
5311 shared_mem_bytes: 0,
5312 };
5313 let ii = idx as i32;
5314 let __s_b = self.gpu.stream();
5315 let mut b = __s_b.launch_builder(&f);
5316 b.arg(buf).arg(map).arg(&ii);
5317 unsafe {
5318 b.launch(cfg)?;
5319 }
5320 Ok(())
5321 }
5322
5323 #[allow(clippy::too_many_arguments)]
5325 pub fn u32_pack2(
5326 &self,
5327 a: &CudaSlice<u32>,
5328 off_a: usize,
5329 n1: usize,
5330 b_in: &CudaSlice<u32>,
5331 n2: usize,
5332 out: &mut CudaSlice<u32>,
5333 ) -> Result<(), Box<dyn std::error::Error>> {
5334 let f = self.func("u32_pack2");
5335 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5336 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5337 let __s_b = self.gpu.stream();
5338 let mut b = __s_b.launch_builder(&f);
5339 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5340 unsafe {
5341 b.launch(cfg)?;
5342 }
5343 Ok(())
5344 }
5345
5346 pub fn moe_w_exscale(
5348 &self,
5349 w: &mut CudaSlice<f32>,
5350 sel: &CudaSlice<i32>,
5351 s: &CudaSlice<f32>,
5352 n: usize,
5353 ) -> Result<(), Box<dyn std::error::Error>> {
5354 let f = self.func("moe_w_exscale");
5355 let cfg = LaunchConfig::for_num_elems(n as u32);
5356 let ni = n as i32;
5357 let __s_b = self.gpu.stream();
5358 let mut b = __s_b.launch_builder(&f);
5359 b.arg(w).arg(sel).arg(s).arg(&ni);
5360 unsafe {
5361 b.launch(cfg)?;
5362 }
5363 Ok(())
5364 }
5365
5366 pub fn moe_w_scale_by_expert(
5369 &self,
5370 w: &mut CudaSlice<f32>,
5371 sel: &CudaSlice<i32>,
5372 macros: &CudaSlice<f32>,
5373 n_expert: usize,
5374 n: usize,
5375 ) -> Result<(), Box<dyn std::error::Error>> {
5376 let f = self.func("moe_w_scale_by_expert");
5377 let cfg = LaunchConfig {
5378 grid_dim: (n.div_ceil(64) as u32, 1, 1),
5379 block_dim: (64, 1, 1),
5380 shared_mem_bytes: 0,
5381 };
5382 let (ne, nn) = (n_expert as i32, n as i32);
5383 let __s_b = self.gpu.stream();
5384 let mut b = __s_b.launch_builder(&f);
5385 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5386 unsafe {
5387 b.launch(cfg)?;
5388 }
5389 Ok(())
5390 }
5391
5392 pub fn moe_gate_up_silu8_dev_q8(
5393 &self,
5394 table: &CudaSlice<u64>,
5395 sel: &cudarc::driver::CudaView<i32>,
5396 aq: &CudaSlice<i8>,
5397 ad: &CudaSlice<f32>,
5398 in_f: usize,
5399 n_ff: usize,
5400 n_used: usize,
5401 n_expert: usize,
5402 qt_g: i32,
5403 qt_u: i32,
5404 rb_g: usize,
5405 rb_u: usize,
5406 macros: &CudaSlice<f32>,
5407 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5408 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5409 let (mode, wpb) = GU.get_or_init(|| {
5410 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5411 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5412 .ok()
5413 .and_then(|v| v.parse().ok())
5414 .unwrap_or(4u32)
5415 .clamp(1, 16);
5416 (mode, wpb)
5417 });
5418 let (mode, wpb) = (mode.as_str(), *wpb);
5419 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5420 let (inf, nff, ne, rbg, rbu) = (
5421 in_f as i32,
5422 n_ff as i32,
5423 n_expert as i32,
5424 rb_g as i64,
5425 rb_u as i64,
5426 );
5427 let (f, cfg) = match mode {
5428 "1" | "2" | "4" => {
5429 let rpw: u32 = mode.parse().unwrap();
5430 let f = self.func(match rpw {
5431 1 => "moe_gate_up_silu8_dev_q8_r1",
5432 2 => "moe_gate_up_silu8_dev_q8_r2",
5433 _ => "moe_gate_up_silu8_dev_q8_r4",
5434 });
5435 let rows_per_block = (rpw * wpb) as usize;
5436 let gx = n_ff.div_ceil(rows_per_block) as u32;
5437 (
5438 f,
5439 LaunchConfig {
5440 grid_dim: (gx, n_used as u32, 1),
5441 block_dim: (32, wpb, 1),
5442 shared_mem_bytes: 0,
5443 },
5444 )
5445 }
5446 "j8" if n_used <= 32 => (
5447 self.func("moe_gate_up_silu8_dev_q8_j8"),
5448 LaunchConfig {
5449 grid_dim: (n_ff as u32, 1, 1),
5450 block_dim: (32, n_used as u32, 1),
5451 shared_mem_bytes: 0,
5452 },
5453 ),
5454 "vsm2" => {
5456 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5457 let sh = (rb_g + rb_u) as u32;
5458 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5459 f.set_attribute(
5460 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5461 sh as i32,
5462 )?;
5463 (
5464 f,
5465 LaunchConfig {
5466 grid_dim: (n_ff as u32, n_used as u32, 1),
5467 block_dim: (32, 1, 1),
5468 shared_mem_bytes: sh,
5469 },
5470 )
5471 }
5472 "vsm" => {
5473 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5474 let sh = (rb_g + rb_u) as u32;
5475 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5476 f.set_attribute(
5477 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5478 sh as i32,
5479 )?;
5480 (
5481 f,
5482 LaunchConfig {
5483 grid_dim: (n_ff as u32, n_used as u32, 1),
5484 block_dim: (32, 1, 1),
5485 shared_mem_bytes: sh,
5486 },
5487 )
5488 }
5489 "sg" => (
5490 self.func("moe_gate_up_silu8_dev_q8_sg"),
5491 LaunchConfig {
5492 grid_dim: (n_ff as u32, n_used as u32, 1),
5493 block_dim: (32, 1, 1),
5494 shared_mem_bytes: 0,
5495 },
5496 ),
5497 "j8sg" if n_used <= 32 => (
5498 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5499 LaunchConfig {
5500 grid_dim: (n_ff as u32, 1, 1),
5501 block_dim: (32, n_used as u32, 1),
5502 shared_mem_bytes: 0,
5503 },
5504 ),
5505 "u64" if in_f == 2048 => (
5506 self.func("moe_gate_up_silu8_dev_q8_u64"),
5507 LaunchConfig {
5508 grid_dim: (n_ff as u32, n_used as u32, 1),
5509 block_dim: (32, 1, 1),
5510 shared_mem_bytes: 0,
5511 },
5512 ),
5513 "gs4" if in_f == 2048 => (
5514 self.func("moe_gate_up_silu8_dev_q8_gs4"),
5515 LaunchConfig {
5516 grid_dim: (n_ff as u32, n_used as u32, 1),
5517 block_dim: (32, 4, 1),
5518 shared_mem_bytes: 0,
5519 },
5520 ),
5521 "v" | "" => (
5523 self.func("moe_gate_up_silu8_dev_q8_v"),
5524 LaunchConfig {
5525 grid_dim: (n_ff as u32, n_used as u32, 1),
5526 block_dim: (32, 1, 1),
5527 shared_mem_bytes: 0,
5528 },
5529 ),
5530 "s2" => (
5531 self.func("moe_gate_up_silu8_dev_q8_s2"),
5532 LaunchConfig {
5533 grid_dim: (n_ff as u32, n_used as u32, 1),
5534 block_dim: (32, 2, 1),
5535 shared_mem_bytes: 0,
5536 },
5537 ),
5538 "s2z" => {
5539 let rz = wpb.min(16); (
5541 self.func("moe_gate_up_silu8_dev_q8_s2z"),
5542 LaunchConfig {
5543 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5544 block_dim: (32, 2, rz),
5545 shared_mem_bytes: 0,
5546 },
5547 )
5548 }
5549 _ => (
5550 self.func("moe_gate_up_silu8_dev_q8"),
5551 LaunchConfig {
5552 grid_dim: (n_ff as u32, n_used as u32, 1),
5553 block_dim: (32, 1, 1),
5554 shared_mem_bytes: 0,
5555 },
5556 ),
5557 };
5558 let __s_b = self.gpu.stream();
5559 let mut b = __s_b.launch_builder(&f);
5560 b.arg(table)
5561 .arg(sel)
5562 .arg(aq)
5563 .arg(ad)
5564 .arg(&mut act)
5565 .arg(&inf)
5566 .arg(&nff)
5567 .arg(&ne)
5568 .arg(&qt_g)
5569 .arg(&qt_u)
5570 .arg(&rbg)
5571 .arg(&rbu)
5572 .arg(macros);
5573 unsafe {
5574 b.launch(cfg)?;
5575 }
5576 Ok(act)
5577 }
5578
5579 #[allow(clippy::too_many_arguments)]
5580 pub fn moe_down8_fma_dev_q8(
5581 &self,
5582 table: &CudaSlice<u64>,
5583 sel: &cudarc::driver::CudaView<i32>,
5584 w: &cudarc::driver::CudaView<f32>,
5585 aq2: &CudaSlice<i8>,
5586 ad2: &CudaSlice<f32>,
5587 dst: &mut cudarc::driver::CudaViewMut<f32>,
5588 in_f: usize,
5589 out_f: usize,
5590 n_used: usize,
5591 n_expert: usize,
5592 qt: i32,
5593 rb: usize,
5594 ) -> Result<(), Box<dyn std::error::Error>> {
5595 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5596 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5597 let (inf, outf, nu, ne, rbi) = (
5598 in_f as i32,
5599 out_f as i32,
5600 n_used as i32,
5601 n_expert as i32,
5602 rb as i64,
5603 );
5604 let (f, cfg) = match mode.as_str() {
5607 m @ ("1" | "2" | "4") if n_used <= 8 => {
5608 let rpw: usize = m.parse().unwrap();
5609 let f = self.func(match rpw {
5610 1 => "moe_down8_fma_dev_q8_w8r1",
5611 2 => "moe_down8_fma_dev_q8_w8r2",
5612 _ => "moe_down8_fma_dev_q8_w8r4",
5613 });
5614 (
5615 f,
5616 LaunchConfig {
5617 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5618 block_dim: (32, n_used as u32, 1),
5619 shared_mem_bytes: 0,
5620 },
5621 )
5622 }
5623 "h2" if in_f == 512 => (
5624 self.func("moe_down8_fma_dev_q8_h2"),
5625 LaunchConfig {
5626 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5627 block_dim: (32, 1, 1),
5628 shared_mem_bytes: 0,
5629 },
5630 ),
5631 "" if in_f == 704 && n_used <= 8 => (
5634 self.func("moe_down8_fma_dev_q8_w8r2"),
5635 LaunchConfig {
5636 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5637 block_dim: (32, n_used as u32, 1),
5638 shared_mem_bytes: 0,
5639 },
5640 ),
5641 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
5645 self.func("moe_down8_fma_dev_q8_w8h2v"),
5646 LaunchConfig {
5647 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5648 block_dim: (32, n_used as u32, 1),
5649 shared_mem_bytes: 0,
5650 },
5651 ),
5652 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
5653 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
5654 LaunchConfig {
5655 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5656 block_dim: (32, n_used as u32, 1),
5657 shared_mem_bytes: 0,
5658 },
5659 ),
5660 "w8h2r2" if in_f == 512 && n_used <= 8 => (
5661 self.func("moe_down8_fma_dev_q8_w8h2r2"),
5662 LaunchConfig {
5663 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5664 block_dim: (32, n_used as u32, 1),
5665 shared_mem_bytes: 0,
5666 },
5667 ),
5668 "w8h2" if in_f == 512 && n_used <= 8 => (
5669 self.func("moe_down8_fma_dev_q8_w8h2"),
5670 LaunchConfig {
5671 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5672 block_dim: (32, n_used as u32, 1),
5673 shared_mem_bytes: 0,
5674 },
5675 ),
5676 _ => (
5677 self.func("moe_down8_fma_dev_q8"),
5678 LaunchConfig {
5679 grid_dim: (out_f as u32, 1, 1),
5680 block_dim: (32, 1, 1),
5681 shared_mem_bytes: 0,
5682 },
5683 ),
5684 };
5685 let __s_b = self.gpu.stream();
5686 let mut b = __s_b.launch_builder(&f);
5687 b.arg(table)
5688 .arg(sel)
5689 .arg(w)
5690 .arg(aq2)
5691 .arg(ad2)
5692 .arg(dst)
5693 .arg(&inf)
5694 .arg(&outf)
5695 .arg(&nu)
5696 .arg(&ne)
5697 .arg(&qt)
5698 .arg(&rbi);
5699 unsafe {
5700 b.launch(cfg)?;
5701 }
5702 Ok(())
5703 }
5704
5705 #[allow(clippy::too_many_arguments)]
5712 pub fn moe_gate_up_silu8_dev_q8_rows(
5713 &self,
5714 table: &CudaSlice<u64>,
5715 sel: &CudaSlice<i32>,
5716 aq: &CudaSlice<i8>,
5717 ad: &CudaSlice<f32>,
5718 t: usize,
5719 in_f: usize,
5720 n_ff: usize,
5721 n_used: usize,
5722 n_expert: usize,
5723 qt_g: i32,
5724 qt_u: i32,
5725 rb_g: usize,
5726 rb_u: usize,
5727 macros: &CudaSlice<f32>,
5728 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5729 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
5730 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5731 let cfg = LaunchConfig {
5732 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5733 block_dim: (32, 1, 1),
5734 shared_mem_bytes: 0,
5735 };
5736 let (inf, nff, ne, nu, rbg, rbu) = (
5737 in_f as i32,
5738 n_ff as i32,
5739 n_expert as i32,
5740 n_used as i32,
5741 rb_g as i64,
5742 rb_u as i64,
5743 );
5744 let __s_b = self.gpu.stream();
5745 let mut b = __s_b.launch_builder(&f);
5746 b.arg(table)
5747 .arg(sel)
5748 .arg(aq)
5749 .arg(ad)
5750 .arg(&mut act)
5751 .arg(&inf)
5752 .arg(&nff)
5753 .arg(&ne)
5754 .arg(&qt_g)
5755 .arg(&qt_u)
5756 .arg(&rbg)
5757 .arg(&rbu)
5758 .arg(&nu)
5759 .arg(macros);
5760 unsafe {
5761 b.launch(cfg)?;
5762 }
5763 Ok(act)
5764 }
5765
5766 #[allow(clippy::too_many_arguments)]
5771 pub fn moe_down8_fma_dev_q8_rows(
5772 &self,
5773 table: &CudaSlice<u64>,
5774 sel: &CudaSlice<i32>,
5775 w: &CudaSlice<f32>,
5776 aq2: &CudaSlice<i8>,
5777 ad2: &CudaSlice<f32>,
5778 dst: &mut CudaSlice<f32>,
5779 t: usize,
5780 in_f: usize,
5781 out_f: usize,
5782 n_used: usize,
5783 n_expert: usize,
5784 qt: i32,
5785 rb: usize,
5786 ) -> Result<(), Box<dyn std::error::Error>> {
5787 assert!(
5788 in_f == 512 && n_used <= 8,
5789 "down rows twin is w8h2v shape-gated"
5790 );
5791 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
5792 let cfg = LaunchConfig {
5793 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
5794 block_dim: (32, n_used as u32, 1),
5795 shared_mem_bytes: 0,
5796 };
5797 let (inf, outf, nu, ne, rbi) = (
5798 in_f as i32,
5799 out_f as i32,
5800 n_used as i32,
5801 n_expert as i32,
5802 rb as i64,
5803 );
5804 let __s_b = self.gpu.stream();
5805 let mut b = __s_b.launch_builder(&f);
5806 b.arg(table)
5807 .arg(sel)
5808 .arg(w)
5809 .arg(aq2)
5810 .arg(ad2)
5811 .arg(dst)
5812 .arg(&inf)
5813 .arg(&outf)
5814 .arg(&nu)
5815 .arg(&ne)
5816 .arg(&qt)
5817 .arg(&rbi);
5818 unsafe {
5819 b.launch(cfg)?;
5820 }
5821 Ok(())
5822 }
5823
5824 #[allow(clippy::too_many_arguments)]
5828 pub fn moe_gate_up_silu8_dev_q8_csr(
5829 &self,
5830 table: &CudaSlice<u64>,
5831 sel: &CudaSlice<i32>,
5832 aq: &CudaSlice<i8>,
5833 ad: &CudaSlice<f32>,
5834 n_pairs: usize,
5835 in_f: usize,
5836 n_ff: usize,
5837 n_used: usize,
5838 n_expert: usize,
5839 qt_g: i32,
5840 qt_u: i32,
5841 rb_g: usize,
5842 rb_u: usize,
5843 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5844 let f = self.func("moe_gate_up_silu8_dev_q8_csr_iq4");
5845 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5846 let cfg = LaunchConfig {
5847 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5848 block_dim: (32, 1, 1),
5849 shared_mem_bytes: 0,
5850 };
5851 let (inf, nff, ne, nu, npi, rbg, rbu) = (
5852 in_f as i32,
5853 n_ff as i32,
5854 n_expert as i32,
5855 n_used as i32,
5856 n_pairs as i32,
5857 rb_g as i64,
5858 rb_u as i64,
5859 );
5860 let __s_b = self.gpu.stream();
5861 let mut b = __s_b.launch_builder(&f);
5862 b.arg(table)
5863 .arg(sel)
5864 .arg(aq)
5865 .arg(ad)
5866 .arg(&mut act)
5867 .arg(&inf)
5868 .arg(&nff)
5869 .arg(&ne)
5870 .arg(&qt_g)
5871 .arg(&qt_u)
5872 .arg(&rbg)
5873 .arg(&rbu)
5874 .arg(&nu)
5875 .arg(&npi);
5876 unsafe {
5877 b.launch(cfg)?;
5878 }
5879 Ok(act)
5880 }
5881
5882 #[allow(clippy::too_many_arguments)]
5886 pub fn moe_down8_fma_dev_q8_variant(
5887 &self,
5888 variant: &str,
5889 table: &CudaSlice<u64>,
5890 sel: &cudarc::driver::CudaView<i32>,
5891 w: &cudarc::driver::CudaView<f32>,
5892 aq2: &CudaSlice<i8>,
5893 ad2: &CudaSlice<f32>,
5894 dst: &mut cudarc::driver::CudaViewMut<f32>,
5895 in_f: usize,
5896 out_f: usize,
5897 n_used: usize,
5898 n_expert: usize,
5899 qt: i32,
5900 rb: usize,
5901 ) -> Result<(), Box<dyn std::error::Error>> {
5902 let (inf, outf, nu, ne, rbi) = (
5903 in_f as i32,
5904 out_f as i32,
5905 n_used as i32,
5906 n_expert as i32,
5907 rb as i64,
5908 );
5909 let (f, cfg) = match variant {
5910 "w8h2" | "w8h2v" => (
5911 self.func(if variant == "w8h2" {
5912 "moe_down8_fma_dev_q8_w8h2"
5913 } else {
5914 "moe_down8_fma_dev_q8_w8h2v"
5915 }),
5916 LaunchConfig {
5917 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5918 block_dim: (32, n_used as u32, 1),
5919 shared_mem_bytes: 0,
5920 },
5921 ),
5922 "w8h2r2" | "w8h2r2v" => (
5923 self.func(if variant == "w8h2r2" {
5924 "moe_down8_fma_dev_q8_w8h2r2"
5925 } else {
5926 "moe_down8_fma_dev_q8_w8h2r2v"
5927 }),
5928 LaunchConfig {
5929 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5930 block_dim: (32, n_used as u32, 1),
5931 shared_mem_bytes: 0,
5932 },
5933 ),
5934 _ => (
5935 self.func("moe_down8_fma_dev_q8"),
5936 LaunchConfig {
5937 grid_dim: (out_f as u32, 1, 1),
5938 block_dim: (32, 1, 1),
5939 shared_mem_bytes: 0,
5940 },
5941 ),
5942 };
5943 let __s_b = self.gpu.stream();
5944 let mut b = __s_b.launch_builder(&f);
5945 b.arg(table)
5946 .arg(sel)
5947 .arg(w)
5948 .arg(aq2)
5949 .arg(ad2)
5950 .arg(dst)
5951 .arg(&inf)
5952 .arg(&outf)
5953 .arg(&nu)
5954 .arg(&ne)
5955 .arg(&qt)
5956 .arg(&rbi);
5957 unsafe {
5958 b.launch(cfg)?;
5959 }
5960 Ok(())
5961 }
5962
5963 #[allow(clippy::too_many_arguments)]
5965 pub fn moe_gate_up_silu8_dev_q8_variant(
5966 &self,
5967 variant: &str,
5968 table: &CudaSlice<u64>,
5969 sel: &cudarc::driver::CudaView<i32>,
5970 aq: &CudaSlice<i8>,
5971 ad: &CudaSlice<f32>,
5972 in_f: usize,
5973 n_ff: usize,
5974 n_used: usize,
5975 n_expert: usize,
5976 qt_g: i32,
5977 qt_u: i32,
5978 rb_g: usize,
5979 rb_u: usize,
5980 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5981 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5982 let (inf, nff, ne, rbg, rbu) = (
5983 in_f as i32,
5984 n_ff as i32,
5985 n_expert as i32,
5986 rb_g as i64,
5987 rb_u as i64,
5988 );
5989 let f = self.func(if variant == "v" {
5990 "moe_gate_up_silu8_dev_q8_v"
5991 } else {
5992 "moe_gate_up_silu8_dev_q8"
5993 });
5994 let cfg = LaunchConfig {
5995 grid_dim: (n_ff as u32, n_used as u32, 1),
5996 block_dim: (32, 1, 1),
5997 shared_mem_bytes: 0,
5998 };
5999 let __s_b = self.gpu.stream();
6000 let mut b = __s_b.launch_builder(&f);
6001 b.arg(table)
6002 .arg(sel)
6003 .arg(aq)
6004 .arg(ad)
6005 .arg(&mut act)
6006 .arg(&inf)
6007 .arg(&nff)
6008 .arg(&ne)
6009 .arg(&qt_g)
6010 .arg(&qt_u)
6011 .arg(&rbg)
6012 .arg(&rbu);
6013 unsafe {
6014 b.launch(cfg)?;
6015 }
6016 Ok(act)
6017 }
6018
6019 pub fn moe_gate_up_silu8_dev(
6020 &self,
6021 table: &CudaSlice<u64>,
6022 sel: &cudarc::driver::CudaView<i32>,
6023 x: &cudarc::driver::CudaView<f32>,
6024 in_f: usize,
6025 n_ff: usize,
6026 n_used: usize,
6027 n_expert: usize,
6028 qt_g: i32,
6029 qt_u: i32,
6030 rb_g: usize,
6031 rb_u: usize,
6032 macros: &CudaSlice<f32>,
6033 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6034 let f = self.func("moe_gate_up_silu8_dev");
6035 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6037 grid_dim: (n_ff as u32, n_used as u32, 1),
6038 block_dim: (256, 1, 1),
6039 shared_mem_bytes: 0,
6040 };
6041 let (inf, nff, ne, rbg, rbu) = (
6042 in_f as i32,
6043 n_ff as i32,
6044 n_expert as i32,
6045 rb_g as i64,
6046 rb_u as i64,
6047 );
6048 let __s_b = self.gpu.stream();
6049 let mut b = __s_b.launch_builder(&f);
6050 b.arg(table)
6051 .arg(sel)
6052 .arg(x)
6053 .arg(&mut act)
6054 .arg(&inf)
6055 .arg(&nff)
6056 .arg(&ne)
6057 .arg(&qt_g)
6058 .arg(&qt_u)
6059 .arg(&rbg)
6060 .arg(&rbu)
6061 .arg(macros);
6062 unsafe {
6063 b.launch(cfg)?;
6064 }
6065 Ok(act)
6066 }
6067
6068 #[allow(clippy::too_many_arguments)]
6071 pub fn moe_down8_fma_dev(
6072 &self,
6073 table: &CudaSlice<u64>,
6074 sel: &cudarc::driver::CudaView<i32>,
6075 w: &cudarc::driver::CudaView<f32>,
6076 act: &CudaSlice<f32>,
6077 dst: &mut cudarc::driver::CudaViewMut<f32>,
6078 in_f: usize,
6079 out_f: usize,
6080 n_used: usize,
6081 n_expert: usize,
6082 qt: i32,
6083 rb: usize,
6084 ) -> Result<(), Box<dyn std::error::Error>> {
6085 let f = self.func("moe_down8_fma_dev");
6086 let cfg = LaunchConfig {
6087 grid_dim: (out_f as u32, 1, 1),
6088 block_dim: (256, 1, 1),
6089 shared_mem_bytes: 0,
6090 };
6091 let (inf, outf, nu, ne, rbv) = (
6092 in_f as i32,
6093 out_f as i32,
6094 n_used as i32,
6095 n_expert as i32,
6096 rb as i64,
6097 );
6098 let __s_b = self.gpu.stream();
6099 let mut b = __s_b.launch_builder(&f);
6100 b.arg(table)
6101 .arg(sel)
6102 .arg(w)
6103 .arg(act)
6104 .arg(dst)
6105 .arg(&inf)
6106 .arg(&outf)
6107 .arg(&nu)
6108 .arg(&ne)
6109 .arg(&qt)
6110 .arg(&rbv);
6111 unsafe {
6112 b.launch(cfg)?;
6113 }
6114 Ok(())
6115 }
6116
6117 pub fn axpy_into(
6119 &self,
6120 src: &CudaSlice<f32>,
6121 alpha: f32,
6122 dst: &mut cudarc::driver::CudaViewMut<f32>,
6123 n: usize,
6124 ) -> Result<(), Box<dyn std::error::Error>> {
6125 let f = self.func("axpy_f32");
6126 let cfg = LaunchConfig::for_num_elems(n as u32);
6127 let (a, ni) = (alpha, n as i32);
6128 let __s_b = self.gpu.stream();
6129 let mut b = __s_b.launch_builder(&f);
6130 b.arg(src).arg(dst).arg(&a).arg(&ni);
6131 unsafe {
6132 b.launch(cfg)?;
6133 }
6134 Ok(())
6135 }
6136
6137 pub fn add_scaled_rows(
6139 &self,
6140 src: &CudaSlice<f32>,
6141 scale: &CudaSlice<f32>,
6142 dst: &mut CudaSlice<f32>,
6143 ncols: usize,
6144 nrows: usize,
6145 ) -> Result<(), Box<dyn std::error::Error>> {
6146 let f = self.func("add_scaled_rows_f32");
6147 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6148 let (nc, nr) = (ncols as i32, nrows as i32);
6149 let __s_b = self.gpu.stream();
6150 let mut b = __s_b.launch_builder(&f);
6151 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6152 unsafe {
6153 b.launch(cfg)?;
6154 }
6155 Ok(())
6156 }
6157
6158 pub fn gather_rows(
6162 &self,
6163 src: &CudaSlice<f32>,
6164 idx: &CudaSlice<i32>,
6165 dst: &mut CudaSlice<f32>,
6166 ncols: usize,
6167 m_e: usize,
6168 ) -> Result<(), Box<dyn std::error::Error>> {
6169 let f = self.func("gather_rows_f32");
6170 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6171 let (nc, me) = (ncols as i32, m_e as i32);
6172 let __s_b = self.gpu.stream();
6173 let mut b = __s_b.launch_builder(&f);
6174 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6175 unsafe {
6176 b.launch(cfg)?;
6177 }
6178 Ok(())
6179 }
6180
6181 pub fn scatter_slot(
6186 &self,
6187 src: &CudaSlice<f32>,
6188 tok_idx: &CudaSlice<i32>,
6189 slot_idx: &CudaSlice<i32>,
6190 weight: &CudaSlice<f32>,
6191 dst: &mut CudaSlice<f32>,
6192 wbuf: &mut CudaSlice<f32>,
6193 ncols: usize,
6194 n_used: usize,
6195 m_e: usize,
6196 ) -> Result<(), Box<dyn std::error::Error>> {
6197 let f = self.func("scatter_add_slot_f32");
6198 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6199 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6200 let __s_b = self.gpu.stream();
6201 let mut b = __s_b.launch_builder(&f);
6202 b.arg(src)
6203 .arg(tok_idx)
6204 .arg(slot_idx)
6205 .arg(weight)
6206 .arg(dst)
6207 .arg(wbuf)
6208 .arg(&nc)
6209 .arg(&nu)
6210 .arg(&me);
6211 unsafe {
6212 b.launch(cfg)?;
6213 }
6214 Ok(())
6215 }
6216
6217 pub fn reduce_slots(
6221 &self,
6222 slots: &CudaSlice<f32>,
6223 wbuf: &CudaSlice<f32>,
6224 dst: &mut CudaSlice<f32>,
6225 ncols: usize,
6226 n_used: usize,
6227 t: usize,
6228 ) -> Result<(), Box<dyn std::error::Error>> {
6229 let f = self.func("reduce_slots_f32");
6230 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6231 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6232 let __s_b = self.gpu.stream();
6233 let mut b = __s_b.launch_builder(&f);
6234 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6235 unsafe {
6236 b.launch(cfg)?;
6237 }
6238 Ok(())
6239 }
6240
6241 pub fn quantize_q8_1_view(
6248 &self,
6249 x: &cudarc::driver::CudaView<f32>,
6250 m: usize,
6251 in_f: usize,
6252 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6253 let f = self.func("quantize_q8_1");
6254 let nblk = in_f / 32;
6255 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6256 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6257 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6258 let (inf, mi) = (in_f as i32, m as i32);
6259 let __s_b = self.gpu.stream();
6260 let mut b = __s_b.launch_builder(&f);
6261 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6262 unsafe {
6263 b.launch(cfg)?;
6264 }
6265 Ok((q, d))
6266 }
6267
6268 pub fn quantize_q8_1(
6269 &self,
6270 x: &CudaSlice<f32>,
6271 m: usize,
6272 in_f: usize,
6273 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6274 let nblk = in_f / 32;
6275 let mut q = self.alloc_uninit::<i8>(m * in_f)?; let mut d = self.alloc_uninit::<f32>(m * nblk)?; let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6279 let (inf, mi) = (in_f as i32, m as i32);
6280 if Self::pdl_on() && Self::pdl_wb_on() {
6281 {
6282 use cudarc::driver::{DevicePtr, DevicePtrMut};
6283 let s = &self.gpu.stream();
6284 let (px, _g0) = x.device_ptr(s);
6285 let (pq, _g1) = q.device_ptr_mut(s);
6286 let (pd, _g2) = d.device_ptr_mut(s);
6287 let mut ps = [
6288 &px as *const _ as *mut std::ffi::c_void,
6289 &pq as *const _ as *mut _,
6290 &pd as *const _ as *mut _,
6291 &inf as *const _ as *mut _,
6292 &mi as *const _ as *mut _,
6293 ];
6294 unsafe {
6295 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6296 }
6297 }
6298 return Ok((q, d));
6299 }
6300 let f = self.func("quantize_q8_1");
6301 let __s_b = self.gpu.stream();
6302 let mut b = __s_b.launch_builder(&f);
6303 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6304 unsafe {
6305 b.launch(cfg)?;
6306 }
6307 Ok((q, d))
6308 }
6309
6310 pub fn quantize_fp4_act(
6314 &self,
6315 x: &CudaSlice<f32>,
6316 m: usize,
6317 in_f: usize,
6318 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6319 let f = self.func("quantize_fp4_act");
6320 let nb16 = in_f / 16;
6321 let mut aq4 = self.alloc_uninit::<u32>(m * (in_f / 8))?; let mut ad4 = self.alloc_uninit::<u8>(m * nb16)?; let cfg = LaunchConfig::for_num_elems((m * nb16) as u32);
6324 let (inf, mi) = (in_f as i32, m as i32);
6325 let __s_b = self.gpu.stream();
6326 let mut b = __s_b.launch_builder(&f);
6327 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6328 unsafe {
6329 b.launch(cfg)?;
6330 }
6331 Ok((aq4, ad4))
6332 }
6333
6334 pub fn qmatvec_gemm_nvfp4_fp4(
6339 &self,
6340 bytes: &CudaSlice<u8>,
6341 x: &CudaSlice<f32>,
6342 m: usize,
6343 in_f: usize,
6344 out_f: usize,
6345 row_bytes: usize,
6346 scale: f32,
6347 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6348 assert!(
6349 in_f % 64 == 0,
6350 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6351 );
6352 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6353 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6354 if scale != 1.0 {
6355 self.scale_inplace(&mut y, scale, m * out_f)?;
6356 }
6357 Ok(y)
6358 }
6359
6360 fn fp4_gemm_launch(
6363 &self,
6364 bytes: &CudaSlice<u8>,
6365 aq4: &CudaSlice<u32>,
6366 ad4: &CudaSlice<u8>,
6367 m: usize,
6368 in_f: usize,
6369 out_f: usize,
6370 row_bytes: usize,
6371 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6372 let f = self.func("qmatvec_gemm_nvfp4_fp4");
6373 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
6375 const BN: u32 = 256;
6376 let cfg = LaunchConfig {
6377 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6378 block_dim: (32, 4, 1),
6379 shared_mem_bytes: 0,
6380 };
6381 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6382 let __s_b = self.gpu.stream();
6383 let mut b = __s_b.launch_builder(&f);
6384 b.arg(bytes)
6385 .arg(aq4)
6386 .arg(ad4)
6387 .arg(&mut y)
6388 .arg(&inf)
6389 .arg(&outf)
6390 .arg(&mi)
6391 .arg(&rb);
6392 unsafe {
6393 b.launch(cfg)?;
6394 }
6395 Ok(y)
6396 }
6397
6398 pub fn qmatvec_gemm_nvfp4_fp4_raw(
6400 &self,
6401 bytes: &CudaSlice<u8>,
6402 x: &CudaSlice<f32>,
6403 m: usize,
6404 in_f: usize,
6405 out_f: usize,
6406 row_bytes: usize,
6407 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6408 assert!(
6409 in_f % 64 == 0,
6410 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6411 );
6412 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6413 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6414 }
6415
6416 pub fn qmatvec_q8_0_fast(
6418 &self,
6419 w: &CudaSlice<u8>,
6420 x: &CudaSlice<f32>,
6421 m: usize,
6422 in_f: usize,
6423 out_f: usize,
6424 row_bytes: usize,
6425 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6426 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6427 let f = self.func("qmatvec_q8_0_dp4a");
6428 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6430 grid_dim: (out_f as u32, m as u32, 1),
6431 block_dim: (128, 1, 1),
6432 shared_mem_bytes: 0,
6433 };
6434 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6435 let __s_b = self.gpu.stream();
6436 let mut b = __s_b.launch_builder(&f);
6437 b.arg(w)
6438 .arg(&aq)
6439 .arg(&ad)
6440 .arg(&mut y)
6441 .arg(&inf)
6442 .arg(&outf)
6443 .arg(&mi)
6444 .arg(&rb);
6445 unsafe {
6446 b.launch(cfg)?;
6447 }
6448 Ok(y)
6449 }
6450
6451 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
6454 &self,
6455 w: &CudaSlice<u8>,
6456 x: &CudaSlice<f32>,
6457 m: usize,
6458 in_f: usize,
6459 out_f: usize,
6460 row_bytes: usize,
6461 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6462 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6463 let f = self.func("qmatvec_q4_K_dp4a");
6464 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6466 grid_dim: (out_f as u32, m as u32, 1),
6467 block_dim: (128, 1, 1),
6468 shared_mem_bytes: 0,
6469 };
6470 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6471 let __s_b = self.gpu.stream();
6472 let mut b = __s_b.launch_builder(&f);
6473 b.arg(w)
6474 .arg(&aq)
6475 .arg(&ad)
6476 .arg(&mut y)
6477 .arg(&inf)
6478 .arg(&outf)
6479 .arg(&mi)
6480 .arg(&rb);
6481 unsafe {
6482 b.launch(cfg)?;
6483 }
6484 Ok(y)
6485 }
6486
6487 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
6490 &self,
6491 w: &CudaSlice<u8>,
6492 x: &CudaSlice<f32>,
6493 m: usize,
6494 in_f: usize,
6495 out_f: usize,
6496 row_bytes: usize,
6497 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6498 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6499 let f = self.func("qmatvec_q6_K_dp4a");
6500 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6502 grid_dim: (out_f as u32, m as u32, 1),
6503 block_dim: (128, 1, 1),
6504 shared_mem_bytes: 0,
6505 };
6506 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6507 let __s_b = self.gpu.stream();
6508 let mut b = __s_b.launch_builder(&f);
6509 b.arg(w)
6510 .arg(&aq)
6511 .arg(&ad)
6512 .arg(&mut y)
6513 .arg(&inf)
6514 .arg(&outf)
6515 .arg(&mi)
6516 .arg(&rb);
6517 unsafe {
6518 b.launch(cfg)?;
6519 }
6520 Ok(y)
6521 }
6522
6523 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
6526 &self,
6527 w: &CudaSlice<u8>,
6528 x: &CudaSlice<f32>,
6529 m: usize,
6530 in_f: usize,
6531 out_f: usize,
6532 row_bytes: usize,
6533 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6534 self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6535 }
6536 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
6539 &self,
6540 w: &CudaSlice<u8>,
6541 x: &CudaSlice<f32>,
6542 m: usize,
6543 in_f: usize,
6544 out_f: usize,
6545 row_bytes: usize,
6546 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6547 self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6548 }
6549 pub fn qmatvec_nvfp4_fast_rp(
6551 &self,
6552 w: &CudaSlice<u8>,
6553 x: &CudaSlice<f32>,
6554 m: usize,
6555 in_f: usize,
6556 out_f: usize,
6557 row_bytes: usize,
6558 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6559 assert!(
6560 in_f % 64 == 0,
6561 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6562 );
6563 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6564 }
6565 pub fn qmatvec_nvfp4_fast(
6567 &self,
6568 w: &CudaSlice<u8>,
6569 x: &CudaSlice<f32>,
6570 m: usize,
6571 in_f: usize,
6572 out_f: usize,
6573 row_bytes: usize,
6574 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6575 assert!(
6578 in_f % 64 == 0,
6579 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6580 );
6581 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6582 }
6583 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
6586 &self,
6587 w: &CudaSlice<u8>,
6588 x: &CudaSlice<f32>,
6589 m: usize,
6590 in_f: usize,
6591 out_f: usize,
6592 row_bytes: usize,
6593 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6594 self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6595 }
6596
6597 fn qmatvec_dp4a_named(
6599 &self,
6600 name: &str,
6601 w: &CudaSlice<u8>,
6602 x: &CudaSlice<f32>,
6603 m: usize,
6604 in_f: usize,
6605 out_f: usize,
6606 row_bytes: usize,
6607 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6608 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6609 let f = self.func(name);
6610 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6612 grid_dim: (out_f as u32, m as u32, 1),
6613 block_dim: (128, 1, 1),
6614 shared_mem_bytes: 0,
6615 };
6616 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6617 let __s_b = self.gpu.stream();
6618 let mut b = __s_b.launch_builder(&f);
6619 b.arg(w)
6620 .arg(&aq)
6621 .arg(&ad)
6622 .arg(&mut y)
6623 .arg(&inf)
6624 .arg(&outf)
6625 .arg(&mi)
6626 .arg(&rb);
6627 unsafe {
6628 b.launch(cfg)?;
6629 }
6630 Ok(y)
6631 }
6632
6633 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6634 Ok(self.gpu.stream().clone_htod(v)?)
6635 }
6636 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6637 Ok(self.gpu.stream().clone_htod(v)?)
6638 }
6639 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6641 Ok(self.gpu.stream().clone_htod(v)?)
6642 }
6643 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
6644 Ok(self.gpu.stream().clone_htod(v)?)
6645 }
6646 pub fn dtoh_view(
6648 &self,
6649 d: &cudarc::driver::CudaView<f32>,
6650 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6651 let v = self.gpu.stream().clone_dtoh(d)?;
6652 self.gpu.stream().synchronize()?;
6653 Ok(v)
6654 }
6655 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6656 let v = self.gpu.stream().clone_dtoh(d)?;
6657 self.gpu.stream().synchronize()?;
6658 Ok(v)
6659 }
6660 pub fn dtoh_pair(
6664 &self,
6665 a: &CudaSlice<f32>,
6666 b: &CudaSlice<f32>,
6667 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
6668 let av = self.gpu.stream().clone_dtoh(a)?;
6669 let bv = self.gpu.stream().clone_dtoh(b)?;
6670 self.gpu.stream().synchronize()?;
6671 Ok((av, bv))
6672 }
6673 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
6675 let v = self.gpu.stream().clone_dtoh(d)?;
6676 self.gpu.stream().synchronize()?;
6677 Ok(v)
6678 }
6679 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6681 let v = self.gpu.stream().clone_dtoh(d)?;
6682 self.gpu.stream().synchronize()?;
6683 Ok(v)
6684 }
6685 pub fn dtoh_u8_view(
6686 &self,
6687 d: &cudarc::driver::CudaView<u8>,
6688 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6689 let v = self.gpu.stream().clone_dtoh(d)?;
6690 self.gpu.stream().synchronize()?;
6691 Ok(v)
6692 }
6693 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6694 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
6695 self.keep_if_capturing(&s);
6696 Ok(s)
6697 }
6698
6699 pub fn prob_of_token_device(
6708 &self,
6709 logits: &CudaSlice<f32>,
6710 tok: &CudaSlice<u32>,
6711 n_vocab: usize,
6712 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6713 let nb = ARGMAX_NB;
6714 let mut part = self.alloc_uninit::<f32>(nb)?;
6715 let mut p = self.alloc_uninit::<f32>(1)?;
6716 let f1 = self.func("prob_of_token_partial_f32");
6717 let cfg1 = LaunchConfig {
6718 grid_dim: (nb as u32, 1, 1),
6719 block_dim: (256, 1, 1),
6720 shared_mem_bytes: 0,
6721 };
6722 let nv = n_vocab as i32;
6723 let __s_b1 = self.gpu.stream();
6724 let mut b1 = __s_b1.launch_builder(&f1);
6725 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6726 unsafe {
6727 b1.launch(cfg1)?;
6728 }
6729 let f2 = self.func("prob_of_token_final_f32");
6730 let cfg2 = LaunchConfig {
6731 grid_dim: (1, 1, 1),
6732 block_dim: (256, 1, 1),
6733 shared_mem_bytes: 0,
6734 };
6735 let nbi = nb as i32;
6736 let __s_b2 = self.gpu.stream();
6737 let mut b2 = __s_b2.launch_builder(&f2);
6738 b2.arg(&part).arg(&mut p).arg(&nbi);
6739 unsafe {
6740 b2.launch(cfg2)?;
6741 }
6742 Ok(p)
6743 }
6744
6745 pub fn prob_of_token_device_col(
6752 &self,
6753 logits: &CudaSlice<f32>,
6754 tok_all: &CudaSlice<u32>,
6755 tok_idx: usize,
6756 p_out: &mut CudaSlice<f32>,
6757 p_idx: usize,
6758 n_vocab: usize,
6759 ) -> Result<(), Box<dyn std::error::Error>> {
6760 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
6761 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
6762 let nb = ARGMAX_NB;
6763 let mut part = self.alloc_uninit::<f32>(nb)?;
6764 let f1 = self.func("prob_of_token_partial_f32");
6765 let cfg1 = LaunchConfig {
6766 grid_dim: (nb as u32, 1, 1),
6767 block_dim: (256, 1, 1),
6768 shared_mem_bytes: 0,
6769 };
6770 let nv = n_vocab as i32;
6771 let __s_b1 = self.gpu.stream();
6772 let mut b1 = __s_b1.launch_builder(&f1);
6773 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
6774 unsafe {
6775 b1.launch(cfg1)?;
6776 }
6777 let f2 = self.func("prob_of_token_final_f32");
6778 let cfg2 = LaunchConfig {
6779 grid_dim: (1, 1, 1),
6780 block_dim: (256, 1, 1),
6781 shared_mem_bytes: 0,
6782 };
6783 let nbi = nb as i32;
6784 let __s_b2 = self.gpu.stream();
6785 let mut b2 = __s_b2.launch_builder(&f2);
6786 b2.arg(&part).arg(&mut p_v).arg(&nbi);
6787 unsafe {
6788 b2.launch(cfg2)?;
6789 }
6790 Ok(())
6791 }
6792
6793 pub fn prob_of_token_device_into(
6794 &self,
6795 logits: &CudaSlice<f32>,
6796 tok: &CudaSlice<u32>,
6797 p_out: &mut CudaSlice<f32>,
6798 n_vocab: usize,
6799 ) -> Result<(), Box<dyn std::error::Error>> {
6800 let nb = ARGMAX_NB;
6801 let mut part = self.alloc_uninit::<f32>(nb)?;
6802 let f1 = self.func("prob_of_token_partial_f32");
6803 let cfg1 = LaunchConfig {
6804 grid_dim: (nb as u32, 1, 1),
6805 block_dim: (256, 1, 1),
6806 shared_mem_bytes: 0,
6807 };
6808 let nv = n_vocab as i32;
6809 let __s_b1 = self.gpu.stream();
6810 let mut b1 = __s_b1.launch_builder(&f1);
6811 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6812 unsafe {
6813 b1.launch(cfg1)?;
6814 }
6815 let f2 = self.func("prob_of_token_final_f32");
6816 let cfg2 = LaunchConfig {
6817 grid_dim: (1, 1, 1),
6818 block_dim: (256, 1, 1),
6819 shared_mem_bytes: 0,
6820 };
6821 let nbi = nb as i32;
6822 let __s_b2 = self.gpu.stream();
6823 let mut b2 = __s_b2.launch_builder(&f2);
6824 b2.arg(&part).arg(p_out).arg(&nbi);
6825 unsafe {
6826 b2.launch(cfg2)?;
6827 }
6828 Ok(())
6829 }
6830
6831 pub fn argmax_token_device(
6832 &self,
6833 logits: &CudaSlice<f32>,
6834 n_vocab: usize,
6835 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6836 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
6837 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
6838 Ok(tok)
6839 }
6840 pub fn argmax_token_device_into(
6847 &self,
6848 logits: &CudaSlice<f32>,
6849 tok: &mut CudaSlice<u32>,
6850 n_vocab: usize,
6851 ) -> Result<(), Box<dyn std::error::Error>> {
6852 let nb = ARGMAX_NB;
6853 let f1 = self.func("argmax_partial_f32");
6854 let f2 = self.func("argmax_final_f32");
6855 let mut guard = self.argmax_partials.lock().unwrap();
6856 if guard.is_none() {
6857 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
6860 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
6861 *guard = Some((pv, pi));
6862 }
6863 let (part_v, part_i) = guard.as_mut().unwrap();
6864 let nv = n_vocab as i32;
6865 let nbi = nb as i32;
6866 let cfg1 = LaunchConfig {
6868 grid_dim: (nb as u32, 1, 1),
6869 block_dim: (256, 1, 1),
6870 shared_mem_bytes: 0,
6871 };
6872 let __s_b1 = self.gpu.stream();
6873 let mut b1 = __s_b1.launch_builder(&f1);
6874 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
6875 unsafe {
6876 b1.launch(cfg1)?;
6877 }
6878 let cfg2 = LaunchConfig {
6880 grid_dim: (1, 1, 1),
6881 block_dim: (256, 1, 1),
6882 shared_mem_bytes: 0,
6883 };
6884 let __s_b2 = self.gpu.stream();
6885 let mut b2 = __s_b2.launch_builder(&f2);
6886 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
6887 unsafe {
6888 b2.launch(cfg2)?;
6889 }
6890 Ok(())
6891 }
6892 pub fn argmax_token_device_col(
6898 &self,
6899 logits: &CudaSlice<f32>,
6900 col: usize,
6901 n_vocab: usize,
6902 toks: &mut CudaSlice<u32>,
6903 out_idx: usize,
6904 ) -> Result<(), Box<dyn std::error::Error>> {
6905 let nb = ARGMAX_NB;
6906 let f1 = self.func("argmax_partial_f32");
6907 let f2 = self.func("argmax_final_f32");
6908 let mut guard = self.argmax_partials.lock().unwrap();
6909 if guard.is_none() {
6910 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
6911 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
6912 *guard = Some((pv, pi));
6913 }
6914 let (part_v, part_i) = guard.as_mut().unwrap();
6915 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
6916 let nv = n_vocab as i32;
6917 let nbi = nb as i32;
6918 let cfg1 = LaunchConfig {
6919 grid_dim: (nb as u32, 1, 1),
6920 block_dim: (256, 1, 1),
6921 shared_mem_bytes: 0,
6922 };
6923 let __s_b1 = self.gpu.stream();
6924 let mut b1 = __s_b1.launch_builder(&f1);
6925 b1.arg(&col_view)
6926 .arg(&mut *part_v)
6927 .arg(&mut *part_i)
6928 .arg(&nv);
6929 unsafe {
6930 b1.launch(cfg1)?;
6931 }
6932 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
6933 let cfg2 = LaunchConfig {
6934 grid_dim: (1, 1, 1),
6935 block_dim: (256, 1, 1),
6936 shared_mem_bytes: 0,
6937 };
6938 let __s_b2 = self.gpu.stream();
6939 let mut b2 = __s_b2.launch_builder(&f2);
6940 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
6941 unsafe {
6942 b2.launch(cfg2)?;
6943 }
6944 Ok(())
6945 }
6946 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6948 Ok(self.gpu.stream().clone_htod(v)?)
6949 }
6950 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6951 let v = self.gpu.stream().clone_dtoh(d)?;
6952 self.gpu.stream().synchronize()?;
6953 Ok(v)
6954 }
6955 pub fn htod_u32_into(
6959 &self,
6960 dst: &mut CudaSlice<u32>,
6961 src: &[u32],
6962 ) -> Result<(), Box<dyn std::error::Error>> {
6963 let mut view = dst.slice_mut(0..src.len());
6964 self.gpu.stream().memcpy_htod(src, &mut view)?;
6965 Ok(())
6966 }
6967
6968 pub fn htod_i32_into(
6971 &self,
6972 dst: &mut CudaSlice<i32>,
6973 src: &[i32],
6974 ) -> Result<(), Box<dyn std::error::Error>> {
6975 let mut view = dst.slice_mut(0..src.len());
6976 self.gpu.stream().memcpy_htod(src, &mut view)?;
6977 Ok(())
6978 }
6979
6980 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6981 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
6982 self.keep_if_capturing(&s);
6983 Ok(s)
6984 }
6985 pub fn embed_gather_device_into(
6988 &self,
6989 embd: &CudaSlice<u8>,
6990 token_d: &CudaSlice<u32>,
6991 x_out: &mut CudaSlice<f32>,
6992 n_embd: usize,
6993 qtype: i32,
6994 row_bytes: usize,
6995 ) -> Result<(), Box<dyn std::error::Error>> {
6996 let f = self.func("embed_gather_u32");
6997 let cfg = LaunchConfig {
6998 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
6999 block_dim: (256, 1, 1),
7000 shared_mem_bytes: 0,
7001 };
7002 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7003 let __s_b = self.gpu.stream();
7004 let mut b = __s_b.launch_builder(&f);
7005 b.arg(embd)
7006 .arg(token_d)
7007 .arg(x_out)
7008 .arg(&ne)
7009 .arg(&qt)
7010 .arg(&rb);
7011 unsafe {
7012 b.launch(cfg)?;
7013 }
7014 Ok(())
7015 }
7016 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7018 let v = self.gpu.stream().clone_dtoh(d)?;
7019 self.gpu.stream().synchronize()?;
7020 Ok(v[0])
7021 }
7022 pub fn i32_set_k(
7029 &self,
7030 dst: &mut CudaSlice<i32>,
7031 v: i32,
7032 ) -> Result<(), Box<dyn std::error::Error>> {
7033 let f = self.func("i32_set_k");
7034 let cfg = LaunchConfig {
7035 grid_dim: (1, 1, 1),
7036 block_dim: (1, 1, 1),
7037 shared_mem_bytes: 0,
7038 };
7039 let idx = 0i32;
7040 let __s_b = self.gpu.stream();
7041 let mut b = __s_b.launch_builder(&f);
7042 b.arg(dst).arg(&v).arg(&idx);
7043 unsafe {
7044 b.launch(cfg)?;
7045 }
7046 Ok(())
7047 }
7048
7049 pub fn set_i32_one(
7050 &self,
7051 d: &mut CudaSlice<i32>,
7052 v: i32,
7053 ) -> Result<(), Box<dyn std::error::Error>> {
7054 self.gpu.stream().memcpy_htod(&[v], d)?;
7055 Ok(())
7056 }
7057 pub fn set_u32_one(
7060 &self,
7061 d: &mut CudaSlice<u32>,
7062 v: u32,
7063 ) -> Result<(), Box<dyn std::error::Error>> {
7064 self.gpu.stream().memcpy_htod(&[v], d)?;
7065 Ok(())
7066 }
7067 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7069 let v = self.gpu.stream().clone_dtoh(d)?;
7070 self.gpu.stream().synchronize()?;
7071 Ok(v[0])
7072 }
7073 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7075 Ok(self.gpu.stream().clone_htod(bytes)?)
7076 }
7077 pub fn embed_gather_device(
7081 &self,
7082 embd: &CudaSlice<u8>,
7083 token_d: &CudaSlice<u32>,
7084 n_embd: usize,
7085 qtype: i32,
7086 row_bytes: usize,
7087 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7088 let f = self.func("embed_gather_u32");
7089 let mut x = self.alloc_uninit::<f32>(n_embd)?;
7090 let cfg = LaunchConfig {
7091 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7092 block_dim: (256, 1, 1),
7093 shared_mem_bytes: 0,
7094 };
7095 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7096 let __s_b = self.gpu.stream();
7097 let mut b = __s_b.launch_builder(&f);
7098 b.arg(embd)
7099 .arg(token_d)
7100 .arg(&mut x)
7101 .arg(&ne)
7102 .arg(&qt)
7103 .arg(&rb);
7104 unsafe {
7105 b.launch(cfg)?;
7106 }
7107 Ok(x)
7108 }
7109
7110 pub fn embed_gather_device_t(
7114 &self,
7115 embd: &CudaSlice<u8>,
7116 tokens: &[u32],
7117 n_embd: usize,
7118 qtype: i32,
7119 row_bytes: usize,
7120 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7121 let t = tokens.len();
7122 let tok_d = self.gpu.stream().clone_htod(tokens)?;
7123 let f = self.func("embed_gather_u32_t");
7124 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7125 let cfg = LaunchConfig {
7126 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7127 block_dim: (256, 1, 1),
7128 shared_mem_bytes: 0,
7129 };
7130 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7131 let __s_b = self.gpu.stream();
7132 let mut b = __s_b.launch_builder(&f);
7133 b.arg(embd)
7134 .arg(&tok_d)
7135 .arg(&mut x)
7136 .arg(&ne)
7137 .arg(&qt)
7138 .arg(&rb)
7139 .arg(&ti);
7140 unsafe {
7141 b.launch(cfg)?;
7142 }
7143 Ok(x)
7144 }
7145
7146 pub fn embed_gather_device_tv(
7151 &self,
7152 embd: &CudaSlice<u8>,
7153 tok_v: &cudarc::driver::CudaView<u32>,
7154 t: usize,
7155 n_embd: usize,
7156 qtype: i32,
7157 row_bytes: usize,
7158 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7159 let f = self.func("embed_gather_u32_t");
7160 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7161 let cfg = LaunchConfig {
7162 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7163 block_dim: (256, 1, 1),
7164 shared_mem_bytes: 0,
7165 };
7166 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7167 let __s_b = self.gpu.stream();
7168 let mut b = __s_b.launch_builder(&f);
7169 b.arg(embd)
7170 .arg(tok_v)
7171 .arg(&mut x)
7172 .arg(&ne)
7173 .arg(&qt)
7174 .arg(&rb)
7175 .arg(&ti);
7176 unsafe {
7177 b.launch(cfg)?;
7178 }
7179 Ok(x)
7180 }
7181
7182 pub fn embed_gather_device_td(
7183 &self,
7184 embd: &CudaSlice<u8>,
7185 tok_d: &CudaSlice<u32>,
7186 t: usize,
7187 n_embd: usize,
7188 qtype: i32,
7189 row_bytes: usize,
7190 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7191 let f = self.func("embed_gather_u32_t");
7192 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7193 let cfg = LaunchConfig {
7194 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7195 block_dim: (256, 1, 1),
7196 shared_mem_bytes: 0,
7197 };
7198 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7199 let __s_b = self.gpu.stream();
7200 let mut b = __s_b.launch_builder(&f);
7201 b.arg(embd)
7202 .arg(tok_d)
7203 .arg(&mut x)
7204 .arg(&ne)
7205 .arg(&qt)
7206 .arg(&rb)
7207 .arg(&ti);
7208 unsafe {
7209 b.launch(cfg)?;
7210 }
7211 Ok(x)
7212 }
7213
7214 #[inline]
7220 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7222 if self
7223 .capture_keep_on
7224 .load(std::sync::atomic::Ordering::Relaxed)
7225 {
7226 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7227 }
7228 }
7229
7230 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7231 &self,
7232 n: usize,
7233 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7234 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7235 {
7239 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7240 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7241 use cudarc::driver::DevicePtrMut;
7243 let n_bytes = s.len() * std::mem::size_of::<T>();
7244 let stream = self.gpu.stream();
7245 let (p_, _g) = s.device_ptr_mut(&stream);
7246 unsafe {
7247 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7248 .result()?;
7249 }
7250 }
7251 }
7252 self.keep_if_capturing(&s);
7253 Ok(s)
7254 }
7255
7256 pub fn uninit_q8_pair(
7261 &self,
7262 n: usize,
7263 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7264 Ok((
7265 self.alloc_uninit::<i8>(n)?,
7266 self.alloc_uninit::<f32>(n / 32)?,
7267 ))
7268 }
7269
7270 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7271 self.alloc_uninit::<f32>(n)
7272 }
7273
7274 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7276 self.alloc_uninit::<i8>(n)
7277 }
7278
7279 #[allow(clippy::too_many_arguments)]
7283 pub fn rms_norm3(
7284 &self,
7285 x: &CudaSlice<f32>,
7286 w0: &CudaSlice<f32>,
7287 w1: &CudaSlice<f32>,
7288 w2: &CudaSlice<f32>,
7289 d0: &mut CudaSlice<f32>,
7290 d1: &mut CudaSlice<f32>,
7291 d2: &mut CudaSlice<f32>,
7292 ncols: usize,
7293 nrows: usize,
7294 eps: f32,
7295 ) -> Result<(), Box<dyn std::error::Error>> {
7296 let f = self.func("rms_norm3_f32");
7297 let cfg = LaunchConfig {
7298 grid_dim: (nrows as u32, 1, 1),
7299 block_dim: (rms_block(), 1, 1),
7300 shared_mem_bytes: 0,
7301 };
7302 let (nc, e) = (ncols as i32, eps);
7303 let __s_b = self.gpu.stream();
7304 let mut b = __s_b.launch_builder(&f);
7305 b.arg(x)
7306 .arg(w0)
7307 .arg(w1)
7308 .arg(w2)
7309 .arg(d0)
7310 .arg(d1)
7311 .arg(d2)
7312 .arg(&nc)
7313 .arg(&e);
7314 unsafe {
7315 b.launch(cfg)?;
7316 }
7317 Ok(())
7318 }
7319
7320 #[allow(clippy::too_many_arguments)]
7322 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7325 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7326 *WARP_ON.get_or_init(|| {
7327 std::env::var("MEMRA_QKVNORM_W")
7328 .map(|v| v != "0")
7329 .unwrap_or(true)
7330 }) && ncols % 4 == 0
7331 && rows >= 64
7332 }
7333
7334 #[allow(clippy::too_many_arguments)]
7337 pub fn rms_norm_qkv_w4b(
7338 &self,
7339 q: &CudaSlice<f32>,
7340 k: &CudaSlice<f32>,
7341 v: &CudaSlice<f32>,
7342 wq: &CudaSlice<f32>,
7343 wk: &CudaSlice<f32>,
7344 wv: &CudaSlice<f32>,
7345 dq: &mut CudaSlice<f32>,
7346 dk: &mut CudaSlice<f32>,
7347 dv: &mut CudaSlice<f32>,
7348 dvb: &mut CudaSlice<u8>,
7349 ncols: usize,
7350 rq: usize,
7351 rk: usize,
7352 eps: f32,
7353 vf16: bool,
7354 ) -> Result<(), Box<dyn std::error::Error>> {
7355 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7356 let f = self.func("rms_norm_qkv_w4b_f32");
7357 let rows = (rq + 2 * rk) as u32;
7358 let cfg = LaunchConfig {
7359 grid_dim: (rows.div_ceil(8), 1, 1),
7360 block_dim: (256, 1, 1),
7361 shared_mem_bytes: 0,
7362 };
7363 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7364 let vf = vf16 as i32;
7365 let __s_b = self.gpu.stream();
7366 let mut b = __s_b.launch_builder(&f);
7367 b.arg(q)
7368 .arg(k)
7369 .arg(v)
7370 .arg(wq)
7371 .arg(wk)
7372 .arg(wv)
7373 .arg(dq)
7374 .arg(dk)
7375 .arg(dv)
7376 .arg(&mut *dvb)
7377 .arg(&nc)
7378 .arg(&rqi)
7379 .arg(&rki)
7380 .arg(&rvi)
7381 .arg(&e)
7382 .arg(&vf);
7383 unsafe {
7384 b.launch(cfg)?;
7385 }
7386 Ok(())
7387 }
7388
7389 pub fn rms_norm_qkv(
7390 &self,
7391 q: &CudaSlice<f32>,
7392 k: &CudaSlice<f32>,
7393 v: &CudaSlice<f32>,
7394 wq: &CudaSlice<f32>,
7395 wk: &CudaSlice<f32>,
7396 wv: &CudaSlice<f32>,
7397 dq: &mut CudaSlice<f32>,
7398 dk: &mut CudaSlice<f32>,
7399 dv: &mut CudaSlice<f32>,
7400 ncols: usize,
7401 rq: usize,
7402 rk: usize,
7403 eps: f32,
7404 ) -> Result<(), Box<dyn std::error::Error>> {
7405 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7409 let warp_on = *WARP_ON.get_or_init(|| {
7410 std::env::var("MEMRA_QKVNORM_W")
7411 .map(|v| v != "0")
7412 .unwrap_or(true)
7413 });
7414 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7417 let f = self.func("rms_norm_qkv_w4_f32");
7418 let rows = (rq + 2 * rk) as u32;
7419 let cfg = LaunchConfig {
7420 grid_dim: (rows.div_ceil(8), 1, 1),
7421 block_dim: (256, 1, 1),
7422 shared_mem_bytes: 0,
7423 };
7424 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7425 let __s_b = self.gpu.stream();
7426 let mut b = __s_b.launch_builder(&f);
7427 b.arg(q)
7428 .arg(k)
7429 .arg(v)
7430 .arg(wq)
7431 .arg(wk)
7432 .arg(wv)
7433 .arg(dq)
7434 .arg(dk)
7435 .arg(dv)
7436 .arg(&nc)
7437 .arg(&rqi)
7438 .arg(&rki)
7439 .arg(&rvi)
7440 .arg(&e);
7441 unsafe {
7442 b.launch(cfg)?;
7443 }
7444 return Ok(());
7445 }
7446 let f = self.func("rms_norm_qkv_f32");
7447 let grid = (rq + 2 * rk) as u32;
7448 let cfg = LaunchConfig {
7449 grid_dim: (grid, 1, 1),
7450 block_dim: (rms_block(), 1, 1),
7451 shared_mem_bytes: 0,
7452 };
7453 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7454 let __s_b = self.gpu.stream();
7455 let mut b = __s_b.launch_builder(&f);
7456 b.arg(q)
7457 .arg(k)
7458 .arg(v)
7459 .arg(wq)
7460 .arg(wk)
7461 .arg(wv)
7462 .arg(dq)
7463 .arg(dk)
7464 .arg(dv)
7465 .arg(&nc)
7466 .arg(&rqi)
7467 .arg(&rki)
7468 .arg(&e);
7469 unsafe {
7470 b.launch(cfg)?;
7471 }
7472 Ok(())
7473 }
7474
7475 #[allow(clippy::too_many_arguments)]
7477 pub fn rms_norm2x(
7478 &self,
7479 a: &CudaSlice<f32>,
7480 bb: &CudaSlice<f32>,
7481 wa: &CudaSlice<f32>,
7482 wb: &CudaSlice<f32>,
7483 da: &mut CudaSlice<f32>,
7484 db: &mut CudaSlice<f32>,
7485 ncols: usize,
7486 nrows: usize,
7487 eps: f32,
7488 ) -> Result<(), Box<dyn std::error::Error>> {
7489 let f = self.func("rms_norm2x_f32");
7490 let cfg = LaunchConfig {
7491 grid_dim: (2 * nrows as u32, 1, 1),
7492 block_dim: (rms_block(), 1, 1),
7493 shared_mem_bytes: 0,
7494 };
7495 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7496 let __s_b = self.gpu.stream();
7497 let mut b = __s_b.launch_builder(&f);
7498 b.arg(a)
7499 .arg(bb)
7500 .arg(wa)
7501 .arg(wb)
7502 .arg(da)
7503 .arg(db)
7504 .arg(&nc)
7505 .arg(&nr)
7506 .arg(&e);
7507 unsafe {
7508 b.launch(cfg)?;
7509 }
7510 Ok(())
7511 }
7512
7513 pub fn softcap(
7515 &self,
7516 y: &mut CudaSlice<f32>,
7517 cap: f32,
7518 n: usize,
7519 ) -> Result<(), Box<dyn std::error::Error>> {
7520 let f = self.func("softcap_f32");
7521 let cfg = LaunchConfig::for_num_elems(n as u32);
7522 let ni = n as i32;
7523 let __s_b = self.gpu.stream();
7524 let mut b = __s_b.launch_builder(&f);
7525 b.arg(y).arg(&cap).arg(&ni);
7526 unsafe {
7527 b.launch(cfg)?;
7528 }
7529 Ok(())
7530 }
7531
7532 pub fn mask_ids_rows(
7535 &self,
7536 y: &mut CudaSlice<f32>,
7537 ids: &CudaSlice<i32>,
7538 n_ids: usize,
7539 n_vocab: usize,
7540 t: usize,
7541 ) -> Result<(), Box<dyn std::error::Error>> {
7542 let f = self.func("mask_ids_rows_f32");
7543 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7544 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7545 let __s_b = self.gpu.stream();
7546 let mut b = __s_b.launch_builder(&f);
7547 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7548 unsafe {
7549 b.launch(cfg)?;
7550 }
7551 Ok(())
7552 }
7553
7554 #[allow(clippy::too_many_arguments)]
7556 pub fn add_scale_rms_norm(
7557 &self,
7558 a: &CudaSlice<f32>,
7559 b_in: &CudaSlice<f32>,
7560 c: f32,
7561 w: &CudaSlice<f32>,
7562 res: &mut CudaSlice<f32>,
7563 dst: &mut CudaSlice<f32>,
7564 ncols: usize,
7565 nrows: usize,
7566 eps: f32,
7567 ) -> Result<(), Box<dyn std::error::Error>> {
7568 let f = self.func("add_scale_rms_norm_f32");
7569 let cfg = LaunchConfig {
7570 grid_dim: (nrows as u32, 1, 1),
7571 block_dim: (rms_block(), 1, 1),
7572 shared_mem_bytes: 0,
7573 };
7574 let (nc, e2) = (ncols as i32, eps);
7575 let __s_b = self.gpu.stream();
7576 let mut b = __s_b.launch_builder(&f);
7577 b.arg(a)
7578 .arg(b_in)
7579 .arg(&c)
7580 .arg(w)
7581 .arg(res)
7582 .arg(dst)
7583 .arg(&nc)
7584 .arg(&e2);
7585 unsafe {
7586 b.launch(cfg)?;
7587 }
7588 Ok(())
7589 }
7590
7591 #[allow(clippy::too_many_arguments)]
7594 pub fn add_scale_rms_norm_q8_1(
7595 &self,
7596 a: &CudaSlice<f32>,
7597 b_in: &CudaSlice<f32>,
7598 c: f32,
7599 w: &CudaSlice<f32>,
7600 res: &mut CudaSlice<f32>,
7601 ncols: usize,
7602 nrows: usize,
7603 eps: f32,
7604 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7605 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7606 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7607 let (nc, e2) = (ncols as i32, eps);
7608 if Self::pdl_on() && Self::pdl_wb_on() {
7609 {
7610 use cudarc::driver::{DevicePtr, DevicePtrMut};
7611 let s = &self.gpu.stream();
7612 let (pa, _g0) = a.device_ptr(s);
7613 let (pb, _g1) = b_in.device_ptr(s);
7614 let (pw, _g2) = w.device_ptr(s);
7615 let (pr, _g3) = res.device_ptr_mut(s);
7616 let (pq, _g4) = out_q.device_ptr_mut(s);
7617 let (pd, _g5) = out_d.device_ptr_mut(s);
7618 let mut ps = [
7619 &pa as *const _ as *mut std::ffi::c_void,
7620 &pb as *const _ as *mut _,
7621 &c as *const _ as *mut _,
7622 &pw as *const _ as *mut _,
7623 &pr as *const _ as *mut _,
7624 &pq as *const _ as *mut _,
7625 &pd as *const _ as *mut _,
7626 &nc as *const _ as *mut _,
7627 &e2 as *const _ as *mut _,
7628 ];
7629 unsafe {
7630 self.launch_pdl(
7631 "add_scale_rms_norm_q8_1",
7632 (nrows as u32, 1, 1),
7633 (rms_block(), 1, 1),
7634 &mut ps,
7635 )?;
7636 }
7637 }
7638 return Ok((out_q, out_d));
7639 }
7640 let f = self.func("add_scale_rms_norm_q8_1");
7641 let cfg = LaunchConfig {
7642 grid_dim: (nrows as u32, 1, 1),
7643 block_dim: (rms_block(), 1, 1),
7644 shared_mem_bytes: 0,
7645 };
7646 let __s_b = self.gpu.stream();
7647 let mut b = __s_b.launch_builder(&f);
7648 b.arg(a)
7649 .arg(b_in)
7650 .arg(&c)
7651 .arg(w)
7652 .arg(res)
7653 .arg(&mut out_q)
7654 .arg(&mut out_d)
7655 .arg(&nc)
7656 .arg(&e2);
7657 unsafe {
7658 b.launch(cfg)?;
7659 }
7660 Ok((out_q, out_d))
7661 }
7662
7663 #[allow(clippy::too_many_arguments)]
7665 pub fn add_scale_rms_norm_q8_1_into(
7666 &self,
7667 a: &CudaSlice<f32>,
7668 b_in: &CudaSlice<f32>,
7669 c: f32,
7670 w: &CudaSlice<f32>,
7671 res: &mut CudaSlice<f32>,
7672 ncols: usize,
7673 nrows: usize,
7674 eps: f32,
7675 out_q: &mut CudaSlice<i8>,
7676 out_d: &mut CudaSlice<f32>,
7677 ) -> Result<(), Box<dyn std::error::Error>> {
7678 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
7679 let (nc, e2) = (ncols as i32, eps);
7680 if Self::pdl_on() && Self::pdl_wb_on() {
7681 use cudarc::driver::{DevicePtr, DevicePtrMut};
7682 let s = &self.gpu.stream();
7683 let (pa, _g0) = a.device_ptr(s);
7684 let (pb, _g1) = b_in.device_ptr(s);
7685 let (pw, _g2) = w.device_ptr(s);
7686 let (pr, _g3) = res.device_ptr_mut(s);
7687 let (pq, _g4) = out_q.device_ptr_mut(s);
7688 let (pd, _g5) = out_d.device_ptr_mut(s);
7689 let mut ps = [
7690 &pa as *const _ as *mut std::ffi::c_void,
7691 &pb as *const _ as *mut _,
7692 &c as *const _ as *mut _,
7693 &pw as *const _ as *mut _,
7694 &pr as *const _ as *mut _,
7695 &pq as *const _ as *mut _,
7696 &pd as *const _ as *mut _,
7697 &nc as *const _ as *mut _,
7698 &e2 as *const _ as *mut _,
7699 ];
7700 unsafe {
7701 self.launch_pdl(
7702 "add_scale_rms_norm_q8_1",
7703 (nrows as u32, 1, 1),
7704 (rms_block(), 1, 1),
7705 &mut ps,
7706 )?;
7707 }
7708 return Ok(());
7709 }
7710 let f = self.func("add_scale_rms_norm_q8_1");
7711 let cfg = LaunchConfig {
7712 grid_dim: (nrows as u32, 1, 1),
7713 block_dim: (rms_block(), 1, 1),
7714 shared_mem_bytes: 0,
7715 };
7716 let __s_b = self.gpu.stream();
7717 let mut b = __s_b.launch_builder(&f);
7718 b.arg(a)
7719 .arg(b_in)
7720 .arg(&c)
7721 .arg(w)
7722 .arg(res)
7723 .arg(&mut *out_q)
7724 .arg(&mut *out_d)
7725 .arg(&nc)
7726 .arg(&e2);
7727 unsafe {
7728 b.launch(cfg)?;
7729 }
7730 Ok(())
7731 }
7732
7733 #[allow(clippy::too_many_arguments)]
7736 pub fn rms_pre_add_scale_rms_norm_q8_1(
7737 &self,
7738 a: &CudaSlice<f32>,
7739 wa: &CudaSlice<f32>,
7740 b_in: &CudaSlice<f32>,
7741 c: f32,
7742 w: &CudaSlice<f32>,
7743 res: &mut CudaSlice<f32>,
7744 ncols: usize,
7745 nrows: usize,
7746 eps: f32,
7747 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7748 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7749 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7750 let (nc, e2) = (ncols as i32, eps);
7751 if Self::pdl_on() {
7752 {
7753 use cudarc::driver::{DevicePtr, DevicePtrMut};
7754 let s = &self.gpu.stream();
7755 let (pa, _g0) = a.device_ptr(s);
7756 let (pwa, _g1) = wa.device_ptr(s);
7757 let (pb, _g2) = b_in.device_ptr(s);
7758 let (pw, _g3) = w.device_ptr(s);
7759 let (pr, _g4) = res.device_ptr_mut(s);
7760 let (pq, _g5) = out_q.device_ptr_mut(s);
7761 let (pd, _g6) = out_d.device_ptr_mut(s);
7762 let mut ps = [
7763 &pa as *const _ as *mut std::ffi::c_void,
7764 &pwa as *const _ as *mut _,
7765 &pb as *const _ as *mut _,
7766 &c as *const _ as *mut _,
7767 &pw as *const _ as *mut _,
7768 &pr as *const _ as *mut _,
7769 &pq as *const _ as *mut _,
7770 &pd as *const _ as *mut _,
7771 &nc as *const _ as *mut _,
7772 &e2 as *const _ as *mut _,
7773 ];
7774 unsafe {
7775 self.launch_pdl(
7776 "rms_pre_add_scale_rms_norm_q8_1",
7777 (nrows as u32, 1, 1),
7778 (rms_block(), 1, 1),
7779 &mut ps,
7780 )?;
7781 }
7782 }
7783 return Ok((out_q, out_d));
7784 }
7785 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
7786 let cfg = LaunchConfig {
7787 grid_dim: (nrows as u32, 1, 1),
7788 block_dim: (rms_block(), 1, 1),
7789 shared_mem_bytes: 0,
7790 };
7791 let __s_b = self.gpu.stream();
7792 let mut b = __s_b.launch_builder(&f);
7793 b.arg(a)
7794 .arg(wa)
7795 .arg(b_in)
7796 .arg(&c)
7797 .arg(w)
7798 .arg(res)
7799 .arg(&mut out_q)
7800 .arg(&mut out_d)
7801 .arg(&nc)
7802 .arg(&e2);
7803 unsafe {
7804 b.launch(cfg)?;
7805 }
7806 Ok((out_q, out_d))
7807 }
7808
7809 pub fn gelu_tanh_mul_q8_1(
7812 &self,
7813 gate: &CudaSlice<f32>,
7814 up: &cudarc::driver::CudaView<f32>,
7815 act: &mut CudaSlice<f32>,
7816 ncols: usize,
7817 nrows: usize,
7818 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7819 debug_assert!(ncols % 128 == 0);
7820 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7821 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7822 let nc = ncols as i32;
7823 if Self::pdl_on() {
7824 {
7825 use cudarc::driver::{DevicePtr, DevicePtrMut};
7826 let s = &self.gpu.stream();
7827 let (pg, _g0) = gate.device_ptr(s);
7828 let (pu, _g1) = up.device_ptr(s);
7829 let (pact, _g2) = act.device_ptr_mut(s);
7830 let (pq, _g3) = out_q.device_ptr_mut(s);
7831 let (pd, _g4) = out_d.device_ptr_mut(s);
7832 let mut ps = [
7833 &pg as *const _ as *mut std::ffi::c_void,
7834 &pu as *const _ as *mut _,
7835 &pact as *const _ as *mut _,
7836 &pq as *const _ as *mut _,
7837 &pd as *const _ as *mut _,
7838 &nc as *const _ as *mut _,
7839 ];
7840 unsafe {
7841 self.launch_pdl(
7842 "gelu_tanh_mul_q8_1",
7843 (nrows as u32, 1, 1),
7844 (rms_block(), 1, 1),
7845 &mut ps,
7846 )?;
7847 }
7848 }
7849 return Ok((out_q, out_d));
7850 }
7851 let f = self.func("gelu_tanh_mul_q8_1");
7852 let cfg = LaunchConfig {
7853 grid_dim: (nrows as u32, 1, 1),
7854 block_dim: (rms_block(), 1, 1),
7855 shared_mem_bytes: 0,
7856 };
7857 let __s_b = self.gpu.stream();
7858 let mut b = __s_b.launch_builder(&f);
7859 b.arg(gate)
7860 .arg(up)
7861 .arg(act)
7862 .arg(&mut out_q)
7863 .arg(&mut out_d)
7864 .arg(&nc);
7865 unsafe {
7866 b.launch(cfg)?;
7867 }
7868 Ok((out_q, out_d))
7869 }
7870
7871 #[allow(clippy::too_many_arguments)]
7873 pub fn gelu_tanh_mul_q8_1_into(
7874 &self,
7875 gate: &CudaSlice<f32>,
7876 up: &cudarc::driver::CudaView<f32>,
7877 act: &mut CudaSlice<f32>,
7878 ncols: usize,
7879 nrows: usize,
7880 out_q: &mut CudaSlice<i8>,
7881 out_d: &mut CudaSlice<f32>,
7882 ) -> Result<(), Box<dyn std::error::Error>> {
7883 debug_assert!(ncols % 128 == 0);
7884 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
7885 let nc = ncols as i32;
7886 if Self::pdl_on() {
7887 use cudarc::driver::{DevicePtr, DevicePtrMut};
7888 let s = &self.gpu.stream();
7889 let (pg, _g0) = gate.device_ptr(s);
7890 let (pu, _g1) = up.device_ptr(s);
7891 let (pact, _g2) = act.device_ptr_mut(s);
7892 let (pq, _g3) = out_q.device_ptr_mut(s);
7893 let (pd, _g4) = out_d.device_ptr_mut(s);
7894 let mut ps = [
7895 &pg as *const _ as *mut std::ffi::c_void,
7896 &pu as *const _ as *mut _,
7897 &pact as *const _ as *mut _,
7898 &pq as *const _ as *mut _,
7899 &pd as *const _ as *mut _,
7900 &nc as *const _ as *mut _,
7901 ];
7902 unsafe {
7903 self.launch_pdl(
7904 "gelu_tanh_mul_q8_1",
7905 (nrows as u32, 1, 1),
7906 (rms_block(), 1, 1),
7907 &mut ps,
7908 )?;
7909 }
7910 return Ok(());
7911 }
7912 let f = self.func("gelu_tanh_mul_q8_1");
7913 let cfg = LaunchConfig {
7914 grid_dim: (nrows as u32, 1, 1),
7915 block_dim: (rms_block(), 1, 1),
7916 shared_mem_bytes: 0,
7917 };
7918 let __s_b = self.gpu.stream();
7919 let mut b = __s_b.launch_builder(&f);
7920 b.arg(gate)
7921 .arg(up)
7922 .arg(&mut *act)
7923 .arg(&mut *out_q)
7924 .arg(&mut *out_d)
7925 .arg(&nc);
7926 unsafe {
7927 b.launch(cfg)?;
7928 }
7929 Ok(())
7930 }
7931
7932 #[allow(clippy::too_many_arguments)]
7934 pub fn add_rms_norm3_q8z(
7935 &self,
7936 a: &CudaSlice<f32>,
7937 b_in: &CudaSlice<f32>,
7938 w0: &CudaSlice<f32>,
7939 w1: &CudaSlice<f32>,
7940 w2: &CudaSlice<f32>,
7941 res: &mut CudaSlice<f32>,
7942 out1: &mut CudaSlice<f32>,
7943 ncols: usize,
7944 nrows: usize,
7945 eps: f32,
7946 ) -> Result<
7947 (
7948 (CudaSlice<i8>, CudaSlice<f32>),
7949 (CudaSlice<i8>, CudaSlice<f32>),
7950 ),
7951 Box<dyn std::error::Error>,
7952 > {
7953 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
7954 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7955 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
7956 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7957 let f = self.func("add_rms_norm3_q8z_f32");
7958 let cfg = LaunchConfig {
7959 grid_dim: (nrows as u32, 1, 1),
7960 block_dim: (rms_block(), 1, 1),
7961 shared_mem_bytes: 0,
7962 };
7963 let (nc, e2) = (ncols as i32, eps);
7964 let __s_b = self.gpu.stream();
7965 let mut b = __s_b.launch_builder(&f);
7966 b.arg(a)
7967 .arg(b_in)
7968 .arg(w0)
7969 .arg(w1)
7970 .arg(w2)
7971 .arg(res)
7972 .arg(&mut q0)
7973 .arg(&mut d0)
7974 .arg(out1)
7975 .arg(&mut q2)
7976 .arg(&mut d2)
7977 .arg(&nc)
7978 .arg(&e2);
7979 unsafe {
7980 b.launch(cfg)?;
7981 }
7982 Ok(((q0, d0), (q2, d2)))
7983 }
7984
7985 #[allow(clippy::too_many_arguments)]
7987 pub fn add_rms_norm3(
7988 &self,
7989 a: &CudaSlice<f32>,
7990 b_in: &CudaSlice<f32>,
7991 w0: &CudaSlice<f32>,
7992 w1: &CudaSlice<f32>,
7993 w2: &CudaSlice<f32>,
7994 res: &mut CudaSlice<f32>,
7995 d0: &mut CudaSlice<f32>,
7996 d1: &mut CudaSlice<f32>,
7997 d2: &mut CudaSlice<f32>,
7998 ncols: usize,
7999 nrows: usize,
8000 eps: f32,
8001 ) -> Result<(), Box<dyn std::error::Error>> {
8002 let f = self.func("add_rms_norm3_f32");
8003 let cfg = LaunchConfig {
8004 grid_dim: (nrows as u32, 1, 1),
8005 block_dim: (rms_block(), 1, 1),
8006 shared_mem_bytes: 0,
8007 };
8008 let (nc, e2) = (ncols as i32, eps);
8009 let __s_b = self.gpu.stream();
8010 let mut b = __s_b.launch_builder(&f);
8011 b.arg(a)
8012 .arg(b_in)
8013 .arg(w0)
8014 .arg(w1)
8015 .arg(w2)
8016 .arg(res)
8017 .arg(d0)
8018 .arg(d1)
8019 .arg(d2)
8020 .arg(&nc)
8021 .arg(&e2);
8022 unsafe {
8023 b.launch(cfg)?;
8024 }
8025 Ok(())
8026 }
8027
8028 pub fn add_scale(
8030 &self,
8031 a: &CudaSlice<f32>,
8032 b_in: &CudaSlice<f32>,
8033 c: f32,
8034 dst: &mut CudaSlice<f32>,
8035 n: usize,
8036 ) -> Result<(), Box<dyn std::error::Error>> {
8037 let f = self.func("add_scale_f32");
8038 let cfg = LaunchConfig::for_num_elems(n as u32);
8039 let ni = n as i32;
8040 let __s_b = self.gpu.stream();
8041 let mut b = __s_b.launch_builder(&f);
8042 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8043 unsafe {
8044 b.launch(cfg)?;
8045 }
8046 Ok(())
8047 }
8048
8049 pub fn layer_norm_bias(
8051 &self,
8052 x: &CudaSlice<f32>,
8053 w: &CudaSlice<f32>,
8054 b: &CudaSlice<f32>,
8055 dst: &mut CudaSlice<f32>,
8056 ncols: usize,
8057 nrows: usize,
8058 eps: f32,
8059 ) -> Result<(), Box<dyn std::error::Error>> {
8060 let f = self.func("layer_norm_bias_f32");
8061 let (nc, e) = (ncols as i32, eps);
8062 let cfg = LaunchConfig {
8063 grid_dim: (nrows as u32, 1, 1),
8064 block_dim: (256, 1, 1),
8065 shared_mem_bytes: 0,
8066 };
8067 let __s_b = self.gpu.stream();
8068 let mut lb = __s_b.launch_builder(&f);
8069 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8070 unsafe {
8071 lb.launch(cfg)?;
8072 }
8073 Ok(())
8074 }
8075
8076 pub fn gelu_tanh(
8078 &self,
8079 x: &CudaSlice<f32>,
8080 dst: &mut CudaSlice<f32>,
8081 n: usize,
8082 ) -> Result<(), Box<dyn std::error::Error>> {
8083 let f = self.func("gelu_tanh_f32");
8084 let ni = n as i64;
8085 let cfg = LaunchConfig {
8086 grid_dim: (n.div_ceil(256) as u32, 1, 1),
8087 block_dim: (256, 1, 1),
8088 shared_mem_bytes: 0,
8089 };
8090 let __s_b = self.gpu.stream();
8091 let mut lb = __s_b.launch_builder(&f);
8092 lb.arg(x).arg(&mut *dst).arg(&ni);
8093 unsafe {
8094 lb.launch(cfg)?;
8095 }
8096 Ok(())
8097 }
8098
8099 pub fn row_softmax(
8101 &self,
8102 x: &mut CudaSlice<f32>,
8103 ncols: usize,
8104 nrows: usize,
8105 ) -> Result<(), Box<dyn std::error::Error>> {
8106 let f = self.func("row_softmax_f32");
8107 let nc = ncols as i32;
8108 let cfg = LaunchConfig {
8109 grid_dim: (nrows as u32, 1, 1),
8110 block_dim: (256, 1, 1),
8111 shared_mem_bytes: 0,
8112 };
8113 let __s_b = self.gpu.stream();
8114 let mut lb = __s_b.launch_builder(&f);
8115 lb.arg(&mut *x).arg(&nc);
8116 unsafe {
8117 lb.launch(cfg)?;
8118 }
8119 Ok(())
8120 }
8121
8122 pub fn rms_norm(
8123 &self,
8124 x: &CudaSlice<f32>,
8125 w: &CudaSlice<f32>,
8126 dst: &mut CudaSlice<f32>,
8127 ncols: usize,
8128 nrows: usize,
8129 eps: f32,
8130 ) -> Result<(), Box<dyn std::error::Error>> {
8131 let (nc, e) = (ncols as i32, eps);
8132 if Self::pdl_on() && Self::pdl_wb_on() {
8133 use cudarc::driver::{DevicePtr, DevicePtrMut};
8134 let s = &self.gpu.stream();
8135 let (px, _g0) = x.device_ptr(s);
8136 let (pw, _g1) = w.device_ptr(s);
8137 let (pd, _g2) = dst.device_ptr_mut(s);
8138 let mut ps = [
8139 &px as *const _ as *mut std::ffi::c_void,
8140 &pw as *const _ as *mut _,
8141 &pd as *const _ as *mut _,
8142 &nc as *const _ as *mut _,
8143 &e as *const _ as *mut _,
8144 ];
8145 unsafe {
8146 self.launch_pdl(
8147 "rms_norm_f32",
8148 (nrows as u32, 1, 1),
8149 (rms_block(), 1, 1),
8150 &mut ps,
8151 )?;
8152 }
8153 return Ok(());
8154 }
8155 let f = self.func("rms_norm_f32");
8156 let cfg = LaunchConfig {
8157 grid_dim: (nrows as u32, 1, 1),
8158 block_dim: (rms_block(), 1, 1),
8159 shared_mem_bytes: 0,
8160 };
8161 let __s_b = self.gpu.stream();
8162 let mut b = __s_b.launch_builder(&f);
8163 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8164 unsafe {
8165 b.launch(cfg)?;
8166 }
8167 Ok(())
8168 }
8169
8170 pub fn rms_norm_decode(
8178 &self,
8179 x: &CudaSlice<f32>,
8180 w: &CudaSlice<f32>,
8181 dst: &mut CudaSlice<f32>,
8182 ncols: usize,
8183 nrows: usize,
8184 eps: f32,
8185 ) -> Result<(), Box<dyn std::error::Error>> {
8186 let f = self.func("rms_norm_f32");
8187 let cfg = LaunchConfig {
8188 grid_dim: (nrows as u32, 1, 1),
8189 block_dim: (1024, 1, 1),
8190 shared_mem_bytes: 0,
8191 };
8192 let (nc, e) = (ncols as i32, eps);
8193 let __s_b = self.gpu.stream();
8194 let mut b = __s_b.launch_builder(&f);
8195 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8196 unsafe {
8197 b.launch(cfg)?;
8198 }
8199 Ok(())
8200 }
8201
8202 pub fn rms_norm_q8_1(
8206 &self,
8207 x: &CudaSlice<f32>,
8208 w: &CudaSlice<f32>,
8209 ncols: usize,
8210 nrows: usize,
8211 eps: f32,
8212 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8213 let nblk = ncols / 32;
8214 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8215 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8216 let (nc, e) = (ncols as i32, eps);
8217 if Self::pdl_on() {
8218 {
8219 use cudarc::driver::{DevicePtr, DevicePtrMut};
8220 let s = &self.gpu.stream();
8221 let (px, _g0) = x.device_ptr(s);
8222 let (pw, _g1) = w.device_ptr(s);
8223 let (pq, _g2) = q.device_ptr_mut(s);
8224 let (pd, _g3) = d.device_ptr_mut(s);
8225 let mut ps = [
8226 &px as *const _ as *mut std::ffi::c_void,
8227 &pw as *const _ as *mut _,
8228 &pq as *const _ as *mut _,
8229 &pd as *const _ as *mut _,
8230 &nc as *const _ as *mut _,
8231 &e as *const _ as *mut _,
8232 ];
8233 unsafe {
8234 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8235 }
8236 }
8237 return Ok((q, d));
8238 }
8239 let f = self.func("rms_norm_q8_1");
8240 let cfg = LaunchConfig {
8243 grid_dim: (nrows as u32, 1, 1),
8244 block_dim: (1024, 1, 1),
8245 shared_mem_bytes: 0,
8246 };
8247 let __s_b = self.gpu.stream();
8248 let mut b = __s_b.launch_builder(&f);
8249 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8250 unsafe {
8251 b.launch(cfg)?;
8252 }
8253 Ok((q, d))
8254 }
8255
8256 pub fn rms_norm_q8_1_into(
8259 &self,
8260 x: &CudaSlice<f32>,
8261 w: &CudaSlice<f32>,
8262 ncols: usize,
8263 nrows: usize,
8264 eps: f32,
8265 q: &mut CudaSlice<i8>,
8266 d: &mut CudaSlice<f32>,
8267 ) -> Result<(), Box<dyn std::error::Error>> {
8268 let nblk = ncols / 32;
8269 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8270 let (nc, e) = (ncols as i32, eps);
8271 if Self::pdl_on() {
8272 use cudarc::driver::{DevicePtr, DevicePtrMut};
8273 let s = &self.gpu.stream();
8274 let (px, _g0) = x.device_ptr(s);
8275 let (pw, _g1) = w.device_ptr(s);
8276 let (pq, _g2) = q.device_ptr_mut(s);
8277 let (pd, _g3) = d.device_ptr_mut(s);
8278 let mut ps = [
8279 &px as *const _ as *mut std::ffi::c_void,
8280 &pw as *const _ as *mut _,
8281 &pq as *const _ as *mut _,
8282 &pd as *const _ as *mut _,
8283 &nc as *const _ as *mut _,
8284 &e as *const _ as *mut _,
8285 ];
8286 unsafe {
8287 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8288 }
8289 return Ok(());
8290 }
8291 let f = self.func("rms_norm_q8_1");
8292 let cfg = LaunchConfig {
8293 grid_dim: (nrows as u32, 1, 1),
8294 block_dim: (1024, 1, 1),
8295 shared_mem_bytes: 0,
8296 };
8297 let __s_b = self.gpu.stream();
8298 let mut b = __s_b.launch_builder(&f);
8299 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8300 unsafe {
8301 b.launch(cfg)?;
8302 }
8303 Ok(())
8304 }
8305
8306 pub fn quantize_q8_1_into(
8308 &self,
8309 x: &CudaSlice<f32>,
8310 m: usize,
8311 in_f: usize,
8312 q: &mut CudaSlice<i8>,
8313 d: &mut CudaSlice<f32>,
8314 ) -> Result<(), Box<dyn std::error::Error>> {
8315 let nblk = in_f / 32;
8316 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8317 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8318 let (inf, mi) = (in_f as i32, m as i32);
8319 if Self::pdl_on() && Self::pdl_wb_on() {
8320 use cudarc::driver::{DevicePtr, DevicePtrMut};
8321 let s = &self.gpu.stream();
8322 let (px, _g0) = x.device_ptr(s);
8323 let (pq, _g1) = q.device_ptr_mut(s);
8324 let (pd, _g2) = d.device_ptr_mut(s);
8325 let mut ps = [
8326 &px as *const _ as *mut std::ffi::c_void,
8327 &pq as *const _ as *mut _,
8328 &pd as *const _ as *mut _,
8329 &inf as *const _ as *mut _,
8330 &mi as *const _ as *mut _,
8331 ];
8332 unsafe {
8333 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8334 }
8335 return Ok(());
8336 }
8337 let f = self.func("quantize_q8_1");
8338 let __s_b = self.gpu.stream();
8339 let mut b = __s_b.launch_builder(&f);
8340 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8341 unsafe {
8342 b.launch(cfg)?;
8343 }
8344 Ok(())
8345 }
8346
8347 pub fn add_rms_norm_q8_1(
8351 &self,
8352 a: &CudaSlice<f32>,
8353 b_in: &CudaSlice<f32>,
8354 w: &CudaSlice<f32>,
8355 res: &mut CudaSlice<f32>,
8356 ncols: usize,
8357 nrows: usize,
8358 eps: f32,
8359 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8360 let nblk = ncols / 32;
8361 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8362 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8363 let f = self.func("add_rms_norm_q8_1");
8364 let cfg = LaunchConfig {
8366 grid_dim: (nrows as u32, 1, 1),
8367 block_dim: (1024, 1, 1),
8368 shared_mem_bytes: 0,
8369 };
8370 let (nc, e) = (ncols as i32, eps);
8371 let __s_bld = self.gpu.stream();
8372 let mut bld = __s_bld.launch_builder(&f);
8373 bld.arg(a)
8374 .arg(b_in)
8375 .arg(w)
8376 .arg(res)
8377 .arg(&mut q)
8378 .arg(&mut d)
8379 .arg(&nc)
8380 .arg(&e);
8381 unsafe {
8382 bld.launch(cfg)?;
8383 }
8384 Ok((q, d))
8385 }
8386
8387 pub fn add_rms_norm(
8391 &self,
8392 a: &CudaSlice<f32>,
8393 b: &CudaSlice<f32>,
8394 w: &CudaSlice<f32>,
8395 res: &mut CudaSlice<f32>,
8396 dst: &mut CudaSlice<f32>,
8397 ncols: usize,
8398 nrows: usize,
8399 eps: f32,
8400 ) -> Result<(), Box<dyn std::error::Error>> {
8401 let (nc, e) = (ncols as i32, eps);
8402 if Self::pdl_on() && Self::pdl_wb_on() {
8403 use cudarc::driver::{DevicePtr, DevicePtrMut};
8404 let s = &self.gpu.stream();
8405 let (pa, _g0) = a.device_ptr(s);
8406 let (pb, _g1) = b.device_ptr(s);
8407 let (pw, _g2) = w.device_ptr(s);
8408 let (pr, _g3) = res.device_ptr_mut(s);
8409 let (pd, _g4) = dst.device_ptr_mut(s);
8410 let mut ps = [
8411 &pa as *const _ as *mut std::ffi::c_void,
8412 &pb as *const _ as *mut _,
8413 &pw as *const _ as *mut _,
8414 &pr as *const _ as *mut _,
8415 &pd as *const _ as *mut _,
8416 &nc as *const _ as *mut _,
8417 &e as *const _ as *mut _,
8418 ];
8419 unsafe {
8420 self.launch_pdl(
8421 "add_rms_norm_f32",
8422 (nrows as u32, 1, 1),
8423 (rms_block(), 1, 1),
8424 &mut ps,
8425 )?;
8426 }
8427 return Ok(());
8428 }
8429 let f = self.func("add_rms_norm_f32");
8430 let cfg = LaunchConfig {
8431 grid_dim: (nrows as u32, 1, 1),
8432 block_dim: (rms_block(), 1, 1),
8433 shared_mem_bytes: 0,
8434 };
8435 let __s_b2 = self.gpu.stream();
8436 let mut b2 = __s_b2.launch_builder(&f);
8437 b2.arg(a)
8438 .arg(b)
8439 .arg(w)
8440 .arg(&mut *res)
8441 .arg(&mut *dst)
8442 .arg(&nc)
8443 .arg(&e);
8444 unsafe {
8445 b2.launch(cfg)?;
8446 }
8447 Ok(())
8448 }
8449
8450 #[allow(clippy::too_many_arguments)]
8453 pub fn rms_pre_add_rms_norm(
8454 &self,
8455 a: &CudaSlice<f32>,
8456 wa: &CudaSlice<f32>,
8457 b: &CudaSlice<f32>,
8458 w: &CudaSlice<f32>,
8459 res: &mut CudaSlice<f32>,
8460 dst: &mut CudaSlice<f32>,
8461 ncols: usize,
8462 nrows: usize,
8463 eps: f32,
8464 ) -> Result<(), Box<dyn std::error::Error>> {
8465 let f = self.func("rms_pre_add_rms_norm_f32");
8466 let cfg = LaunchConfig {
8467 grid_dim: (nrows as u32, 1, 1),
8468 block_dim: (rms_block(), 1, 1),
8469 shared_mem_bytes: 0,
8470 };
8471 let (nc, e) = (ncols as i32, eps);
8472 let __s_b2 = self.gpu.stream();
8473 let mut b2 = __s_b2.launch_builder(&f);
8474 b2.arg(a)
8475 .arg(wa)
8476 .arg(b)
8477 .arg(w)
8478 .arg(&mut *res)
8479 .arg(&mut *dst)
8480 .arg(&nc)
8481 .arg(&e);
8482 unsafe {
8483 b2.launch(cfg)?;
8484 }
8485 Ok(())
8486 }
8487
8488 #[allow(clippy::too_many_arguments)]
8490 pub fn rms_pre_add_rms_norm_q8z(
8491 &self,
8492 a: &CudaSlice<f32>,
8493 wa: &CudaSlice<f32>,
8494 b: &CudaSlice<f32>,
8495 w: &CudaSlice<f32>,
8496 res: &mut CudaSlice<f32>,
8497 dst: &mut CudaSlice<f32>,
8498 ncols: usize,
8499 nrows: usize,
8500 eps: f32,
8501 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8502 debug_assert!(ncols % 128 == 0);
8503 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8504 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8505 let (nc, e) = (ncols as i32, eps);
8506 if Self::pdl_on() {
8507 {
8508 use cudarc::driver::{DevicePtr, DevicePtrMut};
8509 let s = &self.gpu.stream();
8510 let (pa, _g0) = a.device_ptr(s);
8511 let (pwa, _g1) = wa.device_ptr(s);
8512 let (pb, _g2) = b.device_ptr(s);
8513 let (pw, _g3) = w.device_ptr(s);
8514 let (pr, _g4) = res.device_ptr_mut(s);
8515 let (pdst, _g5) = dst.device_ptr_mut(s);
8516 let (pq, _g6) = out_q.device_ptr_mut(s);
8517 let (pd, _g7) = out_d.device_ptr_mut(s);
8518 let mut ps = [
8519 &pa as *const _ as *mut std::ffi::c_void,
8520 &pwa as *const _ as *mut _,
8521 &pb as *const _ as *mut _,
8522 &pw as *const _ as *mut _,
8523 &pr as *const _ as *mut _,
8524 &pdst as *const _ as *mut _,
8525 &pq as *const _ as *mut _,
8526 &pd as *const _ as *mut _,
8527 &nc as *const _ as *mut _,
8528 &e as *const _ as *mut _,
8529 ];
8530 unsafe {
8531 self.launch_pdl(
8532 "rms_pre_add_rms_norm_q8z_f32",
8533 (nrows as u32, 1, 1),
8534 (rms_block(), 1, 1),
8535 &mut ps,
8536 )?;
8537 }
8538 }
8539 return Ok((out_q, out_d));
8540 }
8541 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8542 let cfg = LaunchConfig {
8543 grid_dim: (nrows as u32, 1, 1),
8544 block_dim: (rms_block(), 1, 1),
8545 shared_mem_bytes: 0,
8546 };
8547 let __s_b2 = self.gpu.stream();
8548 let mut b2 = __s_b2.launch_builder(&f);
8549 b2.arg(a)
8550 .arg(wa)
8551 .arg(b)
8552 .arg(w)
8553 .arg(&mut *res)
8554 .arg(&mut *dst)
8555 .arg(&mut out_q)
8556 .arg(&mut out_d)
8557 .arg(&nc)
8558 .arg(&e);
8559 unsafe {
8560 b2.launch(cfg)?;
8561 }
8562 Ok((out_q, out_d))
8563 }
8564
8565 #[allow(clippy::too_many_arguments)]
8569 pub fn rms_pre_add_rms_norm_q8z_into(
8570 &self,
8571 a: &CudaSlice<f32>,
8572 wa: &CudaSlice<f32>,
8573 b: &CudaSlice<f32>,
8574 w: &CudaSlice<f32>,
8575 res: &mut CudaSlice<f32>,
8576 dst: &mut CudaSlice<f32>,
8577 ncols: usize,
8578 nrows: usize,
8579 eps: f32,
8580 out_q: &mut CudaSlice<i8>,
8581 out_d: &mut CudaSlice<f32>,
8582 ) -> Result<(), Box<dyn std::error::Error>> {
8583 debug_assert!(ncols % 128 == 0);
8584 let (nc, e) = (ncols as i32, eps);
8585 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8586 let cfg = LaunchConfig {
8587 grid_dim: (nrows as u32, 1, 1),
8588 block_dim: (rms_block(), 1, 1),
8589 shared_mem_bytes: 0,
8590 };
8591 let __s_b = self.gpu.stream();
8592 let mut b2 = __s_b.launch_builder(&f);
8593 b2.arg(a)
8594 .arg(wa)
8595 .arg(b)
8596 .arg(w)
8597 .arg(&mut *res)
8598 .arg(&mut *dst)
8599 .arg(&mut *out_q)
8600 .arg(&mut *out_d)
8601 .arg(&nc)
8602 .arg(&e);
8603 unsafe {
8604 b2.launch(cfg)?;
8605 }
8606 Ok(())
8607 }
8608
8609 #[allow(clippy::too_many_arguments)]
8612 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8613 &self,
8614 a: &CudaSlice<f32>,
8615 wa: &CudaSlice<f32>,
8616 b_in: &CudaSlice<f32>,
8617 c: f32,
8618 w: &CudaSlice<f32>,
8619 res: &mut CudaSlice<f32>,
8620 ncols: usize,
8621 nrows: usize,
8622 eps: f32,
8623 out_q: &mut CudaSlice<i8>,
8624 out_d: &mut CudaSlice<f32>,
8625 ) -> Result<(), Box<dyn std::error::Error>> {
8626 debug_assert!(ncols % 128 == 0);
8627 let (nc, e2) = (ncols as i32, eps);
8628 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8629 let cfg = LaunchConfig {
8630 grid_dim: (nrows as u32, 1, 1),
8631 block_dim: (rms_block(), 1, 1),
8632 shared_mem_bytes: 0,
8633 };
8634 let __s_b = self.gpu.stream();
8635 let mut b2 = __s_b.launch_builder(&f);
8636 b2.arg(a)
8637 .arg(wa)
8638 .arg(b_in)
8639 .arg(&c)
8640 .arg(w)
8641 .arg(&mut *res)
8642 .arg(&mut *out_q)
8643 .arg(&mut *out_d)
8644 .arg(&nc)
8645 .arg(&e2);
8646 unsafe {
8647 b2.launch(cfg)?;
8648 }
8649 Ok(())
8650 }
8651
8652 pub fn g4_pnfold_on() -> bool {
8660 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8661 *ON.get_or_init(|| {
8662 std::env::var("MEMRA_G4_PNFOLD")
8663 .map(|v| v != "0")
8664 .unwrap_or(true)
8665 })
8666 }
8667
8668 pub fn build_q4_out_concat3(
8672 &self,
8673 w0: &crate::model::GpuTensor,
8674 w1: &crate::model::GpuTensor,
8675 w2: &crate::model::GpuTensor,
8676 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
8677 use crate::model::GpuTensor;
8678 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
8679 match w {
8680 GpuTensor::Quant {
8681 qtype,
8682 row_bytes,
8683 rp,
8684 ..
8685 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
8686 _ => None,
8687 }
8688 };
8689 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
8690 else {
8691 return Ok(None);
8692 };
8693 if rb0 != rb1
8694 || rb0 != rb2
8695 || w0.in_features() != w1.in_features()
8696 || w0.in_features() != w2.in_features()
8697 {
8698 return Ok(None);
8699 }
8700 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
8701 match w {
8702 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
8703 _ => unreachable!(),
8704 }
8705 }
8706 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
8707 let total = rb0 * (o0 + o1 + o2);
8708 let mut cat = self.alloc_u8(total)?;
8709 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
8710 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
8711 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
8712 Ok(Some(GpuTensor::Quant {
8713 bytes: cat,
8714 qtype: QT_Q4_0,
8715 row_bytes: rb0,
8716 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
8717 scale: 1.0,
8718 rp: false,
8719 #[cfg(memra_cutlass)]
8720 cutlass: None,
8721 fp8: None,
8722 blk: None,
8723 rp4: None,
8724 f16: None,
8725 }))
8726 }
8727
8728 #[allow(clippy::too_many_arguments)]
8730 pub fn rms_norm_qkv_rope_cat(
8731 &self,
8732 qkv: &CudaSlice<f32>,
8733 wq: &CudaSlice<f32>,
8734 wk: &CudaSlice<f32>,
8735 wv: &CudaSlice<f32>,
8736 q: &mut CudaSlice<f32>,
8737 k: &mut CudaSlice<f32>,
8738 v: &mut CudaSlice<f32>,
8739 head_dim: usize,
8740 rq: usize,
8741 rk: usize,
8742 pos: &CudaSlice<i32>,
8743 nh_q: usize,
8744 nh_k: usize,
8745 base: f32,
8746 freq_scale: f32,
8747 ff: Option<&CudaSlice<f32>>,
8748 eps: f32,
8749 ) -> Result<(), Box<dyn std::error::Error>> {
8750 let rows = rq + rk + rk;
8751 let theta_scale = base.powf(-2.0 / head_dim as f32);
8752 let (nc, rqi, rki, nhq, nhk) = (
8753 head_dim as i32,
8754 rq as i32,
8755 rk as i32,
8756 nh_q as i32,
8757 nh_k as i32,
8758 );
8759 if Self::pdl_on() {
8760 use cudarc::driver::{DevicePtr, DevicePtrMut};
8761 let s = &self.gpu.stream();
8762 let (pqkv, _g0) = qkv.device_ptr(s);
8763 let (pwq, _g1) = wq.device_ptr(s);
8764 let (pwk, _g2) = wk.device_ptr(s);
8765 let (pwv, _g3) = wv.device_ptr(s);
8766 let (pq, _g4) = q.device_ptr_mut(s);
8767 let (pk, _g5) = k.device_ptr_mut(s);
8768 let (pv, _g6) = v.device_ptr_mut(s);
8769 let (ppos, _g7) = pos.device_ptr(s);
8770 let (pff, _g8) = match ff {
8771 Some(t) => {
8772 let (p, g) = t.device_ptr(s);
8773 (p, Some(g))
8774 }
8775 None => (0, None),
8776 };
8777 let mut ps = [
8778 &pqkv as *const _ as *mut std::ffi::c_void,
8779 &pwq as *const _ as *mut _,
8780 &pwk as *const _ as *mut _,
8781 &pwv as *const _ as *mut _,
8782 &pq as *const _ as *mut _,
8783 &pk as *const _ as *mut _,
8784 &pv as *const _ as *mut _,
8785 &nc as *const _ as *mut _,
8786 &rqi as *const _ as *mut _,
8787 &rki as *const _ as *mut _,
8788 &ppos as *const _ as *mut _,
8789 &nhq as *const _ as *mut _,
8790 &nhk as *const _ as *mut _,
8791 &theta_scale as *const _ as *mut _,
8792 &freq_scale as *const _ as *mut _,
8793 &pff as *const _ as *mut _,
8794 &eps as *const _ as *mut _,
8795 ];
8796 unsafe {
8797 self.launch_pdl(
8798 "rms_norm_qkv_rope_cat_f32",
8799 (rows as u32, 1, 1),
8800 (rms_block(), 1, 1),
8801 &mut ps,
8802 )?;
8803 }
8804 return Ok(());
8805 }
8806 let f = self.func("rms_norm_qkv_rope_cat_f32");
8807 let cfg = LaunchConfig {
8808 grid_dim: (rows as u32, 1, 1),
8809 block_dim: (rms_block(), 1, 1),
8810 shared_mem_bytes: 0,
8811 };
8812 let __s_b = self.gpu.stream();
8813 let mut b = __s_b.launch_builder(&f);
8814 match ff {
8815 Some(t) => {
8816 b.arg(qkv)
8817 .arg(wq)
8818 .arg(wk)
8819 .arg(wv)
8820 .arg(&mut *q)
8821 .arg(&mut *k)
8822 .arg(&mut *v)
8823 .arg(&nc)
8824 .arg(&rqi)
8825 .arg(&rki)
8826 .arg(pos)
8827 .arg(&nhq)
8828 .arg(&nhk)
8829 .arg(&theta_scale)
8830 .arg(&freq_scale)
8831 .arg(t)
8832 .arg(&eps);
8833 unsafe {
8834 b.launch(cfg)?;
8835 }
8836 }
8837 None => {
8838 let null: u64 = 0;
8839 b.arg(qkv)
8840 .arg(wq)
8841 .arg(wk)
8842 .arg(wv)
8843 .arg(&mut *q)
8844 .arg(&mut *k)
8845 .arg(&mut *v)
8846 .arg(&nc)
8847 .arg(&rqi)
8848 .arg(&rki)
8849 .arg(pos)
8850 .arg(&nhq)
8851 .arg(&nhk)
8852 .arg(&theta_scale)
8853 .arg(&freq_scale)
8854 .arg(&null)
8855 .arg(&eps);
8856 unsafe {
8857 b.launch(cfg)?;
8858 }
8859 }
8860 }
8861 Ok(())
8862 }
8863
8864 #[allow(clippy::too_many_arguments)]
8866 pub fn rms_norm_qkv_rope(
8867 &self,
8868 q0: &CudaSlice<f32>,
8869 k0: &CudaSlice<f32>,
8870 v0: &CudaSlice<f32>,
8871 wq: &CudaSlice<f32>,
8872 wk: &CudaSlice<f32>,
8873 wv: &CudaSlice<f32>,
8874 q: &mut CudaSlice<f32>,
8875 k: &mut CudaSlice<f32>,
8876 v: &mut CudaSlice<f32>,
8877 head_dim: usize,
8878 rq: usize,
8879 rk: usize,
8880 pos: &CudaSlice<i32>,
8881 nh_q: usize,
8882 nh_k: usize,
8883 base: f32,
8884 freq_scale: f32,
8885 ff: Option<&CudaSlice<f32>>,
8886 eps: f32,
8887 ) -> Result<(), Box<dyn std::error::Error>> {
8888 let f = self.func("rms_norm_qkv_rope_f32");
8889 let rows = rq + rk + rk; let cfg = LaunchConfig {
8891 grid_dim: (rows as u32, 1, 1),
8892 block_dim: (rms_block(), 1, 1),
8893 shared_mem_bytes: 0,
8894 };
8895 let theta_scale = base.powf(-2.0 / head_dim as f32);
8896 let (nc, rqi, rki, nhq, nhk) = (
8897 head_dim as i32,
8898 rq as i32,
8899 rk as i32,
8900 nh_q as i32,
8901 nh_k as i32,
8902 );
8903 let __s_b = self.gpu.stream();
8904 let mut b = __s_b.launch_builder(&f);
8905 match ff {
8906 Some(t) => {
8907 b.arg(q0)
8908 .arg(k0)
8909 .arg(v0)
8910 .arg(wq)
8911 .arg(wk)
8912 .arg(wv)
8913 .arg(&mut *q)
8914 .arg(&mut *k)
8915 .arg(&mut *v)
8916 .arg(&nc)
8917 .arg(&rqi)
8918 .arg(&rki)
8919 .arg(pos)
8920 .arg(&nhq)
8921 .arg(&nhk)
8922 .arg(&theta_scale)
8923 .arg(&freq_scale)
8924 .arg(t)
8925 .arg(&eps);
8926 unsafe {
8927 b.launch(cfg)?;
8928 }
8929 }
8930 None => {
8931 let null: u64 = 0;
8932 b.arg(q0)
8933 .arg(k0)
8934 .arg(v0)
8935 .arg(wq)
8936 .arg(wk)
8937 .arg(wv)
8938 .arg(&mut *q)
8939 .arg(&mut *k)
8940 .arg(&mut *v)
8941 .arg(&nc)
8942 .arg(&rqi)
8943 .arg(&rki)
8944 .arg(pos)
8945 .arg(&nhq)
8946 .arg(&nhk)
8947 .arg(&theta_scale)
8948 .arg(&freq_scale)
8949 .arg(&null)
8950 .arg(&eps);
8951 unsafe {
8952 b.launch(cfg)?;
8953 }
8954 }
8955 }
8956 Ok(())
8957 }
8958
8959 #[allow(clippy::too_many_arguments)]
8963 pub fn rms_norm_qkv_rope_append_dc(
8964 &self,
8965 q0: &CudaSlice<f32>,
8966 k0: &CudaSlice<f32>,
8967 v0: &CudaSlice<f32>,
8968 wq: &CudaSlice<f32>,
8969 wk: &CudaSlice<f32>,
8970 wv: &CudaSlice<f32>,
8971 q: &mut CudaSlice<f32>,
8972 k: &mut CudaSlice<f32>,
8973 v: &mut CudaSlice<f32>,
8974 head_dim: usize,
8975 rq: usize,
8976 rk: usize,
8977 pos: &CudaSlice<i32>,
8978 nh_q: usize,
8979 nh_k: usize,
8980 base: f32,
8981 freq_scale: f32,
8982 ff: Option<&CudaSlice<f32>>,
8983 eps: f32,
8984 kc: &mut CudaSlice<u8>,
8985 vc: &mut CudaSlice<u8>,
8986 t_dev: &CudaSlice<i32>,
8987 k_tok_bytes: usize,
8988 v_tok_bytes: usize,
8989 g: bool,
8990 ) -> Result<(), Box<dyn std::error::Error>> {
8991 let rows = rq + rk + rk;
8992 let theta_scale = base.powf(-2.0 / head_dim as f32);
8993 let (nc, rqi, rki, nhq, nhk) = (
8994 head_dim as i32,
8995 rq as i32,
8996 rk as i32,
8997 nh_q as i32,
8998 nh_k as i32,
8999 );
9000 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9001 if Self::pdl_on() && Self::pdl_wb_on() {
9002 use cudarc::driver::{DevicePtr, DevicePtrMut};
9003 let s = &self.gpu.stream();
9004 let (p0, _a0) = q0.device_ptr(s);
9005 let (p1, _a1) = k0.device_ptr(s);
9006 let (p2, _a2) = v0.device_ptr(s);
9007 let (pwq, _a3) = wq.device_ptr(s);
9008 let (pwk, _a4) = wk.device_ptr(s);
9009 let (pwv, _a5) = wv.device_ptr(s);
9010 let (pq, _a6) = q.device_ptr_mut(s);
9011 let (pk, _a7) = k.device_ptr_mut(s);
9012 let (pv, _a8) = v.device_ptr_mut(s);
9013 let (pp, _a9) = pos.device_ptr(s);
9014 let pff: u64 = match ff {
9015 Some(t) => {
9016 let (p, _gg) = t.device_ptr(s);
9017 p as u64
9018 }
9019 None => 0,
9020 };
9021 let (pkc, _a10) = kc.device_ptr_mut(s);
9022 let (pvc, _a11) = vc.device_ptr_mut(s);
9023 let (pt, _a12) = t_dev.device_ptr(s);
9024 let mut ps = [
9025 &p0 as *const _ as *mut std::ffi::c_void,
9026 &p1 as *const _ as *mut _,
9027 &p2 as *const _ as *mut _,
9028 &pwq as *const _ as *mut _,
9029 &pwk as *const _ as *mut _,
9030 &pwv as *const _ as *mut _,
9031 &pq as *const _ as *mut _,
9032 &pk as *const _ as *mut _,
9033 &pv as *const _ as *mut _,
9034 &nc as *const _ as *mut _,
9035 &rqi as *const _ as *mut _,
9036 &rki as *const _ as *mut _,
9037 &pp as *const _ as *mut _,
9038 &nhq as *const _ as *mut _,
9039 &nhk as *const _ as *mut _,
9040 &theta_scale as *const _ as *mut _,
9041 &freq_scale as *const _ as *mut _,
9042 &pff as *const _ as *mut _,
9043 &eps as *const _ as *mut _,
9044 &pkc as *const _ as *mut _,
9045 &pvc as *const _ as *mut _,
9046 &pt as *const _ as *mut _,
9047 &ktb as *const _ as *mut _,
9048 &vtb as *const _ as *mut _,
9049 ];
9050 unsafe {
9051 self.launch_pdl_flash(
9052 g,
9053 "rms_norm_qkv_rope_append_dc_f32",
9054 (rows as u32, 1, 1),
9055 (rms_block(), 1, 1),
9056 0,
9057 &mut ps,
9058 )?;
9059 }
9060 return Ok(());
9061 }
9062 let f = if g {
9063 self.func_g("rms_norm_qkv_rope_append_dc_f32")
9064 } else {
9065 self.func("rms_norm_qkv_rope_append_dc_f32")
9066 };
9067 let cfg = LaunchConfig {
9068 grid_dim: (rows as u32, 1, 1),
9069 block_dim: (rms_block(), 1, 1),
9070 shared_mem_bytes: 0,
9071 };
9072 let __s_b = self.gpu.stream();
9073 let mut b = __s_b.launch_builder(&f);
9074 match ff {
9075 Some(t) => {
9076 b.arg(q0)
9077 .arg(k0)
9078 .arg(v0)
9079 .arg(wq)
9080 .arg(wk)
9081 .arg(wv)
9082 .arg(&mut *q)
9083 .arg(&mut *k)
9084 .arg(&mut *v)
9085 .arg(&nc)
9086 .arg(&rqi)
9087 .arg(&rki)
9088 .arg(pos)
9089 .arg(&nhq)
9090 .arg(&nhk)
9091 .arg(&theta_scale)
9092 .arg(&freq_scale)
9093 .arg(t)
9094 .arg(&eps)
9095 .arg(&mut *kc)
9096 .arg(&mut *vc)
9097 .arg(t_dev)
9098 .arg(&ktb)
9099 .arg(&vtb);
9100 unsafe {
9101 b.launch(cfg)?;
9102 }
9103 }
9104 None => {
9105 let null: u64 = 0;
9106 b.arg(q0)
9107 .arg(k0)
9108 .arg(v0)
9109 .arg(wq)
9110 .arg(wk)
9111 .arg(wv)
9112 .arg(&mut *q)
9113 .arg(&mut *k)
9114 .arg(&mut *v)
9115 .arg(&nc)
9116 .arg(&rqi)
9117 .arg(&rki)
9118 .arg(pos)
9119 .arg(&nhq)
9120 .arg(&nhk)
9121 .arg(&theta_scale)
9122 .arg(&freq_scale)
9123 .arg(&null)
9124 .arg(&eps)
9125 .arg(&mut *kc)
9126 .arg(&mut *vc)
9127 .arg(t_dev)
9128 .arg(&ktb)
9129 .arg(&vtb);
9130 unsafe {
9131 b.launch(cfg)?;
9132 }
9133 }
9134 }
9135 Ok(())
9136 }
9137
9138 #[allow(clippy::too_many_arguments)]
9145 pub fn rms_norm_qkv_rope_append(
9146 &self,
9147 q0: &CudaSlice<f32>,
9148 k0: &CudaSlice<f32>,
9149 v0: &CudaSlice<f32>,
9150 wq: &CudaSlice<f32>,
9151 wk: &CudaSlice<f32>,
9152 wv: &CudaSlice<f32>,
9153 q: &mut CudaSlice<f32>,
9154 k: &mut CudaSlice<f32>,
9155 v: &mut CudaSlice<f32>,
9156 head_dim: usize,
9157 rq: usize,
9158 rk: usize,
9159 pos: &CudaSlice<i32>,
9160 nh_q: usize,
9161 nh_k: usize,
9162 base: f32,
9163 freq_scale: f32,
9164 ff: Option<&CudaSlice<f32>>,
9165 eps: f32,
9166 kc: &mut CudaSlice<u8>,
9167 vc: &mut CudaSlice<u8>,
9168 t: usize,
9169 k_tok_bytes: usize,
9170 v_tok_bytes: usize,
9171 g: bool,
9172 ) -> Result<(), Box<dyn std::error::Error>> {
9173 let rows = rq + rk + rk;
9174 let theta_scale = base.powf(-2.0 / head_dim as f32);
9175 let (nc, rqi, rki, nhq, nhk) = (
9176 head_dim as i32,
9177 rq as i32,
9178 rk as i32,
9179 nh_q as i32,
9180 nh_k as i32,
9181 );
9182 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9183 let ti = t as i32;
9184 if Self::pdl_on() && Self::pdl_wb_on() {
9185 use cudarc::driver::{DevicePtr, DevicePtrMut};
9186 let s = &self.gpu.stream();
9187 let (p0, _a0) = q0.device_ptr(s);
9188 let (p1, _a1) = k0.device_ptr(s);
9189 let (p2, _a2) = v0.device_ptr(s);
9190 let (pwq, _a3) = wq.device_ptr(s);
9191 let (pwk, _a4) = wk.device_ptr(s);
9192 let (pwv, _a5) = wv.device_ptr(s);
9193 let (pq, _a6) = q.device_ptr_mut(s);
9194 let (pk, _a7) = k.device_ptr_mut(s);
9195 let (pv, _a8) = v.device_ptr_mut(s);
9196 let (pp, _a9) = pos.device_ptr(s);
9197 let pff: u64 = match ff {
9198 Some(t) => {
9199 let (p, _gg) = t.device_ptr(s);
9200 p as u64
9201 }
9202 None => 0,
9203 };
9204 let (pkc, _a10) = kc.device_ptr_mut(s);
9205 let (pvc, _a11) = vc.device_ptr_mut(s);
9206 let mut ps = [
9207 &p0 as *const _ as *mut std::ffi::c_void,
9208 &p1 as *const _ as *mut _,
9209 &p2 as *const _ as *mut _,
9210 &pwq as *const _ as *mut _,
9211 &pwk as *const _ as *mut _,
9212 &pwv as *const _ as *mut _,
9213 &pq as *const _ as *mut _,
9214 &pk as *const _ as *mut _,
9215 &pv as *const _ as *mut _,
9216 &nc as *const _ as *mut _,
9217 &rqi as *const _ as *mut _,
9218 &rki as *const _ as *mut _,
9219 &pp as *const _ as *mut _,
9220 &nhq as *const _ as *mut _,
9221 &nhk as *const _ as *mut _,
9222 &theta_scale as *const _ as *mut _,
9223 &freq_scale as *const _ as *mut _,
9224 &pff as *const _ as *mut _,
9225 &eps as *const _ as *mut _,
9226 &pkc as *const _ as *mut _,
9227 &pvc as *const _ as *mut _,
9228 &ti as *const _ as *mut _,
9229 &ktb as *const _ as *mut _,
9230 &vtb as *const _ as *mut _,
9231 ];
9232 unsafe {
9233 self.launch_pdl_flash(
9234 g,
9235 "rms_norm_qkv_rope_append_f32",
9236 (rows as u32, 1, 1),
9237 (rms_block(), 1, 1),
9238 0,
9239 &mut ps,
9240 )?;
9241 }
9242 return Ok(());
9243 }
9244 let f = if g {
9245 self.func_g("rms_norm_qkv_rope_append_f32")
9246 } else {
9247 self.func("rms_norm_qkv_rope_append_f32")
9248 };
9249 let cfg = LaunchConfig {
9250 grid_dim: (rows as u32, 1, 1),
9251 block_dim: (rms_block(), 1, 1),
9252 shared_mem_bytes: 0,
9253 };
9254 let __s_b = self.gpu.stream();
9255 let mut b = __s_b.launch_builder(&f);
9256 let null: u64 = 0;
9257 b.arg(q0)
9258 .arg(k0)
9259 .arg(v0)
9260 .arg(wq)
9261 .arg(wk)
9262 .arg(wv)
9263 .arg(&mut *q)
9264 .arg(&mut *k)
9265 .arg(&mut *v)
9266 .arg(&nc)
9267 .arg(&rqi)
9268 .arg(&rki)
9269 .arg(pos)
9270 .arg(&nhq)
9271 .arg(&nhk)
9272 .arg(&theta_scale)
9273 .arg(&freq_scale);
9274 match ff {
9275 Some(t) => {
9276 b.arg(t);
9277 }
9278 None => {
9279 b.arg(&null);
9280 }
9281 }
9282 b.arg(&eps)
9283 .arg(&mut *kc)
9284 .arg(&mut *vc)
9285 .arg(&ti)
9286 .arg(&ktb)
9287 .arg(&vtb);
9288 unsafe {
9289 b.launch(cfg)?;
9290 }
9291 Ok(())
9292 }
9293
9294 pub fn add_q8_1(
9295 &self,
9296 a: &CudaSlice<f32>,
9297 b: &CudaSlice<f32>,
9298 res: &mut CudaSlice<f32>,
9299 ncols: usize,
9300 nrows: usize,
9301 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9302 debug_assert!(ncols % 128 == 0);
9303 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9304 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9305 let f = self.func("add_q8_1_f32");
9306 let cfg = LaunchConfig {
9307 grid_dim: (nrows as u32, 1, 1),
9308 block_dim: (rms_block(), 1, 1),
9309 shared_mem_bytes: 0,
9310 };
9311 let nc = ncols as i32;
9312 let __s_b2 = self.gpu.stream();
9313 let mut b2 = __s_b2.launch_builder(&f);
9314 b2.arg(a)
9315 .arg(b)
9316 .arg(&mut *res)
9317 .arg(&mut out_q)
9318 .arg(&mut out_d)
9319 .arg(&nc);
9320 unsafe {
9321 b2.launch(cfg)?;
9322 }
9323 Ok((out_q, out_d))
9324 }
9325
9326 pub fn rms_pre_add_q8_1(
9330 &self,
9331 a: &CudaSlice<f32>,
9332 wa: &CudaSlice<f32>,
9333 b: &CudaSlice<f32>,
9334 res: &mut CudaSlice<f32>,
9335 ncols: usize,
9336 nrows: usize,
9337 eps: f32,
9338 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9339 debug_assert!(ncols % 128 == 0);
9340 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9341 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9342 let f = self.func("rms_pre_add_q8_1_f32");
9343 let cfg = LaunchConfig {
9344 grid_dim: (nrows as u32, 1, 1),
9345 block_dim: (rms_block(), 1, 1),
9346 shared_mem_bytes: 0,
9347 };
9348 let (nc, ep) = (ncols as i32, eps);
9349 let __s_b2 = self.gpu.stream();
9350 let mut b2 = __s_b2.launch_builder(&f);
9351 b2.arg(a)
9352 .arg(wa)
9353 .arg(b)
9354 .arg(&mut *res)
9355 .arg(&mut out_q)
9356 .arg(&mut out_d)
9357 .arg(&nc)
9358 .arg(&ep);
9359 unsafe {
9360 b2.launch(cfg)?;
9361 }
9362 Ok((out_q, out_d))
9363 }
9364
9365 pub fn l2_v2_on(ncols: usize) -> bool {
9369 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9370 }
9371
9372 pub fn l2_norm_pp(
9373 &self,
9374 x: &CudaSlice<f32>,
9375 dst: &mut CudaSlice<f32>,
9376 dst16: Option<&mut CudaSlice<u8>>,
9377 ncols: usize,
9378 nrows: usize,
9379 eps: f32,
9380 ) -> Result<(), Box<dyn std::error::Error>> {
9381 if Self::l2_v2_on(ncols) {
9382 let f = self.func("l2_norm_pp_v2_f32");
9383 let rows_per_block = 8u32; let cfg = LaunchConfig {
9385 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9386 block_dim: (256, 1, 1),
9387 shared_mem_bytes: 0,
9388 };
9389 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9390 let d16: u64 = match dst16 {
9392 Some(d) => self.addr_u8(d),
9393 None => 0,
9394 };
9395 let __s_b = self.gpu.stream();
9396 let mut b = __s_b.launch_builder(&f);
9397 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9398 unsafe {
9399 b.launch(cfg)?;
9400 }
9401 return Ok(());
9402 }
9403 self.l2_norm(x, dst, ncols, nrows, eps)
9404 }
9405
9406 pub fn l2_norm(
9407 &self,
9408 x: &CudaSlice<f32>,
9409 dst: &mut CudaSlice<f32>,
9410 ncols: usize,
9411 nrows: usize,
9412 eps: f32,
9413 ) -> Result<(), Box<dyn std::error::Error>> {
9414 let f = self.func("l2_norm_f32");
9415 let cfg = LaunchConfig {
9416 grid_dim: (nrows as u32, 1, 1),
9417 block_dim: (256, 1, 1),
9418 shared_mem_bytes: 0,
9419 };
9420 let (nc, e) = (ncols as i32, eps);
9421 let __s_b = self.gpu.stream();
9422 let mut b = __s_b.launch_builder(&f);
9423 b.arg(x).arg(dst).arg(&nc).arg(&e);
9424 unsafe {
9425 b.launch(cfg)?;
9426 }
9427 Ok(())
9428 }
9429
9430 pub fn l2_norm_decode(
9436 &self,
9437 x: &CudaSlice<f32>,
9438 dst: &mut CudaSlice<f32>,
9439 ncols: usize,
9440 nrows: usize,
9441 eps: f32,
9442 ) -> Result<(), Box<dyn std::error::Error>> {
9443 let f = self.func("l2_norm_f32");
9444 let cfg = LaunchConfig {
9445 grid_dim: (nrows as u32, 1, 1),
9446 block_dim: (32, 1, 1),
9447 shared_mem_bytes: 0,
9448 };
9449 let (nc, e) = (ncols as i32, eps);
9450 let __s_b = self.gpu.stream();
9451 let mut b = __s_b.launch_builder(&f);
9452 b.arg(x).arg(dst).arg(&nc).arg(&e);
9453 unsafe {
9454 b.launch(cfg)?;
9455 }
9456 Ok(())
9457 }
9458
9459 pub fn rope_neox(
9461 &self,
9462 x: &mut CudaSlice<f32>,
9463 pos: &CudaSlice<i32>,
9464 head_dim: usize,
9465 n_dims: usize,
9466 n_heads: usize,
9467 n_tokens: usize,
9468 freq_base: f32,
9469 freq_scale: f32,
9470 ) -> Result<(), Box<dyn std::error::Error>> {
9471 let f = self.func("rope_neox_f32");
9472 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9473 let grid = (n_heads * n_tokens) as u32;
9474 let cfg = LaunchConfig {
9475 grid_dim: (grid, 1, 1),
9476 block_dim: ((head_dim / 2) as u32, 1, 1),
9477 shared_mem_bytes: 0,
9478 };
9479 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9480 let __s_b = self.gpu.stream();
9481 let mut b = __s_b.launch_builder(&f);
9482 b.arg(x)
9483 .arg(pos)
9484 .arg(&hd)
9485 .arg(&nd)
9486 .arg(&nh)
9487 .arg(&theta_scale)
9488 .arg(&freq_scale);
9489 unsafe {
9490 b.launch(cfg)?;
9491 }
9492 Ok(())
9493 }
9494
9495 pub fn rope_neox_ff(
9497 &self,
9498 x: &mut CudaSlice<f32>,
9499 pos: &CudaSlice<i32>,
9500 head_dim: usize,
9501 n_dims: usize,
9502 n_heads: usize,
9503 n_tokens: usize,
9504 freq_base: f32,
9505 freq_scale: f32,
9506 ff: &CudaSlice<f32>,
9507 ) -> Result<(), Box<dyn std::error::Error>> {
9508 let f = self.func("rope_neox_ff_f32");
9509 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9510 let grid = (n_heads * n_tokens) as u32;
9511 let cfg = LaunchConfig {
9512 grid_dim: (grid, 1, 1),
9513 block_dim: ((head_dim / 2) as u32, 1, 1),
9514 shared_mem_bytes: 0,
9515 };
9516 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9517 let __s_b = self.gpu.stream();
9518 let mut b = __s_b.launch_builder(&f);
9519 b.arg(x)
9520 .arg(pos)
9521 .arg(&hd)
9522 .arg(&nd)
9523 .arg(&nh)
9524 .arg(&theta_scale)
9525 .arg(&freq_scale)
9526 .arg(ff);
9527 unsafe {
9528 b.launch(cfg)?;
9529 }
9530 Ok(())
9531 }
9532
9533 #[allow(clippy::too_many_arguments)]
9535 pub fn rope_neox2(
9536 &self,
9537 q: &mut CudaSlice<f32>,
9538 k: &mut CudaSlice<f32>,
9539 pos: &CudaSlice<i32>,
9540 head_dim: usize,
9541 n_dims: usize,
9542 nh_q: usize,
9543 nh_k: usize,
9544 n_tokens: usize,
9545 freq_base: f32,
9546 freq_scale: f32,
9547 ff: Option<&CudaSlice<f32>>,
9548 ) -> Result<(), Box<dyn std::error::Error>> {
9549 let f = self.func("rope_neox2_f32");
9550 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9551 let grid = ((nh_q + nh_k) * n_tokens) as u32;
9552 let cfg = LaunchConfig {
9553 grid_dim: (grid, 1, 1),
9554 block_dim: ((head_dim / 2) as u32, 1, 1),
9555 shared_mem_bytes: 0,
9556 };
9557 let (hd, nd, nq, nk, nt) = (
9558 head_dim as i32,
9559 n_dims as i32,
9560 nh_q as i32,
9561 nh_k as i32,
9562 n_tokens as i32,
9563 );
9564 let __s_b = self.gpu.stream();
9565 let mut b = __s_b.launch_builder(&f);
9566 b.arg(q)
9567 .arg(k)
9568 .arg(pos)
9569 .arg(&hd)
9570 .arg(&nd)
9571 .arg(&nq)
9572 .arg(&nk)
9573 .arg(&nt)
9574 .arg(&theta_scale)
9575 .arg(&freq_scale);
9576 match ff {
9577 Some(ffv) => {
9578 b.arg(ffv);
9579 unsafe {
9580 b.launch(cfg)?;
9581 }
9582 }
9583 None => {
9584 let null: u64 = 0;
9585 b.arg(&null);
9586 unsafe {
9587 b.launch(cfg)?;
9588 }
9589 }
9590 }
9591 Ok(())
9592 }
9593
9594 pub fn gelu_tanh_mul(
9596 &self,
9597 gate: &CudaSlice<f32>,
9598 up: &CudaSlice<f32>,
9599 dst: &mut CudaSlice<f32>,
9600 n: usize,
9601 ) -> Result<(), Box<dyn std::error::Error>> {
9602 let f = self.func("gelu_tanh_mul_f32");
9603 let cfg = LaunchConfig::for_num_elems(n as u32);
9604 let ni = n as i32;
9605 let __s_b = self.gpu.stream();
9606 let mut b = __s_b.launch_builder(&f);
9607 b.arg(gate).arg(up).arg(dst).arg(&ni);
9608 unsafe {
9609 b.launch(cfg)?;
9610 }
9611 Ok(())
9612 }
9613
9614 pub fn silu_mul(
9615 &self,
9616 gate: &CudaSlice<f32>,
9617 up: &CudaSlice<f32>,
9618 dst: &mut CudaSlice<f32>,
9619 n: usize,
9620 ) -> Result<(), Box<dyn std::error::Error>> {
9621 let f = self.func("silu_mul_f32");
9622 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9624 let ni = n as i32;
9625 let __s_b = self.gpu.stream();
9626 let mut b = __s_b.launch_builder(&f);
9627 b.arg(gate).arg(up).arg(dst).arg(&ni);
9628 unsafe {
9629 b.launch(cfg)?;
9630 }
9631 Ok(())
9632 }
9633
9634 pub fn silu_mul_f16out(
9637 &self,
9638 gate: &CudaSlice<f32>,
9639 up: &CudaSlice<f32>,
9640 dst: &mut CudaSlice<f32>,
9641 dst16: &mut CudaSlice<u8>,
9642 n: usize,
9643 ) -> Result<(), Box<dyn std::error::Error>> {
9644 let f = self.func("silu_mul_f16out_f32");
9645 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9646 let ni = n as i32;
9647 let __s_b = self.gpu.stream();
9648 let mut b = __s_b.launch_builder(&f);
9649 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
9650 unsafe {
9651 b.launch(cfg)?;
9652 }
9653 Ok(())
9654 }
9655
9656 pub fn silu_mul_scaled(
9663 &self,
9664 gate: &CudaSlice<f32>,
9665 up: &CudaSlice<f32>,
9666 gs: f32,
9667 us: f32,
9668 dst: &mut CudaSlice<f32>,
9669 n: usize,
9670 ) -> Result<(), Box<dyn std::error::Error>> {
9671 let f = self.func("silu_mul_scaled_f32");
9672 let cfg = LaunchConfig::for_num_elems(n as u32);
9673 let ni = n as i32;
9674 let (gsf, usf) = (gs, us);
9675 let __s_b = self.gpu.stream();
9676 let mut b = __s_b.launch_builder(&f);
9677 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
9678 unsafe {
9679 b.launch(cfg)?;
9680 }
9681 Ok(())
9682 }
9683
9684 #[allow(clippy::too_many_arguments)]
9688 pub fn swigluoai_mul_scaled(
9689 &self,
9690 gate: &CudaSlice<f32>,
9691 up: &CudaSlice<f32>,
9692 gs: f32,
9693 us: f32,
9694 alpha: f32,
9695 limit: f32,
9696 dst: &mut CudaSlice<f32>,
9697 n: usize,
9698 ) -> Result<(), Box<dyn std::error::Error>> {
9699 let f = self.func("swigluoai_mul_scaled_f32");
9700 let cfg = LaunchConfig::for_num_elems(n as u32);
9701 let ni = n as i32;
9702 let __s_b = self.gpu.stream();
9703 let mut b = __s_b.launch_builder(&f);
9704 b.arg(gate)
9705 .arg(up)
9706 .arg(&gs)
9707 .arg(&us)
9708 .arg(&alpha)
9709 .arg(&limit)
9710 .arg(dst)
9711 .arg(&ni);
9712 unsafe {
9713 b.launch(cfg)?;
9714 }
9715 Ok(())
9716 }
9717
9718 pub fn silu_mul_scaled_q8_1(
9726 &self,
9727 gate: &CudaSlice<f32>,
9728 up: &CudaSlice<f32>,
9729 gs: f32,
9730 us: f32,
9731 n: usize,
9732 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9733 let f = self.func("silu_mul_scaled_q8_1");
9734 let nblk = n / 32;
9735 let mut aq = self.alloc_uninit::<i8>(n)?; let mut ad = self.alloc_uninit::<f32>(nblk)?; let cfg = LaunchConfig::for_num_elems(n as u32);
9739 let (gsf, usf, ni) = (gs, us, n as i32);
9740 let __s_b = self.gpu.stream();
9741 let mut b = __s_b.launch_builder(&f);
9742 b.arg(gate)
9743 .arg(up)
9744 .arg(&gsf)
9745 .arg(&usf)
9746 .arg(&mut aq)
9747 .arg(&mut ad)
9748 .arg(&ni);
9749 unsafe {
9750 b.launch(cfg)?;
9751 }
9752 Ok((aq, ad))
9753 }
9754
9755 pub fn add(
9756 &self,
9757 a: &CudaSlice<f32>,
9758 b_in: &CudaSlice<f32>,
9759 dst: &mut CudaSlice<f32>,
9760 n: usize,
9761 ) -> Result<(), Box<dyn std::error::Error>> {
9762 let f = self.func("add_f32");
9763 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9765 let ni = n as i32;
9766 let __s_bld = self.gpu.stream();
9767 let mut bld = __s_bld.launch_builder(&f);
9768 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9769 unsafe {
9770 bld.launch(cfg)?;
9771 }
9772 Ok(())
9773 }
9774
9775 pub fn mul(
9776 &self,
9777 a: &CudaSlice<f32>,
9778 b_in: &CudaSlice<f32>,
9779 dst: &mut CudaSlice<f32>,
9780 n: usize,
9781 ) -> Result<(), Box<dyn std::error::Error>> {
9782 let f = self.func("mul_f32");
9783 let cfg = LaunchConfig::for_num_elems(n as u32);
9784 let ni = n as i32;
9785 let __s_bld = self.gpu.stream();
9786 let mut bld = __s_bld.launch_builder(&f);
9787 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9788 unsafe {
9789 bld.launch(cfg)?;
9790 }
9791 Ok(())
9792 }
9793
9794 pub fn matmul(
9797 &self,
9798 w: &crate::model::GpuTensor,
9799 x: &CudaSlice<f32>,
9800 m: usize,
9801 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9802 use crate::model::GpuTensor;
9803 let in_f = w.in_features();
9804 let out_f = w.out_features();
9805 #[allow(non_snake_case)]
9813 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
9816 usize::MAX
9817 } else {
9818 16usize
9819 };
9820
9821 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
9846 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
9847 return Ok(y);
9848 }
9849 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
9856 return Ok(y);
9857 }
9858 if let Some(y) = self.try_f16_gemm(w, x, m)? {
9861 return Ok(y);
9862 }
9863 }
9864 if let GpuTensor::Quant { qtype, .. } = w {
9879 if *qtype == QT_F8_E4M3_BLK {
9880 if m >= GEMM_M_THRESHOLD {
9881 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
9882 return Ok(y);
9883 }
9884 }
9885 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9886 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
9887 return Ok(y);
9888 }
9889 }
9890 }
9891 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
9892 return self.qmatvec_mmq(w, x, m);
9893 }
9894 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
9895 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9896 return self.qmatvec_gemm(w, &aq, &ad, m);
9897 }
9898 if m >= GEMM_M_THRESHOLD {
9901 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
9902 return Ok(y);
9903 }
9904 }
9905 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
9909 if m == 1 && fast {
9914 if let GpuTensor::Quant {
9915 bytes,
9916 qtype,
9917 row_bytes,
9918 rp,
9919 rp4,
9920 scale,
9921 ..
9922 } = w
9923 {
9924 if self.mmvq_supports(*qtype) {
9925 let (bytes, rp) = match rp4 {
9929 Some(m4) => (m4, true),
9930 None => (bytes, *rp),
9931 };
9932 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9933 return self.qmatvec_mmvq(
9934 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
9935 );
9936 }
9937 }
9938 }
9939 if (2..=16).contains(&m)
9955 && fast
9956 && std::env::var("MEMRA_NO_BATCHED").is_err()
9957 && (m <= 4 || Self::b8_enabled())
9958 {
9959 let m_ok = m <= 8
9969 || matches!(w, GpuTensor::Quant { qtype, .. }
9970 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
9971 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
9972 if m_ok {
9973 if let GpuTensor::Quant {
9974 bytes,
9975 qtype,
9976 row_bytes,
9977 rp,
9978 rp4,
9979 ..
9980 } = w
9981 {
9982 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
9983 let (bytes, rp) = match rp4 {
9984 Some(m4) => (m4, true),
9985 None => (bytes, *rp),
9986 };
9987 let mcols = Self::batched_mcols(m);
9988 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
9989 let mut y = self.qmatvec_mmvq_batched(
9990 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
9991 )?;
9992 if let GpuTensor::Quant { scale, .. } = w {
9993 if *scale != 1.0 {
9994 self.scale_inplace(&mut y, *scale, m * out_f)?;
9995 }
9996 }
9997 return Ok(y);
9998 }
9999 }
10000 }
10001 }
10002 if fast {
10008 if let GpuTensor::Quant {
10009 bytes,
10010 qtype,
10011 row_bytes,
10012 scale,
10013 ..
10014 } = w
10015 {
10016 if *qtype == QT_F8_E4M3 {
10017 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10018 return self.qmatvec_mmvq(
10019 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10020 );
10021 }
10022 }
10023 }
10024 let mut y = match w {
10025 GpuTensor::Quant {
10026 bytes,
10027 qtype,
10028 row_bytes,
10029 ..
10030 } if fast && *qtype == QT_Q8_0 => {
10031 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10032 }
10033 GpuTensor::Quant {
10034 bytes,
10035 qtype,
10036 row_bytes,
10037 ..
10038 } if fast && *qtype == QT_Q4_K => {
10039 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10040 }
10041 GpuTensor::Quant {
10042 bytes,
10043 qtype,
10044 row_bytes,
10045 ..
10046 } if fast && *qtype == QT_Q6_K => {
10047 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10048 }
10049 GpuTensor::Quant {
10050 bytes,
10051 qtype,
10052 row_bytes,
10053 ..
10054 } if fast && *qtype == QT_Q5_K => {
10055 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10056 }
10057 GpuTensor::Quant {
10058 bytes,
10059 qtype,
10060 row_bytes,
10061 ..
10062 } if fast && *qtype == QT_Q3_K => {
10063 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10064 }
10065 GpuTensor::Quant {
10066 bytes,
10067 qtype,
10068 row_bytes,
10069 rp,
10070 ..
10071 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10072 if *rp {
10073 "qmatvec_nvfp4_dp4a_rp"
10074 } else {
10075 "qmatvec_nvfp4_dp4a"
10076 },
10077 bytes,
10078 x,
10079 m,
10080 in_f,
10081 out_f,
10082 *row_bytes,
10083 )?,
10084 GpuTensor::Quant {
10088 bytes,
10089 qtype,
10090 row_bytes,
10091 ..
10092 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10093 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10094 }
10095 GpuTensor::Quant {
10100 bytes,
10101 qtype,
10102 row_bytes,
10103 rp,
10104 ..
10105 } =>
10106 {
10109 self.qmatvec(
10110 bytes,
10111 x,
10112 m,
10113 in_f,
10114 out_f,
10115 if *rp && *qtype == QT_NVFP4 {
10116 QT_NVFP4_RP
10117 } else {
10118 *qtype
10119 },
10120 *row_bytes,
10121 )?
10122 }
10123 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10124 GpuTensor::FloatBf16 { data, .. } => {
10127 self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10128 }
10129 };
10130 if let GpuTensor::Quant { scale, .. } = w {
10132 if *scale != 1.0 {
10133 self.scale_inplace(&mut y, *scale, m * out_f)?;
10134 }
10135 }
10136 Ok(y)
10137 }
10138
10139 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10142 use crate::model::GpuTensor;
10143 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10144 return false;
10145 }
10146 match w {
10147 GpuTensor::Quant { qtype, .. } => {
10154 matches!(
10155 *qtype,
10156 QT_Q8_0
10157 | QT_Q4_K
10158 | QT_Q6_K
10159 | QT_Q5_K
10160 | QT_Q3_K
10161 | QT_NVFP4
10162 | QT_F8_E4M3
10163 | QT_F8_E4M3_BLK
10164 | QT_Q4_0
10165 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10166 }
10167 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10168 }
10169 }
10170
10171 pub fn matmul_pre(
10176 &self,
10177 w: &crate::model::GpuTensor,
10178 aq: &CudaSlice<i8>,
10179 ad: &CudaSlice<f32>,
10180 x_fallback: &CudaSlice<f32>,
10181 m: usize,
10182 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10183 use crate::model::GpuTensor;
10184 let x_raw_ok = x_fallback.len() >= m * w.in_features();
10190 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10193 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10194 return Ok(y);
10195 }
10196 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10199 return Ok(y);
10200 }
10201 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10203 return Ok(y);
10204 }
10205 }
10206 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10212 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10213 return Ok(y);
10214 }
10215 }
10216 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10217 return Ok(y);
10218 }
10219 if m >= 16
10224 && w.out_features() >= 128
10225 && self.mmq_supports(w)
10226 && !self.verify_exact_on()
10227 && x_raw_ok
10228 {
10229 return self.qmatvec_mmq(w, x_fallback, m);
10230 }
10231 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10234 if let Some(y) =
10235 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10236 {
10237 return Ok(y);
10238 }
10239 }
10240 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10243 return self.qmatvec_gemm(w, aq, ad, m);
10244 }
10245 if !self.uses_q8_1_fast(w) {
10246 return self.matmul(w, x_fallback, m);
10247 }
10248 let in_f = w.in_features();
10249 let out_f = w.out_features();
10250 let (bytes, qtype, row_bytes, scale, rp) = match w {
10251 GpuTensor::Quant {
10252 bytes,
10253 qtype,
10254 row_bytes,
10255 scale,
10256 rp,
10257 ..
10258 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10259 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10260 };
10261 let (mbytes, mrp) = match w {
10264 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10265 _ => (bytes, rp),
10266 };
10267 if m == 1 && self.mmvq_supports(qtype) {
10271 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10272 }
10273 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10286 && std::env::var("MEMRA_NO_BATCHED").is_err()
10287 && (m <= 4 || Self::b8_enabled())
10288 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10292 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10293 {
10294 let mcols = Self::batched_mcols(m);
10295 return self.qmatvec_mmvq_batched(
10296 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10297 );
10298 }
10299 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10305 let (b2, r2) = if qtype == QT_Q4_0 {
10306 (mbytes, mrp)
10307 } else {
10308 (bytes, rp)
10309 };
10310 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10311 }
10312 let name = match qtype {
10313 QT_Q8_0 => "qmatvec_q8_0_dp4a",
10314 QT_Q4_K => "qmatvec_q4_K_dp4a",
10315 QT_Q6_K => "qmatvec_q6_K_dp4a",
10316 QT_Q5_K => "qmatvec_q5_K_dp4a",
10317 QT_Q3_K => "qmatvec_q3_K_dp4a",
10318 QT_NVFP4 => {
10319 if rp {
10320 "qmatvec_nvfp4_dp4a_rp"
10321 } else {
10322 "qmatvec_nvfp4_dp4a"
10323 }
10324 }
10325 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10326 _ => unreachable!(),
10327 };
10328 let f = self.func(name);
10329 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10331 grid_dim: (out_f as u32, m as u32, 1),
10332 block_dim: (128, 1, 1),
10333 shared_mem_bytes: 0,
10334 };
10335 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10336 let __s_b = self.gpu.stream();
10337 let mut b = __s_b.launch_builder(&f);
10338 b.arg(bytes)
10339 .arg(aq)
10340 .arg(ad)
10341 .arg(&mut y)
10342 .arg(&inf)
10343 .arg(&outf)
10344 .arg(&mi)
10345 .arg(&rb);
10346 unsafe {
10347 b.launch(cfg)?;
10348 }
10349 if scale != 1.0 {
10350 self.scale_inplace(&mut y, scale, m * out_f)?;
10351 }
10352 Ok(y)
10353 }
10354
10355 pub fn matmul_decode_exact(
10363 &self,
10364 w: &crate::model::GpuTensor,
10365 x: &CudaSlice<f32>,
10366 m: usize,
10367 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10368 use crate::model::GpuTensor;
10369 if let GpuTensor::Float { data, .. } = w {
10377 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10378 }
10379 if let GpuTensor::FloatBf16 { data, .. } = w {
10382 let (in_f, out_f) = (w.in_features(), w.out_features());
10383 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10384 }
10385 if !self.uses_q8_1_fast(w) {
10386 return self.matmul(w, x, m);
10387 }
10388 let in_f = w.in_features();
10389 let out_f = w.out_features();
10390 let (bytes, qtype, row_bytes, scale, rp) = match w {
10391 GpuTensor::Quant {
10392 bytes,
10393 qtype,
10394 row_bytes,
10395 scale,
10396 rp,
10397 ..
10398 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10399 _ => return self.matmul(w, x, m),
10400 };
10401 let (bytes, rp) = match w {
10404 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10405 _ => (bytes, rp),
10406 };
10407 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10408 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10412 return Ok(y);
10413 }
10414 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10423 && std::env::var("MEMRA_NO_BATCHED").is_err()
10424 && (m <= 4 || Self::b8_enabled())
10425 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10428 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10429 {
10430 let mcols = Self::batched_mcols(m);
10431 return self.qmatvec_mmvq_batched(
10432 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10433 );
10434 }
10435 if self.mmvq_supports(qtype) {
10436 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10439 }
10440 self.matmul_pre(w, &aq, &ad, x, m)
10443 }
10444
10445 pub fn matmul_decode_exact_pre(
10455 &self,
10456 w: &crate::model::GpuTensor,
10457 aq: &CudaSlice<i8>,
10458 ad: &CudaSlice<f32>,
10459 m: usize,
10460 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10461 use crate::model::GpuTensor;
10462 debug_assert!(
10463 self.uses_q8_1_fast(w),
10464 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10465 );
10466 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10468 return Ok(y);
10469 }
10470 let in_f = w.in_features();
10471 let out_f = w.out_features();
10472 let (bytes, qtype, row_bytes, scale, rp) = match w {
10473 GpuTensor::Quant {
10474 bytes,
10475 qtype,
10476 row_bytes,
10477 scale,
10478 rp,
10479 ..
10480 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10481 _ => {
10482 return Err(
10483 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10484 );
10485 }
10486 };
10487 let (bytes, rp) = match w {
10489 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10490 _ => (bytes, rp),
10491 };
10492 if (2..=16).contains(&m)
10494 && self.batched_supports(qtype)
10495 && self.mmvq_supports(qtype)
10496 && std::env::var("MEMRA_NO_BATCHED").is_err()
10497 && (m <= 4 || Self::b8_enabled())
10498 && (m <= 8
10499 || qtype == QT_Q4_0
10500 || qtype == QT_Q6_K
10501 || qtype == QT_F8_E4M3
10502 || qtype == QT_NVFP4
10503 || qtype == QT_Q4_K
10504 || qtype == QT_Q5_K
10505 || qtype == QT_Q8_0)
10506 {
10507 let mcols = Self::batched_mcols(m);
10508 return self.qmatvec_mmvq_batched(
10509 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10510 );
10511 }
10512 if self.mmvq_supports(qtype) {
10513 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10514 }
10515 let x0 = self.zeros(0)?;
10518 self.matmul_pre(w, aq, ad, &x0, m)
10519 }
10520
10521 pub fn matmul_decode_exact_dual_pre(
10530 &self,
10531 w0: &crate::model::GpuTensor,
10532 w1: &crate::model::GpuTensor,
10533 aq: &CudaSlice<i8>,
10534 ad: &CudaSlice<f32>,
10535 m: usize,
10536 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10537 {
10538 use crate::model::GpuTensor;
10539 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10540 let on = *ON.get_or_init(|| {
10541 std::env::var("MEMRA_SPEC_DUAL_T")
10542 .map(|v| v != "0")
10543 .unwrap_or(true)
10544 });
10545 if !on
10546 || !(2..=7).contains(&m)
10547 || std::env::var("MEMRA_NO_BATCHED").is_ok()
10548 || !self.uses_q8_1_fast(w0)
10549 || !self.uses_q8_1_fast(w1)
10550 {
10551 return Ok(None);
10552 }
10553 if !self.mmvq_supports(QT_NVFP4) {
10558 return Ok(None);
10559 }
10560 let (in_f, out_f) = (w0.in_features(), w0.out_features());
10561 if w1.in_features() != in_f || w1.out_features() != out_f {
10562 return Ok(None);
10563 }
10564 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10565 (
10566 GpuTensor::Quant {
10567 bytes: b0,
10568 qtype: q0,
10569 row_bytes: rb0,
10570 scale: s0,
10571 rp: rp0,
10572 rp4: None,
10573 ..
10574 },
10575 GpuTensor::Quant {
10576 bytes: b1,
10577 qtype: q1,
10578 row_bytes: rb1,
10579 scale: s1,
10580 rp: rp1,
10581 rp4: None,
10582 ..
10583 },
10584 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10585 (b0, b1, *rb0, *s0, *s1, *rp0)
10586 }
10587 _ => return Ok(None),
10588 };
10589 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
10592 {
10593 return Ok(None);
10594 }
10595 let (y0, y1) =
10596 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
10597 Ok(Some(((y0, s0), (y1, s1))))
10598 }
10599
10600 pub fn matmul_decode_exact_dual(
10616 &self,
10617 w0: &crate::model::GpuTensor,
10618 w1: &crate::model::GpuTensor,
10619 x: &CudaSlice<f32>,
10620 m: usize,
10621 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
10622 use crate::model::GpuTensor;
10623 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10624 let on = *ON.get_or_init(|| {
10625 std::env::var("MEMRA_SPEC_DUAL_T")
10626 .map(|v| v != "0")
10627 .unwrap_or(true)
10628 });
10629 if !on
10630 || !(2..=4).contains(&m)
10631 || std::env::var("MEMRA_NO_BATCHED").is_ok()
10632 || !self.uses_q8_1_fast(w0)
10633 || !self.uses_q8_1_fast(w1)
10634 {
10635 return Ok(None);
10636 }
10637 if !self.mmvq_supports(QT_NVFP4) {
10642 return Ok(None);
10643 }
10644 let (in_f, out_f) = (w0.in_features(), w0.out_features());
10645 if w1.in_features() != in_f || w1.out_features() != out_f {
10646 return Ok(None);
10647 }
10648 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10649 (
10650 GpuTensor::Quant {
10651 bytes: b0,
10652 qtype: q0,
10653 row_bytes: rb0,
10654 scale: s0,
10655 rp: rp0,
10656 rp4: None,
10657 ..
10658 },
10659 GpuTensor::Quant {
10660 bytes: b1,
10661 qtype: q1,
10662 row_bytes: rb1,
10663 scale: s1,
10664 rp: rp1,
10665 rp4: None,
10666 ..
10667 },
10668 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10669 (b0, b1, *rb0, *s0, *s1, *rp0)
10670 }
10671 _ => return Ok(None),
10672 };
10673 if std::env::var("MEMRA_DEBUG").is_ok() {
10676 static ONCE: std::sync::Once = std::sync::Once::new();
10677 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
10678 }
10679 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10680 let (y0, y1) =
10681 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
10682 let mut y0 = y0;
10683 let mut y1 = y1;
10684 if s0 != 1.0 {
10685 self.scale_inplace(&mut y0, s0, m * out_f)?;
10686 }
10687 if s1 != 1.0 {
10688 self.scale_inplace(&mut y1, s1, m * out_f)?;
10689 }
10690 Ok(Some((y0, y1)))
10691 }
10692
10693 #[allow(clippy::too_many_arguments)]
10698 pub fn qmatvec_batched_dual_raw(
10699 &self,
10700 b0: &CudaSlice<u8>,
10701 b1: &CudaSlice<u8>,
10702 aq: &CudaSlice<i8>,
10703 ad: &CudaSlice<f32>,
10704 m: usize,
10705 in_f: usize,
10706 out_f: usize,
10707 row_bytes: usize,
10708 rp: bool,
10709 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10710 const ROWS_PER_BLOCK: u32 = 4;
10711 let mcols = Self::batched_mcols(m);
10712 let tiny_rp1 = rp
10715 && mcols == 4
10716 && out_f <= 128
10717 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
10718 let (name, rows_per_block) = if tiny_rp1 {
10719 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
10720 } else {
10721 match (mcols, rp, m) {
10722 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
10723 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
10724 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
10725 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
10726 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
10727 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
10728 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
10729 _ => {
10730 return Err(
10731 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
10732 );
10733 }
10734 }
10735 };
10736 let f = self.func(name);
10737 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
10738 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
10739 let cfg = LaunchConfig {
10740 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
10741 block_dim: (32, ROWS_PER_BLOCK, 1),
10742 shared_mem_bytes: 0,
10743 };
10744 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10745 let __s_b = self.gpu.stream();
10746 let mut b = __s_b.launch_builder(&f);
10747 b.arg(b0)
10748 .arg(b1)
10749 .arg(aq)
10750 .arg(ad)
10751 .arg(&mut y0)
10752 .arg(&mut y1)
10753 .arg(&inf)
10754 .arg(&outf)
10755 .arg(&mi)
10756 .arg(&rb);
10757 unsafe {
10758 b.launch(cfg)?;
10759 }
10760 Ok((y0, y1))
10761 }
10762
10763 pub fn matmul_pre_dual_noscale(
10775 &self,
10776 w0: &crate::model::GpuTensor,
10777 w1: &crate::model::GpuTensor,
10778 aq: &CudaSlice<i8>,
10779 ad: &CudaSlice<f32>,
10780 m: usize,
10781 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10782 {
10783 use crate::model::GpuTensor;
10784 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
10785 return Ok(None);
10786 }
10787 if !self.mmvq_supports(QT_NVFP4) {
10797 return Ok(None);
10798 }
10799 let (in_f, out_f) = (w0.in_features(), w0.out_features());
10800 if w1.in_features() != in_f || w1.out_features() != out_f {
10801 return Ok(None);
10802 }
10803 let no_mirror =
10816 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
10817 if self.q8_ffn_fuse2_on()
10818 && no_mirror(w0)
10819 && no_mirror(w1)
10820 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
10821 {
10822 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
10823 return Ok(Some(((y0, 1.0), (y1, 1.0))));
10824 }
10825 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
10835 let (y0, y1) =
10836 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
10837 return Ok(Some(((y0, p0.3), (y1, p1.3))));
10838 }
10839 let (b0, q0, rb0, s0, rp0) = match w0 {
10840 GpuTensor::Quant {
10841 bytes,
10842 qtype,
10843 row_bytes,
10844 scale,
10845 rp,
10846 ..
10847 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10848 _ => return Ok(None),
10849 };
10850 let (b1, q1, rb1, s1, rp1) = match w1 {
10851 GpuTensor::Quant {
10852 bytes,
10853 qtype,
10854 row_bytes,
10855 scale,
10856 rp,
10857 ..
10858 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10859 _ => return Ok(None),
10860 };
10861 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
10862 return Ok(None);
10863 }
10864 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
10866 let rows_per_block = ROWS_PER_BLOCK * RPW;
10867 let f = self.func(if rp0 {
10868 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
10869 } else {
10870 "qmatvec_nvfp4_mmvq_dual_mr2"
10871 });
10872 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
10873 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
10874 let cfg = LaunchConfig {
10875 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
10876 block_dim: (32, ROWS_PER_BLOCK, 1),
10877 shared_mem_bytes: 0,
10878 };
10879 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
10880 let one = 1.0f32;
10883 let __s_b = self.gpu.stream();
10884 let mut b = __s_b.launch_builder(&f);
10885 b.arg(b0)
10886 .arg(b1)
10887 .arg(aq)
10888 .arg(ad)
10889 .arg(&mut y0)
10890 .arg(&mut y1)
10891 .arg(&inf)
10892 .arg(&outf)
10893 .arg(&mi)
10894 .arg(&rb)
10895 .arg(&one)
10896 .arg(&one);
10897 unsafe {
10898 b.launch(cfg)?;
10899 }
10900 Ok(Some(((y0, s0), (y1, s1))))
10901 }
10902
10903 #[allow(clippy::too_many_arguments)]
10911 pub fn matmul_nvfp4_fused3(
10912 &self,
10913 w0: &crate::model::GpuTensor,
10914 w1: &crate::model::GpuTensor,
10915 w2: &crate::model::GpuTensor,
10916 aq: &CudaSlice<i8>,
10917 ad: &CudaSlice<f32>,
10918 m: usize,
10919 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
10920 {
10921 use crate::model::GpuTensor;
10922 if m != 1
10926 || !self.mmvq_supports(QT_NVFP4)
10927 || !self.uses_q8_1_fast(w0)
10928 || !self.uses_q8_1_fast(w1)
10929 || !self.uses_q8_1_fast(w2)
10930 {
10931 return Ok(None);
10932 }
10933 let unpack = |w: &crate::model::GpuTensor| match w {
10934 GpuTensor::Quant {
10935 bytes,
10936 qtype,
10937 scale,
10938 rp,
10939 ..
10940 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
10941 _ => None,
10942 };
10943 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
10944 return Ok(None);
10945 };
10946 let in_f = w0.in_features();
10947 if w1.in_features() != in_f || w2.in_features() != in_f {
10948 return Ok(None);
10949 }
10950 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
10951 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
10953 let rows_pb = ROWS_PER_BLOCK * RPW;
10954 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
10955 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
10956 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
10957 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
10958 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
10959 let cfg = LaunchConfig {
10960 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
10961 block_dim: (32, ROWS_PER_BLOCK, 1),
10962 shared_mem_bytes: 0,
10963 };
10964 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
10965 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
10968 let __s_b = self.gpu.stream();
10969 let mut b = __s_b.launch_builder(&f);
10970 b.arg(b0)
10971 .arg(b1)
10972 .arg(b2)
10973 .arg(aq)
10974 .arg(ad)
10975 .arg(&mut y0)
10976 .arg(&mut y1)
10977 .arg(&mut y2)
10978 .arg(&inf)
10979 .arg(&oi0)
10980 .arg(&oi1)
10981 .arg(&oi2)
10982 .arg(&mi)
10983 .arg(&p0.1)
10984 .arg(&p1.1)
10985 .arg(&p2.1);
10986 unsafe {
10987 b.launch(cfg)?;
10988 }
10989 Ok(Some((y0, y1, y2)))
10990 }
10991
10992 pub fn matmul_nvfp4_fused2(
11001 &self,
11002 w0: &crate::model::GpuTensor,
11003 w1: &crate::model::GpuTensor,
11004 aq: &CudaSlice<i8>,
11005 ad: &CudaSlice<f32>,
11006 m: usize,
11007 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11008 use crate::model::GpuTensor;
11009 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11010 let off =
11011 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11012 if off
11015 || m != 1
11016 || !self.mmvq_supports(QT_NVFP4)
11017 || !self.uses_q8_1_fast(w0)
11018 || !self.uses_q8_1_fast(w1)
11019 {
11020 return Ok(None);
11021 }
11022 let unpack = |w: &crate::model::GpuTensor| match w {
11023 GpuTensor::Quant {
11024 bytes,
11025 qtype,
11026 scale,
11027 rp,
11028 ..
11029 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11030 _ => None,
11031 };
11032 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11033 return Ok(None);
11034 };
11035 let in_f = w0.in_features();
11036 if w1.in_features() != in_f {
11037 return Ok(None);
11038 }
11039 let (o0, o1) = (w0.out_features(), w1.out_features());
11040 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11042 let rows_pb = ROWS_PER_BLOCK * RPW;
11043 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11044 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11045 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11046 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11047 let cfg = LaunchConfig {
11048 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11049 block_dim: (32, ROWS_PER_BLOCK, 1),
11050 shared_mem_bytes: 0,
11051 };
11052 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11053 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11056 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11059 {
11060 use cudarc::driver::{DevicePtr, DevicePtrMut};
11061 let s = &self.gpu.stream();
11062 let (pw0, _g0) = b0.device_ptr(s);
11063 let (pw1, _g1) = b1.device_ptr(s);
11064 let (paq, _g2) = aq.device_ptr(s);
11065 let (pad, _g3) = ad.device_ptr(s);
11066 let (py0, _g4) = y0.device_ptr_mut(s);
11067 let (py1, _g5) = y1.device_ptr_mut(s);
11068 let (s0, s1) = (p0.1, p1.1);
11069 let mut ps = [
11070 &pw0 as *const _ as *mut std::ffi::c_void,
11071 &pw1 as *const _ as *mut _,
11072 &paq as *const _ as *mut _,
11073 &pad as *const _ as *mut _,
11074 &py0 as *const _ as *mut _,
11075 &py1 as *const _ as *mut _,
11076 &inf as *const _ as *mut _,
11077 &oi0 as *const _ as *mut _,
11078 &oi1 as *const _ as *mut _,
11079 &mi as *const _ as *mut _,
11080 &s0 as *const _ as *mut _,
11081 &s1 as *const _ as *mut _,
11082 ];
11083 unsafe {
11084 self.launch_pdl(
11085 "qmatvec_nvfp4_mmvq_fused2_rp",
11086 cfg.grid_dim,
11087 cfg.block_dim,
11088 &mut ps,
11089 )?;
11090 }
11091 }
11092 return Ok(Some((y0, y1)));
11093 }
11094 let __s_b = self.gpu.stream();
11095 let mut b = __s_b.launch_builder(&f);
11096 b.arg(b0)
11097 .arg(b1)
11098 .arg(aq)
11099 .arg(ad)
11100 .arg(&mut y0)
11101 .arg(&mut y1)
11102 .arg(&inf)
11103 .arg(&oi0)
11104 .arg(&oi1)
11105 .arg(&mi)
11106 .arg(&p0.1)
11107 .arg(&p1.1);
11108 unsafe {
11109 b.launch(cfg)?;
11110 }
11111 Ok(Some((y0, y1)))
11112 }
11113
11114 pub fn matmul_nvfp4_fused2_into(
11119 &self,
11120 w0: &crate::model::GpuTensor,
11121 w1: &crate::model::GpuTensor,
11122 aq: &CudaSlice<i8>,
11123 ad: &CudaSlice<f32>,
11124 y0: &mut CudaSlice<f32>,
11125 y1: &mut CudaSlice<f32>,
11126 ) -> Result<bool, Box<dyn std::error::Error>> {
11127 use crate::model::GpuTensor;
11128 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11129 let off =
11130 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11131 if off
11132 || !self.mmvq_supports(QT_NVFP4)
11133 || !self.uses_q8_1_fast(w0)
11134 || !self.uses_q8_1_fast(w1)
11135 {
11136 return Ok(false);
11137 }
11138 let unpack = |w: &crate::model::GpuTensor| match w {
11139 GpuTensor::Quant {
11140 bytes,
11141 qtype,
11142 scale,
11143 rp,
11144 ..
11145 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11146 _ => None,
11147 };
11148 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11149 return Ok(false);
11150 };
11151 let in_f = w0.in_features();
11152 if w1.in_features() != in_f {
11153 return Ok(false);
11154 }
11155 let (o0, o1) = (w0.out_features(), w1.out_features());
11156 if y0.len() < o0 || y1.len() < o1 {
11157 return Ok(false);
11158 }
11159 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11161 let rows_pb = ROWS_PER_BLOCK * RPW;
11162 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11163 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11164 let cfg = LaunchConfig {
11165 grid_dim: (nb(o0) + nb(o1), 1, 1),
11166 block_dim: (32, ROWS_PER_BLOCK, 1),
11167 shared_mem_bytes: 0,
11168 };
11169 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11170 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11173 let __s_b = self.gpu.stream();
11174 let mut b = __s_b.launch_builder(&f);
11175 b.arg(b0)
11176 .arg(b1)
11177 .arg(aq)
11178 .arg(ad)
11179 .arg(&mut *y0)
11180 .arg(&mut *y1)
11181 .arg(&inf)
11182 .arg(&oi0)
11183 .arg(&oi1)
11184 .arg(&mi)
11185 .arg(&p0.1)
11186 .arg(&p1.1);
11187 unsafe {
11188 b.launch(cfg)?;
11189 }
11190 Ok(true)
11191 }
11192
11193 #[allow(clippy::type_complexity)]
11198 pub fn matmul_nvfp4_fused4(
11199 &self,
11200 w0: &crate::model::GpuTensor,
11201 w1: &crate::model::GpuTensor,
11202 w2: &crate::model::GpuTensor,
11203 w3: &crate::model::GpuTensor,
11204 aq: &CudaSlice<i8>,
11205 ad: &CudaSlice<f32>,
11206 m: usize,
11207 ) -> Result<
11208 Option<(
11209 CudaSlice<f32>,
11210 CudaSlice<f32>,
11211 CudaSlice<f32>,
11212 CudaSlice<f32>,
11213 )>,
11214 Box<dyn std::error::Error>,
11215 > {
11216 use crate::model::GpuTensor;
11217 if m != 1
11219 || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11220 || !self.mmvq_supports(QT_NVFP4)
11221 || !self.uses_q8_1_fast(w0)
11222 || !self.uses_q8_1_fast(w1)
11223 || !self.uses_q8_1_fast(w2)
11224 || !self.uses_q8_1_fast(w3)
11225 {
11226 return Ok(None);
11227 }
11228 let unpack = |w: &crate::model::GpuTensor| match w {
11229 GpuTensor::Quant {
11230 bytes,
11231 qtype,
11232 scale,
11233 rp,
11234 ..
11235 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11236 _ => None,
11237 };
11238 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11239 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11240 else {
11241 return Ok(None);
11242 };
11243 let in_f = w0.in_features();
11244 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11245 return Ok(None);
11246 }
11247 let (o0, o1, o2, o3) = (
11248 w0.out_features(),
11249 w1.out_features(),
11250 w2.out_features(),
11251 w3.out_features(),
11252 );
11253 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11255 let rows_pb = ROWS_PER_BLOCK * RPW;
11256 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11257 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
11258 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11259 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11260 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11261 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11262 let cfg = LaunchConfig {
11263 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
11264 block_dim: (32, ROWS_PER_BLOCK, 1),
11265 shared_mem_bytes: 0,
11266 };
11267 let (inf, oi0, oi1, oi2, oi3, mi) = (
11268 in_f as i32,
11269 o0 as i32,
11270 o1 as i32,
11271 o2 as i32,
11272 o3 as i32,
11273 m as i32,
11274 );
11275 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11278 let __s_b = self.gpu.stream();
11279 let mut b = __s_b.launch_builder(&f);
11280 b.arg(b0)
11281 .arg(b1)
11282 .arg(b2)
11283 .arg(b3)
11284 .arg(aq)
11285 .arg(ad)
11286 .arg(&mut y0)
11287 .arg(&mut y1)
11288 .arg(&mut y2)
11289 .arg(&mut y3)
11290 .arg(&inf)
11291 .arg(&oi0)
11292 .arg(&oi1)
11293 .arg(&oi2)
11294 .arg(&oi3)
11295 .arg(&mi)
11296 .arg(&p0.1)
11297 .arg(&p1.1)
11298 .arg(&p2.1)
11299 .arg(&p3.1);
11300 unsafe {
11301 b.launch(cfg)?;
11302 }
11303 Ok(Some((y0, y1, y2, y3)))
11304 }
11305
11306 pub fn matmul_q8_fused2(
11314 &self,
11315 w0: &crate::model::GpuTensor,
11316 w1: &crate::model::GpuTensor,
11317 aq: &CudaSlice<i8>,
11318 ad: &CudaSlice<f32>,
11319 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11320 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11326 return Ok(Some(self.e4m3_fused2_core(
11327 p0.0,
11328 p1.0,
11329 aq,
11330 ad,
11331 w0.in_features(),
11332 p0.1,
11333 p1.1,
11334 p0.2,
11335 p0.3,
11336 p1.3,
11337 )?));
11338 }
11339 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11340 return Ok(None);
11341 };
11342 Ok(Some(self.q8_fused2_core(
11343 p0.0,
11344 p1.0,
11345 aq,
11346 ad,
11347 w0.in_features(),
11348 p0.1,
11349 p1.1,
11350 p0.2,
11351 )?))
11352 }
11353
11354 #[allow(clippy::too_many_arguments)]
11355 fn q8_fused2_core(
11356 &self,
11357 b0: &CudaSlice<u8>,
11358 b1: &CudaSlice<u8>,
11359 aq: &CudaSlice<i8>,
11360 ad: &CudaSlice<f32>,
11361 in_f: usize,
11362 out0: usize,
11363 out1: usize,
11364 row_bytes: usize,
11365 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11366 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
11368 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
11369 let f = self.func("qmatvec_q8_0_mmvq_fused2");
11370 let mut y0 = self.alloc_uninit::<f32>(out0)?;
11371 let mut y1 = self.alloc_uninit::<f32>(out1)?;
11372 let cfg = LaunchConfig {
11373 grid_dim: (nb0 + nb1, 1, 1),
11374 block_dim: (32, ROWS_PER_BLOCK, 1),
11375 shared_mem_bytes: 0,
11376 };
11377 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
11378 let __s_b = self.gpu.stream();
11379 let mut b = __s_b.launch_builder(&f);
11380 b.arg(b0)
11381 .arg(b1)
11382 .arg(aq)
11383 .arg(ad)
11384 .arg(&mut y0)
11385 .arg(&mut y1)
11386 .arg(&inf)
11387 .arg(&o0)
11388 .arg(&o1)
11389 .arg(&rbl);
11390 unsafe {
11391 b.launch(cfg)?;
11392 }
11393 Ok((y0, y1))
11394 }
11395
11396 pub fn matmul_q8_fused2_x(
11402 &self,
11403 w0: &crate::model::GpuTensor,
11404 w1: &crate::model::GpuTensor,
11405 x: &CudaSlice<f32>,
11406 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11407 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11408 return Ok(None);
11409 }
11410 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11411 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11412 return Ok(Some(self.e4m3_fused2_core(
11413 p0.0,
11414 p1.0,
11415 &aq,
11416 &ad,
11417 w0.in_features(),
11418 p0.1,
11419 p1.1,
11420 p0.2,
11421 p0.3,
11422 p1.3,
11423 )?));
11424 }
11425 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11426 return Ok(None);
11427 };
11428 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11429 Ok(Some(self.q8_fused2_core(
11430 p0.0,
11431 p1.0,
11432 &aq,
11433 &ad,
11434 w0.in_features(),
11435 p0.1,
11436 p1.1,
11437 p0.2,
11438 )?))
11439 }
11440
11441 #[allow(clippy::too_many_arguments)]
11444 pub fn qmatvec_q8_fused2_raw(
11445 &self,
11446 b0: &CudaSlice<u8>,
11447 b1: &CudaSlice<u8>,
11448 x: &CudaSlice<f32>,
11449 in_f: usize,
11450 out0: usize,
11451 out1: usize,
11452 row_bytes: usize,
11453 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11454 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
11455 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
11456 }
11457
11458 pub fn matmul_q4_fused3(
11464 &self,
11465 w0: &crate::model::GpuTensor,
11466 w1: &crate::model::GpuTensor,
11467 w2: &crate::model::GpuTensor,
11468 aq: &CudaSlice<i8>,
11469 ad: &CudaSlice<f32>,
11470 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11471 {
11472 use crate::model::GpuTensor;
11473 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11474 match w {
11475 GpuTensor::Quant {
11476 qtype, row_bytes, ..
11477 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11478 _ => None,
11479 }
11480 };
11481 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11482 return Ok(None);
11483 };
11484 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11485 return Ok(None);
11486 }
11487 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11491 match w {
11492 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11493 Some(m) => (m, true),
11494 None => (bytes, *rp),
11495 },
11496 _ => unreachable!(),
11497 }
11498 }
11499 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11500 if rp0 != rp1 || rp1 != rp2 {
11501 return Ok(None);
11502 }
11503 let rp = rp0;
11504 let rpb: u32 = 4;
11505 let mr1 = rp && Self::q40_mr1_on();
11509 let nb = |o: usize| {
11510 if mr1 {
11511 (o as u32).div_ceil(rpb)
11512 } else {
11513 (o as u32).div_ceil(2).div_ceil(rpb)
11514 }
11515 };
11516 let grid = nb(o0) + nb(o1) + nb(o2);
11517 let mut y0 = self.alloc_uninit::<f32>(o0)?;
11518 let mut y1 = self.alloc_uninit::<f32>(o1)?;
11519 let mut y2 = self.alloc_uninit::<f32>(o2)?;
11520 let f = self.func(if mr1 {
11521 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11522 } else if rp {
11523 "qmatvec_q4_0_mmvq_fused3_rp"
11524 } else {
11525 "qmatvec_q4_0_mmvq_fused3"
11526 });
11527 let cfg = LaunchConfig {
11528 grid_dim: (grid, 1, 1),
11529 block_dim: (32, rpb, 1),
11530 shared_mem_bytes: 0,
11531 };
11532 let inf = w0.in_features() as i32;
11533 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11534 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11535 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11538 {
11539 use cudarc::driver::{DevicePtr, DevicePtrMut};
11540 let s = &self.gpu.stream();
11541 let (p0, _g0) = b0.device_ptr(s);
11542 let (p1, _g1) = b1.device_ptr(s);
11543 let (p2, _g2) = b2.device_ptr(s);
11544 let (paq, _g3) = aq.device_ptr(s);
11545 let (pad, _g4) = ad.device_ptr(s);
11546 let (py0, _g5) = y0.device_ptr_mut(s);
11547 let (py1, _g6) = y1.device_ptr_mut(s);
11548 let (py2, _g7) = y2.device_ptr_mut(s);
11549 let mut ps = [
11550 &p0 as *const _ as *mut std::ffi::c_void,
11551 &p1 as *const _ as *mut _,
11552 &p2 as *const _ as *mut _,
11553 &paq as *const _ as *mut _,
11554 &pad as *const _ as *mut _,
11555 &py0 as *const _ as *mut _,
11556 &py1 as *const _ as *mut _,
11557 &py2 as *const _ as *mut _,
11558 &inf as *const _ as *mut _,
11559 &oo0 as *const _ as *mut _,
11560 &oo1 as *const _ as *mut _,
11561 &oo2 as *const _ as *mut _,
11562 &r0 as *const _ as *mut _,
11563 &r1 as *const _ as *mut _,
11564 &r2 as *const _ as *mut _,
11565 ];
11566 unsafe {
11567 self.launch_pdl(
11568 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11569 (grid, 1, 1),
11570 (32, rpb, 1),
11571 &mut ps,
11572 )?;
11573 }
11574 }
11575 return Ok(Some((y0, y1, y2)));
11576 }
11577 let __s_b = self.gpu.stream();
11578 let mut b = __s_b.launch_builder(&f);
11579 b.arg(b0)
11580 .arg(b1)
11581 .arg(b2)
11582 .arg(aq)
11583 .arg(ad)
11584 .arg(&mut y0)
11585 .arg(&mut y1)
11586 .arg(&mut y2)
11587 .arg(&inf)
11588 .arg(&oo0)
11589 .arg(&oo1)
11590 .arg(&oo2)
11591 .arg(&r0)
11592 .arg(&r1)
11593 .arg(&r2);
11594 unsafe {
11595 b.launch(cfg)?;
11596 }
11597 Ok(Some((y0, y1, y2)))
11598 }
11599
11600 #[allow(clippy::too_many_arguments)]
11603 pub fn matmul_q4_fused3_into(
11604 &self,
11605 w0: &crate::model::GpuTensor,
11606 w1: &crate::model::GpuTensor,
11607 w2: &crate::model::GpuTensor,
11608 aq: &CudaSlice<i8>,
11609 ad: &CudaSlice<f32>,
11610 y0: &mut CudaSlice<f32>,
11611 y1: &mut CudaSlice<f32>,
11612 y2: &mut CudaSlice<f32>,
11613 ) -> Result<bool, Box<dyn std::error::Error>> {
11614 use crate::model::GpuTensor;
11615 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11616 match w {
11617 GpuTensor::Quant {
11618 qtype, row_bytes, ..
11619 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11620 _ => None,
11621 }
11622 };
11623 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11624 return Ok(false);
11625 };
11626 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11627 return Ok(false);
11628 }
11629 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11630 match w {
11631 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11632 Some(m) => (m, true),
11633 None => (bytes, *rp),
11634 },
11635 _ => unreachable!(),
11636 }
11637 }
11638 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11639 if rp0 != rp1 || rp1 != rp2 {
11640 return Ok(false);
11641 }
11642 let rp = rp0;
11643 let rpb: u32 = 4;
11644 let mr1 = rp && Self::q40_mr1_on();
11645 let nb = |o: usize| {
11646 if mr1 {
11647 (o as u32).div_ceil(rpb)
11648 } else {
11649 (o as u32).div_ceil(2).div_ceil(rpb)
11650 }
11651 };
11652 let grid = nb(o0) + nb(o1) + nb(o2);
11653 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
11654 let f = self.func(if mr1 {
11655 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11656 } else if rp {
11657 "qmatvec_q4_0_mmvq_fused3_rp"
11658 } else {
11659 "qmatvec_q4_0_mmvq_fused3"
11660 });
11661 let cfg = LaunchConfig {
11662 grid_dim: (grid, 1, 1),
11663 block_dim: (32, rpb, 1),
11664 shared_mem_bytes: 0,
11665 };
11666 let inf = w0.in_features() as i32;
11667 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11668 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11669 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11671 use cudarc::driver::{DevicePtr, DevicePtrMut};
11672 let s = &self.gpu.stream();
11673 let (p0, _g0) = b0.device_ptr(s);
11674 let (p1, _g1) = b1.device_ptr(s);
11675 let (p2, _g2) = b2.device_ptr(s);
11676 let (paq, _g3) = aq.device_ptr(s);
11677 let (pad, _g4) = ad.device_ptr(s);
11678 let (py0, _g5) = y0.device_ptr_mut(s);
11679 let (py1, _g6) = y1.device_ptr_mut(s);
11680 let (py2, _g7) = y2.device_ptr_mut(s);
11681 let mut ps = [
11682 &p0 as *const _ as *mut std::ffi::c_void,
11683 &p1 as *const _ as *mut _,
11684 &p2 as *const _ as *mut _,
11685 &paq as *const _ as *mut _,
11686 &pad as *const _ as *mut _,
11687 &py0 as *const _ as *mut _,
11688 &py1 as *const _ as *mut _,
11689 &py2 as *const _ as *mut _,
11690 &inf as *const _ as *mut _,
11691 &oo0 as *const _ as *mut _,
11692 &oo1 as *const _ as *mut _,
11693 &oo2 as *const _ as *mut _,
11694 &r0 as *const _ as *mut _,
11695 &r1 as *const _ as *mut _,
11696 &r2 as *const _ as *mut _,
11697 ];
11698 unsafe {
11699 self.launch_pdl(
11700 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11701 (grid, 1, 1),
11702 (32, rpb, 1),
11703 &mut ps,
11704 )?;
11705 }
11706 return Ok(true);
11707 }
11708 let __s_b = self.gpu.stream();
11709 let mut b = __s_b.launch_builder(&f);
11710 b.arg(b0)
11711 .arg(b1)
11712 .arg(b2)
11713 .arg(aq)
11714 .arg(ad)
11715 .arg(&mut *y0)
11716 .arg(&mut *y1)
11717 .arg(&mut *y2)
11718 .arg(&inf)
11719 .arg(&oo0)
11720 .arg(&oo1)
11721 .arg(&oo2)
11722 .arg(&r0)
11723 .arg(&r1)
11724 .arg(&r2);
11725 unsafe {
11726 b.launch(cfg)?;
11727 }
11728 Ok(true)
11729 }
11730
11731 pub fn matmul_q4_fused2(
11733 &self,
11734 w0: &crate::model::GpuTensor,
11735 w1: &crate::model::GpuTensor,
11736 aq: &CudaSlice<i8>,
11737 ad: &CudaSlice<f32>,
11738 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11739 use crate::model::GpuTensor;
11740 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11741 match w {
11742 GpuTensor::Quant {
11743 qtype, row_bytes, ..
11744 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11745 _ => None,
11746 }
11747 };
11748 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
11749 return Ok(None);
11750 };
11751 if w0.in_features() != w1.in_features() {
11752 return Ok(None);
11753 }
11754 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11756 match w {
11757 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11758 Some(m) => (m, true),
11759 None => (bytes, *rp),
11760 },
11761 _ => unreachable!(),
11762 }
11763 }
11764 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
11765 if rp0 != rp1 {
11766 return Ok(None);
11767 }
11768 let rp = rp0;
11769 let rpb: u32 = 4;
11770 let mr1 = rp && Self::q40_mr1_on();
11772 let nb = |o: usize| {
11773 if mr1 {
11774 (o as u32).div_ceil(rpb)
11775 } else {
11776 (o as u32).div_ceil(2).div_ceil(rpb)
11777 }
11778 };
11779 let grid = nb(o0) + nb(o1);
11780 let mut y0 = self.alloc_uninit::<f32>(o0)?;
11781 let mut y1 = self.alloc_uninit::<f32>(o1)?;
11782 let f = self.func(if mr1 {
11783 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
11784 } else if rp {
11785 "qmatvec_q4_0_mmvq_fused2_rp"
11786 } else {
11787 "qmatvec_q4_0_mmvq_fused2"
11788 });
11789 let cfg = LaunchConfig {
11790 grid_dim: (grid, 1, 1),
11791 block_dim: (32, rpb, 1),
11792 shared_mem_bytes: 0,
11793 };
11794 let inf = w0.in_features() as i32;
11795 let (oo0, oo1) = (o0 as i32, o1 as i32);
11796 let (r0, r1) = (rb0 as i64, rb1 as i64);
11797 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11799 {
11800 use cudarc::driver::{DevicePtr, DevicePtrMut};
11801 let s = &self.gpu.stream();
11802 let (p0, _g0) = b0.device_ptr(s);
11803 let (p1, _g1) = b1.device_ptr(s);
11804 let (paq, _g2) = aq.device_ptr(s);
11805 let (pad, _g3) = ad.device_ptr(s);
11806 let (py0, _g4) = y0.device_ptr_mut(s);
11807 let (py1, _g5) = y1.device_ptr_mut(s);
11808 let mut ps = [
11809 &p0 as *const _ as *mut std::ffi::c_void,
11810 &p1 as *const _ as *mut _,
11811 &paq as *const _ as *mut _,
11812 &pad as *const _ as *mut _,
11813 &py0 as *const _ as *mut _,
11814 &py1 as *const _ as *mut _,
11815 &inf as *const _ as *mut _,
11816 &oo0 as *const _ as *mut _,
11817 &oo1 as *const _ as *mut _,
11818 &r0 as *const _ as *mut _,
11819 &r1 as *const _ as *mut _,
11820 ];
11821 unsafe {
11822 self.launch_pdl(
11823 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
11824 (grid, 1, 1),
11825 (32, rpb, 1),
11826 &mut ps,
11827 )?;
11828 }
11829 }
11830 return Ok(Some((y0, y1)));
11831 }
11832 let __s_b = self.gpu.stream();
11833 let mut b = __s_b.launch_builder(&f);
11834 b.arg(b0)
11835 .arg(b1)
11836 .arg(aq)
11837 .arg(ad)
11838 .arg(&mut y0)
11839 .arg(&mut y1)
11840 .arg(&inf)
11841 .arg(&oo0)
11842 .arg(&oo1)
11843 .arg(&r0)
11844 .arg(&r1);
11845 unsafe {
11846 b.launch(cfg)?;
11847 }
11848 Ok(Some((y0, y1)))
11849 }
11850
11851 pub fn matmul_q4_fused2_into(
11853 &self,
11854 w0: &crate::model::GpuTensor,
11855 w1: &crate::model::GpuTensor,
11856 aq: &CudaSlice<i8>,
11857 ad: &CudaSlice<f32>,
11858 y0: &mut CudaSlice<f32>,
11859 y1: &mut CudaSlice<f32>,
11860 ) -> Result<bool, Box<dyn std::error::Error>> {
11861 use crate::model::GpuTensor;
11862 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11863 match w {
11864 GpuTensor::Quant {
11865 qtype, row_bytes, ..
11866 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11867 _ => None,
11868 }
11869 };
11870 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
11871 return Ok(false);
11872 };
11873 if w0.in_features() != w1.in_features() {
11874 return Ok(false);
11875 }
11876 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11877 match w {
11878 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11879 Some(m) => (m, true),
11880 None => (bytes, *rp),
11881 },
11882 _ => unreachable!(),
11883 }
11884 }
11885 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
11886 if rp0 != rp1 {
11887 return Ok(false);
11888 }
11889 let rp = rp0;
11890 let rpb: u32 = 4;
11891 let mr1 = rp && Self::q40_mr1_on();
11892 let nb = |o: usize| {
11893 if mr1 {
11894 (o as u32).div_ceil(rpb)
11895 } else {
11896 (o as u32).div_ceil(2).div_ceil(rpb)
11897 }
11898 };
11899 let grid = nb(o0) + nb(o1);
11900 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
11901 let f = self.func(if mr1 {
11902 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
11903 } else if rp {
11904 "qmatvec_q4_0_mmvq_fused2_rp"
11905 } else {
11906 "qmatvec_q4_0_mmvq_fused2"
11907 });
11908 let cfg = LaunchConfig {
11909 grid_dim: (grid, 1, 1),
11910 block_dim: (32, rpb, 1),
11911 shared_mem_bytes: 0,
11912 };
11913 let inf = w0.in_features() as i32;
11914 let (oo0, oo1) = (o0 as i32, o1 as i32);
11915 let (r0, r1) = (rb0 as i64, rb1 as i64);
11916 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11918 use cudarc::driver::{DevicePtr, DevicePtrMut};
11919 let s = &self.gpu.stream();
11920 let (p0, _g0) = b0.device_ptr(s);
11921 let (p1, _g1) = b1.device_ptr(s);
11922 let (paq, _g2) = aq.device_ptr(s);
11923 let (pad, _g3) = ad.device_ptr(s);
11924 let (py0, _g4) = y0.device_ptr_mut(s);
11925 let (py1, _g5) = y1.device_ptr_mut(s);
11926 let mut ps = [
11927 &p0 as *const _ as *mut std::ffi::c_void,
11928 &p1 as *const _ as *mut _,
11929 &paq as *const _ as *mut _,
11930 &pad as *const _ as *mut _,
11931 &py0 as *const _ as *mut _,
11932 &py1 as *const _ as *mut _,
11933 &inf as *const _ as *mut _,
11934 &oo0 as *const _ as *mut _,
11935 &oo1 as *const _ as *mut _,
11936 &r0 as *const _ as *mut _,
11937 &r1 as *const _ as *mut _,
11938 ];
11939 unsafe {
11940 self.launch_pdl(
11941 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
11942 (grid, 1, 1),
11943 (32, rpb, 1),
11944 &mut ps,
11945 )?;
11946 }
11947 return Ok(true);
11948 }
11949 let __s_b = self.gpu.stream();
11950 let mut b = __s_b.launch_builder(&f);
11951 b.arg(b0)
11952 .arg(b1)
11953 .arg(aq)
11954 .arg(ad)
11955 .arg(&mut *y0)
11956 .arg(&mut *y1)
11957 .arg(&inf)
11958 .arg(&oo0)
11959 .arg(&oo1)
11960 .arg(&r0)
11961 .arg(&r1);
11962 unsafe {
11963 b.launch(cfg)?;
11964 }
11965 Ok(true)
11966 }
11967
11968 pub fn matmul_q4_fused2_batched(
11973 &self,
11974 w0: &crate::model::GpuTensor,
11975 w1: &crate::model::GpuTensor,
11976 aq: &CudaSlice<i8>,
11977 ad: &CudaSlice<f32>,
11978 m: usize,
11979 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11980 use crate::model::GpuTensor;
11981 if m < 2 || m > 8 {
11982 return Ok(None);
11983 }
11984 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11985 match w {
11986 GpuTensor::Quant {
11987 qtype, row_bytes, ..
11988 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11989 _ => None,
11990 }
11991 };
11992 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
11993 return Ok(None);
11994 };
11995 if w0.in_features() != w1.in_features() {
11996 return Ok(None);
11997 }
11998 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11999 match w {
12000 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12001 Some(mr) => (mr, true),
12002 None => (bytes, *rp),
12003 },
12004 _ => unreachable!(),
12005 }
12006 }
12007 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12008 if !rp0 || !rp1 {
12009 return Ok(None);
12010 }
12011 let mcols = Self::batched_mcols(m);
12012 let rpb: u32 = 4;
12013 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12014 let grid = nb(o0) + nb(o1);
12015 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12016 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12017 let f = self.func(match mcols {
12018 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12019 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12020 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12021 });
12022 let cfg = LaunchConfig {
12023 grid_dim: (grid, 1, 1),
12024 block_dim: (32, rpb, 1),
12025 shared_mem_bytes: 0,
12026 };
12027 let inf = w0.in_features() as i32;
12028 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12029 let rb = rb0 as i64;
12030 let __s_b = self.gpu.stream();
12031 let mut b = __s_b.launch_builder(&f);
12032 b.arg(b0)
12033 .arg(b1)
12034 .arg(aq)
12035 .arg(ad)
12036 .arg(&mut y0)
12037 .arg(&mut y1)
12038 .arg(&inf)
12039 .arg(&oo0)
12040 .arg(&oo1)
12041 .arg(&mi)
12042 .arg(&rb);
12043 unsafe {
12044 b.launch(cfg)?;
12045 }
12046 Ok(Some((y0, y1)))
12047 }
12048
12049 #[allow(clippy::too_many_arguments)]
12052 pub fn matmul_q4_fused3_batched(
12053 &self,
12054 w0: &crate::model::GpuTensor,
12055 w1: &crate::model::GpuTensor,
12056 w2: &crate::model::GpuTensor,
12057 aq: &CudaSlice<i8>,
12058 ad: &CudaSlice<f32>,
12059 m: usize,
12060 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12061 {
12062 use crate::model::GpuTensor;
12063 if m < 2 || m > 8 {
12064 return Ok(None);
12065 }
12066 let q4 = |w: &GpuTensor| -> Option<usize> {
12067 match w {
12068 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12069 _ => None,
12070 }
12071 };
12072 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12073 return Ok(None);
12074 };
12075 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12076 return Ok(None);
12077 }
12078 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12079 match w {
12080 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12081 Some(mr) => (mr, true),
12082 None => (bytes, *rp),
12083 },
12084 _ => unreachable!(),
12085 }
12086 }
12087 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12088 if !rp0 || !rp1 || !rp2 {
12089 return Ok(None);
12090 }
12091 let mcols = Self::batched_mcols(m);
12092 let rpb: u32 = 4;
12093 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12094 let grid = nb(o0) + nb(o1) + nb(o2);
12095 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12096 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12097 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12098 let f = self.func(match mcols {
12099 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12100 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12101 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12102 });
12103 let cfg = LaunchConfig {
12104 grid_dim: (grid, 1, 1),
12105 block_dim: (32, rpb, 1),
12106 shared_mem_bytes: 0,
12107 };
12108 let inf = w0.in_features() as i32;
12109 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12110 let rb = 0i64;
12111 let __s_b = self.gpu.stream();
12112 let mut b = __s_b.launch_builder(&f);
12113 b.arg(b0)
12114 .arg(b1)
12115 .arg(b2)
12116 .arg(aq)
12117 .arg(ad)
12118 .arg(&mut y0)
12119 .arg(&mut y1)
12120 .arg(&mut y2)
12121 .arg(&inf)
12122 .arg(&oo0)
12123 .arg(&oo1)
12124 .arg(&oo2)
12125 .arg(&mi)
12126 .arg(&rb);
12127 unsafe {
12128 b.launch(cfg)?;
12129 }
12130 Ok(Some((y0, y1, y2)))
12131 }
12132
12133 pub fn matmul_q8_fused3(
12134 &self,
12135 w0: &crate::model::GpuTensor,
12136 w1: &crate::model::GpuTensor,
12137 w2: &crate::model::GpuTensor,
12138 aq: &CudaSlice<i8>,
12139 ad: &CudaSlice<f32>,
12140 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12141 {
12142 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12145 return Ok(Some(self.e4m3_fused3_core(
12146 p0.0,
12147 p1.0,
12148 p2.0,
12149 aq,
12150 ad,
12151 w0.in_features(),
12152 p0.1,
12153 p1.1,
12154 p2.1,
12155 p0.2,
12156 p0.3,
12157 p1.3,
12158 p2.3,
12159 )?));
12160 }
12161 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12162 return Ok(None);
12163 };
12164 Ok(Some(self.q8_fused3_core(
12165 p0.0,
12166 p1.0,
12167 p2.0,
12168 aq,
12169 ad,
12170 w0.in_features(),
12171 p0.1,
12172 p1.1,
12173 p2.1,
12174 p0.2,
12175 )?))
12176 }
12177
12178 #[allow(clippy::too_many_arguments)]
12179 fn q8_fused3_core(
12180 &self,
12181 b0: &CudaSlice<u8>,
12182 b1: &CudaSlice<u8>,
12183 b2: &CudaSlice<u8>,
12184 aq: &CudaSlice<i8>,
12185 ad: &CudaSlice<f32>,
12186 in_f: usize,
12187 out0: usize,
12188 out1: usize,
12189 out2: usize,
12190 row_bytes: usize,
12191 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12192 const ROWS_PER_BLOCK: u32 = 4;
12193 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12194 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12195 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12196 let f = self.func("qmatvec_q8_0_mmvq_fused3");
12197 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12198 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12199 let mut y2 = self.alloc_uninit::<f32>(out2)?;
12200 let cfg = LaunchConfig {
12201 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12202 block_dim: (32, ROWS_PER_BLOCK, 1),
12203 shared_mem_bytes: 0,
12204 };
12205 let (inf, o0, o1, o2, rbl) = (
12206 in_f as i32,
12207 out0 as i32,
12208 out1 as i32,
12209 out2 as i32,
12210 row_bytes as i64,
12211 );
12212 let __s_b = self.gpu.stream();
12213 let mut b = __s_b.launch_builder(&f);
12214 b.arg(b0)
12215 .arg(b1)
12216 .arg(b2)
12217 .arg(aq)
12218 .arg(ad)
12219 .arg(&mut y0)
12220 .arg(&mut y1)
12221 .arg(&mut y2)
12222 .arg(&inf)
12223 .arg(&o0)
12224 .arg(&o1)
12225 .arg(&o2)
12226 .arg(&rbl);
12227 unsafe {
12228 b.launch(cfg)?;
12229 }
12230 Ok((y0, y1, y2))
12231 }
12232
12233 #[allow(clippy::too_many_arguments)]
12235 pub fn qmatvec_q8_fused3_raw(
12236 &self,
12237 b0: &CudaSlice<u8>,
12238 b1: &CudaSlice<u8>,
12239 b2: &CudaSlice<u8>,
12240 x: &CudaSlice<f32>,
12241 in_f: usize,
12242 out0: usize,
12243 out1: usize,
12244 out2: usize,
12245 row_bytes: usize,
12246 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12247 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12248 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12249 }
12250
12251 pub fn matmul_q8_fused2_t(
12262 &self,
12263 w0: &crate::model::GpuTensor,
12264 w1: &crate::model::GpuTensor,
12265 aq: &CudaSlice<i8>,
12266 ad: &CudaSlice<f32>,
12267 m: usize,
12268 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12269 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12273 return Ok(None);
12274 }
12275 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12278 if m > 4 && !Self::b8_enabled() {
12279 return Ok(None);
12280 }
12281 return Ok(Some(self.e4m3_fused2_t_core(
12282 p0.0,
12283 p1.0,
12284 aq,
12285 ad,
12286 m,
12287 w0.in_features(),
12288 p0.1,
12289 p1.1,
12290 p0.2,
12291 p0.3,
12292 p1.3,
12293 )?));
12294 }
12295 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12296 return Ok(None);
12297 };
12298 Ok(Some(self.q8_fused2_t_core(
12299 p0.0,
12300 p1.0,
12301 aq,
12302 ad,
12303 m,
12304 w0.in_features(),
12305 p0.1,
12306 p1.1,
12307 p0.2,
12308 )?))
12309 }
12310
12311 #[allow(clippy::too_many_arguments)]
12312 fn q8_fused2_t_core(
12313 &self,
12314 b0: &CudaSlice<u8>,
12315 b1: &CudaSlice<u8>,
12316 aq: &CudaSlice<i8>,
12317 ad: &CudaSlice<f32>,
12318 m: usize,
12319 in_f: usize,
12320 out0: usize,
12321 out1: usize,
12322 row_bytes: usize,
12323 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12324 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12326 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12327 let f = self.func(match Self::batched_mcols(m) {
12328 2 => "qmatvec_q8_0_mmvq_fused2_b2",
12329 4 => "qmatvec_q8_0_mmvq_fused2_b4",
12330 _ => "qmatvec_q8_0_mmvq_fused2_b8",
12332 });
12333 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12334 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12335 let cfg = LaunchConfig {
12336 grid_dim: (nb0 + nb1, 1, 1),
12337 block_dim: (32, ROWS_PER_BLOCK, 1),
12338 shared_mem_bytes: 0,
12339 };
12340 let (inf, o0, o1, mi, rbl) = (
12341 in_f as i32,
12342 out0 as i32,
12343 out1 as i32,
12344 m as i32,
12345 row_bytes as i64,
12346 );
12347 let __s_b = self.gpu.stream();
12348 let mut b = __s_b.launch_builder(&f);
12349 b.arg(b0)
12350 .arg(b1)
12351 .arg(aq)
12352 .arg(ad)
12353 .arg(&mut y0)
12354 .arg(&mut y1)
12355 .arg(&inf)
12356 .arg(&o0)
12357 .arg(&o1)
12358 .arg(&mi)
12359 .arg(&rbl);
12360 unsafe {
12361 b.launch(cfg)?;
12362 }
12363 Ok((y0, y1))
12364 }
12365
12366 #[allow(clippy::too_many_arguments)]
12369 pub fn qmatvec_q8_fused2_t_raw(
12370 &self,
12371 b0: &CudaSlice<u8>,
12372 b1: &CudaSlice<u8>,
12373 x: &CudaSlice<f32>,
12374 m: usize,
12375 in_f: usize,
12376 out0: usize,
12377 out1: usize,
12378 row_bytes: usize,
12379 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12380 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12381 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
12382 }
12383
12384 #[allow(clippy::too_many_arguments)]
12387 pub fn matmul_q8_fused3_t(
12388 &self,
12389 w0: &crate::model::GpuTensor,
12390 w1: &crate::model::GpuTensor,
12391 w2: &crate::model::GpuTensor,
12392 aq: &CudaSlice<i8>,
12393 ad: &CudaSlice<f32>,
12394 m: usize,
12395 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12396 {
12397 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12398 return Ok(None);
12399 }
12400 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12401 return Ok(Some(self.e4m3_fused3_t_core(
12402 p0.0,
12403 p1.0,
12404 p2.0,
12405 aq,
12406 ad,
12407 m,
12408 w0.in_features(),
12409 p0.1,
12410 p1.1,
12411 p2.1,
12412 p0.2,
12413 p0.3,
12414 p1.3,
12415 p2.3,
12416 )?));
12417 }
12418 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12419 return Ok(None);
12420 };
12421 Ok(Some(self.q8_fused3_t_core(
12422 p0.0,
12423 p1.0,
12424 p2.0,
12425 aq,
12426 ad,
12427 m,
12428 w0.in_features(),
12429 p0.1,
12430 p1.1,
12431 p2.1,
12432 p0.2,
12433 )?))
12434 }
12435
12436 #[allow(clippy::too_many_arguments)]
12437 fn q8_fused3_t_core(
12438 &self,
12439 b0: &CudaSlice<u8>,
12440 b1: &CudaSlice<u8>,
12441 b2: &CudaSlice<u8>,
12442 aq: &CudaSlice<i8>,
12443 ad: &CudaSlice<f32>,
12444 m: usize,
12445 in_f: usize,
12446 out0: usize,
12447 out1: usize,
12448 out2: usize,
12449 row_bytes: usize,
12450 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12451 const ROWS_PER_BLOCK: u32 = 4;
12452 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12453 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12454 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12455 let f = self.func(if Self::batched_mcols(m) == 2 {
12456 "qmatvec_q8_0_mmvq_fused3_b2"
12457 } else {
12458 "qmatvec_q8_0_mmvq_fused3_b4"
12459 });
12460 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12461 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12462 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
12463 let cfg = LaunchConfig {
12464 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12465 block_dim: (32, ROWS_PER_BLOCK, 1),
12466 shared_mem_bytes: 0,
12467 };
12468 let (inf, o0, o1, o2, mi, rbl) = (
12469 in_f as i32,
12470 out0 as i32,
12471 out1 as i32,
12472 out2 as i32,
12473 m as i32,
12474 row_bytes as i64,
12475 );
12476 let __s_b = self.gpu.stream();
12477 let mut b = __s_b.launch_builder(&f);
12478 b.arg(b0)
12479 .arg(b1)
12480 .arg(b2)
12481 .arg(aq)
12482 .arg(ad)
12483 .arg(&mut y0)
12484 .arg(&mut y1)
12485 .arg(&mut y2)
12486 .arg(&inf)
12487 .arg(&o0)
12488 .arg(&o1)
12489 .arg(&o2)
12490 .arg(&mi)
12491 .arg(&rbl);
12492 unsafe {
12493 b.launch(cfg)?;
12494 }
12495 Ok((y0, y1, y2))
12496 }
12497
12498 #[allow(clippy::too_many_arguments)]
12500 pub fn qmatvec_q8_fused3_t_raw(
12501 &self,
12502 b0: &CudaSlice<u8>,
12503 b1: &CudaSlice<u8>,
12504 b2: &CudaSlice<u8>,
12505 x: &CudaSlice<f32>,
12506 m: usize,
12507 in_f: usize,
12508 out0: usize,
12509 out1: usize,
12510 out2: usize,
12511 row_bytes: usize,
12512 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12513 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12514 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
12515 }
12516
12517 pub fn q8_ffn_fuse2_on(&self) -> bool {
12521 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12522 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
12523 }
12524
12525 #[allow(clippy::type_complexity)]
12531 fn q8_fused_params<'w, const N: usize>(
12532 &self,
12533 ws: &[&'w crate::model::GpuTensor; N],
12534 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
12535 use crate::model::GpuTensor;
12536 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
12537 return None;
12538 }
12539 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
12540 return None;
12541 }
12542 let in_f = ws[0].in_features();
12543 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
12544 for (i, w) in ws.iter().enumerate() {
12545 match w {
12546 GpuTensor::Quant {
12547 bytes,
12548 qtype,
12549 row_bytes,
12550 scale,
12551 ..
12552 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
12553 out[i] = Some((bytes, w.out_features(), *row_bytes))
12554 }
12555 _ => return None,
12556 }
12557 }
12558 Some(out.map(|o| o.unwrap()))
12559 }
12560
12561 pub fn e4m3_dual_on(&self) -> bool {
12564 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12565 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
12566 }
12567
12568 #[allow(clippy::type_complexity)]
12580 fn e4m3_fused_params<'w, const N: usize>(
12581 &self,
12582 ws: &[&'w crate::model::GpuTensor; N],
12583 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
12584 use crate::model::GpuTensor;
12585 if !self.e4m3_dual_on() {
12586 return None;
12587 }
12588 let in_f = ws[0].in_features();
12589 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
12590 for (i, w) in ws.iter().enumerate() {
12591 match w {
12592 GpuTensor::Quant {
12593 bytes,
12594 qtype,
12595 row_bytes,
12596 scale,
12597 rp,
12598 rp4,
12599 ..
12600 } if *qtype == QT_F8_E4M3
12601 && w.in_features() == in_f
12602 && *row_bytes == in_f
12603 && !*rp
12604 && rp4.is_none() =>
12605 {
12606 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
12607 }
12608 _ => return None,
12609 }
12610 }
12611 Some(out.map(|o| o.unwrap()))
12612 }
12613
12614 #[allow(clippy::too_many_arguments)]
12618 fn e4m3_fused2_core(
12619 &self,
12620 b0: &CudaSlice<u8>,
12621 b1: &CudaSlice<u8>,
12622 aq: &CudaSlice<i8>,
12623 ad: &CudaSlice<f32>,
12624 in_f: usize,
12625 out0: usize,
12626 out1: usize,
12627 row_bytes: usize,
12628 ws0: f32,
12629 ws1: f32,
12630 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12631 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12633 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12634 let f = self.func("qmatvec_e4m3_mmvq_fused2");
12635 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12636 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12637 let cfg = LaunchConfig {
12638 grid_dim: (nb0 + nb1, 1, 1),
12639 block_dim: (32, ROWS_PER_BLOCK, 1),
12640 shared_mem_bytes: 0,
12641 };
12642 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12643 let __s_b = self.gpu.stream();
12644 let mut b = __s_b.launch_builder(&f);
12645 b.arg(b0)
12646 .arg(b1)
12647 .arg(aq)
12648 .arg(ad)
12649 .arg(&mut y0)
12650 .arg(&mut y1)
12651 .arg(&inf)
12652 .arg(&o0)
12653 .arg(&o1)
12654 .arg(&rbl)
12655 .arg(&ws0)
12656 .arg(&ws1);
12657 unsafe {
12658 b.launch(cfg)?;
12659 }
12660 Ok((y0, y1))
12661 }
12662
12663 #[allow(clippy::too_many_arguments)]
12665 fn e4m3_fused3_core(
12666 &self,
12667 b0: &CudaSlice<u8>,
12668 b1: &CudaSlice<u8>,
12669 b2: &CudaSlice<u8>,
12670 aq: &CudaSlice<i8>,
12671 ad: &CudaSlice<f32>,
12672 in_f: usize,
12673 out0: usize,
12674 out1: usize,
12675 out2: usize,
12676 row_bytes: usize,
12677 ws0: f32,
12678 ws1: f32,
12679 ws2: f32,
12680 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12681 const ROWS_PER_BLOCK: u32 = 4;
12682 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12683 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12684 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12685 let f = self.func("qmatvec_e4m3_mmvq_fused3");
12686 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12687 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12688 let mut y2 = self.alloc_uninit::<f32>(out2)?;
12689 let cfg = LaunchConfig {
12690 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12691 block_dim: (32, ROWS_PER_BLOCK, 1),
12692 shared_mem_bytes: 0,
12693 };
12694 let (inf, o0, o1, o2, rbl) = (
12695 in_f as i32,
12696 out0 as i32,
12697 out1 as i32,
12698 out2 as i32,
12699 row_bytes as i64,
12700 );
12701 let __s_b = self.gpu.stream();
12702 let mut b = __s_b.launch_builder(&f);
12703 b.arg(b0)
12704 .arg(b1)
12705 .arg(b2)
12706 .arg(aq)
12707 .arg(ad)
12708 .arg(&mut y0)
12709 .arg(&mut y1)
12710 .arg(&mut y2)
12711 .arg(&inf)
12712 .arg(&o0)
12713 .arg(&o1)
12714 .arg(&o2)
12715 .arg(&rbl)
12716 .arg(&ws0)
12717 .arg(&ws1)
12718 .arg(&ws2);
12719 unsafe {
12720 b.launch(cfg)?;
12721 }
12722 Ok((y0, y1, y2))
12723 }
12724
12725 #[allow(clippy::too_many_arguments)]
12729 fn e4m3_fused2_t_core(
12730 &self,
12731 b0: &CudaSlice<u8>,
12732 b1: &CudaSlice<u8>,
12733 aq: &CudaSlice<i8>,
12734 ad: &CudaSlice<f32>,
12735 m: usize,
12736 in_f: usize,
12737 out0: usize,
12738 out1: usize,
12739 row_bytes: usize,
12740 ws0: f32,
12741 ws1: f32,
12742 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12743 const ROWS_PER_BLOCK: u32 = 4;
12744 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12745 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12746 let f = self.func(match Self::batched_mcols(m) {
12747 2 => "qmatvec_e4m3_mmvq_fused2_b2",
12748 4 => "qmatvec_e4m3_mmvq_fused2_b4",
12749 _ => "qmatvec_e4m3_mmvq_fused2_b8",
12750 });
12751 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12752 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12753 let cfg = LaunchConfig {
12754 grid_dim: (nb0 + nb1, 1, 1),
12755 block_dim: (32, ROWS_PER_BLOCK, 1),
12756 shared_mem_bytes: 0,
12757 };
12758 let (inf, o0, o1, mi, rbl) = (
12759 in_f as i32,
12760 out0 as i32,
12761 out1 as i32,
12762 m as i32,
12763 row_bytes as i64,
12764 );
12765 let __s_b = self.gpu.stream();
12766 let mut b = __s_b.launch_builder(&f);
12767 b.arg(b0)
12768 .arg(b1)
12769 .arg(aq)
12770 .arg(ad)
12771 .arg(&mut y0)
12772 .arg(&mut y1)
12773 .arg(&inf)
12774 .arg(&o0)
12775 .arg(&o1)
12776 .arg(&mi)
12777 .arg(&rbl);
12778 unsafe {
12779 b.launch(cfg)?;
12780 }
12781 if ws0 != 1.0 {
12782 self.scale_inplace(&mut y0, ws0, m * out0)?;
12783 }
12784 if ws1 != 1.0 {
12785 self.scale_inplace(&mut y1, ws1, m * out1)?;
12786 }
12787 Ok((y0, y1))
12788 }
12789
12790 #[allow(clippy::too_many_arguments)]
12792 fn e4m3_fused3_t_core(
12793 &self,
12794 b0: &CudaSlice<u8>,
12795 b1: &CudaSlice<u8>,
12796 b2: &CudaSlice<u8>,
12797 aq: &CudaSlice<i8>,
12798 ad: &CudaSlice<f32>,
12799 m: usize,
12800 in_f: usize,
12801 out0: usize,
12802 out1: usize,
12803 out2: usize,
12804 row_bytes: usize,
12805 ws0: f32,
12806 ws1: f32,
12807 ws2: f32,
12808 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12809 const ROWS_PER_BLOCK: u32 = 4;
12810 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12811 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12812 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12813 let f = self.func(if Self::batched_mcols(m) == 2 {
12814 "qmatvec_e4m3_mmvq_fused3_b2"
12815 } else {
12816 "qmatvec_e4m3_mmvq_fused3_b4"
12817 });
12818 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12819 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12820 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
12821 let cfg = LaunchConfig {
12822 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12823 block_dim: (32, ROWS_PER_BLOCK, 1),
12824 shared_mem_bytes: 0,
12825 };
12826 let (inf, o0, o1, o2, mi, rbl) = (
12827 in_f as i32,
12828 out0 as i32,
12829 out1 as i32,
12830 out2 as i32,
12831 m as i32,
12832 row_bytes as i64,
12833 );
12834 let __s_b = self.gpu.stream();
12835 let mut b = __s_b.launch_builder(&f);
12836 b.arg(b0)
12837 .arg(b1)
12838 .arg(b2)
12839 .arg(aq)
12840 .arg(ad)
12841 .arg(&mut y0)
12842 .arg(&mut y1)
12843 .arg(&mut y2)
12844 .arg(&inf)
12845 .arg(&o0)
12846 .arg(&o1)
12847 .arg(&o2)
12848 .arg(&mi)
12849 .arg(&rbl);
12850 unsafe {
12851 b.launch(cfg)?;
12852 }
12853 if ws0 != 1.0 {
12854 self.scale_inplace(&mut y0, ws0, m * out0)?;
12855 }
12856 if ws1 != 1.0 {
12857 self.scale_inplace(&mut y1, ws1, m * out1)?;
12858 }
12859 if ws2 != 1.0 {
12860 self.scale_inplace(&mut y2, ws2, m * out2)?;
12861 }
12862 Ok((y0, y1, y2))
12863 }
12864
12865 pub fn qmatvec_e4m3_blk_mmvq(
12875 &self,
12876 bytes: &CudaSlice<u8>,
12877 aq: &CudaSlice<i8>,
12878 ad: &CudaSlice<f32>,
12879 scales: &CudaSlice<f32>,
12880 m: usize,
12881 in_f: usize,
12882 out_f: usize,
12883 row_bytes: usize,
12884 scale_cols: usize,
12885 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12886 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
12888 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
12889 )?;
12890 Ok(y)
12891 }
12892
12893 #[allow(clippy::too_many_arguments)]
12895 pub fn qmatvec_e4m3_blk_mmvq_into(
12896 &self,
12897 bytes: &CudaSlice<u8>,
12898 aq: &CudaSlice<i8>,
12899 ad: &CudaSlice<f32>,
12900 scales: &CudaSlice<f32>,
12901 m: usize,
12902 in_f: usize,
12903 out_f: usize,
12904 row_bytes: usize,
12905 scale_cols: usize,
12906 y: &mut CudaSlice<f32>,
12907 ) -> Result<(), Box<dyn std::error::Error>> {
12908 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
12910 let cfg = LaunchConfig {
12911 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
12912 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
12915 let (inf, outf, mi, rb, sc) = (
12916 in_f as i32,
12917 out_f as i32,
12918 m as i32,
12919 row_bytes as i64,
12920 scale_cols as i32,
12921 );
12922 let __s_b = self.gpu.stream();
12923 let mut b = __s_b.launch_builder(&f);
12924 b.arg(bytes)
12925 .arg(aq)
12926 .arg(ad)
12927 .arg(scales)
12928 .arg(&mut *y)
12929 .arg(&inf)
12930 .arg(&outf)
12931 .arg(&mi)
12932 .arg(&rb)
12933 .arg(&sc);
12934 unsafe {
12935 b.launch(cfg)?;
12936 }
12937 Ok(())
12938 }
12939
12940 #[allow(clippy::too_many_arguments)]
12946 pub fn qmatvec_e4m3_blk_mmvq_batched(
12947 &self,
12948 bytes: &CudaSlice<u8>,
12949 aq: &CudaSlice<i8>,
12950 ad: &CudaSlice<f32>,
12951 scales: &CudaSlice<f32>,
12952 m: usize,
12953 in_f: usize,
12954 out_f: usize,
12955 row_bytes: usize,
12956 scale_cols: usize,
12957 mcols: usize,
12958 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12959 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
12961 let name = match mcols {
12962 2 => "qmatvec_e4m3_blk_mmvq_b2",
12963 4 => "qmatvec_e4m3_blk_mmvq_b4",
12964 8 => "qmatvec_e4m3_blk_mmvq_b8",
12965 16 => "qmatvec_e4m3_blk_mmvq_b16",
12966 _ => {
12967 return Err(
12968 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
12969 );
12970 }
12971 };
12972 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12973 let f = self.func(name);
12974 let cfg = LaunchConfig {
12975 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
12976 block_dim: (32, ROWS_PER_BLOCK, 1),
12977 shared_mem_bytes: 0,
12978 };
12979 let (inf, outf, mi, rb, sc) = (
12980 in_f as i32,
12981 out_f as i32,
12982 m as i32,
12983 row_bytes as i64,
12984 scale_cols as i32,
12985 );
12986 let __s_b = self.gpu.stream();
12987 let mut b = __s_b.launch_builder(&f);
12988 b.arg(bytes)
12989 .arg(aq)
12990 .arg(ad)
12991 .arg(scales)
12992 .arg(&mut y)
12993 .arg(&inf)
12994 .arg(&outf)
12995 .arg(&mi)
12996 .arg(&rb)
12997 .arg(&sc);
12998 unsafe {
12999 b.launch(cfg)?;
13000 }
13001 Ok(y)
13002 }
13003
13004 #[allow(clippy::too_many_arguments)]
13007 pub fn qmatvec_e4m3_blk_batched_raw(
13008 &self,
13009 bytes: &CudaSlice<u8>,
13010 x: &CudaSlice<f32>,
13011 scales: &CudaSlice<f32>,
13012 m: usize,
13013 in_f: usize,
13014 out_f: usize,
13015 row_bytes: usize,
13016 scale_cols: usize,
13017 mcols: usize,
13018 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13019 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13020 self.qmatvec_e4m3_blk_mmvq_batched(
13021 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13022 )
13023 }
13024
13025 #[allow(clippy::too_many_arguments)]
13028 pub fn qmatvec_e4m3_blk_mmvq_raw(
13029 &self,
13030 bytes: &CudaSlice<u8>,
13031 x: &CudaSlice<f32>,
13032 scales: &CudaSlice<f32>,
13033 m: usize,
13034 in_f: usize,
13035 out_f: usize,
13036 row_bytes: usize,
13037 scale_cols: usize,
13038 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13039 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13040 self.qmatvec_e4m3_blk_mmvq(
13041 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13042 )
13043 }
13044
13045 #[allow(clippy::too_many_arguments)]
13048 pub fn qmatvec_e4m3_fused2_raw(
13049 &self,
13050 b0: &CudaSlice<u8>,
13051 b1: &CudaSlice<u8>,
13052 x: &CudaSlice<f32>,
13053 in_f: usize,
13054 out0: usize,
13055 out1: usize,
13056 row_bytes: usize,
13057 ws0: f32,
13058 ws1: f32,
13059 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13060 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13061 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13062 }
13063
13064 #[allow(clippy::too_many_arguments)]
13065 pub fn qmatvec_e4m3_fused3_raw(
13066 &self,
13067 b0: &CudaSlice<u8>,
13068 b1: &CudaSlice<u8>,
13069 b2: &CudaSlice<u8>,
13070 x: &CudaSlice<f32>,
13071 in_f: usize,
13072 out0: usize,
13073 out1: usize,
13074 out2: usize,
13075 row_bytes: usize,
13076 ws0: f32,
13077 ws1: f32,
13078 ws2: f32,
13079 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13080 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13081 self.e4m3_fused3_core(
13082 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13083 )
13084 }
13085
13086 #[allow(clippy::too_many_arguments)]
13087 pub fn qmatvec_e4m3_fused2_t_raw(
13088 &self,
13089 b0: &CudaSlice<u8>,
13090 b1: &CudaSlice<u8>,
13091 x: &CudaSlice<f32>,
13092 m: usize,
13093 in_f: usize,
13094 out0: usize,
13095 out1: usize,
13096 row_bytes: usize,
13097 ws0: f32,
13098 ws1: f32,
13099 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13100 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13101 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13102 }
13103
13104 #[allow(clippy::too_many_arguments)]
13105 pub fn qmatvec_e4m3_fused3_t_raw(
13106 &self,
13107 b0: &CudaSlice<u8>,
13108 b1: &CudaSlice<u8>,
13109 b2: &CudaSlice<u8>,
13110 x: &CudaSlice<f32>,
13111 m: usize,
13112 in_f: usize,
13113 out0: usize,
13114 out1: usize,
13115 out2: usize,
13116 row_bytes: usize,
13117 ws0: f32,
13118 ws1: f32,
13119 ws2: f32,
13120 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13121 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13122 self.e4m3_fused3_t_core(
13123 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13124 )
13125 }
13126
13127 fn try_e4m3_blk_pre(
13138 &self,
13139 w: &crate::model::GpuTensor,
13140 aq: &CudaSlice<i8>,
13141 ad: &CudaSlice<f32>,
13142 m: usize,
13143 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13144 use crate::model::GpuTensor;
13145 if let GpuTensor::Quant {
13146 bytes,
13147 qtype,
13148 row_bytes,
13149 blk: Some(g),
13150 ..
13151 } = w
13152 {
13153 if *qtype == QT_F8_E4M3_BLK {
13154 if (2..=16).contains(&m)
13160 && std::env::var("MEMRA_NO_BATCHED").is_err()
13161 && (m <= 4 || Self::b8_enabled())
13162 {
13163 let mcols = Self::batched_mcols(m);
13164 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13165 bytes,
13166 aq,
13167 ad,
13168 &g.scales,
13169 m,
13170 w.in_features(),
13171 w.out_features(),
13172 *row_bytes,
13173 g.cols,
13174 mcols,
13175 )?));
13176 }
13177 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13178 bytes,
13179 aq,
13180 ad,
13181 &g.scales,
13182 m,
13183 w.in_features(),
13184 w.out_features(),
13185 *row_bytes,
13186 g.cols,
13187 )?));
13188 }
13189 }
13190 Ok(None)
13191 }
13192
13193 fn try_e4m3_blk_prefill(
13240 &self,
13241 w: &crate::model::GpuTensor,
13242 x: &CudaSlice<f32>,
13243 m: usize,
13244 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13245 use crate::model::GpuTensor;
13246 let GpuTensor::Quant {
13247 bytes,
13248 qtype,
13249 blk: Some(g),
13250 ..
13251 } = w
13252 else {
13253 return Ok(None);
13254 };
13255 if *qtype != QT_F8_E4M3_BLK {
13256 return Ok(None);
13257 }
13258 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
13263 return Ok(Some(y));
13264 }
13265 let (in_f, out_f) = (w.in_features(), w.out_features());
13266 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
13267 let tmp = GpuTensor::Quant {
13268 bytes: slab,
13269 qtype: QT_Q8_0,
13270 row_bytes: in_f / 32 * 34,
13271 ne: vec![in_f as u64, out_f as u64],
13272 scale: 1.0,
13273 rp: false,
13274 #[cfg(memra_cutlass)]
13275 cutlass: None,
13276 fp8: None,
13277 blk: None,
13278 f16: None,
13279 rp4: None,
13280 };
13281 Ok(Some(self.matmul(&tmp, x, m)?))
13283 }
13284
13285 pub fn matmul_pre_noscale(
13286 &self,
13287 w: &crate::model::GpuTensor,
13288 aq: &CudaSlice<i8>,
13289 ad: &CudaSlice<f32>,
13290 m: usize,
13291 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
13292 use crate::model::GpuTensor;
13293 if m == 1 {
13297 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
13298 return Ok(Some((y, 1.0)));
13299 }
13300 }
13301 if m != 1 || !self.uses_q8_1_fast(w) {
13303 return Ok(None);
13304 }
13305 let in_f = w.in_features();
13306 let out_f = w.out_features();
13307 let (bytes, qtype, row_bytes, scale, rp) = match w {
13308 GpuTensor::Quant {
13309 bytes,
13310 qtype,
13311 row_bytes,
13312 scale,
13313 rp,
13314 ..
13315 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13316 _ => return Ok(None),
13317 };
13318 if self.mmvq_supports(qtype) {
13320 let (mbytes, mrp) = match w {
13322 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
13323 _ => (bytes, rp),
13324 };
13325 let y = self.qmatvec_mmvq(
13326 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
13327 )?;
13328 return Ok(Some((y, scale)));
13329 }
13330 let name = match qtype {
13332 QT_Q8_0 => "qmatvec_q8_0_dp4a",
13333 QT_Q4_K => "qmatvec_q4_K_dp4a",
13334 QT_Q6_K => "qmatvec_q6_K_dp4a",
13335 QT_Q5_K => "qmatvec_q5_K_dp4a",
13336 QT_Q3_K => "qmatvec_q3_K_dp4a",
13337 QT_NVFP4 => {
13338 if rp {
13339 "qmatvec_nvfp4_dp4a_rp"
13340 } else {
13341 "qmatvec_nvfp4_dp4a"
13342 }
13343 }
13344 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
13345 _ => return Ok(None),
13346 };
13347 let f = self.func(name);
13348 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13349 let cfg = LaunchConfig {
13350 grid_dim: (out_f as u32, m as u32, 1),
13351 block_dim: (128, 1, 1),
13352 shared_mem_bytes: 0,
13353 };
13354 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13355 let __s_b = self.gpu.stream();
13356 let mut b = __s_b.launch_builder(&f);
13357 b.arg(bytes)
13358 .arg(aq)
13359 .arg(ad)
13360 .arg(&mut y)
13361 .arg(&inf)
13362 .arg(&outf)
13363 .arg(&mi)
13364 .arg(&rb);
13365 unsafe {
13366 b.launch(cfg)?;
13367 }
13368 Ok(Some((y, scale)))
13369 }
13370
13371 pub fn mmvq_supports(&self, qtype: i32) -> bool {
13374 if qtype == QT_F8_E4M3 {
13379 return true;
13380 }
13381 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13382 return false;
13383 }
13384 matches!(
13385 qtype,
13386 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
13387 )
13388 }
13389
13390 pub fn qmatvec_mmvq(
13395 &self,
13396 bytes: &CudaSlice<u8>,
13397 aq: &CudaSlice<i8>,
13398 ad: &CudaSlice<f32>,
13399 m: usize,
13400 in_f: usize,
13401 out_f: usize,
13402 qtype: i32,
13403 row_bytes: usize,
13404 scale: f32,
13405 rp: bool,
13406 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13407 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
13409 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
13410 )?;
13411 Ok(y)
13412 }
13413
13414 #[allow(clippy::too_many_arguments)]
13416 pub fn qmatvec_mmvq_into(
13417 &self,
13418 bytes: &CudaSlice<u8>,
13419 aq: &CudaSlice<i8>,
13420 ad: &CudaSlice<f32>,
13421 m: usize,
13422 in_f: usize,
13423 out_f: usize,
13424 qtype: i32,
13425 row_bytes: usize,
13426 scale: f32,
13427 rp: bool,
13428 y: &mut CudaSlice<f32>,
13429 ) -> Result<(), Box<dyn std::error::Error>> {
13430 debug_assert!(y.len() >= m * out_f);
13431 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
13437 && rp
13438 && m == 1
13439 && out_f >= 64
13440 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
13441 && {
13442 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13443 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
13444 }
13445 {
13446 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
13447 let cfg = LaunchConfig {
13448 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
13449 block_dim: (32, 2, 1),
13450 shared_mem_bytes: 0,
13451 };
13452 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
13453 let __s_b = self.gpu.stream();
13454 let mut b = __s_b.launch_builder(&f);
13455 b.arg(bytes)
13456 .arg(aq)
13457 .arg(ad)
13458 .arg(&mut *y)
13459 .arg(&inf)
13460 .arg(&outf)
13461 .arg(&mi)
13462 .arg(&rb);
13463 unsafe {
13464 b.launch(cfg)?;
13465 }
13466 if scale != 1.0 {
13467 self.scale_inplace(y, scale, out_f)?;
13468 }
13469 return Ok(());
13470 }
13471 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
13480 2
13481 } else {
13482 1
13483 };
13484 if m == 1 && qtype == QT_Q4_0 {
13489 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13490 mr = *Q40MR.get_or_init(|| {
13493 std::env::var("MEMRA_Q40_MR")
13494 .ok()
13495 .and_then(|v| v.parse().ok())
13496 .unwrap_or(1)
13497 });
13498 }
13499 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
13510 let q5_force = q5_mode.as_deref() == Some("2");
13511 let q5_il = qtype == QT_Q5_K
13514 && m == 1
13515 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
13516 if q5_il && !q5_force && out_f > 65536 {
13517 mr = 1;
13518 }
13519 if qtype == QT_Q4_0 && rp && mr != 1 {
13522 mr = 2;
13523 }
13524 if qtype == QT_Q8_0 && rp {
13528 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13529 mr = *Q80MR.get_or_init(|| {
13530 std::env::var("MEMRA_Q80_MR")
13531 .ok()
13532 .and_then(|v| v.parse().ok())
13533 .unwrap_or(1)
13534 });
13535 }
13536 let name = match (qtype, mr, rp) {
13537 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
13538 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
13539 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
13540 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
13541 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
13542 (QT_Q5_K, 2, _) => {
13543 if q5_il {
13544 "qmatvec_q5_K_mmvq_mr2_il"
13545 } else {
13546 "qmatvec_q5_K_mmvq_mr2"
13547 }
13548 }
13549 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
13550 (QT_Q8_0, _, true)
13555 if in_f % 1024 == 0 && {
13556 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13557 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
13558 } =>
13559 {
13560 "qmatvec_q8_0_mmvq_rpca"
13561 }
13562 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
13563 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
13564 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
13568 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
13569 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
13570 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
13571 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
13572 (QT_Q5_K, _, _) => {
13573 if q5_il {
13574 "qmatvec_q5_K_mmvq_il"
13575 } else {
13576 "qmatvec_q5_K_mmvq"
13577 }
13578 }
13579 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
13580 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
13581 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
13582 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
13583 };
13584 let f = self.func(name);
13585 let rows_per_block = ROWS_PER_BLOCK * mr;
13587 let cfg = LaunchConfig {
13588 grid_dim: (
13589 (out_f as u32 + rows_per_block - 1) / rows_per_block,
13590 m as u32,
13591 1,
13592 ),
13593 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
13596 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13597 let __s_b = self.gpu.stream();
13598 let mut b = __s_b.launch_builder(&f);
13599 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
13604 if Self::pdl_on()
13607 && Self::pdl_mmvq_on()
13608 && Self::pdl_nvfp4q8_on()
13609 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
13610 {
13611 use cudarc::driver::{DevicePtr, DevicePtrMut};
13612 let s = &self.gpu.stream();
13613 let (pw, _g0) = bytes.device_ptr(s);
13614 let (paq, _g1) = aq.device_ptr(s);
13615 let (pad, _g2) = ad.device_ptr(s);
13616 let (py, _g3) = y.device_ptr_mut(s);
13617 let mut ps = [
13618 &pw as *const _ as *mut std::ffi::c_void,
13619 &paq as *const _ as *mut _,
13620 &pad as *const _ as *mut _,
13621 &py as *const _ as *mut _,
13622 &inf as *const _ as *mut _,
13623 &outf as *const _ as *mut _,
13624 &mi as *const _ as *mut _,
13625 &rb as *const _ as *mut _,
13626 &scale as *const _ as *mut _,
13627 ];
13628 unsafe {
13629 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13630 }
13631 return Ok(());
13632 }
13633 b.arg(bytes)
13634 .arg(aq)
13635 .arg(ad)
13636 .arg(&mut *y)
13637 .arg(&inf)
13638 .arg(&outf)
13639 .arg(&mi)
13640 .arg(&rb)
13641 .arg(&scale);
13642 unsafe {
13643 b.launch(cfg)?;
13644 }
13645 } else if Self::pdl_on()
13646 && Self::pdl_mmvq_on()
13647 && (matches!(
13648 name,
13649 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
13650 ) || (Self::pdl_nvfp4q8_on()
13651 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
13652 {
13653 {
13657 use cudarc::driver::{DevicePtr, DevicePtrMut};
13658 let s = &self.gpu.stream();
13659 let (pw, _g0) = bytes.device_ptr(s);
13660 let (paq, _g1) = aq.device_ptr(s);
13661 let (pad, _g2) = ad.device_ptr(s);
13662 let (py, _g3) = y.device_ptr_mut(s);
13663 let mut ps = [
13664 &pw as *const _ as *mut std::ffi::c_void,
13665 &paq as *const _ as *mut _,
13666 &pad as *const _ as *mut _,
13667 &py as *const _ as *mut _,
13668 &inf as *const _ as *mut _,
13669 &outf as *const _ as *mut _,
13670 &mi as *const _ as *mut _,
13671 &rb as *const _ as *mut _,
13672 ];
13673 unsafe {
13674 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13675 }
13676 }
13677 if scale != 1.0 {
13678 self.scale_inplace(y, scale, m * out_f)?;
13679 }
13680 } else {
13681 b.arg(bytes)
13682 .arg(aq)
13683 .arg(ad)
13684 .arg(&mut *y)
13685 .arg(&inf)
13686 .arg(&outf)
13687 .arg(&mi)
13688 .arg(&rb);
13689 unsafe {
13690 b.launch(cfg)?;
13691 }
13692 if scale != 1.0 {
13693 self.scale_inplace(y, scale, m * out_f)?;
13694 }
13695 }
13696 Ok(())
13697 }
13698
13699 pub fn qmatvec_mmvq_raw(
13703 &self,
13704 bytes: &CudaSlice<u8>,
13705 x: &CudaSlice<f32>,
13706 m: usize,
13707 in_f: usize,
13708 out_f: usize,
13709 qtype: i32,
13710 row_bytes: usize,
13711 rp: bool,
13712 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13713 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13714 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
13715 }
13716
13717 pub fn batched_supports(&self, qtype: i32) -> bool {
13721 matches!(
13722 qtype,
13723 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
13724 )
13725 }
13726
13727 pub fn iq_fast_enabled() -> bool {
13735 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13736 *ON.get_or_init(|| {
13737 std::env::var("MEMRA_IQ_FAST")
13738 .map(|v| v != "0")
13739 .unwrap_or(true)
13740 })
13741 }
13742
13743 pub fn b8_enabled() -> bool {
13746 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13747 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
13748 }
13749
13750 pub fn batched_mcols(m: usize) -> usize {
13752 if m == 2 {
13753 2
13754 } else if m <= 4 {
13755 4
13756 } else if m <= 8 {
13757 8
13758 } else {
13759 16
13760 }
13761 }
13762
13763 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
13768 Some(match (qtype, mcols) {
13769 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
13770 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
13771 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
13772 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
13778 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
13779 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
13780 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
13781 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
13784 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
13785 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
13786 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
13787 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
13790 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
13791 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
13792 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
13793 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
13794 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
13795 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
13796 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
13797 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
13801 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
13802 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
13803 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
13804 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
13808 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
13809 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
13810 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
13811 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
13812 _ => return None,
13813 })
13814 }
13815
13816 pub fn sm_count(&self) -> i32 {
13851 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
13852 *SMS.get_or_init(|| {
13853 use cudarc::driver::sys::CUdevice_attribute_enum as A;
13854 self.gpu
13855 .ctx
13856 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
13857 .unwrap_or(82)
13858 })
13859 }
13860
13861 pub fn batched_variant(
13862 &self,
13863 _m: usize,
13864 in_f: usize,
13865 out_f: usize,
13866 qtype: i32,
13867 row_bytes: usize,
13868 mcols: usize,
13869 rp: bool,
13870 ) -> &'static str {
13871 if qtype == QT_Q8_0 {
13876 return if rp { "rp" } else { "base" };
13877 }
13878 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
13879 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
13880 Ok("base") => "base",
13881 Ok("pf") => "pf",
13882 Ok("r2") => "r2",
13883 Ok("r2w8") => "r2w8",
13884 Ok("pfr2") => "pfr2",
13885 Ok("ca") => "ca",
13886 Ok("car2") => "car2",
13887 Ok("rp") => "rp",
13890 Ok("rpr2") => "rpr2",
13891 Ok("rpr2w8") => "rpr2w8",
13892 Ok("rpca") => "rpca",
13895 Ok("rpcar2") => "rpcar2",
13896 Ok("rpsc") => "rpsc",
13903 Ok("rpms") => "rpms",
13904 Ok("rpmsc") => "rpmsc",
13905 Ok("rpks") => "rpks",
13906 Ok("rpksc") => "rpksc",
13907 _ => "auto",
13908 });
13909 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
13913 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13918 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
13919 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
13920 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
13921 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
13922 let sms = *SMS.get_or_init(|| {
13923 use cudarc::driver::sys::CUdevice_attribute_enum as A;
13924 self.gpu
13925 .ctx
13926 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
13927 .unwrap_or(82)
13928 });
13929 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
13949 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
13952 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
13953 Ok("base") => "base",
13954 Ok("r2") => "r2",
13955 Ok("r2w8") => "r2w8",
13956 _ => "auto",
13957 });
13958 let variant: &'static str = if qtype == QT_Q4_0 {
13959 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
13963 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
13964 Ok("base") => "base",
13970 Ok("r2") => "r2",
13971 Ok("ms") => "ms",
13972 Ok("sm") => "sm",
13973 Ok("la") => "la",
13974 _ => "auto",
13975 });
13976 let v = if q40 != "auto" {
13977 q40
13978 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
13979 "r2"
13980 } else {
13981 "base"
13982 };
13983 if rp {
13988 match v {
13989 "ms" => "r2ms_rp",
13990 "sm" => "r2sm_rp",
13991 "la" => "r2la_rp",
13992 "r2" => "r2_rp",
13993 _ => "rp",
13994 }
13995 } else if matches!(v, "ms" | "sm" | "la") {
13996 "r2"
13997 } else {
13998 v
13999 }
14000 } else if qtype != QT_NVFP4 && !kq_r2 {
14001 "base"
14002 } else if kq_r2 && rp {
14003 "rp"
14007 } else if kq_r2 {
14008 if kq_bv != "auto" {
14011 if kq_bv == "r2w8" && mcols != 4 {
14012 "r2"
14013 } else {
14014 kq_bv
14015 }
14016 } else if bv != "auto" {
14017 match bv {
14018 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14019 "r2w8" | "rpr2w8" => {
14020 if mcols != 4 {
14021 "r2"
14022 } else {
14023 "r2w8"
14024 }
14025 }
14026 _ => "base", }
14028 } else {
14029 let blocks = (out_f + 7) / 8;
14030 let waves = blocks as f64 / (7 * sms as usize) as f64;
14031 let filled = blocks >= 4 * sms as usize;
14032 let use_r2 = if qtype == QT_Q4_K {
14033 filled
14034 } else {
14035 waves >= 2.0
14036 };
14037 if use_r2 { "r2" } else { "base" }
14038 }
14039 } else if bv != "auto" {
14040 let v = if bv == "r2w8" && mcols == 2 {
14045 "r2"
14046 } else if bv == "ca" && (!ca_ok || mcols == 8) {
14047 "pf"
14048 } else if bv == "car2" && (!ca_ok || mcols == 8) {
14049 "r2"
14050 } else if bv == "pfr2" && mcols == 8 {
14051 "r2"
14052 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14053 "rpr2"
14054 }
14055 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14057 if mcols == 8 { "rpr2w8" } else { "rpr2" }
14058 } else if bv == "rpcar2" && mcols == 2 {
14059 "rpca"
14060 }
14061 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14064 "rpr2"
14065 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14066 "rpr2"
14067 } else {
14068 bv
14069 };
14070 if rp {
14071 match v {
14072 "base" | "pf" | "ca" | "rp" => "rp",
14073 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14074 "r2w8" | "rpr2w8" => {
14075 if mcols == 2 {
14076 "rpr2"
14077 } else {
14078 "rpr2w8"
14079 }
14080 }
14081 other => other, }
14083 } else {
14084 v
14085 }
14086 } else if mcols == 8 {
14087 if rp {
14098 if sc_ok { "rpsc" } else { "rpr2w8" }
14099 } else {
14100 "r2w8"
14101 }
14102 } else if mcols >= 4 {
14103 let blocks = (out_f + 7) / 8;
14107 let r7 = 7 * sms as usize;
14108 let r8 = 8 * sms as usize;
14109 let waves = blocks as f64 / r7 as f64;
14110 let filled = blocks >= 4 * sms as usize;
14111 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14115 if rp { "rpr2w8" } else { "r2w8" }
14119 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14120 if rp { "rpr2" } else { "r2" }
14123 } else {
14124 if rp { "rp" } else { "pf" }
14128 }
14129 } else if in_f >= 6144 {
14130 if rp { "rpr2" } else { "r2" }
14134 } else if rp {
14135 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14140 if sc_ok && waves >= 0.9 && waves <= 1.1 {
14141 "rpsc"
14142 } else {
14143 "rp"
14144 }
14145 } else {
14146 "base"
14147 };
14148 variant
14149 }
14150
14151 pub fn qmatvec_mmvq_batched(
14152 &self,
14153 bytes: &CudaSlice<u8>,
14154 aq: &CudaSlice<i8>,
14155 ad: &CudaSlice<f32>,
14156 m: usize,
14157 in_f: usize,
14158 out_f: usize,
14159 qtype: i32,
14160 row_bytes: usize,
14161 mcols: usize,
14162 scale: f32,
14163 rp: bool,
14164 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14165 const ROWS_PER_BLOCK: u32 = 4;
14166 let forced: Option<&'static str> = {
14171 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14172 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14173 .as_deref()
14174 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14175 };
14176 let variant = match forced {
14177 Some(v) if !rp || v.contains("rp") => v,
14178 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14179 };
14180 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14181 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14182 })?;
14183 let variant = if mcols == 16 {
14187 if rp { "rp" } else { "base" }
14188 } else {
14189 variant
14190 };
14191 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14198 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14199 if b567
14200 && qtype == QT_NVFP4
14201 && rp
14202 && mcols == 8
14203 && (5..=7).contains(&m)
14204 && matches!(variant, "rpsc" | "rpr2w8")
14205 {
14206 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14207 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14209 let cfg = LaunchConfig {
14210 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14211 block_dim: (32, ROWS_PER_BLOCK, 1),
14212 shared_mem_bytes: 0,
14213 };
14214 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14215 let __s_b = self.gpu.stream();
14216 let mut b = __s_b.launch_builder(&f);
14217 b.arg(bytes)
14218 .arg(aq)
14219 .arg(ad)
14220 .arg(&mut y)
14221 .arg(&inf)
14222 .arg(&outf)
14223 .arg(&mi)
14224 .arg(&rb);
14225 unsafe {
14226 b.launch(cfg)?;
14227 }
14228 if scale != 1.0 {
14229 self.scale_inplace(&mut y, scale, m * out_f)?;
14230 }
14231 return Ok(y);
14232 }
14233 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14234 "base" => (base_name.into(), ROWS_PER_BLOCK),
14235 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14236 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14237 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14238 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14242 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14243 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14244 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14245 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14246 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14247 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14248 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
14250 debug_assert!(
14251 !rp || name.contains("_rp"),
14252 "rp weight dispatched to a GGUF-layout kernel"
14253 );
14254 let f = self.func(&name);
14255 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14256 let smem = if name.contains("_r2sm_rp") {
14258 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
14259 } else {
14260 0
14261 };
14262 let cfg = LaunchConfig {
14263 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14264 block_dim: (32, ROWS_PER_BLOCK, 1),
14265 shared_mem_bytes: smem,
14266 };
14267 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14268 let __s_b = self.gpu.stream();
14269 let mut b = __s_b.launch_builder(&f);
14270 b.arg(bytes)
14271 .arg(aq)
14272 .arg(ad)
14273 .arg(&mut y)
14274 .arg(&inf)
14275 .arg(&outf)
14276 .arg(&mi)
14277 .arg(&rb);
14278 unsafe {
14279 b.launch(cfg)?;
14280 }
14281 if scale != 1.0 {
14282 self.scale_inplace(&mut y, scale, m * out_f)?;
14283 }
14284 Ok(y)
14285 }
14286
14287 pub fn qmatvec_batched_raw(
14291 &self,
14292 bytes: &CudaSlice<u8>,
14293 x: &CudaSlice<f32>,
14294 m: usize,
14295 in_f: usize,
14296 out_f: usize,
14297 qtype: i32,
14298 row_bytes: usize,
14299 mcols: usize,
14300 rp: bool,
14301 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14302 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14303 self.qmatvec_mmvq_batched(
14304 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
14305 )
14306 }
14307
14308 pub fn qmatvec_nvfp4_batched_raw(
14310 &self,
14311 bytes: &CudaSlice<u8>,
14312 x: &CudaSlice<f32>,
14313 m: usize,
14314 in_f: usize,
14315 out_f: usize,
14316 row_bytes: usize,
14317 mcols: usize,
14318 rp: bool,
14319 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14320 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
14321 }
14322
14323 fn try_fp4_gemm(
14327 &self,
14328 w: &crate::model::GpuTensor,
14329 x: &CudaSlice<f32>,
14330 m: usize,
14331 in_f: usize,
14332 out_f: usize,
14333 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14334 use crate::model::GpuTensor;
14335 if cfg!(memra_portable_cuda) {
14336 return Ok(None);
14337 }
14338 if std::env::var("MEMRA_FP4").is_err() {
14339 return Ok(None);
14340 }
14341 #[cfg(memra_cutlass)]
14350 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
14351 if let GpuTensor::Quant {
14352 bytes,
14353 qtype,
14354 scale,
14355 row_bytes,
14356 cutlass,
14357 ..
14358 } = w
14359 {
14360 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
14361 if let Some(cw) = cutlass {
14362 let y = self.cutlass_fp4_gemm(
14364 &cw.b_packed,
14365 &cw.sfb_swizzled,
14366 x,
14367 *scale,
14368 m,
14369 out_f,
14370 in_f,
14371 )?;
14372 return Ok(Some(y));
14373 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
14374 let (b_packed, sfb_sw) =
14379 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
14380 let y =
14381 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
14382 return Ok(Some(y));
14383 }
14384 }
14385 }
14386 }
14387 if let GpuTensor::Quant {
14388 bytes,
14389 qtype,
14390 row_bytes,
14391 scale,
14392 rp,
14393 ..
14394 } = w
14395 {
14396 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
14399 let y =
14400 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
14401 return Ok(Some(y));
14402 }
14403 }
14404 Ok(None)
14405 }
14406
14407 pub fn rms_norm_f16out(
14411 &self,
14412 x: &CudaSlice<f32>,
14413 w: &CudaSlice<f32>,
14414 dst: &mut CudaSlice<f32>,
14415 dst16: &mut CudaSlice<u8>,
14416 ncols: usize,
14417 nrows: usize,
14418 eps: f32,
14419 ) -> Result<(), Box<dyn std::error::Error>> {
14420 let f = self.func("rms_norm_f16out_f32");
14421 let cfg = LaunchConfig {
14422 grid_dim: (nrows as u32, 1, 1),
14423 block_dim: (rms_block(), 1, 1),
14424 shared_mem_bytes: 0,
14425 };
14426 let (nc, e) = (ncols as i32, eps);
14427 let __s_b = self.gpu.stream();
14428 let mut b = __s_b.launch_builder(&f);
14429 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
14430 unsafe {
14431 b.launch(cfg)?;
14432 }
14433 Ok(())
14434 }
14435
14436 #[allow(clippy::too_many_arguments)]
14439 pub fn add_rms_norm_f16out(
14440 &self,
14441 a: &CudaSlice<f32>,
14442 b: &CudaSlice<f32>,
14443 w: &CudaSlice<f32>,
14444 res: &mut CudaSlice<f32>,
14445 dst: &mut CudaSlice<f32>,
14446 dst16: &mut CudaSlice<u8>,
14447 ncols: usize,
14448 nrows: usize,
14449 eps: f32,
14450 ) -> Result<(), Box<dyn std::error::Error>> {
14451 let f = self.func("add_rms_norm_f16out_f32");
14452 let cfg = LaunchConfig {
14453 grid_dim: (nrows as u32, 1, 1),
14454 block_dim: (rms_block(), 1, 1),
14455 shared_mem_bytes: 0,
14456 };
14457 let (nc, e) = (ncols as i32, eps);
14458 let __s_lb = self.gpu.stream();
14459 let mut lb = __s_lb.launch_builder(&f);
14460 lb.arg(a)
14461 .arg(b)
14462 .arg(w)
14463 .arg(res)
14464 .arg(dst)
14465 .arg(dst16)
14466 .arg(&nc)
14467 .arg(&e);
14468 unsafe {
14469 lb.launch(cfg)?;
14470 }
14471 Ok(())
14472 }
14473
14474 pub fn matmul_group_xh(
14477 &self,
14478 ws: &[&crate::model::GpuTensor],
14479 x: &CudaSlice<f32>,
14480 xh: &CudaSlice<u8>,
14481 m: usize,
14482 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14483 let mut out = Vec::with_capacity(ws.len());
14484 let in_f = ws[0].in_features();
14485 for w in ws {
14486 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
14487 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
14488 out.push(y);
14489 continue;
14490 }
14491 }
14492 out.push(self.matmul(w, x, m)?);
14493 }
14494 Ok(out)
14495 }
14496
14497 pub fn gdn_pad_mask(
14500 &self,
14501 beta: &mut CudaSlice<f32>,
14502 g_log: &mut CudaSlice<f32>,
14503 len_d: &CudaSlice<i32>,
14504 h: usize,
14505 t: usize,
14506 ) -> Result<(), Box<dyn std::error::Error>> {
14507 let f = self.func("gdn_pad_mask_f32");
14508 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
14509 let (hi, ti) = (h as i32, t as i32);
14510 let __s_b = self.gpu.stream();
14511 let mut b = __s_b.launch_builder(&f);
14512 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
14513 unsafe {
14514 b.launch(cfg)?;
14515 }
14516 Ok(())
14517 }
14518
14519 pub fn row_gather_dev(
14522 &self,
14523 src: &CudaSlice<f32>,
14524 dst: &mut CudaSlice<f32>,
14525 len_d: &CudaSlice<i32>,
14526 ncols: usize,
14527 ) -> Result<(), Box<dyn std::error::Error>> {
14528 let f = self.func("row_gather_dev_f32");
14529 let cfg = LaunchConfig::for_num_elems(ncols as u32);
14530 let nc = ncols as i32;
14531 let __s_b = self.gpu.stream();
14532 let mut b = __s_b.launch_builder(&f);
14533 b.arg(src).arg(dst).arg(len_d).arg(&nc);
14534 unsafe {
14535 b.launch(cfg)?;
14536 }
14537 Ok(())
14538 }
14539
14540 pub fn matmul_group(
14547 &self,
14548 ws: &[&crate::model::GpuTensor],
14549 x: &CudaSlice<f32>,
14550 m: usize,
14551 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14552 use crate::model::GpuTensor;
14553 let mut out = Vec::with_capacity(ws.len());
14554 let any_mirror = ws
14555 .iter()
14556 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
14557 if m >= 16 && any_mirror && !self.verify_exact_on() {
14558 let in_f = ws[0].in_features();
14559 let xh = self.f16_act(x, m * in_f, in_f)?;
14560 for w in ws {
14561 if w.in_features() == in_f {
14562 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
14563 out.push(y);
14564 continue;
14565 }
14566 }
14567 out.push(self.matmul(w, x, m)?);
14568 }
14569 return Ok(out);
14570 }
14571 for w in ws {
14572 out.push(self.matmul(w, x, m)?);
14573 }
14574 Ok(out)
14575 }
14576
14577 pub fn matmul_group_multi(
14584 &self,
14585 ws: &[&crate::model::GpuTensor],
14586 xs: &[&CudaSlice<f32>],
14587 ms: &[usize],
14588 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
14589 assert_eq!(xs.len(), ms.len());
14590 let in_f = ws[0].in_features();
14591 let total: usize = ms.iter().sum();
14592 let mut xcat = self.uninit(total * in_f)?;
14593 let mut off = 0usize;
14594 for (x, &m) in xs.iter().zip(ms) {
14595 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
14596 off += m;
14597 }
14598 let ys = self.matmul_group(ws, &xcat, total)?;
14599 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
14600 for (w, y) in ws.iter().zip(ys) {
14601 let out_f = w.out_features();
14602 let mut off = 0usize;
14603 for (s, &m) in ms.iter().enumerate() {
14604 let mut ys_s = self.uninit(m * out_f)?;
14605 let src = y.slice(off * out_f..(off + m) * out_f);
14606 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
14607 out[s].push(ys_s);
14608 off += m;
14609 }
14610 }
14611 Ok(out)
14612 }
14613
14614 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
14624 use crate::model::GpuTensor;
14625 if !legacy_quant_gemm_allowed(
14626 cfg!(memra_portable_cuda),
14627 cfg!(memra_hopper_mma),
14628 std::env::var_os("MEMRA_NO_GEMM").is_some(),
14629 ) {
14630 return false;
14631 }
14632 match w {
14633 GpuTensor::Quant { qtype, .. } => {
14634 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
14635 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
14636 }
14637 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
14638 }
14639 }
14640
14641 pub fn qmatvec_gemm(
14648 &self,
14649 w: &crate::model::GpuTensor,
14650 aq: &CudaSlice<i8>,
14651 ad: &CudaSlice<f32>,
14652 m: usize,
14653 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14654 use crate::model::GpuTensor;
14655 let in_f = w.in_features();
14656 let out_f = w.out_features();
14657 let (bytes, qtype, row_bytes, scale, rp) = match w {
14658 GpuTensor::Quant {
14659 bytes,
14660 qtype,
14661 row_bytes,
14662 scale,
14663 rp,
14664 ..
14665 } => (bytes, *qtype, *row_bytes, *scale, *rp),
14666 _ => unreachable!("gemm_supports guaranteed Quant"),
14667 };
14668 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
14674 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
14675 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
14676 if scale != 1.0 {
14677 self.scale_inplace(&mut y, scale, m * out_f)?;
14678 }
14679 return Ok(y);
14680 }
14681 }
14682 let name = match qtype {
14683 QT_Q8_0 => "qmatvec_gemm_q8_0",
14684 QT_Q4_K => "qmatvec_gemm_q4_K",
14685 QT_Q4_0 => {
14686 if rp {
14687 "qmatvec_gemm_q4_0_rp"
14688 } else {
14689 "qmatvec_gemm_q4_0"
14690 }
14691 }
14692 QT_Q5_K => "qmatvec_gemm_q5_K",
14693 QT_Q6_K => "qmatvec_gemm_q6_K",
14694 QT_NVFP4 => {
14695 if rp {
14696 "qmatvec_gemm_nvfp4_rp"
14697 } else {
14698 "qmatvec_gemm_nvfp4"
14699 }
14700 }
14701 _ => unreachable!(),
14702 };
14703 let f = self.func(name);
14704 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
14709 let k1_tile = if is_k1 {
14711 k1_launch_override().unwrap_or((128, 128, 8))
14712 } else {
14713 (128, 128, 8)
14714 };
14715 let (bm, bn): (u32, u32) = if is_k1 {
14716 (k1_tile.0, k1_tile.1)
14717 } else {
14718 (64, 256)
14719 };
14720 let warps: u32 = if is_k1 {
14721 k1_tile.2
14722 } else {
14723 match qtype {
14724 QT_NVFP4 => 8,
14725 _ => 4,
14726 }
14727 };
14728 let cfg = LaunchConfig {
14729 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
14730 block_dim: (32, warps, 1),
14731 shared_mem_bytes: 0,
14732 };
14733 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14734 let __s_b = self.gpu.stream();
14735 let mut b = __s_b.launch_builder(&f);
14736 b.arg(bytes)
14737 .arg(aq)
14738 .arg(ad)
14739 .arg(&mut y)
14740 .arg(&inf)
14741 .arg(&outf)
14742 .arg(&mi)
14743 .arg(&rb);
14744 unsafe {
14745 b.launch(cfg)?;
14746 }
14747 if scale != 1.0 {
14748 self.scale_inplace(&mut y, scale, m * out_f)?;
14749 }
14750 Ok(y)
14751 }
14752
14753 pub fn qmatvec_gemm_raw(
14758 &self,
14759 bytes: &CudaSlice<u8>,
14760 x: &CudaSlice<f32>,
14761 m: usize,
14762 in_f: usize,
14763 out_f: usize,
14764 qtype: i32,
14765 row_bytes: usize,
14766 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14767 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14768 let name = match qtype {
14769 QT_Q8_0 => "qmatvec_gemm_q8_0",
14770 QT_Q4_K => "qmatvec_gemm_q4_K",
14771 QT_Q4_0 => "qmatvec_gemm_q4_0",
14772 QT_Q5_K => "qmatvec_gemm_q5_K",
14773 QT_Q6_K => "qmatvec_gemm_q6_K",
14774 QT_NVFP4 => "qmatvec_gemm_nvfp4",
14775 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
14776 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
14777 };
14778 let f = self.func(name);
14779 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let is_k1 = matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q4_0);
14783 let k1_tile = if is_k1 {
14785 k1_launch_override().unwrap_or((128, 128, 8))
14786 } else {
14787 (128, 128, 8)
14788 };
14789 let (bm, bn): (u32, u32) = if is_k1 {
14790 (k1_tile.0, k1_tile.1)
14791 } else {
14792 (64, 256)
14793 };
14794 let warps: u32 = if is_k1 {
14795 k1_tile.2
14796 } else {
14797 match qtype {
14798 QT_NVFP4 | QT_NVFP4_RP => 8,
14799 _ => 4,
14800 }
14801 };
14802 let cfg = LaunchConfig {
14803 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
14804 block_dim: (32, warps, 1),
14805 shared_mem_bytes: 0,
14806 };
14807 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14808 let __s_b = self.gpu.stream();
14809 let mut b = __s_b.launch_builder(&f);
14810 b.arg(bytes)
14811 .arg(&aq)
14812 .arg(&ad)
14813 .arg(&mut y)
14814 .arg(&inf)
14815 .arg(&outf)
14816 .arg(&mi)
14817 .arg(&rb);
14818 unsafe {
14819 b.launch(cfg)?;
14820 }
14821 Ok(y)
14822 }
14823
14824 pub fn qmatvec_gemm_q8_0_wgmma_raw(
14831 &self,
14832 rp4: &CudaSlice<u8>,
14833 aq: &CudaSlice<i8>,
14834 ad: &CudaSlice<f32>,
14835 m: usize,
14836 in_f: usize,
14837 out_f: usize,
14838 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14839 assert!(
14840 out_f % 64 == 0 && in_f % 32 == 0,
14841 "wgmma GEMM needs out_f%64==0, in_f%32==0"
14842 );
14843 let f = self.func("qmatvec_gemm_q8_0_wgmma");
14844 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
14846 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
14847 block_dim: (128, 1, 1),
14848 shared_mem_bytes: 0,
14849 };
14850 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
14851 let __s_b = self.gpu.stream();
14852 let mut b = __s_b.launch_builder(&f);
14853 b.arg(rp4)
14854 .arg(aq)
14855 .arg(ad)
14856 .arg(&mut y)
14857 .arg(&inf)
14858 .arg(&outf)
14859 .arg(&mi);
14860 unsafe {
14861 b.launch(cfg)?;
14862 }
14863 Ok(y)
14864 }
14865
14866 pub fn scale_inplace(
14868 &self,
14869 y: &mut CudaSlice<f32>,
14870 s: f32,
14871 n: usize,
14872 ) -> Result<(), Box<dyn std::error::Error>> {
14873 let f = self.func("scale_f32");
14874 let cfg = LaunchConfig::for_num_elems(n as u32);
14875 let (sf, ni) = (s, n as i32);
14876 let __s_b = self.gpu.stream();
14877 let mut b = __s_b.launch_builder(&f);
14878 b.arg(y).arg(&sf).arg(&ni);
14879 unsafe {
14880 b.launch(cfg)?;
14881 }
14882 Ok(())
14883 }
14884
14885 pub fn bf16_to_f32(
14890 &self,
14891 data: &cudarc::driver::CudaView<'_, u8>,
14892 n: usize,
14893 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14894 let mut out = self.alloc_uninit::<f32>(n)?;
14895 let f = self.func("bf16_to_f32");
14896 let cfg = LaunchConfig::for_num_elems(n as u32);
14897 let ni = n as i32;
14898 let __s_b = self.gpu.stream();
14899 let mut b = __s_b.launch_builder(&f);
14900 b.arg(data).arg(&mut out).arg(&ni);
14901 unsafe {
14902 b.launch(cfg)?;
14903 }
14904 Ok(out)
14905 }
14906
14907 fn linear_bf16_chunked(
14914 &self,
14915 x: &CudaSlice<f32>,
14916 data: &CudaSlice<u8>,
14917 m: usize,
14918 in_f: usize,
14919 out_f: usize,
14920 exact: bool,
14921 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14922 const CHUNK_BYTES: usize = 256 << 20;
14923 let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
14924 if chunk_rows >= out_f {
14925 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
14926 return if exact {
14927 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
14928 } else {
14929 self.linear(x, &wf32, m, in_f, out_f)
14930 };
14931 }
14932 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14933 let mut r0 = 0usize;
14934 while r0 < out_f {
14935 let rows = chunk_rows.min(out_f - r0);
14936 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
14937 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
14938 let yc = if exact {
14939 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
14940 } else {
14941 self.linear(x, &wf32, m, in_f, rows)?
14942 };
14943 for mi in 0..m {
14945 let src = yc.slice(mi * rows..(mi + 1) * rows);
14946 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
14947 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
14948 }
14949 r0 += rows;
14950 }
14951 Ok(y)
14952 }
14953
14954 pub fn linear_decode_exact(
14961 &self,
14962 x: &CudaSlice<f32>,
14963 w: &CudaSlice<f32>,
14964 m_tokens: usize,
14965 in_f: usize,
14966 out_f: usize,
14967 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14968 if m_tokens == 1 {
14969 return self.linear(x, w, 1, in_f, out_f);
14970 }
14971 let xv = self.view(x, m_tokens * in_f);
14972 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
14973 for t in 0..m_tokens {
14974 let row = xv.slice(t * in_f..(t + 1) * in_f);
14975 let mut xr = self.alloc_uninit::<f32>(in_f)?;
14976 self.copy_view_into(&mut xr, 0, &row, in_f)?;
14977 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
14978 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
14979 }
14980 Ok(y)
14981 }
14982
14983 pub fn linear(
14984 &self,
14985 x: &CudaSlice<f32>,
14986 w: &CudaSlice<f32>,
14987 m_tokens: usize,
14988 in_f: usize,
14989 out_f: usize,
14990 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14991 use cudarc::cublaslt::{Matmul, MatmulConfig};
14992 let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; let cfg = MatmulConfig {
14994 transa: true,
14995 transb: false,
14996 transc: false,
14997 m: out_f as u64,
14998 n: m_tokens as u64,
14999 k: in_f as u64,
15000 alpha: 1.0,
15001 lda: in_f as i64,
15002 ldb: in_f as i64,
15003 beta: 0.0,
15004 ldc: out_f as i64,
15005 stride_a: None,
15006 stride_b: None,
15007 stride_c: None,
15008 stride_bias: None,
15009 batch_size: None,
15010 };
15011 unsafe {
15012 self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15013 }
15014 Ok(c)
15015 }
15016
15017 pub fn sdpa_naive(
15019 &self,
15020 q: &CudaSlice<f32>,
15021 k: &CudaSlice<f32>,
15022 v: &CudaSlice<f32>,
15023 o: &mut CudaSlice<f32>,
15024 head_dim: usize,
15025 n_head: usize,
15026 n_head_kv: usize,
15027 t: usize,
15028 t_kv: usize,
15029 scale: f32,
15030 causal: bool,
15031 ) -> Result<(), Box<dyn std::error::Error>> {
15032 let f = self.func("sdpa_naive_f32");
15033 let cfg = LaunchConfig {
15034 grid_dim: (n_head as u32, t as u32, 1),
15035 block_dim: (128, 1, 1),
15036 shared_mem_bytes: (t_kv * 4) as u32,
15037 };
15038 let (hd, nh, nhkv, ti, tkvi, cz) = (
15039 head_dim as i32,
15040 n_head as i32,
15041 n_head_kv as i32,
15042 t as i32,
15043 t_kv as i32,
15044 causal as i32,
15045 );
15046 let __s_b = self.gpu.stream();
15047 let mut b = __s_b.launch_builder(&f);
15048 b.arg(q)
15049 .arg(k)
15050 .arg(v)
15051 .arg(o)
15052 .arg(&hd)
15053 .arg(&nh)
15054 .arg(&nhkv)
15055 .arg(&ti)
15056 .arg(&tkvi)
15057 .arg(&scale)
15058 .arg(&cz);
15059 unsafe {
15060 b.launch(cfg)?;
15061 }
15062 Ok(())
15063 }
15064
15065 #[allow(clippy::too_many_arguments)]
15070 pub fn sdpa_naive_island(
15071 &self,
15072 q: &CudaSlice<f32>,
15073 k: &CudaSlice<f32>,
15074 v: &CudaSlice<f32>,
15075 o: &mut CudaSlice<f32>,
15076 span_id: &CudaSlice<i32>,
15077 head_dim: usize,
15078 n_head: usize,
15079 n_head_kv: usize,
15080 t: usize,
15081 t_kv: usize,
15082 scale: f32,
15083 window: usize,
15084 ) -> Result<(), Box<dyn std::error::Error>> {
15085 let f = self.func("sdpa_naive_island_f32");
15086 let cfg = LaunchConfig {
15087 grid_dim: (n_head as u32, t as u32, 1),
15088 block_dim: (128, 1, 1),
15089 shared_mem_bytes: (t_kv * 4) as u32,
15090 };
15091 let (hd, nh, nhkv, ti, tkvi, wi) = (
15092 head_dim as i32,
15093 n_head as i32,
15094 n_head_kv as i32,
15095 t as i32,
15096 t_kv as i32,
15097 window as i32,
15098 );
15099 let __s_b = self.gpu.stream();
15100 let mut b = __s_b.launch_builder(&f);
15101 b.arg(q)
15102 .arg(k)
15103 .arg(v)
15104 .arg(o)
15105 .arg(span_id)
15106 .arg(&hd)
15107 .arg(&nh)
15108 .arg(&nhkv)
15109 .arg(&ti)
15110 .arg(&tkvi)
15111 .arg(&scale)
15112 .arg(&wi);
15113 unsafe {
15114 b.launch(cfg)?;
15115 }
15116 Ok(())
15117 }
15118
15119 #[allow(clippy::too_many_arguments)]
15121 pub fn sdpa_naive_w(
15122 &self,
15123 q: &CudaSlice<f32>,
15124 k: &CudaSlice<f32>,
15125 v: &CudaSlice<f32>,
15126 o: &mut CudaSlice<f32>,
15127 head_dim: usize,
15128 n_head: usize,
15129 n_head_kv: usize,
15130 t: usize,
15131 t_kv: usize,
15132 scale: f32,
15133 causal: bool,
15134 window: usize,
15135 ) -> Result<(), Box<dyn std::error::Error>> {
15136 let f = self.func("sdpa_naive_w_f32");
15137 let cfg = LaunchConfig {
15138 grid_dim: (n_head as u32, t as u32, 1),
15139 block_dim: (128, 1, 1),
15140 shared_mem_bytes: (t_kv * 4) as u32,
15141 };
15142 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15143 head_dim as i32,
15144 n_head as i32,
15145 n_head_kv as i32,
15146 t as i32,
15147 t_kv as i32,
15148 causal as i32,
15149 window as i32,
15150 );
15151 let __s_b = self.gpu.stream();
15152 let mut b = __s_b.launch_builder(&f);
15153 b.arg(q)
15154 .arg(k)
15155 .arg(v)
15156 .arg(o)
15157 .arg(&hd)
15158 .arg(&nh)
15159 .arg(&nhkv)
15160 .arg(&ti)
15161 .arg(&tkvi)
15162 .arg(&scale)
15163 .arg(&cz)
15164 .arg(&wi);
15165 unsafe {
15166 b.launch(cfg)?;
15167 }
15168 Ok(())
15169 }
15170
15171 pub fn sdpa_naive_view(
15173 &self,
15174 q: &CudaSlice<f32>,
15175 k: &cudarc::driver::CudaView<f32>,
15176 v: &cudarc::driver::CudaView<f32>,
15177 o: &mut CudaSlice<f32>,
15178 head_dim: usize,
15179 n_head: usize,
15180 n_head_kv: usize,
15181 t: usize,
15182 t_kv: usize,
15183 scale: f32,
15184 causal: bool,
15185 ) -> Result<(), Box<dyn std::error::Error>> {
15186 let f = self.func("sdpa_naive_f32");
15187 let cfg = LaunchConfig {
15188 grid_dim: (n_head as u32, t as u32, 1),
15189 block_dim: (128, 1, 1),
15190 shared_mem_bytes: (t_kv * 4) as u32,
15191 };
15192 let (hd, nh, nhkv, ti, tkvi, cz) = (
15193 head_dim as i32,
15194 n_head as i32,
15195 n_head_kv as i32,
15196 t as i32,
15197 t_kv as i32,
15198 causal as i32,
15199 );
15200 let __s_b = self.gpu.stream();
15201 let mut b = __s_b.launch_builder(&f);
15202 b.arg(q)
15203 .arg(k)
15204 .arg(v)
15205 .arg(o)
15206 .arg(&hd)
15207 .arg(&nh)
15208 .arg(&nhkv)
15209 .arg(&ti)
15210 .arg(&tkvi)
15211 .arg(&scale)
15212 .arg(&cz);
15213 unsafe {
15214 b.launch(cfg)?;
15215 }
15216 Ok(())
15217 }
15218
15219 #[allow(clippy::too_many_arguments)]
15227 pub fn fa_dequant_kv_view_f32(
15228 &self,
15229 k: &cudarc::driver::CudaView<u8>,
15230 v: &cudarc::driver::CudaView<u8>,
15231 kf: &mut CudaSlice<f32>,
15232 vf: &mut CudaSlice<f32>,
15233 kv_dim_k: usize,
15234 kv_dim_v: usize,
15235 t_kv: usize,
15236 k_tok_bytes: usize,
15237 v_tok_bytes: usize,
15238 g: bool,
15239 ) -> Result<(), Box<dyn std::error::Error>> {
15240 let f = if g {
15241 self.func_g("fa_dequant_kv_ws_f32")
15242 } else {
15243 self.func("fa_dequant_kv_ws_f32")
15244 };
15245 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
15246 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15247 let cfg = LaunchConfig {
15248 grid_dim: (nblk.max(1), 1, 1),
15249 block_dim: (256, 1, 1),
15250 shared_mem_bytes: 0,
15251 };
15252 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
15253 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15254 let __s_b = self.gpu.stream();
15255 let mut b = __s_b.launch_builder(&f);
15256 b.arg(k)
15257 .arg(v)
15258 .arg(&mut *kf)
15259 .arg(&mut *vf)
15260 .arg(&kdk)
15261 .arg(&kdv)
15262 .arg(&tkvi)
15263 .arg(&ktb)
15264 .arg(&vtb);
15265 unsafe {
15266 b.launch(cfg)?;
15267 }
15268 Ok(())
15269 }
15270
15271 #[allow(clippy::too_many_arguments)]
15272 pub fn sdpa_naive_quantized_view(
15273 &self,
15274 q: &CudaSlice<f32>,
15275 k: &cudarc::driver::CudaView<u8>,
15276 v: &cudarc::driver::CudaView<u8>,
15277 o: &mut CudaSlice<f32>,
15278 head_dim: usize,
15279 n_head: usize,
15280 n_head_kv: usize,
15281 t: usize,
15282 t_kv: usize,
15283 scale: f32,
15284 causal: bool,
15285 k_tok_bytes: usize,
15286 v_tok_bytes: usize,
15287 ) -> Result<(), Box<dyn std::error::Error>> {
15288 let kv_dim = n_head_kv * head_dim;
15289 let mut kf = self.uninit(t_kv * kv_dim)?;
15290 let mut vf = self.uninit(t_kv * kv_dim)?;
15291 let f = self.func("fa_dequant_kv_ws_f32");
15292 let total = (2 * t_kv * kv_dim) as u64;
15293 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15294 let cfg = LaunchConfig {
15295 grid_dim: (nblk.max(1), 1, 1),
15296 block_dim: (256, 1, 1),
15297 shared_mem_bytes: 0,
15298 };
15299 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15300 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15301 let __s_b = self.gpu.stream();
15302 let mut b = __s_b.launch_builder(&f);
15303 b.arg(k)
15304 .arg(v)
15305 .arg(&mut kf)
15306 .arg(&mut vf)
15307 .arg(&kv_dim_i)
15308 .arg(&kv_dim_i)
15309 .arg(&t_kv_i)
15310 .arg(&k_tok_bytes_i)
15311 .arg(&v_tok_bytes_i);
15312 unsafe { b.launch(cfg)? };
15313 self.sdpa_naive(
15314 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15315 )
15316 }
15317
15318 #[allow(clippy::too_many_arguments)]
15330 pub fn sdpa_naive_w_quantized_view(
15331 &self,
15332 q: &CudaSlice<f32>,
15333 k: &cudarc::driver::CudaView<u8>,
15334 v: &cudarc::driver::CudaView<u8>,
15335 o: &mut CudaSlice<f32>,
15336 head_dim: usize,
15337 n_head: usize,
15338 n_head_kv: usize,
15339 t: usize,
15340 t_kv: usize,
15341 scale: f32,
15342 causal: bool,
15343 window: usize,
15344 k_tok_bytes: usize,
15345 v_tok_bytes: usize,
15346 ) -> Result<(), Box<dyn std::error::Error>> {
15347 let kv_dim = n_head_kv * head_dim;
15348 let mut kf = self.uninit(t_kv * kv_dim)?;
15349 let mut vf = self.uninit(t_kv * kv_dim)?;
15350 let f = self.func("fa_dequant_kv_ws_f32");
15351 let total = (2 * t_kv * kv_dim) as u64;
15352 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15353 let cfg = LaunchConfig {
15354 grid_dim: (nblk.max(1), 1, 1),
15355 block_dim: (256, 1, 1),
15356 shared_mem_bytes: 0,
15357 };
15358 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15359 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15360 let __s_b = self.gpu.stream();
15361 let mut b = __s_b.launch_builder(&f);
15362 b.arg(k)
15363 .arg(v)
15364 .arg(&mut kf)
15365 .arg(&mut vf)
15366 .arg(&kv_dim_i)
15367 .arg(&kv_dim_i)
15368 .arg(&t_kv_i)
15369 .arg(&k_tok_bytes_i)
15370 .arg(&v_tok_bytes_i);
15371 unsafe { b.launch(cfg)? };
15372 self.sdpa_naive_w(
15373 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15374 )
15375 }
15376
15377 pub fn fa_prefill(
15381 &self,
15382 q: &CudaSlice<f32>,
15383 k: &CudaSlice<f32>,
15384 v: &CudaSlice<f32>,
15385 o: &mut CudaSlice<f32>,
15386 head_dim: usize,
15387 n_head: usize,
15388 n_head_kv: usize,
15389 t: usize,
15390 t_kv: usize,
15391 scale: f32,
15392 causal: bool,
15393 ) -> Result<(), Box<dyn std::error::Error>> {
15394 if portable_mma_gated() {
15395 return self.sdpa_naive(
15396 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15397 );
15398 }
15399 let fa3_on = head_dim == 256
15407 && causal
15408 && t == t_kv
15409 && match std::env::var("MEMRA_FA3").as_deref() {
15410 Ok("0") => false,
15411 Ok("1") => true,
15412 _ => cfg!(memra_hopper_mma),
15413 };
15414 if fa3_on {
15415 let n = t * n_head * head_dim;
15416 let nkv = t * n_head_kv * head_dim;
15417 let mut q16 = self.alloc_u8_uninit(n * 2)?;
15418 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
15419 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
15420 self.f32_to_bf16_into(q, &mut q16, n)?;
15421 self.f32_to_bf16_into(k, &mut k16, nkv)?;
15422 self.f32_to_bf16_into(v, &mut v16, nkv)?;
15423 let rc = {
15424 use cudarc::driver::{DevicePtr, DevicePtrMut};
15425 let stream = self.gpu.stream();
15426 let (qp, _g1) = q16.device_ptr(&stream);
15427 let (kp, _g2) = k16.device_ptr(&stream);
15428 let (vp, _g3) = v16.device_ptr(&stream);
15429 let (op, _g4) = o.device_ptr_mut(&stream);
15430 unsafe {
15431 memra_fa3_prefill(
15432 qp as *const core::ffi::c_void,
15433 kp as *const core::ffi::c_void,
15434 vp as *const core::ffi::c_void,
15435 op as *mut f32,
15436 t as i32,
15437 n_head as i32,
15438 n_head_kv as i32,
15439 head_dim as i32,
15440 scale,
15441 stream.cu_stream() as *mut core::ffi::c_void,
15442 )
15443 }
15444 };
15445 if rc != 0 {
15446 return Err(format!("memra_fa3_prefill rc={rc}").into());
15447 }
15448 return Ok(());
15449 }
15450 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15455 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
15456 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
15457 const BLOCK_Q: usize = 64;
15458 const BKX: usize = 32;
15459 let f = self.func("fa_prefill_bf16_p1");
15460 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
15461 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
15462 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15463 f.set_attribute(
15464 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15465 shmem as i32,
15466 )?;
15467 let cfg = LaunchConfig {
15468 grid_dim: (
15469 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15470 n_head as u32,
15471 1,
15472 ),
15473 block_dim: (32, 4, 1),
15474 shared_mem_bytes: shmem,
15475 };
15476 let (hd, nh, nhkv, ti, tkvi, cz) = (
15477 head_dim as i32,
15478 n_head as i32,
15479 n_head_kv as i32,
15480 t as i32,
15481 t_kv as i32,
15482 causal as i32,
15483 );
15484 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15485 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15486 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15487 let __s_b = self.gpu.stream();
15488 let mut b = __s_b.launch_builder(&f);
15489 b.arg(&qb)
15490 .arg(&kb)
15491 .arg(&vb)
15492 .arg(o)
15493 .arg(&hd)
15494 .arg(&nh)
15495 .arg(&nhkv)
15496 .arg(&ti)
15497 .arg(&tkvi)
15498 .arg(&scale)
15499 .arg(&cz);
15500 unsafe {
15501 b.launch(cfg)?;
15502 }
15503 return Ok(());
15504 }
15505 const BK: usize = 32;
15511 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
15514 let (block_q, warps, w2_sfx): (usize, u32, &str) =
15515 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
15516 let hd_sfx = fa_hd_suffix(head_dim)?;
15520 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15521 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
15526 let (kb16, vb16) = if bf16kv {
15527 let n = t_kv * n_head_kv * head_dim;
15528 let mut kb = self.alloc_u8_uninit(n * 2)?;
15529 let mut vb = self.alloc_u8_uninit(n * 2)?;
15530 let fcv = self.func("f32_to_bf16_bulk");
15531 let ni = n as i64;
15532 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
15533 let __s_b = self.gpu.stream();
15534 let mut b = __s_b.launch_builder(&fcv);
15535 b.arg(k).arg(&mut kb).arg(&ni);
15536 unsafe {
15537 b.launch(cfgc)?;
15538 }
15539 let __s_b = self.gpu.stream();
15540 let mut b = __s_b.launch_builder(&fcv);
15541 b.arg(v).arg(&mut vb).arg(&ni);
15542 unsafe {
15543 b.launch(cfgc)?;
15544 }
15545 (Some(kb), Some(vb))
15546 } else {
15547 (None, None)
15548 };
15549 let f = self.func(&if bf16kv {
15550 format!("fa_prefill_bf16kv_pp{hd_sfx}")
15551 } else {
15552 format!(
15553 "fa_prefill_f32{}{}{hd_sfx}",
15554 if floor { "" } else { "_pp" },
15555 if floor { "" } else { w2_sfx }
15556 )
15557 });
15558 let kv_stages = if bf16kv { 2 } else { 1 };
15561 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
15562 + 4 * (block_q * BK + 2 * block_q)) as u32;
15563 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15564 f.set_attribute(
15565 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15566 shmem as i32,
15567 )?;
15568 let cfg = LaunchConfig {
15569 grid_dim: (
15570 (t as u32 + block_q as u32 - 1) / block_q as u32,
15571 n_head as u32,
15572 1,
15573 ),
15574 block_dim: (32, warps, 1),
15575 shared_mem_bytes: shmem,
15576 };
15577 let (hd, nh, nhkv, ti, tkvi, cz) = (
15578 head_dim as i32,
15579 n_head as i32,
15580 n_head_kv as i32,
15581 t as i32,
15582 t_kv as i32,
15583 causal as i32,
15584 );
15585 let __s_b = self.gpu.stream();
15586 let mut b = __s_b.launch_builder(&f);
15587 b.arg(q);
15588 match (&kb16, &vb16) {
15589 (Some(kb), Some(vb)) => {
15590 b.arg(kb).arg(vb);
15591 }
15592 _ => {
15593 b.arg(k).arg(v);
15594 }
15595 }
15596 b.arg(o)
15597 .arg(&hd)
15598 .arg(&nh)
15599 .arg(&nhkv)
15600 .arg(&ti)
15601 .arg(&tkvi)
15602 .arg(&scale)
15603 .arg(&cz);
15604 unsafe {
15605 b.launch(cfg)?;
15606 }
15607 Ok(())
15608 }
15609
15610 #[allow(clippy::too_many_arguments)]
15614 pub fn fa_prefill_w(
15615 &self,
15616 q: &CudaSlice<f32>,
15617 k: &CudaSlice<f32>,
15618 v: &CudaSlice<f32>,
15619 o: &mut CudaSlice<f32>,
15620 head_dim: usize,
15621 n_head: usize,
15622 n_head_kv: usize,
15623 t: usize,
15624 t_kv: usize,
15625 scale: f32,
15626 causal: bool,
15627 window: usize,
15628 ) -> Result<(), Box<dyn std::error::Error>> {
15629 if portable_mma_gated() {
15632 return self.sdpa_naive_w(
15633 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15634 );
15635 }
15636 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15640 let faw_f32 =
15641 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
15642 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15643 self.fa_prefill_w_arm(
15644 q,
15645 k,
15646 v,
15647 o,
15648 head_dim,
15649 n_head,
15650 n_head_kv,
15651 t,
15652 t_kv,
15653 scale,
15654 causal,
15655 window,
15656 floor || faw_f32,
15657 floor,
15658 )
15659 }
15660
15661 #[allow(clippy::too_many_arguments)]
15664 pub fn fa_prefill_w_pre(
15665 &self,
15666 qb: &CudaSlice<u8>,
15667 kb: &CudaSlice<u8>,
15668 vb: &CudaSlice<u8>,
15669 o: &mut CudaSlice<f32>,
15670 head_dim: usize,
15671 n_head: usize,
15672 n_head_kv: usize,
15673 t: usize,
15674 t_kv: usize,
15675 scale: f32,
15676 causal: bool,
15677 window: usize,
15678 v_f16: bool,
15679 ) -> Result<(), Box<dyn std::error::Error>> {
15680 const BLOCK_Q: usize = 64;
15681 const BK: usize = 32;
15682 debug_assert_eq!(head_dim, 256);
15683 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
15684 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
15685 if hp {
15686 const BLOCK_QH: usize = 32;
15687 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
15690 let vh: &CudaSlice<u8> = if v_f16 {
15691 vb
15692 } else {
15693 let n = t_kv * n_head_kv * head_dim;
15694 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
15695 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
15696 }
15697 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
15698 vguard.as_ref().unwrap()
15699 };
15700 let f = self.func("fa_prefill_w_bf16_p1h2");
15701 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
15702 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15703 f.set_attribute(
15704 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15705 shmem as i32,
15706 )?;
15707 let cfg = LaunchConfig {
15708 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
15709 block_dim: (32, 4, 1),
15710 shared_mem_bytes: shmem,
15711 };
15712 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15713 head_dim as i32,
15714 n_head as i32,
15715 n_head_kv as i32,
15716 t as i32,
15717 t_kv as i32,
15718 causal as i32,
15719 window as i32,
15720 );
15721 let __s_b = self.gpu.stream();
15722 let mut b = __s_b.launch_builder(&f);
15723 b.arg(qb)
15724 .arg(kb)
15725 .arg(vh)
15726 .arg(o)
15727 .arg(&hd)
15728 .arg(&nh)
15729 .arg(&nhkv)
15730 .arg(&ti)
15731 .arg(&tkvi)
15732 .arg(&scale)
15733 .arg(&cz)
15734 .arg(&wi);
15735 unsafe {
15736 b.launch(cfg)?;
15737 }
15738 return Ok(());
15739 }
15740 let f = self.func("fa_prefill_w_bf16_p1");
15741 let shmem =
15742 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
15743 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15744 f.set_attribute(
15745 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15746 shmem as i32,
15747 )?;
15748 let cfg = LaunchConfig {
15749 grid_dim: (
15750 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15751 n_head as u32,
15752 1,
15753 ),
15754 block_dim: (32, 4, 1),
15755 shared_mem_bytes: shmem,
15756 };
15757 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15758 head_dim as i32,
15759 n_head as i32,
15760 n_head_kv as i32,
15761 t as i32,
15762 t_kv as i32,
15763 causal as i32,
15764 window as i32,
15765 );
15766 let __s_b = self.gpu.stream();
15767 let mut b = __s_b.launch_builder(&f);
15768 b.arg(qb)
15769 .arg(kb)
15770 .arg(vb)
15771 .arg(o)
15772 .arg(&hd)
15773 .arg(&nh)
15774 .arg(&nhkv)
15775 .arg(&ti)
15776 .arg(&tkvi)
15777 .arg(&scale)
15778 .arg(&cz)
15779 .arg(&wi);
15780 unsafe {
15781 b.launch(cfg)?;
15782 }
15783 Ok(())
15784 }
15785
15786 #[allow(clippy::too_many_arguments)]
15788 pub fn fa_prefill_w_arm(
15789 &self,
15790 q: &CudaSlice<f32>,
15791 k: &CudaSlice<f32>,
15792 v: &CudaSlice<f32>,
15793 o: &mut CudaSlice<f32>,
15794 head_dim: usize,
15795 n_head: usize,
15796 n_head_kv: usize,
15797 t: usize,
15798 t_kv: usize,
15799 scale: f32,
15800 causal: bool,
15801 window: usize,
15802 f32_stage: bool,
15803 floor: bool,
15804 ) -> Result<(), Box<dyn std::error::Error>> {
15805 const BLOCK_Q: usize = 64;
15806 const BK: usize = 32;
15807 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
15808 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15812 let p1 = !floor
15813 && !f32_stage
15814 && *P1_ON.get_or_init(|| {
15815 std::env::var("MEMRA_FAW_P1")
15816 .map(|v| v != "0")
15817 .unwrap_or(true)
15818 });
15819 let hp =
15820 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
15821 if hp {
15822 const BLOCK_QH: usize = 32;
15823 let f = self.func("fa_prefill_w_bf16_p1h2");
15824 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
15825 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15826 f.set_attribute(
15827 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15828 shmem as i32,
15829 )?;
15830 let cfg = LaunchConfig {
15831 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
15832 block_dim: (32, 4, 1),
15833 shared_mem_bytes: shmem,
15834 };
15835 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15836 head_dim as i32,
15837 n_head as i32,
15838 n_head_kv as i32,
15839 t as i32,
15840 t_kv as i32,
15841 causal as i32,
15842 window as i32,
15843 );
15844 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15845 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15846 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
15847 let __s_b = self.gpu.stream();
15848 let mut b = __s_b.launch_builder(&f);
15849 b.arg(&qb)
15850 .arg(&kb)
15851 .arg(&vh)
15852 .arg(o)
15853 .arg(&hd)
15854 .arg(&nh)
15855 .arg(&nhkv)
15856 .arg(&ti)
15857 .arg(&tkvi)
15858 .arg(&scale)
15859 .arg(&cz)
15860 .arg(&wi);
15861 unsafe {
15862 b.launch(cfg)?;
15863 }
15864 return Ok(());
15865 }
15866 if p1 {
15867 let f = self.func("fa_prefill_w_bf16_p1");
15868 let shmem =
15869 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
15870 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15871 f.set_attribute(
15872 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15873 shmem as i32,
15874 )?;
15875 let cfg = LaunchConfig {
15876 grid_dim: (
15877 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15878 n_head as u32,
15879 1,
15880 ),
15881 block_dim: (32, 4, 1),
15882 shared_mem_bytes: shmem,
15883 };
15884 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15885 head_dim as i32,
15886 n_head as i32,
15887 n_head_kv as i32,
15888 t as i32,
15889 t_kv as i32,
15890 causal as i32,
15891 window as i32,
15892 );
15893 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15894 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15895 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15896 let __s_b = self.gpu.stream();
15897 let mut b = __s_b.launch_builder(&f);
15898 b.arg(&qb)
15899 .arg(&kb)
15900 .arg(&vb)
15901 .arg(o)
15902 .arg(&hd)
15903 .arg(&nh)
15904 .arg(&nhkv)
15905 .arg(&ti)
15906 .arg(&tkvi)
15907 .arg(&scale)
15908 .arg(&cz)
15909 .arg(&wi);
15910 unsafe {
15911 b.launch(cfg)?;
15912 }
15913 return Ok(());
15914 }
15915 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15918 let g4 = !floor
15919 && !f32_stage
15920 && n_head_kv == 1
15921 && n_head % 4 == 0
15922 && *G4_ON.get_or_init(|| {
15923 std::env::var("MEMRA_FAW_G4")
15924 .map(|v| v != "0")
15925 .unwrap_or(true)
15926 });
15927 if g4 {
15928 const SP_M: usize = 16;
15929 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15932 let o2 = *O2_ON.get_or_init(|| {
15933 std::env::var("MEMRA_FAW_O2")
15934 .map(|v| v != "0")
15935 .unwrap_or(true)
15936 });
15937 let f = self.func(if o2 {
15938 "fa_prefill_w_bf16_g4o2"
15939 } else {
15940 "fa_prefill_w_bf16_g4"
15941 });
15942 let shmem = if o2 {
15943 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
15944 } else {
15945 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
15946 as u32
15947 };
15948 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15949 f.set_attribute(
15950 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15951 shmem as i32,
15952 )?;
15953 let cfg = LaunchConfig {
15954 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
15955 block_dim: (32, 4, 1),
15956 shared_mem_bytes: shmem,
15957 };
15958 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15959 head_dim as i32,
15960 n_head as i32,
15961 n_head_kv as i32,
15962 t as i32,
15963 t_kv as i32,
15964 causal as i32,
15965 window as i32,
15966 );
15967 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15968 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15969 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15970 let __s_b = self.gpu.stream();
15971 let mut b = __s_b.launch_builder(&f);
15972 b.arg(&qb)
15973 .arg(&kb)
15974 .arg(&vb)
15975 .arg(o)
15976 .arg(&hd)
15977 .arg(&nh)
15978 .arg(&nhkv)
15979 .arg(&ti)
15980 .arg(&tkvi)
15981 .arg(&scale)
15982 .arg(&cz)
15983 .arg(&wi);
15984 unsafe {
15985 b.launch(cfg)?;
15986 }
15987 return Ok(());
15988 }
15989 let f = self.func(if floor {
15990 "fa_prefill_w_f32"
15991 } else if f32_stage {
15992 "fa_prefill_w_f32_pp"
15993 } else {
15994 "fa_prefill_w_bf16_pp"
15995 });
15996 let shmem =
15997 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
15998 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15999 f.set_attribute(
16000 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16001 shmem as i32,
16002 )?;
16003 let cfg = LaunchConfig {
16004 grid_dim: (
16005 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16006 n_head as u32,
16007 1,
16008 ),
16009 block_dim: (32, 4, 1),
16010 shared_mem_bytes: shmem,
16011 };
16012 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16013 head_dim as i32,
16014 n_head as i32,
16015 n_head_kv as i32,
16016 t as i32,
16017 t_kv as i32,
16018 causal as i32,
16019 window as i32,
16020 );
16021 if f32_stage {
16022 let __s_b = self.gpu.stream();
16023 let mut b = __s_b.launch_builder(&f);
16024 b.arg(q)
16025 .arg(k)
16026 .arg(v)
16027 .arg(o)
16028 .arg(&hd)
16029 .arg(&nh)
16030 .arg(&nhkv)
16031 .arg(&ti)
16032 .arg(&tkvi)
16033 .arg(&scale)
16034 .arg(&cz)
16035 .arg(&wi);
16036 unsafe {
16037 b.launch(cfg)?;
16038 }
16039 } else {
16040 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16041 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16042 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16043 let __s_b = self.gpu.stream();
16044 let mut b = __s_b.launch_builder(&f);
16045 b.arg(&qb)
16046 .arg(&kb)
16047 .arg(&vb)
16048 .arg(o)
16049 .arg(&hd)
16050 .arg(&nh)
16051 .arg(&nhkv)
16052 .arg(&ti)
16053 .arg(&tkvi)
16054 .arg(&scale)
16055 .arg(&cz)
16056 .arg(&wi);
16057 unsafe {
16058 b.launch(cfg)?;
16059 }
16060 }
16061 Ok(())
16062 }
16063
16064 #[allow(clippy::too_many_arguments)]
16068 pub fn fa_prefill_hd512(
16069 &self,
16070 q: &CudaSlice<f32>,
16071 k: &CudaSlice<f32>,
16072 v: &CudaSlice<f32>,
16073 o: &mut CudaSlice<f32>,
16074 head_dim: usize,
16075 n_head: usize,
16076 n_head_kv: usize,
16077 t: usize,
16078 t_kv: usize,
16079 scale: f32,
16080 causal: bool,
16081 ) -> Result<(), Box<dyn std::error::Error>> {
16082 if portable_mma_gated() {
16084 return self.sdpa_naive(
16085 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16086 );
16087 }
16088 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16094 let f32_stage =
16095 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16096 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16100 let sp = !f32_stage
16101 && *SP_ON.get_or_init(|| {
16102 std::env::var("MEMRA_FA512_SP")
16103 .map(|v| v != "0")
16104 .unwrap_or(true)
16105 });
16106 self.fa_prefill_hd512_arm(
16107 q,
16108 k,
16109 v,
16110 o,
16111 head_dim,
16112 n_head,
16113 n_head_kv,
16114 t,
16115 t_kv,
16116 scale,
16117 causal,
16118 f32_stage,
16119 sp,
16120 sp && fa_f16pv_on(),
16121 )
16122 }
16123
16124 #[allow(clippy::too_many_arguments)]
16126 pub fn fa_prefill_hd512_pre(
16127 &self,
16128 qb: &CudaSlice<u8>,
16129 kb: &CudaSlice<u8>,
16130 vb: &CudaSlice<u8>,
16131 o: &mut CudaSlice<f32>,
16132 head_dim: usize,
16133 n_head: usize,
16134 n_head_kv: usize,
16135 t: usize,
16136 t_kv: usize,
16137 scale: f32,
16138 causal: bool,
16139 v_f16: bool,
16140 ) -> Result<(), Box<dyn std::error::Error>> {
16141 debug_assert_eq!(head_dim, 512);
16142 const SP_M: usize = 16;
16143 const BKS: usize = 32;
16144 let f16pv = fa_f16pv_on();
16148 let nw = if f16pv { fa512_wide_warps() } else { 2 };
16149 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16150 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16151 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16152 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16153 let n = t_kv * n_head_kv * head_dim;
16155 let need = n * 2;
16156 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16157 *vguard = Some(self.alloc_uninit::<u8>(need)?);
16158 }
16159 let dst = vguard.as_mut().unwrap();
16160 self.bf16_to_f16_into(vb, n, dst)?;
16161 vguard.as_ref().unwrap()
16162 } else {
16163 vb
16164 };
16165 let f = self.func(if hp {
16166 "fa_prefill_bf16_hd512_sp16h2"
16167 } else {
16168 match (f16pv, nw) {
16169 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16170 (true, _) => "fa_prefill_bf16_hd512_sp16",
16171 _ => "fa_prefill_bf16_hd512_sp",
16172 }
16173 });
16174 let (nwarp, npart) = if hp {
16175 (4usize, 4usize)
16176 } else if nw > 2 {
16177 (nw, nw)
16178 } else {
16179 (2, 1)
16180 };
16181 let shmem = if hp {
16183 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
16184 as u32
16185 } else {
16186 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16187 + 4 * (npart * SP_M * BKS + SP_M)) as u32
16188 };
16189 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16190 f.set_attribute(
16191 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16192 shmem as i32,
16193 )?;
16194 let grid_y = if hp {
16195 (n_head / 2) as u32
16196 } else {
16197 n_head as u32
16198 };
16199 let cfg = LaunchConfig {
16200 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16201 block_dim: (32, nwarp as u32, 1),
16202 shared_mem_bytes: shmem,
16203 };
16204 let (hd, nh, nhkv, ti, tkvi, cz) = (
16205 head_dim as i32,
16206 n_head as i32,
16207 n_head_kv as i32,
16208 t as i32,
16209 t_kv as i32,
16210 causal as i32,
16211 );
16212 let __s_b = self.gpu.stream();
16213 let mut b = __s_b.launch_builder(&f);
16214 b.arg(qb)
16215 .arg(kb)
16216 .arg(vref)
16217 .arg(o)
16218 .arg(&hd)
16219 .arg(&nh)
16220 .arg(&nhkv)
16221 .arg(&ti)
16222 .arg(&tkvi)
16223 .arg(&scale)
16224 .arg(&cz);
16225 unsafe {
16226 b.launch(cfg)?;
16227 }
16228 Ok(())
16229 }
16230
16231 #[allow(clippy::too_many_arguments)]
16234 pub fn fa_prefill_hd512_arm(
16235 &self,
16236 q: &CudaSlice<f32>,
16237 k: &CudaSlice<f32>,
16238 v: &CudaSlice<f32>,
16239 o: &mut CudaSlice<f32>,
16240 head_dim: usize,
16241 n_head: usize,
16242 n_head_kv: usize,
16243 t: usize,
16244 t_kv: usize,
16245 scale: f32,
16246 causal: bool,
16247 f32_stage: bool,
16248 sp: bool,
16249 f16pv: bool,
16250 ) -> Result<(), Box<dyn std::error::Error>> {
16251 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
16252 if sp && !f32_stage {
16253 const SP_M: usize = 16;
16257 const BKS: usize = 32;
16258 let nw = if f16pv { fa512_wide_warps() } else { 2 };
16259 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16260 let f = self.func(if hp {
16261 "fa_prefill_bf16_hd512_sp16h2"
16262 } else {
16263 match (f16pv, nw) {
16264 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16265 (true, _) => "fa_prefill_bf16_hd512_sp16",
16266 _ => "fa_prefill_bf16_hd512_sp",
16267 }
16268 });
16269 let (nwarp, npart) = if hp {
16270 (4usize, 4usize)
16271 } else if nw > 2 {
16272 (nw, nw)
16273 } else {
16274 (2, 1)
16275 };
16276 let shmem = if hp {
16277 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
16278 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
16279 } else {
16280 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16281 + 4 * (npart * SP_M * BKS + SP_M)) as u32
16282 };
16283 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16284 f.set_attribute(
16285 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16286 shmem as i32,
16287 )?;
16288 let grid_y = if hp {
16289 (n_head / 2) as u32
16290 } else {
16291 n_head as u32
16292 };
16293 let cfg = LaunchConfig {
16294 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16295 block_dim: (32, nwarp as u32, 1),
16296 shared_mem_bytes: shmem,
16297 };
16298 let (hd, nh, nhkv, ti, tkvi, cz) = (
16299 head_dim as i32,
16300 n_head as i32,
16301 n_head_kv as i32,
16302 t as i32,
16303 t_kv as i32,
16304 causal as i32,
16305 );
16306 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16307 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16308 let vb = if f16pv {
16309 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
16310 } else {
16311 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
16312 };
16313 let __s_b = self.gpu.stream();
16314 let mut b = __s_b.launch_builder(&f);
16315 b.arg(&qb)
16316 .arg(&kb)
16317 .arg(&vb)
16318 .arg(o)
16319 .arg(&hd)
16320 .arg(&nh)
16321 .arg(&nhkv)
16322 .arg(&ti)
16323 .arg(&tkvi)
16324 .arg(&scale)
16325 .arg(&cz);
16326 unsafe {
16327 b.launch(cfg)?;
16328 }
16329 return Ok(());
16330 }
16331 const BLOCK_Q: usize = 32;
16332 const BK: usize = 32;
16333 const HALF: usize = 256;
16334 let f = self.func(if f32_stage {
16335 "fa_prefill_f32_hd512"
16336 } else {
16337 "fa_prefill_bf16_hd512"
16338 });
16339 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
16341 + 4 * BLOCK_Q) as u32;
16342 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16343 f.set_attribute(
16344 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16345 shmem as i32,
16346 )?;
16347 let cfg = LaunchConfig {
16348 grid_dim: (
16349 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16350 n_head as u32,
16351 2,
16352 ),
16353 block_dim: (32, 2, 1),
16354 shared_mem_bytes: shmem,
16355 };
16356 let (hd, nh, nhkv, ti, tkvi, cz) = (
16357 head_dim as i32,
16358 n_head as i32,
16359 n_head_kv as i32,
16360 t as i32,
16361 t_kv as i32,
16362 causal as i32,
16363 );
16364 if f32_stage {
16365 let __s_b = self.gpu.stream();
16366 let mut b = __s_b.launch_builder(&f);
16367 b.arg(q)
16368 .arg(k)
16369 .arg(v)
16370 .arg(o)
16371 .arg(&hd)
16372 .arg(&nh)
16373 .arg(&nhkv)
16374 .arg(&ti)
16375 .arg(&tkvi)
16376 .arg(&scale)
16377 .arg(&cz);
16378 unsafe {
16379 b.launch(cfg)?;
16380 }
16381 } else {
16382 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16383 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16384 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16385 let __s_b = self.gpu.stream();
16386 let mut b = __s_b.launch_builder(&f);
16387 b.arg(&qb)
16388 .arg(&kb)
16389 .arg(&vb)
16390 .arg(o)
16391 .arg(&hd)
16392 .arg(&nh)
16393 .arg(&nhkv)
16394 .arg(&ti)
16395 .arg(&tkvi)
16396 .arg(&scale)
16397 .arg(&cz);
16398 unsafe {
16399 b.launch(cfg)?;
16400 }
16401 }
16402 Ok(())
16403 }
16404
16405 #[allow(clippy::too_many_arguments)]
16409 pub fn rope_neox2_bf16e(
16410 &self,
16411 q: &mut CudaSlice<f32>,
16412 k: &mut CudaSlice<f32>,
16413 qb: &mut CudaSlice<u8>,
16414 kb: &mut CudaSlice<u8>,
16415 pos: &CudaSlice<i32>,
16416 head_dim: usize,
16417 n_dims: usize,
16418 nh_q: usize,
16419 nh_k: usize,
16420 n_tokens: usize,
16421 base: f32,
16422 freq_scale: f32,
16423 ff: Option<&CudaSlice<f32>>,
16424 ) -> Result<(), Box<dyn std::error::Error>> {
16425 let f = self.func("rope_neox2_bf16e_f32");
16426 let rows = ((nh_q + nh_k) * n_tokens) as u32;
16427 let cfg = LaunchConfig {
16428 grid_dim: (rows, 1, 1),
16429 block_dim: ((head_dim / 2) as u32, 1, 1),
16430 shared_mem_bytes: 0,
16431 };
16432 let theta_scale = base.powf(-2.0 / n_dims as f32);
16433 let (hd, nd, nhq, nhk, nt) = (
16434 head_dim as i32,
16435 n_dims as i32,
16436 nh_q as i32,
16437 nh_k as i32,
16438 n_tokens as i32,
16439 );
16440 let __s_b = self.gpu.stream();
16441 let mut b = __s_b.launch_builder(&f);
16442 match ff {
16443 Some(t) => {
16444 b.arg(&mut *q)
16445 .arg(&mut *k)
16446 .arg(&mut *qb)
16447 .arg(&mut *kb)
16448 .arg(pos)
16449 .arg(&hd)
16450 .arg(&nd)
16451 .arg(&nhq)
16452 .arg(&nhk)
16453 .arg(&nt)
16454 .arg(&theta_scale)
16455 .arg(&freq_scale)
16456 .arg(t);
16457 unsafe {
16458 b.launch(cfg)?;
16459 }
16460 }
16461 None => {
16462 let null: u64 = 0;
16463 b.arg(&mut *q)
16464 .arg(&mut *k)
16465 .arg(&mut *qb)
16466 .arg(&mut *kb)
16467 .arg(pos)
16468 .arg(&hd)
16469 .arg(&nd)
16470 .arg(&nhq)
16471 .arg(&nhk)
16472 .arg(&nt)
16473 .arg(&theta_scale)
16474 .arg(&freq_scale)
16475 .arg(&null);
16476 unsafe {
16477 b.launch(cfg)?;
16478 }
16479 }
16480 }
16481 Ok(())
16482 }
16483
16484 pub fn f32_to_bf16(
16487 &self,
16488 x: &CudaSlice<f32>,
16489 n: usize,
16490 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16491 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
16492 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16493 let f = self.func("f32_to_bf16_flat");
16494 let n_i = n as i64;
16495 let cfg = LaunchConfig {
16496 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16497 block_dim: (256, 1, 1),
16498 shared_mem_bytes: 0,
16499 };
16500 let __s_b = self.gpu.stream();
16501 let mut b = __s_b.launch_builder(&f);
16502 b.arg(x).arg(&mut y).arg(&n_i);
16503 unsafe {
16504 b.launch(cfg)?;
16505 }
16506 Ok(y)
16507 }
16508
16509 pub fn f32_to_f16(
16510 &self,
16511 x: &CudaSlice<f32>,
16512 n: usize,
16513 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16514 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
16515 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16516 let f = self.func("f32_to_f16_flat");
16517 let n_i = n as i64;
16518 let cfg = LaunchConfig {
16519 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16520 block_dim: (256, 1, 1),
16521 shared_mem_bytes: 0,
16522 };
16523 let __s_b = self.gpu.stream();
16524 let mut b = __s_b.launch_builder(&f);
16525 b.arg(x).arg(&mut y).arg(&n_i);
16526 unsafe {
16527 b.launch(cfg)?;
16528 }
16529 Ok(y)
16530 }
16531
16532 pub fn bf16_to_f16(
16534 &self,
16535 xb: &CudaSlice<u8>,
16536 n: usize,
16537 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16538 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16539 self.bf16_to_f16_into(xb, n, &mut y)?;
16540 Ok(y)
16541 }
16542
16543 pub fn bf16_to_f16_into(
16545 &self,
16546 xb: &CudaSlice<u8>,
16547 n: usize,
16548 y: &mut CudaSlice<u8>,
16549 ) -> Result<(), Box<dyn std::error::Error>> {
16550 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
16551 assert!(y.len() >= n * 2);
16552 let f = self.func("bf16_to_f16_flat");
16553 let n2 = (n / 2) as i64;
16554 let cfg = LaunchConfig {
16555 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
16556 block_dim: (256, 1, 1),
16557 shared_mem_bytes: 0,
16558 };
16559 let __s_b = self.gpu.stream();
16560 let mut b = __s_b.launch_builder(&f);
16561 b.arg(xb).arg(y).arg(&n2);
16562 unsafe {
16563 b.launch(cfg)?;
16564 }
16565 Ok(())
16566 }
16567
16568 #[allow(clippy::too_many_arguments)]
16573 pub fn fa_prefill_vl8(
16574 &self,
16575 seqs: &[FaSeqVl],
16576 head_dim: usize,
16577 n_head: usize,
16578 n_head_kv: usize,
16579 scale: f32,
16580 ) -> Result<(), Box<dyn std::error::Error>> {
16581 const BK: usize = 32;
16582 let b = seqs.len();
16583 assert!(b >= 1 && b <= 8);
16584 let mut packed = [FaSeqVl::default(); 8];
16585 packed[..b].copy_from_slice(seqs);
16586 let v = FaVl8(packed);
16587 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16588 let ept = (n_head_kv * head_dim) as i32;
16589 {
16590 let f = self.func("fa_mirror_vl");
16591 let max_n = (max_t as i64) * ept as i64;
16592 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
16593 for which in 0..2i32 {
16594 let cfg = LaunchConfig {
16595 grid_dim: (blocks, 1, b as u32),
16596 block_dim: (256, 1, 1),
16597 shared_mem_bytes: 0,
16598 };
16599 let __s_lb = self.gpu.stream();
16600 let mut lb = __s_lb.launch_builder(&f);
16601 lb.arg(&v).arg(&ept).arg(&which);
16602 unsafe {
16603 lb.launch(cfg)?;
16604 }
16605 }
16606 }
16607 let hd_sfx = fa_hd_suffix(head_dim)?;
16608 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
16609 let block_q = 64usize;
16610 let kv_stages = 2usize;
16611 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16612 + 4 * (block_q * BK + 2 * block_q)) as u32;
16613 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16614 f.set_attribute(
16615 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16616 shmem as i32,
16617 )?;
16618 let cfg = LaunchConfig {
16619 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
16620 block_dim: (32, 4, 1),
16621 shared_mem_bytes: shmem,
16622 };
16623 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16624 let __s_lb = self.gpu.stream();
16625 let mut lb = __s_lb.launch_builder(&f);
16626 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
16627 unsafe {
16628 lb.launch(cfg)?;
16629 }
16630 Ok(())
16631 }
16632
16633 #[allow(clippy::too_many_arguments)]
16637 pub fn attn_pre_vl8(
16638 &self,
16639 seqs: &[AttnPreVl],
16640 wq: &CudaSlice<f32>,
16641 wk: &CudaSlice<f32>,
16642 head_dim: usize,
16643 rope_dims: usize,
16644 n_head: usize,
16645 n_head_kv: usize,
16646 eps: f32,
16647 freq_base: f32,
16648 freq_scale: f32,
16649 kv_dim_k: usize,
16650 kv_dim_v: usize,
16651 k_tok_bytes: usize,
16652 v_tok_bytes: usize,
16653 ) -> Result<(), Box<dyn std::error::Error>> {
16654 let b = seqs.len();
16655 assert!(b >= 1 && b <= 8);
16656 let mut packed = [AttnPreVl::default(); 8];
16657 packed[..b].copy_from_slice(seqs);
16658 let v = AttnPreVl8(packed);
16659 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16660 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16661 {
16662 let f = self.func("q_gate_split_vl");
16663 let n = max_t * (n_head * head_dim) as u32;
16664 let cfg = LaunchConfig {
16665 grid_dim: (n.div_ceil(256), 1, b as u32),
16666 block_dim: (256, 1, 1),
16667 shared_mem_bytes: 0,
16668 };
16669 let __s_lb = self.gpu.stream();
16670 let mut lb = __s_lb.launch_builder(&f);
16671 lb.arg(&v).arg(&hd).arg(&nh);
16672 unsafe {
16673 lb.launch(cfg)?;
16674 }
16675 }
16676 {
16677 let f = self.func("attn_rms_vl");
16678 let cfg = LaunchConfig {
16679 grid_dim: (max_t * n_head as u32, 2, b as u32),
16680 block_dim: (rms_block(), 1, 1),
16681 shared_mem_bytes: 0,
16682 };
16683 let __s_lb = self.gpu.stream();
16684 let mut lb = __s_lb.launch_builder(&f);
16685 lb.arg(&v)
16686 .arg(wq)
16687 .arg(wk)
16688 .arg(&hd)
16689 .arg(&nh)
16690 .arg(&nhkv)
16691 .arg(&eps);
16692 unsafe {
16693 lb.launch(cfg)?;
16694 }
16695 }
16696 {
16697 let f = self.func("attn_rope_vl");
16698 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
16699 let nd = rope_dims as i32;
16700 let cfg = LaunchConfig {
16701 grid_dim: (max_t * n_head as u32, 2, b as u32),
16702 block_dim: ((head_dim / 2) as u32, 1, 1),
16703 shared_mem_bytes: 0,
16704 };
16705 let __s_lb = self.gpu.stream();
16706 let mut lb = __s_lb.launch_builder(&f);
16707 lb.arg(&v)
16708 .arg(&hd)
16709 .arg(&nd)
16710 .arg(&nh)
16711 .arg(&nhkv)
16712 .arg(&theta_scale)
16713 .arg(&freq_scale);
16714 unsafe {
16715 lb.launch(cfg)?;
16716 }
16717 }
16718 {
16719 let f = self.func("append_kv_vl");
16720 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
16721 let cfg = LaunchConfig {
16722 grid_dim: (nblk, max_t, b as u32),
16723 block_dim: (32, 1, 1),
16724 shared_mem_bytes: 0,
16725 };
16726 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
16727 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16728 let __s_lb = self.gpu.stream();
16729 let mut lb = __s_lb.launch_builder(&f);
16730 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
16731 unsafe {
16732 lb.launch(cfg)?;
16733 }
16734 }
16735 Ok(())
16736 }
16737
16738 pub fn fa_prefill_view(
16743 &self,
16744 q: &CudaSlice<f32>,
16745 k: &cudarc::driver::CudaView<u8>,
16746 v: &cudarc::driver::CudaView<u8>,
16747 o: &mut CudaSlice<f32>,
16748 head_dim: usize,
16749 n_head: usize,
16750 n_head_kv: usize,
16751 t: usize,
16752 t_kv: usize,
16753 scale: f32,
16754 causal: bool,
16755 k_tok_bytes: usize,
16756 v_tok_bytes: usize,
16757 g: bool,
16758 ) -> Result<(), Box<dyn std::error::Error>> {
16759 if portable_mma_gated() {
16760 return self.sdpa_naive_quantized_view(
16761 q,
16762 k,
16763 v,
16764 o,
16765 head_dim,
16766 n_head,
16767 n_head_kv,
16768 t,
16769 t_kv,
16770 scale,
16771 causal,
16772 k_tok_bytes,
16773 v_tok_bytes,
16774 );
16775 }
16776 const BLOCK_Q: usize = 64;
16777 const BK: usize = 32;
16778 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
16781 let f = if g {
16782 self.func_g(&name)
16783 } else {
16784 self.func(&name)
16785 };
16786 let shmem =
16787 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16788 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16789 f.set_attribute(
16790 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16791 shmem as i32,
16792 )?;
16793 let cfg = LaunchConfig {
16794 grid_dim: (
16795 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16796 n_head as u32,
16797 1,
16798 ),
16799 block_dim: (32, 4, 1),
16800 shared_mem_bytes: shmem,
16801 };
16802 let (hd, nh, nhkv, ti, tkvi, cz) = (
16803 head_dim as i32,
16804 n_head as i32,
16805 n_head_kv as i32,
16806 t as i32,
16807 t_kv as i32,
16808 causal as i32,
16809 );
16810 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16811 let __s_b = self.gpu.stream();
16812 let mut b = __s_b.launch_builder(&f);
16813 b.arg(q)
16814 .arg(k)
16815 .arg(v)
16816 .arg(o)
16817 .arg(&hd)
16818 .arg(&nh)
16819 .arg(&nhkv)
16820 .arg(&ti)
16821 .arg(&tkvi)
16822 .arg(&scale)
16823 .arg(&cz)
16824 .arg(&ktb)
16825 .arg(&vtb);
16826 unsafe {
16827 b.launch(cfg)?;
16828 }
16829 Ok(())
16830 }
16831
16832 #[allow(clippy::too_many_arguments)]
16842 pub fn fa_prefill_view_ws(
16843 &self,
16844 q: &CudaSlice<f32>,
16845 k: &cudarc::driver::CudaView<u8>,
16846 v: &cudarc::driver::CudaView<u8>,
16847 o: &mut CudaSlice<f32>,
16848 head_dim: usize,
16849 n_head: usize,
16850 n_head_kv: usize,
16851 t: usize,
16852 t_kv: usize,
16853 scale: f32,
16854 causal: bool,
16855 k_tok_bytes: usize,
16856 v_tok_bytes: usize,
16857 g: bool,
16858 ) -> Result<(), Box<dyn std::error::Error>> {
16859 if portable_mma_gated() {
16860 return self.sdpa_naive_quantized_view(
16861 q,
16862 k,
16863 v,
16864 o,
16865 head_dim,
16866 n_head,
16867 n_head_kv,
16868 t,
16869 t_kv,
16870 scale,
16871 causal,
16872 k_tok_bytes,
16873 v_tok_bytes,
16874 );
16875 }
16876 const BLOCK_Q: usize = 64;
16877 const BK: usize = 32;
16878 let kv_dim_k = n_head_kv * head_dim;
16879 let kv_dim_v = n_head_kv * head_dim;
16880 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
16882 let mut guard = self.prime_deqw_ws.lock().unwrap();
16884 let need_grow = match guard.as_ref() {
16885 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
16886 None => true,
16887 };
16888 if need_grow {
16889 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
16890 let (ck, cv) = guard
16891 .as_ref()
16892 .map(|(a, b)| (a.len(), b.len()))
16893 .unwrap_or((0, 0));
16894 *guard = Some((
16895 self.alloc_u8(grow(ck, k_ws_bytes))?,
16896 self.alloc_u8(grow(cv, v_ws_bytes))?,
16897 ));
16898 }
16899 let (kw, vw) = guard.as_mut().unwrap();
16900 {
16902 let f = if g {
16904 self.func_g("fa_dequant_kv_ws_bf16")
16905 } else {
16906 self.func("fa_dequant_kv_ws_bf16")
16907 };
16908 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
16909 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16910 let cfg = LaunchConfig {
16911 grid_dim: (nblk.max(1), 1, 1),
16912 block_dim: (256, 1, 1),
16913 shared_mem_bytes: 0,
16914 };
16915 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
16916 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16917 let __s_b = self.gpu.stream();
16918 let mut b = __s_b.launch_builder(&f);
16919 b.arg(k)
16920 .arg(v)
16921 .arg(&mut *kw)
16922 .arg(&mut *vw)
16923 .arg(&kdk)
16924 .arg(&kdv)
16925 .arg(&tkvi)
16926 .arg(&ktb)
16927 .arg(&vtb);
16928 unsafe {
16929 b.launch(cfg)?;
16930 }
16931 }
16932 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
16940 .map(|v| v != "0")
16941 .unwrap_or(true);
16942 {
16943 let hd_sfx = fa_hd_suffix(head_dim)?;
16944 let f = self.func(&format!(
16945 "fa_prefill_qw{}{hd_sfx}",
16946 if db { "_db" } else { "" }
16947 ));
16948 let shmem = if db {
16949 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
16951 } else {
16952 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
16953 };
16954 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16955 f.set_attribute(
16956 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16957 shmem as i32,
16958 )?;
16959 let cfg = LaunchConfig {
16960 grid_dim: (
16961 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16962 n_head as u32,
16963 1,
16964 ),
16965 block_dim: (32, 4, 1),
16966 shared_mem_bytes: shmem,
16967 };
16968 let (hd, nh, nhkv, ti, tkvi, cz) = (
16969 head_dim as i32,
16970 n_head as i32,
16971 n_head_kv as i32,
16972 t as i32,
16973 t_kv as i32,
16974 causal as i32,
16975 );
16976 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
16977 let __s_b = self.gpu.stream();
16978 let mut b = __s_b.launch_builder(&f);
16979 b.arg(q)
16980 .arg(&*kw)
16981 .arg(&*vw)
16982 .arg(o)
16983 .arg(&hd)
16984 .arg(&nh)
16985 .arg(&nhkv)
16986 .arg(&ti)
16987 .arg(&tkvi)
16988 .arg(&scale)
16989 .arg(&cz)
16990 .arg(&kdk)
16991 .arg(&kdv);
16992 unsafe {
16993 b.launch(cfg)?;
16994 }
16995 }
16996 Ok(())
16997 }
16998
16999 #[allow(clippy::too_many_arguments)]
17015 pub fn fa_prefill_view_ws_w_hd128(
17016 &self,
17017 q: &CudaSlice<f32>,
17018 k: &cudarc::driver::CudaView<u8>,
17019 v: &cudarc::driver::CudaView<u8>,
17020 o: &mut CudaSlice<f32>,
17021 head_dim: usize,
17022 n_head: usize,
17023 n_head_kv: usize,
17024 t: usize,
17025 t_kv: usize,
17026 scale: f32,
17027 causal: bool,
17028 window: usize,
17029 k_tok_bytes: usize,
17030 v_tok_bytes: usize,
17031 ) -> Result<(), Box<dyn std::error::Error>> {
17032 assert_eq!(
17033 head_dim, 128,
17034 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17035 );
17036 if portable_mma_gated() {
17037 return self.sdpa_naive_w_quantized_view(
17038 q,
17039 k,
17040 v,
17041 o,
17042 head_dim,
17043 n_head,
17044 n_head_kv,
17045 t,
17046 t_kv,
17047 scale,
17048 causal,
17049 window,
17050 k_tok_bytes,
17051 v_tok_bytes,
17052 );
17053 }
17054 const BLOCK_Q: usize = 64;
17055 const BK: usize = 32;
17056 let kv_dim_k = n_head_kv * head_dim;
17057 let kv_dim_v = n_head_kv * head_dim;
17058 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17060 let mut guard = self.prime_deqw_ws.lock().unwrap();
17061 let need_grow = match guard.as_ref() {
17062 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17063 None => true,
17064 };
17065 if need_grow {
17066 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17067 let (ck, cv) = guard
17068 .as_ref()
17069 .map(|(a, b)| (a.len(), b.len()))
17070 .unwrap_or((0, 0));
17071 *guard = Some((
17072 self.alloc_u8(grow(ck, k_ws_bytes))?,
17073 self.alloc_u8(grow(cv, v_ws_bytes))?,
17074 ));
17075 }
17076 let (kw, vw) = guard.as_mut().unwrap();
17077 {
17080 let f = self.func("fa_dequant_kv_ws_bf16");
17081 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17082 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17083 let cfg = LaunchConfig {
17084 grid_dim: (nblk.max(1), 1, 1),
17085 block_dim: (256, 1, 1),
17086 shared_mem_bytes: 0,
17087 };
17088 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17089 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17090 let __s_b = self.gpu.stream();
17091 let mut b = __s_b.launch_builder(&f);
17092 b.arg(k)
17093 .arg(v)
17094 .arg(&mut *kw)
17095 .arg(&mut *vw)
17096 .arg(&kdk)
17097 .arg(&kdv)
17098 .arg(&tkvi)
17099 .arg(&ktb)
17100 .arg(&vtb);
17101 unsafe {
17102 b.launch(cfg)?;
17103 }
17104 }
17105 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17107 .map(|v| v != "0")
17108 .unwrap_or(true);
17109 {
17110 let f = self.func(if db {
17111 "fa_prefill_qw_db_w_hd128"
17112 } else {
17113 "fa_prefill_qw_w_hd128"
17114 });
17115 let shmem = if db {
17116 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17117 } else {
17118 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17119 };
17120 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17121 f.set_attribute(
17122 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17123 shmem as i32,
17124 )?;
17125 let cfg = LaunchConfig {
17126 grid_dim: (
17127 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17128 n_head as u32,
17129 1,
17130 ),
17131 block_dim: (32, 4, 1),
17132 shared_mem_bytes: shmem,
17133 };
17134 let (hd, nh, nhkv, ti, tkvi, cz) = (
17135 head_dim as i32,
17136 n_head as i32,
17137 n_head_kv as i32,
17138 t as i32,
17139 t_kv as i32,
17140 causal as i32,
17141 );
17142 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17143 let __s_b = self.gpu.stream();
17144 let mut b = __s_b.launch_builder(&f);
17145 b.arg(q)
17146 .arg(&*kw)
17147 .arg(&*vw)
17148 .arg(o)
17149 .arg(&hd)
17150 .arg(&nh)
17151 .arg(&nhkv)
17152 .arg(&ti)
17153 .arg(&tkvi)
17154 .arg(&scale)
17155 .arg(&cz)
17156 .arg(&kdk)
17157 .arg(&kdv)
17158 .arg(&wnd);
17159 unsafe {
17160 b.launch(cfg)?;
17161 }
17162 }
17163 Ok(())
17164 }
17165
17166 pub fn fa_decode(
17170 &self,
17171 q: &CudaSlice<f32>,
17172 k: &cudarc::driver::CudaView<u8>,
17173 v: &cudarc::driver::CudaView<u8>,
17174 o: &mut CudaSlice<f32>,
17175 head_dim: usize,
17176 n_head: usize,
17177 n_head_kv: usize,
17178 t_kv: usize,
17179 scale: f32,
17180 k_tok_bytes: usize,
17181 v_tok_bytes: usize,
17182 ) -> Result<(), Box<dyn std::error::Error>> {
17183 self.fa_decode_kvmod(
17184 q,
17185 k,
17186 v,
17187 o,
17188 head_dim,
17189 n_head,
17190 n_head_kv,
17191 t_kv,
17192 scale,
17193 k_tok_bytes,
17194 v_tok_bytes,
17195 false,
17196 )
17197 }
17198
17199 #[allow(clippy::too_many_arguments)]
17203 #[allow(clippy::too_many_arguments)]
17207 #[allow(clippy::too_many_arguments)]
17208 fn fa_decode_scalar_unified(
17209 &self,
17210 q: &cudarc::driver::CudaView<f32>,
17211 k: &cudarc::driver::CudaView<u8>,
17212 v: &cudarc::driver::CudaView<u8>,
17213 o: &mut cudarc::driver::CudaViewMut<f32>,
17214 head_dim: usize,
17215 n_head: usize,
17216 n_head_kv: usize,
17217 t_kv_host: usize,
17218 t_kv_dev: Option<&CudaSlice<i32>>,
17219 scale: f32,
17220 n_splits: usize,
17221 split_keys: usize,
17222 k_tok_bytes: usize,
17223 v_tok_bytes: usize,
17224 g: bool,
17225 part_o: &mut CudaSlice<f32>,
17226 part_m: &mut CudaSlice<f32>,
17227 part_l: &mut CudaSlice<f32>,
17228 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
17229 ) -> Result<(), Box<dyn std::error::Error>> {
17230 let f = if g {
17231 self.func_g("fa_decode_f32")
17232 } else {
17233 self.fa_func("fa_decode_f32", head_dim)
17234 };
17235 let cfg = LaunchConfig {
17236 grid_dim: (n_head as u32, n_splits as u32, 1),
17237 block_dim: (head_dim as u32, 1, 1),
17238 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
17239 };
17240 let (hd, nh, nhkv, nsp) = (
17241 head_dim as i32,
17242 n_head as i32,
17243 n_head_kv as i32,
17244 n_splits as i32,
17245 );
17246 let (ktb, vtb, tkvi, ski) = (
17247 k_tok_bytes as i64,
17248 v_tok_bytes as i64,
17249 t_kv_host as i32,
17250 split_keys as i32,
17251 );
17252 let __s_b = self.gpu.stream();
17253 let mut b = __s_b.launch_builder(&f);
17254 match t_kv_dev {
17255 Some(d) => {
17256 b.arg(q)
17257 .arg(k)
17258 .arg(v)
17259 .arg(&mut *part_o)
17260 .arg(&mut *part_m)
17261 .arg(&mut *part_l)
17262 .arg(&hd)
17263 .arg(&nh)
17264 .arg(&nhkv)
17265 .arg(&tkvi)
17266 .arg(d)
17267 .arg(&scale)
17268 .arg(&nsp)
17269 .arg(&ski)
17270 .arg(&ktb)
17271 .arg(&vtb);
17272 unsafe {
17273 b.launch(cfg)?;
17274 }
17275 }
17276 None => {
17277 let null: u64 = 0;
17278 b.arg(q)
17279 .arg(k)
17280 .arg(v)
17281 .arg(&mut *part_o)
17282 .arg(&mut *part_m)
17283 .arg(&mut *part_l)
17284 .arg(&hd)
17285 .arg(&nh)
17286 .arg(&nhkv)
17287 .arg(&tkvi)
17288 .arg(&null)
17289 .arg(&scale)
17290 .arg(&nsp)
17291 .arg(&ski)
17292 .arg(&ktb)
17293 .arg(&vtb);
17294 unsafe {
17295 b.launch(cfg)?;
17296 }
17297 }
17298 }
17299 let cfg2 = LaunchConfig {
17300 grid_dim: (n_head as u32, 1, 1),
17301 block_dim: (head_dim as u32, 1, 1),
17302 shared_mem_bytes: 0,
17303 };
17304 if let Some((oq, od)) = q8_out {
17305 let fc = if g {
17307 self.func_g("fa_decode_combine_q8_1")
17308 } else {
17309 self.fa_func("fa_decode_combine_q8_1", head_dim)
17310 };
17311 let __s_b2 = self.gpu.stream();
17312 let mut b2 = __s_b2.launch_builder(&fc);
17313 b2.arg(&*part_o)
17314 .arg(&*part_m)
17315 .arg(&*part_l)
17316 .arg(oq)
17317 .arg(od)
17318 .arg(&hd)
17319 .arg(&nh)
17320 .arg(&nsp);
17321 unsafe {
17322 b2.launch(cfg2)?;
17323 }
17324 return Ok(());
17325 }
17326 let fc = if g {
17327 self.func_g("fa_decode_combine_f32")
17328 } else {
17329 self.fa_func("fa_decode_combine_f32", head_dim)
17330 };
17331 let __s_b2 = self.gpu.stream();
17332 let mut b2 = __s_b2.launch_builder(&fc);
17333 b2.arg(&*part_o)
17334 .arg(&*part_m)
17335 .arg(&*part_l)
17336 .arg(o)
17337 .arg(&hd)
17338 .arg(&nh)
17339 .arg(&nsp);
17340 unsafe {
17341 b2.launch(cfg2)?;
17342 }
17343 Ok(())
17344 }
17345
17346 pub fn fa_decode_kvmod(
17347 &self,
17348 q: &CudaSlice<f32>,
17349 k: &cudarc::driver::CudaView<u8>,
17350 v: &cudarc::driver::CudaView<u8>,
17351 o: &mut CudaSlice<f32>,
17352 head_dim: usize,
17353 n_head: usize,
17354 n_head_kv: usize,
17355 t_kv: usize,
17356 scale: f32,
17357 k_tok_bytes: usize,
17358 v_tok_bytes: usize,
17359 g: bool,
17360 ) -> Result<(), Box<dyn std::error::Error>> {
17361 let q_view = q.as_view();
17362 let mut o_view = o.as_view_mut();
17363 self.fa_decode_kvmod_view(
17364 &q_view,
17365 k,
17366 v,
17367 &mut o_view,
17368 head_dim,
17369 n_head,
17370 n_head_kv,
17371 t_kv,
17372 scale,
17373 k_tok_bytes,
17374 v_tok_bytes,
17375 g,
17376 )
17377 }
17378
17379 #[allow(clippy::too_many_arguments)]
17384 pub fn fa_decode_kvmod_view(
17385 &self,
17386 q: &cudarc::driver::CudaView<f32>,
17387 k: &cudarc::driver::CudaView<u8>,
17388 v: &cudarc::driver::CudaView<u8>,
17389 o: &mut cudarc::driver::CudaViewMut<f32>,
17390 head_dim: usize,
17391 n_head: usize,
17392 n_head_kv: usize,
17393 t_kv: usize,
17394 scale: f32,
17395 k_tok_bytes: usize,
17396 v_tok_bytes: usize,
17397 g: bool,
17398 ) -> Result<(), Box<dyn std::error::Error>> {
17399 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
17420 if g && head_dim == 256 && !fa_v4_at(t_kv) {
17424 fa_vec = false;
17425 }
17426 let sp = fa_split_keys(t_kv, n_head_kv);
17427 let n_splits = if fa_vec {
17428 ((t_kv + sp - 1) / sp).max(1)
17429 } else {
17430 ((t_kv + 255) / 256).max(1)
17431 };
17432 let o_len = n_head * n_splits * head_dim;
17433 let ml_len = n_head * n_splits;
17434 let mut part_guard = self.fa_part_pool.lock().unwrap();
17435 if part_guard
17436 .as_ref()
17437 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17438 .unwrap_or(true)
17439 {
17440 let old = part_guard.take();
17451 let (co, cm) = old
17452 .as_ref()
17453 .map(|pp| (pp.0.len(), pp.1.len()))
17454 .unwrap_or((0, 0));
17455 if let Some(old) = old {
17456 self.fa_part_retired.lock().unwrap().push(old);
17457 }
17458 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17459 eprintln!(
17460 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17461 co, o_len, cm, ml_len
17462 );
17463 }
17464 *part_guard = Some((
17465 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
17466 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17467 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17468 ));
17469 }
17470 let pg = part_guard.as_mut().unwrap();
17471 self.gpu
17472 .stream()
17473 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
17474 self.gpu
17475 .stream()
17476 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
17477 self.gpu
17478 .stream()
17479 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
17480 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
17481 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
17482 let (hd, nh, nhkv, tkvi, nsp) = (
17483 head_dim as i32,
17484 n_head as i32,
17485 n_head_kv as i32,
17486 t_kv as i32,
17487 n_splits as i32,
17488 );
17489 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17490 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
17494 let fa512_min = fa512_min_tkv();
17499 let deep = fa_vec
17502 && head_dim == 256
17503 && fa_v4_at(t_kv)
17504 && !g
17505 && fa_deep_at(t_kv)
17506 && !matches!(fa_v4_mode(), "noB3" | "stage");
17507 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
17508 let gqa = (n_head / n_head_kv).max(1) as u32;
17511 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
17512 (
17513 fv,
17514 LaunchConfig {
17515 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17516 block_dim: (32, gqa, 1),
17517 shared_mem_bytes: 0,
17518 },
17519 )
17520 } else if fa_vec && head_dim <= 256 {
17521 let gqa = (n_head / n_head_kv).max(1) as u32;
17522 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17533 let smem_tkv = *SMEM_TKV.get_or_init(|| {
17534 std::env::var("MEMRA_FA_SMEM_TKV")
17535 .ok()
17536 .and_then(|v| v.parse().ok())
17537 .unwrap_or_else(|| {
17538 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
17539 })
17540 });
17541 if fa_v4_at(t_kv) && head_dim == 256 {
17542 let v4name = match fa_v4_mode() {
17546 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
17549 _ => "fa_decode_vec_q_v4",
17550 };
17551 let fv = if g {
17552 self.func_g(v4name)
17553 } else {
17554 self.func(v4name)
17555 };
17556 let shmem = (if deep { 12160 } else { 11520 }
17559 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
17560 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17561 fv.set_attribute(
17562 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17563 shmem as i32,
17564 )?;
17565 (
17566 fv,
17567 LaunchConfig {
17568 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17569 block_dim: (32, gqa, 1),
17570 shared_mem_bytes: shmem,
17571 },
17572 )
17573 } else if fa_v3_active(head_dim) {
17574 let fv = if g {
17577 self.func_g("fa_decode_vec_q_v3")
17578 } else {
17579 self.func("fa_decode_vec_q_v3")
17580 };
17581 let shmem = (32 * head_dim * 2) as u32; (
17583 fv,
17584 LaunchConfig {
17585 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17586 block_dim: (32, gqa, 1),
17587 shared_mem_bytes: shmem,
17588 },
17589 )
17590 } else if fa_v2_on() {
17591 let fv = if g {
17595 self.func_g("fa_decode_vec_q_v2")
17596 } else {
17597 self.func("fa_decode_vec_q_v2")
17598 };
17599 let shmem = (2 * 32 * head_dim * 2) as u32; (
17601 fv,
17602 LaunchConfig {
17603 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17604 block_dim: (32, gqa, 1),
17605 shared_mem_bytes: shmem,
17606 },
17607 )
17608 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
17609 {
17610 let fv = if g {
17614 self.func_g("fa_decode_vec_q_smem")
17615 } else {
17616 self.func("fa_decode_vec_q_smem")
17617 };
17618 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
17620 fv.set_attribute(
17621 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17622 shmem as i32,
17623 )?;
17624 (
17625 fv,
17626 LaunchConfig {
17627 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17628 block_dim: (32, gqa, 1),
17629 shared_mem_bytes: shmem,
17630 },
17631 )
17632 } else {
17633 let fv = if g {
17636 self.func_g("fa_decode_vec_q")
17637 } else {
17638 self.func("fa_decode_vec_q")
17639 };
17640 (
17641 fv,
17642 LaunchConfig {
17643 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17644 block_dim: (32, gqa, 1),
17645 shared_mem_bytes: 0,
17646 },
17647 )
17648 }
17649 } else {
17650 return self.fa_decode_scalar_unified(
17653 q,
17654 k,
17655 v,
17656 o,
17657 head_dim,
17658 n_head,
17659 n_head_kv,
17660 t_kv,
17661 None,
17662 scale,
17663 n_splits,
17664 if fa_vec { sp } else { 256 },
17665 k_tok_bytes,
17666 v_tok_bytes,
17667 g,
17668 part_o,
17669 part_m,
17670 part_l,
17671 None,
17672 );
17673 };
17674 let __s_b = self.gpu.stream();
17675 let mut b = __s_b.launch_builder(&f);
17676 b.arg(q)
17677 .arg(k)
17678 .arg(v)
17679 .arg(&mut *part_o)
17680 .arg(&mut *part_m)
17681 .arg(&mut *part_l)
17682 .arg(&hd)
17683 .arg(&nh)
17684 .arg(&nhkv)
17685 .arg(&tkvi)
17686 .arg(&scale)
17687 .arg(&nsp)
17688 .arg(&ktb)
17689 .arg(&vtb);
17690 unsafe {
17691 b.launch(cfg)?;
17692 }
17693 let (fc, cfg2) = (
17696 if g {
17697 self.func_g("fa_decode_combine_f32")
17698 } else {
17699 self.fa_func("fa_decode_combine_f32", head_dim)
17700 },
17701 LaunchConfig {
17702 grid_dim: (n_head as u32, 1, 1),
17703 block_dim: (head_dim as u32, 1, 1),
17704 shared_mem_bytes: 0,
17705 },
17706 );
17707 let __s_b2 = self.gpu.stream();
17708 let mut b2 = __s_b2.launch_builder(&fc);
17709 b2.arg(&*part_o)
17710 .arg(&*part_m)
17711 .arg(&*part_l)
17712 .arg(o)
17713 .arg(&hd)
17714 .arg(&nh)
17715 .arg(&nsp);
17716 unsafe {
17717 b2.launch(cfg2)?;
17718 }
17719 Ok(())
17720 }
17721
17722 #[allow(clippy::too_many_arguments)]
17733 pub fn fa_decode_batch_seqs_v4(
17734 &self,
17735 q: &CudaSlice<f32>,
17736 kv_ptrs: &cudarc::driver::CudaView<u64>,
17737 pos_seq: &CudaSlice<i32>,
17738 o: &mut CudaSlice<f32>,
17739 head_dim: usize,
17740 n_head: usize,
17741 n_head_kv: usize,
17742 b_n: usize,
17743 t_kv_max: usize,
17744 scale: f32,
17745 split_keys: usize,
17746 k_tok_bytes: usize,
17747 v_tok_bytes: usize,
17748 ) -> Result<(), Box<dyn std::error::Error>> {
17749 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
17750 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
17751 let o_len = b_n * n_head * n_splits_max * head_dim;
17752 let ml_len = b_n * n_head * n_splits_max;
17753 let mut part_guard = self.fa_part_pool.lock().unwrap();
17754 if part_guard
17755 .as_ref()
17756 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17757 .unwrap_or(true)
17758 {
17759 let old = part_guard.take();
17770 let (co, cm) = old
17771 .as_ref()
17772 .map(|pp| (pp.0.len(), pp.1.len()))
17773 .unwrap_or((0, 0));
17774 if let Some(old) = old {
17775 self.fa_part_retired.lock().unwrap().push(old);
17776 }
17777 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17778 eprintln!(
17779 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17780 co, o_len, cm, ml_len
17781 );
17782 }
17783 *part_guard = Some((
17784 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
17785 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17786 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17787 ));
17788 }
17789 let pg = part_guard.as_mut().unwrap();
17790 self.gpu
17791 .stream()
17792 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
17793 self.gpu
17794 .stream()
17795 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
17796 self.gpu
17797 .stream()
17798 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
17799 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
17800 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17801 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
17802 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17803 let gqa = (n_head / n_head_kv).max(1) as u32;
17804 let f = self.func("fa_decode_vec_q_seqs_v4");
17805 let shmem = (11520 + 32 * head_dim * 2) as u32;
17807 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17808 f.set_attribute(
17809 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17810 shmem as i32,
17811 )?;
17812 let cfg = LaunchConfig {
17813 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
17814 block_dim: (32, gqa, 1),
17815 shared_mem_bytes: shmem,
17816 };
17817 {
17818 let __s_b = self.gpu.stream();
17819 let mut b = __s_b.launch_builder(&f);
17820 b.arg(q)
17821 .arg(kv_ptrs)
17822 .arg(pos_seq)
17823 .arg(&mut *part_o)
17824 .arg(&mut *part_m)
17825 .arg(&mut *part_l)
17826 .arg(&hd)
17827 .arg(&nh)
17828 .arg(&nhkv)
17829 .arg(&scale)
17830 .arg(&nspm)
17831 .arg(&spk)
17832 .arg(&ktb)
17833 .arg(&vtb);
17834 unsafe {
17835 b.launch(cfg)?;
17836 }
17837 }
17838 let fc = self.func("fa_decode_combine_seqs");
17839 let cfg2 = LaunchConfig {
17840 grid_dim: (n_head as u32, b_n as u32, 1),
17841 block_dim: (head_dim as u32, 1, 1),
17842 shared_mem_bytes: 0,
17843 };
17844 let __s_b2 = self.gpu.stream();
17845 let mut b2 = __s_b2.launch_builder(&fc);
17846 b2.arg(&*part_o)
17847 .arg(&*part_m)
17848 .arg(&*part_l)
17849 .arg(o)
17850 .arg(&hd)
17851 .arg(&nh)
17852 .arg(pos_seq)
17853 .arg(&nspm)
17854 .arg(&spk);
17855 unsafe {
17856 b2.launch(cfg2)?;
17857 }
17858 Ok(())
17859 }
17860
17861 #[allow(clippy::too_many_arguments)]
17868 pub fn append_kv_quantized_seqs(
17869 &self,
17870 k_rows: &CudaSlice<f32>,
17871 v_rows: &CudaSlice<f32>,
17872 kv_ptrs: &cudarc::driver::CudaView<u64>,
17873 pos_seq: &CudaSlice<i32>,
17874 b_n: usize,
17875 kv_dim_k: usize,
17876 kv_dim_v: usize,
17877 k_tok_bytes: usize,
17878 v_tok_bytes: usize,
17879 ) -> Result<(), Box<dyn std::error::Error>> {
17880 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
17881 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
17882 let cfg = LaunchConfig {
17883 grid_dim: (nblk, b_n as u32, 1),
17884 block_dim: (32, 1, 1),
17885 shared_mem_bytes: 0,
17886 };
17887 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17888 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17889 let __s_b = self.gpu.stream();
17890 let mut b = __s_b.launch_builder(&f);
17891 b.arg(k_rows)
17892 .arg(v_rows)
17893 .arg(kv_ptrs)
17894 .arg(pos_seq)
17895 .arg(&kdk)
17896 .arg(&kdv)
17897 .arg(&ktb)
17898 .arg(&vtb);
17899 unsafe {
17900 b.launch(cfg)?;
17901 }
17902 Ok(())
17903 }
17904
17905 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
17911 std::env::var("MEMRA_NO_FA_VEC").is_err()
17912 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
17913 && base_len + 1 >= fa_vec_min_tkv()
17914 && head_dim <= 256
17915 && head_dim % 32 == 0
17916 }
17917
17918 #[allow(clippy::too_many_arguments)]
17927 pub fn fa_decode_rows(
17928 &self,
17929 q: &CudaSlice<f32>,
17930 k: &cudarc::driver::CudaView<u8>,
17931 v: &cudarc::driver::CudaView<u8>,
17932 o: &mut CudaSlice<f32>,
17933 head_dim: usize,
17934 n_head: usize,
17935 n_head_kv: usize,
17936 base_len: usize,
17937 t: usize,
17938 scale: f32,
17939 k_tok_bytes: usize,
17940 v_tok_bytes: usize,
17941 base_dev: Option<(&CudaSlice<i32>, i32)>,
17945 kv_shared: bool,
17948 g: bool,
17952 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
17955 ) -> Result<(), Box<dyn std::error::Error>> {
17956 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
17957 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
17964 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17965 let v = *SP512.get_or_init(|| {
17968 std::env::var("MEMRA_FA_SP512")
17969 .ok()
17970 .and_then(|x| x.parse().ok())
17971 .unwrap_or(0)
17972 });
17973 sp = if v >= 8 {
17974 v
17975 } else {
17976 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
17977 };
17978 }
17979 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17980 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17981 let gqa = (n_head / n_head_kv).max(1) as u32;
17982 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
17993 groups.push((0, t, sp));
17994 } else {
17995 let mut r0 = 0usize;
17996 while r0 < t {
17997 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
17998 let mut r1 = r0 + 1;
17999 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18000 r1 += 1;
18001 }
18002 groups.push((r0, r1 - r0, sp_g));
18003 r0 = r1;
18004 }
18005 }
18006 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18010 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18011 std::env::var("MEMRA_FA_SMEM_TKV")
18012 .ok()
18013 .and_then(|v| v.parse().ok())
18014 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18015 });
18016 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18017 let v3 = fa_v3_active(head_dim);
18018 let smem_rows =
18019 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18020 let _ = kv_shared;
18025 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18028 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18042 let tb512 = head_dim == 512
18044 && sp <= 32
18045 && n_head / n_head_kv.max(1) <= 16
18046 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18047 let fname = if tb512 {
18048 "fa_decode_vec_q_rows_v4_512_tb"
18049 } else if i2 {
18050 "fa_decode_vec_q_rows_dpl16_i2"
18051 } else if head_dim == 512 {
18052 "fa_decode_vec_q_rows_dpl16"
18053 }
18054 else if v4 {
18056 "fa_decode_vec_q_rows_v4"
18057 } else if v3 {
18058 "fa_decode_vec_q_rows_v3"
18059 } else if fa_v2_on() {
18060 "fa_decode_vec_q_rows_v2"
18061 } else if smem_rows {
18062 "fa_decode_vec_q_rows_smem"
18063 } else {
18064 "fa_decode_vec_q_rows"
18065 };
18066 let f = if head_dim == 512 {
18067 self.fa_func(fname, head_dim)
18068 } else if g {
18069 self.func_g(if smem_rows {
18077 "fa_decode_vec_q_rows"
18078 } else {
18079 fname
18080 })
18081 } else {
18082 self.func(fname)
18083 };
18084 let shmem = if tb512 {
18085 let gk = Self::gkv_on();
18087 let sh =
18088 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18089 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18090 f.set_attribute(
18091 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18092 sh as i32,
18093 )?;
18094 sh
18095 } else if v4 || v3 || smem_rows || fa_v2_on() {
18096 let sh = (if v4 {
18098 11520 + 32 * head_dim * if g { 1 } else { 2 }
18099 } else if v3 {
18100 32 * head_dim * 2
18101 } else {
18102 2 * 32 * head_dim * 2
18103 }) as u32;
18104 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18105 f.set_attribute(
18106 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18107 sh as i32,
18108 )?;
18109 sh
18110 } else {
18111 0
18112 };
18113 for &(r0, t_g, sp_g) in &groups {
18117 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18118 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18119 let base_i = (base_len + r0) as i32;
18120 let o_len = t_g * n_head * n_splits_g * head_dim;
18121 let ml_len = t_g * n_head * n_splits_g;
18122 let mut part_guard = self.fa_part_pool.lock().unwrap();
18123 if part_guard
18124 .as_ref()
18125 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18126 .unwrap_or(true)
18127 {
18128 let old = part_guard.take();
18139 let (co, cm) = old
18140 .as_ref()
18141 .map(|pp| (pp.0.len(), pp.1.len()))
18142 .unwrap_or((0, 0));
18143 if let Some(old) = old {
18144 self.fa_part_retired.lock().unwrap().push(old);
18145 }
18146 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18147 eprintln!(
18148 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18149 co, o_len, cm, ml_len
18150 );
18151 }
18152 *part_guard = Some((
18153 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18154 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18155 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18156 ));
18157 }
18158 let pg = part_guard.as_mut().unwrap();
18159 self.gpu
18160 .stream()
18161 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18162 self.gpu
18163 .stream()
18164 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18165 self.gpu
18166 .stream()
18167 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18168 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18169 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18170 let qv = self.view(q, t * n_head * head_dim);
18171 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18172 let cfg = LaunchConfig {
18173 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
18174 block_dim: (32, gqa, 1),
18175 shared_mem_bytes: shmem,
18176 };
18177 {
18178 let __s_b = self.gpu.stream();
18179 let mut b = __s_b.launch_builder(&f);
18180 if tb512 {
18181 let (bd, plus) =
18183 base_dev.expect("hd512 rows twin requires a device base counter");
18184 let plus_g = plus + r0 as i32;
18185 let nr = t_g as i32;
18186 if Self::pdl_on() && Self::pdl_wb_on() {
18187 use cudarc::driver::{DevicePtr, DevicePtrMut};
18189 let s = &self.gpu.stream();
18190 let (pq, _b0) = q_g.device_ptr(s);
18191 let (pk, _b1) = k.device_ptr(s);
18192 let (pv, _b2) = v.device_ptr(s);
18193 let (po, _b3) = part_o.device_ptr_mut(s);
18194 let (pm, _b4) = part_m.device_ptr_mut(s);
18195 let (pl, _b5) = part_l.device_ptr_mut(s);
18196 let (pb, _b6) = bd.device_ptr(s);
18197 let mut ps = [
18198 &pq as *const _ as *mut std::ffi::c_void,
18199 &pk as *const _ as *mut _,
18200 &pv as *const _ as *mut _,
18201 &po as *const _ as *mut _,
18202 &pm as *const _ as *mut _,
18203 &pl as *const _ as *mut _,
18204 &hd as *const _ as *mut _,
18205 &nh as *const _ as *mut _,
18206 &nhkv as *const _ as *mut _,
18207 &pb as *const _ as *mut _,
18208 &plus_g as *const _ as *mut _,
18209 &scale as *const _ as *mut _,
18210 &nspm as *const _ as *mut _,
18211 &spk as *const _ as *mut _,
18212 &ktb as *const _ as *mut _,
18213 &vtb as *const _ as *mut _,
18214 &nr as *const _ as *mut _,
18215 ];
18216 unsafe {
18217 self.launch_pdl_flash(
18218 Self::gkv_on(),
18219 "fa_decode_vec_q_rows_v4_512_tb",
18220 (n_head_kv as u32, n_splits_g as u32, 1),
18221 (32, gqa, 1),
18222 shmem,
18223 &mut ps,
18224 )?;
18225 }
18226 } else {
18227 let cfg_tb = LaunchConfig {
18228 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
18229 block_dim: (32, gqa, 1),
18230 shared_mem_bytes: shmem,
18231 };
18232 b.arg(&q_g)
18233 .arg(k)
18234 .arg(v)
18235 .arg(&mut *part_o)
18236 .arg(&mut *part_m)
18237 .arg(&mut *part_l)
18238 .arg(&hd)
18239 .arg(&nh)
18240 .arg(&nhkv)
18241 .arg(bd)
18242 .arg(&plus_g)
18243 .arg(&scale)
18244 .arg(&nspm)
18245 .arg(&spk)
18246 .arg(&ktb)
18247 .arg(&vtb)
18248 .arg(&nr);
18249 unsafe {
18250 b.launch(cfg_tb)?;
18251 }
18252 }
18253 } else if head_dim == 512 {
18254 let (bd, plus) =
18255 base_dev.expect("hd512 rows twin requires a device base counter");
18256 let plus_g = plus + r0 as i32;
18257 b.arg(&q_g)
18258 .arg(k)
18259 .arg(v)
18260 .arg(&mut *part_o)
18261 .arg(&mut *part_m)
18262 .arg(&mut *part_l)
18263 .arg(&hd)
18264 .arg(&nh)
18265 .arg(&nhkv)
18266 .arg(bd)
18267 .arg(&plus_g)
18268 .arg(&scale)
18269 .arg(&nspm)
18270 .arg(&spk)
18271 .arg(&ktb)
18272 .arg(&vtb);
18273 unsafe {
18274 b.launch(cfg)?;
18275 }
18276 } else {
18277 b.arg(&q_g)
18278 .arg(k)
18279 .arg(v)
18280 .arg(&mut *part_o)
18281 .arg(&mut *part_m)
18282 .arg(&mut *part_l)
18283 .arg(&hd)
18284 .arg(&nh)
18285 .arg(&nhkv)
18286 .arg(&base_i)
18287 .arg(&scale)
18288 .arg(&nspm)
18289 .arg(&spk)
18290 .arg(&ktb)
18291 .arg(&vtb);
18292 unsafe {
18293 b.launch(cfg)?;
18294 }
18295 }
18296 }
18297 let cfg2 = LaunchConfig {
18298 grid_dim: (n_head as u32, t_g as u32, 1),
18299 block_dim: (head_dim as u32, 1, 1),
18300 shared_mem_bytes: 0,
18301 };
18302 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18303 if head_dim == 512 {
18304 let (bd, plus) = base_dev.unwrap();
18307 let plus_g = plus + r0 as i32;
18308 if let Some((oq, od)) = q8_out.as_mut() {
18309 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
18311 if Self::pdl_on() && Self::pdl_wb_on() {
18312 use cudarc::driver::{DevicePtr, DevicePtrMut};
18314 let s = &self.gpu.stream();
18315 let (po, _g0) = part_o.device_ptr(s);
18316 let (pm, _g1) = part_m.device_ptr(s);
18317 let (pl, _g2) = part_l.device_ptr(s);
18318 let (pq, _g3) = oq.device_ptr_mut(s);
18319 let (pd, _g4) = od.device_ptr_mut(s);
18320 let (pb, _g5) = bd.device_ptr(s);
18321 let mut ps = [
18322 &po as *const _ as *mut std::ffi::c_void,
18323 &pm as *const _ as *mut _,
18324 &pl as *const _ as *mut _,
18325 &pq as *const _ as *mut _,
18326 &pd as *const _ as *mut _,
18327 &hd as *const _ as *mut _,
18328 &nh as *const _ as *mut _,
18329 &pb as *const _ as *mut _,
18330 &plus_g as *const _ as *mut _,
18331 &nspm as *const _ as *mut _,
18332 &spk as *const _ as *mut _,
18333 ];
18334 unsafe {
18335 self.launch_pdl_flash(
18336 Self::gkv_on(),
18337 "fa_decode_combine_rows_dc_q8_1",
18338 cfg2.grid_dim,
18339 cfg2.block_dim,
18340 0,
18341 &mut ps,
18342 )?;
18343 }
18344 continue;
18345 }
18346 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
18347 let __s_b2 = self.gpu.stream();
18348 let mut b2 = __s_b2.launch_builder(&fc);
18349 b2.arg(&*part_o)
18350 .arg(&*part_m)
18351 .arg(&*part_l)
18352 .arg(&mut **oq)
18353 .arg(&mut **od)
18354 .arg(&hd)
18355 .arg(&nh)
18356 .arg(bd)
18357 .arg(&plus_g)
18358 .arg(&nspm)
18359 .arg(&spk);
18360 unsafe {
18361 b2.launch(cfg2)?;
18362 }
18363 continue;
18364 }
18365 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
18366 let __s_b2 = self.gpu.stream();
18367 let mut b2 = __s_b2.launch_builder(&fc);
18368 b2.arg(&*part_o)
18369 .arg(&*part_m)
18370 .arg(&*part_l)
18371 .arg(&mut o_g)
18372 .arg(&hd)
18373 .arg(&nh)
18374 .arg(bd)
18375 .arg(&plus_g)
18376 .arg(&nspm)
18377 .arg(&spk);
18378 unsafe {
18379 b2.launch(cfg2)?;
18380 }
18381 } else {
18382 assert!(
18385 q8_out.is_none(),
18386 "rows q8 emit requires the hd512 dc combine"
18387 );
18388 let fc = self.func("fa_decode_combine_rows");
18389 let __s_b2 = self.gpu.stream();
18390 let mut b2 = __s_b2.launch_builder(&fc);
18391 b2.arg(&*part_o)
18392 .arg(&*part_m)
18393 .arg(&*part_l)
18394 .arg(&mut o_g)
18395 .arg(&hd)
18396 .arg(&nh)
18397 .arg(&base_i)
18398 .arg(&nspm)
18399 .arg(&spk);
18400 unsafe {
18401 b2.launch(cfg2)?;
18402 }
18403 }
18404 }
18405 Ok(())
18406 }
18407
18408 #[allow(clippy::too_many_arguments)]
18412 pub fn fa_decode_rows_w(
18413 &self,
18414 q: &CudaSlice<f32>,
18415 k: &cudarc::driver::CudaView<u8>,
18416 v: &cudarc::driver::CudaView<u8>,
18417 o: &mut CudaSlice<f32>,
18418 head_dim: usize,
18419 n_head: usize,
18420 n_head_kv: usize,
18421 base_dev: &CudaSlice<i32>,
18422 base_plus: i32,
18423 t: usize,
18424 scale: f32,
18425 window: usize,
18426 k_tok_bytes: usize,
18427 v_tok_bytes: usize,
18428 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18429 ) -> Result<(), Box<dyn std::error::Error>> {
18430 debug_assert!(head_dim == 256);
18435 let sp = {
18443 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18444 let v = *SPW.get_or_init(|| {
18445 std::env::var("MEMRA_FA_SPW")
18446 .ok()
18447 .and_then(|x| x.parse().ok())
18448 .unwrap_or(0)
18449 });
18450 if v >= 8 {
18451 v
18452 } else {
18453 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18454 }
18455 };
18456 let n_splits_max = (window + sp - 1) / sp;
18457 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18458 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
18459 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18460 let gqa = (n_head / n_head_kv).max(1) as u32;
18461 let o_len = t * n_head * n_splits_max * head_dim;
18462 let ml_len = t * n_head * n_splits_max;
18463 let mut part_guard = self.fa_part_pool.lock().unwrap();
18464 if part_guard
18465 .as_ref()
18466 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18467 .unwrap_or(true)
18468 {
18469 let old = part_guard.take();
18480 let (co, cm) = old
18481 .as_ref()
18482 .map(|pp| (pp.0.len(), pp.1.len()))
18483 .unwrap_or((0, 0));
18484 if let Some(old) = old {
18485 self.fa_part_retired.lock().unwrap().push(old);
18486 }
18487 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18488 eprintln!(
18489 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18490 co, o_len, cm, ml_len
18491 );
18492 }
18493 *part_guard = Some((
18494 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18495 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18496 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18497 ));
18498 }
18499 let pg = part_guard.as_mut().unwrap();
18500 self.gpu
18501 .stream()
18502 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18503 self.gpu
18504 .stream()
18505 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18506 self.gpu
18507 .stream()
18508 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18509 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18510 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18516 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
18517 std::env::var("MEMRA_FA_SMEM_TKV")
18518 .ok()
18519 .and_then(|v| v.parse().ok())
18520 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18521 });
18522 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18528 let wg = Self::wkv_on();
18533 let sp2 =
18536 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
18537 if sp2 {
18538 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18539 if Self::pdl_on() && Self::pdl_wb_on() {
18540 use cudarc::driver::{DevicePtr, DevicePtrMut};
18542 let s = &self.gpu.stream();
18543 let (pq, _b0) = q.device_ptr(s);
18544 let (pk, _b1) = k.device_ptr(s);
18545 let (pv, _b2) = v.device_ptr(s);
18546 let (po, _b3) = part_o.device_ptr_mut(s);
18547 let (pm, _b4) = part_m.device_ptr_mut(s);
18548 let (pl, _b5) = part_l.device_ptr_mut(s);
18549 let (pb, _b6) = base_dev.device_ptr(s);
18550 let mut ps = [
18551 &pq as *const _ as *mut std::ffi::c_void,
18552 &pk as *const _ as *mut _,
18553 &pv as *const _ as *mut _,
18554 &po as *const _ as *mut _,
18555 &pm as *const _ as *mut _,
18556 &pl as *const _ as *mut _,
18557 &hd as *const _ as *mut _,
18558 &nh as *const _ as *mut _,
18559 &nhkv as *const _ as *mut _,
18560 &pb as *const _ as *mut _,
18561 &base_plus as *const _ as *mut _,
18562 &scale as *const _ as *mut _,
18563 &nspm as *const _ as *mut _,
18564 &spk as *const _ as *mut _,
18565 &ktb as *const _ as *mut _,
18566 &vtb as *const _ as *mut _,
18567 &wini as *const _ as *mut _,
18568 ];
18569 unsafe {
18570 self.launch_pdl_flash(
18571 wg,
18572 "fa_decode_vec_q_rows_v4_w_sp",
18573 (n_head_kv as u32, n_splits_max as u32, t as u32),
18574 (32, gqa + 1, 1),
18575 sh,
18576 &mut ps,
18577 )?;
18578 }
18579 } else {
18580 let f = if wg {
18581 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
18582 } else {
18583 self.func("fa_decode_vec_q_rows_v4_w_sp")
18584 };
18585 f.set_attribute(
18586 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18587 sh as i32,
18588 )?;
18589 let cfg = LaunchConfig {
18590 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18591 block_dim: (32, gqa + 1, 1),
18592 shared_mem_bytes: sh,
18593 };
18594 let __s_b = self.gpu.stream();
18595 let mut b = __s_b.launch_builder(&f);
18596 b.arg(q)
18597 .arg(k)
18598 .arg(v)
18599 .arg(&mut *part_o)
18600 .arg(&mut *part_m)
18601 .arg(&mut *part_l)
18602 .arg(&hd)
18603 .arg(&nh)
18604 .arg(&nhkv)
18605 .arg(base_dev)
18606 .arg(&base_plus)
18607 .arg(&scale)
18608 .arg(&nspm)
18609 .arg(&spk)
18610 .arg(&ktb)
18611 .arg(&vtb)
18612 .arg(&wini);
18613 unsafe {
18614 b.launch(cfg)?;
18615 }
18616 }
18617 } else {
18618 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
18619 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18621 use cudarc::driver::{DevicePtr, DevicePtrMut};
18622 let s = &self.gpu.stream();
18623 let (pq, _b0) = q.device_ptr(s);
18624 let (pk, _b1) = k.device_ptr(s);
18625 let (pv, _b2) = v.device_ptr(s);
18626 let (po, _b3) = part_o.device_ptr_mut(s);
18627 let (pm, _b4) = part_m.device_ptr_mut(s);
18628 let (pl, _b5) = part_l.device_ptr_mut(s);
18629 let (pb, _b6) = base_dev.device_ptr(s);
18630 let mut ps = [
18631 &pq as *const _ as *mut std::ffi::c_void,
18632 &pk as *const _ as *mut _,
18633 &pv as *const _ as *mut _,
18634 &po as *const _ as *mut _,
18635 &pm as *const _ as *mut _,
18636 &pl as *const _ as *mut _,
18637 &hd as *const _ as *mut _,
18638 &nh as *const _ as *mut _,
18639 &nhkv as *const _ as *mut _,
18640 &pb as *const _ as *mut _,
18641 &base_plus as *const _ as *mut _,
18642 &scale as *const _ as *mut _,
18643 &nspm as *const _ as *mut _,
18644 &spk as *const _ as *mut _,
18645 &ktb as *const _ as *mut _,
18646 &vtb as *const _ as *mut _,
18647 &wini as *const _ as *mut _,
18648 ];
18649 unsafe {
18650 self.launch_pdl_flash(
18651 wg,
18652 "fa_decode_vec_q_rows_v4_w",
18653 (n_head_kv as u32, n_splits_max as u32, t as u32),
18654 (32, gqa, 1),
18655 sh,
18656 &mut ps,
18657 )?;
18658 }
18659 } else {
18660 let pick = |name: &str| {
18661 if wg {
18662 self.func_g(name)
18663 } else {
18664 self.func(name)
18665 }
18666 };
18667 let (f, sh) = if fa_v4_at(window) {
18668 let f = pick("fa_decode_vec_q_rows_v4_w");
18669 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
18670 } else if smem_tkv > 0 && window >= smem_tkv {
18671 (
18674 pick("fa_decode_vec_q_rows_smem_w"),
18675 (2 * 32 * head_dim * 2) as u32,
18676 )
18677 } else {
18678 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
18679 };
18680 f.set_attribute(
18681 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18682 sh as i32,
18683 )?;
18684 let cfg = LaunchConfig {
18685 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18686 block_dim: (32, gqa, 1),
18687 shared_mem_bytes: sh,
18688 };
18689 let __s_b = self.gpu.stream();
18690 let mut b = __s_b.launch_builder(&f);
18691 b.arg(q)
18692 .arg(k)
18693 .arg(v)
18694 .arg(&mut *part_o)
18695 .arg(&mut *part_m)
18696 .arg(&mut *part_l)
18697 .arg(&hd)
18698 .arg(&nh)
18699 .arg(&nhkv)
18700 .arg(base_dev)
18701 .arg(&base_plus)
18702 .arg(&scale)
18703 .arg(&nspm)
18704 .arg(&spk)
18705 .arg(&ktb)
18706 .arg(&vtb)
18707 .arg(&wini);
18708 unsafe {
18709 b.launch(cfg)?;
18710 }
18711 }
18712 }
18713 let cfg2 = LaunchConfig {
18714 grid_dim: (n_head as u32, t as u32, 1),
18715 block_dim: (head_dim as u32, 1, 1),
18716 shared_mem_bytes: 0,
18717 };
18718 if let Some((oq, od)) = q8_out {
18719 if Self::pdl_on() && Self::pdl_wb_on() {
18722 use cudarc::driver::{DevicePtr, DevicePtrMut};
18724 let s = &self.gpu.stream();
18725 let (po, _g0) = part_o.device_ptr(s);
18726 let (pm, _g1) = part_m.device_ptr(s);
18727 let (pl, _g2) = part_l.device_ptr(s);
18728 let (pq, _g3) = oq.device_ptr_mut(s);
18729 let (pd, _g4) = od.device_ptr_mut(s);
18730 let mut ps = [
18731 &po as *const _ as *mut std::ffi::c_void,
18732 &pm as *const _ as *mut _,
18733 &pl as *const _ as *mut _,
18734 &pq as *const _ as *mut _,
18735 &pd as *const _ as *mut _,
18736 &hd as *const _ as *mut _,
18737 &nh as *const _ as *mut _,
18738 &nspm as *const _ as *mut _,
18739 &spk as *const _ as *mut _,
18740 &wini as *const _ as *mut _,
18741 ];
18742 unsafe {
18743 self.launch_pdl_flash(
18744 wg,
18745 "fa_decode_combine_rows_w_q8_1",
18746 cfg2.grid_dim,
18747 cfg2.block_dim,
18748 0,
18749 &mut ps,
18750 )?;
18751 }
18752 return Ok(());
18753 }
18754 let fc = if wg {
18755 self.func_g("fa_decode_combine_rows_w_q8_1")
18756 } else {
18757 self.func("fa_decode_combine_rows_w_q8_1")
18758 };
18759 let __s_b2 = self.gpu.stream();
18760 let mut b2 = __s_b2.launch_builder(&fc);
18761 b2.arg(&*part_o)
18762 .arg(&*part_m)
18763 .arg(&*part_l)
18764 .arg(oq)
18765 .arg(od)
18766 .arg(&hd)
18767 .arg(&nh)
18768 .arg(&nspm)
18769 .arg(&spk)
18770 .arg(&wini);
18771 unsafe {
18772 b2.launch(cfg2)?;
18773 }
18774 return Ok(());
18775 }
18776 let fc = if wg {
18777 self.func_g("fa_decode_combine_rows_w")
18778 } else {
18779 self.func("fa_decode_combine_rows_w")
18780 };
18781 let __s_b2 = self.gpu.stream();
18782 let mut b2 = __s_b2.launch_builder(&fc);
18783 b2.arg(&*part_o)
18784 .arg(&*part_m)
18785 .arg(&*part_l)
18786 .arg(o)
18787 .arg(&hd)
18788 .arg(&nh)
18789 .arg(&nspm)
18790 .arg(&spk)
18791 .arg(&wini);
18792 unsafe {
18793 b2.launch(cfg2)?;
18794 }
18795 Ok(())
18796 }
18797
18798 #[allow(clippy::too_many_arguments)]
18804 pub fn fa_decode_rows_dc(
18805 &self,
18806 q: &CudaSlice<f32>,
18807 k: &cudarc::driver::CudaView<u8>,
18808 v: &cudarc::driver::CudaView<u8>,
18809 o: &mut CudaSlice<f32>,
18810 head_dim: usize,
18811 n_head: usize,
18812 n_head_kv: usize,
18813 base_dev: &CudaSlice<i32>,
18814 t_kv_upper: usize,
18815 t: usize,
18816 scale: f32,
18817 k_tok_bytes: usize,
18818 v_tok_bytes: usize,
18819 base_plus: i32,
18820 g: bool,
18821 ) -> Result<(), Box<dyn std::error::Error>> {
18822 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
18823 assert!(
18824 v4 || fa_v3_active(head_dim),
18825 "stream fa rows requires the v3 or v4 lane"
18826 );
18827 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
18828 if v4 {
18829 let sp = fa_split_keys(t_kv_upper, n_head_kv);
18830 let n_splits_max = (t_kv_upper + sp - 1) / sp;
18831 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18832 let (nspm, spk) = (n_splits_max as i32, sp as i32);
18833 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18834 let gqa = (n_head / n_head_kv).max(1) as u32;
18835 let o_len = t * n_head * n_splits_max * head_dim;
18836 let ml_len = t * n_head * n_splits_max;
18837 let mut part_guard = self.fa_part_pool.lock().unwrap();
18838 if part_guard
18839 .as_ref()
18840 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18841 .unwrap_or(true)
18842 {
18843 let old = part_guard.take();
18854 let (co, cm) = old
18855 .as_ref()
18856 .map(|pp| (pp.0.len(), pp.1.len()))
18857 .unwrap_or((0, 0));
18858 if let Some(old) = old {
18859 self.fa_part_retired.lock().unwrap().push(old);
18860 }
18861 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18862 eprintln!(
18863 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18864 co, o_len, cm, ml_len
18865 );
18866 }
18867 *part_guard = Some((
18868 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18869 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18870 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18871 ));
18872 }
18873 let pg = part_guard.as_mut().unwrap();
18874 self.gpu
18875 .stream()
18876 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18877 self.gpu
18878 .stream()
18879 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18880 self.gpu
18881 .stream()
18882 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18883 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18884 let f = if g {
18885 self.func_g("fa_decode_vec_q_rows_v4_dc")
18886 } else {
18887 self.func("fa_decode_vec_q_rows_v4_dc")
18888 };
18889 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
18890 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18891 f.set_attribute(
18892 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18893 sh as i32,
18894 )?;
18895 let cfg = LaunchConfig {
18896 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18897 block_dim: (32, gqa, 1),
18898 shared_mem_bytes: sh,
18899 };
18900 let __s_b = self.gpu.stream();
18901 let mut b = __s_b.launch_builder(&f);
18902 b.arg(q)
18903 .arg(k)
18904 .arg(v)
18905 .arg(&mut *part_o)
18906 .arg(&mut *part_m)
18907 .arg(&mut *part_l)
18908 .arg(&hd)
18909 .arg(&nh)
18910 .arg(&nhkv)
18911 .arg(base_dev)
18912 .arg(&base_plus)
18913 .arg(&scale)
18914 .arg(&nspm)
18915 .arg(&spk)
18916 .arg(&ktb)
18917 .arg(&vtb);
18918 unsafe {
18919 b.launch(cfg)?;
18920 }
18921 let fc = self.func("fa_decode_combine_rows_dc");
18922 let cfg2 = LaunchConfig {
18923 grid_dim: (n_head as u32, t as u32, 1),
18924 block_dim: (head_dim as u32, 1, 1),
18925 shared_mem_bytes: 0,
18926 };
18927 let __s_b2 = self.gpu.stream();
18928 let mut b2 = __s_b2.launch_builder(&fc);
18929 b2.arg(&*part_o)
18930 .arg(&*part_m)
18931 .arg(&*part_l)
18932 .arg(o)
18933 .arg(&hd)
18934 .arg(&nh)
18935 .arg(base_dev)
18936 .arg(&base_plus)
18937 .arg(&nspm)
18938 .arg(&spk);
18939 unsafe {
18940 b2.launch(cfg2)?;
18941 }
18942 return Ok(());
18943 }
18944 let sp = fa_split_keys(t_kv_upper, n_head_kv);
18945 let n_splits_max = (t_kv_upper + sp - 1) / sp;
18946 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18947 let (nspm, spk) = (n_splits_max as i32, sp as i32);
18948 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18949 let gqa = (n_head / n_head_kv).max(1) as u32;
18950 let o_len = t * n_head * n_splits_max * head_dim;
18951 let ml_len = t * n_head * n_splits_max;
18952 let mut part_guard = self.fa_part_pool.lock().unwrap();
18953 if part_guard
18954 .as_ref()
18955 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18956 .unwrap_or(true)
18957 {
18958 let old = part_guard.take();
18969 let (co, cm) = old
18970 .as_ref()
18971 .map(|pp| (pp.0.len(), pp.1.len()))
18972 .unwrap_or((0, 0));
18973 if let Some(old) = old {
18974 self.fa_part_retired.lock().unwrap().push(old);
18975 }
18976 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18977 eprintln!(
18978 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18979 co, o_len, cm, ml_len
18980 );
18981 }
18982 *part_guard = Some((
18983 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18984 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18985 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18986 ));
18987 }
18988 let pg = part_guard.as_mut().unwrap();
18989 self.gpu
18990 .stream()
18991 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18992 self.gpu
18993 .stream()
18994 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18995 self.gpu
18996 .stream()
18997 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18998 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18999 let f = self.func("fa_decode_vec_q_rows_v3_dc");
19000 let sh = (32 * head_dim * 2) as u32;
19001 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19002 f.set_attribute(
19003 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19004 sh as i32,
19005 )?;
19006 let cfg = LaunchConfig {
19007 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19008 block_dim: (32, gqa, 1),
19009 shared_mem_bytes: sh,
19010 };
19011 let __s_b = self.gpu.stream();
19012 let mut b = __s_b.launch_builder(&f);
19013 b.arg(q)
19014 .arg(k)
19015 .arg(v)
19016 .arg(&mut *part_o)
19017 .arg(&mut *part_m)
19018 .arg(&mut *part_l)
19019 .arg(&hd)
19020 .arg(&nh)
19021 .arg(&nhkv)
19022 .arg(base_dev)
19023 .arg(&scale)
19024 .arg(&nspm)
19025 .arg(&spk)
19026 .arg(&ktb)
19027 .arg(&vtb);
19028 unsafe {
19029 b.launch(cfg)?;
19030 }
19031 let fc = self.func("fa_decode_combine_rows_dc");
19032 let cfg2 = LaunchConfig {
19033 grid_dim: (n_head as u32, t as u32, 1),
19034 block_dim: (head_dim as u32, 1, 1),
19035 shared_mem_bytes: 0,
19036 };
19037 let plus0 = 0i32;
19038 let __s_b2 = self.gpu.stream();
19039 let mut b2 = __s_b2.launch_builder(&fc);
19040 b2.arg(&*part_o)
19041 .arg(&*part_m)
19042 .arg(&*part_l)
19043 .arg(o)
19044 .arg(&hd)
19045 .arg(&nh)
19046 .arg(base_dev)
19047 .arg(&plus0)
19048 .arg(&nspm)
19049 .arg(&spk);
19050 unsafe {
19051 b2.launch(cfg2)?;
19052 }
19053 Ok(())
19054 }
19055
19056 pub fn fa_decode_dc(
19067 &self,
19068 q: &CudaSlice<f32>,
19069 k: &cudarc::driver::CudaView<u8>,
19070 v: &cudarc::driver::CudaView<u8>,
19071 o: &mut CudaSlice<f32>,
19072 head_dim: usize,
19073 n_head: usize,
19074 n_head_kv: usize,
19075 t_kv_dev: &CudaSlice<i32>,
19076 bucket_max: usize,
19077 scale: f32,
19078 k_tok_bytes: usize,
19079 v_tok_bytes: usize,
19080 g: bool,
19081 ) -> Result<(), Box<dyn std::error::Error>> {
19082 self.fa_decode_dc_q8(
19083 q,
19084 k,
19085 v,
19086 o,
19087 head_dim,
19088 n_head,
19089 n_head_kv,
19090 t_kv_dev,
19091 bucket_max,
19092 scale,
19093 k_tok_bytes,
19094 v_tok_bytes,
19095 g,
19096 None,
19097 )
19098 }
19099
19100 #[allow(clippy::too_many_arguments)]
19103 pub fn fa_decode_dc_q8(
19104 &self,
19105 q: &CudaSlice<f32>,
19106 k: &cudarc::driver::CudaView<u8>,
19107 v: &cudarc::driver::CudaView<u8>,
19108 o: &mut CudaSlice<f32>,
19109 head_dim: usize,
19110 n_head: usize,
19111 n_head_kv: usize,
19112 t_kv_dev: &CudaSlice<i32>,
19113 bucket_max: usize,
19114 scale: f32,
19115 k_tok_bytes: usize,
19116 v_tok_bytes: usize,
19117 g: bool,
19118 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19119 ) -> Result<(), Box<dyn std::error::Error>> {
19120 let mut fa_vec =
19128 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19129 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19130 fa_vec = false;
19131 } let sp = fa_split_keys(bucket_max, n_head_kv);
19133 let n_splits = if fa_vec {
19134 ((bucket_max + sp - 1) / sp).max(1)
19135 } else {
19136 ((bucket_max + 255) / 256).max(1)
19137 };
19138 let o_len = n_head * n_splits * head_dim;
19139 let ml_len = n_head * n_splits;
19140 let mut part_guard = self.fa_part_pool.lock().unwrap();
19141 if part_guard
19142 .as_ref()
19143 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19144 .unwrap_or(true)
19145 {
19146 let old = part_guard.take();
19157 let (co, cm) = old
19158 .as_ref()
19159 .map(|pp| (pp.0.len(), pp.1.len()))
19160 .unwrap_or((0, 0));
19161 if let Some(old) = old {
19162 self.fa_part_retired.lock().unwrap().push(old);
19163 }
19164 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19165 eprintln!(
19166 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19167 co, o_len, cm, ml_len
19168 );
19169 }
19170 *part_guard = Some((
19171 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19172 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19173 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19174 ));
19175 }
19176 let pg = part_guard.as_mut().unwrap();
19177 self.gpu
19178 .stream()
19179 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19180 self.gpu
19181 .stream()
19182 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19183 self.gpu
19184 .stream()
19185 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19186 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19187 let (hd, nh, nhkv, nsp) = (
19188 head_dim as i32,
19189 n_head as i32,
19190 n_head_kv as i32,
19191 n_splits as i32,
19192 );
19193 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19194 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
19195 let deep = fa_vec
19198 && head_dim == 256
19199 && fa_v4_at(bucket_max)
19200 && !g
19201 && fa_deep_at(bucket_max)
19202 && !matches!(fa_v4_mode(), "noB3" | "stage");
19203 let (f, cfg) = if fa_vec
19204 && head_dim == 512
19205 && bucket_max >= {
19206 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19207 *FA512_MIN_DC.get_or_init(|| {
19208 std::env::var("MEMRA_FA512_MIN")
19209 .ok()
19210 .and_then(|v| v.parse().ok())
19211 .unwrap_or(512)
19212 })
19213 } {
19214 let gqa = (n_head / n_head_kv).max(1) as u32;
19216 (
19217 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
19218 LaunchConfig {
19219 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19220 block_dim: (32, gqa, 1),
19221 shared_mem_bytes: 0,
19222 },
19223 )
19224 } else if fa_vec && head_dim == 512 {
19225 let q_view = q.as_view();
19228 let mut o_view = o.as_view_mut();
19229 return self.fa_decode_scalar_unified(
19230 &q_view,
19231 k,
19232 v,
19233 &mut o_view,
19234 head_dim,
19235 n_head,
19236 n_head_kv,
19237 0,
19238 Some(t_kv_dev),
19239 scale,
19240 n_splits,
19241 sp,
19242 k_tok_bytes,
19243 v_tok_bytes,
19244 g,
19245 &mut *part_o,
19246 &mut *part_m,
19247 &mut *part_l,
19248 q8_out,
19249 );
19250 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
19251 let gqa = (n_head / n_head_kv).max(1) as u32;
19254 let fv = if g {
19255 self.func_g("fa_decode_vec_q_v4_dc")
19256 } else if deep {
19257 self.func("fa_decode_vec_q_v4_deep_dc")
19258 } else {
19259 self.func("fa_decode_vec_q_v4_dc")
19260 };
19261 let shmem =
19262 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19263 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19264 fv.set_attribute(
19265 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19266 shmem as i32,
19267 )?;
19268 (
19269 fv,
19270 LaunchConfig {
19271 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19272 block_dim: (32, gqa, 1),
19273 shared_mem_bytes: shmem,
19274 },
19275 )
19276 } else if fa_vec && fa_v3_active(head_dim) {
19277 let gqa = (n_head / n_head_kv).max(1) as u32;
19280 let fv = if g {
19281 self.func_g("fa_decode_vec_q_v3_dc")
19282 } else {
19283 self.func("fa_decode_vec_q_v3_dc")
19284 };
19285 let shmem = (32 * head_dim * 2) as u32; (
19287 fv,
19288 LaunchConfig {
19289 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19290 block_dim: (32, gqa, 1),
19291 shared_mem_bytes: shmem,
19292 },
19293 )
19294 } else if fa_vec && fa_v2_on() {
19295 let gqa = (n_head / n_head_kv).max(1) as u32;
19299 let fv = if g {
19300 self.func_g("fa_decode_vec_q_v2_dc")
19301 } else {
19302 self.func("fa_decode_vec_q_v2_dc")
19303 };
19304 let shmem = (2 * 32 * head_dim * 2) as u32; (
19306 fv,
19307 LaunchConfig {
19308 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19309 block_dim: (32, gqa, 1),
19310 shared_mem_bytes: shmem,
19311 },
19312 )
19313 } else if fa_vec {
19314 let gqa = (n_head / n_head_kv).max(1) as u32;
19315 let fv = if g {
19317 self.func_g("fa_decode_vec_q_dc")
19318 } else {
19319 self.func("fa_decode_vec_q_dc")
19320 };
19321 (
19322 fv,
19323 LaunchConfig {
19324 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19325 block_dim: (32, gqa, 1),
19326 shared_mem_bytes: 0,
19327 },
19328 )
19329 } else {
19330 let q_view = q.as_view();
19331 let mut o_view = o.as_view_mut();
19332 return self.fa_decode_scalar_unified(
19333 &q_view,
19334 k,
19335 v,
19336 &mut o_view,
19337 head_dim,
19338 n_head,
19339 n_head_kv,
19340 0,
19341 Some(t_kv_dev),
19342 scale,
19343 n_splits,
19344 if fa_vec { sp } else { 256 },
19345 k_tok_bytes,
19346 v_tok_bytes,
19347 g,
19348 &mut *part_o,
19349 &mut *part_m,
19350 &mut *part_l,
19351 q8_out,
19352 );
19353 };
19354 let ski = sp as i32; let __s_b = self.gpu.stream();
19356 let mut b = __s_b.launch_builder(&f);
19357 b.arg(q)
19358 .arg(k)
19359 .arg(v)
19360 .arg(&mut *part_o)
19361 .arg(&mut *part_m)
19362 .arg(&mut *part_l)
19363 .arg(&hd)
19364 .arg(&nh)
19365 .arg(&nhkv)
19366 .arg(t_kv_dev)
19367 .arg(&scale)
19368 .arg(&nsp)
19369 .arg(&ski)
19370 .arg(&ktb)
19371 .arg(&vtb);
19372 unsafe {
19373 b.launch(cfg)?;
19374 }
19375 let cfg2 = LaunchConfig {
19376 grid_dim: (n_head as u32, 1, 1),
19377 block_dim: (head_dim as u32, 1, 1),
19378 shared_mem_bytes: 0,
19379 };
19380 if let Some((oq, od)) = q8_out {
19381 let fc = if g {
19382 self.func_g("fa_decode_combine_q8_1")
19383 } else {
19384 self.fa_func("fa_decode_combine_q8_1", head_dim)
19385 };
19386 let __s_b2 = self.gpu.stream();
19387 let mut b2 = __s_b2.launch_builder(&fc);
19388 b2.arg(&*part_o)
19389 .arg(&*part_m)
19390 .arg(&*part_l)
19391 .arg(oq)
19392 .arg(od)
19393 .arg(&hd)
19394 .arg(&nh)
19395 .arg(&nsp);
19396 unsafe {
19397 b2.launch(cfg2)?;
19398 }
19399 return Ok(());
19400 }
19401 let fc = if g {
19402 self.func_g("fa_decode_combine_f32")
19403 } else {
19404 self.fa_func("fa_decode_combine_f32", head_dim)
19405 };
19406 let __s_b2 = self.gpu.stream();
19407 let mut b2 = __s_b2.launch_builder(&fc);
19408 b2.arg(&*part_o)
19409 .arg(&*part_m)
19410 .arg(&*part_l)
19411 .arg(o)
19412 .arg(&hd)
19413 .arg(&nh)
19414 .arg(&nsp);
19415 unsafe {
19416 b2.launch(cfg2)?;
19417 }
19418 Ok(())
19419 }
19420
19421 pub fn fa_geom_eager(
19427 &self,
19428 t_kv: usize,
19429 head_dim: usize,
19430 n_head_kv: usize,
19431 g: bool,
19432 ) -> (bool, usize) {
19433 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
19437 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
19443 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
19444 if g && head_dim == 256 && !fa_v4_at(t_kv) {
19450 fa_vec = false;
19451 }
19452 let sp = fa_split_keys(t_kv, n_head_kv);
19453 let n_splits = if fa_vec {
19454 ((t_kv + sp - 1) / sp).max(1)
19455 } else {
19456 ((t_kv + 255) / 256).max(1)
19457 };
19458 (fa_vec, n_splits)
19459 }
19460
19461 pub fn fa_bucket_key(
19467 &self,
19468 t_kv: usize,
19469 head_dim: usize,
19470 n_head_kv: usize,
19471 g: bool,
19472 ) -> (bool, usize) {
19473 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
19474 }
19475
19476 pub fn capture_graph_retained<F>(
19488 &self,
19489 step: F,
19490 ) -> Result<
19491 (
19492 cudarc::driver::CudaGraph,
19493 Vec<Box<dyn std::any::Any + Send>>,
19494 ),
19495 Box<dyn std::error::Error>,
19496 >
19497 where
19498 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19499 {
19500 use cudarc::driver::sys::CUgraphInstantiate_flags;
19501 self.capture_graph_retained_flags(
19502 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19503 step,
19504 )
19505 }
19506
19507 pub fn capture_graph_retained_flags<F>(
19512 &self,
19513 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
19514 mut step: F,
19515 ) -> Result<
19516 (
19517 cudarc::driver::CudaGraph,
19518 Vec<Box<dyn std::any::Any + Send>>,
19519 ),
19520 Box<dyn std::error::Error>,
19521 >
19522 where
19523 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19524 {
19525 use cudarc::driver::sys::CUstreamCaptureMode;
19526 self.capture_keep.lock().unwrap().clear();
19534 let was_tracking = self.gpu.ctx.is_event_tracking();
19535 if was_tracking {
19536 unsafe {
19537 self.gpu.ctx.disable_event_tracking();
19538 }
19539 }
19540 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19541 self.capture_keep_on
19542 .store(true, std::sync::atomic::Ordering::Relaxed);
19543 let w = (|| {
19544 step(self)?;
19545 step(self)
19546 })();
19547 self.capture_keep_on
19548 .store(false, std::sync::atomic::Ordering::Relaxed);
19549 w?;
19550 self.gpu.stream().synchronize()?;
19551 self.gpu
19552 .stream()
19553 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19554 let r = step(self);
19555 let g = self.gpu.stream().end_capture(flags);
19556 r?;
19557 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19558 graph.upload()?;
19559 Ok(graph)
19560 };
19561 let result = run();
19562 self.capture_keep_on
19563 .store(false, std::sync::atomic::Ordering::Relaxed);
19564 if was_tracking {
19565 unsafe {
19566 self.gpu.ctx.enable_event_tracking();
19567 }
19568 }
19569 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
19570 Ok((result?, keeper))
19571 }
19572
19573 pub fn capture_graph<F>(
19574 &self,
19575 mut step: F,
19576 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
19577 where
19578 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19579 {
19580 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
19581 let was_tracking = self.gpu.ctx.is_event_tracking();
19589 if was_tracking {
19590 unsafe {
19591 self.gpu.ctx.disable_event_tracking();
19592 }
19593 }
19594 let iflag = {
19601 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
19602 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
19603 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
19606 Ok("priority") => {
19607 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
19608 }
19609 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19610 })
19611 };
19612 let ct = {
19619 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19620 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
19621 };
19622 let warmups = {
19645 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19646 *W.get_or_init(|| {
19647 std::env::var("MEMRA_GRAPH_WARMUPS")
19648 .ok()
19649 .and_then(|v| v.parse().ok())
19650 .filter(|n| *n >= 1)
19651 .unwrap_or(1)
19652 })
19653 };
19654 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19655 let t_w = std::time::Instant::now();
19656 for _ in 0..warmups {
19658 step(self)?;
19659 }
19660 self.gpu.stream().synchronize()?;
19661 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
19662 let t_c = std::time::Instant::now();
19664 self.gpu
19665 .stream()
19666 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19667 let r = step(self);
19670 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
19671 let t_i = std::time::Instant::now();
19672 let g = self.gpu.stream().end_capture(iflag);
19673 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
19674 r?;
19675 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19676 let t_u = std::time::Instant::now();
19677 graph.upload()?;
19678 if ct {
19679 println!(
19680 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
19681 instantiate {ms_inst:.2} ms upload {:.2} ms",
19682 t_u.elapsed().as_secs_f64() * 1e3
19683 );
19684 }
19685 Ok(graph)
19686 };
19687 let result = run();
19688 if was_tracking {
19689 unsafe {
19690 self.gpu.ctx.enable_event_tracking();
19691 }
19692 }
19693 result
19694 }
19695
19696 pub fn gdn_scan_s128_view(
19698 &self,
19699 q: &CudaSlice<f32>,
19700 k: &CudaSlice<f32>,
19701 v: &CudaSlice<f32>,
19702 g: &CudaSlice<f32>,
19703 beta: &CudaSlice<f32>,
19704 state_in: &cudarc::driver::CudaView<f32>,
19705 state_out: &mut cudarc::driver::CudaViewMut<f32>,
19706 o: &mut CudaSlice<f32>,
19707 n_head: usize,
19708 t: usize,
19709 scale: f32,
19710 ) -> Result<(), Box<dyn std::error::Error>> {
19711 let f = self.func("gdn_scan_s128");
19712 const S_V: u32 = 128;
19713 const WARP: u32 = 32;
19714 const COLS: u32 = 4;
19715 let cfg = LaunchConfig {
19716 grid_dim: (n_head as u32, 1, S_V / COLS),
19717 block_dim: (WARP, COLS, 1),
19718 shared_mem_bytes: 0,
19719 };
19720 let (h, ti) = (n_head as i32, t as i32);
19721 let __s_b = self.gpu.stream();
19722 let mut b = __s_b.launch_builder(&f);
19723 b.arg(q)
19724 .arg(k)
19725 .arg(v)
19726 .arg(g)
19727 .arg(beta)
19728 .arg(state_in)
19729 .arg(state_out)
19730 .arg(o)
19731 .arg(&h)
19732 .arg(&ti)
19733 .arg(&scale);
19734 unsafe {
19735 b.launch(cfg)?;
19736 }
19737 Ok(())
19738 }
19739
19740 pub fn ssm_conv1d_view(
19742 &self,
19743 x: &cudarc::driver::CudaView<f32>,
19744 w: &CudaSlice<f32>,
19745 y: &mut CudaSlice<f32>,
19746 conv_dim: usize,
19747 t: usize,
19748 d_conv: usize,
19749 silu: bool,
19750 ) -> Result<(), Box<dyn std::error::Error>> {
19751 let f = self.func("ssm_conv1d_silu_f32");
19752 let cfg = LaunchConfig {
19754 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
19755 block_dim: (256, 1, 1),
19756 shared_mem_bytes: 0,
19757 };
19758 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
19759 let __s_b = self.gpu.stream();
19760 let mut b = __s_b.launch_builder(&f);
19761 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
19762 unsafe {
19763 b.launch(cfg)?;
19764 }
19765 Ok(())
19766 }
19767
19768 pub fn ssm_conv1d_tm(
19775 &self,
19776 qkv_tm: &CudaSlice<f32>,
19777 w: &CudaSlice<f32>,
19778 y: &mut CudaSlice<f32>,
19779 conv_dim: usize,
19780 t: usize,
19781 d_conv: usize,
19782 ) -> Result<(), Box<dyn std::error::Error>> {
19783 let f = self.func("ssm_conv1d_tm_f32");
19784 let cfg = LaunchConfig {
19785 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
19786 block_dim: (256, 1, 1),
19787 shared_mem_bytes: 0,
19788 };
19789 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19790 let __s_b = self.gpu.stream();
19791 let mut b = __s_b.launch_builder(&f);
19792 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
19793 unsafe {
19794 b.launch(cfg)?;
19795 }
19796 Ok(())
19797 }
19798
19799 pub fn ssm_conv1d_tm_state(
19807 &self,
19808 qkv_tm: &CudaSlice<f32>,
19809 conv_state: &mut CudaSlice<f32>,
19810 w: &CudaSlice<f32>,
19811 y: &mut CudaSlice<f32>,
19812 conv_dim: usize,
19813 t: usize,
19814 d_conv: usize,
19815 ) -> Result<(), Box<dyn std::error::Error>> {
19816 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
19817 }
19818
19819 #[allow(clippy::too_many_arguments)]
19822 pub fn ssm_conv1d_tm_state_pad(
19823 &self,
19824 qkv_tm: &CudaSlice<f32>,
19825 conv_state: &mut CudaSlice<f32>,
19826 w: &CudaSlice<f32>,
19827 y: &mut CudaSlice<f32>,
19828 conv_dim: usize,
19829 t: usize,
19830 d_conv: usize,
19831 pad_len: Option<&CudaSlice<i32>>,
19832 ) -> Result<(), Box<dyn std::error::Error>> {
19833 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
19834 let ring_old = if t < d_conv - 1 {
19838 Some(self.clone_dtod(conv_state)?)
19839 } else {
19840 None
19841 };
19842 {
19843 let f = self.func("ssm_conv1d_tm_state_f32");
19844 let cfg = LaunchConfig {
19845 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
19846 block_dim: (256, 1, 1),
19847 shared_mem_bytes: 0,
19848 };
19849 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19850 let __s_b = self.gpu.stream();
19851 let mut b = __s_b.launch_builder(&f);
19852 b.arg(qkv_tm)
19853 .arg(&*conv_state)
19854 .arg(w)
19855 .arg(y)
19856 .arg(&cd)
19857 .arg(&ti)
19858 .arg(&dc);
19859 unsafe {
19860 b.launch(cfg)?;
19861 }
19862 }
19863 match (ring_old, pad_len) {
19864 (None, Some(len_d)) => {
19865 let f = self.func("ssm_conv_ring_update_dev_f32");
19866 let n = conv_dim * (d_conv - 1);
19867 let cfg = LaunchConfig::for_num_elems(n as u32);
19868 let (cd, dc) = (conv_dim as i32, d_conv as i32);
19869 let __s_b = self.gpu.stream();
19870 let mut b = __s_b.launch_builder(&f);
19871 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
19872 unsafe {
19873 b.launch(cfg)?;
19874 }
19875 }
19876 (None, None) => {
19877 let f = self.func("ssm_conv_ring_update_f32");
19878 let n = conv_dim * (d_conv - 1);
19879 let cfg = LaunchConfig::for_num_elems(n as u32);
19880 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19881 let __s_b = self.gpu.stream();
19882 let mut b = __s_b.launch_builder(&f);
19883 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
19884 unsafe {
19885 b.launch(cfg)?;
19886 }
19887 }
19888 (Some(old), _) => {
19889 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
19890 }
19891 }
19892 Ok(())
19893 }
19894
19895 pub fn ssm_conv1d_tm_state_pad_v(
19897 &self,
19898 qkv_tm: &cudarc::driver::CudaView<f32>,
19899 conv_state: &mut CudaSlice<f32>,
19900 w: &CudaSlice<f32>,
19901 y: &mut CudaSlice<f32>,
19902 conv_dim: usize,
19903 t: usize,
19904 d_conv: usize,
19905 pad_len: Option<&CudaSlice<i32>>,
19906 ) -> Result<(), Box<dyn std::error::Error>> {
19907 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
19908 let ring_old = if t < d_conv - 1 {
19912 Some(self.clone_dtod(conv_state)?)
19913 } else {
19914 None
19915 };
19916 {
19917 let f = self.func("ssm_conv1d_tm_state_f32");
19918 let cfg = LaunchConfig {
19919 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
19920 block_dim: (256, 1, 1),
19921 shared_mem_bytes: 0,
19922 };
19923 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19924 let __s_b = self.gpu.stream();
19925 let mut b = __s_b.launch_builder(&f);
19926 b.arg(qkv_tm)
19927 .arg(&*conv_state)
19928 .arg(w)
19929 .arg(y)
19930 .arg(&cd)
19931 .arg(&ti)
19932 .arg(&dc);
19933 unsafe {
19934 b.launch(cfg)?;
19935 }
19936 }
19937 match (ring_old, pad_len) {
19938 (None, Some(len_d)) => {
19939 let f = self.func("ssm_conv_ring_update_dev_f32");
19940 let n = conv_dim * (d_conv - 1);
19941 let cfg = LaunchConfig::for_num_elems(n as u32);
19942 let (cd, dc) = (conv_dim as i32, d_conv as i32);
19943 let __s_b = self.gpu.stream();
19944 let mut b = __s_b.launch_builder(&f);
19945 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
19946 unsafe {
19947 b.launch(cfg)?;
19948 }
19949 }
19950 (None, None) => {
19951 let f = self.func("ssm_conv_ring_update_f32");
19952 let n = conv_dim * (d_conv - 1);
19953 let cfg = LaunchConfig::for_num_elems(n as u32);
19954 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
19955 let __s_b = self.gpu.stream();
19956 let mut b = __s_b.launch_builder(&f);
19957 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
19958 unsafe {
19959 b.launch(cfg)?;
19960 }
19961 }
19962 (Some(_), _) => unreachable!(
19963 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
19964 ),
19965 }
19966 Ok(())
19967 }
19968
19969 pub fn ssm_conv_ring_rebuild(
19974 &self,
19975 qkv_tm: &CudaSlice<f32>,
19976 ring_old: &CudaSlice<f32>,
19977 conv_state: &mut CudaSlice<f32>,
19978 conv_dim: usize,
19979 tc: usize,
19980 d_conv: usize,
19981 ) -> Result<(), Box<dyn std::error::Error>> {
19982 let f = self.func("ssm_conv_ring_rebuild_f32");
19983 let n = conv_dim * (d_conv - 1);
19984 let cfg = LaunchConfig::for_num_elems(n as u32);
19985 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
19986 let __s_b = self.gpu.stream();
19987 let mut b = __s_b.launch_builder(&f);
19988 b.arg(qkv_tm)
19989 .arg(ring_old)
19990 .arg(conv_state)
19991 .arg(&cd)
19992 .arg(&ti)
19993 .arg(&dc);
19994 unsafe {
19995 b.launch(cfg)?;
19996 }
19997 Ok(())
19998 }
19999
20000 #[allow(clippy::too_many_arguments)]
20005 pub fn gdn_prep_decode(
20006 &self,
20007 conv_out: &CudaSlice<f32>,
20008 beta_raw: &CudaSlice<f32>,
20009 alpha: &CudaSlice<f32>,
20010 dt_bias: &CudaSlice<f32>,
20011 a: &CudaSlice<f32>,
20012 q_l2: &mut CudaSlice<f32>,
20013 k_l2: &mut CudaSlice<f32>,
20014 v_g: &mut CudaSlice<f32>,
20015 beta: &mut CudaSlice<f32>,
20016 g_log: &mut CudaSlice<f32>,
20017 d_state: usize,
20018 num_v: usize,
20019 num_k: usize,
20020 key_dim: usize,
20021 eps: f32,
20022 ) -> Result<(), Box<dyn std::error::Error>> {
20023 let f = self.func("gdn_prep_decode_f32");
20024 let cfg = LaunchConfig {
20025 grid_dim: (num_v as u32, 1, 1),
20026 block_dim: (32, 4, 1),
20027 shared_mem_bytes: 0,
20028 };
20029 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20030 let __s_b = self.gpu.stream();
20031 let mut b = __s_b.launch_builder(&f);
20032 b.arg(conv_out)
20033 .arg(beta_raw)
20034 .arg(alpha)
20035 .arg(dt_bias)
20036 .arg(a)
20037 .arg(q_l2)
20038 .arg(k_l2)
20039 .arg(v_g)
20040 .arg(beta)
20041 .arg(g_log)
20042 .arg(&ds)
20043 .arg(&nv)
20044 .arg(&nk)
20045 .arg(&kd)
20046 .arg(&eps);
20047 unsafe {
20048 b.launch(cfg)?;
20049 }
20050 Ok(())
20051 }
20052
20053 #[allow(clippy::too_many_arguments)]
20057 pub fn ssm_conv1d_gdn(
20058 &self,
20059 qkv_tm: &CudaSlice<f32>,
20060 w: &CudaSlice<f32>,
20061 q_g: &mut CudaSlice<f32>,
20062 k_g: &mut CudaSlice<f32>,
20063 v_g: &mut CudaSlice<f32>,
20064 conv_dim: usize,
20065 t: usize,
20066 d_conv: usize,
20067 d_state: usize,
20068 num_v: usize,
20069 num_k: usize,
20070 key_dim: usize,
20071 ) -> Result<(), Box<dyn std::error::Error>> {
20072 let f = self.func("ssm_conv1d_gdn_f32");
20073 let cfg = LaunchConfig {
20074 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20075 block_dim: (256, 1, 1),
20076 shared_mem_bytes: 0,
20077 };
20078 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20079 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20080 let __s_b = self.gpu.stream();
20081 let mut b = __s_b.launch_builder(&f);
20082 b.arg(qkv_tm)
20083 .arg(w)
20084 .arg(q_g)
20085 .arg(k_g)
20086 .arg(v_g)
20087 .arg(&cd)
20088 .arg(&ti)
20089 .arg(&dc)
20090 .arg(&ds)
20091 .arg(&nv)
20092 .arg(&nk)
20093 .arg(&kd);
20094 unsafe {
20095 b.launch(cfg)?;
20096 }
20097 Ok(())
20098 }
20099
20100 pub fn ssm_conv1d(
20101 &self,
20102 x: &CudaSlice<f32>,
20103 w: &CudaSlice<f32>,
20104 y: &mut CudaSlice<f32>,
20105 conv_dim: usize,
20106 t: usize,
20107 d_conv: usize,
20108 silu: bool,
20109 ) -> Result<(), Box<dyn std::error::Error>> {
20110 let f = self.func("ssm_conv1d_silu_f32");
20111 let cfg = LaunchConfig {
20112 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20113 block_dim: (256, 1, 1),
20114 shared_mem_bytes: 0,
20115 };
20116 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20117 let __s_b = self.gpu.stream();
20118 let mut b = __s_b.launch_builder(&f);
20119 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20120 unsafe {
20121 b.launch(cfg)?;
20122 }
20123 Ok(())
20124 }
20125
20126 pub fn gdn_scan_s128(
20129 &self,
20130 q: &CudaSlice<f32>,
20131 k: &CudaSlice<f32>,
20132 v: &CudaSlice<f32>,
20133 g: &CudaSlice<f32>,
20134 beta: &CudaSlice<f32>,
20135 state_in: &CudaSlice<f32>,
20136 state_out: &mut CudaSlice<f32>,
20137 o: &mut CudaSlice<f32>,
20138 n_head: usize,
20139 t: usize,
20140 scale: f32,
20141 ) -> Result<(), Box<dyn std::error::Error>> {
20142 let f = self.func("gdn_scan_s128");
20143 const S_V: u32 = 128;
20144 const WARP: u32 = 32;
20145 const COLS_PER_BLOCK: u32 = 4;
20146 let cfg = LaunchConfig {
20147 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20148 block_dim: (WARP, COLS_PER_BLOCK, 1),
20149 shared_mem_bytes: 0,
20150 };
20151 let (h, ti) = (n_head as i32, t as i32);
20152 let __s_b = self.gpu.stream();
20153 let mut b = __s_b.launch_builder(&f);
20154 b.arg(q)
20155 .arg(k)
20156 .arg(v)
20157 .arg(g)
20158 .arg(beta)
20159 .arg(state_in)
20160 .arg(state_out)
20161 .arg(o)
20162 .arg(&h)
20163 .arg(&ti)
20164 .arg(&scale);
20165 unsafe {
20166 b.launch(cfg)?;
20167 }
20168 Ok(())
20169 }
20170
20171 #[allow(clippy::too_many_arguments)]
20176 pub fn ssm_conv1d_fused_decode_b(
20177 &self,
20178 qkv_cols: &CudaSlice<f32>,
20179 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20180 w: &CudaSlice<f32>,
20181 conv_outs: &mut CudaSlice<f32>,
20182 conv_dim: usize,
20183 d_conv: usize,
20184 b_n: usize,
20185 ) -> Result<(), Box<dyn std::error::Error>> {
20186 let f = self.func("ssm_conv1d_fused_decode_b_f32");
20187 let cfg = LaunchConfig {
20188 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20189 block_dim: (256, 1, 1),
20190 shared_mem_bytes: 0,
20191 };
20192 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20193 let __s_b = self.gpu.stream();
20194 let mut b = __s_b.launch_builder(&f);
20195 b.arg(qkv_cols)
20196 .arg(conv_state_ptrs)
20197 .arg(w)
20198 .arg(conv_outs)
20199 .arg(&cd)
20200 .arg(&dc);
20201 unsafe {
20202 b.launch(cfg)?;
20203 }
20204 Ok(())
20205 }
20206
20207 #[allow(clippy::too_many_arguments)]
20208 pub fn gdn_prep_decode_b(
20209 &self,
20210 conv_outs: &CudaSlice<f32>,
20211 beta_raws: &CudaSlice<f32>,
20212 alphas: &CudaSlice<f32>,
20213 dt_bias: &CudaSlice<f32>,
20214 a: &CudaSlice<f32>,
20215 q_l2: &mut CudaSlice<f32>,
20216 k_l2: &mut CudaSlice<f32>,
20217 v_g: &mut CudaSlice<f32>,
20218 beta: &mut CudaSlice<f32>,
20219 g_log: &mut CudaSlice<f32>,
20220 d_state: usize,
20221 num_v: usize,
20222 num_k: usize,
20223 key_dim: usize,
20224 eps: f32,
20225 conv_dim: usize,
20226 b_n: usize,
20227 ) -> Result<(), Box<dyn std::error::Error>> {
20228 let f = self.func("gdn_prep_decode_b_f32");
20229 let cfg = LaunchConfig {
20230 grid_dim: (num_v as u32, 1, b_n as u32),
20231 block_dim: (32, 4, 1),
20232 shared_mem_bytes: 0,
20233 };
20234 let (ds, nv, nk, kd, cd) = (
20235 d_state as i32,
20236 num_v as i32,
20237 num_k as i32,
20238 key_dim as i32,
20239 conv_dim as i32,
20240 );
20241 let __s_b = self.gpu.stream();
20242 let mut b = __s_b.launch_builder(&f);
20243 b.arg(conv_outs)
20244 .arg(beta_raws)
20245 .arg(alphas)
20246 .arg(dt_bias)
20247 .arg(a)
20248 .arg(q_l2)
20249 .arg(k_l2)
20250 .arg(v_g)
20251 .arg(beta)
20252 .arg(g_log)
20253 .arg(&ds)
20254 .arg(&nv)
20255 .arg(&nk)
20256 .arg(&kd)
20257 .arg(&eps)
20258 .arg(&cd);
20259 unsafe {
20260 b.launch(cfg)?;
20261 }
20262 Ok(())
20263 }
20264
20265 #[allow(clippy::too_many_arguments)]
20266 pub fn gdn_scan_s128_batched(
20267 &self,
20268 q: &CudaSlice<f32>,
20269 k: &CudaSlice<f32>,
20270 v: &CudaSlice<f32>,
20271 g: &CudaSlice<f32>,
20272 beta: &CudaSlice<f32>,
20273 state_in_ptrs: &cudarc::driver::CudaView<u64>,
20274 state_out_ptrs: &cudarc::driver::CudaView<u64>,
20275 o: &mut CudaSlice<f32>,
20276 n_head: usize,
20277 b_n: usize,
20278 scale: f32,
20279 ) -> Result<(), Box<dyn std::error::Error>> {
20280 let f = self.func("gdn_scan_s128_b");
20281 const S_V: u32 = 128;
20282 const WARP: u32 = 32;
20283 const COLS_PER_BLOCK: u32 = 4;
20284 let cfg = LaunchConfig {
20285 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20286 block_dim: (WARP, COLS_PER_BLOCK, 1),
20287 shared_mem_bytes: 0,
20288 };
20289 let h = n_head as i32;
20290 let __s_b = self.gpu.stream();
20291 let mut b = __s_b.launch_builder(&f);
20292 b.arg(q)
20293 .arg(k)
20294 .arg(v)
20295 .arg(g)
20296 .arg(beta)
20297 .arg(state_in_ptrs)
20298 .arg(state_out_ptrs)
20299 .arg(o)
20300 .arg(&h)
20301 .arg(&scale);
20302 unsafe {
20303 b.launch(cfg)?;
20304 }
20305 Ok(())
20306 }
20307
20308 #[allow(clippy::too_many_arguments)]
20314 pub fn ssm_conv1d_fused_decode_b_view(
20315 &self,
20316 qkv_cols: &cudarc::driver::CudaView<f32>,
20317 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20318 w: &CudaSlice<f32>,
20319 conv_outs: &mut CudaSlice<f32>,
20320 conv_dim: usize,
20321 d_conv: usize,
20322 b_n: usize,
20323 ) -> Result<(), Box<dyn std::error::Error>> {
20324 let f = self.func("ssm_conv1d_fused_decode_b_f32");
20325 let cfg = LaunchConfig {
20326 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20327 block_dim: (256, 1, 1),
20328 shared_mem_bytes: 0,
20329 };
20330 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20331 let __s_b = self.gpu.stream();
20332 let mut b = __s_b.launch_builder(&f);
20333 b.arg(qkv_cols)
20334 .arg(conv_state_ptrs)
20335 .arg(w)
20336 .arg(conv_outs)
20337 .arg(&cd)
20338 .arg(&dc);
20339 unsafe {
20340 b.launch(cfg)?;
20341 }
20342 Ok(())
20343 }
20344
20345 #[allow(clippy::too_many_arguments)]
20346 pub fn gdn_prep_decode_b_view(
20347 &self,
20348 conv_outs: &CudaSlice<f32>,
20349 beta_raws: &cudarc::driver::CudaView<f32>,
20350 alphas: &cudarc::driver::CudaView<f32>,
20351 dt_bias: &CudaSlice<f32>,
20352 a: &CudaSlice<f32>,
20353 q_l2: &mut CudaSlice<f32>,
20354 k_l2: &mut CudaSlice<f32>,
20355 v_g: &mut CudaSlice<f32>,
20356 beta: &mut CudaSlice<f32>,
20357 g_log: &mut CudaSlice<f32>,
20358 d_state: usize,
20359 num_v: usize,
20360 num_k: usize,
20361 key_dim: usize,
20362 eps: f32,
20363 conv_dim: usize,
20364 b_n: usize,
20365 ) -> Result<(), Box<dyn std::error::Error>> {
20366 let f = self.func("gdn_prep_decode_b_f32");
20367 let cfg = LaunchConfig {
20368 grid_dim: (num_v as u32, 1, b_n as u32),
20369 block_dim: (32, 4, 1),
20370 shared_mem_bytes: 0,
20371 };
20372 let (ds, nv, nk, kd, cd) = (
20373 d_state as i32,
20374 num_v as i32,
20375 num_k as i32,
20376 key_dim as i32,
20377 conv_dim as i32,
20378 );
20379 let __s_b = self.gpu.stream();
20380 let mut b = __s_b.launch_builder(&f);
20381 b.arg(conv_outs)
20382 .arg(beta_raws)
20383 .arg(alphas)
20384 .arg(dt_bias)
20385 .arg(a)
20386 .arg(q_l2)
20387 .arg(k_l2)
20388 .arg(v_g)
20389 .arg(beta)
20390 .arg(g_log)
20391 .arg(&ds)
20392 .arg(&nv)
20393 .arg(&nk)
20394 .arg(&kd)
20395 .arg(&eps)
20396 .arg(&cd);
20397 unsafe {
20398 b.launch(cfg)?;
20399 }
20400 Ok(())
20401 }
20402
20403 #[allow(clippy::too_many_arguments)]
20404 pub fn gdn_scan_s128_batched_view(
20405 &self,
20406 q: &CudaSlice<f32>,
20407 k: &CudaSlice<f32>,
20408 v: &CudaSlice<f32>,
20409 g: &CudaSlice<f32>,
20410 beta: &CudaSlice<f32>,
20411 state_in_ptrs: &cudarc::driver::CudaView<u64>,
20412 state_out_ptrs: &cudarc::driver::CudaView<u64>,
20413 o: &mut cudarc::driver::CudaViewMut<f32>,
20414 n_head: usize,
20415 b_n: usize,
20416 scale: f32,
20417 ) -> Result<(), Box<dyn std::error::Error>> {
20418 let f = self.func("gdn_scan_s128_b");
20419 const S_V: u32 = 128;
20420 const WARP: u32 = 32;
20421 const COLS_PER_BLOCK: u32 = 4;
20422 let cfg = LaunchConfig {
20423 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20424 block_dim: (WARP, COLS_PER_BLOCK, 1),
20425 shared_mem_bytes: 0,
20426 };
20427 let h = n_head as i32;
20428 let __s_b = self.gpu.stream();
20429 let mut b = __s_b.launch_builder(&f);
20430 b.arg(q)
20431 .arg(k)
20432 .arg(v)
20433 .arg(g)
20434 .arg(beta)
20435 .arg(state_in_ptrs)
20436 .arg(state_out_ptrs)
20437 .arg(o)
20438 .arg(&h)
20439 .arg(&scale);
20440 unsafe {
20441 b.launch(cfg)?;
20442 }
20443 Ok(())
20444 }
20445
20446 pub fn gdn_chunked_enabled() -> bool {
20455 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20456 *E.get_or_init(|| {
20457 std::env::var("MEMRA_GDN_CHUNKED")
20458 .map(|v| v != "0")
20459 .unwrap_or(true)
20460 })
20461 }
20462
20463 pub fn gdn_chunk_size() -> usize {
20468 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20469 *C.get_or_init(|| {
20470 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
20471 .ok()
20472 .and_then(|v| v.parse().ok())
20473 .unwrap_or(32);
20474 c.clamp(32, 128) / 32 * 32
20475 })
20476 }
20477
20478 #[allow(clippy::too_many_arguments)]
20483 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
20486 #[allow(clippy::too_many_arguments)]
20487 pub fn gdn_chunk_k123(
20488 &self,
20489 q: &CudaSlice<f32>,
20490 k: &CudaSlice<f32>,
20491 v: &CudaSlice<f32>,
20492 g: &CudaSlice<f32>,
20493 beta: &CudaSlice<f32>,
20494 wb16: Option<&mut CudaSlice<u8>>,
20495 n_head: usize,
20496 t: usize,
20497 c: usize,
20498 hk: usize,
20499 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
20500 ) -> Result<
20501 (
20502 CudaSlice<f32>,
20503 CudaSlice<f32>,
20504 CudaSlice<f32>,
20505 CudaSlice<f32>,
20506 ),
20507 Box<dyn std::error::Error>,
20508 > {
20509 const D: usize = 128;
20510 let h = n_head;
20511 let nc = (t + c - 1) / c;
20512 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
20513 let mut gcum = self.uninit(t * h)?;
20514 let mut a = self.uninit(nc * h * c * c)?;
20515 let mut p = self.uninit(nc * h * c * c)?;
20516 let mut u = self.uninit(nc * h * c * D)?;
20517 let mut w = self.uninit(nc * h * c * D)?;
20518 {
20519 let f = self.func("gdn_chunk_cumgate_f32");
20521 let cfg = LaunchConfig {
20522 grid_dim: (nc as u32, h as u32, 1),
20523 block_dim: (32, 1, 1),
20524 shared_mem_bytes: 0,
20525 };
20526 let __s_b = self.gpu.stream();
20527 let mut b = __s_b.launch_builder(&f);
20528 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
20529 unsafe {
20530 b.launch(cfg)?;
20531 }
20532 }
20533 if let Some((qb, kb, pb)) = k2w {
20534 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
20537 let f = self.func("gdn_k2_wgmma");
20538 let cfg = LaunchConfig {
20539 grid_dim: (nc as u32, h as u32, 1),
20540 block_dim: (128, 1, 1),
20541 shared_mem_bytes: 0,
20542 };
20543 let hki = hk as i32;
20544 let __s_b = self.gpu.stream();
20545 let mut b = __s_b.launch_builder(&f);
20546 b.arg(qb)
20547 .arg(kb)
20548 .arg(&gcum)
20549 .arg(beta)
20550 .arg(&mut a)
20551 .arg(&mut *pb)
20552 .arg(&hi)
20553 .arg(&ti)
20554 .arg(&ci)
20555 .arg(&hki);
20556 unsafe {
20557 b.launch(cfg)?;
20558 }
20559 } else if c <= 64 && !portable_mma_gated() {
20560 let f = self.func("gdn_chunk_attn_f32");
20562 let jt = ((c + 31) / 32) as u32;
20563 let cfg = LaunchConfig {
20564 grid_dim: (nc as u32, h as u32, jt),
20565 block_dim: (256, 1, 1),
20566 shared_mem_bytes: 0,
20567 };
20568 let hki = hk as i32;
20569 let __s_b = self.gpu.stream();
20570 let mut b = __s_b.launch_builder(&f);
20571 b.arg(q)
20572 .arg(k)
20573 .arg(&gcum)
20574 .arg(beta)
20575 .arg(&mut a)
20576 .arg(&mut p)
20577 .arg(&hi)
20578 .arg(&ti)
20579 .arg(&ci)
20580 .arg(&hki);
20581 unsafe {
20582 b.launch(cfg)?;
20583 }
20584 } else {
20585 assert!(
20587 hk == h,
20588 "generic K2 is broadcast-only (de-broadcast rides C==32)"
20589 );
20590 let f = self.func("gdn_chunk_attn_g_f32");
20591 let cfg = LaunchConfig {
20592 grid_dim: (nc as u32, h as u32, 1),
20593 block_dim: (32, 8, 1),
20594 shared_mem_bytes: 0,
20595 };
20596 let __s_b = self.gpu.stream();
20597 let mut b = __s_b.launch_builder(&f);
20598 b.arg(q)
20599 .arg(k)
20600 .arg(&gcum)
20601 .arg(beta)
20602 .arg(&mut a)
20603 .arg(&mut p)
20604 .arg(&hi)
20605 .arg(&ti)
20606 .arg(&ci);
20607 unsafe {
20608 b.launch(cfg)?;
20609 }
20610 }
20611 {
20612 let cfg = LaunchConfig {
20614 grid_dim: (nc as u32, h as u32, 1),
20615 block_dim: (256, 1, 1),
20616 shared_mem_bytes: 0,
20617 };
20618 match c {
20619 32 | 64 => {
20620 let f = self.func(if c == 32 {
20621 "gdn_chunk_solve32_f32"
20622 } else {
20623 "gdn_chunk_solve64_f32"
20624 });
20625 let wb: u64 = match wb16 {
20627 Some(d) => self.addr_u8(d),
20628 None => 0,
20629 };
20630 let hki = hk as i32;
20631 let __s_b = self.gpu.stream();
20632 let mut b = __s_b.launch_builder(&f);
20633 b.arg(v)
20634 .arg(k)
20635 .arg(&a)
20636 .arg(&gcum)
20637 .arg(&mut u)
20638 .arg(&mut w)
20639 .arg(&wb)
20640 .arg(&hi)
20641 .arg(&ti)
20642 .arg(&hki);
20643 unsafe {
20644 b.launch(cfg)?;
20645 }
20646 }
20647 _ => {
20648 assert!(hk == h, "generic K3 is broadcast-only");
20649 let f = self.func("gdn_chunk_solve_f32");
20650 let __s_b = self.gpu.stream();
20651 let mut b = __s_b.launch_builder(&f);
20652 b.arg(v)
20653 .arg(k)
20654 .arg(&a)
20655 .arg(&gcum)
20656 .arg(&mut u)
20657 .arg(&mut w)
20658 .arg(&hi)
20659 .arg(&ti)
20660 .arg(&ci);
20661 unsafe {
20662 b.launch(cfg)?;
20663 }
20664 }
20665 }
20666 }
20667 Ok((gcum, p, u, w))
20668 }
20669
20670 pub fn gdn_db_on() -> bool {
20674 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
20675 }
20676
20677 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
20680 !portable_mma_gated()
20681 && c == 32
20682 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
20683 Ok("1") => true,
20684 Ok("0") => false,
20685 _ => cfg!(memra_hopper_mma),
20686 }
20687 }
20688
20689 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
20692 self.gdn_mma_enabled(c)
20693 && match std::env::var("MEMRA_GDN_WGMMA").as_deref() {
20694 Ok("0") => false,
20695 Ok("1") => true,
20696 _ => cfg!(memra_hopper_mma),
20697 }
20698 }
20699
20700 #[allow(clippy::too_many_arguments)]
20705 pub fn ssm_conv1d_gdn_state_pad(
20706 &self,
20707 qkv_tm: &cudarc::driver::CudaView<f32>,
20708 conv_state: &mut CudaSlice<f32>,
20709 w: &CudaSlice<f32>,
20710 q_g: &mut CudaSlice<f32>,
20711 k_g: &mut CudaSlice<f32>,
20712 v_g: &mut CudaSlice<f32>,
20713 conv_dim: usize,
20714 t: usize,
20715 d_conv: usize,
20716 d_state: usize,
20717 num_v: usize,
20718 num_k: usize,
20719 key_dim: usize,
20720 hk: usize,
20721 pad_len: Option<&CudaSlice<i32>>,
20722 ) -> Result<(), Box<dyn std::error::Error>> {
20723 assert!(
20724 t >= d_conv - 1,
20725 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
20726 );
20727 {
20728 let f = self.func("ssm_conv1d_gdn_state_f32");
20729 let cfg = LaunchConfig {
20730 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20731 block_dim: (256, 1, 1),
20732 shared_mem_bytes: 0,
20733 };
20734 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20735 let (ds, nv, nk, kd, hki) = (
20736 d_state as i32,
20737 num_v as i32,
20738 num_k as i32,
20739 key_dim as i32,
20740 hk as i32,
20741 );
20742 let __s_b = self.gpu.stream();
20743 let mut b = __s_b.launch_builder(&f);
20744 b.arg(qkv_tm)
20745 .arg(&*conv_state)
20746 .arg(w)
20747 .arg(q_g)
20748 .arg(k_g)
20749 .arg(v_g)
20750 .arg(&cd)
20751 .arg(&ti)
20752 .arg(&dc)
20753 .arg(&ds)
20754 .arg(&nv)
20755 .arg(&nk)
20756 .arg(&kd)
20757 .arg(&hki);
20758 unsafe {
20759 b.launch(cfg)?;
20760 }
20761 }
20762 match pad_len {
20763 Some(len_d) => {
20764 let f = self.func("ssm_conv_ring_update_dev_f32");
20765 let n = conv_dim * (d_conv - 1);
20766 let cfg = LaunchConfig::for_num_elems(n as u32);
20767 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20768 let __s_b = self.gpu.stream();
20769 let mut b = __s_b.launch_builder(&f);
20770 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20771 unsafe {
20772 b.launch(cfg)?;
20773 }
20774 }
20775 None => {
20776 let f = self.func("ssm_conv_ring_update_f32");
20777 let n = conv_dim * (d_conv - 1);
20778 let cfg = LaunchConfig::for_num_elems(n as u32);
20779 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20780 let __s_b = self.gpu.stream();
20781 let mut b = __s_b.launch_builder(&f);
20782 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20783 unsafe {
20784 b.launch(cfg)?;
20785 }
20786 }
20787 }
20788 Ok(())
20789 }
20790
20791 pub fn gdn_chunk_alloc(
20795 &self,
20796 n_head: usize,
20797 t: usize,
20798 c: usize,
20799 hk: usize,
20800 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
20801 const D: usize = 128;
20802 assert!(
20803 c == 32,
20804 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
20805 );
20806 let h = n_head;
20807 let nc = (t + c - 1) / c;
20808 Ok(GdnChunkBufs {
20809 gcum: self.uninit(t * h)?,
20810 a: self.uninit(nc * h * c * c)?,
20811 p: self.uninit(nc * h * c * c)?,
20812 u: self.uninit(nc * h * c * D)?,
20813 w: self.uninit(nc * h * c * D)?,
20814 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
20815 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
20816 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
20817 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
20818 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
20819 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
20820 o: self.uninit(D * h * t)?,
20821 t,
20822 nc,
20823 })
20824 }
20825
20826 pub fn f32_to_bf16_v(
20828 &self,
20829 x: &cudarc::driver::CudaView<f32>,
20830 dst: &mut CudaSlice<u8>,
20831 n: usize,
20832 ) -> Result<(), Box<dyn std::error::Error>> {
20833 let f = self.func("f32_to_bf16_bulk");
20834 let ni = n as i64;
20835 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20836 let __s_b = self.gpu.stream();
20837 let mut b = __s_b.launch_builder(&f);
20838 b.arg(x).arg(dst).arg(&ni);
20839 unsafe {
20840 b.launch(cfg)?;
20841 }
20842 Ok(())
20843 }
20844
20845 pub fn f32_to_bf16_into(
20847 &self,
20848 x: &CudaSlice<f32>,
20849 dst: &mut CudaSlice<u8>,
20850 n: usize,
20851 ) -> Result<(), Box<dyn std::error::Error>> {
20852 let f = self.func("f32_to_bf16_bulk");
20853 let ni = n as i64;
20854 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20855 let __s_b = self.gpu.stream();
20856 let mut b = __s_b.launch_builder(&f);
20857 b.arg(x).arg(dst).arg(&ni);
20858 unsafe {
20859 b.launch(cfg)?;
20860 }
20861 Ok(())
20862 }
20863
20864 pub fn gdn_chunk_k123_vl8(
20867 &self,
20868 seqs: &[GdnSeqVl],
20869 n_head: usize,
20870 hk: usize,
20871 wq: Option<&GdnWVl8>,
20872 ) -> Result<(), Box<dyn std::error::Error>> {
20873 let b = seqs.len();
20874 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
20875 let mut packed = [GdnSeqVl::default(); 8];
20876 packed[..b].copy_from_slice(seqs);
20877 let v = GdnVl8(packed);
20878 let (hi, ci) = (n_head as i32, 32i32);
20879 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
20880 {
20881 let f = self.func("gdn_chunk_cumgate_vl");
20882 let cfg = LaunchConfig {
20883 grid_dim: (max_nc, n_head as u32, b as u32),
20884 block_dim: (32, 1, 1),
20885 shared_mem_bytes: 0,
20886 };
20887 let __s_lb = self.gpu.stream();
20888 let mut lb = __s_lb.launch_builder(&f);
20889 lb.arg(&v).arg(&hi).arg(&ci);
20890 unsafe {
20891 lb.launch(cfg)?;
20892 }
20893 }
20894 let hki = hk as i32;
20895 if let Some(w) = wq {
20896 let f = self.func("gdn_k2_wgmma_vl");
20898 let cfg = LaunchConfig {
20899 grid_dim: (max_nc, n_head as u32, b as u32),
20900 block_dim: (128, 1, 1),
20901 shared_mem_bytes: 0,
20902 };
20903 let __s_lb = self.gpu.stream();
20904 let mut lb = __s_lb.launch_builder(&f);
20905 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
20906 unsafe {
20907 lb.launch(cfg)?;
20908 }
20909 } else {
20910 let f = self.func("gdn_chunk_attn_vl");
20911 let cfg = LaunchConfig {
20912 grid_dim: (max_nc, n_head as u32, b as u32),
20913 block_dim: (256, 1, 1),
20914 shared_mem_bytes: 0,
20915 };
20916 let __s_lb = self.gpu.stream();
20917 let mut lb = __s_lb.launch_builder(&f);
20918 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
20919 unsafe {
20920 lb.launch(cfg)?;
20921 }
20922 }
20923 {
20924 let f = self.func("gdn_chunk_solve32_vl");
20925 let cfg = LaunchConfig {
20926 grid_dim: (max_nc, n_head as u32, b as u32),
20927 block_dim: (256, 1, 1),
20928 shared_mem_bytes: 0,
20929 };
20930 let __s_lb = self.gpu.stream();
20931 let mut lb = __s_lb.launch_builder(&f);
20932 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
20933 unsafe {
20934 lb.launch(cfg)?;
20935 }
20936 }
20937 Ok(())
20938 }
20939
20940 #[allow(clippy::too_many_arguments)]
20944 pub fn gdn_prep_vl8(
20945 &self,
20946 seqs: &[GdnPrepVl],
20947 conv_w: &CudaSlice<f32>,
20948 dt_bias: &CudaSlice<f32>,
20949 a: &CudaSlice<f32>,
20950 conv_dim: usize,
20951 d_conv: usize,
20952 d_state: usize,
20953 num_v: usize,
20954 num_k: usize,
20955 key_dim: usize,
20956 hk: usize,
20957 eps: f32,
20958 ) -> Result<(), Box<dyn std::error::Error>> {
20959 let b = seqs.len();
20960 assert!(b >= 1 && b <= 8);
20961 let mut packed = [GdnPrepVl::default(); 8];
20962 packed[..b].copy_from_slice(seqs);
20963 let v = GdnPrepVl8(packed);
20964 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20965 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
20966 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
20967 assert!(
20968 conv_fuse || hk == num_v,
20969 "de-broadcast requires the fused conv"
20970 );
20971 if conv_fuse {
20972 let f = self.func("ssm_conv1d_gdn_state_vl");
20973 let cfg = LaunchConfig {
20974 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
20975 block_dim: (256, 1, 1),
20976 shared_mem_bytes: 0,
20977 };
20978 let (dsi, nvi, nki, kdi, hki) = (
20979 d_state as i32,
20980 num_v as i32,
20981 num_k as i32,
20982 key_dim as i32,
20983 hk as i32,
20984 );
20985 let __s_lb = self.gpu.stream();
20986 let mut lb = __s_lb.launch_builder(&f);
20987 lb.arg(&v)
20988 .arg(conv_w)
20989 .arg(&cdi)
20990 .arg(&dci)
20991 .arg(&dsi)
20992 .arg(&nvi)
20993 .arg(&nki)
20994 .arg(&kdi)
20995 .arg(&hki);
20996 unsafe {
20997 lb.launch(cfg)?;
20998 }
20999 } else {
21000 let f = self.func("ssm_conv1d_tm_state_vl");
21001 let cfg = LaunchConfig {
21002 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21003 block_dim: (256, 1, 1),
21004 shared_mem_bytes: 0,
21005 };
21006 let __s_lb = self.gpu.stream();
21007 let mut lb = __s_lb.launch_builder(&f);
21008 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21009 unsafe {
21010 lb.launch(cfg)?;
21011 }
21012 }
21013 {
21014 let f = self.func("ssm_conv_ring_update_vl");
21015 let n = (conv_dim * (d_conv - 1)) as u32;
21016 let cfg = LaunchConfig {
21017 grid_dim: (n.div_ceil(256), 1, b as u32),
21018 block_dim: (256, 1, 1),
21019 shared_mem_bytes: 0,
21020 };
21021 let __s_lb = self.gpu.stream();
21022 let mut lb = __s_lb.launch_builder(&f);
21023 lb.arg(&v).arg(&cdi).arg(&dci);
21024 unsafe {
21025 lb.launch(cfg)?;
21026 }
21027 }
21028 if !conv_fuse {
21029 let f = self.func("qkv_to_gdn_repack_vl");
21030 let n = max_t * (num_v * d_state) as u32;
21031 let cfg = LaunchConfig {
21032 grid_dim: (n.div_ceil(256), 1, b as u32),
21033 block_dim: (256, 1, 1),
21034 shared_mem_bytes: 0,
21035 };
21036 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21037 let __s_lb = self.gpu.stream();
21038 let mut lb = __s_lb.launch_builder(&f);
21039 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21040 unsafe {
21041 lb.launch(cfg)?;
21042 }
21043 }
21044 if Self::l2_v2_on(d_state) {
21045 let f = self.func("gdn_l2_v2_vl");
21046 let cfg = LaunchConfig {
21047 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21048 block_dim: (256, 1, 1),
21049 shared_mem_bytes: 0,
21050 };
21051 let (dsi, nvi) = (d_state as i32, hk as i32);
21052 let __s_lb = self.gpu.stream();
21053 let mut lb = __s_lb.launch_builder(&f);
21054 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21055 unsafe {
21056 lb.launch(cfg)?;
21057 }
21058 } else {
21059 let f = self.func("gdn_l2_vl");
21060 let cfg = LaunchConfig {
21061 grid_dim: (max_t * hk as u32, 2, b as u32),
21062 block_dim: (256, 1, 1),
21063 shared_mem_bytes: 0,
21064 };
21065 let (dsi, nvi) = (d_state as i32, hk as i32);
21066 let __s_lb = self.gpu.stream();
21067 let mut lb = __s_lb.launch_builder(&f);
21068 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21069 unsafe {
21070 lb.launch(cfg)?;
21071 }
21072 }
21073 {
21074 let f = self.func("gdn_gate_prep_vl");
21075 let n = max_t * num_v as u32;
21076 let cfg = LaunchConfig {
21077 grid_dim: (n.div_ceil(256), 1, b as u32),
21078 block_dim: (256, 1, 1),
21079 shared_mem_bytes: 0,
21080 };
21081 let nvi = num_v as i32;
21082 let __s_lb = self.gpu.stream();
21083 let mut lb = __s_lb.launch_builder(&f);
21084 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21085 unsafe {
21086 lb.launch(cfg)?;
21087 }
21088 }
21089 Ok(())
21090 }
21091
21092 pub fn gdn_mirror_vl8(
21094 &self,
21095 seqs: &[GdnSeqVl],
21096 n_head: usize,
21097 which: i32,
21098 hk: usize,
21099 ) -> Result<(), Box<dyn std::error::Error>> {
21100 let b = seqs.len();
21101 assert!(b >= 1 && b <= 8);
21102 let mut packed = [GdnSeqVl::default(); 8];
21103 packed[..b].copy_from_slice(seqs);
21104 let v = GdnVl8(packed);
21105 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21106 let max_n = seqs
21107 .iter()
21108 .map(|s| {
21109 if which == 0 {
21110 s.t as i64 * ept as i64
21111 } else {
21112 s.nc as i64 * ept as i64 * 32
21113 }
21114 })
21115 .max()
21116 .unwrap();
21117 let f = self.func("gdn_mirror_vl");
21118 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21119 let cfg = LaunchConfig {
21120 grid_dim: (blocks, 1, b as u32),
21121 block_dim: (256, 1, 1),
21122 shared_mem_bytes: 0,
21123 };
21124 let __s_lb = self.gpu.stream();
21125 let mut lb = __s_lb.launch_builder(&f);
21126 lb.arg(&v).arg(&ept).arg(&which);
21127 unsafe {
21128 lb.launch(cfg)?;
21129 }
21130 Ok(())
21131 }
21132
21133 pub fn gdn_tail_vl8(
21135 &self,
21136 seqs: &[GdnPrepVl],
21137 norm_w: &CudaSlice<f32>,
21138 d_state: usize,
21139 num_v: usize,
21140 eps: f32,
21141 ) -> Result<(), Box<dyn std::error::Error>> {
21142 let b = seqs.len();
21143 assert!(b >= 1 && b <= 8);
21144 let mut packed = [GdnPrepVl::default(); 8];
21145 packed[..b].copy_from_slice(seqs);
21146 let v = GdnPrepVl8(packed);
21147 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21148 let f = self.func("gated_rmsnorm_f16out_vl");
21149 let cfg = LaunchConfig {
21151 grid_dim: (max_t * num_v as u32, 1, b as u32),
21152 block_dim: (128, 1, 1),
21153 shared_mem_bytes: 0,
21154 };
21155 let (dsi, nvi) = (d_state as i32, num_v as i32);
21156 let __s_lb = self.gpu.stream();
21157 let mut lb = __s_lb.launch_builder(&f);
21158 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21159 unsafe {
21160 lb.launch(cfg)?;
21161 }
21162 Ok(())
21163 }
21164
21165 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
21168 use cudarc::driver::DevicePtr;
21169 let s = self.gpu.stream();
21170 let (p, _g) = x.device_ptr(&s);
21171 p as u64
21172 }
21173 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
21174 use cudarc::driver::DevicePtrMut;
21175 let s = self.gpu.stream();
21176 let (p, _g) = x.device_ptr_mut(&s);
21177 p as u64
21178 }
21179 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
21180 use cudarc::driver::DevicePtr;
21181 let s = self.gpu.stream();
21182 let (p, _g) = x.device_ptr(&s);
21183 p as u64
21184 }
21185 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
21186 use cudarc::driver::DevicePtr;
21187 let s = self.gpu.stream();
21188 let (p, _g) = x.device_ptr(&s);
21189 p as u64
21190 }
21191
21192 pub fn gdn_chunk_vl8(
21196 &self,
21197 seqs: &[GdnSeqVl],
21198 n_head: usize,
21199 scale: f32,
21200 hk: usize,
21201 wq: Option<&GdnWVl8>,
21202 ) -> Result<(), Box<dyn std::error::Error>> {
21203 const NSPLIT: u32 = 4;
21204 let b = seqs.len();
21205 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
21206 let mut packed = [GdnSeqVl::default(); 8];
21207 packed[..b].copy_from_slice(seqs);
21208 let v = GdnVl8(packed);
21209 let (hi, ci) = (n_head as i32, 32i32);
21210 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21211 let hki = hk as i32;
21212 if let Some(w) = wq {
21213 let f = self.func("gdn_k45_wgmma_vl");
21215 let cfg = LaunchConfig {
21216 grid_dim: (n_head as u32, NSPLIT, b as u32),
21217 block_dim: (256, 1, 1),
21218 shared_mem_bytes: 0,
21219 };
21220 let __s_lb = self.gpu.stream();
21221 let mut lb = __s_lb.launch_builder(&f);
21222 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
21223 unsafe {
21224 lb.launch(cfg)?;
21225 }
21226 let _ = max_nc;
21227 return Ok(());
21228 }
21229 {
21230 let f = self.func("gdn_chunk_state_mma_vl");
21231 let cfg = LaunchConfig {
21232 grid_dim: (n_head as u32, NSPLIT, b as u32),
21233 block_dim: (256, 1, 1),
21234 shared_mem_bytes: 0,
21235 };
21236 let __s_lb = self.gpu.stream();
21237 let mut lb = __s_lb.launch_builder(&f);
21238 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21239 unsafe {
21240 lb.launch(cfg)?;
21241 }
21242 }
21243 {
21244 let f = self.func("gdn_chunk_output_mma_vl");
21245 let cfg = LaunchConfig {
21246 grid_dim: (max_nc, n_head as u32, b as u32),
21247 block_dim: (256, 1, 1),
21248 shared_mem_bytes: 0,
21249 };
21250 let __s_lb = self.gpu.stream();
21251 let mut lb = __s_lb.launch_builder(&f);
21252 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
21253 unsafe {
21254 lb.launch(cfg)?;
21255 }
21256 }
21257 Ok(())
21258 }
21259 pub fn gdn_scan_chunked(
21260 &self,
21261 q: &CudaSlice<f32>,
21262 k: &CudaSlice<f32>,
21263 v: &CudaSlice<f32>,
21264 g: &CudaSlice<f32>,
21265 beta: &CudaSlice<f32>,
21266 kb16_pre: Option<&CudaSlice<u8>>,
21267 qb16_pre: Option<&CudaSlice<u8>>,
21268 state_in: &CudaSlice<f32>,
21269 state_out: &mut CudaSlice<f32>,
21270 o: &mut CudaSlice<f32>,
21271 n_head: usize,
21272 t: usize,
21273 scale: f32,
21274 c: usize,
21275 hk: usize,
21276 ) -> Result<(), Box<dyn std::error::Error>> {
21277 const D: usize = 128;
21278 const NSPLIT: u32 = 4;
21279 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
21280 let h = n_head;
21281 let nc = (t + c - 1) / c;
21282 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21283 let gdn_mma_pre = !portable_mma_gated()
21287 && c == 32
21288 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21289 Ok("1") => true,
21290 Ok("0") => false,
21291 _ => cfg!(memra_hopper_mma),
21292 };
21293 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
21294 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
21295 } else {
21296 None
21297 };
21298 let gdn_wgmma_pre = gdn_mma_pre
21302 && match std::env::var("MEMRA_GDN_WGMMA").as_deref() {
21303 Ok("0") => false,
21304 Ok("1") => true,
21305 _ => cfg!(memra_hopper_mma),
21306 };
21307 let nk = t * hk * D;
21308 let mut kb16_local: Option<CudaSlice<u8>> = None;
21309 if gdn_mma_pre && kb16_pre.is_none() {
21310 let mut kb = self.alloc_u8_uninit(nk * 2)?;
21311 let f = self.func("f32_to_bf16_bulk");
21312 let n2 = nk as i64;
21313 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21314 let __s_b = self.gpu.stream();
21315 let mut b = __s_b.launch_builder(&f);
21316 b.arg(k).arg(&mut kb).arg(&n2);
21317 unsafe {
21318 b.launch(cfg2)?;
21319 }
21320 kb16_local = Some(kb);
21321 }
21322 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
21323 if let Some(kb) = kb16_pre {
21324 assert!(kb.len() >= nk * 2, "kb16_pre too small");
21325 }
21326 let mut qb16: Option<CudaSlice<u8>> = None;
21327 let mut pb16: Option<CudaSlice<u8>> = None;
21328 if gdn_wgmma_pre {
21329 if qb16_pre.is_none() {
21332 let mut qb = self.alloc_u8_uninit(nk * 2)?;
21333 let f = self.func("f32_to_bf16_bulk");
21334 let n2 = nk as i64;
21335 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21336 let __s_b = self.gpu.stream();
21337 let mut b = __s_b.launch_builder(&f);
21338 b.arg(q).arg(&mut qb).arg(&n2);
21339 unsafe {
21340 b.launch(cfg2)?;
21341 }
21342 qb16 = Some(qb);
21343 } else if let Some(qb) = qb16_pre {
21344 assert!(qb.len() >= nk * 2, "qb16_pre too small");
21345 }
21346 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
21347 }
21348 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
21349 let k2w = if gdn_wgmma_pre {
21350 Some((
21351 *qb16_ref0.as_ref().unwrap(),
21352 *kb16_ref0.as_ref().unwrap(),
21353 pb16.as_mut().unwrap(),
21354 ))
21355 } else {
21356 None
21357 };
21358 let (gcum, p, u, w) =
21359 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
21360 let _ = &w;
21361 let mut y = self.uninit(nc * h * c * D)?;
21362 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
21374 && c == 32
21375 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21376 Ok("1") => true,
21377 Ok("0") => false,
21378 _ => cfg!(memra_hopper_mma),
21379 };
21380 if gdn_mma {
21381 let wb16 = wb16_pre
21382 .take()
21383 .expect("mma path pre-allocates wb16 (K3 store fold)");
21384 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
21385 if gdn_wgmma_pre {
21397 let qb16 = qb16_ref0.unwrap();
21399 let pb16 = pb16.as_ref().unwrap();
21400 {
21401 let f = self.func("gdn_k45_wgmma");
21402 let cfg = LaunchConfig {
21403 grid_dim: (h as u32, 4, 1),
21404 block_dim: (256, 1, 1),
21405 shared_mem_bytes: 0,
21406 };
21407 let hki = hk as i32;
21408 let __s_b = self.gpu.stream();
21409 let mut b = __s_b.launch_builder(&f);
21410 b.arg(kb16_ref)
21411 .arg(&gcum)
21412 .arg(beta)
21413 .arg(&u)
21414 .arg(&wb16)
21415 .arg(qb16)
21416 .arg(pb16)
21417 .arg(o)
21418 .arg(&scale)
21419 .arg(state_in)
21420 .arg(&mut *state_out)
21421 .arg(&hi)
21422 .arg(&ti)
21423 .arg(&ci)
21424 .arg(&hki);
21425 unsafe {
21426 b.launch(cfg)?;
21427 }
21428 }
21429 return Ok(());
21430 }
21431 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
21435 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
21436 {
21437 let f = self.func("gdn_chunk_state_mma");
21438 let cfg = LaunchConfig {
21439 grid_dim: (h as u32, NSPLIT, 1),
21440 block_dim: (256, 1, 1),
21441 shared_mem_bytes: 0,
21442 };
21443 let hki = hk as i32;
21444 let __s_b = self.gpu.stream();
21445 let mut b = __s_b.launch_builder(&f);
21446 b.arg(kb16_ref)
21447 .arg(&gcum)
21448 .arg(beta)
21449 .arg(&u)
21450 .arg(&wb16)
21451 .arg(&mut y16)
21452 .arg(&mut ssnap16)
21453 .arg(state_in)
21454 .arg(&mut *state_out)
21455 .arg(&hi)
21456 .arg(&ti)
21457 .arg(&ci)
21458 .arg(&hki);
21459 unsafe {
21460 b.launch(cfg)?;
21461 }
21462 }
21463 {
21464 let f = self.func("gdn_chunk_output_mma");
21466 let jt = ((c + 31) / 32) as u32;
21467 let cfg = LaunchConfig {
21468 grid_dim: (nc as u32, h as u32, jt),
21469 block_dim: (256, 1, 1),
21470 shared_mem_bytes: 0,
21471 };
21472 let hki = hk as i32;
21473 let __s_b = self.gpu.stream();
21474 let mut b = __s_b.launch_builder(&f);
21475 b.arg(q)
21476 .arg(&gcum)
21477 .arg(&p)
21478 .arg(&y16)
21479 .arg(&ssnap16)
21480 .arg(o)
21481 .arg(&hi)
21482 .arg(&ti)
21483 .arg(&ci)
21484 .arg(&scale)
21485 .arg(&hki);
21486 unsafe {
21487 b.launch(cfg)?;
21488 }
21489 }
21490 return Ok(());
21491 }
21492 {
21493 let f = self.func("gdn_chunk_state_f32");
21495 let cfg = LaunchConfig {
21496 grid_dim: (h as u32, NSPLIT, 1),
21497 block_dim: (256, 1, 1),
21498 shared_mem_bytes: 0,
21499 };
21500 let __s_b = self.gpu.stream();
21501 let mut b = __s_b.launch_builder(&f);
21502 b.arg(k)
21503 .arg(&gcum)
21504 .arg(beta)
21505 .arg(&u)
21506 .arg(&w)
21507 .arg(&mut y)
21508 .arg(&mut ssnap)
21509 .arg(state_in)
21510 .arg(&mut *state_out)
21511 .arg(&hi)
21512 .arg(&ti)
21513 .arg(&ci);
21514 unsafe {
21515 b.launch(cfg)?;
21516 }
21517 }
21518 {
21519 let f = self.func("gdn_chunk_output_f32");
21521 let jt = ((c + 31) / 32) as u32;
21522 let cfg = LaunchConfig {
21523 grid_dim: (nc as u32, h as u32, jt),
21524 block_dim: (256, 1, 1),
21525 shared_mem_bytes: 0,
21526 };
21527 let __s_b = self.gpu.stream();
21528 let mut b = __s_b.launch_builder(&f);
21529 b.arg(q)
21530 .arg(&gcum)
21531 .arg(&p)
21532 .arg(&y)
21533 .arg(&ssnap)
21534 .arg(o)
21535 .arg(&hi)
21536 .arg(&ti)
21537 .arg(&ci)
21538 .arg(&scale);
21539 unsafe {
21540 b.launch(cfg)?;
21541 }
21542 }
21543 Ok(())
21544 }
21545
21546 #[allow(clippy::too_many_arguments)]
21555 #[allow(clippy::too_many_arguments)]
21556 pub fn gdn_scan_prefill(
21557 &self,
21558 q: &CudaSlice<f32>,
21559 k: &CudaSlice<f32>,
21560 v: &CudaSlice<f32>,
21561 g: &CudaSlice<f32>,
21562 beta: &CudaSlice<f32>,
21563 kb16_pre: Option<&CudaSlice<u8>>,
21564 qb16_pre: Option<&CudaSlice<u8>>,
21565 state_in: &CudaSlice<f32>,
21566 state_out: &mut CudaSlice<f32>,
21567 o: &mut CudaSlice<f32>,
21568 n_head: usize,
21569 t: usize,
21570 scale: f32,
21571 hk: usize,
21572 ) -> Result<(), Box<dyn std::error::Error>> {
21573 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
21574 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
21575 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
21576 }
21577 if Self::gdn_chunked_enabled() && t >= 16 {
21578 self.gdn_scan_chunked(
21579 q,
21580 k,
21581 v,
21582 g,
21583 beta,
21584 kb16_pre,
21585 qb16_pre,
21586 state_in,
21587 state_out,
21588 o,
21589 n_head,
21590 t,
21591 scale,
21592 Self::gdn_chunk_size(),
21593 hk,
21594 )
21595 } else {
21596 assert!(
21597 hk == n_head,
21598 "s128 scan is broadcast-only (prep guarantees by predicate)"
21599 );
21600 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
21601 }
21602 }
21603
21604 #[allow(clippy::too_many_arguments)]
21606 fn gdn_scan_diff(
21607 &self,
21608 q: &CudaSlice<f32>,
21609 k: &CudaSlice<f32>,
21610 v: &CudaSlice<f32>,
21611 g: &CudaSlice<f32>,
21612 beta: &CudaSlice<f32>,
21613 state_in: &CudaSlice<f32>,
21614 state_out: &mut CudaSlice<f32>,
21615 o: &mut CudaSlice<f32>,
21616 n_head: usize,
21617 t: usize,
21618 scale: f32,
21619 ) -> Result<(), Box<dyn std::error::Error>> {
21620 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
21621 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
21622 let mut o_c = self.uninit(o.len())?;
21623 let mut st_c = self.uninit(state_out.len())?;
21624 self.gdn_scan_chunked(
21625 q,
21626 k,
21627 v,
21628 g,
21629 beta,
21630 None,
21631 None,
21632 state_in,
21633 &mut st_c,
21634 &mut o_c,
21635 n_head,
21636 t,
21637 scale,
21638 Self::gdn_chunk_size(),
21639 n_head,
21640 )?;
21641 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
21642 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
21643 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
21644 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
21645 let mut max_abs = 0f32;
21646 let mut max_rel = 0f32;
21647 let mut sum_rel = 0f64;
21648 for (x, y) in a.iter().zip(b) {
21649 let ad = (x - y).abs();
21650 let rel = ad / x.abs().max(y.abs()).max(1e-3);
21651 if ad > max_abs {
21652 max_abs = ad;
21653 }
21654 if rel > max_rel {
21655 max_rel = rel;
21656 }
21657 sum_rel += rel as f64;
21658 }
21659 (max_abs, max_rel, sum_rel / a.len() as f64)
21660 };
21661 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
21662 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
21663 println!(
21664 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
21665 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
21666 Self::gdn_chunk_size()
21667 );
21668 Ok(())
21669 }
21670
21671 pub fn gdn_glog(
21673 &self,
21674 alpha: &CudaSlice<f32>,
21675 dt_bias: &CudaSlice<f32>,
21676 a: &CudaSlice<f32>,
21677 g_log: &mut CudaSlice<f32>,
21678 n_head: usize,
21679 t: usize,
21680 ) -> Result<(), Box<dyn std::error::Error>> {
21681 let f = self.func("gdn_glog_f32");
21682 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
21683 let (h, ti) = (n_head as i32, t as i32);
21684 let __s_b = self.gpu.stream();
21685 let mut b = __s_b.launch_builder(&f);
21686 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
21687 unsafe {
21688 b.launch(cfg)?;
21689 }
21690 Ok(())
21691 }
21692
21693 pub fn sigmoid_v(
21696 &self,
21697 x: &cudarc::driver::CudaView<f32>,
21698 y: &mut CudaSlice<f32>,
21699 n: usize,
21700 ) -> Result<(), Box<dyn std::error::Error>> {
21701 let f = self.func("sigmoid_f32");
21702 let cfg = LaunchConfig::for_num_elems(n as u32);
21703 let ni = n as i32;
21704 let __s_b = self.gpu.stream();
21705 let mut b = __s_b.launch_builder(&f);
21706 b.arg(x).arg(y).arg(&ni);
21707 unsafe {
21708 b.launch(cfg)?;
21709 }
21710 Ok(())
21711 }
21712
21713 pub fn gdn_glog_v(
21714 &self,
21715 alpha: &cudarc::driver::CudaView<f32>,
21716 dt_bias: &CudaSlice<f32>,
21717 a: &CudaSlice<f32>,
21718 g_log: &mut CudaSlice<f32>,
21719 n_head: usize,
21720 t: usize,
21721 ) -> Result<(), Box<dyn std::error::Error>> {
21722 let f = self.func("gdn_glog_f32");
21723 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
21724 let (h, ti) = (n_head as i32, t as i32);
21725 let __s_b = self.gpu.stream();
21726 let mut b = __s_b.launch_builder(&f);
21727 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
21728 unsafe {
21729 b.launch(cfg)?;
21730 }
21731 Ok(())
21732 }
21733
21734 pub fn sigmoid(
21735 &self,
21736 x: &CudaSlice<f32>,
21737 y: &mut CudaSlice<f32>,
21738 n: usize,
21739 ) -> Result<(), Box<dyn std::error::Error>> {
21740 let f = self.func("sigmoid_f32");
21741 let cfg = LaunchConfig::for_num_elems(n as u32);
21742 let ni = n as i32;
21743 let __s_b = self.gpu.stream();
21744 let mut b = __s_b.launch_builder(&f);
21745 b.arg(x).arg(y).arg(&ni);
21746 unsafe {
21747 b.launch(cfg)?;
21748 }
21749 Ok(())
21750 }
21751
21752 pub fn sig_mul_f16out(
21755 &self,
21756 a: &CudaSlice<f32>,
21757 g: &CudaSlice<f32>,
21758 dst: &mut CudaSlice<f32>,
21759 dst16: &mut CudaSlice<u8>,
21760 n: usize,
21761 ) -> Result<(), Box<dyn std::error::Error>> {
21762 let f = self.func("sig_mul_f16out_f32");
21763 let cfg = LaunchConfig::for_num_elems(n as u32);
21764 let ni = n as i32;
21765 let __s_b = self.gpu.stream();
21766 let mut b = __s_b.launch_builder(&f);
21767 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
21768 unsafe {
21769 b.launch(cfg)?;
21770 }
21771 Ok(())
21772 }
21773
21774 #[allow(clippy::too_many_arguments)]
21783 pub fn attn_head_gate(
21784 &self,
21785 a: &CudaSlice<f32>,
21786 g: &CudaSlice<f32>,
21787 dst: &mut CudaSlice<f32>,
21788 dst16: Option<&mut CudaSlice<u8>>,
21789 head_dim: usize,
21790 n_head: usize,
21791 t: usize,
21792 ) -> Result<(), Box<dyn std::error::Error>> {
21793 let f = self.func("attn_head_gate_f32");
21794 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
21795 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
21796 let d16: u64 = match dst16 {
21798 Some(d) => self.addr_u8(d),
21799 None => 0,
21800 };
21801 let __s_b = self.gpu.stream();
21802 let mut b = __s_b.launch_builder(&f);
21803 b.arg(a)
21804 .arg(g)
21805 .arg(dst)
21806 .arg(&d16)
21807 .arg(&hd)
21808 .arg(&nh)
21809 .arg(&ti);
21810 unsafe {
21811 b.launch(cfg)?;
21812 }
21813 Ok(())
21814 }
21815
21816 #[allow(clippy::too_many_arguments)]
21825 pub fn swiglu_clamped_mul_scaled(
21826 &self,
21827 gate: &CudaSlice<f32>,
21828 up: &CudaSlice<f32>,
21829 gs: f32,
21830 us: f32,
21831 limit: f32,
21832 dst: &mut CudaSlice<f32>,
21833 n: usize,
21834 ) -> Result<(), Box<dyn std::error::Error>> {
21835 debug_assert!(
21836 limit > 1e-6,
21837 "swiglu_clamped needs a live limit; use silu_mul_scaled"
21838 );
21839 let f = self.func("swiglu_clamped_mul_scaled_f32");
21840 let cfg = LaunchConfig::for_num_elems(n as u32);
21841 let ni = n as i32;
21842 let __s_b = self.gpu.stream();
21843 let mut b = __s_b.launch_builder(&f);
21844 b.arg(gate)
21845 .arg(up)
21846 .arg(&gs)
21847 .arg(&us)
21848 .arg(&limit)
21849 .arg(dst)
21850 .arg(&ni);
21851 unsafe {
21852 b.launch(cfg)?;
21853 }
21854 Ok(())
21855 }
21856
21857 pub fn gated_rmsnorm(
21859 &self,
21860 o: &CudaSlice<f32>,
21861 w: &CudaSlice<f32>,
21862 z: &CudaSlice<f32>,
21863 dst: &mut CudaSlice<f32>,
21864 ncols: usize,
21865 nrows: usize,
21866 eps: f32,
21867 ) -> Result<(), Box<dyn std::error::Error>> {
21868 let f = self.func("gated_rmsnorm_f32");
21869 let cfg = LaunchConfig {
21870 grid_dim: (nrows as u32, 1, 1),
21871 block_dim: (128, 1, 1),
21872 shared_mem_bytes: 0,
21873 };
21874 let (nc, e) = (ncols as i32, eps);
21875 let __s_b = self.gpu.stream();
21876 let mut b = __s_b.launch_builder(&f);
21877 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
21878 unsafe {
21879 b.launch(cfg)?;
21880 }
21881 Ok(())
21882 }
21883
21884 pub fn gated_rmsnorm_f16out(
21887 &self,
21888 o: &CudaSlice<f32>,
21889 w: &CudaSlice<f32>,
21890 z: &CudaSlice<f32>,
21891 dst: &mut CudaSlice<f32>,
21892 dst16: &mut CudaSlice<u8>,
21893 ncols: usize,
21894 nrows: usize,
21895 eps: f32,
21896 ) -> Result<(), Box<dyn std::error::Error>> {
21897 let f = self.func("gated_rmsnorm_f16out_f32");
21898 let cfg = LaunchConfig {
21900 grid_dim: (nrows as u32, 1, 1),
21901 block_dim: (128, 1, 1),
21902 shared_mem_bytes: 0,
21903 };
21904 let (nc, e) = (ncols as i32, eps);
21905 let __s_b = self.gpu.stream();
21906 let mut b = __s_b.launch_builder(&f);
21907 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
21908 unsafe {
21909 b.launch(cfg)?;
21910 }
21911 Ok(())
21912 }
21913
21914 #[allow(clippy::too_many_arguments)]
21918 pub fn add_rms_norm_zq8(
21919 &self,
21920 a: &CudaSlice<f32>,
21921 b_in: &CudaSlice<f32>,
21922 w: &CudaSlice<f32>,
21923 res: &mut CudaSlice<f32>,
21924 z: &mut CudaSlice<f32>,
21925 ncols: usize,
21926 nrows: usize,
21927 eps: f32,
21928 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21929 assert!(ncols % 32 == 0);
21930 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
21931 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
21932 let f = self.func("add_rms_norm_zq8");
21933 let cfg = LaunchConfig {
21934 grid_dim: (nrows as u32, 1, 1),
21935 block_dim: (1024, 1, 1),
21936 shared_mem_bytes: 0,
21937 };
21938 let (nc, ep) = (ncols as i32, eps);
21939 let __s_b = self.gpu.stream();
21940 let mut b = __s_b.launch_builder(&f);
21941 b.arg(a)
21942 .arg(b_in)
21943 .arg(w)
21944 .arg(res)
21945 .arg(z)
21946 .arg(&mut q)
21947 .arg(&mut d)
21948 .arg(&nc)
21949 .arg(&ep);
21950 unsafe {
21951 b.launch(cfg)?;
21952 }
21953 Ok((q, d))
21954 }
21955
21956 pub fn gated_rmsnorm_zv(
21961 &self,
21962 o: &CudaSlice<f32>,
21963 w: &CudaSlice<f32>,
21964 z: &cudarc::driver::CudaView<f32>,
21965 dst: &mut CudaSlice<f32>,
21966 ncols: usize,
21967 nrows: usize,
21968 eps: f32,
21969 ) -> Result<(), Box<dyn std::error::Error>> {
21970 let f = self.func("gated_rmsnorm_f32");
21971 let cfg = LaunchConfig {
21972 grid_dim: (nrows as u32, 1, 1),
21973 block_dim: (128, 1, 1),
21974 shared_mem_bytes: 0,
21975 };
21976 let (nc, e) = (ncols as i32, eps);
21977 let __s_b = self.gpu.stream();
21978 let mut b = __s_b.launch_builder(&f);
21979 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
21980 unsafe {
21981 b.launch(cfg)?;
21982 }
21983 Ok(())
21984 }
21985
21986 pub fn gated_rmsnorm_f16out_zv(
21987 &self,
21988 o: &CudaSlice<f32>,
21989 w: &CudaSlice<f32>,
21990 z: &cudarc::driver::CudaView<f32>,
21991 dst: &mut CudaSlice<f32>,
21992 dst16: &mut CudaSlice<u8>,
21993 ncols: usize,
21994 nrows: usize,
21995 eps: f32,
21996 ) -> Result<(), Box<dyn std::error::Error>> {
21997 let f = self.func("gated_rmsnorm_f16out_f32");
21998 let cfg = LaunchConfig {
22000 grid_dim: (nrows as u32, 1, 1),
22001 block_dim: (128, 1, 1),
22002 shared_mem_bytes: 0,
22003 };
22004 let (nc, e) = (ncols as i32, eps);
22005 let __s_b = self.gpu.stream();
22006 let mut b = __s_b.launch_builder(&f);
22007 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22008 unsafe {
22009 b.launch(cfg)?;
22010 }
22011 Ok(())
22012 }
22013
22014 pub fn gated_rmsnorm_q8_1(
22015 &self,
22016 o: &CudaSlice<f32>,
22017 w: &CudaSlice<f32>,
22018 z: &CudaSlice<f32>,
22019 ncols: usize,
22020 nrows: usize,
22021 eps: f32,
22022 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22023 assert!(ncols % 32 == 0);
22024 let f = self.func("gated_rmsnorm_q8_1");
22025 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22026 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22027 let cfg = LaunchConfig {
22028 grid_dim: (nrows as u32, 1, 1),
22029 block_dim: (128, 1, 1),
22030 shared_mem_bytes: 0,
22031 };
22032 let (nc, ep) = (ncols as i32, eps);
22033 let __s_b = self.gpu.stream();
22034 let mut b = __s_b.launch_builder(&f);
22035 b.arg(o)
22036 .arg(w)
22037 .arg(z)
22038 .arg(&mut out_q)
22039 .arg(&mut out_d)
22040 .arg(&nc)
22041 .arg(&ep);
22042 unsafe {
22043 b.launch(cfg)?;
22044 }
22045 Ok((out_q, out_d))
22046 }
22047
22048 pub fn transpose(
22050 &self,
22051 inp: &CudaSlice<f32>,
22052 rows: usize,
22053 cols: usize,
22054 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22055 let f = self.func("transpose_f32");
22056 let mut out = self.zeros(rows * cols)?;
22057 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22058 let (r, c) = (rows as i32, cols as i32);
22059 let __s_b = self.gpu.stream();
22060 let mut b = __s_b.launch_builder(&f);
22061 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22062 unsafe {
22063 b.launch(cfg)?;
22064 }
22065 Ok(out)
22066 }
22067
22068 pub fn repeat_heads(
22070 &self,
22071 inp: &CudaSlice<f32>,
22072 out: &mut CudaSlice<f32>,
22073 head_dim: usize,
22074 n_in: usize,
22075 n_out: usize,
22076 t: usize,
22077 ) -> Result<(), Box<dyn std::error::Error>> {
22078 let f = self.func("repeat_heads_f32");
22079 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22080 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22081 let __s_b = self.gpu.stream();
22082 let mut b = __s_b.launch_builder(&f);
22083 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22084 unsafe {
22085 b.launch(cfg)?;
22086 }
22087 Ok(())
22088 }
22089
22090 pub fn q_gate_split(
22093 &self,
22094 qf: &CudaSlice<f32>,
22095 q_out: &mut CudaSlice<f32>,
22096 gate_out: &mut CudaSlice<f32>,
22097 head_dim: usize,
22098 n_head: usize,
22099 t: usize,
22100 ) -> Result<(), Box<dyn std::error::Error>> {
22101 let f = self.func("q_gate_split_f32");
22102 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22103 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22104 let __s_b = self.gpu.stream();
22105 let mut b = __s_b.launch_builder(&f);
22106 b.arg(qf)
22107 .arg(q_out)
22108 .arg(gate_out)
22109 .arg(&hd)
22110 .arg(&nh)
22111 .arg(&ti);
22112 unsafe {
22113 b.launch(cfg)?;
22114 }
22115 Ok(())
22116 }
22117
22118 pub fn qkv_to_gdn_repack(
22122 &self,
22123 conv_out: &CudaSlice<f32>,
22124 q_g: &mut CudaSlice<f32>,
22125 k_g: &mut CudaSlice<f32>,
22126 v_g: &mut CudaSlice<f32>,
22127 d_state: usize,
22128 num_v: usize,
22129 num_k: usize,
22130 key_dim: usize,
22131 t: usize,
22132 ) -> Result<(), Box<dyn std::error::Error>> {
22133 let f = self.func("qkv_to_gdn_repack_f32");
22134 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22135 let (ds, nv, nk, kd, ti) = (
22136 d_state as i32,
22137 num_v as i32,
22138 num_k as i32,
22139 key_dim as i32,
22140 t as i32,
22141 );
22142 let __s_b = self.gpu.stream();
22143 let mut b = __s_b.launch_builder(&f);
22144 b.arg(conv_out)
22145 .arg(q_g)
22146 .arg(k_g)
22147 .arg(v_g)
22148 .arg(&ds)
22149 .arg(&nv)
22150 .arg(&nk)
22151 .arg(&kd)
22152 .arg(&ti);
22153 unsafe {
22154 b.launch(cfg)?;
22155 }
22156 Ok(())
22157 }
22158
22159 pub fn conv_left_pad(
22162 &self,
22163 src: &CudaSlice<f32>,
22164 dst: &mut CudaSlice<f32>,
22165 conv_dim: usize,
22166 t: usize,
22167 pad: usize,
22168 ) -> Result<(), Box<dyn std::error::Error>> {
22169 let f = self.func("conv_left_pad_f32");
22170 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
22171 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
22172 let __s_b = self.gpu.stream();
22173 let mut b = __s_b.launch_builder(&f);
22174 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
22175 unsafe {
22176 b.launch(cfg)?;
22177 }
22178 Ok(())
22179 }
22180
22181 pub fn conv_assemble_and_roll(
22185 &self,
22186 qkv_col: &CudaSlice<f32>,
22187 conv_state: &mut CudaSlice<f32>,
22188 conv_in: &mut CudaSlice<f32>,
22189 conv_dim: usize,
22190 pad: usize,
22191 ) -> Result<(), Box<dyn std::error::Error>> {
22192 let f = self.func("conv_assemble_and_roll_f32");
22193 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22194 let (cd, p) = (conv_dim as i32, pad as i32);
22195 let __s_b = self.gpu.stream();
22196 let mut b = __s_b.launch_builder(&f);
22197 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
22198 unsafe {
22199 b.launch(cfg)?;
22200 }
22201 Ok(())
22202 }
22203
22204 pub fn ssm_conv1d_fused_decode(
22210 &self,
22211 qkv_col: &CudaSlice<f32>,
22212 conv_state: &mut CudaSlice<f32>,
22213 w: &CudaSlice<f32>,
22214 conv_out: &mut CudaSlice<f32>,
22215 conv_dim: usize,
22216 d_conv: usize,
22217 ) -> Result<(), Box<dyn std::error::Error>> {
22218 let f = self.func("ssm_conv1d_fused_decode_f32");
22219 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22220 let (cd, dc) = (conv_dim as i32, d_conv as i32);
22221 let __s_b = self.gpu.stream();
22222 let mut b = __s_b.launch_builder(&f);
22223 b.arg(qkv_col)
22224 .arg(conv_state)
22225 .arg(w)
22226 .arg(conv_out)
22227 .arg(&cd)
22228 .arg(&dc);
22229 unsafe {
22230 b.launch(cfg)?;
22231 }
22232 Ok(())
22233 }
22234
22235 pub fn slice_range(
22238 &self,
22239 src: &CudaSlice<f32>,
22240 start: usize,
22241 len: usize,
22242 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22243 let host = self.gpu.stream().clone_dtoh(src)?;
22244 self.gpu.stream().synchronize()?;
22245 Ok(self.htod(&host[start..start + len])?)
22246 }
22247}
22248
22249#[cfg(test)]
22250mod target_dispatch_tests {
22251 use super::legacy_quant_gemm_allowed;
22252
22253 #[test]
22254 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
22255 assert!(legacy_quant_gemm_allowed(false, false, false));
22257 assert!(!legacy_quant_gemm_allowed(false, false, true));
22258 assert!(!legacy_quant_gemm_allowed(true, false, false));
22260 assert!(!legacy_quant_gemm_allowed(true, false, true));
22261 assert!(legacy_quant_gemm_allowed(true, true, false));
22263 assert!(!legacy_quant_gemm_allowed(true, true, true));
22264 }
22265
22266 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
22267 #[test]
22268 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
22269 assert!(!legacy_quant_gemm_allowed(
22270 cfg!(memra_portable_cuda),
22271 cfg!(memra_hopper_mma),
22272 false
22273 ));
22274 }
22275
22276 #[cfg(memra_hopper_mma)]
22277 #[test]
22278 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
22279 assert!(legacy_quant_gemm_allowed(
22280 cfg!(memra_portable_cuda),
22281 cfg!(memra_hopper_mma),
22282 false
22283 ));
22284 assert!(super::portable_mma_gated() == false);
22285 }
22286}
22287
22288impl memra_kv::KvDev for Engine {
22291 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22292 Engine::zeros(self, n)
22293 }
22294 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22295 Engine::uninit(self, n)
22296 }
22297 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
22298 Engine::alloc_u8(self, n)
22299 }
22300 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
22301 Engine::htod_i32(self, v)
22302 }
22303 fn clone_dtod(
22304 &self,
22305 src: &CudaSlice<f32>,
22306 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22307 Engine::clone_dtod(self, src)
22308 }
22309 fn copy_into(
22310 &self,
22311 dst: &mut CudaSlice<f32>,
22312 off: usize,
22313 src: &CudaSlice<f32>,
22314 len: usize,
22315 ) -> Result<(), Box<dyn std::error::Error>> {
22316 Engine::copy_into(self, dst, off, src, len)
22317 }
22318 fn set_i32_one(
22319 &self,
22320 d: &mut CudaSlice<i32>,
22321 v: i32,
22322 ) -> Result<(), Box<dyn std::error::Error>> {
22323 Engine::set_i32_one(self, d, v)
22324 }
22325}