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 dsv4_ffi;
242pub mod dsv4_gpu;
243pub mod f16_ffi;
244pub mod fp8_ffi;
245pub mod mmq_ffi;
246pub mod moe_cache;
247pub mod prime_graph;
248pub mod spill;
249mod spill_pread;
250
251const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
258const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
259const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
260const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
261const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
262const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
263const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
265
266fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
272 assert!(
273 !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
274 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
275 );
276 match std::env::var("MEMRA_GEMM_FATBIN") {
277 Ok(path) => std::borrow::Cow::Owned(
278 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
279 ),
280 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
281 }
282}
283
284pub(crate) const fn portable_mma_gated() -> bool {
291 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
292}
293
294pub(crate) const fn gdn_mma_default_on() -> bool {
303 cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
304}
305
306const fn konst_eq(a: &str, b: &str) -> bool {
308 let (a, b) = (a.as_bytes(), b.as_bytes());
309 if a.len() != b.len() {
310 return false;
311 }
312 let mut i = 0;
313 while i < a.len() {
314 if a[i] != b[i] {
315 return false;
316 }
317 i += 1;
318 }
319 true
320}
321
322const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
327 (!portable_cuda || hopper_mma) && !no_gemm
328}
329
330const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
338const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
339const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
340const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
341const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
342
343pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
346
347fn flash_fatbin_bytes() -> &'static [u8] {
349 match kv_cache_formats() {
350 ("q8_0", "q5_1") => FLASH_FATBIN,
351 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
352 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
353 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
354 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
355 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
356 other => unreachable!("kv_cache_formats returned {other:?}"),
357 }
358}
359
360fn k1_launch_override() -> Option<(u32, u32, u32)> {
367 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
368 *K1.get_or_init(|| {
369 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
370 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
371 match p.as_slice() {
372 [bm, bn, w] => Some((*bm, *bn, *w)),
373 _ => None,
374 }
375 })
376}
377
378pub(crate) fn wgmma_gemm_enabled() -> bool {
385 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
386 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
387}
388
389pub const FA_VEC_MIN_TKV: usize = 96;
404pub fn fa_vec_min_tkv() -> usize {
408 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
409 *V.get_or_init(|| {
410 std::env::var("MEMRA_FA_VEC_MIN")
411 .ok()
412 .and_then(|v| v.parse().ok())
413 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
414 })
415}
416
417pub fn fa_f16pv_on() -> bool {
428 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
429 *ON.get_or_init(|| {
430 std::env::var("MEMRA_FA_F16PV")
431 .map(|v| v != "0")
432 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
433 })
434}
435
436pub fn fa512_hp_on() -> bool {
440 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
441 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
442}
443
444pub fn faw_hp_on() -> bool {
448 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
449 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
450}
451
452pub fn fa512_wide_warps() -> usize {
456 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
457 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
458 Ok("1") => 4,
459 _ => 2,
460 })
461}
462
463pub fn fa512_min_tkv() -> usize {
466 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
467 *FA512_MIN.get_or_init(|| {
468 std::env::var("MEMRA_FA512_MIN")
469 .ok()
470 .and_then(|v| v.parse().ok())
471 .unwrap_or(512)
472 })
473}
474pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
478 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
479pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
483pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
489 std::sync::atomic::AtomicBool::new(false);
490pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
497 std::sync::atomic::AtomicBool::new(true);
498pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
499 std::sync::atomic::AtomicUsize::new(16);
500pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
505pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
507pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
512pub use memra_kv::KV_FP8_FORCE;
515pub(crate) fn rms_block() -> u32 {
516 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
517 *V.get_or_init(|| {
518 std::env::var("MEMRA_RMS_BLOCK")
519 .ok()
520 .and_then(|v| v.parse().ok())
521 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
522 })
523}
524
525pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
526 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
527 if let Some(forced) = *S.get_or_init(|| {
528 std::env::var("MEMRA_FA_SPLIT")
529 .ok()
530 .and_then(|v| v.parse().ok())
531 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
532 }) {
533 return forced;
534 }
535 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
553 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
554 {
555 return if t_kv <= 8192 {
556 16
557 } else if t_kv <= 16384 {
558 64
559 } else {
560 128
561 };
562 }
563 let big_rig = fa_sm_count() >= 128;
564 if big_rig {
565 let _ = n_head_kv;
566 if t_kv <= 2048 {
567 16
568 } else if t_kv <= 16384 {
569 64
570 } else {
571 128
572 }
573 } else if n_head_kv <= 4 {
574 if t_kv <= 512 {
595 8
596 } else if t_kv <= 16384 {
597 64
598 } else {
599 128
600 }
601 } else {
602 if t_kv <= 8192 {
603 32
604 } else if t_kv <= 16384 {
605 64
606 } else {
607 128
608 }
609 }
610}
611
612fn fa_sm_count() -> i32 {
615 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
616 *N.get_or_init(|| {
617 cudarc::driver::result::init().ok();
618 cudarc::driver::result::device::get(0)
619 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
620 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
621 .unwrap_or(82)
622 })
623}
624
625fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
629 match head_dim {
630 256 => Ok(""),
631 128 => Ok("_hd128"),
632 d => Err(format!(
633 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
634 callers must gate to sdpa_naive"
635 )
636 .into()),
637 }
638}
639
640pub const QT_Q8_0: i32 = 0;
642pub const QT_Q4_K: i32 = 1;
643pub const QT_Q6_K: i32 = 2;
644pub const QT_Q5_K: i32 = 3;
645pub const QT_Q3_K: i32 = 4;
646pub const QT_IQ4_XS: i32 = 5;
647pub const QT_IQ3_S: i32 = 6;
648pub const QT_NVFP4: i32 = 7;
649pub const QT_F8_E4M3: i32 = 10;
655pub const QT_NVFP4_RP: i32 = 9;
658pub const QT_F32: i32 = 8;
660pub const QT_BF16: i32 = 11;
661pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
666pub const QT_F8_E4M3_BLK: i32 = 14;
682
683pub struct Engine {
685 pub gpu: memra_runtime::Gpu,
686 module: Arc<CudaModule>,
687 hybrid: Arc<CudaModule>,
688 qmatvec: Arc<CudaModule>,
689 flash: Arc<CudaModule>,
690 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
694 gemm: Arc<CudaModule>,
695 router: Arc<CudaModule>,
696 sample: Arc<CudaModule>,
698 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
702 moe_cache_layout: Mutex<Option<Vec<usize>>>,
706 capture_keep_on: std::sync::atomic::AtomicBool,
712 verify_exact: std::sync::atomic::AtomicBool,
717 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
718 pub copy_stream: Arc<CudaStream>,
720 #[cfg(memra_cutlass)]
727 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
728 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
732 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
735 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
739 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
743 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
745 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
746 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
751 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
756 router_stage: Mutex<Option<PinnedStage>>,
760}
761
762fn fa_v2_on() -> bool {
772 std::env::var("MEMRA_FA_V2")
778 .map(|v| v != "0")
779 .unwrap_or(true)
780}
781
782fn fa_v3_on() -> bool {
790 std::env::var("MEMRA_FA_V3")
794 .map(|v| v != "0")
795 .unwrap_or(true)
796}
797
798fn fa_v4_mode() -> &'static str {
803 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
804 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
805}
806fn fa_v4_on() -> bool {
807 fa_v4_mode() != "0"
808} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
817 std::sync::atomic::AtomicUsize::new(1024);
818pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
819 std::sync::atomic::AtomicUsize::new(usize::MAX);
820pub fn fa_v4_at_pub(t_kv: usize) -> bool {
821 fa_v4_at(t_kv)
822}
823fn fa_v4_at(t_kv: usize) -> bool {
824 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
825 let mx = *M.get_or_init(|| {
826 std::env::var("MEMRA_FA_V4_MAX")
827 .ok()
828 .and_then(|v| v.parse().ok())
829 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
830 });
831 fa_v4_on() && t_kv < mx
832}
833pub const FA_DEEP_MIN_DEFAULT: usize = 0;
847fn fa_deep_at(t_kv: usize) -> bool {
848 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
849 return false;
850 }
851 let min = std::env::var("MEMRA_FA_DEEP_MIN")
852 .ok()
853 .and_then(|v| v.parse().ok())
854 .unwrap_or(FA_DEEP_MIN_DEFAULT);
855 t_kv >= min
856}
857pub fn fa_deep_at_pub(t_kv: usize) -> bool {
859 fa_deep_at(t_kv)
860}
861
862fn fa_v3_active(head_dim: usize) -> bool {
863 fa_v3_on()
866 && head_dim % 128 == 0
867 && kv_cache_formats() == ("q8_0", "q5_1")
868 && !Engine::kv_fp8_on()
869}
870
871pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
879 std::env::var("MEMRA_NO_FA_VEC").is_err()
880 && t_kv >= fa_vec_min_tkv()
881 && head_dim == 256
882 && fa_v4_at(t_kv)
883 && !matches!(fa_v4_mode(), "noB3" | "stage")
884 && !Engine::kv_fp8_on()
885}
886pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
888 fa_split_keys(t_kv, n_head_kv)
889}
890
891struct PinnedStage {
896 ptr: *mut u8,
897 cap: usize,
898}
899unsafe impl Send for PinnedStage {}
900impl PinnedStage {
901 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
902 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
903 Ok(PinnedStage { ptr, cap })
904 }
905}
906impl Drop for PinnedStage {
907 fn drop(&mut self) {
908 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
909 }
910}
911
912pub const ARGMAX_NB: usize = 256;
915
916pub(crate) use memra_fa3_vl as fa3_vl_raw;
918
919unsafe extern "C" {
920 fn memra_fa3_prefill(
922 q16: *const core::ffi::c_void,
923 k16: *const core::ffi::c_void,
924 v16: *const core::ffi::c_void,
925 o: *mut f32,
926 t: i32,
927 h: i32,
928 hkv: i32,
929 d: i32,
930 scale: f32,
931 stream: *mut core::ffi::c_void,
932 ) -> i32;
933 pub(crate) fn memra_fa3_vl(
935 q16s: *const *const core::ffi::c_void,
936 k16s: *const *const core::ffi::c_void,
937 v16s: *const *const core::ffi::c_void,
938 os: *const *mut f32,
939 ts: *const i32,
940 b: i32,
941 h: i32,
942 hkv: i32,
943 d: i32,
944 scale: f32,
945 stream: *mut core::ffi::c_void,
946 ) -> i32;
947}
948
949#[repr(C)]
954#[derive(Clone, Copy)]
955pub struct WPtr8(pub [u64; 8]);
956unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
957
958#[repr(C)]
963#[derive(Clone, Copy, Default)]
964pub struct GdnSeqVl {
965 pub kb16: u64,
966 pub gcum: u64,
967 pub beta: u64,
968 pub u: u64,
969 pub wb16: u64,
970 pub y: u64,
971 pub ssnap: u64,
972 pub state_in: u64,
973 pub state_out: u64,
974 pub q: u64,
975 pub p: u64,
976 pub o: u64,
977 pub k: u64,
978 pub v: u64,
979 pub g: u64,
980 pub a: u64,
981 pub w: u64,
982 pub t: i32,
983 pub nc: i32,
984}
985unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
986#[repr(C)]
987#[derive(Clone, Copy)]
988pub struct GdnVl8(pub [GdnSeqVl; 8]);
989unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
990
991#[repr(C)]
994#[derive(Clone, Copy, Default)]
995pub struct GdnWVl {
996 pub qb16: u64,
997 pub pb16: u64,
998}
999unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1000#[repr(C)]
1001#[derive(Clone, Copy)]
1002pub struct GdnWVl8(pub [GdnWVl; 8]);
1003unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1004
1005#[repr(C)]
1007#[derive(Clone, Copy, Default)]
1008pub struct GdnPrepVl {
1009 pub qkv: u64,
1010 pub conv_state: u64,
1011 pub conv_out: u64,
1012 pub q_g: u64,
1013 pub k_g: u64,
1014 pub v_g: u64,
1015 pub q_l2: u64,
1016 pub k_l2: u64,
1017 pub beta_raw: u64,
1018 pub alpha: u64,
1019 pub beta: u64,
1020 pub g_log: u64,
1021 pub o: u64,
1022 pub z: u64,
1023 pub gn: u64,
1024 pub gn16: u64,
1025 pub kb16: u64,
1026 pub qb16: u64,
1027 pub t: i32,
1028 pub pad: i32,
1029}
1030unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1031#[repr(C)]
1032#[derive(Clone, Copy)]
1033pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1034unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1035
1036#[repr(C)]
1038#[derive(Clone, Copy, Default)]
1039pub struct FaSeqVl {
1040 pub q: u64,
1041 pub k16: u64,
1042 pub v16: u64,
1043 pub o: u64,
1044 pub kf: u64,
1045 pub vf: u64,
1046 pub t: i32,
1047 pub pad: i32,
1048}
1049unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1050#[repr(C)]
1051#[derive(Clone, Copy)]
1052pub struct FaVl8(pub [FaSeqVl; 8]);
1053unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1054
1055#[repr(C)]
1057#[derive(Clone, Copy, Default)]
1058pub struct AttnPreVl {
1059 pub qf: u64,
1060 pub kf: u64,
1061 pub vf: u64,
1062 pub q: u64,
1063 pub gate: u64,
1064 pub qn: u64,
1065 pub kn: u64,
1066 pub kc: u64,
1067 pub vc: u64,
1068 pub t: i32,
1069 pub pad: i32,
1070}
1071unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1072#[repr(C)]
1073#[derive(Clone, Copy)]
1074pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1075unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1076
1077pub struct GdnChunkBufs {
1080 pub gcum: CudaSlice<f32>,
1081 pub a: CudaSlice<f32>,
1082 pub p: CudaSlice<f32>,
1083 pub u: CudaSlice<f32>,
1084 pub w: CudaSlice<f32>,
1085 pub kb16: CudaSlice<u8>,
1086 pub wb16: CudaSlice<u8>,
1087 pub y16: CudaSlice<u8>,
1088 pub ssnap16: CudaSlice<u8>,
1089 pub qb16: CudaSlice<u8>,
1090 pub pb16: CudaSlice<u8>,
1091 pub o: CudaSlice<f32>,
1092 pub t: usize,
1093 pub nc: usize,
1094}
1095
1096#[repr(C)]
1098#[derive(Clone, Copy)]
1099pub struct F32x8(pub [f32; 8]);
1100unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1101
1102pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1106
1107impl Engine {
1108 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1109 let gpu = memra_runtime::Gpu::new(ordinal)?;
1110 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1114 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1115 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1116 .and_then(|d| unsafe {
1117 Ok((
1118 cudarc::driver::result::device::get_attribute(
1119 d,
1120 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1121 )?,
1122 cudarc::driver::result::device::get_attribute(
1123 d,
1124 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1125 )?,
1126 ))
1127 })
1128 .unwrap_or((0, 0));
1129 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1130 let ok = matches!(
1131 (built, maj, min),
1132 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1133 );
1134 if !ok {
1135 return Err(format!(
1136 "memra was built for sm_{built} but device {ordinal} reports compute \
1137 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1138 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1139 )
1140 .into());
1141 }
1142 }
1143 unsafe {
1148 use cudarc::driver::sys;
1149 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1150 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1151 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1152 let mut thresh: u64 = u64::MAX;
1153 let _ = sys::cuMemPoolSetAttribute(
1154 pool,
1155 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1156 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1157 );
1158 }
1159 }
1160 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1161 let hybrid = gpu
1162 .ctx
1163 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1164 let qmatvec = gpu
1165 .ctx
1166 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1167 let flash = gpu
1168 .ctx
1169 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1170 let gemm = gpu
1171 .ctx
1172 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1173 let router = gpu
1174 .ctx
1175 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1176 let sample = gpu
1177 .ctx
1178 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1179 let copy_stream = gpu.ctx.new_stream()?;
1180 if std::env::var("MEMRA_EVT")
1196 .map(|v| v == "1")
1197 .unwrap_or(false)
1198 {
1199 } else {
1201 unsafe {
1202 gpu.ctx.disable_event_tracking();
1203 }
1204 }
1205 Ok(Self {
1206 gpu,
1207 module,
1208 hybrid,
1209 qmatvec,
1210 flash,
1211 flash_g: std::sync::OnceLock::new(),
1212 gemm,
1213 router,
1214 sample,
1215 moe_cache: Mutex::new(None),
1216 moe_cache_layout: Mutex::new(None),
1217 copy_stream,
1218 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1219 verify_exact: std::sync::atomic::AtomicBool::new(false),
1220 capture_keep: Mutex::new(Vec::new()),
1221 argmax_partials: Mutex::new(None),
1222 prime_deqw_ws: Mutex::new(None),
1223 router_stage: Mutex::new(None),
1224 fp8_scratch: Mutex::new(None),
1225 fa_vf16_scratch: Mutex::new(None),
1226 fa_part_pool: Mutex::new(None),
1227 fa_part_retired: Mutex::new(Vec::new()),
1228 fn_cache: Mutex::new(Default::default()),
1229 f16_scratch: Mutex::new(None),
1230 #[cfg(memra_cutlass)]
1231 cutlass_scratch: Mutex::new(None),
1232 })
1233 }
1234
1235 pub fn ctx(&self) -> &Arc<CudaContext> {
1236 &self.gpu.ctx
1237 }
1238
1239 pub fn pool_cached_bytes(&self) -> usize {
1257 let (reserved, used) = self.pool_reserved_used();
1258 reserved.saturating_sub(used)
1259 }
1260
1261 pub fn pool_reserved_used(&self) -> (usize, usize) {
1268 use cudarc::driver::sys;
1269 unsafe {
1270 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1271 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1272 != sys::CUresult::CUDA_SUCCESS
1273 {
1274 return (0, 0);
1275 }
1276 let (mut reserved, mut used) = (0u64, 0u64);
1277 if sys::cuMemPoolGetAttribute(
1278 pool,
1279 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1280 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1281 ) != sys::CUresult::CUDA_SUCCESS
1282 {
1283 return (0, 0);
1284 }
1285 if sys::cuMemPoolGetAttribute(
1286 pool,
1287 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1288 &mut used as *mut u64 as *mut core::ffi::c_void,
1289 ) != sys::CUresult::CUDA_SUCCESS
1290 {
1291 return (0, 0);
1292 }
1293 (reserved as usize, used as usize)
1294 }
1295 }
1296
1297 pub fn stream(&self) -> Arc<CudaStream> {
1300 self.gpu.stream()
1301 }
1302 pub fn gkv_on() -> bool {
1305 memra_kv::gkv_on()
1306 }
1307
1308 pub fn wkv_on() -> bool {
1320 memra_kv::wkv_on()
1321 }
1322
1323 pub fn kv_fp8_on() -> bool {
1329 memra_kv::kv_fp8_on()
1330 }
1331
1332 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1335 if head_dim == 512 && Self::gkv_on() {
1336 self.func_g(name)
1337 } else {
1338 self.func(name)
1339 }
1340 }
1341
1342 fn func_g(&self, name: &str) -> CudaFunction {
1346 let m = self.flash_g.get_or_init(|| {
1347 self.gpu
1348 .ctx
1349 .load_module(cudarc::nvrtc::Ptx::from_binary(
1350 FLASH_FATBIN_KF8VF8.to_vec(),
1351 ))
1352 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1353 });
1354 let key = format!("g:{name}");
1355 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1356 return f.clone();
1357 }
1358 let f = match m.load_function(name) {
1359 Ok(f) => f,
1360 Err(_) => self.func(name),
1361 };
1362 self.fn_cache.lock().unwrap().insert(key, f.clone());
1363 f
1364 }
1365
1366 fn func(&self, name: &str) -> CudaFunction {
1367 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1370 return f.clone();
1371 }
1372 let f = self
1373 .module
1374 .load_function(name)
1375 .or_else(|_| self.hybrid.load_function(name))
1376 .or_else(|_| self.qmatvec.load_function(name))
1377 .or_else(|_| self.flash.load_function(name))
1378 .or_else(|_| self.gemm.load_function(name))
1379 .or_else(|_| self.router.load_function(name))
1380 .or_else(|_| self.sample.load_function(name))
1381 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1382 self.fn_cache
1383 .lock()
1384 .unwrap()
1385 .insert(name.to_string(), f.clone());
1386 f
1387 }
1388
1389 pub fn scatter_trim_logits(
1392 &self,
1393 src: &CudaSlice<f32>,
1394 d2t: &CudaSlice<u32>,
1395 dst: &mut CudaSlice<f32>,
1396 d_vocab: usize,
1397 n_vocab: usize,
1398 ) -> Result<(), Box<dyn std::error::Error>> {
1399 let f1 = self.func("scatter_trim_logits_f32");
1400 let f2 = self.func("scatter_trim_logits_pass2_f32");
1401 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1402 let cfg1 = LaunchConfig {
1403 grid_dim: (256, 1, 1),
1404 block_dim: (256, 1, 1),
1405 shared_mem_bytes: 0,
1406 };
1407 let __s_b1 = self.gpu.stream();
1408 let mut b1 = __s_b1.launch_builder(&f1);
1409 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1410 unsafe {
1411 b1.launch(cfg1)?;
1412 }
1413 let cfg2 = LaunchConfig {
1414 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1415 block_dim: (256, 1, 1),
1416 shared_mem_bytes: 0,
1417 };
1418 let __s_b2 = self.gpu.stream();
1419 let mut b2 = __s_b2.launch_builder(&f2);
1420 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1421 unsafe {
1422 b2.launch(cfg2)?;
1423 }
1424 Ok(())
1425 }
1426
1427 #[allow(clippy::too_many_arguments)]
1433 pub fn filter_stats(
1434 &self,
1435 x: &CudaSlice<f32>,
1436 row_stride: usize,
1437 rows: &CudaSlice<i32>,
1438 out_th: &mut CudaSlice<f32>,
1439 out_z: &mut CudaSlice<f32>,
1440 out_max: &mut CudaSlice<f32>,
1441 n: usize,
1442 nrow: usize,
1443 temp: f32,
1444 top_k: i32,
1445 top_p: f32,
1446 min_p: f32,
1447 ) -> Result<(), Box<dyn std::error::Error>> {
1448 let f = self.func("filter_stats_f32");
1455 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1456 let cfg = LaunchConfig {
1457 grid_dim: (nrow as u32, 1, 1),
1458 block_dim: (1024, 1, 1),
1459 shared_mem_bytes: 0,
1460 };
1461 let __s_b = self.gpu.stream();
1462 let mut b = __s_b.launch_builder(&f);
1463 b.arg(x)
1464 .arg(&rs)
1465 .arg(rows)
1466 .arg(&mut *out_th)
1467 .arg(&mut *out_z)
1468 .arg(&mut *out_max)
1469 .arg(&ni)
1470 .arg(&nr)
1471 .arg(&temp)
1472 .arg(&top_k)
1473 .arg(&top_p)
1474 .arg(&min_p);
1475 unsafe {
1476 b.launch(cfg)?;
1477 }
1478 Ok(())
1479 }
1480
1481 #[allow(clippy::too_many_arguments)]
1483 pub fn softmax_gather_filtered(
1484 &self,
1485 x: &CudaSlice<f32>,
1486 row_stride: usize,
1487 ids: &CudaSlice<u32>,
1488 rows: &CudaSlice<i32>,
1489 th: &CudaSlice<f32>,
1490 z: &CudaSlice<f32>,
1491 out: &mut CudaSlice<f32>,
1492 n: usize,
1493 npair: usize,
1494 temp: f32,
1495 ) -> Result<(), Box<dyn std::error::Error>> {
1496 let f = self.func("softmax_gather_filtered_f32");
1497 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1498 let cfg = LaunchConfig {
1499 grid_dim: (npair as u32, 1, 1),
1500 block_dim: (256, 1, 1),
1501 shared_mem_bytes: 0,
1502 };
1503 let __s_b = self.gpu.stream();
1504 let mut b = __s_b.launch_builder(&f);
1505 b.arg(x)
1506 .arg(&rs)
1507 .arg(ids)
1508 .arg(rows)
1509 .arg(th)
1510 .arg(z)
1511 .arg(&mut *out)
1512 .arg(&ni)
1513 .arg(&np)
1514 .arg(&temp);
1515 unsafe {
1516 b.launch(cfg)?;
1517 }
1518 Ok(())
1519 }
1520
1521 #[allow(clippy::too_many_arguments)]
1523 pub fn residual_sample_filtered(
1524 &self,
1525 p: &CudaSlice<f32>,
1526 q: Option<&CudaSlice<f32>>,
1527 n: usize,
1528 temp: f32,
1529 seed: u64,
1530 stream_pos: u32,
1531 p_stats: (f32, f32, f32),
1532 q_stats: (f32, f32, f32),
1533 out_tok: &mut CudaSlice<u32>,
1534 ) -> Result<(), Box<dyn std::error::Error>> {
1535 let f = self.func("residual_sample_filtered_f32");
1536 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1537 let has_q: i32 = q.is_some() as i32;
1538 let qbuf = q.unwrap_or(p);
1539 let (pm, pth, pz) = p_stats;
1540 let (qm, qth, qz) = q_stats;
1541 let cfg = LaunchConfig {
1542 grid_dim: (1, 1, 1),
1543 block_dim: (1024, 1, 1),
1544 shared_mem_bytes: 0,
1545 };
1546 let __s_b = self.gpu.stream();
1547 let mut b = __s_b.launch_builder(&f);
1548 b.arg(p)
1549 .arg(qbuf)
1550 .arg(&has_q)
1551 .arg(&ni)
1552 .arg(&temp)
1553 .arg(&slo)
1554 .arg(&shi)
1555 .arg(&stream_pos)
1556 .arg(&pm)
1557 .arg(&pth)
1558 .arg(&pz)
1559 .arg(&qm)
1560 .arg(&qth)
1561 .arg(&qz)
1562 .arg(&mut *out_tok);
1563 unsafe {
1564 b.launch(cfg)?;
1565 }
1566 Ok(())
1567 }
1568
1569 #[allow(clippy::too_many_arguments)]
1571 pub fn gumbel_perturb_filtered(
1572 &self,
1573 x: &CudaSlice<f32>,
1574 y: &mut CudaSlice<f32>,
1575 n: usize,
1576 seed: u64,
1577 stream_pos: u32,
1578 temp: f32,
1579 row_max: f32,
1580 th: f32,
1581 ) -> Result<(), Box<dyn std::error::Error>> {
1582 let f = self.func("gumbel_perturb_filtered_f32");
1583 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1584 let cfg = LaunchConfig {
1585 grid_dim: (n.div_ceil(256) as u32, 1, 1),
1586 block_dim: (256, 1, 1),
1587 shared_mem_bytes: 0,
1588 };
1589 let __s_b = self.gpu.stream();
1590 let mut b = __s_b.launch_builder(&f);
1591 b.arg(x)
1592 .arg(&mut *y)
1593 .arg(&ni)
1594 .arg(&slo)
1595 .arg(&shi)
1596 .arg(&stream_pos)
1597 .arg(&temp)
1598 .arg(&row_max)
1599 .arg(&th);
1600 unsafe {
1601 b.launch(cfg)?;
1602 }
1603 Ok(())
1604 }
1605
1606 #[allow(clippy::too_many_arguments)]
1610 pub fn penalize_logits(
1611 &self,
1612 x: &mut CudaSlice<f32>,
1613 hist: &CudaSlice<u32>,
1614 n_hist: usize,
1615 rep: f32,
1616 freq: f32,
1617 present: f32,
1618 n: usize,
1619 ) -> Result<(), Box<dyn std::error::Error>> {
1620 if n_hist == 0 {
1621 return Ok(());
1622 }
1623 let f = self.func("penalize_logits_f32");
1624 let (nh, ni) = (n_hist as i32, n as i32);
1625 let cfg = LaunchConfig {
1626 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1627 block_dim: (128, 1, 1),
1628 shared_mem_bytes: 0,
1629 };
1630 let __s_b = self.gpu.stream();
1631 let mut b = __s_b.launch_builder(&f);
1632 b.arg(&mut *x)
1633 .arg(hist)
1634 .arg(&nh)
1635 .arg(&rep)
1636 .arg(&freq)
1637 .arg(&present)
1638 .arg(&ni);
1639 unsafe {
1640 b.launch(cfg)?;
1641 }
1642 Ok(())
1643 }
1644
1645 #[allow(clippy::too_many_arguments)]
1647 pub fn penalize_logits_rows(
1648 &self,
1649 x: &mut CudaSlice<f32>,
1650 hist: &CudaSlice<u32>,
1651 n_hist: usize,
1652 rep: f32,
1653 freq: f32,
1654 present: f32,
1655 n: usize,
1656 nrow: usize,
1657 ) -> Result<(), Box<dyn std::error::Error>> {
1658 if n_hist == 0 || nrow == 0 {
1659 return Ok(());
1660 }
1661 let f = self.func("penalize_logits_rows_f32");
1662 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1663 let cfg = LaunchConfig {
1664 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1665 block_dim: (128, 1, 1),
1666 shared_mem_bytes: 0,
1667 };
1668 let __s_b = self.gpu.stream();
1669 let mut b = __s_b.launch_builder(&f);
1670 b.arg(&mut *x)
1671 .arg(hist)
1672 .arg(&nh)
1673 .arg(&rep)
1674 .arg(&freq)
1675 .arg(&present)
1676 .arg(&ni)
1677 .arg(&nr);
1678 unsafe {
1679 b.launch(cfg)?;
1680 }
1681 Ok(())
1682 }
1683
1684 pub fn wpf_level() -> u32 {
1692 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1693 *ON.get_or_init(|| {
1694 std::env::var("MEMRA_WPF")
1695 .ok()
1696 .and_then(|v| v.parse().ok())
1697 .unwrap_or(1)
1698 })
1699 }
1700
1701 pub fn set_verify_exact(&self, on: bool) {
1713 self.verify_exact
1714 .store(on, std::sync::atomic::Ordering::Relaxed);
1715 }
1716 pub(crate) fn verify_exact_on(&self) -> bool {
1717 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1718 }
1719
1720 pub fn qkv_append_on() -> bool {
1723 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1724 *ON.get_or_init(|| {
1725 std::env::var("MEMRA_QKV_APPEND")
1726 .map(|v| v != "0")
1727 .unwrap_or(true)
1728 })
1729 }
1730
1731 pub fn pdl_wb_on() -> bool {
1734 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1735 *ON.get_or_init(|| {
1736 std::env::var("MEMRA_PDL_WB")
1737 .map(|v| v != "0")
1738 .unwrap_or(true)
1739 })
1740 }
1741
1742 pub fn pdl_mmvq_on() -> bool {
1746 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1747 *ON.get_or_init(|| {
1748 std::env::var("MEMRA_PDL_MMVQ")
1749 .map(|v| v != "0")
1750 .unwrap_or(true)
1751 })
1752 }
1753
1754 pub fn pdl_on() -> bool {
1755 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1756 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1757 }
1758
1759 pub fn pdl_nvfp4q8_on() -> bool {
1765 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1766 *ON.get_or_init(|| {
1767 std::env::var("MEMRA_PDL_NVFP4")
1768 .map(|v| v != "0")
1769 .unwrap_or(true)
1770 })
1771 }
1772
1773 fn q40_mr1_on() -> bool {
1779 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1780 match *Q40MR.get_or_init(|| {
1781 std::env::var("MEMRA_Q40_MR")
1782 .ok()
1783 .and_then(|v| v.parse().ok())
1784 }) {
1785 Some(v) => v == 1,
1786 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1787 }
1788 }
1789
1790 fn pdl_func_flash(
1795 &self,
1796 g: bool,
1797 name: &'static str,
1798 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1799 use cudarc::driver::sys as cu;
1800 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1807 std::sync::Mutex::new(None);
1808 static FNS: std::sync::Mutex<
1809 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1810 > = std::sync::Mutex::new(None);
1811 let ctx_key = self.ctx().cu_ctx() as usize;
1812 if let Some(&f) = FNS
1813 .lock()
1814 .unwrap()
1815 .get_or_insert_with(Default::default)
1816 .get(&(ctx_key, g, name))
1817 {
1818 return Ok(f as cu::CUfunction);
1819 }
1820 let module = {
1821 let mut mods = MODS.lock().unwrap();
1822 let map = mods.get_or_insert_with(Default::default);
1823 match map.get(&(ctx_key, g)) {
1824 Some(&m) => m,
1825 None => {
1826 let m = self.pdl_load_module_in_ctx(if g {
1827 FLASH_FATBIN_KF8VF8
1828 } else {
1829 FLASH_FATBIN
1830 })?;
1831 map.insert((ctx_key, g), m);
1832 m
1833 }
1834 }
1835 };
1836 let cname = std::ffi::CString::new(name)?;
1837 let mut f: cu::CUfunction = std::ptr::null_mut();
1838 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1839 if r != cu::CUresult::CUDA_SUCCESS {
1840 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
1841 }
1842 FNS.lock()
1843 .unwrap()
1844 .get_or_insert_with(Default::default)
1845 .insert((ctx_key, g, name), f as usize);
1846 Ok(f)
1847 }
1848
1849 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1854 use cudarc::driver::sys as cu;
1855 let mut prev: cu::CUcontext = std::ptr::null_mut();
1856 unsafe {
1857 cu::cuCtxGetCurrent(&mut prev).result()?;
1858 }
1859 self.ctx().bind_to_thread()?;
1860 let mut m: cu::CUmodule = std::ptr::null_mut();
1861 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1862 let restore = if prev.is_null() {
1863 cu::CUresult::CUDA_SUCCESS
1864 } else {
1865 unsafe { cu::cuCtxSetCurrent(prev) }
1866 };
1867 if r != cu::CUresult::CUDA_SUCCESS {
1868 return Err(format!("pdl module load: {r:?}").into());
1869 }
1870 if restore != cu::CUresult::CUDA_SUCCESS {
1871 return Err(format!("pdl module load: ctx restore {restore:?}").into());
1872 }
1873 Ok(m as usize)
1874 }
1875
1876 fn pdl_func(
1877 &self,
1878 name: &'static str,
1879 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1880 use cudarc::driver::sys as cu;
1881 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1884 std::sync::Mutex::new(None);
1885 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1888 std::sync::Mutex::new(None);
1889 static FNS: std::sync::Mutex<
1890 Option<std::collections::HashMap<(usize, &'static str), usize>>,
1891 > = std::sync::Mutex::new(None);
1892 let ctx_key = self.ctx().cu_ctx() as usize;
1893 if let Some(&f) = FNS
1894 .lock()
1895 .unwrap()
1896 .get_or_insert_with(Default::default)
1897 .get(&(ctx_key, name))
1898 {
1899 return Ok(f as cu::CUfunction);
1900 }
1901 let module = {
1902 let mut mods = MODULES.lock().unwrap();
1903 let map = mods.get_or_insert_with(Default::default);
1904 match map.get(&ctx_key) {
1905 Some(&m) => m,
1906 None => {
1907 let m = self.pdl_load_module_in_ctx(FATBIN)?;
1908 map.insert(ctx_key, m);
1909 m
1910 }
1911 }
1912 };
1913 let cname = std::ffi::CString::new(name)?;
1914 let mut f: cu::CUfunction = std::ptr::null_mut();
1915 let mut r =
1916 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1917 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
1918 let qmodule = {
1919 let mut mods = QMODULES.lock().unwrap();
1920 let map = mods.get_or_insert_with(Default::default);
1921 match map.get(&ctx_key) {
1922 Some(&m) => m,
1923 None => {
1924 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
1925 map.insert(ctx_key, m);
1926 m
1927 }
1928 }
1929 };
1930 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
1931 }
1932 if r != cu::CUresult::CUDA_SUCCESS {
1933 return Err(format!("pdl_func {name}: {r:?}").into());
1934 }
1935 FNS.lock()
1936 .unwrap()
1937 .get_or_insert_with(Default::default)
1938 .insert((ctx_key, name), f as usize);
1939 Ok(f)
1940 }
1941
1942 unsafe fn launch_pdl_flash(
1954 &self,
1955 g: bool,
1956 name: &'static str,
1957 grid: (u32, u32, u32),
1958 block: (u32, u32, u32),
1959 smem: u32,
1960 params: &mut [*mut std::ffi::c_void],
1961 ) -> Result<(), Box<dyn std::error::Error>> {
1962 use cudarc::driver::sys as cu;
1963 let f = self.pdl_func_flash(g, name)?;
1964 if smem > 0 {
1965 let r =
1967 unsafe {
1968 cu::cuFuncSetAttribute(f,
1969 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
1970 smem as i32)
1971 };
1972 if r != cu::CUresult::CUDA_SUCCESS {
1973 return Err(format!("pdl smem attr {name}: {r:?}").into());
1974 }
1975 }
1976 let mut attr = cu::CUlaunchAttribute {
1977 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1978 pad: [0; 4],
1979 value: cu::CUlaunchAttributeValue {
1980 programmaticStreamSerializationAllowed: 1,
1981 },
1982 };
1983 let cfg = cu::CUlaunchConfig {
1984 gridDimX: grid.0,
1985 gridDimY: grid.1,
1986 gridDimZ: grid.2,
1987 blockDimX: block.0,
1988 blockDimY: block.1,
1989 blockDimZ: block.2,
1990 sharedMemBytes: smem,
1991 hStream: self.gpu.stream().cu_stream(),
1992 attrs: &mut attr,
1993 numAttrs: 1,
1994 };
1995 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1996 if r != cu::CUresult::CUDA_SUCCESS {
1997 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
1998 }
1999 Ok(())
2000 }
2001
2002 unsafe fn launch_pdl(
2003 &self,
2004 name: &'static str,
2005 grid: (u32, u32, u32),
2006 block: (u32, u32, u32),
2007 params: &mut [*mut std::ffi::c_void],
2008 ) -> Result<(), Box<dyn std::error::Error>> {
2009 use cudarc::driver::sys as cu;
2010 let f = self.pdl_func(name)?;
2011 let mut attr = cu::CUlaunchAttribute {
2012 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2013 pad: [0; 4],
2014 value: cu::CUlaunchAttributeValue {
2015 programmaticStreamSerializationAllowed: 1,
2016 },
2017 };
2018 let cfg = cu::CUlaunchConfig {
2019 gridDimX: grid.0,
2020 gridDimY: grid.1,
2021 gridDimZ: grid.2,
2022 blockDimX: block.0,
2023 blockDimY: block.1,
2024 blockDimZ: block.2,
2025 sharedMemBytes: 0,
2026 hStream: self.gpu.stream().cu_stream(),
2027 attrs: &mut attr,
2028 numAttrs: 1,
2029 };
2030 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2031 if r != cu::CUresult::CUDA_SUCCESS {
2032 return Err(format!("launch_pdl {name}: {r:?}").into());
2033 }
2034 Ok(())
2035 }
2036
2037 pub fn prefetch_weight_l2(
2040 &self,
2041 w: &crate::model::GpuTensor,
2042 ) -> Result<(), Box<dyn std::error::Error>> {
2043 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2044 let p = rp4.as_ref().unwrap_or(bytes);
2045 self.prefetch_l2(p, p.len())?;
2046 }
2047 Ok(())
2048 }
2049
2050 pub fn gather_row_bf16(
2053 &self,
2054 table: &CudaSlice<u8>,
2055 tok: &CudaSlice<u32>,
2056 idx: usize,
2057 dst: &mut CudaSlice<f32>,
2058 ncols: usize,
2059 ) -> Result<(), Box<dyn std::error::Error>> {
2060 let f = self.func("gather_row_bf16_f32");
2061 let cfg = LaunchConfig {
2062 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2063 block_dim: (256, 1, 1),
2064 shared_mem_bytes: 0,
2065 };
2066 let (nc, ix) = (ncols as i32, idx as i32);
2067 let __s_b = self.gpu.stream();
2068 let mut b = __s_b.launch_builder(&f);
2069 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2070 unsafe {
2071 b.launch(cfg)?;
2072 }
2073 Ok(())
2074 }
2075
2076 pub fn add_row_inplace(
2078 &self,
2079 logits: &mut CudaSlice<f32>,
2080 bias: &CudaSlice<f32>,
2081 n: usize,
2082 row_off: usize,
2083 ) -> Result<(), Box<dyn std::error::Error>> {
2084 let f = self.func("add_row_inplace_f32");
2085 let cfg = LaunchConfig {
2086 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2087 block_dim: (256, 1, 1),
2088 shared_mem_bytes: 0,
2089 };
2090 let (ni, off) = (n as i32, row_off as i64);
2091 let __s_b = self.gpu.stream();
2092 let mut b = __s_b.launch_builder(&f);
2093 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2094 unsafe {
2095 b.launch(cfg)?;
2096 }
2097 Ok(())
2098 }
2099
2100 pub fn prefetch_l2(
2102 &self,
2103 p: &CudaSlice<u8>,
2104 n: usize,
2105 ) -> Result<(), Box<dyn std::error::Error>> {
2106 let f = self.func("prefetch_l2_bytes");
2107 let lines = n.div_ceil(128);
2108 let ni = n as i64;
2109 let cfg = LaunchConfig {
2110 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2111 block_dim: (256, 1, 1),
2112 shared_mem_bytes: 0,
2113 };
2114 let __s_b = self.gpu.stream();
2115 let mut b = __s_b.launch_builder(&f);
2116 b.arg(p).arg(&ni);
2117 unsafe {
2118 b.launch(cfg)?;
2119 }
2120 Ok(())
2121 }
2122
2123 pub fn router_gemv(
2126 &self,
2127 w: &CudaSlice<f32>,
2128 x: &CudaSlice<f32>,
2129 n_embd: usize,
2130 n_experts: usize,
2131 t: usize,
2132 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2133 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2139 Ok("0") => false,
2140 Ok(_) => true,
2141 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2142 };
2143 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2153 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2154 }
2155
2156 pub fn router_gemv_form(
2159 &self,
2160 w: &CudaSlice<f32>,
2161 x: &CudaSlice<f32>,
2162 n_embd: usize,
2163 n_experts: usize,
2164 t: usize,
2165 w8: bool,
2166 batch: bool,
2167 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2168 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2169 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2170 let f = if batch {
2171 self.func("router_gemv_f32_w8_batch")
2172 } else if w8 {
2173 self.func("router_gemv_f32_w8")
2174 } else {
2175 self.func("router_gemv_f32")
2176 };
2177 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2178 let cfg = if batch {
2179 LaunchConfig {
2180 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2181 block_dim: (32, 8, 1),
2182 shared_mem_bytes: 0,
2183 }
2184 } else {
2185 LaunchConfig {
2186 grid_dim: (n_experts as u32, t as u32, 1),
2187 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2188 shared_mem_bytes: 0,
2189 }
2190 };
2191 let __s_b = self.gpu.stream();
2192 let mut b = __s_b.launch_builder(&f);
2193 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2194 unsafe {
2195 b.launch(cfg)?;
2196 }
2197 Ok(y)
2198 }
2199
2200 pub fn rows_permute(
2202 &self,
2203 src: &CudaSlice<f32>,
2204 idx: &CudaSlice<i32>,
2205 nrows: usize,
2206 ncols: usize,
2207 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2208 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2209 let f = self.func("rows_permute_f32");
2210 let (nc, nr) = (ncols as i32, nrows as i32);
2211 let cfg = LaunchConfig {
2212 grid_dim: (nrows as u32, 1, 1),
2213 block_dim: (256, 1, 1),
2214 shared_mem_bytes: 0,
2215 };
2216 let __s_b = self.gpu.stream();
2217 let mut b = __s_b.launch_builder(&f);
2218 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2219 unsafe {
2220 b.launch(cfg)?;
2221 }
2222 Ok(dst)
2223 }
2224
2225 pub fn sigmoid_dot_rows(
2230 &self,
2231 x: &CudaSlice<f32>,
2232 w: &CudaSlice<f32>,
2233 n_embd: usize,
2234 t: usize,
2235 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2236 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2239 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2240 let gs = self.linear(x, w, t, n_embd, 1)?;
2241 let mut g = self.uninit(t)?;
2242 self.sigmoid(&gs, &mut g, t)?;
2243 return Ok(g);
2244 }
2245 let mut g = self.alloc_uninit::<f32>(t)?;
2251 let f = self.func("sigmoid_dot_rows_f32");
2252 let (ne, ti) = (n_embd as i32, t as i32);
2253 let cfg = LaunchConfig {
2254 grid_dim: (t as u32, 1, 1),
2255 block_dim: (32, 8, 1),
2256 shared_mem_bytes: 0,
2257 };
2258 let __s_b = self.gpu.stream();
2259 let mut b = __s_b.launch_builder(&f);
2260 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2261 unsafe {
2262 b.launch(cfg)?;
2263 }
2264 Ok(g)
2265 }
2266
2267 pub fn spec_rollback_stream(
2269 &self,
2270 len_ptrs: &CudaSlice<u64>,
2271 pos_start: &CudaSlice<i32>,
2272 acc: &CudaSlice<u32>,
2273 base: usize,
2274 n_rows: usize,
2275 ) -> Result<(), Box<dyn std::error::Error>> {
2276 let f = self.func("spec_rollback_stream");
2277 let (b, nr) = (base as i32, n_rows as i32);
2278 let cfg = LaunchConfig {
2279 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2280 block_dim: (64, 1, 1),
2281 shared_mem_bytes: 0,
2282 };
2283 let __s_bl = self.gpu.stream();
2284 let mut bl = __s_bl.launch_builder(&f);
2285 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2286 unsafe {
2287 bl.launch(cfg)?;
2288 }
2289 Ok(())
2290 }
2291
2292 pub fn plain_tok_ring(
2294 &self,
2295 vam: &CudaSlice<u32>,
2296 pos_start: &CudaSlice<i32>,
2297 base: usize,
2298 ring: &mut CudaSlice<u32>,
2299 ) -> Result<(), Box<dyn std::error::Error>> {
2300 let f = self.func("plain_tok_ring");
2301 let (b, cap) = (base as i32, ring.len() as i32);
2302 let cfg = LaunchConfig {
2303 grid_dim: (1, 1, 1),
2304 block_dim: (32, 1, 1),
2305 shared_mem_bytes: 0,
2306 };
2307 let __s_bl = self.gpu.stream();
2308 let mut bl = __s_bl.launch_builder(&f);
2309 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2310 unsafe {
2311 bl.launch(cfg)?;
2312 }
2313 Ok(())
2314 }
2315
2316 pub fn spec_ring_commit(
2318 &self,
2319 vtok: &CudaSlice<u32>,
2320 acc: &CudaSlice<u32>,
2321 brk: &CudaSlice<u32>,
2322 ring: &mut CudaSlice<u32>,
2323 pend: &mut CudaSlice<u32>,
2324 ) -> Result<(), Box<dyn std::error::Error>> {
2325 let f = self.func("spec_ring_commit");
2326 let cfg = LaunchConfig {
2327 grid_dim: (1, 1, 1),
2328 block_dim: (32, 1, 1),
2329 shared_mem_bytes: 0,
2330 };
2331 let __s_b = self.gpu.stream();
2332 let mut b = __s_b.launch_builder(&f);
2333 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2334 unsafe {
2335 b.launch(cfg)?;
2336 }
2337 Ok(())
2338 }
2339 pub fn i32_copy_add(
2340 &self,
2341 src: &CudaSlice<i32>,
2342 dst: &mut CudaSlice<i32>,
2343 delta: i32,
2344 ) -> Result<(), Box<dyn std::error::Error>> {
2345 let f = self.func("i32_copy_add");
2346 let cfg = LaunchConfig {
2347 grid_dim: (1, 1, 1),
2348 block_dim: (32, 1, 1),
2349 shared_mem_bytes: 0,
2350 };
2351 let __s_b = self.gpu.stream();
2352 let mut b = __s_b.launch_builder(&f);
2353 b.arg(src).arg(dst).arg(&delta);
2354 unsafe {
2355 b.launch(cfg)?;
2356 }
2357 Ok(())
2358 }
2359 pub fn u32_copy(
2360 &self,
2361 src: &CudaSlice<u32>,
2362 dst: &mut CudaSlice<u32>,
2363 ) -> Result<(), Box<dyn std::error::Error>> {
2364 let f = self.func("u32_copy");
2365 let cfg = LaunchConfig {
2366 grid_dim: (1, 1, 1),
2367 block_dim: (32, 1, 1),
2368 shared_mem_bytes: 0,
2369 };
2370 let __s_b = self.gpu.stream();
2371 let mut b = __s_b.launch_builder(&f);
2372 b.arg(src).arg(dst);
2373 unsafe {
2374 b.launch(cfg)?;
2375 }
2376 Ok(())
2377 }
2378
2379 pub fn spec_adapt_k(
2383 &self,
2384 acc: &CudaSlice<u32>,
2385 brk: &mut CudaSlice<u32>,
2386 floor: usize,
2387 cap: usize,
2388 ) -> Result<(), Box<dyn std::error::Error>> {
2389 let f = self.func("spec_adapt_k");
2390 let (fl, cp) = (floor as i32, cap as i32);
2391 let cfg = LaunchConfig {
2392 grid_dim: (1, 1, 1),
2393 block_dim: (32, 1, 1),
2394 shared_mem_bytes: 0,
2395 };
2396 let __s_b = self.gpu.stream();
2397 let mut b = __s_b.launch_builder(&f);
2398 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2399 unsafe {
2400 b.launch(cfg)?;
2401 }
2402 Ok(())
2403 }
2404
2405 pub fn spec_accept_greedy_dc(
2407 &self,
2408 preds: &CudaSlice<u32>,
2409 vtok: &CudaSlice<u32>,
2410 last_pred: &CudaSlice<u32>,
2411 brk: &CudaSlice<u32>,
2412 out: &mut CudaSlice<u32>,
2413 ) -> Result<(), Box<dyn std::error::Error>> {
2414 let f = self.func("spec_accept_greedy_dc");
2415 let cfg = LaunchConfig {
2416 grid_dim: (1, 1, 1),
2417 block_dim: (32, 1, 1),
2418 shared_mem_bytes: 0,
2419 };
2420 let __s_b = self.gpu.stream();
2421 let mut b = __s_b.launch_builder(&f);
2422 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2423 unsafe {
2424 b.launch(cfg)?;
2425 }
2426 Ok(())
2427 }
2428
2429 pub fn pos_iota(
2431 &self,
2432 pos0: &CudaSlice<i32>,
2433 out: &mut CudaSlice<i32>,
2434 t: usize,
2435 ) -> Result<(), Box<dyn std::error::Error>> {
2436 let f = self.func("pos_iota_i32");
2437 let ti = t as i32;
2438 let cfg = LaunchConfig {
2439 grid_dim: (1, 1, 1),
2440 block_dim: (t.max(1) as u32, 1, 1),
2441 shared_mem_bytes: 0,
2442 };
2443 let __s_b = self.gpu.stream();
2444 let mut b = __s_b.launch_builder(&f);
2445 b.arg(pos0).arg(out).arg(&ti);
2446 unsafe {
2447 b.launch(cfg)?;
2448 }
2449 Ok(())
2450 }
2451 #[allow(clippy::too_many_arguments)]
2452 pub fn append_kv_quantized_rows_dc(
2453 &self,
2454 k_rows: &CudaSlice<f32>,
2455 v_rows: &CudaSlice<f32>,
2456 kc: &mut CudaSlice<u8>,
2457 vc: &mut CudaSlice<u8>,
2458 t0_dev: &CudaSlice<i32>,
2459 t: usize,
2460 kv_dim_k: usize,
2461 kv_dim_v: usize,
2462 k_tok_bytes: usize,
2463 v_tok_bytes: usize,
2464 g: bool,
2465 ) -> Result<(), Box<dyn std::error::Error>> {
2466 let f = if g {
2467 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2468 } else {
2469 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2470 };
2471 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2472 let cfg = LaunchConfig {
2473 grid_dim: (nblk, t as u32, 1),
2474 block_dim: (32, 1, 1),
2475 shared_mem_bytes: 0,
2476 };
2477 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2478 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2479 let __s_b = self.gpu.stream();
2480 let mut b = __s_b.launch_builder(&f);
2481 b.arg(k_rows)
2482 .arg(v_rows)
2483 .arg(kc)
2484 .arg(vc)
2485 .arg(t0_dev)
2486 .arg(&kdk)
2487 .arg(&kdv)
2488 .arg(&ktb)
2489 .arg(&vtb);
2490 unsafe {
2491 b.launch(cfg)?;
2492 }
2493 Ok(())
2494 }
2495
2496 #[allow(clippy::too_many_arguments)]
2499 pub fn append_kv_quantized_row_dc_inc(
2500 &self,
2501 k_row: &CudaSlice<f32>,
2502 v_row: &CudaSlice<f32>,
2503 kc: &mut CudaSlice<u8>,
2504 vc: &mut CudaSlice<u8>,
2505 t0_dev: &mut CudaSlice<i32>,
2506 kv_dim_k: usize,
2507 kv_dim_v: usize,
2508 k_tok_bytes: usize,
2509 v_tok_bytes: usize,
2510 g: bool,
2511 ) -> Result<(), Box<dyn std::error::Error>> {
2512 let f = if g {
2513 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2514 } else {
2515 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2516 };
2517 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2518 let cfg = LaunchConfig {
2519 grid_dim: (1, 1, 1),
2520 block_dim: (nthreads, 1, 1),
2521 shared_mem_bytes: 0,
2522 };
2523 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2524 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2525 let __s_b = self.gpu.stream();
2526 let mut b = __s_b.launch_builder(&f);
2527 b.arg(k_row)
2528 .arg(v_row)
2529 .arg(kc)
2530 .arg(vc)
2531 .arg(t0_dev)
2532 .arg(&kdk)
2533 .arg(&kdv)
2534 .arg(&ktb)
2535 .arg(&vtb);
2536 unsafe {
2537 b.launch(cfg)?;
2538 }
2539 Ok(())
2540 }
2541
2542 pub fn pack_tok_p(
2544 &self,
2545 tok: &CudaSlice<u32>,
2546 p: &CudaSlice<f32>,
2547 out: &mut CudaSlice<u32>,
2548 slot: usize,
2549 ) -> Result<(), Box<dyn std::error::Error>> {
2550 let f = self.func("pack_tok_p");
2551 let sl = slot as i32;
2552 let cfg = LaunchConfig {
2553 grid_dim: (1, 1, 1),
2554 block_dim: (32, 1, 1),
2555 shared_mem_bytes: 0,
2556 };
2557 let __s_b = self.gpu.stream();
2558 let mut b = __s_b.launch_builder(&f);
2559 b.arg(tok).arg(p).arg(out).arg(&sl);
2560 unsafe {
2561 b.launch(cfg)?;
2562 }
2563 Ok(())
2564 }
2565 pub fn tok_map_u32(
2566 &self,
2567 tok: &mut CudaSlice<u32>,
2568 map: &CudaSlice<u32>,
2569 ) -> Result<(), Box<dyn std::error::Error>> {
2570 let f = self.func("tok_map_u32");
2571 let cfg = LaunchConfig {
2572 grid_dim: (1, 1, 1),
2573 block_dim: (32, 1, 1),
2574 shared_mem_bytes: 0,
2575 };
2576 let __s_b = self.gpu.stream();
2577 let mut b = __s_b.launch_builder(&f);
2578 b.arg(tok).arg(map);
2579 unsafe {
2580 b.launch(cfg)?;
2581 }
2582 Ok(())
2583 }
2584
2585 #[allow(clippy::too_many_arguments)]
2587 pub fn spec_assemble_verify(
2588 &self,
2589 tokp: &CudaSlice<u32>,
2590 pend: &CudaSlice<u32>,
2591 d2t: Option<&CudaSlice<u32>>,
2592 vtok: &mut CudaSlice<u32>,
2593 brk: &mut CudaSlice<u32>,
2594 p_min: f32,
2595 k: usize,
2596 pmin0: bool,
2597 ) -> Result<(), Box<dyn std::error::Error>> {
2598 let f = self.func("spec_assemble_verify");
2599 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2600 let cfg = LaunchConfig {
2601 grid_dim: (1, 1, 1),
2602 block_dim: (32, 1, 1),
2603 shared_mem_bytes: 0,
2604 };
2605 let __s_b = self.gpu.stream();
2606 let mut b = __s_b.launch_builder(&f);
2607 match d2t {
2608 Some(m) => {
2609 b.arg(tokp)
2610 .arg(pend)
2611 .arg(m)
2612 .arg(vtok)
2613 .arg(brk)
2614 .arg(&p_min)
2615 .arg(&ki)
2616 .arg(&pm);
2617 unsafe {
2618 b.launch(cfg)?;
2619 }
2620 }
2621 None => {
2622 let null: u64 = 0;
2623 b.arg(tokp)
2624 .arg(pend)
2625 .arg(&null)
2626 .arg(vtok)
2627 .arg(brk)
2628 .arg(&p_min)
2629 .arg(&ki)
2630 .arg(&pm);
2631 unsafe {
2632 b.launch(cfg)?;
2633 }
2634 }
2635 }
2636 Ok(())
2637 }
2638
2639 #[allow(clippy::too_many_arguments)]
2641 pub fn ssm_conv_ring_rebuild_dc(
2642 &self,
2643 qkv_tm: &CudaSlice<f32>,
2644 ring_old: &CudaSlice<f32>,
2645 conv_state: &mut CudaSlice<f32>,
2646 conv_dim: usize,
2647 acc: &CudaSlice<u32>,
2648 base: usize,
2649 t_v: usize,
2650 d_conv: usize,
2651 ) -> Result<(), Box<dyn std::error::Error>> {
2652 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2653 let n = conv_dim * (d_conv - 1);
2654 let cfg = LaunchConfig::for_num_elems(n as u32);
2655 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2656 let __s_b = self.gpu.stream();
2657 let mut b = __s_b.launch_builder(&f);
2658 b.arg(qkv_tm)
2659 .arg(ring_old)
2660 .arg(conv_state)
2661 .arg(&cd)
2662 .arg(acc)
2663 .arg(&b0)
2664 .arg(&tv)
2665 .arg(&dc);
2666 unsafe {
2667 b.launch(cfg)?;
2668 }
2669 Ok(())
2670 }
2671 #[allow(clippy::too_many_arguments)]
2672 pub fn gdn_scan_s128_dc(
2673 &self,
2674 q: &CudaSlice<f32>,
2675 k: &CudaSlice<f32>,
2676 v: &CudaSlice<f32>,
2677 g: &CudaSlice<f32>,
2678 beta: &CudaSlice<f32>,
2679 state_in: &CudaSlice<f32>,
2680 state_out: &mut CudaSlice<f32>,
2681 o: &mut CudaSlice<f32>,
2682 n_head: usize,
2683 acc: &CudaSlice<u32>,
2684 base: usize,
2685 t_v: usize,
2686 scale: f32,
2687 ) -> Result<(), Box<dyn std::error::Error>> {
2688 let f = self.func("gdn_scan_s128_dc");
2689 const S_V: u32 = 128;
2690 const WARP: u32 = 32;
2691 const COLS_PER_BLOCK: u32 = 4;
2692 let cfg = LaunchConfig {
2693 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2694 block_dim: (WARP, COLS_PER_BLOCK, 1),
2695 shared_mem_bytes: 0,
2696 };
2697 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2698 let __s_b = self.gpu.stream();
2699 let mut b = __s_b.launch_builder(&f);
2700 b.arg(q)
2701 .arg(k)
2702 .arg(v)
2703 .arg(g)
2704 .arg(beta)
2705 .arg(state_in)
2706 .arg(state_out)
2707 .arg(o)
2708 .arg(&h)
2709 .arg(acc)
2710 .arg(&b0)
2711 .arg(&tv)
2712 .arg(&scale);
2713 unsafe {
2714 b.launch(cfg)?;
2715 }
2716 Ok(())
2717 }
2718
2719 pub fn spec_rollback_kv(
2721 &self,
2722 len_ptrs: &CudaSlice<u64>,
2723 saved: &CudaSlice<i32>,
2724 acc: &CudaSlice<u32>,
2725 base: usize,
2726 n_layer: usize,
2727 ) -> Result<(), Box<dyn std::error::Error>> {
2728 let f = self.func("spec_rollback_kv");
2729 let (b, nl) = (base as i32, n_layer as i32);
2730 let cfg = LaunchConfig {
2731 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2732 block_dim: (64, 1, 1),
2733 shared_mem_bytes: 0,
2734 };
2735 let __s_bl = self.gpu.stream();
2736 let mut bl = __s_bl.launch_builder(&f);
2737 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
2738 unsafe {
2739 bl.launch(cfg)?;
2740 }
2741 Ok(())
2742 }
2743
2744 pub fn spec_fork_valid(
2746 &self,
2747 acc: &CudaSlice<u32>,
2748 optimistic_pending: u32,
2749 valid: &mut CudaSlice<u32>,
2750 ) -> Result<(), Box<dyn std::error::Error>> {
2751 let f = self.func("spec_fork_valid");
2752 let cfg = LaunchConfig {
2753 grid_dim: (1, 1, 1),
2754 block_dim: (1, 1, 1),
2755 shared_mem_bytes: 0,
2756 };
2757 let __s_bl = self.gpu.stream();
2758 let mut bl = __s_bl.launch_builder(&f);
2759 bl.arg(acc).arg(&optimistic_pending).arg(valid);
2760 unsafe {
2761 bl.launch(cfg)?;
2762 }
2763 Ok(())
2764 }
2765
2766 pub fn spec_fork_reconcile_kv(
2768 &self,
2769 len_ptrs: &CudaSlice<u64>,
2770 saved: &CudaSlice<i32>,
2771 acc: &CudaSlice<u32>,
2772 valid: &CudaSlice<u32>,
2773 base: usize,
2774 n_layer: usize,
2775 ) -> Result<(), Box<dyn std::error::Error>> {
2776 let f = self.func("spec_fork_reconcile_kv");
2777 let (b, nl) = (base as i32, n_layer as i32);
2778 let cfg = LaunchConfig {
2779 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2780 block_dim: (64, 1, 1),
2781 shared_mem_bytes: 0,
2782 };
2783 let __s_bl = self.gpu.stream();
2784 let mut bl = __s_bl.launch_builder(&f);
2785 bl.arg(len_ptrs)
2786 .arg(saved)
2787 .arg(acc)
2788 .arg(valid)
2789 .arg(&b)
2790 .arg(&nl);
2791 unsafe {
2792 bl.launch(cfg)?;
2793 }
2794 Ok(())
2795 }
2796
2797 pub fn spec_fork_restore_f32(
2799 &self,
2800 snapshot: &CudaSlice<f32>,
2801 state: &mut CudaSlice<f32>,
2802 valid: &CudaSlice<u32>,
2803 ) -> Result<(), Box<dyn std::error::Error>> {
2804 assert_eq!(
2805 snapshot.len(),
2806 state.len(),
2807 "fork recurrent snapshot shape mismatch"
2808 );
2809 let f = self.func("spec_fork_restore_f32");
2810 let n = state.len() as i32;
2811 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
2812 let cfg = LaunchConfig {
2813 grid_dim: (blocks, 1, 1),
2814 block_dim: (256, 1, 1),
2815 shared_mem_bytes: 0,
2816 };
2817 let __s_bl = self.gpu.stream();
2818 let mut bl = __s_bl.launch_builder(&f);
2819 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
2820 unsafe {
2821 bl.launch(cfg)?;
2822 }
2823 Ok(())
2824 }
2825
2826 pub fn spec_seed_gather(
2829 &self,
2830 vx: &CudaSlice<f32>,
2831 fill_prev: &CudaSlice<f32>,
2832 acc: &CudaSlice<u32>,
2833 h_seed: &mut CudaSlice<f32>,
2834 base: usize,
2835 n_embd: usize,
2836 ) -> Result<(), Box<dyn std::error::Error>> {
2837 let f = self.func("spec_seed_gather");
2838 let (b, ne) = (base as i32, n_embd as i32);
2839 let cfg = LaunchConfig {
2840 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
2841 block_dim: (256, 1, 1),
2842 shared_mem_bytes: 0,
2843 };
2844 let __s_bl = self.gpu.stream();
2845 let mut bl = __s_bl.launch_builder(&f);
2846 bl.arg(vx)
2847 .arg(fill_prev)
2848 .arg(acc)
2849 .arg(h_seed)
2850 .arg(&b)
2851 .arg(&ne);
2852 unsafe {
2853 bl.launch(cfg)?;
2854 }
2855 Ok(())
2856 }
2857
2858 pub fn spec_accept_greedy(
2860 &self,
2861 preds: &CudaSlice<u32>,
2862 draft: &CudaSlice<u32>,
2863 last_pred: u32,
2864 base: usize,
2865 k_round: usize,
2866 out: &mut CudaSlice<u32>,
2867 ) -> Result<(), Box<dyn std::error::Error>> {
2868 let f = self.func("spec_accept_greedy");
2869 let (b, k) = (base as i32, k_round as i32);
2870 let cfg = LaunchConfig {
2871 grid_dim: (1, 1, 1),
2872 block_dim: (32, 1, 1),
2873 shared_mem_bytes: 0,
2874 };
2875 let __s_bl = self.gpu.stream();
2876 let mut bl = __s_bl.launch_builder(&f);
2877 bl.arg(preds)
2878 .arg(draft)
2879 .arg(&last_pred)
2880 .arg(&b)
2881 .arg(&k)
2882 .arg(out);
2883 unsafe {
2884 bl.launch(cfg)?;
2885 }
2886 Ok(())
2887 }
2888
2889 pub fn gumbel_perturb(
2896 &self,
2897 x: &CudaSlice<f32>,
2898 y: &mut CudaSlice<f32>,
2899 n: usize,
2900 seed: u64,
2901 stream_pos: u32,
2902 temp: f32,
2903 ) -> Result<(), Box<dyn std::error::Error>> {
2904 let f = self.func("gumbel_perturb_f32");
2905 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2906 let cfg = LaunchConfig {
2907 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2908 block_dim: (256, 1, 1),
2909 shared_mem_bytes: 0,
2910 };
2911 let __s_b = self.gpu.stream();
2912 let mut b = __s_b.launch_builder(&f);
2913 b.arg(x)
2914 .arg(&mut *y)
2915 .arg(&ni)
2916 .arg(&slo)
2917 .arg(&shi)
2918 .arg(&stream_pos)
2919 .arg(&temp);
2920 unsafe {
2921 b.launch(cfg)?;
2922 }
2923 Ok(())
2924 }
2925
2926 pub fn mask_logits_col(
2934 &self,
2935 logits: &mut CudaSlice<f32>,
2936 mask: &CudaSlice<u32>,
2937 col: usize,
2938 n: usize,
2939 mask_words: usize,
2940 ) -> Result<(), Box<dyn std::error::Error>> {
2941 let f = self.func("mask_logits_f32");
2942 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
2943 let cfg = LaunchConfig {
2944 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
2945 block_dim: (256, 1, 1),
2946 shared_mem_bytes: 0,
2947 };
2948 let __s_b = self.gpu.stream();
2949 let mut b = __s_b.launch_builder(&f);
2950 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
2951 unsafe {
2952 b.launch(cfg)?;
2953 }
2954 Ok(())
2955 }
2956
2957 pub fn gumbel_perturb_col(
2964 &self,
2965 x: &CudaSlice<f32>,
2966 col: usize,
2967 y: &mut CudaSlice<f32>,
2968 n: usize,
2969 seed: u64,
2970 stream_pos: u32,
2971 temp: f32,
2972 ) -> Result<(), Box<dyn std::error::Error>> {
2973 let f = self.func("gumbel_perturb_f32");
2974 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2975 let col_view = x.slice(col * n..(col + 1) * n);
2976 let cfg = LaunchConfig {
2977 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2978 block_dim: (256, 1, 1),
2979 shared_mem_bytes: 0,
2980 };
2981 let __s_b = self.gpu.stream();
2982 let mut b = __s_b.launch_builder(&f);
2983 b.arg(&col_view)
2984 .arg(&mut *y)
2985 .arg(&ni)
2986 .arg(&slo)
2987 .arg(&shi)
2988 .arg(&stream_pos)
2989 .arg(&temp);
2990 unsafe {
2991 b.launch(cfg)?;
2992 }
2993 Ok(())
2994 }
2995
2996 #[allow(clippy::too_many_arguments)]
3002 pub fn gumbel_perturb_filtered_col(
3003 &self,
3004 x: &CudaSlice<f32>,
3005 col: usize,
3006 y: &mut CudaSlice<f32>,
3007 n: usize,
3008 seed: u64,
3009 stream_pos: u32,
3010 temp: f32,
3011 stat_max: &CudaSlice<f32>,
3012 stat_th: &CudaSlice<f32>,
3013 stat_idx: usize,
3014 ) -> Result<(), Box<dyn std::error::Error>> {
3015 let f = self.func("gumbel_perturb_filtered_col_f32");
3016 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3017 let (ci, si) = (col as i32, stat_idx as i32);
3018 let cfg = LaunchConfig {
3019 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3020 block_dim: (256, 1, 1),
3021 shared_mem_bytes: 0,
3022 };
3023 let __s_b = self.gpu.stream();
3024 let mut b = __s_b.launch_builder(&f);
3025 b.arg(x)
3026 .arg(&ci)
3027 .arg(&mut *y)
3028 .arg(&ni)
3029 .arg(&slo)
3030 .arg(&shi)
3031 .arg(&stream_pos)
3032 .arg(&temp)
3033 .arg(stat_max)
3034 .arg(stat_th)
3035 .arg(&si);
3036 unsafe {
3037 b.launch(cfg)?;
3038 }
3039 Ok(())
3040 }
3041
3042 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3047 let f = self.func("memra_sctr_inc");
3048 let cfg = LaunchConfig {
3049 grid_dim: (1, 1, 1),
3050 block_dim: (1, 1, 1),
3051 shared_mem_bytes: 0,
3052 };
3053 let __s_b = self.gpu.stream();
3054 let mut b = __s_b.launch_builder(&f);
3055 b.arg(&mut *ctr);
3056 unsafe {
3057 b.launch(cfg)?;
3058 }
3059 Ok(())
3060 }
3061
3062 pub fn gumbel_perturb_ctr(
3067 &self,
3068 x: &CudaSlice<f32>,
3069 y: &mut CudaSlice<f32>,
3070 n: usize,
3071 seed: u64,
3072 ctr: &CudaSlice<u32>,
3073 temp: f32,
3074 ) -> Result<(), Box<dyn std::error::Error>> {
3075 let f = self.func("gumbel_perturb_ctr_f32");
3076 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3077 let cfg = LaunchConfig {
3078 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3079 block_dim: (256, 1, 1),
3080 shared_mem_bytes: 0,
3081 };
3082 let __s_b = self.gpu.stream();
3083 let mut b = __s_b.launch_builder(&f);
3084 b.arg(x)
3085 .arg(&mut *y)
3086 .arg(&ni)
3087 .arg(&slo)
3088 .arg(&shi)
3089 .arg(ctr)
3090 .arg(&temp);
3091 unsafe {
3092 b.launch(cfg)?;
3093 }
3094 Ok(())
3095 }
3096
3097 pub fn softmax_gather(
3101 &self,
3102 x: &CudaSlice<f32>,
3103 row_stride: usize,
3104 ids: &CudaSlice<u32>,
3105 rows: &CudaSlice<i32>,
3106 out: &mut CudaSlice<f32>,
3107 n: usize,
3108 npair: usize,
3109 temp: f32,
3110 ) -> Result<(), Box<dyn std::error::Error>> {
3111 let f = self.func("softmax_gather_f32");
3112 let (ni, rs) = (n as i32, row_stride as i64);
3113 let np = npair as i32;
3114 let cfg = LaunchConfig {
3115 grid_dim: (npair as u32, 1, 1),
3116 block_dim: (256, 1, 1),
3117 shared_mem_bytes: 0,
3118 };
3119 let __s_b = self.gpu.stream();
3120 let mut b = __s_b.launch_builder(&f);
3121 b.arg(x)
3122 .arg(&rs)
3123 .arg(ids)
3124 .arg(rows)
3125 .arg(&mut *out)
3126 .arg(&ni)
3127 .arg(&np)
3128 .arg(&temp);
3129 unsafe {
3130 b.launch(cfg)?;
3131 }
3132 Ok(())
3133 }
3134
3135 pub fn residual_sample(
3139 &self,
3140 p: &CudaSlice<f32>,
3141 q: Option<&CudaSlice<f32>>,
3142 n: usize,
3143 temp: f32,
3144 seed: u64,
3145 stream_pos: u32,
3146 out_tok: &mut CudaSlice<u32>,
3147 ) -> Result<(), Box<dyn std::error::Error>> {
3148 let f = self.func("residual_sample_f32");
3149 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3150 let nth = 1024u32;
3151 let cfg = LaunchConfig {
3152 grid_dim: (1, 1, 1),
3153 block_dim: (nth, 1, 1),
3154 shared_mem_bytes: 0,
3155 };
3156 let has_q: i32 = q.is_some() as i32;
3157 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3159 let mut b = __s_b.launch_builder(&f);
3160 b.arg(p)
3161 .arg(qbuf)
3162 .arg(&has_q)
3163 .arg(&ni)
3164 .arg(&temp)
3165 .arg(&slo)
3166 .arg(&shi)
3167 .arg(&stream_pos)
3168 .arg(&mut *out_tok);
3169 unsafe {
3170 b.launch(cfg)?;
3171 }
3172 Ok(())
3173 }
3174
3175 pub fn with_moe_cache<R>(
3180 &self,
3181 max_block_bytes: usize,
3182 f: impl FnOnce(
3183 &mut crate::moe_cache::MoeSlotCache,
3184 &Engine,
3185 ) -> Result<R, Box<dyn std::error::Error>>,
3186 ) -> Result<R, Box<dyn std::error::Error>> {
3187 let mut guard = self.moe_cache.lock().unwrap();
3188 if guard.is_none() {
3189 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3190 }
3191 let cache = guard.as_mut().unwrap();
3192 f(cache, self)
3193 }
3194
3195 pub fn freeze_moe_cache(&self) {
3198 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3199 cache.freeze();
3200 }
3201 }
3202
3203 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3206 self.moe_cache
3207 .lock()
3208 .unwrap()
3209 .as_ref()
3210 .map(crate::moe_cache::MoeSlotCache::export_residency)
3211 }
3212
3213 pub(crate) fn moe_cache_frozen(&self) -> bool {
3214 self.moe_cache
3215 .lock()
3216 .unwrap()
3217 .as_ref()
3218 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3219 }
3220
3221 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3228 crate::cpu_experts::configured()
3229 && self.moe_cache_frozen()
3230 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3231 }
3232
3233 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3235 assert!(
3236 self.moe_cache.lock().unwrap().is_none(),
3237 "MoE cache layout configured after cache construction"
3238 );
3239 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3240 }
3241
3242 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3243 self.moe_cache_layout.lock().unwrap().clone()
3244 }
3245
3246 pub fn moe_cache_enabled() -> bool {
3248 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3249 }
3250
3251 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3254 let guard = self.moe_cache.lock().unwrap();
3255 guard
3256 .as_ref()
3257 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3258 }
3259
3260 pub fn cpu_expert_stats(
3264 &self,
3265 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3266 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3267 }
3268
3269 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3272 crate::cpu_experts::predictor_stats()
3273 }
3274
3275 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3276 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3277 }
3278
3279 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3282 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3283 }
3284
3285 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3288 let guard = self.moe_cache.lock().unwrap();
3289 guard
3290 .as_ref()
3291 .and_then(|cache| cache.pread_stats())
3292 .map(|stats| {
3293 (
3294 stats.reads,
3295 stats.bytes,
3296 stats.read_errors,
3297 stats.short_reads,
3298 stats.fallbacks,
3299 stats.buffer_waits,
3300 stats.ring_full,
3301 )
3302 })
3303 }
3304
3305 pub fn spill_config_fallbacks(&self) -> u64 {
3307 crate::spill_pread::config_fallbacks()
3308 }
3309
3310 pub fn moe_cache_reset_counters(&self) {
3312 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3313 c.reset_counters();
3314 }
3315 }
3316
3317 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3318 Ok(self.gpu.stream().clone_htod(v)?)
3319 }
3320
3321 pub fn htod_bytes_padded(
3325 &self,
3326 v: &[u8],
3327 pad: usize,
3328 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3329 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3330 {
3331 let mut view = d.slice_mut(0..v.len());
3332 self.gpu.stream().memcpy_htod(v, &mut view)?;
3333 }
3334 Ok(d)
3335 }
3336
3337 pub fn copy_into(
3339 &self,
3340 dst: &mut CudaSlice<f32>,
3341 off: usize,
3342 src: &CudaSlice<f32>,
3343 len: usize,
3344 ) -> Result<(), Box<dyn std::error::Error>> {
3345 let mut view = dst.slice_mut(off..off + len);
3346 self.gpu
3347 .stream()
3348 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3349 Ok(())
3350 }
3351
3352 pub fn copy_u8_into(
3355 &self,
3356 dst: &mut CudaSlice<u8>,
3357 off: usize,
3358 src: &CudaSlice<u8>,
3359 len: usize,
3360 ) -> Result<(), Box<dyn std::error::Error>> {
3361 let mut view = dst.slice_mut(off..off + len);
3362 self.gpu
3363 .stream()
3364 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3365 Ok(())
3366 }
3367
3368 pub fn copy_u8_range_into(
3370 &self,
3371 dst: &mut CudaSlice<u8>,
3372 dst_off: usize,
3373 src: &CudaSlice<u8>,
3374 src_off: usize,
3375 len: usize,
3376 ) -> Result<(), Box<dyn std::error::Error>> {
3377 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3378 self.gpu
3379 .stream()
3380 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3381 Ok(())
3382 }
3383
3384 pub fn prepare_kv_append(
3388 &self,
3389 kv: &mut crate::cache::KvLayer,
3390 retain_from: usize,
3391 append_rows: usize,
3392 ) -> Result<usize, Box<dyn std::error::Error>> {
3393 let Some(plan) = kv
3394 .ring
3395 .as_ref()
3396 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3397 .transpose()?
3398 else {
3399 return Ok(kv.len);
3400 };
3401 match plan {
3402 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3403 crate::cache::KvRingAppend::Rebase {
3404 src_row,
3405 keep_rows,
3406 new_base,
3407 write_row,
3408 } => {
3409 if keep_rows > 0 {
3410 let k_len = keep_rows * kv.k_tok_bytes;
3411 let v_len = keep_rows * kv.v_tok_bytes;
3412 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3413 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3414 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3415 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3416 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3417 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3418 }
3419 kv.ring.as_mut().unwrap().apply_rebase(new_base);
3420 Ok(write_row)
3421 }
3422 }
3423 }
3424
3425 pub fn htod_u8_into(
3428 &self,
3429 dst: &mut CudaSlice<u8>,
3430 off: usize,
3431 src: &[u8],
3432 ) -> Result<(), Box<dyn std::error::Error>> {
3433 let mut view = dst.slice_mut(off..off + src.len());
3434 self.gpu.stream().memcpy_htod(src, &mut view)?;
3435 Ok(())
3436 }
3437
3438 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3439 b.slice(0..len)
3440 }
3441
3442 pub fn view_u8_range<'a>(
3445 &self,
3446 b: &'a CudaSlice<u8>,
3447 start: usize,
3448 end: usize,
3449 ) -> cudarc::driver::CudaView<'a, u8> {
3450 b.slice(start..end)
3451 }
3452 pub fn view_u8<'a>(
3453 &self,
3454 b: &'a CudaSlice<u8>,
3455 len: usize,
3456 ) -> cudarc::driver::CudaView<'a, u8> {
3457 b.slice(0..len)
3458 }
3459
3460 pub fn append_kv_quantized(
3464 &self,
3465 k_row: &CudaSlice<f32>,
3466 v_row: &CudaSlice<f32>,
3467 kc: &mut CudaSlice<u8>,
3468 vc: &mut CudaSlice<u8>,
3469 t: usize,
3470 kv_dim_k: usize,
3471 kv_dim_v: usize,
3472 k_tok_bytes: usize,
3473 v_tok_bytes: usize,
3474 g: bool,
3475 ) -> Result<(), Box<dyn std::error::Error>> {
3476 let f = if g {
3477 self.func_g("append_quantize_kv_q8_0_q5_1")
3478 } else {
3479 self.func("append_quantize_kv_q8_0_q5_1")
3480 };
3481 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3482 let cfg = LaunchConfig {
3483 grid_dim: (nblk, 1, 1),
3484 block_dim: (32, 1, 1),
3485 shared_mem_bytes: 0,
3486 };
3487 let (ti, kdk, kdv) = (t as i32, 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 let __s_b = self.gpu.stream();
3490 let mut b = __s_b.launch_builder(&f);
3491 b.arg(k_row)
3492 .arg(v_row)
3493 .arg(kc)
3494 .arg(vc)
3495 .arg(&ti)
3496 .arg(&kdk)
3497 .arg(&kdv)
3498 .arg(&ktb)
3499 .arg(&vtb);
3500 unsafe {
3501 b.launch(cfg)?;
3502 }
3503 Ok(())
3504 }
3505
3506 pub fn append_kv_quantized_dc(
3510 &self,
3511 k_row: &CudaSlice<f32>,
3512 v_row: &CudaSlice<f32>,
3513 kc: &mut CudaSlice<u8>,
3514 vc: &mut CudaSlice<u8>,
3515 t_dev: &CudaSlice<i32>,
3516 kv_dim_k: usize,
3517 kv_dim_v: usize,
3518 k_tok_bytes: usize,
3519 v_tok_bytes: usize,
3520 g: bool,
3521 ) -> Result<(), Box<dyn std::error::Error>> {
3522 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3523 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3524 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3525 if Self::pdl_on() && Self::pdl_wb_on() {
3527 use cudarc::driver::{DevicePtr, DevicePtrMut};
3528 let s = &self.gpu.stream();
3529 let (pk, _g0) = k_row.device_ptr(s);
3530 let (pv, _g1) = v_row.device_ptr(s);
3531 let (pkc, _g2) = kc.device_ptr_mut(s);
3532 let (pvc, _g3) = vc.device_ptr_mut(s);
3533 let (pt, _g4) = t_dev.device_ptr(s);
3534 let mut ps = [
3535 &pk as *const _ as *mut std::ffi::c_void,
3536 &pv as *const _ as *mut _,
3537 &pkc as *const _ as *mut _,
3538 &pvc as *const _ as *mut _,
3539 &pt as *const _ as *mut _,
3540 &kdk as *const _ as *mut _,
3541 &kdv as *const _ as *mut _,
3542 &ktb as *const _ as *mut _,
3543 &vtb as *const _ as *mut _,
3544 ];
3545 unsafe {
3546 self.launch_pdl_flash(
3547 g,
3548 "append_quantize_kv_q8_0_q5_1_dc",
3549 (nblk, 1, 1),
3550 (32, 1, 1),
3551 0,
3552 &mut ps,
3553 )?;
3554 }
3555 return Ok(());
3556 }
3557 let f = if g {
3558 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3559 } else {
3560 self.func("append_quantize_kv_q8_0_q5_1_dc")
3561 };
3562 let cfg = LaunchConfig {
3563 grid_dim: (nblk, 1, 1),
3564 block_dim: (32, 1, 1),
3565 shared_mem_bytes: 0,
3566 };
3567 let __s_b = self.gpu.stream();
3568 let mut b = __s_b.launch_builder(&f);
3569 b.arg(k_row)
3570 .arg(v_row)
3571 .arg(kc)
3572 .arg(vc)
3573 .arg(t_dev)
3574 .arg(&kdk)
3575 .arg(&kdv)
3576 .arg(&ktb)
3577 .arg(&vtb);
3578 unsafe {
3579 b.launch(cfg)?;
3580 }
3581 Ok(())
3582 }
3583
3584 #[allow(clippy::too_many_arguments)]
3591 pub fn append_kv_quantized_rows(
3592 &self,
3593 k_rows: &CudaSlice<f32>,
3594 v_rows: &CudaSlice<f32>,
3595 kc: &mut CudaSlice<u8>,
3596 vc: &mut CudaSlice<u8>,
3597 t0: usize,
3598 t: usize,
3599 kv_dim_k: usize,
3600 kv_dim_v: usize,
3601 k_tok_bytes: usize,
3602 v_tok_bytes: usize,
3603 g: bool,
3604 ) -> Result<(), Box<dyn std::error::Error>> {
3605 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3606 for i in 0..t {
3607 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3608 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3609 self.append_kv_quantized_view(
3610 &k_row,
3611 &v_row,
3612 kc,
3613 vc,
3614 t0 + i,
3615 kv_dim_k,
3616 kv_dim_v,
3617 k_tok_bytes,
3618 v_tok_bytes,
3619 g,
3620 )?;
3621 }
3622 return Ok(());
3623 }
3624 let f = if g {
3625 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3626 } else {
3627 self.func("append_quantize_kv_q8_0_q5_1_rows")
3628 };
3629 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3630 let cfg = LaunchConfig {
3631 grid_dim: (nblk, t as u32, 1),
3632 block_dim: (32, 1, 1),
3633 shared_mem_bytes: 0,
3634 };
3635 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3636 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3637 let __s_b = self.gpu.stream();
3638 let mut b = __s_b.launch_builder(&f);
3639 b.arg(k_rows)
3640 .arg(v_rows)
3641 .arg(kc)
3642 .arg(vc)
3643 .arg(&t0i)
3644 .arg(&kdk)
3645 .arg(&kdv)
3646 .arg(&ktb)
3647 .arg(&vtb);
3648 unsafe {
3649 b.launch(cfg)?;
3650 }
3651 Ok(())
3652 }
3653
3654 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3658 let f = self.func("inc_i32");
3659 let cfg = LaunchConfig {
3660 grid_dim: (1, 1, 1),
3661 block_dim: (1, 1, 1),
3662 shared_mem_bytes: 0,
3663 };
3664 let __s_b = self.gpu.stream();
3665 let mut b = __s_b.launch_builder(&f);
3666 b.arg(p);
3667 unsafe {
3668 b.launch(cfg)?;
3669 }
3670 Ok(())
3671 }
3672
3673 pub fn append_kv_quantized_view(
3676 &self,
3677 k_row: &cudarc::driver::CudaView<f32>,
3678 v_row: &cudarc::driver::CudaView<f32>,
3679 kc: &mut CudaSlice<u8>,
3680 vc: &mut CudaSlice<u8>,
3681 t: usize,
3682 kv_dim_k: usize,
3683 kv_dim_v: usize,
3684 k_tok_bytes: usize,
3685 v_tok_bytes: usize,
3686 g: bool,
3687 ) -> Result<(), Box<dyn std::error::Error>> {
3688 let f = if g {
3689 self.func_g("append_quantize_kv_q8_0_q5_1")
3690 } else {
3691 self.func("append_quantize_kv_q8_0_q5_1")
3692 };
3693 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3694 let cfg = LaunchConfig {
3695 grid_dim: (nblk, 1, 1),
3696 block_dim: (32, 1, 1),
3697 shared_mem_bytes: 0,
3698 };
3699 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3700 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3701 let __s_b = self.gpu.stream();
3702 let mut b = __s_b.launch_builder(&f);
3703 b.arg(k_row)
3704 .arg(v_row)
3705 .arg(kc)
3706 .arg(vc)
3707 .arg(&ti)
3708 .arg(&kdk)
3709 .arg(&kdv)
3710 .arg(&ktb)
3711 .arg(&vtb);
3712 unsafe {
3713 b.launch(cfg)?;
3714 }
3715 Ok(())
3716 }
3717
3718 pub fn copy_view_into(
3721 &self,
3722 dst: &mut CudaSlice<f32>,
3723 off: usize,
3724 src: &cudarc::driver::CudaView<f32>,
3725 len: usize,
3726 ) -> Result<(), Box<dyn std::error::Error>> {
3727 let mut view = dst.slice_mut(off..off + len);
3728 self.gpu
3729 .stream()
3730 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3731 Ok(())
3732 }
3733
3734 pub fn clone_dtod(
3738 &self,
3739 src: &CudaSlice<f32>,
3740 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3741 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
3742 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
3743 Ok(dst)
3744 }
3745
3746 pub fn dtod_copy_view(
3749 &self,
3750 src: &cudarc::driver::CudaView<f32>,
3751 dst: &mut CudaSlice<f32>,
3752 ) -> Result<(), Box<dyn std::error::Error>> {
3753 self.gpu.stream().memcpy_dtod(src, dst)?;
3754 Ok(())
3755 }
3756
3757 pub fn dtod_copy_view_i8(
3759 &self,
3760 src: &cudarc::driver::CudaView<i8>,
3761 dst: &mut CudaSlice<i8>,
3762 ) -> Result<(), Box<dyn std::error::Error>> {
3763 self.gpu.stream().memcpy_dtod(src, dst)?;
3764 Ok(())
3765 }
3766
3767 pub fn dtod_copy_into(
3769 &self,
3770 src: &CudaSlice<f32>,
3771 dst: &mut CudaSlice<f32>,
3772 offset: usize,
3773 ) -> Result<(), Box<dyn std::error::Error>> {
3774 let n = src.len();
3775 let mut dv = dst.slice_mut(offset..offset + n);
3776 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
3777 Ok(())
3778 }
3779
3780 pub fn copy_batch_uniform_f32(
3786 &self,
3787 table: &CudaSlice<u64>,
3788 n: usize,
3789 words: usize,
3790 ) -> Result<(), Box<dyn std::error::Error>> {
3791 if n == 0 || words == 0 {
3792 return Ok(());
3793 }
3794 debug_assert!(
3795 table.len() >= 2 * n,
3796 "pointer table must hold n srcs + n dsts"
3797 );
3798 let f = self.func("copy_batch_uniform_f32");
3799 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
3802 let (ni, wi) = (n as i32, words as i32);
3803 let cfg = LaunchConfig {
3804 grid_dim: (chunks, n as u32, 1),
3805 block_dim: (256, 1, 1),
3806 shared_mem_bytes: 0,
3807 };
3808 let __s = self.gpu.stream();
3809 let mut b = __s.launch_builder(&f);
3810 b.arg(table).arg(&ni).arg(&wi);
3811 unsafe {
3812 b.launch(cfg)?;
3813 }
3814 Ok(())
3815 }
3816
3817 pub fn htod_u64_into(
3820 &self,
3821 v: &[u64],
3822 dst: &mut CudaSlice<u64>,
3823 ) -> Result<(), Box<dyn std::error::Error>> {
3824 let mut view = dst.slice_mut(0..v.len());
3825 self.gpu.stream().memcpy_htod(v, &mut view)?;
3826 Ok(())
3827 }
3828
3829 pub fn copy_indirect_src_f32(
3834 &self,
3835 src_entry: &cudarc::driver::CudaView<u64>,
3836 dst: &mut CudaSlice<f32>,
3837 dst_off: usize,
3838 words: usize,
3839 ) -> Result<(), Box<dyn std::error::Error>> {
3840 let f = self.func("copy_indirect_src_f32");
3841 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
3842 let wi = words as i32;
3843 let cfg = LaunchConfig {
3844 grid_dim: (chunks, 1, 1),
3845 block_dim: (256, 1, 1),
3846 shared_mem_bytes: 0,
3847 };
3848 let mut dv = dst.slice_mut(dst_off..dst_off + words);
3849 let __s = self.gpu.stream();
3850 let mut b = __s.launch_builder(&f);
3851 b.arg(src_entry).arg(&mut dv).arg(&wi);
3852 unsafe {
3853 b.launch(cfg)?;
3854 }
3855 Ok(())
3856 }
3857
3858 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
3860 self.alloc_uninit::<i8>(n)
3861 }
3862
3863 pub fn qmatvec(
3865 &self,
3866 w: &CudaSlice<u8>,
3867 x: &CudaSlice<f32>,
3868 m: usize,
3869 in_f: usize,
3870 out_f: usize,
3871 qtype: i32,
3872 row_bytes: usize,
3873 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3874 let f = self.func("qmatvec_f32");
3875 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
3877 grid_dim: (out_f as u32, m as u32, 1),
3878 block_dim: (256, 1, 1),
3879 shared_mem_bytes: 0,
3880 };
3881 let (inf, outf, mi, qt, rb) =
3882 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
3883 let __s_b = self.gpu.stream();
3884 let mut b = __s_b.launch_builder(&f);
3885 b.arg(w)
3886 .arg(x)
3887 .arg(&mut y)
3888 .arg(&inf)
3889 .arg(&outf)
3890 .arg(&mi)
3891 .arg(&qt)
3892 .arg(&rb);
3893 unsafe {
3894 b.launch(cfg)?;
3895 }
3896 Ok(y)
3897 }
3898
3899 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3901 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
3902 self.keep_if_capturing(&s);
3903 Ok(s)
3904 }
3905
3906 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3910 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
3911 self.keep_if_capturing(&s);
3912 Ok(s)
3913 }
3914
3915 pub fn memset_zeros_view(
3918 &self,
3919 dst: &mut cudarc::driver::CudaViewMut<f32>,
3920 ) -> Result<(), Box<dyn std::error::Error>> {
3921 self.gpu.stream().memset_zeros(dst)?;
3922 Ok(())
3923 }
3924
3925 pub fn stage_expert(
3931 &self,
3932 host_bytes: &[u8],
3933 scratch: &mut CudaSlice<u8>,
3934 off: usize,
3935 ) -> Result<(), Box<dyn std::error::Error>> {
3936 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
3939 }
3940
3941 pub fn moe_router_topk(
3947 &self,
3948 logits: &CudaSlice<f32>,
3949 t: usize,
3950 n_expert: usize,
3951 n_used: usize,
3952 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3953 let f = self.func("moe_router_topk_f32");
3954 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 {
3957 grid_dim: (t as u32, 1, 1),
3958 block_dim: (n_expert as u32, 1, 1),
3959 shared_mem_bytes: 0,
3960 };
3961 let (ne, nu) = (n_expert as i32, n_used as i32);
3962 let __s_b = self.gpu.stream();
3963 let mut b = __s_b.launch_builder(&f);
3964 b.arg(logits)
3965 .arg(&mut sel_idx)
3966 .arg(&mut sel_w)
3967 .arg(&ne)
3968 .arg(&nu);
3969 unsafe {
3970 b.launch(cfg)?;
3971 }
3972 Ok((sel_idx, sel_w))
3973 }
3974
3975 pub fn moe_router_topk_scaled(
3978 &self,
3979 logits: &CudaSlice<f32>,
3980 t: usize,
3981 n_expert: usize,
3982 n_used: usize,
3983 ex_scale: &CudaSlice<f32>,
3984 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3985 let f = self.func("moe_router_topk_scaled_f32");
3990 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
3991 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
3992 let cfg = LaunchConfig {
3993 grid_dim: (t as u32, 1, 1),
3994 block_dim: (n_expert as u32, 1, 1),
3995 shared_mem_bytes: 0,
3996 };
3997 let (ne, nu) = (n_expert as i32, n_used as i32);
3998 let __s_b = self.gpu.stream();
3999 let mut b = __s_b.launch_builder(&f);
4000 b.arg(logits)
4001 .arg(&mut sel_idx)
4002 .arg(&mut sel_w)
4003 .arg(&ne)
4004 .arg(&nu)
4005 .arg(ex_scale);
4006 unsafe {
4007 b.launch(cfg)?;
4008 }
4009 Ok((sel_idx, sel_w))
4010 }
4011
4012 pub fn moe_router_topk_host(
4020 &self,
4021 logits: &CudaSlice<f32>,
4022 t: usize,
4023 n_expert: usize,
4024 n_used: usize,
4025 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4026 let f = self.func("moe_router_topk_f32");
4027 let n = t * n_used;
4028 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4029 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4030 let cfg = LaunchConfig {
4031 grid_dim: (t as u32, 1, 1),
4032 block_dim: (n_expert as u32, 1, 1),
4033 shared_mem_bytes: 0,
4034 };
4035 let (ne, nu) = (n_expert as i32, n_used as i32);
4036 let __s_b = self.gpu.stream();
4037 let mut b = __s_b.launch_builder(&f);
4038 b.arg(logits)
4039 .arg(&mut sel_idx)
4040 .arg(&mut sel_w)
4041 .arg(&ne)
4042 .arg(&nu);
4043 unsafe {
4044 b.launch(cfg)?;
4045 }
4046 let bytes = n * 8;
4048 let mut guard = self.router_stage.lock().unwrap();
4049 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4050 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4051 }
4052 let stage = guard.as_mut().unwrap();
4053 let (si, sw) = unsafe {
4054 (
4055 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4056 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4057 )
4058 };
4059 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()))
4063 }
4064
4065 #[allow(clippy::too_many_arguments)]
4069 pub fn moe_router_sigmoid_topk(
4070 &self,
4071 logits: &CudaSlice<f32>,
4072 t: usize,
4073 n_expert: usize,
4074 n_used: usize,
4075 active_count: usize,
4076 correction_bias: &CudaSlice<f32>,
4077 active: &CudaSlice<u8>,
4078 scaling_factor: f32,
4079 route_norm: bool,
4080 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4081 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4082 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4083 return Err(format!(
4084 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4085 )
4086 .into());
4087 }
4088 if logits.len() < t * n_expert
4089 || correction_bias.len() != n_expert
4090 || active.len() != n_expert
4091 {
4092 return Err(format!(
4093 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4094 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4095 ).into());
4096 }
4097 let f = self.func("moe_router_sigmoid_topk_f32");
4098 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4099 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4100 let threads = n_expert.div_ceil(32) * 32;
4101 let cfg = LaunchConfig {
4102 grid_dim: (t as u32, 1, 1),
4103 block_dim: (threads as u32, 1, 1),
4104 shared_mem_bytes: 0,
4105 };
4106 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4107 let __s_b = self.gpu.stream();
4108 let mut b = __s_b.launch_builder(&f);
4109 b.arg(logits)
4110 .arg(correction_bias)
4111 .arg(active)
4112 .arg(&mut sel_idx)
4113 .arg(&mut sel_w)
4114 .arg(&ne)
4115 .arg(&nu)
4116 .arg(&scaling_factor)
4117 .arg(&rn);
4118 unsafe {
4119 b.launch(cfg)?;
4120 }
4121 Ok((sel_idx, sel_w))
4122 }
4123
4124 #[allow(clippy::too_many_arguments)]
4127 pub fn moe_router_sigmoid_topk_host(
4128 &self,
4129 logits: &CudaSlice<f32>,
4130 t: usize,
4131 n_expert: usize,
4132 n_used: usize,
4133 active_count: usize,
4134 correction_bias: &CudaSlice<f32>,
4135 active: &CudaSlice<u8>,
4136 scaling_factor: f32,
4137 route_norm: bool,
4138 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4139 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4140 logits,
4141 t,
4142 n_expert,
4143 n_used,
4144 active_count,
4145 correction_bias,
4146 active,
4147 scaling_factor,
4148 route_norm,
4149 )?;
4150 let n = t * n_used;
4151 let bytes = n * 8;
4152 let mut guard = self.router_stage.lock().unwrap();
4153 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4154 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4155 }
4156 let stage = guard.as_mut().unwrap();
4157 let (si, sw) = unsafe {
4158 (
4159 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4160 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4161 )
4162 };
4163 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4164 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4165 self.gpu.stream().synchronize()?;
4166 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4167 }
4168
4169 pub fn stage_expert_async(
4173 &self,
4174 host_bytes: &[u8],
4175 scratch: &mut CudaSlice<u8>,
4176 off: usize,
4177 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4178 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4179 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4180 Ok(self.copy_stream.record_event(None)?)
4181 }
4182
4183 pub fn compute_wait(
4185 &self,
4186 ev: &cudarc::driver::CudaEvent,
4187 ) -> Result<(), Box<dyn std::error::Error>> {
4188 self.gpu.stream().wait(ev)?;
4189 Ok(())
4190 }
4191
4192 pub fn qmatvec_view(
4197 &self,
4198 w: &CudaSlice<u8>,
4199 range: std::ops::Range<usize>,
4200 x: &cudarc::driver::CudaView<f32>,
4201 m: usize,
4202 in_f: usize,
4203 out_f: usize,
4204 qtype: i32,
4205 row_bytes: usize,
4206 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4207 let f = self.func("qmatvec_f32");
4208 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4211 grid_dim: (out_f as u32, m as u32, 1),
4212 block_dim: (256, 1, 1),
4213 shared_mem_bytes: 0,
4214 };
4215 let (inf, outf, mi, qt, rb) =
4216 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4217 let __s_b = self.gpu.stream();
4218 let mut b = __s_b.launch_builder(&f);
4219 b.arg(&wv)
4220 .arg(x)
4221 .arg(&mut y)
4222 .arg(&inf)
4223 .arg(&outf)
4224 .arg(&mi)
4225 .arg(&qt)
4226 .arg(&rb);
4227 unsafe {
4228 b.launch(cfg)?;
4229 }
4230 Ok(y)
4231 }
4232
4233 #[allow(clippy::too_many_arguments)]
4240 pub fn moe_gate_up_silu8_q8(
4244 &self,
4245 gp: WPtr8,
4246 up: WPtr8,
4247 aq: &CudaSlice<i8>,
4248 ad: &CudaSlice<f32>,
4249 in_f: usize,
4250 n_ff: usize,
4251 n_used: usize,
4252 qt_g: i32,
4253 qt_u: i32,
4254 rb_g: usize,
4255 rb_u: usize,
4256 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4257 let f = self.func("moe_gate_up_silu8_q8");
4258 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4259 let cfg = LaunchConfig {
4260 grid_dim: (n_ff as u32, n_used as u32, 1),
4261 block_dim: (32, 1, 1),
4262 shared_mem_bytes: 0,
4263 };
4264 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4265 let __s_b = self.gpu.stream();
4266 let mut b = __s_b.launch_builder(&f);
4267 b.arg(&gp)
4268 .arg(&up)
4269 .arg(aq)
4270 .arg(ad)
4271 .arg(&mut act)
4272 .arg(&inf)
4273 .arg(&nff)
4274 .arg(&qt_g)
4275 .arg(&qt_u)
4276 .arg(&rbg)
4277 .arg(&rbu);
4278 unsafe {
4279 b.launch(cfg)?;
4280 }
4281 Ok(act)
4282 }
4283
4284 #[allow(clippy::too_many_arguments)]
4285 pub fn moe_down8_fma_q8(
4286 &self,
4287 dp: WPtr8,
4288 w: F32x8,
4289 aq2: &CudaSlice<i8>,
4290 ad2: &CudaSlice<f32>,
4291 dst: &mut cudarc::driver::CudaViewMut<f32>,
4292 in_f: usize,
4293 out_f: usize,
4294 n_used: usize,
4295 qt: i32,
4296 rb: usize,
4297 ) -> Result<(), Box<dyn std::error::Error>> {
4298 let f = self.func("moe_down8_fma_q8");
4299 let cfg = LaunchConfig {
4300 grid_dim: (out_f as u32, 1, 1),
4301 block_dim: (32, 1, 1),
4302 shared_mem_bytes: 0,
4303 };
4304 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4305 let __s_b = self.gpu.stream();
4306 let mut b = __s_b.launch_builder(&f);
4307 b.arg(&dp)
4308 .arg(&w)
4309 .arg(aq2)
4310 .arg(ad2)
4311 .arg(dst)
4312 .arg(&inf)
4313 .arg(&outf)
4314 .arg(&nu)
4315 .arg(&qt)
4316 .arg(&rbi);
4317 unsafe {
4318 b.launch(cfg)?;
4319 }
4320 Ok(())
4321 }
4322
4323 pub fn qmatvec_expert_q8(
4325 &self,
4326 w: &CudaSlice<u8>,
4327 range: std::ops::Range<usize>,
4328 aq: &CudaSlice<i8>,
4329 ad: &CudaSlice<f32>,
4330 m: usize,
4331 in_f: usize,
4332 out_f: usize,
4333 qtype: i32,
4334 row_bytes: usize,
4335 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4336 let f = self.func("qmatvec_expert_q8");
4337 let wv = w.slice(range);
4338 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4339 const ROWS: u32 = 4; let cfg = LaunchConfig {
4341 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4342 block_dim: (32, ROWS, 1),
4343 shared_mem_bytes: 0,
4344 };
4345 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4346 let __s_b = self.gpu.stream();
4347 let mut b = __s_b.launch_builder(&f);
4348 b.arg(&wv)
4349 .arg(aq)
4350 .arg(ad)
4351 .arg(&mut y)
4352 .arg(&inf)
4353 .arg(&outf)
4354 .arg(&mi)
4355 .arg(&qtype)
4356 .arg(&rbi);
4357 unsafe {
4358 b.launch(cfg)?;
4359 }
4360 Ok(y)
4361 }
4362
4363 pub fn moe_gate_up_silu8(
4364 &self,
4365 gp: WPtr8,
4366 up: WPtr8,
4367 x: &cudarc::driver::CudaView<f32>,
4368 in_f: usize,
4369 n_ff: usize,
4370 n_used: usize,
4371 qt_g: i32,
4372 qt_u: i32,
4373 rb_g: usize,
4374 rb_u: usize,
4375 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4376 let f = self.func("moe_gate_up_silu8_f32");
4377 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
4379 grid_dim: (n_ff as u32, n_used as u32, 1),
4380 block_dim: (256, 1, 1),
4381 shared_mem_bytes: 0,
4382 };
4383 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4384 let __s_b = self.gpu.stream();
4385 let mut b = __s_b.launch_builder(&f);
4386 b.arg(&gp)
4387 .arg(&up)
4388 .arg(x)
4389 .arg(&mut act)
4390 .arg(&inf)
4391 .arg(&nff)
4392 .arg(&qt_g)
4393 .arg(&qt_u)
4394 .arg(&rbg)
4395 .arg(&rbu);
4396 unsafe {
4397 b.launch(cfg)?;
4398 }
4399 Ok(act)
4400 }
4401
4402 #[allow(clippy::too_many_arguments)]
4408 pub fn moe_down8_fma_into(
4409 &self,
4410 dp: WPtr8,
4411 w: F32x8,
4412 act: &CudaSlice<f32>,
4413 dst: &mut cudarc::driver::CudaViewMut<f32>,
4414 in_f: usize,
4415 out_f: usize,
4416 n_used: usize,
4417 qt: i32,
4418 rb: usize,
4419 ) -> Result<(), Box<dyn std::error::Error>> {
4420 let f = self.func("moe_down8_fma_f32");
4421 let cfg = LaunchConfig {
4422 grid_dim: (out_f as u32, 1, 1),
4423 block_dim: (256, 1, 1),
4424 shared_mem_bytes: 0,
4425 };
4426 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4427 let __s_b = self.gpu.stream();
4428 let mut b = __s_b.launch_builder(&f);
4429 b.arg(&dp)
4430 .arg(&w)
4431 .arg(act)
4432 .arg(dst)
4433 .arg(&inf)
4434 .arg(&outf)
4435 .arg(&nu)
4436 .arg(&qt)
4437 .arg(&rbv);
4438 unsafe {
4439 b.launch(cfg)?;
4440 }
4441 Ok(())
4442 }
4443
4444 #[allow(clippy::too_many_arguments)]
4449 #[allow(clippy::too_many_arguments)]
4464 #[allow(clippy::too_many_arguments)]
4466 pub fn moe_pairs_matvec_q8(
4467 &self,
4468 table: &CudaSlice<u64>,
4469 proj: i32,
4470 pair_tok: &CudaSlice<i32>,
4471 pair_ex: &CudaSlice<i32>,
4472 aq: &CudaSlice<i8>,
4473 ad: &CudaSlice<f32>,
4474 in_f: usize,
4475 out_f: usize,
4476 n_expert: usize,
4477 n_pairs: usize,
4478 qtype: i32,
4479 row_bytes: usize,
4480 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4481 let f = self.func("moe_pairs_matvec_q8");
4482 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4483 const ROWS: u32 = 4;
4484 let cfg = LaunchConfig {
4485 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4486 block_dim: (32, ROWS, 1),
4487 shared_mem_bytes: 0,
4488 };
4489 let (inf, outf, ne, np, rbi) = (
4490 in_f as i32,
4491 out_f as i32,
4492 n_expert as i32,
4493 n_pairs as i32,
4494 row_bytes as i64,
4495 );
4496 let __s_b = self.gpu.stream();
4497 let mut b = __s_b.launch_builder(&f);
4498 b.arg(table)
4499 .arg(&proj)
4500 .arg(pair_tok)
4501 .arg(pair_ex)
4502 .arg(aq)
4503 .arg(ad)
4504 .arg(&mut y)
4505 .arg(&inf)
4506 .arg(&outf)
4507 .arg(&ne)
4508 .arg(&np)
4509 .arg(&qtype)
4510 .arg(&rbi);
4511 unsafe {
4512 b.launch(cfg)?;
4513 }
4514 Ok(y)
4515 }
4516
4517 #[allow(clippy::too_many_arguments)]
4519 pub fn moe_pairs_matvec_q8_em(
4520 &self,
4521 table: &CudaSlice<u64>,
4522 proj: i32,
4523 ex_ids: &CudaSlice<i32>,
4524 ex_off: &CudaSlice<i32>,
4525 ex_pairs: &CudaSlice<i32>,
4526 pair_tok: &CudaSlice<i32>,
4527 aq: &CudaSlice<i8>,
4528 ad: &CudaSlice<f32>,
4529 in_f: usize,
4530 out_f: usize,
4531 n_expert: usize,
4532 n_active: usize,
4533 n_pairs: usize,
4534 qtype: i32,
4535 row_bytes: usize,
4536 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4537 let f = self.func("moe_pairs_matvec_q8_em");
4538 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4539 const ROWS: u32 = 4;
4540 let cfg = LaunchConfig {
4541 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4542 block_dim: (32, ROWS, 1),
4543 shared_mem_bytes: 0,
4544 };
4545 let (inf, outf, ne, na, rbi) = (
4546 in_f as i32,
4547 out_f as i32,
4548 n_expert as i32,
4549 n_active as i32,
4550 row_bytes as i64,
4551 );
4552 let __s_b = self.gpu.stream();
4553 let mut b = __s_b.launch_builder(&f);
4554 b.arg(table)
4555 .arg(&proj)
4556 .arg(ex_ids)
4557 .arg(ex_off)
4558 .arg(ex_pairs)
4559 .arg(pair_tok)
4560 .arg(aq)
4561 .arg(ad)
4562 .arg(&mut y)
4563 .arg(&inf)
4564 .arg(&outf)
4565 .arg(&ne)
4566 .arg(&na)
4567 .arg(&qtype)
4568 .arg(&rbi);
4569 unsafe {
4570 b.launch(cfg)?;
4571 }
4572 Ok(y)
4573 }
4574
4575 #[allow(clippy::too_many_arguments)]
4578 pub fn moe_pairs_matvec_q8_dec(
4579 &self,
4580 table: &CudaSlice<u64>,
4581 proj: i32,
4582 ex_ids: &CudaSlice<i32>,
4583 ex_off: &CudaSlice<i32>,
4584 ex_pairs: &CudaSlice<i32>,
4585 pair_tok: &CudaSlice<i32>,
4586 aq: &CudaSlice<i8>,
4587 ad: &CudaSlice<f32>,
4588 in_f: usize,
4589 out_f: usize,
4590 n_expert: usize,
4591 n_active: usize,
4592 n_pairs: usize,
4593 qtype: i32,
4594 row_bytes: usize,
4595 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4596 let f = self.func("moe_pairs_matvec_q8_dec");
4597 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4598 const ROWS: u32 = 4;
4599 let cfg = LaunchConfig {
4600 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4601 block_dim: (32, ROWS, 1),
4602 shared_mem_bytes: 0,
4603 };
4604 let (inf, outf, ne, na, rbi) = (
4605 in_f as i32,
4606 out_f as i32,
4607 n_expert as i32,
4608 n_active as i32,
4609 row_bytes as i64,
4610 );
4611 let __s_b = self.gpu.stream();
4612 let mut b = __s_b.launch_builder(&f);
4613 b.arg(table)
4614 .arg(&proj)
4615 .arg(ex_ids)
4616 .arg(ex_off)
4617 .arg(ex_pairs)
4618 .arg(pair_tok)
4619 .arg(aq)
4620 .arg(ad)
4621 .arg(&mut y)
4622 .arg(&inf)
4623 .arg(&outf)
4624 .arg(&ne)
4625 .arg(&na)
4626 .arg(&qtype)
4627 .arg(&rbi);
4628 unsafe {
4629 b.launch(cfg)?;
4630 }
4631 Ok(y)
4632 }
4633
4634 pub fn moe_pairs_gelu_mul(
4635 &self,
4636 gate: &CudaSlice<f32>,
4637 up: &CudaSlice<f32>,
4638 n: usize,
4639 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4640 let f = self.func("moe_pairs_gelu_mul");
4641 let mut act = self.alloc_uninit::<f32>(n)?;
4642 let cfg = LaunchConfig::for_num_elems(n as u32);
4643 let nl = n as i64;
4644 let __s_b = self.gpu.stream();
4645 let mut b = __s_b.launch_builder(&f);
4646 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4647 unsafe {
4648 b.launch(cfg)?;
4649 }
4650 Ok(act)
4651 }
4652
4653 pub fn moe_pairs_silu_mul(
4654 &self,
4655 gate: &CudaSlice<f32>,
4656 up: &CudaSlice<f32>,
4657 n: usize,
4658 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4659 let f = self.func("moe_pairs_silu_mul");
4660 let mut act = self.alloc_uninit::<f32>(n)?;
4661 let cfg = LaunchConfig::for_num_elems(n as u32);
4662 let nl = n as i64;
4663 let __s_b = self.gpu.stream();
4664 let mut b = __s_b.launch_builder(&f);
4665 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4666 unsafe {
4667 b.launch(cfg)?;
4668 }
4669 Ok(act)
4670 }
4671
4672 #[allow(clippy::too_many_arguments)]
4673 pub fn moe_pairs_scatter(
4674 &self,
4675 y_down: &CudaSlice<f32>,
4676 pair_w: &CudaSlice<f32>,
4677 tok_pair_off: &CudaSlice<i32>,
4678 tok_pair_ids: &CudaSlice<i32>,
4679 moe_out: &mut CudaSlice<f32>,
4680 t: usize,
4681 n_embd: usize,
4682 ) -> Result<(), Box<dyn std::error::Error>> {
4683 let f = self.func("moe_pairs_scatter");
4684 let cfg = LaunchConfig {
4685 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4686 block_dim: (256, 1, 1),
4687 shared_mem_bytes: 0,
4688 };
4689 let ne = n_embd as i32;
4690 let __s_b = self.gpu.stream();
4691 let mut b = __s_b.launch_builder(&f);
4692 b.arg(y_down)
4693 .arg(pair_w)
4694 .arg(tok_pair_off)
4695 .arg(tok_pair_ids)
4696 .arg(moe_out)
4697 .arg(&ne);
4698 unsafe {
4699 b.launch(cfg)?;
4700 }
4701 Ok(())
4702 }
4703
4704 #[allow(clippy::too_many_arguments)]
4708 pub fn moe_gate_up_gelu8_dev_q8(
4709 &self,
4710 table: &CudaSlice<u64>,
4711 sel: &cudarc::driver::CudaView<i32>,
4712 aq: &CudaSlice<i8>,
4713 ad: &CudaSlice<f32>,
4714 in_f: usize,
4715 n_ff: usize,
4716 n_used: usize,
4717 n_expert: usize,
4718 qt_g: i32,
4719 qt_u: i32,
4720 rb_g: usize,
4721 rb_u: usize,
4722 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4723 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4724 let (inf, nff, ne, rbg, rbu) = (
4725 in_f as i32,
4726 n_ff as i32,
4727 n_expert as i32,
4728 rb_g as i64,
4729 rb_u as i64,
4730 );
4731 let f = self.func("moe_gate_up_gelu8_dev_q8");
4732 let cfg = LaunchConfig {
4733 grid_dim: (n_ff as u32, n_used as u32, 1),
4734 block_dim: (32, 1, 1),
4735 shared_mem_bytes: 0,
4736 };
4737 let __s_b = self.gpu.stream();
4738 let mut b = __s_b.launch_builder(&f);
4739 b.arg(table)
4740 .arg(sel)
4741 .arg(aq)
4742 .arg(ad)
4743 .arg(&mut act)
4744 .arg(&inf)
4745 .arg(&nff)
4746 .arg(&ne)
4747 .arg(&qt_g)
4748 .arg(&qt_u)
4749 .arg(&rbg)
4750 .arg(&rbu);
4751 unsafe {
4752 b.launch(cfg)?;
4753 }
4754 Ok(act)
4755 }
4756
4757 #[allow(clippy::too_many_arguments)]
4759 pub fn moe_gate_up_gelu8_dev_q8_rows(
4760 &self,
4761 table: &CudaSlice<u64>,
4762 sel: &CudaSlice<i32>,
4763 aq: &CudaSlice<i8>,
4764 ad: &CudaSlice<f32>,
4765 t: usize,
4766 in_f: usize,
4767 n_ff: usize,
4768 n_used: usize,
4769 n_expert: usize,
4770 qt_g: i32,
4771 qt_u: i32,
4772 rb_g: usize,
4773 rb_u: usize,
4774 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4775 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
4776 let (inf, nff, ne, rbg, rbu, nu) = (
4777 in_f as i32,
4778 n_ff as i32,
4779 n_expert as i32,
4780 rb_g as i64,
4781 rb_u as i64,
4782 n_used as i32,
4783 );
4784 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
4785 let cfg = LaunchConfig {
4786 grid_dim: (n_ff as u32, n_used as u32, t as u32),
4787 block_dim: (32, 1, 1),
4788 shared_mem_bytes: 0,
4789 };
4790 let __s_b = self.gpu.stream();
4791 let mut b = __s_b.launch_builder(&f);
4792 b.arg(table)
4793 .arg(sel)
4794 .arg(aq)
4795 .arg(ad)
4796 .arg(&mut act)
4797 .arg(&inf)
4798 .arg(&nff)
4799 .arg(&ne)
4800 .arg(&qt_g)
4801 .arg(&qt_u)
4802 .arg(&rbg)
4803 .arg(&rbu)
4804 .arg(&nu);
4805 unsafe {
4806 b.launch(cfg)?;
4807 }
4808 Ok(act)
4809 }
4810
4811 #[allow(clippy::too_many_arguments)]
4813 pub fn moe_gate_up_gelu8_dev_q8_csr(
4814 &self,
4815 table: &CudaSlice<u64>,
4816 sel: &CudaSlice<i32>,
4817 aq: &CudaSlice<i8>,
4818 ad: &CudaSlice<f32>,
4819 n_pairs: usize,
4820 in_f: usize,
4821 n_ff: usize,
4822 n_used: usize,
4823 n_expert: usize,
4824 qt_g: i32,
4825 qt_u: i32,
4826 rb_g: usize,
4827 rb_u: usize,
4828 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4829 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
4830 let (inf, nff, ne, rbg, rbu, nu, npi) = (
4831 in_f as i32,
4832 n_ff as i32,
4833 n_expert as i32,
4834 rb_g as i64,
4835 rb_u as i64,
4836 n_used as i32,
4837 n_pairs as i32,
4838 );
4839 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
4840 let cfg = LaunchConfig {
4841 grid_dim: (n_ff as u32, n_pairs as u32, 1),
4842 block_dim: (32, 1, 1),
4843 shared_mem_bytes: 0,
4844 };
4845 let __s_b = self.gpu.stream();
4846 let mut b = __s_b.launch_builder(&f);
4847 b.arg(table)
4848 .arg(sel)
4849 .arg(aq)
4850 .arg(ad)
4851 .arg(&mut act)
4852 .arg(&inf)
4853 .arg(&nff)
4854 .arg(&ne)
4855 .arg(&qt_g)
4856 .arg(&qt_u)
4857 .arg(&rbg)
4858 .arg(&rbu)
4859 .arg(&nu)
4860 .arg(&npi);
4861 unsafe {
4862 b.launch(cfg)?;
4863 }
4864 Ok(act)
4865 }
4866
4867 #[allow(clippy::too_many_arguments)]
4869 pub fn moe_down8_fma_dev_q8_rows_g(
4870 &self,
4871 table: &CudaSlice<u64>,
4872 sel: &CudaSlice<i32>,
4873 w: &CudaSlice<f32>,
4874 aq2: &CudaSlice<i8>,
4875 ad2: &CudaSlice<f32>,
4876 dst: &mut CudaSlice<f32>,
4877 t: usize,
4878 in_f: usize,
4879 out_f: usize,
4880 n_used: usize,
4881 n_expert: usize,
4882 qt: i32,
4883 rb: usize,
4884 ) -> Result<(), Box<dyn std::error::Error>> {
4885 let (inf, outf, nu, ne, rbi) = (
4886 in_f as i32,
4887 out_f as i32,
4888 n_used as i32,
4889 n_expert as i32,
4890 rb as i64,
4891 );
4892 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
4896 let f = self.func(if step_b1_w8 {
4897 "moe_down8_fma_dev_q8_rows_w8"
4898 } else {
4899 "moe_down8_fma_dev_q8_rows_g"
4900 });
4901 let cfg = LaunchConfig {
4902 grid_dim: (out_f as u32, 1, t as u32),
4903 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
4904 shared_mem_bytes: 0,
4905 };
4906 let __s_b = self.gpu.stream();
4907 let mut b = __s_b.launch_builder(&f);
4908 b.arg(table)
4909 .arg(sel)
4910 .arg(w)
4911 .arg(aq2)
4912 .arg(ad2)
4913 .arg(dst)
4914 .arg(&inf)
4915 .arg(&outf)
4916 .arg(&nu)
4917 .arg(&ne)
4918 .arg(&qt)
4919 .arg(&rbi);
4920 unsafe {
4921 b.launch(cfg)?;
4922 }
4923 Ok(())
4924 }
4925
4926 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
4930 let (out_f, in_f) = (2048usize, 2816usize);
4931 let nblk = in_f / 32;
4932 let mut seed = 0x9E3779B97F4A7C15u64;
4933 let mut rng = move || {
4934 seed = seed
4935 .wrapping_mul(6364136223846793005)
4936 .wrapping_add(1442695040888963407);
4937 (seed >> 33) as u8
4938 };
4939 let mut w = vec![0u8; out_f * nblk * 18];
4940 for b in w.iter_mut() {
4941 *b = rng();
4942 }
4943 for r in 0..out_f {
4944 for g in 0..nblk {
4945 let off = (r * nblk + g) * 18;
4946 w[off] = 0x00;
4947 w[off + 1] = 0x2C; }
4949 }
4950 let qplane = out_f * nblk * 16;
4951 let mut wrp = vec![0u8; w.len()];
4952 for r in 0..out_f {
4953 for g in 0..nblk {
4954 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
4955 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
4956 .copy_from_slice(&src[0..2]);
4957 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
4958 }
4959 }
4960 let w_d = self.htod_bytes(&w)?;
4961 let wrp_d = self.htod_bytes(&wrp)?;
4962 let mut aq = vec![0i8; m * in_f];
4963 for v in aq.iter_mut() {
4964 *v = rng() as i8;
4965 }
4966 let aq_d = self.htod_i8(&aq)?;
4967 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
4968 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
4969 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
4970 const RPB: u32 = 4;
4971 let cfg = LaunchConfig {
4972 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
4973 block_dim: (32, RPB, 1),
4974 shared_mem_bytes: 0,
4975 };
4976 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
4977 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
4978 let fb = self.func("qmatvec_q4_0_mmvq_b4");
4979 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
4980 {
4981 let __s_b = self.gpu.stream();
4982 let mut b = __s_b.launch_builder(&fb);
4983 b.arg(&w_d)
4984 .arg(&aq_d)
4985 .arg(&ad_d)
4986 .arg(&mut y0)
4987 .arg(&inf)
4988 .arg(&outf)
4989 .arg(&mi)
4990 .arg(&rb);
4991 unsafe {
4992 b.launch(cfg)?;
4993 }
4994 let __s_b = self.gpu.stream();
4995 let mut b = __s_b.launch_builder(&fr);
4996 b.arg(&wrp_d)
4997 .arg(&aq_d)
4998 .arg(&ad_d)
4999 .arg(&mut y1)
5000 .arg(&inf)
5001 .arg(&outf)
5002 .arg(&mi)
5003 .arg(&qp);
5004 unsafe {
5005 b.launch(cfg)?;
5006 }
5007 }
5008 self.gpu.stream().synchronize()?;
5009 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5010 let nd = h0
5011 .iter()
5012 .zip(&h1)
5013 .filter(|(a, b)| a.to_bits() != b.to_bits())
5014 .count();
5015 if nd != 0 {
5016 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5017 }
5018 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5019 self.gpu.stream().synchronize()?;
5020 let t0 = std::time::Instant::now();
5021 for _ in 0..500 {
5022 if rp {
5023 let __s_b = self.gpu.stream();
5024 let mut b = __s_b.launch_builder(&fr);
5025 b.arg(&wrp_d)
5026 .arg(&aq_d)
5027 .arg(&ad_d)
5028 .arg(&mut y1)
5029 .arg(&inf)
5030 .arg(&outf)
5031 .arg(&mi)
5032 .arg(&qp);
5033 unsafe {
5034 b.launch(cfg)?;
5035 }
5036 } else {
5037 let __s_b = self.gpu.stream();
5038 let mut b = __s_b.launch_builder(&fb);
5039 b.arg(&w_d)
5040 .arg(&aq_d)
5041 .arg(&ad_d)
5042 .arg(&mut y0)
5043 .arg(&inf)
5044 .arg(&outf)
5045 .arg(&mi)
5046 .arg(&rb);
5047 unsafe {
5048 b.launch(cfg)?;
5049 }
5050 }
5051 }
5052 self.gpu.stream().synchronize()?;
5053 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5054 };
5055 let _ = time(false)?;
5056 let _ = time(true)?; Ok((time(false)?, time(true)?))
5058 }
5059
5060 pub fn build_q4_rp4(
5065 &self,
5066 t: &mut crate::model::GpuTensor,
5067 ) -> Result<(), Box<dyn std::error::Error>> {
5068 use crate::model::GpuTensor;
5069 let GpuTensor::Quant {
5070 bytes,
5071 qtype,
5072 row_bytes,
5073 ne,
5074 rp4,
5075 ..
5076 } = t
5077 else {
5078 return Ok(());
5079 };
5080 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5081 return Ok(());
5082 }
5083 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5084 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5085 return Ok(());
5086 }
5087 let nblk = in_f / 32;
5088 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5089 let f = self.func("q4_0_split_rp_build");
5090 let n = (out_f * nblk) as i32;
5091 let cfg = LaunchConfig {
5092 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5093 block_dim: (256, 1, 1),
5094 shared_mem_bytes: 0,
5095 };
5096 let (of, nb) = (out_f as i32, nblk as i32);
5097 let _ = n;
5098 let __s_b = self.gpu.stream();
5099 let mut b = __s_b.launch_builder(&f);
5100 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5101 unsafe {
5102 b.launch(cfg)?;
5103 }
5104 *rp4 = Some(dst);
5105 Ok(())
5106 }
5107
5108 pub fn build_q8_rp4(
5113 &self,
5114 t: &mut crate::model::GpuTensor,
5115 ) -> Result<(), Box<dyn std::error::Error>> {
5116 use crate::model::GpuTensor;
5117 let GpuTensor::Quant {
5118 bytes,
5119 qtype,
5120 row_bytes,
5121 ne,
5122 rp4,
5123 ..
5124 } = t
5125 else {
5126 return Ok(());
5127 };
5128 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5129 return Ok(());
5130 }
5131 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5132 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5133 return Ok(());
5134 }
5135 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5136 Ok(())
5137 }
5138
5139 pub fn build_q8_rp4_raw(
5142 &self,
5143 bytes: &CudaSlice<u8>,
5144 in_f: usize,
5145 out_f: usize,
5146 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5147 assert!(in_f % 32 == 0);
5148 let nblk = in_f / 32;
5149 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5150 let f = self.func("q8_0_split_rp_build");
5151 let cfg = LaunchConfig {
5152 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5153 block_dim: (256, 1, 1),
5154 shared_mem_bytes: 0,
5155 };
5156 let (of, nb) = (out_f as i32, nblk as i32);
5157 let __s_b = self.gpu.stream();
5158 let mut b = __s_b.launch_builder(&f);
5159 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5160 unsafe {
5161 b.launch(cfg)?;
5162 }
5163 Ok(dst)
5164 }
5165
5166 pub fn build_q4k_rp4(
5174 &self,
5175 t: &mut crate::model::GpuTensor,
5176 ) -> Result<(), Box<dyn std::error::Error>> {
5177 use crate::model::GpuTensor;
5178 let GpuTensor::Quant {
5179 bytes,
5180 qtype,
5181 row_bytes,
5182 ne,
5183 rp4,
5184 ..
5185 } = t
5186 else {
5187 return Ok(());
5188 };
5189 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5190 return Ok(());
5191 }
5192 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5193 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5194 return Ok(());
5195 }
5196 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5197 Ok(())
5198 }
5199
5200 pub fn build_q6k_rp4(
5201 &self,
5202 t: &mut crate::model::GpuTensor,
5203 ) -> Result<(), Box<dyn std::error::Error>> {
5204 use crate::model::GpuTensor;
5205 let GpuTensor::Quant {
5206 bytes,
5207 qtype,
5208 row_bytes,
5209 ne,
5210 rp4,
5211 ..
5212 } = t
5213 else {
5214 return Ok(());
5215 };
5216 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5217 return Ok(());
5218 }
5219 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5220 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5221 return Ok(());
5222 }
5223 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5224 Ok(())
5225 }
5226
5227 pub fn build_kq_rp4_raw(
5229 &self,
5230 bytes: &CudaSlice<u8>,
5231 in_f: usize,
5232 out_f: usize,
5233 qtype: i32,
5234 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5235 assert!(in_f % 256 == 0);
5236 let nsbk = in_f / 256;
5237 let (sb_bytes, kname) = match qtype {
5238 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5239 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5240 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5241 };
5242 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5243 let f = self.func(kname);
5244 let cfg = LaunchConfig {
5245 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5246 block_dim: (256, 1, 1),
5247 shared_mem_bytes: 0,
5248 };
5249 let (of, nb) = (out_f as i32, nsbk as i32);
5250 let __s_b = self.gpu.stream();
5251 let mut b = __s_b.launch_builder(&f);
5252 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5253 unsafe {
5254 b.launch(cfg)?;
5255 }
5256 Ok(dst)
5257 }
5258
5259 pub fn kqrp_enabled() -> bool {
5263 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5264 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5265 Ok("0") => false,
5266 Ok(_) => true,
5267 Err(_) => cfg!(memra_hopper_mma),
5268 })
5269 }
5270
5271 pub fn build_q4_rp_swap(
5277 &self,
5278 t: &mut crate::model::GpuTensor,
5279 ) -> Result<bool, Box<dyn std::error::Error>> {
5280 use crate::model::GpuTensor;
5281 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5291 return Ok(false);
5292 }
5293 self.build_q4_rp4(t)?;
5294 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5296 return Ok(false);
5297 };
5298 match rp4.take() {
5299 Some(split) => {
5300 *bytes = split; *rp = true;
5302 Ok(true)
5303 }
5304 None => Ok(false),
5305 }
5306 }
5307
5308 pub fn q4rp_enabled() -> bool {
5310 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5311 *ON.get_or_init(|| {
5312 std::env::var("MEMRA_Q4RP")
5313 .map(|v| v != "0")
5314 .unwrap_or(true)
5315 })
5316 }
5317
5318 pub fn copy_rows_strided(
5321 &self,
5322 src: &CudaSlice<f32>,
5323 dst: &mut CudaSlice<f32>,
5324 row_elems: usize,
5325 n_rows: usize,
5326 src_stride: usize,
5327 src_off: usize,
5328 ) -> Result<(), Box<dyn std::error::Error>> {
5329 let f = self.func("copy_rows_strided_f32");
5330 let cfg = LaunchConfig {
5331 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5332 block_dim: (256, 1, 1),
5333 shared_mem_bytes: 0,
5334 };
5335 let (re, nr) = (row_elems as i32, n_rows as i32);
5336 let (st, off) = (src_stride as i64, src_off as i64);
5337 let __s_b = self.gpu.stream();
5338 let mut b = __s_b.launch_builder(&f);
5339 b.arg(src)
5340 .arg(&mut *dst)
5341 .arg(&re)
5342 .arg(&nr)
5343 .arg(&st)
5344 .arg(&off);
5345 unsafe {
5346 b.launch(cfg)?;
5347 }
5348 Ok(())
5349 }
5350
5351 pub fn u32_set_k(
5353 &self,
5354 dst: &mut CudaSlice<u32>,
5355 v: u32,
5356 idx: usize,
5357 ) -> Result<(), Box<dyn std::error::Error>> {
5358 let f = self.func("u32_set_k");
5359 let cfg = LaunchConfig {
5360 grid_dim: (1, 1, 1),
5361 block_dim: (1, 1, 1),
5362 shared_mem_bytes: 0,
5363 };
5364 let ii = idx as i32;
5365 let __s_b = self.gpu.stream();
5366 let mut b = __s_b.launch_builder(&f);
5367 b.arg(dst).arg(&v).arg(&ii);
5368 unsafe {
5369 b.launch(cfg)?;
5370 }
5371 Ok(())
5372 }
5373
5374 pub fn i32_add_k(
5376 &self,
5377 d: &mut CudaSlice<i32>,
5378 v: i32,
5379 ) -> Result<(), Box<dyn std::error::Error>> {
5380 let f = self.func("i32_add_k");
5381 let cfg = LaunchConfig {
5382 grid_dim: (1, 1, 1),
5383 block_dim: (32, 1, 1),
5384 shared_mem_bytes: 0,
5385 };
5386 let __s_b = self.gpu.stream();
5387 let mut b = __s_b.launch_builder(&f);
5388 b.arg(d).arg(&v);
5389 unsafe {
5390 b.launch(cfg)?;
5391 }
5392 Ok(())
5393 }
5394
5395 pub fn i32_iota_from(
5397 &self,
5398 ctr: &CudaSlice<i32>,
5399 dst: &mut CudaSlice<i32>,
5400 n: usize,
5401 ) -> Result<(), Box<dyn std::error::Error>> {
5402 let f = self.func("i32_iota_from");
5403 let cfg = LaunchConfig::for_num_elems(n as u32);
5404 let ni = n as i32;
5405 let __s_b = self.gpu.stream();
5406 let mut b = __s_b.launch_builder(&f);
5407 b.arg(ctr).arg(dst).arg(&ni);
5408 unsafe {
5409 b.launch(cfg)?;
5410 }
5411 Ok(())
5412 }
5413
5414 pub fn u32_map_k(
5416 &self,
5417 buf: &mut CudaSlice<u32>,
5418 map: &CudaSlice<u32>,
5419 idx: usize,
5420 ) -> Result<(), Box<dyn std::error::Error>> {
5421 let f = self.func("u32_map_k");
5422 let cfg = LaunchConfig {
5423 grid_dim: (1, 1, 1),
5424 block_dim: (1, 1, 1),
5425 shared_mem_bytes: 0,
5426 };
5427 let ii = idx as i32;
5428 let __s_b = self.gpu.stream();
5429 let mut b = __s_b.launch_builder(&f);
5430 b.arg(buf).arg(map).arg(&ii);
5431 unsafe {
5432 b.launch(cfg)?;
5433 }
5434 Ok(())
5435 }
5436
5437 #[allow(clippy::too_many_arguments)]
5439 pub fn u32_pack2(
5440 &self,
5441 a: &CudaSlice<u32>,
5442 off_a: usize,
5443 n1: usize,
5444 b_in: &CudaSlice<u32>,
5445 n2: usize,
5446 out: &mut CudaSlice<u32>,
5447 ) -> Result<(), Box<dyn std::error::Error>> {
5448 let f = self.func("u32_pack2");
5449 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5450 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5451 let __s_b = self.gpu.stream();
5452 let mut b = __s_b.launch_builder(&f);
5453 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5454 unsafe {
5455 b.launch(cfg)?;
5456 }
5457 Ok(())
5458 }
5459
5460 pub fn moe_w_exscale(
5462 &self,
5463 w: &mut CudaSlice<f32>,
5464 sel: &CudaSlice<i32>,
5465 s: &CudaSlice<f32>,
5466 n: usize,
5467 ) -> Result<(), Box<dyn std::error::Error>> {
5468 let f = self.func("moe_w_exscale");
5469 let cfg = LaunchConfig::for_num_elems(n as u32);
5470 let ni = n as i32;
5471 let __s_b = self.gpu.stream();
5472 let mut b = __s_b.launch_builder(&f);
5473 b.arg(w).arg(sel).arg(s).arg(&ni);
5474 unsafe {
5475 b.launch(cfg)?;
5476 }
5477 Ok(())
5478 }
5479
5480 pub fn moe_w_scale_by_expert(
5483 &self,
5484 w: &mut CudaSlice<f32>,
5485 sel: &CudaSlice<i32>,
5486 macros: &CudaSlice<f32>,
5487 n_expert: usize,
5488 n: usize,
5489 ) -> Result<(), Box<dyn std::error::Error>> {
5490 let f = self.func("moe_w_scale_by_expert");
5491 let cfg = LaunchConfig {
5492 grid_dim: (n.div_ceil(64) as u32, 1, 1),
5493 block_dim: (64, 1, 1),
5494 shared_mem_bytes: 0,
5495 };
5496 let (ne, nn) = (n_expert as i32, n as i32);
5497 let __s_b = self.gpu.stream();
5498 let mut b = __s_b.launch_builder(&f);
5499 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5500 unsafe {
5501 b.launch(cfg)?;
5502 }
5503 Ok(())
5504 }
5505
5506 pub fn moe_gate_up_silu8_dev_q8(
5507 &self,
5508 table: &CudaSlice<u64>,
5509 sel: &cudarc::driver::CudaView<i32>,
5510 aq: &CudaSlice<i8>,
5511 ad: &CudaSlice<f32>,
5512 in_f: usize,
5513 n_ff: usize,
5514 n_used: usize,
5515 n_expert: usize,
5516 qt_g: i32,
5517 qt_u: i32,
5518 rb_g: usize,
5519 rb_u: usize,
5520 macros: &CudaSlice<f32>,
5521 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5522 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5523 let (mode, wpb) = GU.get_or_init(|| {
5524 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5525 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5526 .ok()
5527 .and_then(|v| v.parse().ok())
5528 .unwrap_or(4u32)
5529 .clamp(1, 16);
5530 (mode, wpb)
5531 });
5532 let (mode, wpb) = (mode.as_str(), *wpb);
5533 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5534 let (inf, nff, ne, rbg, rbu) = (
5535 in_f as i32,
5536 n_ff as i32,
5537 n_expert as i32,
5538 rb_g as i64,
5539 rb_u as i64,
5540 );
5541 let (f, cfg) = match mode {
5542 "1" | "2" | "4" => {
5543 let rpw: u32 = mode.parse().unwrap();
5544 let f = self.func(match rpw {
5545 1 => "moe_gate_up_silu8_dev_q8_r1",
5546 2 => "moe_gate_up_silu8_dev_q8_r2",
5547 _ => "moe_gate_up_silu8_dev_q8_r4",
5548 });
5549 let rows_per_block = (rpw * wpb) as usize;
5550 let gx = n_ff.div_ceil(rows_per_block) as u32;
5551 (
5552 f,
5553 LaunchConfig {
5554 grid_dim: (gx, n_used as u32, 1),
5555 block_dim: (32, wpb, 1),
5556 shared_mem_bytes: 0,
5557 },
5558 )
5559 }
5560 "j8" if n_used <= 32 => (
5561 self.func("moe_gate_up_silu8_dev_q8_j8"),
5562 LaunchConfig {
5563 grid_dim: (n_ff as u32, 1, 1),
5564 block_dim: (32, n_used as u32, 1),
5565 shared_mem_bytes: 0,
5566 },
5567 ),
5568 "vsm2" => {
5570 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5571 let sh = (rb_g + rb_u) as u32;
5572 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5573 f.set_attribute(
5574 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5575 sh as i32,
5576 )?;
5577 (
5578 f,
5579 LaunchConfig {
5580 grid_dim: (n_ff as u32, n_used as u32, 1),
5581 block_dim: (32, 1, 1),
5582 shared_mem_bytes: sh,
5583 },
5584 )
5585 }
5586 "vsm" => {
5587 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5588 let sh = (rb_g + rb_u) as u32;
5589 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5590 f.set_attribute(
5591 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5592 sh as i32,
5593 )?;
5594 (
5595 f,
5596 LaunchConfig {
5597 grid_dim: (n_ff as u32, n_used as u32, 1),
5598 block_dim: (32, 1, 1),
5599 shared_mem_bytes: sh,
5600 },
5601 )
5602 }
5603 "sg" => (
5604 self.func("moe_gate_up_silu8_dev_q8_sg"),
5605 LaunchConfig {
5606 grid_dim: (n_ff as u32, n_used as u32, 1),
5607 block_dim: (32, 1, 1),
5608 shared_mem_bytes: 0,
5609 },
5610 ),
5611 "j8sg" if n_used <= 32 => (
5612 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5613 LaunchConfig {
5614 grid_dim: (n_ff as u32, 1, 1),
5615 block_dim: (32, n_used as u32, 1),
5616 shared_mem_bytes: 0,
5617 },
5618 ),
5619 "u64" if in_f == 2048 => (
5620 self.func("moe_gate_up_silu8_dev_q8_u64"),
5621 LaunchConfig {
5622 grid_dim: (n_ff as u32, n_used as u32, 1),
5623 block_dim: (32, 1, 1),
5624 shared_mem_bytes: 0,
5625 },
5626 ),
5627 "gs4" if in_f == 2048 => (
5628 self.func("moe_gate_up_silu8_dev_q8_gs4"),
5629 LaunchConfig {
5630 grid_dim: (n_ff as u32, n_used as u32, 1),
5631 block_dim: (32, 4, 1),
5632 shared_mem_bytes: 0,
5633 },
5634 ),
5635 "v" | "" => (
5637 self.func("moe_gate_up_silu8_dev_q8_v"),
5638 LaunchConfig {
5639 grid_dim: (n_ff as u32, n_used as u32, 1),
5640 block_dim: (32, 1, 1),
5641 shared_mem_bytes: 0,
5642 },
5643 ),
5644 "s2" => (
5645 self.func("moe_gate_up_silu8_dev_q8_s2"),
5646 LaunchConfig {
5647 grid_dim: (n_ff as u32, n_used as u32, 1),
5648 block_dim: (32, 2, 1),
5649 shared_mem_bytes: 0,
5650 },
5651 ),
5652 "s2z" => {
5653 let rz = wpb.min(16); (
5655 self.func("moe_gate_up_silu8_dev_q8_s2z"),
5656 LaunchConfig {
5657 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5658 block_dim: (32, 2, rz),
5659 shared_mem_bytes: 0,
5660 },
5661 )
5662 }
5663 _ => (
5664 self.func("moe_gate_up_silu8_dev_q8"),
5665 LaunchConfig {
5666 grid_dim: (n_ff as u32, n_used as u32, 1),
5667 block_dim: (32, 1, 1),
5668 shared_mem_bytes: 0,
5669 },
5670 ),
5671 };
5672 let __s_b = self.gpu.stream();
5673 let mut b = __s_b.launch_builder(&f);
5674 b.arg(table)
5675 .arg(sel)
5676 .arg(aq)
5677 .arg(ad)
5678 .arg(&mut act)
5679 .arg(&inf)
5680 .arg(&nff)
5681 .arg(&ne)
5682 .arg(&qt_g)
5683 .arg(&qt_u)
5684 .arg(&rbg)
5685 .arg(&rbu)
5686 .arg(macros);
5687 unsafe {
5688 b.launch(cfg)?;
5689 }
5690 Ok(act)
5691 }
5692
5693 #[allow(clippy::too_many_arguments)]
5694 pub fn moe_down8_fma_dev_q8(
5695 &self,
5696 table: &CudaSlice<u64>,
5697 sel: &cudarc::driver::CudaView<i32>,
5698 w: &cudarc::driver::CudaView<f32>,
5699 aq2: &CudaSlice<i8>,
5700 ad2: &CudaSlice<f32>,
5701 dst: &mut cudarc::driver::CudaViewMut<f32>,
5702 in_f: usize,
5703 out_f: usize,
5704 n_used: usize,
5705 n_expert: usize,
5706 qt: i32,
5707 rb: usize,
5708 ) -> Result<(), Box<dyn std::error::Error>> {
5709 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5710 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5711 let (inf, outf, nu, ne, rbi) = (
5712 in_f as i32,
5713 out_f as i32,
5714 n_used as i32,
5715 n_expert as i32,
5716 rb as i64,
5717 );
5718 let (f, cfg) = match mode.as_str() {
5721 m @ ("1" | "2" | "4") if n_used <= 8 => {
5722 let rpw: usize = m.parse().unwrap();
5723 let f = self.func(match rpw {
5724 1 => "moe_down8_fma_dev_q8_w8r1",
5725 2 => "moe_down8_fma_dev_q8_w8r2",
5726 _ => "moe_down8_fma_dev_q8_w8r4",
5727 });
5728 (
5729 f,
5730 LaunchConfig {
5731 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5732 block_dim: (32, n_used as u32, 1),
5733 shared_mem_bytes: 0,
5734 },
5735 )
5736 }
5737 "h2" if in_f == 512 => (
5738 self.func("moe_down8_fma_dev_q8_h2"),
5739 LaunchConfig {
5740 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5741 block_dim: (32, 1, 1),
5742 shared_mem_bytes: 0,
5743 },
5744 ),
5745 "" if in_f == 704 && n_used <= 8 => (
5748 self.func("moe_down8_fma_dev_q8_w8r2"),
5749 LaunchConfig {
5750 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5751 block_dim: (32, n_used as u32, 1),
5752 shared_mem_bytes: 0,
5753 },
5754 ),
5755 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
5759 self.func("moe_down8_fma_dev_q8_w8h2v"),
5760 LaunchConfig {
5761 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5762 block_dim: (32, n_used as u32, 1),
5763 shared_mem_bytes: 0,
5764 },
5765 ),
5766 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
5767 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
5768 LaunchConfig {
5769 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5770 block_dim: (32, n_used as u32, 1),
5771 shared_mem_bytes: 0,
5772 },
5773 ),
5774 "w8h2r2" if in_f == 512 && n_used <= 8 => (
5775 self.func("moe_down8_fma_dev_q8_w8h2r2"),
5776 LaunchConfig {
5777 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5778 block_dim: (32, n_used as u32, 1),
5779 shared_mem_bytes: 0,
5780 },
5781 ),
5782 "w8h2" if in_f == 512 && n_used <= 8 => (
5783 self.func("moe_down8_fma_dev_q8_w8h2"),
5784 LaunchConfig {
5785 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5786 block_dim: (32, n_used as u32, 1),
5787 shared_mem_bytes: 0,
5788 },
5789 ),
5790 _ => (
5791 self.func("moe_down8_fma_dev_q8"),
5792 LaunchConfig {
5793 grid_dim: (out_f as u32, 1, 1),
5794 block_dim: (32, 1, 1),
5795 shared_mem_bytes: 0,
5796 },
5797 ),
5798 };
5799 let __s_b = self.gpu.stream();
5800 let mut b = __s_b.launch_builder(&f);
5801 b.arg(table)
5802 .arg(sel)
5803 .arg(w)
5804 .arg(aq2)
5805 .arg(ad2)
5806 .arg(dst)
5807 .arg(&inf)
5808 .arg(&outf)
5809 .arg(&nu)
5810 .arg(&ne)
5811 .arg(&qt)
5812 .arg(&rbi);
5813 unsafe {
5814 b.launch(cfg)?;
5815 }
5816 Ok(())
5817 }
5818
5819 #[allow(clippy::too_many_arguments)]
5826 pub fn moe_gate_up_silu8_dev_q8_rows(
5827 &self,
5828 table: &CudaSlice<u64>,
5829 sel: &CudaSlice<i32>,
5830 aq: &CudaSlice<i8>,
5831 ad: &CudaSlice<f32>,
5832 t: usize,
5833 in_f: usize,
5834 n_ff: usize,
5835 n_used: usize,
5836 n_expert: usize,
5837 qt_g: i32,
5838 qt_u: i32,
5839 rb_g: usize,
5840 rb_u: usize,
5841 macros: &CudaSlice<f32>,
5842 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5843 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
5844 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5845 let cfg = LaunchConfig {
5846 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5847 block_dim: (32, 1, 1),
5848 shared_mem_bytes: 0,
5849 };
5850 let (inf, nff, ne, nu, rbg, rbu) = (
5851 in_f as i32,
5852 n_ff as i32,
5853 n_expert as i32,
5854 n_used as i32,
5855 rb_g as i64,
5856 rb_u as i64,
5857 );
5858 let __s_b = self.gpu.stream();
5859 let mut b = __s_b.launch_builder(&f);
5860 b.arg(table)
5861 .arg(sel)
5862 .arg(aq)
5863 .arg(ad)
5864 .arg(&mut act)
5865 .arg(&inf)
5866 .arg(&nff)
5867 .arg(&ne)
5868 .arg(&qt_g)
5869 .arg(&qt_u)
5870 .arg(&rbg)
5871 .arg(&rbu)
5872 .arg(&nu)
5873 .arg(macros);
5874 unsafe {
5875 b.launch(cfg)?;
5876 }
5877 Ok(act)
5878 }
5879
5880 #[allow(clippy::too_many_arguments)]
5885 pub fn moe_down8_fma_dev_q8_rows(
5886 &self,
5887 table: &CudaSlice<u64>,
5888 sel: &CudaSlice<i32>,
5889 w: &CudaSlice<f32>,
5890 aq2: &CudaSlice<i8>,
5891 ad2: &CudaSlice<f32>,
5892 dst: &mut CudaSlice<f32>,
5893 t: usize,
5894 in_f: usize,
5895 out_f: usize,
5896 n_used: usize,
5897 n_expert: usize,
5898 qt: i32,
5899 rb: usize,
5900 ) -> Result<(), Box<dyn std::error::Error>> {
5901 assert!(
5902 in_f == 512 && n_used <= 8,
5903 "down rows twin is w8h2v shape-gated"
5904 );
5905 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
5906 let cfg = LaunchConfig {
5907 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
5908 block_dim: (32, n_used as u32, 1),
5909 shared_mem_bytes: 0,
5910 };
5911 let (inf, outf, nu, ne, rbi) = (
5912 in_f as i32,
5913 out_f as i32,
5914 n_used as i32,
5915 n_expert as i32,
5916 rb as i64,
5917 );
5918 let __s_b = self.gpu.stream();
5919 let mut b = __s_b.launch_builder(&f);
5920 b.arg(table)
5921 .arg(sel)
5922 .arg(w)
5923 .arg(aq2)
5924 .arg(ad2)
5925 .arg(dst)
5926 .arg(&inf)
5927 .arg(&outf)
5928 .arg(&nu)
5929 .arg(&ne)
5930 .arg(&qt)
5931 .arg(&rbi);
5932 unsafe {
5933 b.launch(cfg)?;
5934 }
5935 Ok(())
5936 }
5937
5938 #[allow(clippy::too_many_arguments)]
5942 pub fn moe_gate_up_silu8_dev_q8_csr(
5943 &self,
5944 table: &CudaSlice<u64>,
5945 sel: &CudaSlice<i32>,
5946 aq: &CudaSlice<i8>,
5947 ad: &CudaSlice<f32>,
5948 n_pairs: usize,
5949 in_f: usize,
5950 n_ff: usize,
5951 n_used: usize,
5952 n_expert: usize,
5953 qt_g: i32,
5954 qt_u: i32,
5955 rb_g: usize,
5956 rb_u: usize,
5957 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5958 let f = if qt_g == crate::QT_NVFP4 {
5961 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
5962 } else {
5963 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
5964 };
5965 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5966 let cfg = LaunchConfig {
5967 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5968 block_dim: (32, 1, 1),
5969 shared_mem_bytes: 0,
5970 };
5971 let (inf, nff, ne, nu, npi, rbg, rbu) = (
5972 in_f as i32,
5973 n_ff as i32,
5974 n_expert as i32,
5975 n_used as i32,
5976 n_pairs as i32,
5977 rb_g as i64,
5978 rb_u as i64,
5979 );
5980 let __s_b = self.gpu.stream();
5981 let mut b = __s_b.launch_builder(&f);
5982 b.arg(table)
5983 .arg(sel)
5984 .arg(aq)
5985 .arg(ad)
5986 .arg(&mut act)
5987 .arg(&inf)
5988 .arg(&nff)
5989 .arg(&ne)
5990 .arg(&qt_g)
5991 .arg(&qt_u)
5992 .arg(&rbg)
5993 .arg(&rbu)
5994 .arg(&nu)
5995 .arg(&npi);
5996 unsafe {
5997 b.launch(cfg)?;
5998 }
5999 Ok(act)
6000 }
6001
6002 #[allow(clippy::too_many_arguments)]
6006 pub fn moe_down8_fma_dev_q8_variant(
6007 &self,
6008 variant: &str,
6009 table: &CudaSlice<u64>,
6010 sel: &cudarc::driver::CudaView<i32>,
6011 w: &cudarc::driver::CudaView<f32>,
6012 aq2: &CudaSlice<i8>,
6013 ad2: &CudaSlice<f32>,
6014 dst: &mut cudarc::driver::CudaViewMut<f32>,
6015 in_f: usize,
6016 out_f: usize,
6017 n_used: usize,
6018 n_expert: usize,
6019 qt: i32,
6020 rb: usize,
6021 ) -> Result<(), Box<dyn std::error::Error>> {
6022 let (inf, outf, nu, ne, rbi) = (
6023 in_f as i32,
6024 out_f as i32,
6025 n_used as i32,
6026 n_expert as i32,
6027 rb as i64,
6028 );
6029 let (f, cfg) = match variant {
6030 "w8h2" | "w8h2v" => (
6031 self.func(if variant == "w8h2" {
6032 "moe_down8_fma_dev_q8_w8h2"
6033 } else {
6034 "moe_down8_fma_dev_q8_w8h2v"
6035 }),
6036 LaunchConfig {
6037 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6038 block_dim: (32, n_used as u32, 1),
6039 shared_mem_bytes: 0,
6040 },
6041 ),
6042 "w8h2r2" | "w8h2r2v" => (
6043 self.func(if variant == "w8h2r2" {
6044 "moe_down8_fma_dev_q8_w8h2r2"
6045 } else {
6046 "moe_down8_fma_dev_q8_w8h2r2v"
6047 }),
6048 LaunchConfig {
6049 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6050 block_dim: (32, n_used as u32, 1),
6051 shared_mem_bytes: 0,
6052 },
6053 ),
6054 _ => (
6055 self.func("moe_down8_fma_dev_q8"),
6056 LaunchConfig {
6057 grid_dim: (out_f as u32, 1, 1),
6058 block_dim: (32, 1, 1),
6059 shared_mem_bytes: 0,
6060 },
6061 ),
6062 };
6063 let __s_b = self.gpu.stream();
6064 let mut b = __s_b.launch_builder(&f);
6065 b.arg(table)
6066 .arg(sel)
6067 .arg(w)
6068 .arg(aq2)
6069 .arg(ad2)
6070 .arg(dst)
6071 .arg(&inf)
6072 .arg(&outf)
6073 .arg(&nu)
6074 .arg(&ne)
6075 .arg(&qt)
6076 .arg(&rbi);
6077 unsafe {
6078 b.launch(cfg)?;
6079 }
6080 Ok(())
6081 }
6082
6083 #[allow(clippy::too_many_arguments)]
6085 pub fn moe_gate_up_silu8_dev_q8_variant(
6086 &self,
6087 variant: &str,
6088 table: &CudaSlice<u64>,
6089 sel: &cudarc::driver::CudaView<i32>,
6090 aq: &CudaSlice<i8>,
6091 ad: &CudaSlice<f32>,
6092 in_f: usize,
6093 n_ff: usize,
6094 n_used: usize,
6095 n_expert: usize,
6096 qt_g: i32,
6097 qt_u: i32,
6098 rb_g: usize,
6099 rb_u: usize,
6100 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6101 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6102 let (inf, nff, ne, rbg, rbu) = (
6103 in_f as i32,
6104 n_ff as i32,
6105 n_expert as i32,
6106 rb_g as i64,
6107 rb_u as i64,
6108 );
6109 let f = self.func(if variant == "v" {
6110 "moe_gate_up_silu8_dev_q8_v"
6111 } else {
6112 "moe_gate_up_silu8_dev_q8"
6113 });
6114 let cfg = LaunchConfig {
6115 grid_dim: (n_ff as u32, n_used as u32, 1),
6116 block_dim: (32, 1, 1),
6117 shared_mem_bytes: 0,
6118 };
6119 let __s_b = self.gpu.stream();
6120 let mut b = __s_b.launch_builder(&f);
6121 b.arg(table)
6122 .arg(sel)
6123 .arg(aq)
6124 .arg(ad)
6125 .arg(&mut act)
6126 .arg(&inf)
6127 .arg(&nff)
6128 .arg(&ne)
6129 .arg(&qt_g)
6130 .arg(&qt_u)
6131 .arg(&rbg)
6132 .arg(&rbu);
6133 unsafe {
6134 b.launch(cfg)?;
6135 }
6136 Ok(act)
6137 }
6138
6139 pub fn moe_gate_up_silu8_dev(
6140 &self,
6141 table: &CudaSlice<u64>,
6142 sel: &cudarc::driver::CudaView<i32>,
6143 x: &cudarc::driver::CudaView<f32>,
6144 in_f: usize,
6145 n_ff: usize,
6146 n_used: usize,
6147 n_expert: usize,
6148 qt_g: i32,
6149 qt_u: i32,
6150 rb_g: usize,
6151 rb_u: usize,
6152 macros: &CudaSlice<f32>,
6153 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6154 let f = self.func("moe_gate_up_silu8_dev");
6155 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6157 grid_dim: (n_ff as u32, n_used as u32, 1),
6158 block_dim: (256, 1, 1),
6159 shared_mem_bytes: 0,
6160 };
6161 let (inf, nff, ne, rbg, rbu) = (
6162 in_f as i32,
6163 n_ff as i32,
6164 n_expert as i32,
6165 rb_g as i64,
6166 rb_u as i64,
6167 );
6168 let __s_b = self.gpu.stream();
6169 let mut b = __s_b.launch_builder(&f);
6170 b.arg(table)
6171 .arg(sel)
6172 .arg(x)
6173 .arg(&mut act)
6174 .arg(&inf)
6175 .arg(&nff)
6176 .arg(&ne)
6177 .arg(&qt_g)
6178 .arg(&qt_u)
6179 .arg(&rbg)
6180 .arg(&rbu)
6181 .arg(macros);
6182 unsafe {
6183 b.launch(cfg)?;
6184 }
6185 Ok(act)
6186 }
6187
6188 #[allow(clippy::too_many_arguments)]
6191 pub fn moe_down8_fma_dev(
6192 &self,
6193 table: &CudaSlice<u64>,
6194 sel: &cudarc::driver::CudaView<i32>,
6195 w: &cudarc::driver::CudaView<f32>,
6196 act: &CudaSlice<f32>,
6197 dst: &mut cudarc::driver::CudaViewMut<f32>,
6198 in_f: usize,
6199 out_f: usize,
6200 n_used: usize,
6201 n_expert: usize,
6202 qt: i32,
6203 rb: usize,
6204 ) -> Result<(), Box<dyn std::error::Error>> {
6205 let f = self.func("moe_down8_fma_dev");
6206 let cfg = LaunchConfig {
6207 grid_dim: (out_f as u32, 1, 1),
6208 block_dim: (256, 1, 1),
6209 shared_mem_bytes: 0,
6210 };
6211 let (inf, outf, nu, ne, rbv) = (
6212 in_f as i32,
6213 out_f as i32,
6214 n_used as i32,
6215 n_expert as i32,
6216 rb as i64,
6217 );
6218 let __s_b = self.gpu.stream();
6219 let mut b = __s_b.launch_builder(&f);
6220 b.arg(table)
6221 .arg(sel)
6222 .arg(w)
6223 .arg(act)
6224 .arg(dst)
6225 .arg(&inf)
6226 .arg(&outf)
6227 .arg(&nu)
6228 .arg(&ne)
6229 .arg(&qt)
6230 .arg(&rbv);
6231 unsafe {
6232 b.launch(cfg)?;
6233 }
6234 Ok(())
6235 }
6236
6237 pub fn axpy_into(
6239 &self,
6240 src: &CudaSlice<f32>,
6241 alpha: f32,
6242 dst: &mut cudarc::driver::CudaViewMut<f32>,
6243 n: usize,
6244 ) -> Result<(), Box<dyn std::error::Error>> {
6245 let f = self.func("axpy_f32");
6246 let cfg = LaunchConfig::for_num_elems(n as u32);
6247 let (a, ni) = (alpha, n as i32);
6248 let __s_b = self.gpu.stream();
6249 let mut b = __s_b.launch_builder(&f);
6250 b.arg(src).arg(dst).arg(&a).arg(&ni);
6251 unsafe {
6252 b.launch(cfg)?;
6253 }
6254 Ok(())
6255 }
6256
6257 pub fn add_scaled_rows(
6259 &self,
6260 src: &CudaSlice<f32>,
6261 scale: &CudaSlice<f32>,
6262 dst: &mut CudaSlice<f32>,
6263 ncols: usize,
6264 nrows: usize,
6265 ) -> Result<(), Box<dyn std::error::Error>> {
6266 let f = self.func("add_scaled_rows_f32");
6267 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6268 let (nc, nr) = (ncols as i32, nrows as i32);
6269 let __s_b = self.gpu.stream();
6270 let mut b = __s_b.launch_builder(&f);
6271 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6272 unsafe {
6273 b.launch(cfg)?;
6274 }
6275 Ok(())
6276 }
6277
6278 pub fn gather_rows(
6282 &self,
6283 src: &CudaSlice<f32>,
6284 idx: &CudaSlice<i32>,
6285 dst: &mut CudaSlice<f32>,
6286 ncols: usize,
6287 m_e: usize,
6288 ) -> Result<(), Box<dyn std::error::Error>> {
6289 let f = self.func("gather_rows_f32");
6290 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6291 let (nc, me) = (ncols as i32, m_e as i32);
6292 let __s_b = self.gpu.stream();
6293 let mut b = __s_b.launch_builder(&f);
6294 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6295 unsafe {
6296 b.launch(cfg)?;
6297 }
6298 Ok(())
6299 }
6300
6301 pub fn scatter_slot(
6306 &self,
6307 src: &CudaSlice<f32>,
6308 tok_idx: &CudaSlice<i32>,
6309 slot_idx: &CudaSlice<i32>,
6310 weight: &CudaSlice<f32>,
6311 dst: &mut CudaSlice<f32>,
6312 wbuf: &mut CudaSlice<f32>,
6313 ncols: usize,
6314 n_used: usize,
6315 m_e: usize,
6316 ) -> Result<(), Box<dyn std::error::Error>> {
6317 let f = self.func("scatter_add_slot_f32");
6318 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6319 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6320 let __s_b = self.gpu.stream();
6321 let mut b = __s_b.launch_builder(&f);
6322 b.arg(src)
6323 .arg(tok_idx)
6324 .arg(slot_idx)
6325 .arg(weight)
6326 .arg(dst)
6327 .arg(wbuf)
6328 .arg(&nc)
6329 .arg(&nu)
6330 .arg(&me);
6331 unsafe {
6332 b.launch(cfg)?;
6333 }
6334 Ok(())
6335 }
6336
6337 pub fn reduce_slots(
6341 &self,
6342 slots: &CudaSlice<f32>,
6343 wbuf: &CudaSlice<f32>,
6344 dst: &mut CudaSlice<f32>,
6345 ncols: usize,
6346 n_used: usize,
6347 t: usize,
6348 ) -> Result<(), Box<dyn std::error::Error>> {
6349 let f = self.func("reduce_slots_f32");
6350 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6351 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6352 let __s_b = self.gpu.stream();
6353 let mut b = __s_b.launch_builder(&f);
6354 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6355 unsafe {
6356 b.launch(cfg)?;
6357 }
6358 Ok(())
6359 }
6360
6361 pub fn quantize_q8_1_view(
6368 &self,
6369 x: &cudarc::driver::CudaView<f32>,
6370 m: usize,
6371 in_f: usize,
6372 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6373 let f = self.func("quantize_q8_1");
6374 let nblk = in_f / 32;
6375 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6376 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6377 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6378 let (inf, mi) = (in_f as i32, m as i32);
6379 let __s_b = self.gpu.stream();
6380 let mut b = __s_b.launch_builder(&f);
6381 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6382 unsafe {
6383 b.launch(cfg)?;
6384 }
6385 Ok((q, d))
6386 }
6387
6388 pub fn quantize_q8_1(
6389 &self,
6390 x: &CudaSlice<f32>,
6391 m: usize,
6392 in_f: usize,
6393 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6394 let nblk = in_f / 32;
6395 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);
6399 let (inf, mi) = (in_f as i32, m as i32);
6400 if Self::pdl_on() && Self::pdl_wb_on() {
6401 {
6402 use cudarc::driver::{DevicePtr, DevicePtrMut};
6403 let s = &self.gpu.stream();
6404 let (px, _g0) = x.device_ptr(s);
6405 let (pq, _g1) = q.device_ptr_mut(s);
6406 let (pd, _g2) = d.device_ptr_mut(s);
6407 let mut ps = [
6408 &px as *const _ as *mut std::ffi::c_void,
6409 &pq as *const _ as *mut _,
6410 &pd as *const _ as *mut _,
6411 &inf as *const _ as *mut _,
6412 &mi as *const _ as *mut _,
6413 ];
6414 unsafe {
6415 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6416 }
6417 }
6418 return Ok((q, d));
6419 }
6420 let f = self.func("quantize_q8_1");
6421 let __s_b = self.gpu.stream();
6422 let mut b = __s_b.launch_builder(&f);
6423 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6424 unsafe {
6425 b.launch(cfg)?;
6426 }
6427 Ok((q, d))
6428 }
6429
6430 pub fn quantize_fp4_act(
6434 &self,
6435 x: &CudaSlice<f32>,
6436 m: usize,
6437 in_f: usize,
6438 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6439 let f = self.func("quantize_fp4_act");
6440 let nb16 = in_f / 16;
6441 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);
6444 let (inf, mi) = (in_f as i32, m as i32);
6445 let __s_b = self.gpu.stream();
6446 let mut b = __s_b.launch_builder(&f);
6447 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6448 unsafe {
6449 b.launch(cfg)?;
6450 }
6451 Ok((aq4, ad4))
6452 }
6453
6454 pub fn qmatvec_gemm_nvfp4_fp4(
6459 &self,
6460 bytes: &CudaSlice<u8>,
6461 x: &CudaSlice<f32>,
6462 m: usize,
6463 in_f: usize,
6464 out_f: usize,
6465 row_bytes: usize,
6466 scale: f32,
6467 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6468 assert!(
6469 in_f % 64 == 0,
6470 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6471 );
6472 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6473 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6474 if scale != 1.0 {
6475 self.scale_inplace(&mut y, scale, m * out_f)?;
6476 }
6477 Ok(y)
6478 }
6479
6480 fn fp4_gemm_launch(
6483 &self,
6484 bytes: &CudaSlice<u8>,
6485 aq4: &CudaSlice<u32>,
6486 ad4: &CudaSlice<u8>,
6487 m: usize,
6488 in_f: usize,
6489 out_f: usize,
6490 row_bytes: usize,
6491 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6492 let f = self.func("qmatvec_gemm_nvfp4_fp4");
6493 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
6495 const BN: u32 = 256;
6496 let cfg = LaunchConfig {
6497 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6498 block_dim: (32, 4, 1),
6499 shared_mem_bytes: 0,
6500 };
6501 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6502 let __s_b = self.gpu.stream();
6503 let mut b = __s_b.launch_builder(&f);
6504 b.arg(bytes)
6505 .arg(aq4)
6506 .arg(ad4)
6507 .arg(&mut y)
6508 .arg(&inf)
6509 .arg(&outf)
6510 .arg(&mi)
6511 .arg(&rb);
6512 unsafe {
6513 b.launch(cfg)?;
6514 }
6515 Ok(y)
6516 }
6517
6518 pub fn qmatvec_gemm_nvfp4_fp4_raw(
6520 &self,
6521 bytes: &CudaSlice<u8>,
6522 x: &CudaSlice<f32>,
6523 m: usize,
6524 in_f: usize,
6525 out_f: usize,
6526 row_bytes: usize,
6527 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6528 assert!(
6529 in_f % 64 == 0,
6530 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6531 );
6532 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6533 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6534 }
6535
6536 pub fn qmatvec_q8_0_fast(
6538 &self,
6539 w: &CudaSlice<u8>,
6540 x: &CudaSlice<f32>,
6541 m: usize,
6542 in_f: usize,
6543 out_f: usize,
6544 row_bytes: usize,
6545 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6546 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6547 let f = self.func("qmatvec_q8_0_dp4a");
6548 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6550 grid_dim: (out_f as u32, m as u32, 1),
6551 block_dim: (128, 1, 1),
6552 shared_mem_bytes: 0,
6553 };
6554 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6555 let __s_b = self.gpu.stream();
6556 let mut b = __s_b.launch_builder(&f);
6557 b.arg(w)
6558 .arg(&aq)
6559 .arg(&ad)
6560 .arg(&mut y)
6561 .arg(&inf)
6562 .arg(&outf)
6563 .arg(&mi)
6564 .arg(&rb);
6565 unsafe {
6566 b.launch(cfg)?;
6567 }
6568 Ok(y)
6569 }
6570
6571 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
6574 &self,
6575 w: &CudaSlice<u8>,
6576 x: &CudaSlice<f32>,
6577 m: usize,
6578 in_f: usize,
6579 out_f: usize,
6580 row_bytes: usize,
6581 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6582 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6583 let f = self.func("qmatvec_q4_K_dp4a");
6584 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6586 grid_dim: (out_f as u32, m as u32, 1),
6587 block_dim: (128, 1, 1),
6588 shared_mem_bytes: 0,
6589 };
6590 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6591 let __s_b = self.gpu.stream();
6592 let mut b = __s_b.launch_builder(&f);
6593 b.arg(w)
6594 .arg(&aq)
6595 .arg(&ad)
6596 .arg(&mut y)
6597 .arg(&inf)
6598 .arg(&outf)
6599 .arg(&mi)
6600 .arg(&rb);
6601 unsafe {
6602 b.launch(cfg)?;
6603 }
6604 Ok(y)
6605 }
6606
6607 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
6610 &self,
6611 w: &CudaSlice<u8>,
6612 x: &CudaSlice<f32>,
6613 m: usize,
6614 in_f: usize,
6615 out_f: usize,
6616 row_bytes: usize,
6617 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6618 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6619 let f = self.func("qmatvec_q6_K_dp4a");
6620 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6622 grid_dim: (out_f as u32, m as u32, 1),
6623 block_dim: (128, 1, 1),
6624 shared_mem_bytes: 0,
6625 };
6626 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6627 let __s_b = self.gpu.stream();
6628 let mut b = __s_b.launch_builder(&f);
6629 b.arg(w)
6630 .arg(&aq)
6631 .arg(&ad)
6632 .arg(&mut y)
6633 .arg(&inf)
6634 .arg(&outf)
6635 .arg(&mi)
6636 .arg(&rb);
6637 unsafe {
6638 b.launch(cfg)?;
6639 }
6640 Ok(y)
6641 }
6642
6643 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
6646 &self,
6647 w: &CudaSlice<u8>,
6648 x: &CudaSlice<f32>,
6649 m: usize,
6650 in_f: usize,
6651 out_f: usize,
6652 row_bytes: usize,
6653 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6654 self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6655 }
6656 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
6659 &self,
6660 w: &CudaSlice<u8>,
6661 x: &CudaSlice<f32>,
6662 m: usize,
6663 in_f: usize,
6664 out_f: usize,
6665 row_bytes: usize,
6666 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6667 self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6668 }
6669 pub fn qmatvec_nvfp4_fast_rp(
6671 &self,
6672 w: &CudaSlice<u8>,
6673 x: &CudaSlice<f32>,
6674 m: usize,
6675 in_f: usize,
6676 out_f: usize,
6677 row_bytes: usize,
6678 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6679 assert!(
6680 in_f % 64 == 0,
6681 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6682 );
6683 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6684 }
6685 pub fn qmatvec_nvfp4_fast(
6687 &self,
6688 w: &CudaSlice<u8>,
6689 x: &CudaSlice<f32>,
6690 m: usize,
6691 in_f: usize,
6692 out_f: usize,
6693 row_bytes: usize,
6694 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6695 assert!(
6698 in_f % 64 == 0,
6699 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6700 );
6701 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6702 }
6703 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
6706 &self,
6707 w: &CudaSlice<u8>,
6708 x: &CudaSlice<f32>,
6709 m: usize,
6710 in_f: usize,
6711 out_f: usize,
6712 row_bytes: usize,
6713 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6714 self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6715 }
6716
6717 fn qmatvec_dp4a_named(
6719 &self,
6720 name: &str,
6721 w: &CudaSlice<u8>,
6722 x: &CudaSlice<f32>,
6723 m: usize,
6724 in_f: usize,
6725 out_f: usize,
6726 row_bytes: usize,
6727 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6728 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6729 let f = self.func(name);
6730 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6732 grid_dim: (out_f as u32, m as u32, 1),
6733 block_dim: (128, 1, 1),
6734 shared_mem_bytes: 0,
6735 };
6736 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6737 let __s_b = self.gpu.stream();
6738 let mut b = __s_b.launch_builder(&f);
6739 b.arg(w)
6740 .arg(&aq)
6741 .arg(&ad)
6742 .arg(&mut y)
6743 .arg(&inf)
6744 .arg(&outf)
6745 .arg(&mi)
6746 .arg(&rb);
6747 unsafe {
6748 b.launch(cfg)?;
6749 }
6750 Ok(y)
6751 }
6752
6753 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6754 Ok(self.gpu.stream().clone_htod(v)?)
6755 }
6756 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6757 Ok(self.gpu.stream().clone_htod(v)?)
6758 }
6759 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6761 Ok(self.gpu.stream().clone_htod(v)?)
6762 }
6763 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
6764 Ok(self.gpu.stream().clone_htod(v)?)
6765 }
6766 pub fn dtoh_view(
6768 &self,
6769 d: &cudarc::driver::CudaView<f32>,
6770 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6771 let v = self.gpu.stream().clone_dtoh(d)?;
6772 self.gpu.stream().synchronize()?;
6773 Ok(v)
6774 }
6775 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6776 let v = self.gpu.stream().clone_dtoh(d)?;
6777 self.gpu.stream().synchronize()?;
6778 Ok(v)
6779 }
6780 pub fn dtoh_pair(
6784 &self,
6785 a: &CudaSlice<f32>,
6786 b: &CudaSlice<f32>,
6787 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
6788 let av = self.gpu.stream().clone_dtoh(a)?;
6789 let bv = self.gpu.stream().clone_dtoh(b)?;
6790 self.gpu.stream().synchronize()?;
6791 Ok((av, bv))
6792 }
6793 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
6795 let v = self.gpu.stream().clone_dtoh(d)?;
6796 self.gpu.stream().synchronize()?;
6797 Ok(v)
6798 }
6799 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6801 let v = self.gpu.stream().clone_dtoh(d)?;
6802 self.gpu.stream().synchronize()?;
6803 Ok(v)
6804 }
6805 pub fn dtoh_u8_view(
6806 &self,
6807 d: &cudarc::driver::CudaView<u8>,
6808 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6809 let v = self.gpu.stream().clone_dtoh(d)?;
6810 self.gpu.stream().synchronize()?;
6811 Ok(v)
6812 }
6813 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6814 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
6815 self.keep_if_capturing(&s);
6816 Ok(s)
6817 }
6818
6819 pub fn prob_of_token_device(
6828 &self,
6829 logits: &CudaSlice<f32>,
6830 tok: &CudaSlice<u32>,
6831 n_vocab: usize,
6832 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6833 let nb = ARGMAX_NB;
6834 let mut part = self.alloc_uninit::<f32>(nb)?;
6835 let mut p = self.alloc_uninit::<f32>(1)?;
6836 let f1 = self.func("prob_of_token_partial_f32");
6837 let cfg1 = LaunchConfig {
6838 grid_dim: (nb as u32, 1, 1),
6839 block_dim: (256, 1, 1),
6840 shared_mem_bytes: 0,
6841 };
6842 let nv = n_vocab as i32;
6843 let __s_b1 = self.gpu.stream();
6844 let mut b1 = __s_b1.launch_builder(&f1);
6845 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6846 unsafe {
6847 b1.launch(cfg1)?;
6848 }
6849 let f2 = self.func("prob_of_token_final_f32");
6850 let cfg2 = LaunchConfig {
6851 grid_dim: (1, 1, 1),
6852 block_dim: (256, 1, 1),
6853 shared_mem_bytes: 0,
6854 };
6855 let nbi = nb as i32;
6856 let __s_b2 = self.gpu.stream();
6857 let mut b2 = __s_b2.launch_builder(&f2);
6858 b2.arg(&part).arg(&mut p).arg(&nbi);
6859 unsafe {
6860 b2.launch(cfg2)?;
6861 }
6862 Ok(p)
6863 }
6864
6865 pub fn prob_of_token_device_col(
6872 &self,
6873 logits: &CudaSlice<f32>,
6874 tok_all: &CudaSlice<u32>,
6875 tok_idx: usize,
6876 p_out: &mut CudaSlice<f32>,
6877 p_idx: usize,
6878 n_vocab: usize,
6879 ) -> Result<(), Box<dyn std::error::Error>> {
6880 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
6881 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
6882 let nb = ARGMAX_NB;
6883 let mut part = self.alloc_uninit::<f32>(nb)?;
6884 let f1 = self.func("prob_of_token_partial_f32");
6885 let cfg1 = LaunchConfig {
6886 grid_dim: (nb as u32, 1, 1),
6887 block_dim: (256, 1, 1),
6888 shared_mem_bytes: 0,
6889 };
6890 let nv = n_vocab as i32;
6891 let __s_b1 = self.gpu.stream();
6892 let mut b1 = __s_b1.launch_builder(&f1);
6893 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
6894 unsafe {
6895 b1.launch(cfg1)?;
6896 }
6897 let f2 = self.func("prob_of_token_final_f32");
6898 let cfg2 = LaunchConfig {
6899 grid_dim: (1, 1, 1),
6900 block_dim: (256, 1, 1),
6901 shared_mem_bytes: 0,
6902 };
6903 let nbi = nb as i32;
6904 let __s_b2 = self.gpu.stream();
6905 let mut b2 = __s_b2.launch_builder(&f2);
6906 b2.arg(&part).arg(&mut p_v).arg(&nbi);
6907 unsafe {
6908 b2.launch(cfg2)?;
6909 }
6910 Ok(())
6911 }
6912
6913 pub fn prob_of_token_device_into(
6914 &self,
6915 logits: &CudaSlice<f32>,
6916 tok: &CudaSlice<u32>,
6917 p_out: &mut CudaSlice<f32>,
6918 n_vocab: usize,
6919 ) -> Result<(), Box<dyn std::error::Error>> {
6920 let nb = ARGMAX_NB;
6921 let mut part = self.alloc_uninit::<f32>(nb)?;
6922 let f1 = self.func("prob_of_token_partial_f32");
6923 let cfg1 = LaunchConfig {
6924 grid_dim: (nb as u32, 1, 1),
6925 block_dim: (256, 1, 1),
6926 shared_mem_bytes: 0,
6927 };
6928 let nv = n_vocab as i32;
6929 let __s_b1 = self.gpu.stream();
6930 let mut b1 = __s_b1.launch_builder(&f1);
6931 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6932 unsafe {
6933 b1.launch(cfg1)?;
6934 }
6935 let f2 = self.func("prob_of_token_final_f32");
6936 let cfg2 = LaunchConfig {
6937 grid_dim: (1, 1, 1),
6938 block_dim: (256, 1, 1),
6939 shared_mem_bytes: 0,
6940 };
6941 let nbi = nb as i32;
6942 let __s_b2 = self.gpu.stream();
6943 let mut b2 = __s_b2.launch_builder(&f2);
6944 b2.arg(&part).arg(p_out).arg(&nbi);
6945 unsafe {
6946 b2.launch(cfg2)?;
6947 }
6948 Ok(())
6949 }
6950
6951 pub fn argmax_token_device(
6952 &self,
6953 logits: &CudaSlice<f32>,
6954 n_vocab: usize,
6955 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6956 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
6957 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
6958 Ok(tok)
6959 }
6960 pub fn argmax_token_device_into(
6967 &self,
6968 logits: &CudaSlice<f32>,
6969 tok: &mut CudaSlice<u32>,
6970 n_vocab: usize,
6971 ) -> Result<(), Box<dyn std::error::Error>> {
6972 let nb = ARGMAX_NB;
6973 let f1 = self.func("argmax_partial_f32");
6974 let f2 = self.func("argmax_final_f32");
6975 let mut guard = self.argmax_partials.lock().unwrap();
6976 if guard.is_none() {
6977 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
6980 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
6981 *guard = Some((pv, pi));
6982 }
6983 let (part_v, part_i) = guard.as_mut().unwrap();
6984 let nv = n_vocab as i32;
6985 let nbi = nb as i32;
6986 let cfg1 = LaunchConfig {
6988 grid_dim: (nb as u32, 1, 1),
6989 block_dim: (256, 1, 1),
6990 shared_mem_bytes: 0,
6991 };
6992 let __s_b1 = self.gpu.stream();
6993 let mut b1 = __s_b1.launch_builder(&f1);
6994 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
6995 unsafe {
6996 b1.launch(cfg1)?;
6997 }
6998 let cfg2 = LaunchConfig {
7000 grid_dim: (1, 1, 1),
7001 block_dim: (256, 1, 1),
7002 shared_mem_bytes: 0,
7003 };
7004 let __s_b2 = self.gpu.stream();
7005 let mut b2 = __s_b2.launch_builder(&f2);
7006 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
7007 unsafe {
7008 b2.launch(cfg2)?;
7009 }
7010 Ok(())
7011 }
7012 pub fn argmax_token_device_col(
7018 &self,
7019 logits: &CudaSlice<f32>,
7020 col: usize,
7021 n_vocab: usize,
7022 toks: &mut CudaSlice<u32>,
7023 out_idx: usize,
7024 ) -> Result<(), Box<dyn std::error::Error>> {
7025 let nb = ARGMAX_NB;
7026 let f1 = self.func("argmax_partial_f32");
7027 let f2 = self.func("argmax_final_f32");
7028 let mut guard = self.argmax_partials.lock().unwrap();
7029 if guard.is_none() {
7030 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7031 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7032 *guard = Some((pv, pi));
7033 }
7034 let (part_v, part_i) = guard.as_mut().unwrap();
7035 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
7036 let nv = n_vocab as i32;
7037 let nbi = nb as i32;
7038 let cfg1 = LaunchConfig {
7039 grid_dim: (nb as u32, 1, 1),
7040 block_dim: (256, 1, 1),
7041 shared_mem_bytes: 0,
7042 };
7043 let __s_b1 = self.gpu.stream();
7044 let mut b1 = __s_b1.launch_builder(&f1);
7045 b1.arg(&col_view)
7046 .arg(&mut *part_v)
7047 .arg(&mut *part_i)
7048 .arg(&nv);
7049 unsafe {
7050 b1.launch(cfg1)?;
7051 }
7052 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
7053 let cfg2 = LaunchConfig {
7054 grid_dim: (1, 1, 1),
7055 block_dim: (256, 1, 1),
7056 shared_mem_bytes: 0,
7057 };
7058 let __s_b2 = self.gpu.stream();
7059 let mut b2 = __s_b2.launch_builder(&f2);
7060 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
7061 unsafe {
7062 b2.launch(cfg2)?;
7063 }
7064 Ok(())
7065 }
7066 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7068 Ok(self.gpu.stream().clone_htod(v)?)
7069 }
7070 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7071 let v = self.gpu.stream().clone_dtoh(d)?;
7072 self.gpu.stream().synchronize()?;
7073 Ok(v)
7074 }
7075 pub fn htod_u32_into(
7079 &self,
7080 dst: &mut CudaSlice<u32>,
7081 src: &[u32],
7082 ) -> Result<(), Box<dyn std::error::Error>> {
7083 let mut view = dst.slice_mut(0..src.len());
7084 self.gpu.stream().memcpy_htod(src, &mut view)?;
7085 Ok(())
7086 }
7087
7088 pub fn htod_i32_into(
7091 &self,
7092 dst: &mut CudaSlice<i32>,
7093 src: &[i32],
7094 ) -> Result<(), Box<dyn std::error::Error>> {
7095 let mut view = dst.slice_mut(0..src.len());
7096 self.gpu.stream().memcpy_htod(src, &mut view)?;
7097 Ok(())
7098 }
7099
7100 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7101 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
7102 self.keep_if_capturing(&s);
7103 Ok(s)
7104 }
7105 pub fn embed_gather_device_into(
7108 &self,
7109 embd: &CudaSlice<u8>,
7110 token_d: &CudaSlice<u32>,
7111 x_out: &mut CudaSlice<f32>,
7112 n_embd: usize,
7113 qtype: i32,
7114 row_bytes: usize,
7115 ) -> Result<(), Box<dyn std::error::Error>> {
7116 let f = self.func("embed_gather_u32");
7117 let cfg = LaunchConfig {
7118 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7119 block_dim: (256, 1, 1),
7120 shared_mem_bytes: 0,
7121 };
7122 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7123 let __s_b = self.gpu.stream();
7124 let mut b = __s_b.launch_builder(&f);
7125 b.arg(embd)
7126 .arg(token_d)
7127 .arg(x_out)
7128 .arg(&ne)
7129 .arg(&qt)
7130 .arg(&rb);
7131 unsafe {
7132 b.launch(cfg)?;
7133 }
7134 Ok(())
7135 }
7136 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7138 let v = self.gpu.stream().clone_dtoh(d)?;
7139 self.gpu.stream().synchronize()?;
7140 Ok(v[0])
7141 }
7142 pub fn i32_set_k(
7149 &self,
7150 dst: &mut CudaSlice<i32>,
7151 v: i32,
7152 ) -> Result<(), Box<dyn std::error::Error>> {
7153 let f = self.func("i32_set_k");
7154 let cfg = LaunchConfig {
7155 grid_dim: (1, 1, 1),
7156 block_dim: (1, 1, 1),
7157 shared_mem_bytes: 0,
7158 };
7159 let idx = 0i32;
7160 let __s_b = self.gpu.stream();
7161 let mut b = __s_b.launch_builder(&f);
7162 b.arg(dst).arg(&v).arg(&idx);
7163 unsafe {
7164 b.launch(cfg)?;
7165 }
7166 Ok(())
7167 }
7168
7169 pub fn set_i32_one(
7170 &self,
7171 d: &mut CudaSlice<i32>,
7172 v: i32,
7173 ) -> Result<(), Box<dyn std::error::Error>> {
7174 self.gpu.stream().memcpy_htod(&[v], d)?;
7175 Ok(())
7176 }
7177 pub fn set_u32_one(
7180 &self,
7181 d: &mut CudaSlice<u32>,
7182 v: u32,
7183 ) -> Result<(), Box<dyn std::error::Error>> {
7184 self.gpu.stream().memcpy_htod(&[v], d)?;
7185 Ok(())
7186 }
7187 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7189 let v = self.gpu.stream().clone_dtoh(d)?;
7190 self.gpu.stream().synchronize()?;
7191 Ok(v[0])
7192 }
7193 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7195 Ok(self.gpu.stream().clone_htod(bytes)?)
7196 }
7197 pub fn embed_gather_device(
7201 &self,
7202 embd: &CudaSlice<u8>,
7203 token_d: &CudaSlice<u32>,
7204 n_embd: usize,
7205 qtype: i32,
7206 row_bytes: usize,
7207 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7208 let f = self.func("embed_gather_u32");
7209 let mut x = self.alloc_uninit::<f32>(n_embd)?;
7210 let cfg = LaunchConfig {
7211 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7212 block_dim: (256, 1, 1),
7213 shared_mem_bytes: 0,
7214 };
7215 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7216 let __s_b = self.gpu.stream();
7217 let mut b = __s_b.launch_builder(&f);
7218 b.arg(embd)
7219 .arg(token_d)
7220 .arg(&mut x)
7221 .arg(&ne)
7222 .arg(&qt)
7223 .arg(&rb);
7224 unsafe {
7225 b.launch(cfg)?;
7226 }
7227 Ok(x)
7228 }
7229
7230 pub fn embed_gather_device_t(
7234 &self,
7235 embd: &CudaSlice<u8>,
7236 tokens: &[u32],
7237 n_embd: usize,
7238 qtype: i32,
7239 row_bytes: usize,
7240 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7241 let t = tokens.len();
7242 let tok_d = self.gpu.stream().clone_htod(tokens)?;
7243 let f = self.func("embed_gather_u32_t");
7244 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7245 let cfg = LaunchConfig {
7246 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7247 block_dim: (256, 1, 1),
7248 shared_mem_bytes: 0,
7249 };
7250 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7251 let __s_b = self.gpu.stream();
7252 let mut b = __s_b.launch_builder(&f);
7253 b.arg(embd)
7254 .arg(&tok_d)
7255 .arg(&mut x)
7256 .arg(&ne)
7257 .arg(&qt)
7258 .arg(&rb)
7259 .arg(&ti);
7260 unsafe {
7261 b.launch(cfg)?;
7262 }
7263 Ok(x)
7264 }
7265
7266 pub fn embed_gather_device_tv(
7271 &self,
7272 embd: &CudaSlice<u8>,
7273 tok_v: &cudarc::driver::CudaView<u32>,
7274 t: usize,
7275 n_embd: usize,
7276 qtype: i32,
7277 row_bytes: usize,
7278 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7279 let f = self.func("embed_gather_u32_t");
7280 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7281 let cfg = LaunchConfig {
7282 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7283 block_dim: (256, 1, 1),
7284 shared_mem_bytes: 0,
7285 };
7286 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7287 let __s_b = self.gpu.stream();
7288 let mut b = __s_b.launch_builder(&f);
7289 b.arg(embd)
7290 .arg(tok_v)
7291 .arg(&mut x)
7292 .arg(&ne)
7293 .arg(&qt)
7294 .arg(&rb)
7295 .arg(&ti);
7296 unsafe {
7297 b.launch(cfg)?;
7298 }
7299 Ok(x)
7300 }
7301
7302 pub fn embed_gather_device_td(
7303 &self,
7304 embd: &CudaSlice<u8>,
7305 tok_d: &CudaSlice<u32>,
7306 t: usize,
7307 n_embd: usize,
7308 qtype: i32,
7309 row_bytes: usize,
7310 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7311 let f = self.func("embed_gather_u32_t");
7312 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7313 let cfg = LaunchConfig {
7314 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7315 block_dim: (256, 1, 1),
7316 shared_mem_bytes: 0,
7317 };
7318 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7319 let __s_b = self.gpu.stream();
7320 let mut b = __s_b.launch_builder(&f);
7321 b.arg(embd)
7322 .arg(tok_d)
7323 .arg(&mut x)
7324 .arg(&ne)
7325 .arg(&qt)
7326 .arg(&rb)
7327 .arg(&ti);
7328 unsafe {
7329 b.launch(cfg)?;
7330 }
7331 Ok(x)
7332 }
7333
7334 #[inline]
7340 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7342 if self
7343 .capture_keep_on
7344 .load(std::sync::atomic::Ordering::Relaxed)
7345 {
7346 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7347 }
7348 }
7349
7350 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7351 &self,
7352 n: usize,
7353 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7354 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7355 {
7359 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7360 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7361 use cudarc::driver::DevicePtrMut;
7363 let n_bytes = s.len() * std::mem::size_of::<T>();
7364 let stream = self.gpu.stream();
7365 let (p_, _g) = s.device_ptr_mut(&stream);
7366 unsafe {
7367 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7368 .result()?;
7369 }
7370 }
7371 }
7372 self.keep_if_capturing(&s);
7373 Ok(s)
7374 }
7375
7376 pub fn uninit_q8_pair(
7381 &self,
7382 n: usize,
7383 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7384 Ok((
7385 self.alloc_uninit::<i8>(n)?,
7386 self.alloc_uninit::<f32>(n / 32)?,
7387 ))
7388 }
7389
7390 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7391 self.alloc_uninit::<f32>(n)
7392 }
7393
7394 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7396 self.alloc_uninit::<i8>(n)
7397 }
7398
7399 #[allow(clippy::too_many_arguments)]
7403 pub fn rms_norm3(
7404 &self,
7405 x: &CudaSlice<f32>,
7406 w0: &CudaSlice<f32>,
7407 w1: &CudaSlice<f32>,
7408 w2: &CudaSlice<f32>,
7409 d0: &mut CudaSlice<f32>,
7410 d1: &mut CudaSlice<f32>,
7411 d2: &mut CudaSlice<f32>,
7412 ncols: usize,
7413 nrows: usize,
7414 eps: f32,
7415 ) -> Result<(), Box<dyn std::error::Error>> {
7416 let f = self.func("rms_norm3_f32");
7417 let cfg = LaunchConfig {
7418 grid_dim: (nrows as u32, 1, 1),
7419 block_dim: (rms_block(), 1, 1),
7420 shared_mem_bytes: 0,
7421 };
7422 let (nc, e) = (ncols as i32, eps);
7423 let __s_b = self.gpu.stream();
7424 let mut b = __s_b.launch_builder(&f);
7425 b.arg(x)
7426 .arg(w0)
7427 .arg(w1)
7428 .arg(w2)
7429 .arg(d0)
7430 .arg(d1)
7431 .arg(d2)
7432 .arg(&nc)
7433 .arg(&e);
7434 unsafe {
7435 b.launch(cfg)?;
7436 }
7437 Ok(())
7438 }
7439
7440 #[allow(clippy::too_many_arguments)]
7442 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7445 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7446 *WARP_ON.get_or_init(|| {
7447 std::env::var("MEMRA_QKVNORM_W")
7448 .map(|v| v != "0")
7449 .unwrap_or(true)
7450 }) && ncols % 4 == 0
7451 && rows >= 64
7452 }
7453
7454 #[allow(clippy::too_many_arguments)]
7457 pub fn rms_norm_qkv_w4b(
7458 &self,
7459 q: &CudaSlice<f32>,
7460 k: &CudaSlice<f32>,
7461 v: &CudaSlice<f32>,
7462 wq: &CudaSlice<f32>,
7463 wk: &CudaSlice<f32>,
7464 wv: &CudaSlice<f32>,
7465 dq: &mut CudaSlice<f32>,
7466 dk: &mut CudaSlice<f32>,
7467 dv: &mut CudaSlice<f32>,
7468 dvb: &mut CudaSlice<u8>,
7469 ncols: usize,
7470 rq: usize,
7471 rk: usize,
7472 eps: f32,
7473 vf16: bool,
7474 ) -> Result<(), Box<dyn std::error::Error>> {
7475 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7476 let f = self.func("rms_norm_qkv_w4b_f32");
7477 let rows = (rq + 2 * rk) as u32;
7478 let cfg = LaunchConfig {
7479 grid_dim: (rows.div_ceil(8), 1, 1),
7480 block_dim: (256, 1, 1),
7481 shared_mem_bytes: 0,
7482 };
7483 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7484 let vf = vf16 as i32;
7485 let __s_b = self.gpu.stream();
7486 let mut b = __s_b.launch_builder(&f);
7487 b.arg(q)
7488 .arg(k)
7489 .arg(v)
7490 .arg(wq)
7491 .arg(wk)
7492 .arg(wv)
7493 .arg(dq)
7494 .arg(dk)
7495 .arg(dv)
7496 .arg(&mut *dvb)
7497 .arg(&nc)
7498 .arg(&rqi)
7499 .arg(&rki)
7500 .arg(&rvi)
7501 .arg(&e)
7502 .arg(&vf);
7503 unsafe {
7504 b.launch(cfg)?;
7505 }
7506 Ok(())
7507 }
7508
7509 pub fn rms_norm_qkv(
7510 &self,
7511 q: &CudaSlice<f32>,
7512 k: &CudaSlice<f32>,
7513 v: &CudaSlice<f32>,
7514 wq: &CudaSlice<f32>,
7515 wk: &CudaSlice<f32>,
7516 wv: &CudaSlice<f32>,
7517 dq: &mut CudaSlice<f32>,
7518 dk: &mut CudaSlice<f32>,
7519 dv: &mut CudaSlice<f32>,
7520 ncols: usize,
7521 rq: usize,
7522 rk: usize,
7523 eps: f32,
7524 ) -> Result<(), Box<dyn std::error::Error>> {
7525 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7529 let warp_on = *WARP_ON.get_or_init(|| {
7530 std::env::var("MEMRA_QKVNORM_W")
7531 .map(|v| v != "0")
7532 .unwrap_or(true)
7533 });
7534 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7537 let f = self.func("rms_norm_qkv_w4_f32");
7538 let rows = (rq + 2 * rk) as u32;
7539 let cfg = LaunchConfig {
7540 grid_dim: (rows.div_ceil(8), 1, 1),
7541 block_dim: (256, 1, 1),
7542 shared_mem_bytes: 0,
7543 };
7544 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7545 let __s_b = self.gpu.stream();
7546 let mut b = __s_b.launch_builder(&f);
7547 b.arg(q)
7548 .arg(k)
7549 .arg(v)
7550 .arg(wq)
7551 .arg(wk)
7552 .arg(wv)
7553 .arg(dq)
7554 .arg(dk)
7555 .arg(dv)
7556 .arg(&nc)
7557 .arg(&rqi)
7558 .arg(&rki)
7559 .arg(&rvi)
7560 .arg(&e);
7561 unsafe {
7562 b.launch(cfg)?;
7563 }
7564 return Ok(());
7565 }
7566 let f = self.func("rms_norm_qkv_f32");
7567 let grid = (rq + 2 * rk) as u32;
7568 let cfg = LaunchConfig {
7569 grid_dim: (grid, 1, 1),
7570 block_dim: (rms_block(), 1, 1),
7571 shared_mem_bytes: 0,
7572 };
7573 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7574 let __s_b = self.gpu.stream();
7575 let mut b = __s_b.launch_builder(&f);
7576 b.arg(q)
7577 .arg(k)
7578 .arg(v)
7579 .arg(wq)
7580 .arg(wk)
7581 .arg(wv)
7582 .arg(dq)
7583 .arg(dk)
7584 .arg(dv)
7585 .arg(&nc)
7586 .arg(&rqi)
7587 .arg(&rki)
7588 .arg(&e);
7589 unsafe {
7590 b.launch(cfg)?;
7591 }
7592 Ok(())
7593 }
7594
7595 #[allow(clippy::too_many_arguments)]
7597 pub fn rms_norm2x(
7598 &self,
7599 a: &CudaSlice<f32>,
7600 bb: &CudaSlice<f32>,
7601 wa: &CudaSlice<f32>,
7602 wb: &CudaSlice<f32>,
7603 da: &mut CudaSlice<f32>,
7604 db: &mut CudaSlice<f32>,
7605 ncols: usize,
7606 nrows: usize,
7607 eps: f32,
7608 ) -> Result<(), Box<dyn std::error::Error>> {
7609 let f = self.func("rms_norm2x_f32");
7610 let cfg = LaunchConfig {
7611 grid_dim: (2 * nrows as u32, 1, 1),
7612 block_dim: (rms_block(), 1, 1),
7613 shared_mem_bytes: 0,
7614 };
7615 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7616 let __s_b = self.gpu.stream();
7617 let mut b = __s_b.launch_builder(&f);
7618 b.arg(a)
7619 .arg(bb)
7620 .arg(wa)
7621 .arg(wb)
7622 .arg(da)
7623 .arg(db)
7624 .arg(&nc)
7625 .arg(&nr)
7626 .arg(&e);
7627 unsafe {
7628 b.launch(cfg)?;
7629 }
7630 Ok(())
7631 }
7632
7633 pub fn softcap(
7635 &self,
7636 y: &mut CudaSlice<f32>,
7637 cap: f32,
7638 n: usize,
7639 ) -> Result<(), Box<dyn std::error::Error>> {
7640 let f = self.func("softcap_f32");
7641 let cfg = LaunchConfig::for_num_elems(n as u32);
7642 let ni = n as i32;
7643 let __s_b = self.gpu.stream();
7644 let mut b = __s_b.launch_builder(&f);
7645 b.arg(y).arg(&cap).arg(&ni);
7646 unsafe {
7647 b.launch(cfg)?;
7648 }
7649 Ok(())
7650 }
7651
7652 pub fn mask_ids_rows(
7655 &self,
7656 y: &mut CudaSlice<f32>,
7657 ids: &CudaSlice<i32>,
7658 n_ids: usize,
7659 n_vocab: usize,
7660 t: usize,
7661 ) -> Result<(), Box<dyn std::error::Error>> {
7662 let f = self.func("mask_ids_rows_f32");
7663 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7664 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7665 let __s_b = self.gpu.stream();
7666 let mut b = __s_b.launch_builder(&f);
7667 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7668 unsafe {
7669 b.launch(cfg)?;
7670 }
7671 Ok(())
7672 }
7673
7674 #[allow(clippy::too_many_arguments)]
7676 pub fn add_scale_rms_norm(
7677 &self,
7678 a: &CudaSlice<f32>,
7679 b_in: &CudaSlice<f32>,
7680 c: f32,
7681 w: &CudaSlice<f32>,
7682 res: &mut CudaSlice<f32>,
7683 dst: &mut CudaSlice<f32>,
7684 ncols: usize,
7685 nrows: usize,
7686 eps: f32,
7687 ) -> Result<(), Box<dyn std::error::Error>> {
7688 let f = self.func("add_scale_rms_norm_f32");
7689 let cfg = LaunchConfig {
7690 grid_dim: (nrows as u32, 1, 1),
7691 block_dim: (rms_block(), 1, 1),
7692 shared_mem_bytes: 0,
7693 };
7694 let (nc, e2) = (ncols as i32, eps);
7695 let __s_b = self.gpu.stream();
7696 let mut b = __s_b.launch_builder(&f);
7697 b.arg(a)
7698 .arg(b_in)
7699 .arg(&c)
7700 .arg(w)
7701 .arg(res)
7702 .arg(dst)
7703 .arg(&nc)
7704 .arg(&e2);
7705 unsafe {
7706 b.launch(cfg)?;
7707 }
7708 Ok(())
7709 }
7710
7711 #[allow(clippy::too_many_arguments)]
7714 pub fn add_scale_rms_norm_q8_1(
7715 &self,
7716 a: &CudaSlice<f32>,
7717 b_in: &CudaSlice<f32>,
7718 c: f32,
7719 w: &CudaSlice<f32>,
7720 res: &mut CudaSlice<f32>,
7721 ncols: usize,
7722 nrows: usize,
7723 eps: f32,
7724 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7725 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7726 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7727 let (nc, e2) = (ncols as i32, eps);
7728 if Self::pdl_on() && Self::pdl_wb_on() {
7729 {
7730 use cudarc::driver::{DevicePtr, DevicePtrMut};
7731 let s = &self.gpu.stream();
7732 let (pa, _g0) = a.device_ptr(s);
7733 let (pb, _g1) = b_in.device_ptr(s);
7734 let (pw, _g2) = w.device_ptr(s);
7735 let (pr, _g3) = res.device_ptr_mut(s);
7736 let (pq, _g4) = out_q.device_ptr_mut(s);
7737 let (pd, _g5) = out_d.device_ptr_mut(s);
7738 let mut ps = [
7739 &pa as *const _ as *mut std::ffi::c_void,
7740 &pb as *const _ as *mut _,
7741 &c as *const _ as *mut _,
7742 &pw as *const _ as *mut _,
7743 &pr as *const _ as *mut _,
7744 &pq as *const _ as *mut _,
7745 &pd as *const _ as *mut _,
7746 &nc as *const _ as *mut _,
7747 &e2 as *const _ as *mut _,
7748 ];
7749 unsafe {
7750 self.launch_pdl(
7751 "add_scale_rms_norm_q8_1",
7752 (nrows as u32, 1, 1),
7753 (rms_block(), 1, 1),
7754 &mut ps,
7755 )?;
7756 }
7757 }
7758 return Ok((out_q, out_d));
7759 }
7760 let f = self.func("add_scale_rms_norm_q8_1");
7761 let cfg = LaunchConfig {
7762 grid_dim: (nrows as u32, 1, 1),
7763 block_dim: (rms_block(), 1, 1),
7764 shared_mem_bytes: 0,
7765 };
7766 let __s_b = self.gpu.stream();
7767 let mut b = __s_b.launch_builder(&f);
7768 b.arg(a)
7769 .arg(b_in)
7770 .arg(&c)
7771 .arg(w)
7772 .arg(res)
7773 .arg(&mut out_q)
7774 .arg(&mut out_d)
7775 .arg(&nc)
7776 .arg(&e2);
7777 unsafe {
7778 b.launch(cfg)?;
7779 }
7780 Ok((out_q, out_d))
7781 }
7782
7783 #[allow(clippy::too_many_arguments)]
7785 pub fn add_scale_rms_norm_q8_1_into(
7786 &self,
7787 a: &CudaSlice<f32>,
7788 b_in: &CudaSlice<f32>,
7789 c: f32,
7790 w: &CudaSlice<f32>,
7791 res: &mut CudaSlice<f32>,
7792 ncols: usize,
7793 nrows: usize,
7794 eps: f32,
7795 out_q: &mut CudaSlice<i8>,
7796 out_d: &mut CudaSlice<f32>,
7797 ) -> Result<(), Box<dyn std::error::Error>> {
7798 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
7799 let (nc, e2) = (ncols as i32, eps);
7800 if Self::pdl_on() && Self::pdl_wb_on() {
7801 use cudarc::driver::{DevicePtr, DevicePtrMut};
7802 let s = &self.gpu.stream();
7803 let (pa, _g0) = a.device_ptr(s);
7804 let (pb, _g1) = b_in.device_ptr(s);
7805 let (pw, _g2) = w.device_ptr(s);
7806 let (pr, _g3) = res.device_ptr_mut(s);
7807 let (pq, _g4) = out_q.device_ptr_mut(s);
7808 let (pd, _g5) = out_d.device_ptr_mut(s);
7809 let mut ps = [
7810 &pa as *const _ as *mut std::ffi::c_void,
7811 &pb as *const _ as *mut _,
7812 &c as *const _ as *mut _,
7813 &pw as *const _ as *mut _,
7814 &pr as *const _ as *mut _,
7815 &pq as *const _ as *mut _,
7816 &pd as *const _ as *mut _,
7817 &nc as *const _ as *mut _,
7818 &e2 as *const _ as *mut _,
7819 ];
7820 unsafe {
7821 self.launch_pdl(
7822 "add_scale_rms_norm_q8_1",
7823 (nrows as u32, 1, 1),
7824 (rms_block(), 1, 1),
7825 &mut ps,
7826 )?;
7827 }
7828 return Ok(());
7829 }
7830 let f = self.func("add_scale_rms_norm_q8_1");
7831 let cfg = LaunchConfig {
7832 grid_dim: (nrows as u32, 1, 1),
7833 block_dim: (rms_block(), 1, 1),
7834 shared_mem_bytes: 0,
7835 };
7836 let __s_b = self.gpu.stream();
7837 let mut b = __s_b.launch_builder(&f);
7838 b.arg(a)
7839 .arg(b_in)
7840 .arg(&c)
7841 .arg(w)
7842 .arg(res)
7843 .arg(&mut *out_q)
7844 .arg(&mut *out_d)
7845 .arg(&nc)
7846 .arg(&e2);
7847 unsafe {
7848 b.launch(cfg)?;
7849 }
7850 Ok(())
7851 }
7852
7853 #[allow(clippy::too_many_arguments)]
7856 pub fn rms_pre_add_scale_rms_norm_q8_1(
7857 &self,
7858 a: &CudaSlice<f32>,
7859 wa: &CudaSlice<f32>,
7860 b_in: &CudaSlice<f32>,
7861 c: f32,
7862 w: &CudaSlice<f32>,
7863 res: &mut CudaSlice<f32>,
7864 ncols: usize,
7865 nrows: usize,
7866 eps: f32,
7867 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7868 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7869 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7870 let (nc, e2) = (ncols as i32, eps);
7871 if Self::pdl_on() {
7872 {
7873 use cudarc::driver::{DevicePtr, DevicePtrMut};
7874 let s = &self.gpu.stream();
7875 let (pa, _g0) = a.device_ptr(s);
7876 let (pwa, _g1) = wa.device_ptr(s);
7877 let (pb, _g2) = b_in.device_ptr(s);
7878 let (pw, _g3) = w.device_ptr(s);
7879 let (pr, _g4) = res.device_ptr_mut(s);
7880 let (pq, _g5) = out_q.device_ptr_mut(s);
7881 let (pd, _g6) = out_d.device_ptr_mut(s);
7882 let mut ps = [
7883 &pa as *const _ as *mut std::ffi::c_void,
7884 &pwa as *const _ as *mut _,
7885 &pb as *const _ as *mut _,
7886 &c as *const _ as *mut _,
7887 &pw as *const _ as *mut _,
7888 &pr as *const _ as *mut _,
7889 &pq as *const _ as *mut _,
7890 &pd as *const _ as *mut _,
7891 &nc as *const _ as *mut _,
7892 &e2 as *const _ as *mut _,
7893 ];
7894 unsafe {
7895 self.launch_pdl(
7896 "rms_pre_add_scale_rms_norm_q8_1",
7897 (nrows as u32, 1, 1),
7898 (rms_block(), 1, 1),
7899 &mut ps,
7900 )?;
7901 }
7902 }
7903 return Ok((out_q, out_d));
7904 }
7905 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
7906 let cfg = LaunchConfig {
7907 grid_dim: (nrows as u32, 1, 1),
7908 block_dim: (rms_block(), 1, 1),
7909 shared_mem_bytes: 0,
7910 };
7911 let __s_b = self.gpu.stream();
7912 let mut b = __s_b.launch_builder(&f);
7913 b.arg(a)
7914 .arg(wa)
7915 .arg(b_in)
7916 .arg(&c)
7917 .arg(w)
7918 .arg(res)
7919 .arg(&mut out_q)
7920 .arg(&mut out_d)
7921 .arg(&nc)
7922 .arg(&e2);
7923 unsafe {
7924 b.launch(cfg)?;
7925 }
7926 Ok((out_q, out_d))
7927 }
7928
7929 pub fn gelu_tanh_mul_q8_1(
7932 &self,
7933 gate: &CudaSlice<f32>,
7934 up: &cudarc::driver::CudaView<f32>,
7935 act: &mut CudaSlice<f32>,
7936 ncols: usize,
7937 nrows: usize,
7938 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7939 debug_assert!(ncols % 128 == 0);
7940 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7941 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7942 let nc = ncols as i32;
7943 if Self::pdl_on() {
7944 {
7945 use cudarc::driver::{DevicePtr, DevicePtrMut};
7946 let s = &self.gpu.stream();
7947 let (pg, _g0) = gate.device_ptr(s);
7948 let (pu, _g1) = up.device_ptr(s);
7949 let (pact, _g2) = act.device_ptr_mut(s);
7950 let (pq, _g3) = out_q.device_ptr_mut(s);
7951 let (pd, _g4) = out_d.device_ptr_mut(s);
7952 let mut ps = [
7953 &pg as *const _ as *mut std::ffi::c_void,
7954 &pu as *const _ as *mut _,
7955 &pact as *const _ as *mut _,
7956 &pq as *const _ as *mut _,
7957 &pd as *const _ as *mut _,
7958 &nc as *const _ as *mut _,
7959 ];
7960 unsafe {
7961 self.launch_pdl(
7962 "gelu_tanh_mul_q8_1",
7963 (nrows as u32, 1, 1),
7964 (rms_block(), 1, 1),
7965 &mut ps,
7966 )?;
7967 }
7968 }
7969 return Ok((out_q, out_d));
7970 }
7971 let f = self.func("gelu_tanh_mul_q8_1");
7972 let cfg = LaunchConfig {
7973 grid_dim: (nrows as u32, 1, 1),
7974 block_dim: (rms_block(), 1, 1),
7975 shared_mem_bytes: 0,
7976 };
7977 let __s_b = self.gpu.stream();
7978 let mut b = __s_b.launch_builder(&f);
7979 b.arg(gate)
7980 .arg(up)
7981 .arg(act)
7982 .arg(&mut out_q)
7983 .arg(&mut out_d)
7984 .arg(&nc);
7985 unsafe {
7986 b.launch(cfg)?;
7987 }
7988 Ok((out_q, out_d))
7989 }
7990
7991 #[allow(clippy::too_many_arguments)]
7993 pub fn gelu_tanh_mul_q8_1_into(
7994 &self,
7995 gate: &CudaSlice<f32>,
7996 up: &cudarc::driver::CudaView<f32>,
7997 act: &mut CudaSlice<f32>,
7998 ncols: usize,
7999 nrows: usize,
8000 out_q: &mut CudaSlice<i8>,
8001 out_d: &mut CudaSlice<f32>,
8002 ) -> Result<(), Box<dyn std::error::Error>> {
8003 debug_assert!(ncols % 128 == 0);
8004 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8005 let nc = ncols as i32;
8006 if Self::pdl_on() {
8007 use cudarc::driver::{DevicePtr, DevicePtrMut};
8008 let s = &self.gpu.stream();
8009 let (pg, _g0) = gate.device_ptr(s);
8010 let (pu, _g1) = up.device_ptr(s);
8011 let (pact, _g2) = act.device_ptr_mut(s);
8012 let (pq, _g3) = out_q.device_ptr_mut(s);
8013 let (pd, _g4) = out_d.device_ptr_mut(s);
8014 let mut ps = [
8015 &pg as *const _ as *mut std::ffi::c_void,
8016 &pu as *const _ as *mut _,
8017 &pact as *const _ as *mut _,
8018 &pq as *const _ as *mut _,
8019 &pd as *const _ as *mut _,
8020 &nc as *const _ as *mut _,
8021 ];
8022 unsafe {
8023 self.launch_pdl(
8024 "gelu_tanh_mul_q8_1",
8025 (nrows as u32, 1, 1),
8026 (rms_block(), 1, 1),
8027 &mut ps,
8028 )?;
8029 }
8030 return Ok(());
8031 }
8032 let f = self.func("gelu_tanh_mul_q8_1");
8033 let cfg = LaunchConfig {
8034 grid_dim: (nrows as u32, 1, 1),
8035 block_dim: (rms_block(), 1, 1),
8036 shared_mem_bytes: 0,
8037 };
8038 let __s_b = self.gpu.stream();
8039 let mut b = __s_b.launch_builder(&f);
8040 b.arg(gate)
8041 .arg(up)
8042 .arg(&mut *act)
8043 .arg(&mut *out_q)
8044 .arg(&mut *out_d)
8045 .arg(&nc);
8046 unsafe {
8047 b.launch(cfg)?;
8048 }
8049 Ok(())
8050 }
8051
8052 #[allow(clippy::too_many_arguments)]
8054 pub fn add_rms_norm3_q8z(
8055 &self,
8056 a: &CudaSlice<f32>,
8057 b_in: &CudaSlice<f32>,
8058 w0: &CudaSlice<f32>,
8059 w1: &CudaSlice<f32>,
8060 w2: &CudaSlice<f32>,
8061 res: &mut CudaSlice<f32>,
8062 out1: &mut CudaSlice<f32>,
8063 ncols: usize,
8064 nrows: usize,
8065 eps: f32,
8066 ) -> Result<
8067 (
8068 (CudaSlice<i8>, CudaSlice<f32>),
8069 (CudaSlice<i8>, CudaSlice<f32>),
8070 ),
8071 Box<dyn std::error::Error>,
8072 > {
8073 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
8074 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8075 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
8076 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8077 let f = self.func("add_rms_norm3_q8z_f32");
8078 let cfg = LaunchConfig {
8079 grid_dim: (nrows as u32, 1, 1),
8080 block_dim: (rms_block(), 1, 1),
8081 shared_mem_bytes: 0,
8082 };
8083 let (nc, e2) = (ncols as i32, eps);
8084 let __s_b = self.gpu.stream();
8085 let mut b = __s_b.launch_builder(&f);
8086 b.arg(a)
8087 .arg(b_in)
8088 .arg(w0)
8089 .arg(w1)
8090 .arg(w2)
8091 .arg(res)
8092 .arg(&mut q0)
8093 .arg(&mut d0)
8094 .arg(out1)
8095 .arg(&mut q2)
8096 .arg(&mut d2)
8097 .arg(&nc)
8098 .arg(&e2);
8099 unsafe {
8100 b.launch(cfg)?;
8101 }
8102 Ok(((q0, d0), (q2, d2)))
8103 }
8104
8105 #[allow(clippy::too_many_arguments)]
8107 pub fn add_rms_norm3(
8108 &self,
8109 a: &CudaSlice<f32>,
8110 b_in: &CudaSlice<f32>,
8111 w0: &CudaSlice<f32>,
8112 w1: &CudaSlice<f32>,
8113 w2: &CudaSlice<f32>,
8114 res: &mut CudaSlice<f32>,
8115 d0: &mut CudaSlice<f32>,
8116 d1: &mut CudaSlice<f32>,
8117 d2: &mut CudaSlice<f32>,
8118 ncols: usize,
8119 nrows: usize,
8120 eps: f32,
8121 ) -> Result<(), Box<dyn std::error::Error>> {
8122 let f = self.func("add_rms_norm3_f32");
8123 let cfg = LaunchConfig {
8124 grid_dim: (nrows as u32, 1, 1),
8125 block_dim: (rms_block(), 1, 1),
8126 shared_mem_bytes: 0,
8127 };
8128 let (nc, e2) = (ncols as i32, eps);
8129 let __s_b = self.gpu.stream();
8130 let mut b = __s_b.launch_builder(&f);
8131 b.arg(a)
8132 .arg(b_in)
8133 .arg(w0)
8134 .arg(w1)
8135 .arg(w2)
8136 .arg(res)
8137 .arg(d0)
8138 .arg(d1)
8139 .arg(d2)
8140 .arg(&nc)
8141 .arg(&e2);
8142 unsafe {
8143 b.launch(cfg)?;
8144 }
8145 Ok(())
8146 }
8147
8148 pub fn add_scale(
8150 &self,
8151 a: &CudaSlice<f32>,
8152 b_in: &CudaSlice<f32>,
8153 c: f32,
8154 dst: &mut CudaSlice<f32>,
8155 n: usize,
8156 ) -> Result<(), Box<dyn std::error::Error>> {
8157 let f = self.func("add_scale_f32");
8158 let cfg = LaunchConfig::for_num_elems(n as u32);
8159 let ni = n as i32;
8160 let __s_b = self.gpu.stream();
8161 let mut b = __s_b.launch_builder(&f);
8162 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8163 unsafe {
8164 b.launch(cfg)?;
8165 }
8166 Ok(())
8167 }
8168
8169 pub fn layer_norm_bias(
8171 &self,
8172 x: &CudaSlice<f32>,
8173 w: &CudaSlice<f32>,
8174 b: &CudaSlice<f32>,
8175 dst: &mut CudaSlice<f32>,
8176 ncols: usize,
8177 nrows: usize,
8178 eps: f32,
8179 ) -> Result<(), Box<dyn std::error::Error>> {
8180 let f = self.func("layer_norm_bias_f32");
8181 let (nc, e) = (ncols as i32, eps);
8182 let cfg = LaunchConfig {
8183 grid_dim: (nrows as u32, 1, 1),
8184 block_dim: (256, 1, 1),
8185 shared_mem_bytes: 0,
8186 };
8187 let __s_b = self.gpu.stream();
8188 let mut lb = __s_b.launch_builder(&f);
8189 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8190 unsafe {
8191 lb.launch(cfg)?;
8192 }
8193 Ok(())
8194 }
8195
8196 pub fn gelu_tanh(
8198 &self,
8199 x: &CudaSlice<f32>,
8200 dst: &mut CudaSlice<f32>,
8201 n: usize,
8202 ) -> Result<(), Box<dyn std::error::Error>> {
8203 let f = self.func("gelu_tanh_f32");
8204 let ni = n as i64;
8205 let cfg = LaunchConfig {
8206 grid_dim: (n.div_ceil(256) as u32, 1, 1),
8207 block_dim: (256, 1, 1),
8208 shared_mem_bytes: 0,
8209 };
8210 let __s_b = self.gpu.stream();
8211 let mut lb = __s_b.launch_builder(&f);
8212 lb.arg(x).arg(&mut *dst).arg(&ni);
8213 unsafe {
8214 lb.launch(cfg)?;
8215 }
8216 Ok(())
8217 }
8218
8219 pub fn row_softmax(
8221 &self,
8222 x: &mut CudaSlice<f32>,
8223 ncols: usize,
8224 nrows: usize,
8225 ) -> Result<(), Box<dyn std::error::Error>> {
8226 let f = self.func("row_softmax_f32");
8227 let nc = ncols as i32;
8228 let cfg = LaunchConfig {
8229 grid_dim: (nrows as u32, 1, 1),
8230 block_dim: (256, 1, 1),
8231 shared_mem_bytes: 0,
8232 };
8233 let __s_b = self.gpu.stream();
8234 let mut lb = __s_b.launch_builder(&f);
8235 lb.arg(&mut *x).arg(&nc);
8236 unsafe {
8237 lb.launch(cfg)?;
8238 }
8239 Ok(())
8240 }
8241
8242 pub fn rms_norm(
8243 &self,
8244 x: &CudaSlice<f32>,
8245 w: &CudaSlice<f32>,
8246 dst: &mut CudaSlice<f32>,
8247 ncols: usize,
8248 nrows: usize,
8249 eps: f32,
8250 ) -> Result<(), Box<dyn std::error::Error>> {
8251 let (nc, e) = (ncols as i32, eps);
8252 if Self::pdl_on() && Self::pdl_wb_on() {
8253 use cudarc::driver::{DevicePtr, DevicePtrMut};
8254 let s = &self.gpu.stream();
8255 let (px, _g0) = x.device_ptr(s);
8256 let (pw, _g1) = w.device_ptr(s);
8257 let (pd, _g2) = dst.device_ptr_mut(s);
8258 let mut ps = [
8259 &px as *const _ as *mut std::ffi::c_void,
8260 &pw as *const _ as *mut _,
8261 &pd as *const _ as *mut _,
8262 &nc as *const _ as *mut _,
8263 &e as *const _ as *mut _,
8264 ];
8265 unsafe {
8266 self.launch_pdl(
8267 "rms_norm_f32",
8268 (nrows as u32, 1, 1),
8269 (rms_block(), 1, 1),
8270 &mut ps,
8271 )?;
8272 }
8273 return Ok(());
8274 }
8275 let f = self.func("rms_norm_f32");
8276 let cfg = LaunchConfig {
8277 grid_dim: (nrows as u32, 1, 1),
8278 block_dim: (rms_block(), 1, 1),
8279 shared_mem_bytes: 0,
8280 };
8281 let __s_b = self.gpu.stream();
8282 let mut b = __s_b.launch_builder(&f);
8283 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8284 unsafe {
8285 b.launch(cfg)?;
8286 }
8287 Ok(())
8288 }
8289
8290 pub fn rms_norm_decode(
8298 &self,
8299 x: &CudaSlice<f32>,
8300 w: &CudaSlice<f32>,
8301 dst: &mut CudaSlice<f32>,
8302 ncols: usize,
8303 nrows: usize,
8304 eps: f32,
8305 ) -> Result<(), Box<dyn std::error::Error>> {
8306 let f = self.func("rms_norm_f32");
8307 let cfg = LaunchConfig {
8308 grid_dim: (nrows as u32, 1, 1),
8309 block_dim: (1024, 1, 1),
8310 shared_mem_bytes: 0,
8311 };
8312 let (nc, e) = (ncols as i32, eps);
8313 let __s_b = self.gpu.stream();
8314 let mut b = __s_b.launch_builder(&f);
8315 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8316 unsafe {
8317 b.launch(cfg)?;
8318 }
8319 Ok(())
8320 }
8321
8322 pub fn rms_norm_q8_1(
8326 &self,
8327 x: &CudaSlice<f32>,
8328 w: &CudaSlice<f32>,
8329 ncols: usize,
8330 nrows: usize,
8331 eps: f32,
8332 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8333 let nblk = ncols / 32;
8334 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8335 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8336 let (nc, e) = (ncols as i32, eps);
8337 if Self::pdl_on() {
8338 {
8339 use cudarc::driver::{DevicePtr, DevicePtrMut};
8340 let s = &self.gpu.stream();
8341 let (px, _g0) = x.device_ptr(s);
8342 let (pw, _g1) = w.device_ptr(s);
8343 let (pq, _g2) = q.device_ptr_mut(s);
8344 let (pd, _g3) = d.device_ptr_mut(s);
8345 let mut ps = [
8346 &px as *const _ as *mut std::ffi::c_void,
8347 &pw as *const _ as *mut _,
8348 &pq as *const _ as *mut _,
8349 &pd as *const _ as *mut _,
8350 &nc as *const _ as *mut _,
8351 &e as *const _ as *mut _,
8352 ];
8353 unsafe {
8354 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8355 }
8356 }
8357 return Ok((q, d));
8358 }
8359 let f = self.func("rms_norm_q8_1");
8360 let cfg = LaunchConfig {
8363 grid_dim: (nrows as u32, 1, 1),
8364 block_dim: (1024, 1, 1),
8365 shared_mem_bytes: 0,
8366 };
8367 let __s_b = self.gpu.stream();
8368 let mut b = __s_b.launch_builder(&f);
8369 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8370 unsafe {
8371 b.launch(cfg)?;
8372 }
8373 Ok((q, d))
8374 }
8375
8376 pub fn rms_norm_q8_1_into(
8379 &self,
8380 x: &CudaSlice<f32>,
8381 w: &CudaSlice<f32>,
8382 ncols: usize,
8383 nrows: usize,
8384 eps: f32,
8385 q: &mut CudaSlice<i8>,
8386 d: &mut CudaSlice<f32>,
8387 ) -> Result<(), Box<dyn std::error::Error>> {
8388 let nblk = ncols / 32;
8389 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8390 let (nc, e) = (ncols as i32, eps);
8391 if Self::pdl_on() {
8392 use cudarc::driver::{DevicePtr, DevicePtrMut};
8393 let s = &self.gpu.stream();
8394 let (px, _g0) = x.device_ptr(s);
8395 let (pw, _g1) = w.device_ptr(s);
8396 let (pq, _g2) = q.device_ptr_mut(s);
8397 let (pd, _g3) = d.device_ptr_mut(s);
8398 let mut ps = [
8399 &px as *const _ as *mut std::ffi::c_void,
8400 &pw as *const _ as *mut _,
8401 &pq as *const _ as *mut _,
8402 &pd as *const _ as *mut _,
8403 &nc as *const _ as *mut _,
8404 &e as *const _ as *mut _,
8405 ];
8406 unsafe {
8407 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8408 }
8409 return Ok(());
8410 }
8411 let f = self.func("rms_norm_q8_1");
8412 let cfg = LaunchConfig {
8413 grid_dim: (nrows as u32, 1, 1),
8414 block_dim: (1024, 1, 1),
8415 shared_mem_bytes: 0,
8416 };
8417 let __s_b = self.gpu.stream();
8418 let mut b = __s_b.launch_builder(&f);
8419 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8420 unsafe {
8421 b.launch(cfg)?;
8422 }
8423 Ok(())
8424 }
8425
8426 pub fn quantize_q8_1_into(
8428 &self,
8429 x: &CudaSlice<f32>,
8430 m: usize,
8431 in_f: usize,
8432 q: &mut CudaSlice<i8>,
8433 d: &mut CudaSlice<f32>,
8434 ) -> Result<(), Box<dyn std::error::Error>> {
8435 let nblk = in_f / 32;
8436 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8437 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8438 let (inf, mi) = (in_f as i32, m as i32);
8439 if Self::pdl_on() && Self::pdl_wb_on() {
8440 use cudarc::driver::{DevicePtr, DevicePtrMut};
8441 let s = &self.gpu.stream();
8442 let (px, _g0) = x.device_ptr(s);
8443 let (pq, _g1) = q.device_ptr_mut(s);
8444 let (pd, _g2) = d.device_ptr_mut(s);
8445 let mut ps = [
8446 &px as *const _ as *mut std::ffi::c_void,
8447 &pq as *const _ as *mut _,
8448 &pd as *const _ as *mut _,
8449 &inf as *const _ as *mut _,
8450 &mi as *const _ as *mut _,
8451 ];
8452 unsafe {
8453 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8454 }
8455 return Ok(());
8456 }
8457 let f = self.func("quantize_q8_1");
8458 let __s_b = self.gpu.stream();
8459 let mut b = __s_b.launch_builder(&f);
8460 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8461 unsafe {
8462 b.launch(cfg)?;
8463 }
8464 Ok(())
8465 }
8466
8467 pub fn add_rms_norm_q8_1(
8471 &self,
8472 a: &CudaSlice<f32>,
8473 b_in: &CudaSlice<f32>,
8474 w: &CudaSlice<f32>,
8475 res: &mut CudaSlice<f32>,
8476 ncols: usize,
8477 nrows: usize,
8478 eps: f32,
8479 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8480 let nblk = ncols / 32;
8481 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8482 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8483 let f = self.func("add_rms_norm_q8_1");
8484 let cfg = LaunchConfig {
8486 grid_dim: (nrows as u32, 1, 1),
8487 block_dim: (1024, 1, 1),
8488 shared_mem_bytes: 0,
8489 };
8490 let (nc, e) = (ncols as i32, eps);
8491 let __s_bld = self.gpu.stream();
8492 let mut bld = __s_bld.launch_builder(&f);
8493 bld.arg(a)
8494 .arg(b_in)
8495 .arg(w)
8496 .arg(res)
8497 .arg(&mut q)
8498 .arg(&mut d)
8499 .arg(&nc)
8500 .arg(&e);
8501 unsafe {
8502 bld.launch(cfg)?;
8503 }
8504 Ok((q, d))
8505 }
8506
8507 pub fn add_rms_norm(
8511 &self,
8512 a: &CudaSlice<f32>,
8513 b: &CudaSlice<f32>,
8514 w: &CudaSlice<f32>,
8515 res: &mut CudaSlice<f32>,
8516 dst: &mut CudaSlice<f32>,
8517 ncols: usize,
8518 nrows: usize,
8519 eps: f32,
8520 ) -> Result<(), Box<dyn std::error::Error>> {
8521 let (nc, e) = (ncols as i32, eps);
8522 if Self::pdl_on() && Self::pdl_wb_on() {
8523 use cudarc::driver::{DevicePtr, DevicePtrMut};
8524 let s = &self.gpu.stream();
8525 let (pa, _g0) = a.device_ptr(s);
8526 let (pb, _g1) = b.device_ptr(s);
8527 let (pw, _g2) = w.device_ptr(s);
8528 let (pr, _g3) = res.device_ptr_mut(s);
8529 let (pd, _g4) = dst.device_ptr_mut(s);
8530 let mut ps = [
8531 &pa as *const _ as *mut std::ffi::c_void,
8532 &pb as *const _ as *mut _,
8533 &pw as *const _ as *mut _,
8534 &pr as *const _ as *mut _,
8535 &pd as *const _ as *mut _,
8536 &nc as *const _ as *mut _,
8537 &e as *const _ as *mut _,
8538 ];
8539 unsafe {
8540 self.launch_pdl(
8541 "add_rms_norm_f32",
8542 (nrows as u32, 1, 1),
8543 (rms_block(), 1, 1),
8544 &mut ps,
8545 )?;
8546 }
8547 return Ok(());
8548 }
8549 let f = self.func("add_rms_norm_f32");
8550 let cfg = LaunchConfig {
8551 grid_dim: (nrows as u32, 1, 1),
8552 block_dim: (rms_block(), 1, 1),
8553 shared_mem_bytes: 0,
8554 };
8555 let __s_b2 = self.gpu.stream();
8556 let mut b2 = __s_b2.launch_builder(&f);
8557 b2.arg(a)
8558 .arg(b)
8559 .arg(w)
8560 .arg(&mut *res)
8561 .arg(&mut *dst)
8562 .arg(&nc)
8563 .arg(&e);
8564 unsafe {
8565 b2.launch(cfg)?;
8566 }
8567 Ok(())
8568 }
8569
8570 #[allow(clippy::too_many_arguments)]
8573 pub fn rms_pre_add_rms_norm(
8574 &self,
8575 a: &CudaSlice<f32>,
8576 wa: &CudaSlice<f32>,
8577 b: &CudaSlice<f32>,
8578 w: &CudaSlice<f32>,
8579 res: &mut CudaSlice<f32>,
8580 dst: &mut CudaSlice<f32>,
8581 ncols: usize,
8582 nrows: usize,
8583 eps: f32,
8584 ) -> Result<(), Box<dyn std::error::Error>> {
8585 let f = self.func("rms_pre_add_rms_norm_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 (nc, e) = (ncols as i32, eps);
8592 let __s_b2 = self.gpu.stream();
8593 let mut b2 = __s_b2.launch_builder(&f);
8594 b2.arg(a)
8595 .arg(wa)
8596 .arg(b)
8597 .arg(w)
8598 .arg(&mut *res)
8599 .arg(&mut *dst)
8600 .arg(&nc)
8601 .arg(&e);
8602 unsafe {
8603 b2.launch(cfg)?;
8604 }
8605 Ok(())
8606 }
8607
8608 #[allow(clippy::too_many_arguments)]
8610 pub fn rms_pre_add_rms_norm_q8z(
8611 &self,
8612 a: &CudaSlice<f32>,
8613 wa: &CudaSlice<f32>,
8614 b: &CudaSlice<f32>,
8615 w: &CudaSlice<f32>,
8616 res: &mut CudaSlice<f32>,
8617 dst: &mut CudaSlice<f32>,
8618 ncols: usize,
8619 nrows: usize,
8620 eps: f32,
8621 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8622 debug_assert!(ncols % 128 == 0);
8623 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8624 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8625 let (nc, e) = (ncols as i32, eps);
8626 if Self::pdl_on() {
8627 {
8628 use cudarc::driver::{DevicePtr, DevicePtrMut};
8629 let s = &self.gpu.stream();
8630 let (pa, _g0) = a.device_ptr(s);
8631 let (pwa, _g1) = wa.device_ptr(s);
8632 let (pb, _g2) = b.device_ptr(s);
8633 let (pw, _g3) = w.device_ptr(s);
8634 let (pr, _g4) = res.device_ptr_mut(s);
8635 let (pdst, _g5) = dst.device_ptr_mut(s);
8636 let (pq, _g6) = out_q.device_ptr_mut(s);
8637 let (pd, _g7) = out_d.device_ptr_mut(s);
8638 let mut ps = [
8639 &pa as *const _ as *mut std::ffi::c_void,
8640 &pwa as *const _ as *mut _,
8641 &pb as *const _ as *mut _,
8642 &pw as *const _ as *mut _,
8643 &pr as *const _ as *mut _,
8644 &pdst as *const _ as *mut _,
8645 &pq as *const _ as *mut _,
8646 &pd as *const _ as *mut _,
8647 &nc as *const _ as *mut _,
8648 &e as *const _ as *mut _,
8649 ];
8650 unsafe {
8651 self.launch_pdl(
8652 "rms_pre_add_rms_norm_q8z_f32",
8653 (nrows as u32, 1, 1),
8654 (rms_block(), 1, 1),
8655 &mut ps,
8656 )?;
8657 }
8658 }
8659 return Ok((out_q, out_d));
8660 }
8661 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8662 let cfg = LaunchConfig {
8663 grid_dim: (nrows as u32, 1, 1),
8664 block_dim: (rms_block(), 1, 1),
8665 shared_mem_bytes: 0,
8666 };
8667 let __s_b2 = self.gpu.stream();
8668 let mut b2 = __s_b2.launch_builder(&f);
8669 b2.arg(a)
8670 .arg(wa)
8671 .arg(b)
8672 .arg(w)
8673 .arg(&mut *res)
8674 .arg(&mut *dst)
8675 .arg(&mut out_q)
8676 .arg(&mut out_d)
8677 .arg(&nc)
8678 .arg(&e);
8679 unsafe {
8680 b2.launch(cfg)?;
8681 }
8682 Ok((out_q, out_d))
8683 }
8684
8685 #[allow(clippy::too_many_arguments)]
8689 pub fn rms_pre_add_rms_norm_q8z_into(
8690 &self,
8691 a: &CudaSlice<f32>,
8692 wa: &CudaSlice<f32>,
8693 b: &CudaSlice<f32>,
8694 w: &CudaSlice<f32>,
8695 res: &mut CudaSlice<f32>,
8696 dst: &mut CudaSlice<f32>,
8697 ncols: usize,
8698 nrows: usize,
8699 eps: f32,
8700 out_q: &mut CudaSlice<i8>,
8701 out_d: &mut CudaSlice<f32>,
8702 ) -> Result<(), Box<dyn std::error::Error>> {
8703 debug_assert!(ncols % 128 == 0);
8704 let (nc, e) = (ncols as i32, eps);
8705 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8706 let cfg = LaunchConfig {
8707 grid_dim: (nrows as u32, 1, 1),
8708 block_dim: (rms_block(), 1, 1),
8709 shared_mem_bytes: 0,
8710 };
8711 let __s_b = self.gpu.stream();
8712 let mut b2 = __s_b.launch_builder(&f);
8713 b2.arg(a)
8714 .arg(wa)
8715 .arg(b)
8716 .arg(w)
8717 .arg(&mut *res)
8718 .arg(&mut *dst)
8719 .arg(&mut *out_q)
8720 .arg(&mut *out_d)
8721 .arg(&nc)
8722 .arg(&e);
8723 unsafe {
8724 b2.launch(cfg)?;
8725 }
8726 Ok(())
8727 }
8728
8729 #[allow(clippy::too_many_arguments)]
8732 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8733 &self,
8734 a: &CudaSlice<f32>,
8735 wa: &CudaSlice<f32>,
8736 b_in: &CudaSlice<f32>,
8737 c: f32,
8738 w: &CudaSlice<f32>,
8739 res: &mut CudaSlice<f32>,
8740 ncols: usize,
8741 nrows: usize,
8742 eps: f32,
8743 out_q: &mut CudaSlice<i8>,
8744 out_d: &mut CudaSlice<f32>,
8745 ) -> Result<(), Box<dyn std::error::Error>> {
8746 debug_assert!(ncols % 128 == 0);
8747 let (nc, e2) = (ncols as i32, eps);
8748 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8749 let cfg = LaunchConfig {
8750 grid_dim: (nrows as u32, 1, 1),
8751 block_dim: (rms_block(), 1, 1),
8752 shared_mem_bytes: 0,
8753 };
8754 let __s_b = self.gpu.stream();
8755 let mut b2 = __s_b.launch_builder(&f);
8756 b2.arg(a)
8757 .arg(wa)
8758 .arg(b_in)
8759 .arg(&c)
8760 .arg(w)
8761 .arg(&mut *res)
8762 .arg(&mut *out_q)
8763 .arg(&mut *out_d)
8764 .arg(&nc)
8765 .arg(&e2);
8766 unsafe {
8767 b2.launch(cfg)?;
8768 }
8769 Ok(())
8770 }
8771
8772 pub fn g4_pnfold_on() -> bool {
8780 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8781 *ON.get_or_init(|| {
8782 std::env::var("MEMRA_G4_PNFOLD")
8783 .map(|v| v != "0")
8784 .unwrap_or(true)
8785 })
8786 }
8787
8788 pub fn build_q4_out_concat3(
8792 &self,
8793 w0: &crate::model::GpuTensor,
8794 w1: &crate::model::GpuTensor,
8795 w2: &crate::model::GpuTensor,
8796 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
8797 use crate::model::GpuTensor;
8798 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
8799 match w {
8800 GpuTensor::Quant {
8801 qtype,
8802 row_bytes,
8803 rp,
8804 ..
8805 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
8806 _ => None,
8807 }
8808 };
8809 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
8810 else {
8811 return Ok(None);
8812 };
8813 if rb0 != rb1
8814 || rb0 != rb2
8815 || w0.in_features() != w1.in_features()
8816 || w0.in_features() != w2.in_features()
8817 {
8818 return Ok(None);
8819 }
8820 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
8821 match w {
8822 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
8823 _ => unreachable!(),
8824 }
8825 }
8826 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
8827 let total = rb0 * (o0 + o1 + o2);
8828 let mut cat = self.alloc_u8(total)?;
8829 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
8830 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
8831 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
8832 Ok(Some(GpuTensor::Quant {
8833 bytes: cat,
8834 qtype: QT_Q4_0,
8835 row_bytes: rb0,
8836 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
8837 scale: 1.0,
8838 rp: false,
8839 #[cfg(memra_cutlass)]
8840 cutlass: None,
8841 fp8: None,
8842 blk: None,
8843 rp4: None,
8844 f16: None,
8845 }))
8846 }
8847
8848 fn full_width_rope_only(
8866 kernel: &str,
8867 n_rot: usize,
8868 head_dim: usize,
8869 ) -> Result<(), Box<dyn std::error::Error>> {
8870 if n_rot == head_dim {
8871 return Ok(());
8872 }
8873 Err(format!(
8874 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
8875 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
8876 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
8877 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
8878 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
8879 )
8880 .into())
8881 }
8882
8883 #[allow(clippy::too_many_arguments)]
8887 pub fn rms_norm_qkv_rope_cat(
8888 &self,
8889 qkv: &CudaSlice<f32>,
8890 wq: &CudaSlice<f32>,
8891 wk: &CudaSlice<f32>,
8892 wv: &CudaSlice<f32>,
8893 q: &mut CudaSlice<f32>,
8894 k: &mut CudaSlice<f32>,
8895 v: &mut CudaSlice<f32>,
8896 head_dim: usize,
8897 n_rot: usize,
8898 rq: usize,
8899 rk: usize,
8900 pos: &CudaSlice<i32>,
8901 nh_q: usize,
8902 nh_k: usize,
8903 base: f32,
8904 freq_scale: f32,
8905 ff: Option<&CudaSlice<f32>>,
8906 eps: f32,
8907 ) -> Result<(), Box<dyn std::error::Error>> {
8908 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
8909 let rows = rq + rk + rk;
8910 let theta_scale = base.powf(-2.0 / head_dim as f32);
8911 let (nc, rqi, rki, nhq, nhk) = (
8912 head_dim as i32,
8913 rq as i32,
8914 rk as i32,
8915 nh_q as i32,
8916 nh_k as i32,
8917 );
8918 if Self::pdl_on() {
8919 use cudarc::driver::{DevicePtr, DevicePtrMut};
8920 let s = &self.gpu.stream();
8921 let (pqkv, _g0) = qkv.device_ptr(s);
8922 let (pwq, _g1) = wq.device_ptr(s);
8923 let (pwk, _g2) = wk.device_ptr(s);
8924 let (pwv, _g3) = wv.device_ptr(s);
8925 let (pq, _g4) = q.device_ptr_mut(s);
8926 let (pk, _g5) = k.device_ptr_mut(s);
8927 let (pv, _g6) = v.device_ptr_mut(s);
8928 let (ppos, _g7) = pos.device_ptr(s);
8929 let (pff, _g8) = match ff {
8930 Some(t) => {
8931 let (p, g) = t.device_ptr(s);
8932 (p, Some(g))
8933 }
8934 None => (0, None),
8935 };
8936 let mut ps = [
8937 &pqkv as *const _ as *mut std::ffi::c_void,
8938 &pwq as *const _ as *mut _,
8939 &pwk as *const _ as *mut _,
8940 &pwv as *const _ as *mut _,
8941 &pq as *const _ as *mut _,
8942 &pk as *const _ as *mut _,
8943 &pv as *const _ as *mut _,
8944 &nc as *const _ as *mut _,
8945 &rqi as *const _ as *mut _,
8946 &rki as *const _ as *mut _,
8947 &ppos as *const _ as *mut _,
8948 &nhq as *const _ as *mut _,
8949 &nhk as *const _ as *mut _,
8950 &theta_scale as *const _ as *mut _,
8951 &freq_scale as *const _ as *mut _,
8952 &pff as *const _ as *mut _,
8953 &eps as *const _ as *mut _,
8954 ];
8955 unsafe {
8956 self.launch_pdl(
8957 "rms_norm_qkv_rope_cat_f32",
8958 (rows as u32, 1, 1),
8959 (rms_block(), 1, 1),
8960 &mut ps,
8961 )?;
8962 }
8963 return Ok(());
8964 }
8965 let f = self.func("rms_norm_qkv_rope_cat_f32");
8966 let cfg = LaunchConfig {
8967 grid_dim: (rows as u32, 1, 1),
8968 block_dim: (rms_block(), 1, 1),
8969 shared_mem_bytes: 0,
8970 };
8971 let __s_b = self.gpu.stream();
8972 let mut b = __s_b.launch_builder(&f);
8973 match ff {
8974 Some(t) => {
8975 b.arg(qkv)
8976 .arg(wq)
8977 .arg(wk)
8978 .arg(wv)
8979 .arg(&mut *q)
8980 .arg(&mut *k)
8981 .arg(&mut *v)
8982 .arg(&nc)
8983 .arg(&rqi)
8984 .arg(&rki)
8985 .arg(pos)
8986 .arg(&nhq)
8987 .arg(&nhk)
8988 .arg(&theta_scale)
8989 .arg(&freq_scale)
8990 .arg(t)
8991 .arg(&eps);
8992 unsafe {
8993 b.launch(cfg)?;
8994 }
8995 }
8996 None => {
8997 let null: u64 = 0;
8998 b.arg(qkv)
8999 .arg(wq)
9000 .arg(wk)
9001 .arg(wv)
9002 .arg(&mut *q)
9003 .arg(&mut *k)
9004 .arg(&mut *v)
9005 .arg(&nc)
9006 .arg(&rqi)
9007 .arg(&rki)
9008 .arg(pos)
9009 .arg(&nhq)
9010 .arg(&nhk)
9011 .arg(&theta_scale)
9012 .arg(&freq_scale)
9013 .arg(&null)
9014 .arg(&eps);
9015 unsafe {
9016 b.launch(cfg)?;
9017 }
9018 }
9019 }
9020 Ok(())
9021 }
9022
9023 #[allow(clippy::too_many_arguments)]
9027 pub fn rms_norm_qkv_rope(
9028 &self,
9029 q0: &CudaSlice<f32>,
9030 k0: &CudaSlice<f32>,
9031 v0: &CudaSlice<f32>,
9032 wq: &CudaSlice<f32>,
9033 wk: &CudaSlice<f32>,
9034 wv: &CudaSlice<f32>,
9035 q: &mut CudaSlice<f32>,
9036 k: &mut CudaSlice<f32>,
9037 v: &mut CudaSlice<f32>,
9038 head_dim: usize,
9039 n_rot: usize,
9040 rq: usize,
9041 rk: usize,
9042 pos: &CudaSlice<i32>,
9043 nh_q: usize,
9044 nh_k: usize,
9045 base: f32,
9046 freq_scale: f32,
9047 ff: Option<&CudaSlice<f32>>,
9048 eps: f32,
9049 ) -> Result<(), Box<dyn std::error::Error>> {
9050 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
9051 let f = self.func("rms_norm_qkv_rope_f32");
9052 let rows = rq + rk + rk; let cfg = LaunchConfig {
9054 grid_dim: (rows as u32, 1, 1),
9055 block_dim: (rms_block(), 1, 1),
9056 shared_mem_bytes: 0,
9057 };
9058 let theta_scale = base.powf(-2.0 / head_dim as f32);
9059 let (nc, rqi, rki, nhq, nhk) = (
9060 head_dim as i32,
9061 rq as i32,
9062 rk as i32,
9063 nh_q as i32,
9064 nh_k as i32,
9065 );
9066 let __s_b = self.gpu.stream();
9067 let mut b = __s_b.launch_builder(&f);
9068 match ff {
9069 Some(t) => {
9070 b.arg(q0)
9071 .arg(k0)
9072 .arg(v0)
9073 .arg(wq)
9074 .arg(wk)
9075 .arg(wv)
9076 .arg(&mut *q)
9077 .arg(&mut *k)
9078 .arg(&mut *v)
9079 .arg(&nc)
9080 .arg(&rqi)
9081 .arg(&rki)
9082 .arg(pos)
9083 .arg(&nhq)
9084 .arg(&nhk)
9085 .arg(&theta_scale)
9086 .arg(&freq_scale)
9087 .arg(t)
9088 .arg(&eps);
9089 unsafe {
9090 b.launch(cfg)?;
9091 }
9092 }
9093 None => {
9094 let null: u64 = 0;
9095 b.arg(q0)
9096 .arg(k0)
9097 .arg(v0)
9098 .arg(wq)
9099 .arg(wk)
9100 .arg(wv)
9101 .arg(&mut *q)
9102 .arg(&mut *k)
9103 .arg(&mut *v)
9104 .arg(&nc)
9105 .arg(&rqi)
9106 .arg(&rki)
9107 .arg(pos)
9108 .arg(&nhq)
9109 .arg(&nhk)
9110 .arg(&theta_scale)
9111 .arg(&freq_scale)
9112 .arg(&null)
9113 .arg(&eps);
9114 unsafe {
9115 b.launch(cfg)?;
9116 }
9117 }
9118 }
9119 Ok(())
9120 }
9121
9122 #[allow(clippy::too_many_arguments)]
9128 pub fn rms_norm_qkv_rope_append_dc(
9129 &self,
9130 q0: &CudaSlice<f32>,
9131 k0: &CudaSlice<f32>,
9132 v0: &CudaSlice<f32>,
9133 wq: &CudaSlice<f32>,
9134 wk: &CudaSlice<f32>,
9135 wv: &CudaSlice<f32>,
9136 q: &mut CudaSlice<f32>,
9137 k: &mut CudaSlice<f32>,
9138 v: &mut CudaSlice<f32>,
9139 head_dim: usize,
9140 n_rot: usize,
9141 rq: usize,
9142 rk: usize,
9143 pos: &CudaSlice<i32>,
9144 nh_q: usize,
9145 nh_k: usize,
9146 base: f32,
9147 freq_scale: f32,
9148 ff: Option<&CudaSlice<f32>>,
9149 eps: f32,
9150 kc: &mut CudaSlice<u8>,
9151 vc: &mut CudaSlice<u8>,
9152 t_dev: &CudaSlice<i32>,
9153 k_tok_bytes: usize,
9154 v_tok_bytes: usize,
9155 g: bool,
9156 ) -> Result<(), Box<dyn std::error::Error>> {
9157 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
9158 let rows = rq + rk + rk;
9159 let theta_scale = base.powf(-2.0 / head_dim as f32);
9160 let (nc, rqi, rki, nhq, nhk) = (
9161 head_dim as i32,
9162 rq as i32,
9163 rk as i32,
9164 nh_q as i32,
9165 nh_k as i32,
9166 );
9167 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9168 if Self::pdl_on() && Self::pdl_wb_on() {
9169 use cudarc::driver::{DevicePtr, DevicePtrMut};
9170 let s = &self.gpu.stream();
9171 let (p0, _a0) = q0.device_ptr(s);
9172 let (p1, _a1) = k0.device_ptr(s);
9173 let (p2, _a2) = v0.device_ptr(s);
9174 let (pwq, _a3) = wq.device_ptr(s);
9175 let (pwk, _a4) = wk.device_ptr(s);
9176 let (pwv, _a5) = wv.device_ptr(s);
9177 let (pq, _a6) = q.device_ptr_mut(s);
9178 let (pk, _a7) = k.device_ptr_mut(s);
9179 let (pv, _a8) = v.device_ptr_mut(s);
9180 let (pp, _a9) = pos.device_ptr(s);
9181 let pff: u64 = match ff {
9182 Some(t) => {
9183 let (p, _gg) = t.device_ptr(s);
9184 p as u64
9185 }
9186 None => 0,
9187 };
9188 let (pkc, _a10) = kc.device_ptr_mut(s);
9189 let (pvc, _a11) = vc.device_ptr_mut(s);
9190 let (pt, _a12) = t_dev.device_ptr(s);
9191 let mut ps = [
9192 &p0 as *const _ as *mut std::ffi::c_void,
9193 &p1 as *const _ as *mut _,
9194 &p2 as *const _ as *mut _,
9195 &pwq as *const _ as *mut _,
9196 &pwk as *const _ as *mut _,
9197 &pwv as *const _ as *mut _,
9198 &pq as *const _ as *mut _,
9199 &pk as *const _ as *mut _,
9200 &pv as *const _ as *mut _,
9201 &nc as *const _ as *mut _,
9202 &rqi as *const _ as *mut _,
9203 &rki as *const _ as *mut _,
9204 &pp as *const _ as *mut _,
9205 &nhq as *const _ as *mut _,
9206 &nhk as *const _ as *mut _,
9207 &theta_scale as *const _ as *mut _,
9208 &freq_scale as *const _ as *mut _,
9209 &pff as *const _ as *mut _,
9210 &eps as *const _ as *mut _,
9211 &pkc as *const _ as *mut _,
9212 &pvc as *const _ as *mut _,
9213 &pt as *const _ as *mut _,
9214 &ktb as *const _ as *mut _,
9215 &vtb as *const _ as *mut _,
9216 ];
9217 unsafe {
9218 self.launch_pdl_flash(
9219 g,
9220 "rms_norm_qkv_rope_append_dc_f32",
9221 (rows as u32, 1, 1),
9222 (rms_block(), 1, 1),
9223 0,
9224 &mut ps,
9225 )?;
9226 }
9227 return Ok(());
9228 }
9229 let f = if g {
9230 self.func_g("rms_norm_qkv_rope_append_dc_f32")
9231 } else {
9232 self.func("rms_norm_qkv_rope_append_dc_f32")
9233 };
9234 let cfg = LaunchConfig {
9235 grid_dim: (rows as u32, 1, 1),
9236 block_dim: (rms_block(), 1, 1),
9237 shared_mem_bytes: 0,
9238 };
9239 let __s_b = self.gpu.stream();
9240 let mut b = __s_b.launch_builder(&f);
9241 match ff {
9242 Some(t) => {
9243 b.arg(q0)
9244 .arg(k0)
9245 .arg(v0)
9246 .arg(wq)
9247 .arg(wk)
9248 .arg(wv)
9249 .arg(&mut *q)
9250 .arg(&mut *k)
9251 .arg(&mut *v)
9252 .arg(&nc)
9253 .arg(&rqi)
9254 .arg(&rki)
9255 .arg(pos)
9256 .arg(&nhq)
9257 .arg(&nhk)
9258 .arg(&theta_scale)
9259 .arg(&freq_scale)
9260 .arg(t)
9261 .arg(&eps)
9262 .arg(&mut *kc)
9263 .arg(&mut *vc)
9264 .arg(t_dev)
9265 .arg(&ktb)
9266 .arg(&vtb);
9267 unsafe {
9268 b.launch(cfg)?;
9269 }
9270 }
9271 None => {
9272 let null: u64 = 0;
9273 b.arg(q0)
9274 .arg(k0)
9275 .arg(v0)
9276 .arg(wq)
9277 .arg(wk)
9278 .arg(wv)
9279 .arg(&mut *q)
9280 .arg(&mut *k)
9281 .arg(&mut *v)
9282 .arg(&nc)
9283 .arg(&rqi)
9284 .arg(&rki)
9285 .arg(pos)
9286 .arg(&nhq)
9287 .arg(&nhk)
9288 .arg(&theta_scale)
9289 .arg(&freq_scale)
9290 .arg(&null)
9291 .arg(&eps)
9292 .arg(&mut *kc)
9293 .arg(&mut *vc)
9294 .arg(t_dev)
9295 .arg(&ktb)
9296 .arg(&vtb);
9297 unsafe {
9298 b.launch(cfg)?;
9299 }
9300 }
9301 }
9302 Ok(())
9303 }
9304
9305 #[allow(clippy::too_many_arguments)]
9313 pub fn rms_norm_qkv_rope_append(
9314 &self,
9315 q0: &CudaSlice<f32>,
9316 k0: &CudaSlice<f32>,
9317 v0: &CudaSlice<f32>,
9318 wq: &CudaSlice<f32>,
9319 wk: &CudaSlice<f32>,
9320 wv: &CudaSlice<f32>,
9321 q: &mut CudaSlice<f32>,
9322 k: &mut CudaSlice<f32>,
9323 v: &mut CudaSlice<f32>,
9324 head_dim: usize,
9325 n_rot: usize,
9326 rq: usize,
9327 rk: usize,
9328 pos: &CudaSlice<i32>,
9329 nh_q: usize,
9330 nh_k: usize,
9331 base: f32,
9332 freq_scale: f32,
9333 ff: Option<&CudaSlice<f32>>,
9334 eps: f32,
9335 kc: &mut CudaSlice<u8>,
9336 vc: &mut CudaSlice<u8>,
9337 t: usize,
9338 k_tok_bytes: usize,
9339 v_tok_bytes: usize,
9340 g: bool,
9341 ) -> Result<(), Box<dyn std::error::Error>> {
9342 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
9343 let rows = rq + rk + rk;
9344 let theta_scale = base.powf(-2.0 / head_dim as f32);
9345 let (nc, rqi, rki, nhq, nhk) = (
9346 head_dim as i32,
9347 rq as i32,
9348 rk as i32,
9349 nh_q as i32,
9350 nh_k as i32,
9351 );
9352 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9353 let ti = t as i32;
9354 if Self::pdl_on() && Self::pdl_wb_on() {
9355 use cudarc::driver::{DevicePtr, DevicePtrMut};
9356 let s = &self.gpu.stream();
9357 let (p0, _a0) = q0.device_ptr(s);
9358 let (p1, _a1) = k0.device_ptr(s);
9359 let (p2, _a2) = v0.device_ptr(s);
9360 let (pwq, _a3) = wq.device_ptr(s);
9361 let (pwk, _a4) = wk.device_ptr(s);
9362 let (pwv, _a5) = wv.device_ptr(s);
9363 let (pq, _a6) = q.device_ptr_mut(s);
9364 let (pk, _a7) = k.device_ptr_mut(s);
9365 let (pv, _a8) = v.device_ptr_mut(s);
9366 let (pp, _a9) = pos.device_ptr(s);
9367 let pff: u64 = match ff {
9368 Some(t) => {
9369 let (p, _gg) = t.device_ptr(s);
9370 p as u64
9371 }
9372 None => 0,
9373 };
9374 let (pkc, _a10) = kc.device_ptr_mut(s);
9375 let (pvc, _a11) = vc.device_ptr_mut(s);
9376 let mut ps = [
9377 &p0 as *const _ as *mut std::ffi::c_void,
9378 &p1 as *const _ as *mut _,
9379 &p2 as *const _ as *mut _,
9380 &pwq as *const _ as *mut _,
9381 &pwk as *const _ as *mut _,
9382 &pwv as *const _ as *mut _,
9383 &pq as *const _ as *mut _,
9384 &pk as *const _ as *mut _,
9385 &pv as *const _ as *mut _,
9386 &nc as *const _ as *mut _,
9387 &rqi as *const _ as *mut _,
9388 &rki as *const _ as *mut _,
9389 &pp as *const _ as *mut _,
9390 &nhq as *const _ as *mut _,
9391 &nhk as *const _ as *mut _,
9392 &theta_scale as *const _ as *mut _,
9393 &freq_scale as *const _ as *mut _,
9394 &pff as *const _ as *mut _,
9395 &eps as *const _ as *mut _,
9396 &pkc as *const _ as *mut _,
9397 &pvc as *const _ as *mut _,
9398 &ti as *const _ as *mut _,
9399 &ktb as *const _ as *mut _,
9400 &vtb as *const _ as *mut _,
9401 ];
9402 unsafe {
9403 self.launch_pdl_flash(
9404 g,
9405 "rms_norm_qkv_rope_append_f32",
9406 (rows as u32, 1, 1),
9407 (rms_block(), 1, 1),
9408 0,
9409 &mut ps,
9410 )?;
9411 }
9412 return Ok(());
9413 }
9414 let f = if g {
9415 self.func_g("rms_norm_qkv_rope_append_f32")
9416 } else {
9417 self.func("rms_norm_qkv_rope_append_f32")
9418 };
9419 let cfg = LaunchConfig {
9420 grid_dim: (rows as u32, 1, 1),
9421 block_dim: (rms_block(), 1, 1),
9422 shared_mem_bytes: 0,
9423 };
9424 let __s_b = self.gpu.stream();
9425 let mut b = __s_b.launch_builder(&f);
9426 let null: u64 = 0;
9427 b.arg(q0)
9428 .arg(k0)
9429 .arg(v0)
9430 .arg(wq)
9431 .arg(wk)
9432 .arg(wv)
9433 .arg(&mut *q)
9434 .arg(&mut *k)
9435 .arg(&mut *v)
9436 .arg(&nc)
9437 .arg(&rqi)
9438 .arg(&rki)
9439 .arg(pos)
9440 .arg(&nhq)
9441 .arg(&nhk)
9442 .arg(&theta_scale)
9443 .arg(&freq_scale);
9444 match ff {
9445 Some(t) => {
9446 b.arg(t);
9447 }
9448 None => {
9449 b.arg(&null);
9450 }
9451 }
9452 b.arg(&eps)
9453 .arg(&mut *kc)
9454 .arg(&mut *vc)
9455 .arg(&ti)
9456 .arg(&ktb)
9457 .arg(&vtb);
9458 unsafe {
9459 b.launch(cfg)?;
9460 }
9461 Ok(())
9462 }
9463
9464 pub fn add_q8_1(
9465 &self,
9466 a: &CudaSlice<f32>,
9467 b: &CudaSlice<f32>,
9468 res: &mut CudaSlice<f32>,
9469 ncols: usize,
9470 nrows: usize,
9471 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9472 debug_assert!(ncols % 128 == 0);
9473 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9474 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9475 let f = self.func("add_q8_1_f32");
9476 let cfg = LaunchConfig {
9477 grid_dim: (nrows as u32, 1, 1),
9478 block_dim: (rms_block(), 1, 1),
9479 shared_mem_bytes: 0,
9480 };
9481 let nc = ncols as i32;
9482 let __s_b2 = self.gpu.stream();
9483 let mut b2 = __s_b2.launch_builder(&f);
9484 b2.arg(a)
9485 .arg(b)
9486 .arg(&mut *res)
9487 .arg(&mut out_q)
9488 .arg(&mut out_d)
9489 .arg(&nc);
9490 unsafe {
9491 b2.launch(cfg)?;
9492 }
9493 Ok((out_q, out_d))
9494 }
9495
9496 pub fn rms_pre_add_q8_1(
9500 &self,
9501 a: &CudaSlice<f32>,
9502 wa: &CudaSlice<f32>,
9503 b: &CudaSlice<f32>,
9504 res: &mut CudaSlice<f32>,
9505 ncols: usize,
9506 nrows: usize,
9507 eps: f32,
9508 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9509 debug_assert!(ncols % 128 == 0);
9510 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9511 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9512 let f = self.func("rms_pre_add_q8_1_f32");
9513 let cfg = LaunchConfig {
9514 grid_dim: (nrows as u32, 1, 1),
9515 block_dim: (rms_block(), 1, 1),
9516 shared_mem_bytes: 0,
9517 };
9518 let (nc, ep) = (ncols as i32, eps);
9519 let __s_b2 = self.gpu.stream();
9520 let mut b2 = __s_b2.launch_builder(&f);
9521 b2.arg(a)
9522 .arg(wa)
9523 .arg(b)
9524 .arg(&mut *res)
9525 .arg(&mut out_q)
9526 .arg(&mut out_d)
9527 .arg(&nc)
9528 .arg(&ep);
9529 unsafe {
9530 b2.launch(cfg)?;
9531 }
9532 Ok((out_q, out_d))
9533 }
9534
9535 pub fn l2_v2_on(ncols: usize) -> bool {
9539 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9540 }
9541
9542 pub fn l2_norm_pp(
9543 &self,
9544 x: &CudaSlice<f32>,
9545 dst: &mut CudaSlice<f32>,
9546 dst16: Option<&mut CudaSlice<u8>>,
9547 ncols: usize,
9548 nrows: usize,
9549 eps: f32,
9550 ) -> Result<(), Box<dyn std::error::Error>> {
9551 if Self::l2_v2_on(ncols) {
9552 let f = self.func("l2_norm_pp_v2_f32");
9553 let rows_per_block = 8u32; let cfg = LaunchConfig {
9555 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9556 block_dim: (256, 1, 1),
9557 shared_mem_bytes: 0,
9558 };
9559 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9560 let d16: u64 = match dst16 {
9562 Some(d) => self.addr_u8(d),
9563 None => 0,
9564 };
9565 let __s_b = self.gpu.stream();
9566 let mut b = __s_b.launch_builder(&f);
9567 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9568 unsafe {
9569 b.launch(cfg)?;
9570 }
9571 return Ok(());
9572 }
9573 self.l2_norm(x, dst, ncols, nrows, eps)
9574 }
9575
9576 pub fn l2_norm(
9577 &self,
9578 x: &CudaSlice<f32>,
9579 dst: &mut CudaSlice<f32>,
9580 ncols: usize,
9581 nrows: usize,
9582 eps: f32,
9583 ) -> Result<(), Box<dyn std::error::Error>> {
9584 let f = self.func("l2_norm_f32");
9585 let cfg = LaunchConfig {
9586 grid_dim: (nrows as u32, 1, 1),
9587 block_dim: (256, 1, 1),
9588 shared_mem_bytes: 0,
9589 };
9590 let (nc, e) = (ncols as i32, eps);
9591 let __s_b = self.gpu.stream();
9592 let mut b = __s_b.launch_builder(&f);
9593 b.arg(x).arg(dst).arg(&nc).arg(&e);
9594 unsafe {
9595 b.launch(cfg)?;
9596 }
9597 Ok(())
9598 }
9599
9600 pub fn l2_norm_decode(
9606 &self,
9607 x: &CudaSlice<f32>,
9608 dst: &mut CudaSlice<f32>,
9609 ncols: usize,
9610 nrows: usize,
9611 eps: f32,
9612 ) -> Result<(), Box<dyn std::error::Error>> {
9613 let f = self.func("l2_norm_f32");
9614 let cfg = LaunchConfig {
9615 grid_dim: (nrows as u32, 1, 1),
9616 block_dim: (32, 1, 1),
9617 shared_mem_bytes: 0,
9618 };
9619 let (nc, e) = (ncols as i32, eps);
9620 let __s_b = self.gpu.stream();
9621 let mut b = __s_b.launch_builder(&f);
9622 b.arg(x).arg(dst).arg(&nc).arg(&e);
9623 unsafe {
9624 b.launch(cfg)?;
9625 }
9626 Ok(())
9627 }
9628
9629 pub fn rope_neox(
9631 &self,
9632 x: &mut CudaSlice<f32>,
9633 pos: &CudaSlice<i32>,
9634 head_dim: usize,
9635 n_dims: usize,
9636 n_heads: usize,
9637 n_tokens: usize,
9638 freq_base: f32,
9639 freq_scale: f32,
9640 ) -> Result<(), Box<dyn std::error::Error>> {
9641 let f = self.func("rope_neox_f32");
9642 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9643 let grid = (n_heads * n_tokens) as u32;
9644 let cfg = LaunchConfig {
9645 grid_dim: (grid, 1, 1),
9646 block_dim: ((head_dim / 2) as u32, 1, 1),
9647 shared_mem_bytes: 0,
9648 };
9649 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9650 let __s_b = self.gpu.stream();
9651 let mut b = __s_b.launch_builder(&f);
9652 b.arg(x)
9653 .arg(pos)
9654 .arg(&hd)
9655 .arg(&nd)
9656 .arg(&nh)
9657 .arg(&theta_scale)
9658 .arg(&freq_scale);
9659 unsafe {
9660 b.launch(cfg)?;
9661 }
9662 Ok(())
9663 }
9664
9665 pub fn rope_neox_ff(
9667 &self,
9668 x: &mut CudaSlice<f32>,
9669 pos: &CudaSlice<i32>,
9670 head_dim: usize,
9671 n_dims: usize,
9672 n_heads: usize,
9673 n_tokens: usize,
9674 freq_base: f32,
9675 freq_scale: f32,
9676 ff: &CudaSlice<f32>,
9677 ) -> Result<(), Box<dyn std::error::Error>> {
9678 let f = self.func("rope_neox_ff_f32");
9679 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9680 let grid = (n_heads * n_tokens) as u32;
9681 let cfg = LaunchConfig {
9682 grid_dim: (grid, 1, 1),
9683 block_dim: ((head_dim / 2) as u32, 1, 1),
9684 shared_mem_bytes: 0,
9685 };
9686 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9687 let __s_b = self.gpu.stream();
9688 let mut b = __s_b.launch_builder(&f);
9689 b.arg(x)
9690 .arg(pos)
9691 .arg(&hd)
9692 .arg(&nd)
9693 .arg(&nh)
9694 .arg(&theta_scale)
9695 .arg(&freq_scale)
9696 .arg(ff);
9697 unsafe {
9698 b.launch(cfg)?;
9699 }
9700 Ok(())
9701 }
9702
9703 #[allow(clippy::too_many_arguments)]
9705 pub fn rope_neox2(
9706 &self,
9707 q: &mut CudaSlice<f32>,
9708 k: &mut CudaSlice<f32>,
9709 pos: &CudaSlice<i32>,
9710 head_dim: usize,
9711 n_dims: usize,
9712 nh_q: usize,
9713 nh_k: usize,
9714 n_tokens: usize,
9715 freq_base: f32,
9716 freq_scale: f32,
9717 ff: Option<&CudaSlice<f32>>,
9718 ) -> Result<(), Box<dyn std::error::Error>> {
9719 let f = self.func("rope_neox2_f32");
9720 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9721 let grid = ((nh_q + nh_k) * n_tokens) as u32;
9722 let cfg = LaunchConfig {
9723 grid_dim: (grid, 1, 1),
9724 block_dim: ((head_dim / 2) as u32, 1, 1),
9725 shared_mem_bytes: 0,
9726 };
9727 let (hd, nd, nq, nk, nt) = (
9728 head_dim as i32,
9729 n_dims as i32,
9730 nh_q as i32,
9731 nh_k as i32,
9732 n_tokens as i32,
9733 );
9734 let __s_b = self.gpu.stream();
9735 let mut b = __s_b.launch_builder(&f);
9736 b.arg(q)
9737 .arg(k)
9738 .arg(pos)
9739 .arg(&hd)
9740 .arg(&nd)
9741 .arg(&nq)
9742 .arg(&nk)
9743 .arg(&nt)
9744 .arg(&theta_scale)
9745 .arg(&freq_scale);
9746 match ff {
9747 Some(ffv) => {
9748 b.arg(ffv);
9749 unsafe {
9750 b.launch(cfg)?;
9751 }
9752 }
9753 None => {
9754 let null: u64 = 0;
9755 b.arg(&null);
9756 unsafe {
9757 b.launch(cfg)?;
9758 }
9759 }
9760 }
9761 Ok(())
9762 }
9763
9764 pub fn gelu_tanh_mul(
9766 &self,
9767 gate: &CudaSlice<f32>,
9768 up: &CudaSlice<f32>,
9769 dst: &mut CudaSlice<f32>,
9770 n: usize,
9771 ) -> Result<(), Box<dyn std::error::Error>> {
9772 let f = self.func("gelu_tanh_mul_f32");
9773 let cfg = LaunchConfig::for_num_elems(n as u32);
9774 let ni = n as i32;
9775 let __s_b = self.gpu.stream();
9776 let mut b = __s_b.launch_builder(&f);
9777 b.arg(gate).arg(up).arg(dst).arg(&ni);
9778 unsafe {
9779 b.launch(cfg)?;
9780 }
9781 Ok(())
9782 }
9783
9784 pub fn silu_mul(
9785 &self,
9786 gate: &CudaSlice<f32>,
9787 up: &CudaSlice<f32>,
9788 dst: &mut CudaSlice<f32>,
9789 n: usize,
9790 ) -> Result<(), Box<dyn std::error::Error>> {
9791 let f = self.func("silu_mul_f32");
9792 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9794 let ni = n as i32;
9795 let __s_b = self.gpu.stream();
9796 let mut b = __s_b.launch_builder(&f);
9797 b.arg(gate).arg(up).arg(dst).arg(&ni);
9798 unsafe {
9799 b.launch(cfg)?;
9800 }
9801 Ok(())
9802 }
9803
9804 pub fn silu_mul_f16out(
9807 &self,
9808 gate: &CudaSlice<f32>,
9809 up: &CudaSlice<f32>,
9810 dst: &mut CudaSlice<f32>,
9811 dst16: &mut CudaSlice<u8>,
9812 n: usize,
9813 ) -> Result<(), Box<dyn std::error::Error>> {
9814 let f = self.func("silu_mul_f16out_f32");
9815 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9816 let ni = n as i32;
9817 let __s_b = self.gpu.stream();
9818 let mut b = __s_b.launch_builder(&f);
9819 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
9820 unsafe {
9821 b.launch(cfg)?;
9822 }
9823 Ok(())
9824 }
9825
9826 pub fn silu_mul_scaled(
9833 &self,
9834 gate: &CudaSlice<f32>,
9835 up: &CudaSlice<f32>,
9836 gs: f32,
9837 us: f32,
9838 dst: &mut CudaSlice<f32>,
9839 n: usize,
9840 ) -> Result<(), Box<dyn std::error::Error>> {
9841 let f = self.func("silu_mul_scaled_f32");
9842 let cfg = LaunchConfig::for_num_elems(n as u32);
9843 let ni = n as i32;
9844 let (gsf, usf) = (gs, us);
9845 let __s_b = self.gpu.stream();
9846 let mut b = __s_b.launch_builder(&f);
9847 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
9848 unsafe {
9849 b.launch(cfg)?;
9850 }
9851 Ok(())
9852 }
9853
9854 #[allow(clippy::too_many_arguments)]
9858 pub fn swigluoai_mul_scaled(
9859 &self,
9860 gate: &CudaSlice<f32>,
9861 up: &CudaSlice<f32>,
9862 gs: f32,
9863 us: f32,
9864 alpha: f32,
9865 limit: f32,
9866 dst: &mut CudaSlice<f32>,
9867 n: usize,
9868 ) -> Result<(), Box<dyn std::error::Error>> {
9869 let f = self.func("swigluoai_mul_scaled_f32");
9870 let cfg = LaunchConfig::for_num_elems(n as u32);
9871 let ni = n as i32;
9872 let __s_b = self.gpu.stream();
9873 let mut b = __s_b.launch_builder(&f);
9874 b.arg(gate)
9875 .arg(up)
9876 .arg(&gs)
9877 .arg(&us)
9878 .arg(&alpha)
9879 .arg(&limit)
9880 .arg(dst)
9881 .arg(&ni);
9882 unsafe {
9883 b.launch(cfg)?;
9884 }
9885 Ok(())
9886 }
9887
9888 pub fn silu_mul_scaled_q8_1(
9896 &self,
9897 gate: &CudaSlice<f32>,
9898 up: &CudaSlice<f32>,
9899 gs: f32,
9900 us: f32,
9901 n: usize,
9902 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9903 let f = self.func("silu_mul_scaled_q8_1");
9904 let nblk = n / 32;
9905 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);
9909 let (gsf, usf, ni) = (gs, us, n as i32);
9910 let __s_b = self.gpu.stream();
9911 let mut b = __s_b.launch_builder(&f);
9912 b.arg(gate)
9913 .arg(up)
9914 .arg(&gsf)
9915 .arg(&usf)
9916 .arg(&mut aq)
9917 .arg(&mut ad)
9918 .arg(&ni);
9919 unsafe {
9920 b.launch(cfg)?;
9921 }
9922 Ok((aq, ad))
9923 }
9924
9925 pub fn add(
9926 &self,
9927 a: &CudaSlice<f32>,
9928 b_in: &CudaSlice<f32>,
9929 dst: &mut CudaSlice<f32>,
9930 n: usize,
9931 ) -> Result<(), Box<dyn std::error::Error>> {
9932 let f = self.func("add_f32");
9933 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9935 let ni = n as i32;
9936 let __s_bld = self.gpu.stream();
9937 let mut bld = __s_bld.launch_builder(&f);
9938 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9939 unsafe {
9940 bld.launch(cfg)?;
9941 }
9942 Ok(())
9943 }
9944
9945 pub fn mul(
9946 &self,
9947 a: &CudaSlice<f32>,
9948 b_in: &CudaSlice<f32>,
9949 dst: &mut CudaSlice<f32>,
9950 n: usize,
9951 ) -> Result<(), Box<dyn std::error::Error>> {
9952 let f = self.func("mul_f32");
9953 let cfg = LaunchConfig::for_num_elems(n as u32);
9954 let ni = n as i32;
9955 let __s_bld = self.gpu.stream();
9956 let mut bld = __s_bld.launch_builder(&f);
9957 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9958 unsafe {
9959 bld.launch(cfg)?;
9960 }
9961 Ok(())
9962 }
9963
9964 pub fn matmul(
9967 &self,
9968 w: &crate::model::GpuTensor,
9969 x: &CudaSlice<f32>,
9970 m: usize,
9971 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9972 use crate::model::GpuTensor;
9973 let in_f = w.in_features();
9974 let out_f = w.out_features();
9975 #[allow(non_snake_case)]
9983 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
9986 usize::MAX
9987 } else {
9988 16usize
9989 };
9990
9991 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
10016 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
10017 return Ok(y);
10018 }
10019 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
10026 return Ok(y);
10027 }
10028 if let Some(y) = self.try_f16_gemm(w, x, m)? {
10031 return Ok(y);
10032 }
10033 }
10034 if let GpuTensor::Quant { qtype, .. } = w {
10049 if *qtype == QT_F8_E4M3_BLK {
10050 if m >= GEMM_M_THRESHOLD {
10051 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
10052 return Ok(y);
10053 }
10054 }
10055 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10056 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10057 return Ok(y);
10058 }
10059 }
10060 }
10061 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
10062 return self.qmatvec_mmq(w, x, m);
10063 }
10064 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
10065 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10066 return self.qmatvec_gemm(w, &aq, &ad, m);
10067 }
10068 if m >= GEMM_M_THRESHOLD {
10071 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
10072 return Ok(y);
10073 }
10074 }
10075 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
10079 if m == 1 && fast {
10084 if let GpuTensor::Quant {
10085 bytes,
10086 qtype,
10087 row_bytes,
10088 rp,
10089 rp4,
10090 scale,
10091 ..
10092 } = w
10093 {
10094 if self.mmvq_supports(*qtype) {
10095 let (bytes, rp) = match rp4 {
10099 Some(m4) => (m4, true),
10100 None => (bytes, *rp),
10101 };
10102 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10103 return self.qmatvec_mmvq(
10104 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
10105 );
10106 }
10107 }
10108 }
10109 if (2..=16).contains(&m)
10125 && fast
10126 && std::env::var("MEMRA_NO_BATCHED").is_err()
10127 && (m <= 4 || Self::b8_enabled())
10128 {
10129 let m_ok = m <= 8
10139 || matches!(w, GpuTensor::Quant { qtype, .. }
10140 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
10141 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
10142 if m_ok {
10143 if let GpuTensor::Quant {
10144 bytes,
10145 qtype,
10146 row_bytes,
10147 rp,
10148 rp4,
10149 ..
10150 } = w
10151 {
10152 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
10153 let (bytes, rp) = match rp4 {
10154 Some(m4) => (m4, true),
10155 None => (bytes, *rp),
10156 };
10157 let mcols = Self::batched_mcols(m);
10158 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10159 let mut y = self.qmatvec_mmvq_batched(
10160 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
10161 )?;
10162 if let GpuTensor::Quant { scale, .. } = w {
10163 if *scale != 1.0 {
10164 self.scale_inplace(&mut y, *scale, m * out_f)?;
10165 }
10166 }
10167 return Ok(y);
10168 }
10169 }
10170 }
10171 }
10172 if fast {
10178 if let GpuTensor::Quant {
10179 bytes,
10180 qtype,
10181 row_bytes,
10182 scale,
10183 ..
10184 } = w
10185 {
10186 if *qtype == QT_F8_E4M3 {
10187 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10188 return self.qmatvec_mmvq(
10189 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10190 );
10191 }
10192 }
10193 }
10194 let mut y = match w {
10195 GpuTensor::Quant {
10196 bytes,
10197 qtype,
10198 row_bytes,
10199 ..
10200 } if fast && *qtype == QT_Q8_0 => {
10201 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10202 }
10203 GpuTensor::Quant {
10204 bytes,
10205 qtype,
10206 row_bytes,
10207 ..
10208 } if fast && *qtype == QT_Q4_K => {
10209 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10210 }
10211 GpuTensor::Quant {
10212 bytes,
10213 qtype,
10214 row_bytes,
10215 ..
10216 } if fast && *qtype == QT_Q6_K => {
10217 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10218 }
10219 GpuTensor::Quant {
10220 bytes,
10221 qtype,
10222 row_bytes,
10223 ..
10224 } if fast && *qtype == QT_Q5_K => {
10225 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10226 }
10227 GpuTensor::Quant {
10228 bytes,
10229 qtype,
10230 row_bytes,
10231 ..
10232 } if fast && *qtype == QT_Q3_K => {
10233 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10234 }
10235 GpuTensor::Quant {
10236 bytes,
10237 qtype,
10238 row_bytes,
10239 rp,
10240 ..
10241 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10242 if *rp {
10243 "qmatvec_nvfp4_dp4a_rp"
10244 } else {
10245 "qmatvec_nvfp4_dp4a"
10246 },
10247 bytes,
10248 x,
10249 m,
10250 in_f,
10251 out_f,
10252 *row_bytes,
10253 )?,
10254 GpuTensor::Quant {
10258 bytes,
10259 qtype,
10260 row_bytes,
10261 ..
10262 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10263 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10264 }
10265 GpuTensor::Quant {
10270 bytes,
10271 qtype,
10272 row_bytes,
10273 rp,
10274 ..
10275 } =>
10276 {
10279 self.qmatvec(
10280 bytes,
10281 x,
10282 m,
10283 in_f,
10284 out_f,
10285 if *rp && *qtype == QT_NVFP4 {
10286 QT_NVFP4_RP
10287 } else {
10288 *qtype
10289 },
10290 *row_bytes,
10291 )?
10292 }
10293 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10294 GpuTensor::FloatBf16 { data, .. } => {
10297 self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10298 }
10299 };
10300 if let GpuTensor::Quant { scale, .. } = w {
10302 if *scale != 1.0 {
10303 self.scale_inplace(&mut y, *scale, m * out_f)?;
10304 }
10305 }
10306 Ok(y)
10307 }
10308
10309 pub fn stage_a_raw_needed() -> bool {
10319 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10320 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
10321 }
10322
10323 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10326 use crate::model::GpuTensor;
10327 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10328 return false;
10329 }
10330 match w {
10331 GpuTensor::Quant { qtype, .. } => {
10338 matches!(
10339 *qtype,
10340 QT_Q8_0
10341 | QT_Q4_K
10342 | QT_Q6_K
10343 | QT_Q5_K
10344 | QT_Q3_K
10345 | QT_NVFP4
10346 | QT_F8_E4M3
10347 | QT_F8_E4M3_BLK
10348 | QT_Q4_0
10349 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10350 }
10351 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10352 }
10353 }
10354
10355 pub fn matmul_pre(
10360 &self,
10361 w: &crate::model::GpuTensor,
10362 aq: &CudaSlice<i8>,
10363 ad: &CudaSlice<f32>,
10364 x_fallback: &CudaSlice<f32>,
10365 m: usize,
10366 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10367 use crate::model::GpuTensor;
10368 let x_raw_ok = x_fallback.len() >= m * w.in_features();
10374 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10377 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10378 return Ok(y);
10379 }
10380 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10383 return Ok(y);
10384 }
10385 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10387 return Ok(y);
10388 }
10389 }
10390 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10396 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10397 return Ok(y);
10398 }
10399 }
10400 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10401 return Ok(y);
10402 }
10403 if m >= 16
10408 && w.out_features() >= 128
10409 && self.mmq_supports(w)
10410 && !self.verify_exact_on()
10411 && x_raw_ok
10412 {
10413 return self.qmatvec_mmq(w, x_fallback, m);
10414 }
10415 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10418 if let Some(y) =
10419 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10420 {
10421 return Ok(y);
10422 }
10423 }
10424 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10427 return self.qmatvec_gemm(w, aq, ad, m);
10428 }
10429 if !self.uses_q8_1_fast(w) {
10448 if !x_raw_ok {
10449 return Err(format!(
10450 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
10451 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
10452 dropped the f32, so there is nothing to fall back to — pass the real f32 \
10453 activation (see Engine::rms_norm_decode, which is bit-identical to \
10454 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
10455 x_fallback.len(),
10456 m,
10457 w.in_features(),
10458 m * w.in_features()
10459 )
10460 .into());
10461 }
10462 return self.matmul(w, x_fallback, m);
10463 }
10464 let in_f = w.in_features();
10465 let out_f = w.out_features();
10466 let (bytes, qtype, row_bytes, scale, rp) = match w {
10467 GpuTensor::Quant {
10468 bytes,
10469 qtype,
10470 row_bytes,
10471 scale,
10472 rp,
10473 ..
10474 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10475 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10476 };
10477 let (mbytes, mrp) = match w {
10480 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10481 _ => (bytes, rp),
10482 };
10483 if m == 1 && self.mmvq_supports(qtype) {
10487 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10488 }
10489 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10502 && std::env::var("MEMRA_NO_BATCHED").is_err()
10503 && (m <= 4 || Self::b8_enabled())
10504 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10508 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10509 {
10510 let mcols = Self::batched_mcols(m);
10511 return self.qmatvec_mmvq_batched(
10512 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10513 );
10514 }
10515 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10521 let (b2, r2) = if qtype == QT_Q4_0 {
10522 (mbytes, mrp)
10523 } else {
10524 (bytes, rp)
10525 };
10526 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10527 }
10528 let name = match qtype {
10529 QT_Q8_0 => "qmatvec_q8_0_dp4a",
10530 QT_Q4_K => "qmatvec_q4_K_dp4a",
10531 QT_Q6_K => "qmatvec_q6_K_dp4a",
10532 QT_Q5_K => "qmatvec_q5_K_dp4a",
10533 QT_Q3_K => "qmatvec_q3_K_dp4a",
10534 QT_NVFP4 => {
10535 if rp {
10536 "qmatvec_nvfp4_dp4a_rp"
10537 } else {
10538 "qmatvec_nvfp4_dp4a"
10539 }
10540 }
10541 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10542 _ => unreachable!(),
10543 };
10544 let f = self.func(name);
10545 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10547 grid_dim: (out_f as u32, m as u32, 1),
10548 block_dim: (128, 1, 1),
10549 shared_mem_bytes: 0,
10550 };
10551 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10552 let __s_b = self.gpu.stream();
10553 let mut b = __s_b.launch_builder(&f);
10554 b.arg(bytes)
10555 .arg(aq)
10556 .arg(ad)
10557 .arg(&mut y)
10558 .arg(&inf)
10559 .arg(&outf)
10560 .arg(&mi)
10561 .arg(&rb);
10562 unsafe {
10563 b.launch(cfg)?;
10564 }
10565 if scale != 1.0 {
10566 self.scale_inplace(&mut y, scale, m * out_f)?;
10567 }
10568 Ok(y)
10569 }
10570
10571 pub fn matmul_decode_exact(
10579 &self,
10580 w: &crate::model::GpuTensor,
10581 x: &CudaSlice<f32>,
10582 m: usize,
10583 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10584 use crate::model::GpuTensor;
10585 if let GpuTensor::Float { data, .. } = w {
10593 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10594 }
10595 if let GpuTensor::FloatBf16 { data, .. } = w {
10598 let (in_f, out_f) = (w.in_features(), w.out_features());
10599 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10600 }
10601 if !self.uses_q8_1_fast(w) {
10602 return self.matmul(w, x, m);
10603 }
10604 let in_f = w.in_features();
10605 let out_f = w.out_features();
10606 let (bytes, qtype, row_bytes, scale, rp) = match w {
10607 GpuTensor::Quant {
10608 bytes,
10609 qtype,
10610 row_bytes,
10611 scale,
10612 rp,
10613 ..
10614 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10615 _ => return self.matmul(w, x, m),
10616 };
10617 let (bytes, rp) = match w {
10620 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10621 _ => (bytes, rp),
10622 };
10623 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10624 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10628 return Ok(y);
10629 }
10630 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10639 && std::env::var("MEMRA_NO_BATCHED").is_err()
10640 && (m <= 4 || Self::b8_enabled())
10641 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10644 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10645 {
10646 let mcols = Self::batched_mcols(m);
10647 return self.qmatvec_mmvq_batched(
10648 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10649 );
10650 }
10651 if self.mmvq_supports(qtype) {
10652 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10655 }
10656 self.matmul_pre(w, &aq, &ad, x, m)
10659 }
10660
10661 pub fn matmul_decode_exact_pre(
10671 &self,
10672 w: &crate::model::GpuTensor,
10673 aq: &CudaSlice<i8>,
10674 ad: &CudaSlice<f32>,
10675 m: usize,
10676 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10677 use crate::model::GpuTensor;
10678 debug_assert!(
10679 self.uses_q8_1_fast(w),
10680 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10681 );
10682 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10684 return Ok(y);
10685 }
10686 let in_f = w.in_features();
10687 let out_f = w.out_features();
10688 let (bytes, qtype, row_bytes, scale, rp) = match w {
10689 GpuTensor::Quant {
10690 bytes,
10691 qtype,
10692 row_bytes,
10693 scale,
10694 rp,
10695 ..
10696 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10697 _ => {
10698 return Err(
10699 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10700 );
10701 }
10702 };
10703 let (bytes, rp) = match w {
10705 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10706 _ => (bytes, rp),
10707 };
10708 if (2..=16).contains(&m)
10710 && self.batched_supports(qtype)
10711 && self.mmvq_supports(qtype)
10712 && std::env::var("MEMRA_NO_BATCHED").is_err()
10713 && (m <= 4 || Self::b8_enabled())
10714 && (m <= 8
10715 || qtype == QT_Q4_0
10716 || qtype == QT_Q6_K
10717 || qtype == QT_F8_E4M3
10718 || qtype == QT_NVFP4
10719 || qtype == QT_Q4_K
10720 || qtype == QT_Q5_K
10721 || qtype == QT_Q8_0)
10722 {
10723 let mcols = Self::batched_mcols(m);
10724 return self.qmatvec_mmvq_batched(
10725 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10726 );
10727 }
10728 if self.mmvq_supports(qtype) {
10729 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10730 }
10731 let x0 = self.zeros(0)?;
10734 self.matmul_pre(w, aq, ad, &x0, m)
10735 }
10736
10737 pub fn matmul_decode_exact_dual_pre(
10746 &self,
10747 w0: &crate::model::GpuTensor,
10748 w1: &crate::model::GpuTensor,
10749 aq: &CudaSlice<i8>,
10750 ad: &CudaSlice<f32>,
10751 m: usize,
10752 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10753 {
10754 use crate::model::GpuTensor;
10755 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10756 let on = *ON.get_or_init(|| {
10757 std::env::var("MEMRA_SPEC_DUAL_T")
10758 .map(|v| v != "0")
10759 .unwrap_or(true)
10760 });
10761 if !on
10762 || !(2..=7).contains(&m)
10763 || std::env::var("MEMRA_NO_BATCHED").is_ok()
10764 || !self.uses_q8_1_fast(w0)
10765 || !self.uses_q8_1_fast(w1)
10766 {
10767 return Ok(None);
10768 }
10769 if !self.mmvq_supports(QT_NVFP4) {
10774 return Ok(None);
10775 }
10776 let (in_f, out_f) = (w0.in_features(), w0.out_features());
10777 if w1.in_features() != in_f || w1.out_features() != out_f {
10778 return Ok(None);
10779 }
10780 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10781 (
10782 GpuTensor::Quant {
10783 bytes: b0,
10784 qtype: q0,
10785 row_bytes: rb0,
10786 scale: s0,
10787 rp: rp0,
10788 rp4: None,
10789 ..
10790 },
10791 GpuTensor::Quant {
10792 bytes: b1,
10793 qtype: q1,
10794 row_bytes: rb1,
10795 scale: s1,
10796 rp: rp1,
10797 rp4: None,
10798 ..
10799 },
10800 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10801 (b0, b1, *rb0, *s0, *s1, *rp0)
10802 }
10803 _ => return Ok(None),
10804 };
10805 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
10808 {
10809 return Ok(None);
10810 }
10811 let (y0, y1) =
10812 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
10813 Ok(Some(((y0, s0), (y1, s1))))
10814 }
10815
10816 pub fn matmul_decode_exact_dual(
10832 &self,
10833 w0: &crate::model::GpuTensor,
10834 w1: &crate::model::GpuTensor,
10835 x: &CudaSlice<f32>,
10836 m: usize,
10837 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
10838 use crate::model::GpuTensor;
10839 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10840 let on = *ON.get_or_init(|| {
10841 std::env::var("MEMRA_SPEC_DUAL_T")
10842 .map(|v| v != "0")
10843 .unwrap_or(true)
10844 });
10845 if !on
10846 || !(2..=4).contains(&m)
10847 || std::env::var("MEMRA_NO_BATCHED").is_ok()
10848 || !self.uses_q8_1_fast(w0)
10849 || !self.uses_q8_1_fast(w1)
10850 {
10851 return Ok(None);
10852 }
10853 if !self.mmvq_supports(QT_NVFP4) {
10858 return Ok(None);
10859 }
10860 let (in_f, out_f) = (w0.in_features(), w0.out_features());
10861 if w1.in_features() != in_f || w1.out_features() != out_f {
10862 return Ok(None);
10863 }
10864 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10865 (
10866 GpuTensor::Quant {
10867 bytes: b0,
10868 qtype: q0,
10869 row_bytes: rb0,
10870 scale: s0,
10871 rp: rp0,
10872 rp4: None,
10873 ..
10874 },
10875 GpuTensor::Quant {
10876 bytes: b1,
10877 qtype: q1,
10878 row_bytes: rb1,
10879 scale: s1,
10880 rp: rp1,
10881 rp4: None,
10882 ..
10883 },
10884 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10885 (b0, b1, *rb0, *s0, *s1, *rp0)
10886 }
10887 _ => return Ok(None),
10888 };
10889 if std::env::var("MEMRA_DEBUG").is_ok() {
10892 static ONCE: std::sync::Once = std::sync::Once::new();
10893 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
10894 }
10895 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10896 let (y0, y1) =
10897 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
10898 let mut y0 = y0;
10899 let mut y1 = y1;
10900 if s0 != 1.0 {
10901 self.scale_inplace(&mut y0, s0, m * out_f)?;
10902 }
10903 if s1 != 1.0 {
10904 self.scale_inplace(&mut y1, s1, m * out_f)?;
10905 }
10906 Ok(Some((y0, y1)))
10907 }
10908
10909 #[allow(clippy::too_many_arguments)]
10914 pub fn qmatvec_batched_dual_raw(
10915 &self,
10916 b0: &CudaSlice<u8>,
10917 b1: &CudaSlice<u8>,
10918 aq: &CudaSlice<i8>,
10919 ad: &CudaSlice<f32>,
10920 m: usize,
10921 in_f: usize,
10922 out_f: usize,
10923 row_bytes: usize,
10924 rp: bool,
10925 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10926 const ROWS_PER_BLOCK: u32 = 4;
10927 let mcols = Self::batched_mcols(m);
10928 let tiny_rp1 = rp
10931 && mcols == 4
10932 && out_f <= 128
10933 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
10934 let (name, rows_per_block) = if tiny_rp1 {
10935 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
10936 } else {
10937 match (mcols, rp, m) {
10938 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
10939 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
10940 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
10941 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
10942 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
10943 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
10944 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
10945 _ => {
10946 return Err(
10947 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
10948 );
10949 }
10950 }
10951 };
10952 let f = self.func(name);
10953 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
10954 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
10955 let cfg = LaunchConfig {
10956 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
10957 block_dim: (32, ROWS_PER_BLOCK, 1),
10958 shared_mem_bytes: 0,
10959 };
10960 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10961 let __s_b = self.gpu.stream();
10962 let mut b = __s_b.launch_builder(&f);
10963 b.arg(b0)
10964 .arg(b1)
10965 .arg(aq)
10966 .arg(ad)
10967 .arg(&mut y0)
10968 .arg(&mut y1)
10969 .arg(&inf)
10970 .arg(&outf)
10971 .arg(&mi)
10972 .arg(&rb);
10973 unsafe {
10974 b.launch(cfg)?;
10975 }
10976 Ok((y0, y1))
10977 }
10978
10979 pub fn matmul_pre_dual_noscale(
10991 &self,
10992 w0: &crate::model::GpuTensor,
10993 w1: &crate::model::GpuTensor,
10994 aq: &CudaSlice<i8>,
10995 ad: &CudaSlice<f32>,
10996 m: usize,
10997 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10998 {
10999 use crate::model::GpuTensor;
11000 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11001 return Ok(None);
11002 }
11003 if !self.mmvq_supports(QT_NVFP4) {
11013 return Ok(None);
11014 }
11015 let (in_f, out_f) = (w0.in_features(), w0.out_features());
11016 if w1.in_features() != in_f || w1.out_features() != out_f {
11017 return Ok(None);
11018 }
11019 let no_mirror =
11032 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
11033 if self.q8_ffn_fuse2_on()
11034 && no_mirror(w0)
11035 && no_mirror(w1)
11036 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
11037 {
11038 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
11039 return Ok(Some(((y0, 1.0), (y1, 1.0))));
11040 }
11041 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11051 let (y0, y1) =
11052 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
11053 return Ok(Some(((y0, p0.3), (y1, p1.3))));
11054 }
11055 let (b0, q0, rb0, s0, rp0) = match w0 {
11056 GpuTensor::Quant {
11057 bytes,
11058 qtype,
11059 row_bytes,
11060 scale,
11061 rp,
11062 ..
11063 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11064 _ => return Ok(None),
11065 };
11066 let (b1, q1, rb1, s1, rp1) = match w1 {
11067 GpuTensor::Quant {
11068 bytes,
11069 qtype,
11070 row_bytes,
11071 scale,
11072 rp,
11073 ..
11074 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11075 _ => return Ok(None),
11076 };
11077 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
11078 return Ok(None);
11079 }
11080 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11082 let rows_per_block = ROWS_PER_BLOCK * RPW;
11083 let f = self.func(if rp0 {
11084 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
11085 } else {
11086 "qmatvec_nvfp4_mmvq_dual_mr2"
11087 });
11088 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
11089 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
11090 let cfg = LaunchConfig {
11091 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11092 block_dim: (32, ROWS_PER_BLOCK, 1),
11093 shared_mem_bytes: 0,
11094 };
11095 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
11096 let one = 1.0f32;
11099 let __s_b = self.gpu.stream();
11100 let mut b = __s_b.launch_builder(&f);
11101 b.arg(b0)
11102 .arg(b1)
11103 .arg(aq)
11104 .arg(ad)
11105 .arg(&mut y0)
11106 .arg(&mut y1)
11107 .arg(&inf)
11108 .arg(&outf)
11109 .arg(&mi)
11110 .arg(&rb)
11111 .arg(&one)
11112 .arg(&one);
11113 unsafe {
11114 b.launch(cfg)?;
11115 }
11116 Ok(Some(((y0, s0), (y1, s1))))
11117 }
11118
11119 #[allow(clippy::too_many_arguments)]
11127 pub fn matmul_nvfp4_fused3(
11128 &self,
11129 w0: &crate::model::GpuTensor,
11130 w1: &crate::model::GpuTensor,
11131 w2: &crate::model::GpuTensor,
11132 aq: &CudaSlice<i8>,
11133 ad: &CudaSlice<f32>,
11134 m: usize,
11135 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11136 {
11137 use crate::model::GpuTensor;
11138 if m != 1
11142 || !self.mmvq_supports(QT_NVFP4)
11143 || !self.uses_q8_1_fast(w0)
11144 || !self.uses_q8_1_fast(w1)
11145 || !self.uses_q8_1_fast(w2)
11146 {
11147 return Ok(None);
11148 }
11149 let unpack = |w: &crate::model::GpuTensor| match w {
11150 GpuTensor::Quant {
11151 bytes,
11152 qtype,
11153 scale,
11154 rp,
11155 ..
11156 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11157 _ => None,
11158 };
11159 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
11160 return Ok(None);
11161 };
11162 let in_f = w0.in_features();
11163 if w1.in_features() != in_f || w2.in_features() != in_f {
11164 return Ok(None);
11165 }
11166 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
11167 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11169 let rows_pb = ROWS_PER_BLOCK * RPW;
11170 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11171 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
11172 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11173 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11174 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11175 let cfg = LaunchConfig {
11176 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
11177 block_dim: (32, ROWS_PER_BLOCK, 1),
11178 shared_mem_bytes: 0,
11179 };
11180 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
11181 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
11184 let __s_b = self.gpu.stream();
11185 let mut b = __s_b.launch_builder(&f);
11186 b.arg(b0)
11187 .arg(b1)
11188 .arg(b2)
11189 .arg(aq)
11190 .arg(ad)
11191 .arg(&mut y0)
11192 .arg(&mut y1)
11193 .arg(&mut y2)
11194 .arg(&inf)
11195 .arg(&oi0)
11196 .arg(&oi1)
11197 .arg(&oi2)
11198 .arg(&mi)
11199 .arg(&p0.1)
11200 .arg(&p1.1)
11201 .arg(&p2.1);
11202 unsafe {
11203 b.launch(cfg)?;
11204 }
11205 Ok(Some((y0, y1, y2)))
11206 }
11207
11208 pub fn matmul_nvfp4_fused2(
11217 &self,
11218 w0: &crate::model::GpuTensor,
11219 w1: &crate::model::GpuTensor,
11220 aq: &CudaSlice<i8>,
11221 ad: &CudaSlice<f32>,
11222 m: usize,
11223 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11224 use crate::model::GpuTensor;
11225 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11226 let off =
11227 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11228 if off
11231 || m != 1
11232 || !self.mmvq_supports(QT_NVFP4)
11233 || !self.uses_q8_1_fast(w0)
11234 || !self.uses_q8_1_fast(w1)
11235 {
11236 return Ok(None);
11237 }
11238 let unpack = |w: &crate::model::GpuTensor| match w {
11239 GpuTensor::Quant {
11240 bytes,
11241 qtype,
11242 scale,
11243 rp,
11244 ..
11245 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11246 _ => None,
11247 };
11248 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11249 return Ok(None);
11250 };
11251 let in_f = w0.in_features();
11252 if w1.in_features() != in_f {
11253 return Ok(None);
11254 }
11255 let (o0, o1) = (w0.out_features(), w1.out_features());
11256 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11258 let rows_pb = ROWS_PER_BLOCK * RPW;
11259 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11260 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11261 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11262 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11263 let cfg = LaunchConfig {
11264 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11265 block_dim: (32, ROWS_PER_BLOCK, 1),
11266 shared_mem_bytes: 0,
11267 };
11268 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11269 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11272 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11275 {
11276 use cudarc::driver::{DevicePtr, DevicePtrMut};
11277 let s = &self.gpu.stream();
11278 let (pw0, _g0) = b0.device_ptr(s);
11279 let (pw1, _g1) = b1.device_ptr(s);
11280 let (paq, _g2) = aq.device_ptr(s);
11281 let (pad, _g3) = ad.device_ptr(s);
11282 let (py0, _g4) = y0.device_ptr_mut(s);
11283 let (py1, _g5) = y1.device_ptr_mut(s);
11284 let (s0, s1) = (p0.1, p1.1);
11285 let mut ps = [
11286 &pw0 as *const _ as *mut std::ffi::c_void,
11287 &pw1 as *const _ as *mut _,
11288 &paq as *const _ as *mut _,
11289 &pad as *const _ as *mut _,
11290 &py0 as *const _ as *mut _,
11291 &py1 as *const _ as *mut _,
11292 &inf as *const _ as *mut _,
11293 &oi0 as *const _ as *mut _,
11294 &oi1 as *const _ as *mut _,
11295 &mi as *const _ as *mut _,
11296 &s0 as *const _ as *mut _,
11297 &s1 as *const _ as *mut _,
11298 ];
11299 unsafe {
11300 self.launch_pdl(
11301 "qmatvec_nvfp4_mmvq_fused2_rp",
11302 cfg.grid_dim,
11303 cfg.block_dim,
11304 &mut ps,
11305 )?;
11306 }
11307 }
11308 return Ok(Some((y0, y1)));
11309 }
11310 let __s_b = self.gpu.stream();
11311 let mut b = __s_b.launch_builder(&f);
11312 b.arg(b0)
11313 .arg(b1)
11314 .arg(aq)
11315 .arg(ad)
11316 .arg(&mut y0)
11317 .arg(&mut y1)
11318 .arg(&inf)
11319 .arg(&oi0)
11320 .arg(&oi1)
11321 .arg(&mi)
11322 .arg(&p0.1)
11323 .arg(&p1.1);
11324 unsafe {
11325 b.launch(cfg)?;
11326 }
11327 Ok(Some((y0, y1)))
11328 }
11329
11330 pub fn matmul_nvfp4_fused2_into(
11335 &self,
11336 w0: &crate::model::GpuTensor,
11337 w1: &crate::model::GpuTensor,
11338 aq: &CudaSlice<i8>,
11339 ad: &CudaSlice<f32>,
11340 y0: &mut CudaSlice<f32>,
11341 y1: &mut CudaSlice<f32>,
11342 ) -> Result<bool, Box<dyn std::error::Error>> {
11343 use crate::model::GpuTensor;
11344 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11345 let off =
11346 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11347 if off
11348 || !self.mmvq_supports(QT_NVFP4)
11349 || !self.uses_q8_1_fast(w0)
11350 || !self.uses_q8_1_fast(w1)
11351 {
11352 return Ok(false);
11353 }
11354 let unpack = |w: &crate::model::GpuTensor| match w {
11355 GpuTensor::Quant {
11356 bytes,
11357 qtype,
11358 scale,
11359 rp,
11360 ..
11361 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11362 _ => None,
11363 };
11364 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11365 return Ok(false);
11366 };
11367 let in_f = w0.in_features();
11368 if w1.in_features() != in_f {
11369 return Ok(false);
11370 }
11371 let (o0, o1) = (w0.out_features(), w1.out_features());
11372 if y0.len() < o0 || y1.len() < o1 {
11373 return Ok(false);
11374 }
11375 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11377 let rows_pb = ROWS_PER_BLOCK * RPW;
11378 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11379 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11380 let cfg = LaunchConfig {
11381 grid_dim: (nb(o0) + nb(o1), 1, 1),
11382 block_dim: (32, ROWS_PER_BLOCK, 1),
11383 shared_mem_bytes: 0,
11384 };
11385 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11386 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11389 let __s_b = self.gpu.stream();
11390 let mut b = __s_b.launch_builder(&f);
11391 b.arg(b0)
11392 .arg(b1)
11393 .arg(aq)
11394 .arg(ad)
11395 .arg(&mut *y0)
11396 .arg(&mut *y1)
11397 .arg(&inf)
11398 .arg(&oi0)
11399 .arg(&oi1)
11400 .arg(&mi)
11401 .arg(&p0.1)
11402 .arg(&p1.1);
11403 unsafe {
11404 b.launch(cfg)?;
11405 }
11406 Ok(true)
11407 }
11408
11409 #[allow(clippy::type_complexity)]
11414 pub fn matmul_nvfp4_fused4(
11415 &self,
11416 w0: &crate::model::GpuTensor,
11417 w1: &crate::model::GpuTensor,
11418 w2: &crate::model::GpuTensor,
11419 w3: &crate::model::GpuTensor,
11420 aq: &CudaSlice<i8>,
11421 ad: &CudaSlice<f32>,
11422 m: usize,
11423 ) -> Result<
11424 Option<(
11425 CudaSlice<f32>,
11426 CudaSlice<f32>,
11427 CudaSlice<f32>,
11428 CudaSlice<f32>,
11429 )>,
11430 Box<dyn std::error::Error>,
11431 > {
11432 use crate::model::GpuTensor;
11433 if m != 1
11435 || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11436 || !self.mmvq_supports(QT_NVFP4)
11437 || !self.uses_q8_1_fast(w0)
11438 || !self.uses_q8_1_fast(w1)
11439 || !self.uses_q8_1_fast(w2)
11440 || !self.uses_q8_1_fast(w3)
11441 {
11442 return Ok(None);
11443 }
11444 let unpack = |w: &crate::model::GpuTensor| match w {
11445 GpuTensor::Quant {
11446 bytes,
11447 qtype,
11448 scale,
11449 rp,
11450 ..
11451 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11452 _ => None,
11453 };
11454 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11455 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11456 else {
11457 return Ok(None);
11458 };
11459 let in_f = w0.in_features();
11460 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11461 return Ok(None);
11462 }
11463 let (o0, o1, o2, o3) = (
11464 w0.out_features(),
11465 w1.out_features(),
11466 w2.out_features(),
11467 w3.out_features(),
11468 );
11469 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11471 let rows_pb = ROWS_PER_BLOCK * RPW;
11472 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11473 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
11474 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11475 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11476 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11477 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11478 let cfg = LaunchConfig {
11479 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
11480 block_dim: (32, ROWS_PER_BLOCK, 1),
11481 shared_mem_bytes: 0,
11482 };
11483 let (inf, oi0, oi1, oi2, oi3, mi) = (
11484 in_f as i32,
11485 o0 as i32,
11486 o1 as i32,
11487 o2 as i32,
11488 o3 as i32,
11489 m as i32,
11490 );
11491 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11494 let __s_b = self.gpu.stream();
11495 let mut b = __s_b.launch_builder(&f);
11496 b.arg(b0)
11497 .arg(b1)
11498 .arg(b2)
11499 .arg(b3)
11500 .arg(aq)
11501 .arg(ad)
11502 .arg(&mut y0)
11503 .arg(&mut y1)
11504 .arg(&mut y2)
11505 .arg(&mut y3)
11506 .arg(&inf)
11507 .arg(&oi0)
11508 .arg(&oi1)
11509 .arg(&oi2)
11510 .arg(&oi3)
11511 .arg(&mi)
11512 .arg(&p0.1)
11513 .arg(&p1.1)
11514 .arg(&p2.1)
11515 .arg(&p3.1);
11516 unsafe {
11517 b.launch(cfg)?;
11518 }
11519 Ok(Some((y0, y1, y2, y3)))
11520 }
11521
11522 pub fn matmul_q8_fused2(
11530 &self,
11531 w0: &crate::model::GpuTensor,
11532 w1: &crate::model::GpuTensor,
11533 aq: &CudaSlice<i8>,
11534 ad: &CudaSlice<f32>,
11535 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11536 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11542 return Ok(Some(self.e4m3_fused2_core(
11543 p0.0,
11544 p1.0,
11545 aq,
11546 ad,
11547 w0.in_features(),
11548 p0.1,
11549 p1.1,
11550 p0.2,
11551 p0.3,
11552 p1.3,
11553 )?));
11554 }
11555 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11556 return Ok(None);
11557 };
11558 Ok(Some(self.q8_fused2_core(
11559 p0.0,
11560 p1.0,
11561 aq,
11562 ad,
11563 w0.in_features(),
11564 p0.1,
11565 p1.1,
11566 p0.2,
11567 )?))
11568 }
11569
11570 #[allow(clippy::too_many_arguments)]
11571 fn q8_fused2_core(
11572 &self,
11573 b0: &CudaSlice<u8>,
11574 b1: &CudaSlice<u8>,
11575 aq: &CudaSlice<i8>,
11576 ad: &CudaSlice<f32>,
11577 in_f: usize,
11578 out0: usize,
11579 out1: usize,
11580 row_bytes: usize,
11581 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11582 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
11584 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
11585 let f = self.func("qmatvec_q8_0_mmvq_fused2");
11586 let mut y0 = self.alloc_uninit::<f32>(out0)?;
11587 let mut y1 = self.alloc_uninit::<f32>(out1)?;
11588 let cfg = LaunchConfig {
11589 grid_dim: (nb0 + nb1, 1, 1),
11590 block_dim: (32, ROWS_PER_BLOCK, 1),
11591 shared_mem_bytes: 0,
11592 };
11593 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
11594 let __s_b = self.gpu.stream();
11595 let mut b = __s_b.launch_builder(&f);
11596 b.arg(b0)
11597 .arg(b1)
11598 .arg(aq)
11599 .arg(ad)
11600 .arg(&mut y0)
11601 .arg(&mut y1)
11602 .arg(&inf)
11603 .arg(&o0)
11604 .arg(&o1)
11605 .arg(&rbl);
11606 unsafe {
11607 b.launch(cfg)?;
11608 }
11609 Ok((y0, y1))
11610 }
11611
11612 pub fn matmul_q8_fused2_x(
11618 &self,
11619 w0: &crate::model::GpuTensor,
11620 w1: &crate::model::GpuTensor,
11621 x: &CudaSlice<f32>,
11622 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11623 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11624 return Ok(None);
11625 }
11626 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11627 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11628 return Ok(Some(self.e4m3_fused2_core(
11629 p0.0,
11630 p1.0,
11631 &aq,
11632 &ad,
11633 w0.in_features(),
11634 p0.1,
11635 p1.1,
11636 p0.2,
11637 p0.3,
11638 p1.3,
11639 )?));
11640 }
11641 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11642 return Ok(None);
11643 };
11644 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11645 Ok(Some(self.q8_fused2_core(
11646 p0.0,
11647 p1.0,
11648 &aq,
11649 &ad,
11650 w0.in_features(),
11651 p0.1,
11652 p1.1,
11653 p0.2,
11654 )?))
11655 }
11656
11657 #[allow(clippy::too_many_arguments)]
11660 pub fn qmatvec_q8_fused2_raw(
11661 &self,
11662 b0: &CudaSlice<u8>,
11663 b1: &CudaSlice<u8>,
11664 x: &CudaSlice<f32>,
11665 in_f: usize,
11666 out0: usize,
11667 out1: usize,
11668 row_bytes: usize,
11669 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11670 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
11671 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
11672 }
11673
11674 pub fn matmul_q4_fused3(
11680 &self,
11681 w0: &crate::model::GpuTensor,
11682 w1: &crate::model::GpuTensor,
11683 w2: &crate::model::GpuTensor,
11684 aq: &CudaSlice<i8>,
11685 ad: &CudaSlice<f32>,
11686 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11687 {
11688 use crate::model::GpuTensor;
11689 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11690 match w {
11691 GpuTensor::Quant {
11692 qtype, row_bytes, ..
11693 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11694 _ => None,
11695 }
11696 };
11697 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11698 return Ok(None);
11699 };
11700 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11701 return Ok(None);
11702 }
11703 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11707 match w {
11708 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11709 Some(m) => (m, true),
11710 None => (bytes, *rp),
11711 },
11712 _ => unreachable!(),
11713 }
11714 }
11715 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11716 if rp0 != rp1 || rp1 != rp2 {
11717 return Ok(None);
11718 }
11719 let rp = rp0;
11720 let rpb: u32 = 4;
11721 let mr1 = rp && Self::q40_mr1_on();
11725 let nb = |o: usize| {
11726 if mr1 {
11727 (o as u32).div_ceil(rpb)
11728 } else {
11729 (o as u32).div_ceil(2).div_ceil(rpb)
11730 }
11731 };
11732 let grid = nb(o0) + nb(o1) + nb(o2);
11733 let mut y0 = self.alloc_uninit::<f32>(o0)?;
11734 let mut y1 = self.alloc_uninit::<f32>(o1)?;
11735 let mut y2 = self.alloc_uninit::<f32>(o2)?;
11736 let f = self.func(if mr1 {
11737 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11738 } else if rp {
11739 "qmatvec_q4_0_mmvq_fused3_rp"
11740 } else {
11741 "qmatvec_q4_0_mmvq_fused3"
11742 });
11743 let cfg = LaunchConfig {
11744 grid_dim: (grid, 1, 1),
11745 block_dim: (32, rpb, 1),
11746 shared_mem_bytes: 0,
11747 };
11748 let inf = w0.in_features() as i32;
11749 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11750 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11751 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11754 {
11755 use cudarc::driver::{DevicePtr, DevicePtrMut};
11756 let s = &self.gpu.stream();
11757 let (p0, _g0) = b0.device_ptr(s);
11758 let (p1, _g1) = b1.device_ptr(s);
11759 let (p2, _g2) = b2.device_ptr(s);
11760 let (paq, _g3) = aq.device_ptr(s);
11761 let (pad, _g4) = ad.device_ptr(s);
11762 let (py0, _g5) = y0.device_ptr_mut(s);
11763 let (py1, _g6) = y1.device_ptr_mut(s);
11764 let (py2, _g7) = y2.device_ptr_mut(s);
11765 let mut ps = [
11766 &p0 as *const _ as *mut std::ffi::c_void,
11767 &p1 as *const _ as *mut _,
11768 &p2 as *const _ as *mut _,
11769 &paq as *const _ as *mut _,
11770 &pad as *const _ as *mut _,
11771 &py0 as *const _ as *mut _,
11772 &py1 as *const _ as *mut _,
11773 &py2 as *const _ as *mut _,
11774 &inf as *const _ as *mut _,
11775 &oo0 as *const _ as *mut _,
11776 &oo1 as *const _ as *mut _,
11777 &oo2 as *const _ as *mut _,
11778 &r0 as *const _ as *mut _,
11779 &r1 as *const _ as *mut _,
11780 &r2 as *const _ as *mut _,
11781 ];
11782 unsafe {
11783 self.launch_pdl(
11784 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11785 (grid, 1, 1),
11786 (32, rpb, 1),
11787 &mut ps,
11788 )?;
11789 }
11790 }
11791 return Ok(Some((y0, y1, y2)));
11792 }
11793 let __s_b = self.gpu.stream();
11794 let mut b = __s_b.launch_builder(&f);
11795 b.arg(b0)
11796 .arg(b1)
11797 .arg(b2)
11798 .arg(aq)
11799 .arg(ad)
11800 .arg(&mut y0)
11801 .arg(&mut y1)
11802 .arg(&mut y2)
11803 .arg(&inf)
11804 .arg(&oo0)
11805 .arg(&oo1)
11806 .arg(&oo2)
11807 .arg(&r0)
11808 .arg(&r1)
11809 .arg(&r2);
11810 unsafe {
11811 b.launch(cfg)?;
11812 }
11813 Ok(Some((y0, y1, y2)))
11814 }
11815
11816 #[allow(clippy::too_many_arguments)]
11819 pub fn matmul_q4_fused3_into(
11820 &self,
11821 w0: &crate::model::GpuTensor,
11822 w1: &crate::model::GpuTensor,
11823 w2: &crate::model::GpuTensor,
11824 aq: &CudaSlice<i8>,
11825 ad: &CudaSlice<f32>,
11826 y0: &mut CudaSlice<f32>,
11827 y1: &mut CudaSlice<f32>,
11828 y2: &mut CudaSlice<f32>,
11829 ) -> Result<bool, Box<dyn std::error::Error>> {
11830 use crate::model::GpuTensor;
11831 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11832 match w {
11833 GpuTensor::Quant {
11834 qtype, row_bytes, ..
11835 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11836 _ => None,
11837 }
11838 };
11839 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11840 return Ok(false);
11841 };
11842 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11843 return Ok(false);
11844 }
11845 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11846 match w {
11847 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11848 Some(m) => (m, true),
11849 None => (bytes, *rp),
11850 },
11851 _ => unreachable!(),
11852 }
11853 }
11854 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11855 if rp0 != rp1 || rp1 != rp2 {
11856 return Ok(false);
11857 }
11858 let rp = rp0;
11859 let rpb: u32 = 4;
11860 let mr1 = rp && Self::q40_mr1_on();
11861 let nb = |o: usize| {
11862 if mr1 {
11863 (o as u32).div_ceil(rpb)
11864 } else {
11865 (o as u32).div_ceil(2).div_ceil(rpb)
11866 }
11867 };
11868 let grid = nb(o0) + nb(o1) + nb(o2);
11869 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
11870 let f = self.func(if mr1 {
11871 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11872 } else if rp {
11873 "qmatvec_q4_0_mmvq_fused3_rp"
11874 } else {
11875 "qmatvec_q4_0_mmvq_fused3"
11876 });
11877 let cfg = LaunchConfig {
11878 grid_dim: (grid, 1, 1),
11879 block_dim: (32, rpb, 1),
11880 shared_mem_bytes: 0,
11881 };
11882 let inf = w0.in_features() as i32;
11883 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11884 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11885 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11887 use cudarc::driver::{DevicePtr, DevicePtrMut};
11888 let s = &self.gpu.stream();
11889 let (p0, _g0) = b0.device_ptr(s);
11890 let (p1, _g1) = b1.device_ptr(s);
11891 let (p2, _g2) = b2.device_ptr(s);
11892 let (paq, _g3) = aq.device_ptr(s);
11893 let (pad, _g4) = ad.device_ptr(s);
11894 let (py0, _g5) = y0.device_ptr_mut(s);
11895 let (py1, _g6) = y1.device_ptr_mut(s);
11896 let (py2, _g7) = y2.device_ptr_mut(s);
11897 let mut ps = [
11898 &p0 as *const _ as *mut std::ffi::c_void,
11899 &p1 as *const _ as *mut _,
11900 &p2 as *const _ as *mut _,
11901 &paq as *const _ as *mut _,
11902 &pad as *const _ as *mut _,
11903 &py0 as *const _ as *mut _,
11904 &py1 as *const _ as *mut _,
11905 &py2 as *const _ as *mut _,
11906 &inf as *const _ as *mut _,
11907 &oo0 as *const _ as *mut _,
11908 &oo1 as *const _ as *mut _,
11909 &oo2 as *const _ as *mut _,
11910 &r0 as *const _ as *mut _,
11911 &r1 as *const _ as *mut _,
11912 &r2 as *const _ as *mut _,
11913 ];
11914 unsafe {
11915 self.launch_pdl(
11916 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11917 (grid, 1, 1),
11918 (32, rpb, 1),
11919 &mut ps,
11920 )?;
11921 }
11922 return Ok(true);
11923 }
11924 let __s_b = self.gpu.stream();
11925 let mut b = __s_b.launch_builder(&f);
11926 b.arg(b0)
11927 .arg(b1)
11928 .arg(b2)
11929 .arg(aq)
11930 .arg(ad)
11931 .arg(&mut *y0)
11932 .arg(&mut *y1)
11933 .arg(&mut *y2)
11934 .arg(&inf)
11935 .arg(&oo0)
11936 .arg(&oo1)
11937 .arg(&oo2)
11938 .arg(&r0)
11939 .arg(&r1)
11940 .arg(&r2);
11941 unsafe {
11942 b.launch(cfg)?;
11943 }
11944 Ok(true)
11945 }
11946
11947 pub fn matmul_q4_fused2(
11949 &self,
11950 w0: &crate::model::GpuTensor,
11951 w1: &crate::model::GpuTensor,
11952 aq: &CudaSlice<i8>,
11953 ad: &CudaSlice<f32>,
11954 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11955 use crate::model::GpuTensor;
11956 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11957 match w {
11958 GpuTensor::Quant {
11959 qtype, row_bytes, ..
11960 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11961 _ => None,
11962 }
11963 };
11964 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
11965 return Ok(None);
11966 };
11967 if w0.in_features() != w1.in_features() {
11968 return Ok(None);
11969 }
11970 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11972 match w {
11973 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11974 Some(m) => (m, true),
11975 None => (bytes, *rp),
11976 },
11977 _ => unreachable!(),
11978 }
11979 }
11980 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
11981 if rp0 != rp1 {
11982 return Ok(None);
11983 }
11984 let rp = rp0;
11985 let rpb: u32 = 4;
11986 let mr1 = rp && Self::q40_mr1_on();
11988 let nb = |o: usize| {
11989 if mr1 {
11990 (o as u32).div_ceil(rpb)
11991 } else {
11992 (o as u32).div_ceil(2).div_ceil(rpb)
11993 }
11994 };
11995 let grid = nb(o0) + nb(o1);
11996 let mut y0 = self.alloc_uninit::<f32>(o0)?;
11997 let mut y1 = self.alloc_uninit::<f32>(o1)?;
11998 let f = self.func(if mr1 {
11999 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12000 } else if rp {
12001 "qmatvec_q4_0_mmvq_fused2_rp"
12002 } else {
12003 "qmatvec_q4_0_mmvq_fused2"
12004 });
12005 let cfg = LaunchConfig {
12006 grid_dim: (grid, 1, 1),
12007 block_dim: (32, rpb, 1),
12008 shared_mem_bytes: 0,
12009 };
12010 let inf = w0.in_features() as i32;
12011 let (oo0, oo1) = (o0 as i32, o1 as i32);
12012 let (r0, r1) = (rb0 as i64, rb1 as i64);
12013 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12015 {
12016 use cudarc::driver::{DevicePtr, DevicePtrMut};
12017 let s = &self.gpu.stream();
12018 let (p0, _g0) = b0.device_ptr(s);
12019 let (p1, _g1) = b1.device_ptr(s);
12020 let (paq, _g2) = aq.device_ptr(s);
12021 let (pad, _g3) = ad.device_ptr(s);
12022 let (py0, _g4) = y0.device_ptr_mut(s);
12023 let (py1, _g5) = y1.device_ptr_mut(s);
12024 let mut ps = [
12025 &p0 as *const _ as *mut std::ffi::c_void,
12026 &p1 as *const _ as *mut _,
12027 &paq as *const _ as *mut _,
12028 &pad as *const _ as *mut _,
12029 &py0 as *const _ as *mut _,
12030 &py1 as *const _ as *mut _,
12031 &inf as *const _ as *mut _,
12032 &oo0 as *const _ as *mut _,
12033 &oo1 as *const _ as *mut _,
12034 &r0 as *const _ as *mut _,
12035 &r1 as *const _ as *mut _,
12036 ];
12037 unsafe {
12038 self.launch_pdl(
12039 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12040 (grid, 1, 1),
12041 (32, rpb, 1),
12042 &mut ps,
12043 )?;
12044 }
12045 }
12046 return Ok(Some((y0, y1)));
12047 }
12048 let __s_b = self.gpu.stream();
12049 let mut b = __s_b.launch_builder(&f);
12050 b.arg(b0)
12051 .arg(b1)
12052 .arg(aq)
12053 .arg(ad)
12054 .arg(&mut y0)
12055 .arg(&mut y1)
12056 .arg(&inf)
12057 .arg(&oo0)
12058 .arg(&oo1)
12059 .arg(&r0)
12060 .arg(&r1);
12061 unsafe {
12062 b.launch(cfg)?;
12063 }
12064 Ok(Some((y0, y1)))
12065 }
12066
12067 pub fn matmul_q4_fused2_into(
12069 &self,
12070 w0: &crate::model::GpuTensor,
12071 w1: &crate::model::GpuTensor,
12072 aq: &CudaSlice<i8>,
12073 ad: &CudaSlice<f32>,
12074 y0: &mut CudaSlice<f32>,
12075 y1: &mut CudaSlice<f32>,
12076 ) -> Result<bool, Box<dyn std::error::Error>> {
12077 use crate::model::GpuTensor;
12078 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12079 match w {
12080 GpuTensor::Quant {
12081 qtype, row_bytes, ..
12082 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12083 _ => None,
12084 }
12085 };
12086 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12087 return Ok(false);
12088 };
12089 if w0.in_features() != w1.in_features() {
12090 return Ok(false);
12091 }
12092 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12093 match w {
12094 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12095 Some(m) => (m, true),
12096 None => (bytes, *rp),
12097 },
12098 _ => unreachable!(),
12099 }
12100 }
12101 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12102 if rp0 != rp1 {
12103 return Ok(false);
12104 }
12105 let rp = rp0;
12106 let rpb: u32 = 4;
12107 let mr1 = rp && Self::q40_mr1_on();
12108 let nb = |o: usize| {
12109 if mr1 {
12110 (o as u32).div_ceil(rpb)
12111 } else {
12112 (o as u32).div_ceil(2).div_ceil(rpb)
12113 }
12114 };
12115 let grid = nb(o0) + nb(o1);
12116 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
12117 let f = self.func(if mr1 {
12118 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12119 } else if rp {
12120 "qmatvec_q4_0_mmvq_fused2_rp"
12121 } else {
12122 "qmatvec_q4_0_mmvq_fused2"
12123 });
12124 let cfg = LaunchConfig {
12125 grid_dim: (grid, 1, 1),
12126 block_dim: (32, rpb, 1),
12127 shared_mem_bytes: 0,
12128 };
12129 let inf = w0.in_features() as i32;
12130 let (oo0, oo1) = (o0 as i32, o1 as i32);
12131 let (r0, r1) = (rb0 as i64, rb1 as i64);
12132 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12134 use cudarc::driver::{DevicePtr, DevicePtrMut};
12135 let s = &self.gpu.stream();
12136 let (p0, _g0) = b0.device_ptr(s);
12137 let (p1, _g1) = b1.device_ptr(s);
12138 let (paq, _g2) = aq.device_ptr(s);
12139 let (pad, _g3) = ad.device_ptr(s);
12140 let (py0, _g4) = y0.device_ptr_mut(s);
12141 let (py1, _g5) = y1.device_ptr_mut(s);
12142 let mut ps = [
12143 &p0 as *const _ as *mut std::ffi::c_void,
12144 &p1 as *const _ as *mut _,
12145 &paq as *const _ as *mut _,
12146 &pad as *const _ as *mut _,
12147 &py0 as *const _ as *mut _,
12148 &py1 as *const _ as *mut _,
12149 &inf as *const _ as *mut _,
12150 &oo0 as *const _ as *mut _,
12151 &oo1 as *const _ as *mut _,
12152 &r0 as *const _ as *mut _,
12153 &r1 as *const _ as *mut _,
12154 ];
12155 unsafe {
12156 self.launch_pdl(
12157 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12158 (grid, 1, 1),
12159 (32, rpb, 1),
12160 &mut ps,
12161 )?;
12162 }
12163 return Ok(true);
12164 }
12165 let __s_b = self.gpu.stream();
12166 let mut b = __s_b.launch_builder(&f);
12167 b.arg(b0)
12168 .arg(b1)
12169 .arg(aq)
12170 .arg(ad)
12171 .arg(&mut *y0)
12172 .arg(&mut *y1)
12173 .arg(&inf)
12174 .arg(&oo0)
12175 .arg(&oo1)
12176 .arg(&r0)
12177 .arg(&r1);
12178 unsafe {
12179 b.launch(cfg)?;
12180 }
12181 Ok(true)
12182 }
12183
12184 pub fn matmul_q4_fused2_batched(
12189 &self,
12190 w0: &crate::model::GpuTensor,
12191 w1: &crate::model::GpuTensor,
12192 aq: &CudaSlice<i8>,
12193 ad: &CudaSlice<f32>,
12194 m: usize,
12195 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12196 use crate::model::GpuTensor;
12197 if m < 2 || m > 8 {
12198 return Ok(None);
12199 }
12200 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12201 match w {
12202 GpuTensor::Quant {
12203 qtype, row_bytes, ..
12204 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12205 _ => None,
12206 }
12207 };
12208 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
12209 return Ok(None);
12210 };
12211 if w0.in_features() != w1.in_features() {
12212 return Ok(None);
12213 }
12214 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12215 match w {
12216 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12217 Some(mr) => (mr, true),
12218 None => (bytes, *rp),
12219 },
12220 _ => unreachable!(),
12221 }
12222 }
12223 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12224 if !rp0 || !rp1 {
12225 return Ok(None);
12226 }
12227 let mcols = Self::batched_mcols(m);
12228 let rpb: u32 = 4;
12229 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12230 let grid = nb(o0) + nb(o1);
12231 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12232 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12233 let f = self.func(match mcols {
12234 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12235 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12236 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12237 });
12238 let cfg = LaunchConfig {
12239 grid_dim: (grid, 1, 1),
12240 block_dim: (32, rpb, 1),
12241 shared_mem_bytes: 0,
12242 };
12243 let inf = w0.in_features() as i32;
12244 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12245 let rb = rb0 as i64;
12246 let __s_b = self.gpu.stream();
12247 let mut b = __s_b.launch_builder(&f);
12248 b.arg(b0)
12249 .arg(b1)
12250 .arg(aq)
12251 .arg(ad)
12252 .arg(&mut y0)
12253 .arg(&mut y1)
12254 .arg(&inf)
12255 .arg(&oo0)
12256 .arg(&oo1)
12257 .arg(&mi)
12258 .arg(&rb);
12259 unsafe {
12260 b.launch(cfg)?;
12261 }
12262 Ok(Some((y0, y1)))
12263 }
12264
12265 #[allow(clippy::too_many_arguments)]
12268 pub fn matmul_q4_fused3_batched(
12269 &self,
12270 w0: &crate::model::GpuTensor,
12271 w1: &crate::model::GpuTensor,
12272 w2: &crate::model::GpuTensor,
12273 aq: &CudaSlice<i8>,
12274 ad: &CudaSlice<f32>,
12275 m: usize,
12276 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12277 {
12278 use crate::model::GpuTensor;
12279 if m < 2 || m > 8 {
12280 return Ok(None);
12281 }
12282 let q4 = |w: &GpuTensor| -> Option<usize> {
12283 match w {
12284 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12285 _ => None,
12286 }
12287 };
12288 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12289 return Ok(None);
12290 };
12291 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12292 return Ok(None);
12293 }
12294 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12295 match w {
12296 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12297 Some(mr) => (mr, true),
12298 None => (bytes, *rp),
12299 },
12300 _ => unreachable!(),
12301 }
12302 }
12303 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12304 if !rp0 || !rp1 || !rp2 {
12305 return Ok(None);
12306 }
12307 let mcols = Self::batched_mcols(m);
12308 let rpb: u32 = 4;
12309 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12310 let grid = nb(o0) + nb(o1) + nb(o2);
12311 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12312 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12313 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12314 let f = self.func(match mcols {
12315 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12316 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12317 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12318 });
12319 let cfg = LaunchConfig {
12320 grid_dim: (grid, 1, 1),
12321 block_dim: (32, rpb, 1),
12322 shared_mem_bytes: 0,
12323 };
12324 let inf = w0.in_features() as i32;
12325 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12326 let rb = 0i64;
12327 let __s_b = self.gpu.stream();
12328 let mut b = __s_b.launch_builder(&f);
12329 b.arg(b0)
12330 .arg(b1)
12331 .arg(b2)
12332 .arg(aq)
12333 .arg(ad)
12334 .arg(&mut y0)
12335 .arg(&mut y1)
12336 .arg(&mut y2)
12337 .arg(&inf)
12338 .arg(&oo0)
12339 .arg(&oo1)
12340 .arg(&oo2)
12341 .arg(&mi)
12342 .arg(&rb);
12343 unsafe {
12344 b.launch(cfg)?;
12345 }
12346 Ok(Some((y0, y1, y2)))
12347 }
12348
12349 pub fn matmul_q8_fused3(
12350 &self,
12351 w0: &crate::model::GpuTensor,
12352 w1: &crate::model::GpuTensor,
12353 w2: &crate::model::GpuTensor,
12354 aq: &CudaSlice<i8>,
12355 ad: &CudaSlice<f32>,
12356 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12357 {
12358 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12361 return Ok(Some(self.e4m3_fused3_core(
12362 p0.0,
12363 p1.0,
12364 p2.0,
12365 aq,
12366 ad,
12367 w0.in_features(),
12368 p0.1,
12369 p1.1,
12370 p2.1,
12371 p0.2,
12372 p0.3,
12373 p1.3,
12374 p2.3,
12375 )?));
12376 }
12377 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12378 return Ok(None);
12379 };
12380 Ok(Some(self.q8_fused3_core(
12381 p0.0,
12382 p1.0,
12383 p2.0,
12384 aq,
12385 ad,
12386 w0.in_features(),
12387 p0.1,
12388 p1.1,
12389 p2.1,
12390 p0.2,
12391 )?))
12392 }
12393
12394 #[allow(clippy::too_many_arguments)]
12395 fn q8_fused3_core(
12396 &self,
12397 b0: &CudaSlice<u8>,
12398 b1: &CudaSlice<u8>,
12399 b2: &CudaSlice<u8>,
12400 aq: &CudaSlice<i8>,
12401 ad: &CudaSlice<f32>,
12402 in_f: usize,
12403 out0: usize,
12404 out1: usize,
12405 out2: usize,
12406 row_bytes: usize,
12407 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12408 const ROWS_PER_BLOCK: u32 = 4;
12409 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12410 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12411 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12412 let f = self.func("qmatvec_q8_0_mmvq_fused3");
12413 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12414 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12415 let mut y2 = self.alloc_uninit::<f32>(out2)?;
12416 let cfg = LaunchConfig {
12417 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12418 block_dim: (32, ROWS_PER_BLOCK, 1),
12419 shared_mem_bytes: 0,
12420 };
12421 let (inf, o0, o1, o2, rbl) = (
12422 in_f as i32,
12423 out0 as i32,
12424 out1 as i32,
12425 out2 as i32,
12426 row_bytes as i64,
12427 );
12428 let __s_b = self.gpu.stream();
12429 let mut b = __s_b.launch_builder(&f);
12430 b.arg(b0)
12431 .arg(b1)
12432 .arg(b2)
12433 .arg(aq)
12434 .arg(ad)
12435 .arg(&mut y0)
12436 .arg(&mut y1)
12437 .arg(&mut y2)
12438 .arg(&inf)
12439 .arg(&o0)
12440 .arg(&o1)
12441 .arg(&o2)
12442 .arg(&rbl);
12443 unsafe {
12444 b.launch(cfg)?;
12445 }
12446 Ok((y0, y1, y2))
12447 }
12448
12449 #[allow(clippy::too_many_arguments)]
12451 pub fn qmatvec_q8_fused3_raw(
12452 &self,
12453 b0: &CudaSlice<u8>,
12454 b1: &CudaSlice<u8>,
12455 b2: &CudaSlice<u8>,
12456 x: &CudaSlice<f32>,
12457 in_f: usize,
12458 out0: usize,
12459 out1: usize,
12460 out2: usize,
12461 row_bytes: usize,
12462 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12463 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12464 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12465 }
12466
12467 pub fn matmul_q8_fused2_t(
12478 &self,
12479 w0: &crate::model::GpuTensor,
12480 w1: &crate::model::GpuTensor,
12481 aq: &CudaSlice<i8>,
12482 ad: &CudaSlice<f32>,
12483 m: usize,
12484 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12485 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12489 return Ok(None);
12490 }
12491 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12494 if m > 4 && !Self::b8_enabled() {
12495 return Ok(None);
12496 }
12497 return Ok(Some(self.e4m3_fused2_t_core(
12498 p0.0,
12499 p1.0,
12500 aq,
12501 ad,
12502 m,
12503 w0.in_features(),
12504 p0.1,
12505 p1.1,
12506 p0.2,
12507 p0.3,
12508 p1.3,
12509 )?));
12510 }
12511 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12512 return Ok(None);
12513 };
12514 Ok(Some(self.q8_fused2_t_core(
12515 p0.0,
12516 p1.0,
12517 aq,
12518 ad,
12519 m,
12520 w0.in_features(),
12521 p0.1,
12522 p1.1,
12523 p0.2,
12524 )?))
12525 }
12526
12527 #[allow(clippy::too_many_arguments)]
12528 fn q8_fused2_t_core(
12529 &self,
12530 b0: &CudaSlice<u8>,
12531 b1: &CudaSlice<u8>,
12532 aq: &CudaSlice<i8>,
12533 ad: &CudaSlice<f32>,
12534 m: usize,
12535 in_f: usize,
12536 out0: usize,
12537 out1: usize,
12538 row_bytes: usize,
12539 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12540 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12542 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12543 let f = self.func(match Self::batched_mcols(m) {
12544 2 => "qmatvec_q8_0_mmvq_fused2_b2",
12545 4 => "qmatvec_q8_0_mmvq_fused2_b4",
12546 _ => "qmatvec_q8_0_mmvq_fused2_b8",
12548 });
12549 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12550 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12551 let cfg = LaunchConfig {
12552 grid_dim: (nb0 + nb1, 1, 1),
12553 block_dim: (32, ROWS_PER_BLOCK, 1),
12554 shared_mem_bytes: 0,
12555 };
12556 let (inf, o0, o1, mi, rbl) = (
12557 in_f as i32,
12558 out0 as i32,
12559 out1 as i32,
12560 m as i32,
12561 row_bytes as i64,
12562 );
12563 let __s_b = self.gpu.stream();
12564 let mut b = __s_b.launch_builder(&f);
12565 b.arg(b0)
12566 .arg(b1)
12567 .arg(aq)
12568 .arg(ad)
12569 .arg(&mut y0)
12570 .arg(&mut y1)
12571 .arg(&inf)
12572 .arg(&o0)
12573 .arg(&o1)
12574 .arg(&mi)
12575 .arg(&rbl);
12576 unsafe {
12577 b.launch(cfg)?;
12578 }
12579 Ok((y0, y1))
12580 }
12581
12582 #[allow(clippy::too_many_arguments)]
12585 pub fn qmatvec_q8_fused2_t_raw(
12586 &self,
12587 b0: &CudaSlice<u8>,
12588 b1: &CudaSlice<u8>,
12589 x: &CudaSlice<f32>,
12590 m: usize,
12591 in_f: usize,
12592 out0: usize,
12593 out1: usize,
12594 row_bytes: usize,
12595 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12596 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12597 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
12598 }
12599
12600 #[allow(clippy::too_many_arguments)]
12603 pub fn matmul_q8_fused3_t(
12604 &self,
12605 w0: &crate::model::GpuTensor,
12606 w1: &crate::model::GpuTensor,
12607 w2: &crate::model::GpuTensor,
12608 aq: &CudaSlice<i8>,
12609 ad: &CudaSlice<f32>,
12610 m: usize,
12611 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12612 {
12613 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12614 return Ok(None);
12615 }
12616 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12617 return Ok(Some(self.e4m3_fused3_t_core(
12618 p0.0,
12619 p1.0,
12620 p2.0,
12621 aq,
12622 ad,
12623 m,
12624 w0.in_features(),
12625 p0.1,
12626 p1.1,
12627 p2.1,
12628 p0.2,
12629 p0.3,
12630 p1.3,
12631 p2.3,
12632 )?));
12633 }
12634 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12635 return Ok(None);
12636 };
12637 Ok(Some(self.q8_fused3_t_core(
12638 p0.0,
12639 p1.0,
12640 p2.0,
12641 aq,
12642 ad,
12643 m,
12644 w0.in_features(),
12645 p0.1,
12646 p1.1,
12647 p2.1,
12648 p0.2,
12649 )?))
12650 }
12651
12652 #[allow(clippy::too_many_arguments)]
12653 fn q8_fused3_t_core(
12654 &self,
12655 b0: &CudaSlice<u8>,
12656 b1: &CudaSlice<u8>,
12657 b2: &CudaSlice<u8>,
12658 aq: &CudaSlice<i8>,
12659 ad: &CudaSlice<f32>,
12660 m: usize,
12661 in_f: usize,
12662 out0: usize,
12663 out1: usize,
12664 out2: usize,
12665 row_bytes: usize,
12666 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12667 const ROWS_PER_BLOCK: u32 = 4;
12668 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12669 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12670 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12671 let f = self.func(if Self::batched_mcols(m) == 2 {
12672 "qmatvec_q8_0_mmvq_fused3_b2"
12673 } else {
12674 "qmatvec_q8_0_mmvq_fused3_b4"
12675 });
12676 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12677 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12678 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
12679 let cfg = LaunchConfig {
12680 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12681 block_dim: (32, ROWS_PER_BLOCK, 1),
12682 shared_mem_bytes: 0,
12683 };
12684 let (inf, o0, o1, o2, mi, rbl) = (
12685 in_f as i32,
12686 out0 as i32,
12687 out1 as i32,
12688 out2 as i32,
12689 m as i32,
12690 row_bytes as i64,
12691 );
12692 let __s_b = self.gpu.stream();
12693 let mut b = __s_b.launch_builder(&f);
12694 b.arg(b0)
12695 .arg(b1)
12696 .arg(b2)
12697 .arg(aq)
12698 .arg(ad)
12699 .arg(&mut y0)
12700 .arg(&mut y1)
12701 .arg(&mut y2)
12702 .arg(&inf)
12703 .arg(&o0)
12704 .arg(&o1)
12705 .arg(&o2)
12706 .arg(&mi)
12707 .arg(&rbl);
12708 unsafe {
12709 b.launch(cfg)?;
12710 }
12711 Ok((y0, y1, y2))
12712 }
12713
12714 #[allow(clippy::too_many_arguments)]
12716 pub fn qmatvec_q8_fused3_t_raw(
12717 &self,
12718 b0: &CudaSlice<u8>,
12719 b1: &CudaSlice<u8>,
12720 b2: &CudaSlice<u8>,
12721 x: &CudaSlice<f32>,
12722 m: usize,
12723 in_f: usize,
12724 out0: usize,
12725 out1: usize,
12726 out2: usize,
12727 row_bytes: usize,
12728 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12729 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12730 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
12731 }
12732
12733 pub fn q8_ffn_fuse2_on(&self) -> bool {
12737 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12738 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
12739 }
12740
12741 #[allow(clippy::type_complexity)]
12747 fn q8_fused_params<'w, const N: usize>(
12748 &self,
12749 ws: &[&'w crate::model::GpuTensor; N],
12750 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
12751 use crate::model::GpuTensor;
12752 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
12753 return None;
12754 }
12755 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
12756 return None;
12757 }
12758 let in_f = ws[0].in_features();
12759 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
12760 for (i, w) in ws.iter().enumerate() {
12761 match w {
12762 GpuTensor::Quant {
12763 bytes,
12764 qtype,
12765 row_bytes,
12766 scale,
12767 ..
12768 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
12769 out[i] = Some((bytes, w.out_features(), *row_bytes))
12770 }
12771 _ => return None,
12772 }
12773 }
12774 Some(out.map(|o| o.unwrap()))
12775 }
12776
12777 pub fn e4m3_dual_on(&self) -> bool {
12780 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12781 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
12782 }
12783
12784 #[allow(clippy::type_complexity)]
12796 fn e4m3_fused_params<'w, const N: usize>(
12797 &self,
12798 ws: &[&'w crate::model::GpuTensor; N],
12799 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
12800 use crate::model::GpuTensor;
12801 if !self.e4m3_dual_on() {
12802 return None;
12803 }
12804 let in_f = ws[0].in_features();
12805 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
12806 for (i, w) in ws.iter().enumerate() {
12807 match w {
12808 GpuTensor::Quant {
12809 bytes,
12810 qtype,
12811 row_bytes,
12812 scale,
12813 rp,
12814 rp4,
12815 ..
12816 } if *qtype == QT_F8_E4M3
12817 && w.in_features() == in_f
12818 && *row_bytes == in_f
12819 && !*rp
12820 && rp4.is_none() =>
12821 {
12822 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
12823 }
12824 _ => return None,
12825 }
12826 }
12827 Some(out.map(|o| o.unwrap()))
12828 }
12829
12830 #[allow(clippy::too_many_arguments)]
12834 fn e4m3_fused2_core(
12835 &self,
12836 b0: &CudaSlice<u8>,
12837 b1: &CudaSlice<u8>,
12838 aq: &CudaSlice<i8>,
12839 ad: &CudaSlice<f32>,
12840 in_f: usize,
12841 out0: usize,
12842 out1: usize,
12843 row_bytes: usize,
12844 ws0: f32,
12845 ws1: f32,
12846 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12847 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12849 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12850 let f = self.func("qmatvec_e4m3_mmvq_fused2");
12851 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12852 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12853 let cfg = LaunchConfig {
12854 grid_dim: (nb0 + nb1, 1, 1),
12855 block_dim: (32, ROWS_PER_BLOCK, 1),
12856 shared_mem_bytes: 0,
12857 };
12858 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12859 let __s_b = self.gpu.stream();
12860 let mut b = __s_b.launch_builder(&f);
12861 b.arg(b0)
12862 .arg(b1)
12863 .arg(aq)
12864 .arg(ad)
12865 .arg(&mut y0)
12866 .arg(&mut y1)
12867 .arg(&inf)
12868 .arg(&o0)
12869 .arg(&o1)
12870 .arg(&rbl)
12871 .arg(&ws0)
12872 .arg(&ws1);
12873 unsafe {
12874 b.launch(cfg)?;
12875 }
12876 Ok((y0, y1))
12877 }
12878
12879 #[allow(clippy::too_many_arguments)]
12881 fn e4m3_fused3_core(
12882 &self,
12883 b0: &CudaSlice<u8>,
12884 b1: &CudaSlice<u8>,
12885 b2: &CudaSlice<u8>,
12886 aq: &CudaSlice<i8>,
12887 ad: &CudaSlice<f32>,
12888 in_f: usize,
12889 out0: usize,
12890 out1: usize,
12891 out2: usize,
12892 row_bytes: usize,
12893 ws0: f32,
12894 ws1: f32,
12895 ws2: f32,
12896 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12897 const ROWS_PER_BLOCK: u32 = 4;
12898 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12899 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12900 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12901 let f = self.func("qmatvec_e4m3_mmvq_fused3");
12902 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12903 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12904 let mut y2 = self.alloc_uninit::<f32>(out2)?;
12905 let cfg = LaunchConfig {
12906 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12907 block_dim: (32, ROWS_PER_BLOCK, 1),
12908 shared_mem_bytes: 0,
12909 };
12910 let (inf, o0, o1, o2, rbl) = (
12911 in_f as i32,
12912 out0 as i32,
12913 out1 as i32,
12914 out2 as i32,
12915 row_bytes as i64,
12916 );
12917 let __s_b = self.gpu.stream();
12918 let mut b = __s_b.launch_builder(&f);
12919 b.arg(b0)
12920 .arg(b1)
12921 .arg(b2)
12922 .arg(aq)
12923 .arg(ad)
12924 .arg(&mut y0)
12925 .arg(&mut y1)
12926 .arg(&mut y2)
12927 .arg(&inf)
12928 .arg(&o0)
12929 .arg(&o1)
12930 .arg(&o2)
12931 .arg(&rbl)
12932 .arg(&ws0)
12933 .arg(&ws1)
12934 .arg(&ws2);
12935 unsafe {
12936 b.launch(cfg)?;
12937 }
12938 Ok((y0, y1, y2))
12939 }
12940
12941 #[allow(clippy::too_many_arguments)]
12945 fn e4m3_fused2_t_core(
12946 &self,
12947 b0: &CudaSlice<u8>,
12948 b1: &CudaSlice<u8>,
12949 aq: &CudaSlice<i8>,
12950 ad: &CudaSlice<f32>,
12951 m: usize,
12952 in_f: usize,
12953 out0: usize,
12954 out1: usize,
12955 row_bytes: usize,
12956 ws0: f32,
12957 ws1: f32,
12958 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12959 const ROWS_PER_BLOCK: u32 = 4;
12960 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12961 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12962 let f = self.func(match Self::batched_mcols(m) {
12963 2 => "qmatvec_e4m3_mmvq_fused2_b2",
12964 4 => "qmatvec_e4m3_mmvq_fused2_b4",
12965 _ => "qmatvec_e4m3_mmvq_fused2_b8",
12966 });
12967 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12968 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12969 let cfg = LaunchConfig {
12970 grid_dim: (nb0 + nb1, 1, 1),
12971 block_dim: (32, ROWS_PER_BLOCK, 1),
12972 shared_mem_bytes: 0,
12973 };
12974 let (inf, o0, o1, mi, rbl) = (
12975 in_f as i32,
12976 out0 as i32,
12977 out1 as i32,
12978 m as i32,
12979 row_bytes as i64,
12980 );
12981 let __s_b = self.gpu.stream();
12982 let mut b = __s_b.launch_builder(&f);
12983 b.arg(b0)
12984 .arg(b1)
12985 .arg(aq)
12986 .arg(ad)
12987 .arg(&mut y0)
12988 .arg(&mut y1)
12989 .arg(&inf)
12990 .arg(&o0)
12991 .arg(&o1)
12992 .arg(&mi)
12993 .arg(&rbl);
12994 unsafe {
12995 b.launch(cfg)?;
12996 }
12997 if ws0 != 1.0 {
12998 self.scale_inplace(&mut y0, ws0, m * out0)?;
12999 }
13000 if ws1 != 1.0 {
13001 self.scale_inplace(&mut y1, ws1, m * out1)?;
13002 }
13003 Ok((y0, y1))
13004 }
13005
13006 #[allow(clippy::too_many_arguments)]
13008 fn e4m3_fused3_t_core(
13009 &self,
13010 b0: &CudaSlice<u8>,
13011 b1: &CudaSlice<u8>,
13012 b2: &CudaSlice<u8>,
13013 aq: &CudaSlice<i8>,
13014 ad: &CudaSlice<f32>,
13015 m: usize,
13016 in_f: usize,
13017 out0: usize,
13018 out1: usize,
13019 out2: usize,
13020 row_bytes: usize,
13021 ws0: f32,
13022 ws1: f32,
13023 ws2: f32,
13024 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13025 const ROWS_PER_BLOCK: u32 = 4;
13026 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13027 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13028 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13029 let f = self.func(if Self::batched_mcols(m) == 2 {
13030 "qmatvec_e4m3_mmvq_fused3_b2"
13031 } else {
13032 "qmatvec_e4m3_mmvq_fused3_b4"
13033 });
13034 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13035 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13036 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13037 let cfg = LaunchConfig {
13038 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13039 block_dim: (32, ROWS_PER_BLOCK, 1),
13040 shared_mem_bytes: 0,
13041 };
13042 let (inf, o0, o1, o2, mi, rbl) = (
13043 in_f as i32,
13044 out0 as i32,
13045 out1 as i32,
13046 out2 as i32,
13047 m as i32,
13048 row_bytes as i64,
13049 );
13050 let __s_b = self.gpu.stream();
13051 let mut b = __s_b.launch_builder(&f);
13052 b.arg(b0)
13053 .arg(b1)
13054 .arg(b2)
13055 .arg(aq)
13056 .arg(ad)
13057 .arg(&mut y0)
13058 .arg(&mut y1)
13059 .arg(&mut y2)
13060 .arg(&inf)
13061 .arg(&o0)
13062 .arg(&o1)
13063 .arg(&o2)
13064 .arg(&mi)
13065 .arg(&rbl);
13066 unsafe {
13067 b.launch(cfg)?;
13068 }
13069 if ws0 != 1.0 {
13070 self.scale_inplace(&mut y0, ws0, m * out0)?;
13071 }
13072 if ws1 != 1.0 {
13073 self.scale_inplace(&mut y1, ws1, m * out1)?;
13074 }
13075 if ws2 != 1.0 {
13076 self.scale_inplace(&mut y2, ws2, m * out2)?;
13077 }
13078 Ok((y0, y1, y2))
13079 }
13080
13081 pub fn qmatvec_e4m3_blk_mmvq(
13091 &self,
13092 bytes: &CudaSlice<u8>,
13093 aq: &CudaSlice<i8>,
13094 ad: &CudaSlice<f32>,
13095 scales: &CudaSlice<f32>,
13096 m: usize,
13097 in_f: usize,
13098 out_f: usize,
13099 row_bytes: usize,
13100 scale_cols: usize,
13101 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13102 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
13104 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
13105 )?;
13106 Ok(y)
13107 }
13108
13109 #[allow(clippy::too_many_arguments)]
13111 pub fn qmatvec_e4m3_blk_mmvq_into(
13112 &self,
13113 bytes: &CudaSlice<u8>,
13114 aq: &CudaSlice<i8>,
13115 ad: &CudaSlice<f32>,
13116 scales: &CudaSlice<f32>,
13117 m: usize,
13118 in_f: usize,
13119 out_f: usize,
13120 row_bytes: usize,
13121 scale_cols: usize,
13122 y: &mut CudaSlice<f32>,
13123 ) -> Result<(), Box<dyn std::error::Error>> {
13124 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
13126 let cfg = LaunchConfig {
13127 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
13128 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
13131 let (inf, outf, mi, rb, sc) = (
13132 in_f as i32,
13133 out_f as i32,
13134 m as i32,
13135 row_bytes as i64,
13136 scale_cols as i32,
13137 );
13138 let __s_b = self.gpu.stream();
13139 let mut b = __s_b.launch_builder(&f);
13140 b.arg(bytes)
13141 .arg(aq)
13142 .arg(ad)
13143 .arg(scales)
13144 .arg(&mut *y)
13145 .arg(&inf)
13146 .arg(&outf)
13147 .arg(&mi)
13148 .arg(&rb)
13149 .arg(&sc);
13150 unsafe {
13151 b.launch(cfg)?;
13152 }
13153 Ok(())
13154 }
13155
13156 #[allow(clippy::too_many_arguments)]
13162 pub fn qmatvec_e4m3_blk_mmvq_batched(
13163 &self,
13164 bytes: &CudaSlice<u8>,
13165 aq: &CudaSlice<i8>,
13166 ad: &CudaSlice<f32>,
13167 scales: &CudaSlice<f32>,
13168 m: usize,
13169 in_f: usize,
13170 out_f: usize,
13171 row_bytes: usize,
13172 scale_cols: usize,
13173 mcols: usize,
13174 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13175 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
13177 let name = match mcols {
13178 2 => "qmatvec_e4m3_blk_mmvq_b2",
13179 4 => "qmatvec_e4m3_blk_mmvq_b4",
13180 8 => "qmatvec_e4m3_blk_mmvq_b8",
13181 16 => "qmatvec_e4m3_blk_mmvq_b16",
13182 _ => {
13183 return Err(
13184 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
13185 );
13186 }
13187 };
13188 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13189 let f = self.func(name);
13190 let cfg = LaunchConfig {
13191 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
13192 block_dim: (32, ROWS_PER_BLOCK, 1),
13193 shared_mem_bytes: 0,
13194 };
13195 let (inf, outf, mi, rb, sc) = (
13196 in_f as i32,
13197 out_f as i32,
13198 m as i32,
13199 row_bytes as i64,
13200 scale_cols as i32,
13201 );
13202 let __s_b = self.gpu.stream();
13203 let mut b = __s_b.launch_builder(&f);
13204 b.arg(bytes)
13205 .arg(aq)
13206 .arg(ad)
13207 .arg(scales)
13208 .arg(&mut y)
13209 .arg(&inf)
13210 .arg(&outf)
13211 .arg(&mi)
13212 .arg(&rb)
13213 .arg(&sc);
13214 unsafe {
13215 b.launch(cfg)?;
13216 }
13217 Ok(y)
13218 }
13219
13220 #[allow(clippy::too_many_arguments)]
13223 pub fn qmatvec_e4m3_blk_batched_raw(
13224 &self,
13225 bytes: &CudaSlice<u8>,
13226 x: &CudaSlice<f32>,
13227 scales: &CudaSlice<f32>,
13228 m: usize,
13229 in_f: usize,
13230 out_f: usize,
13231 row_bytes: usize,
13232 scale_cols: usize,
13233 mcols: usize,
13234 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13235 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13236 self.qmatvec_e4m3_blk_mmvq_batched(
13237 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13238 )
13239 }
13240
13241 #[allow(clippy::too_many_arguments)]
13244 pub fn qmatvec_e4m3_blk_mmvq_raw(
13245 &self,
13246 bytes: &CudaSlice<u8>,
13247 x: &CudaSlice<f32>,
13248 scales: &CudaSlice<f32>,
13249 m: usize,
13250 in_f: usize,
13251 out_f: usize,
13252 row_bytes: usize,
13253 scale_cols: usize,
13254 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13255 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13256 self.qmatvec_e4m3_blk_mmvq(
13257 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13258 )
13259 }
13260
13261 #[allow(clippy::too_many_arguments)]
13264 pub fn qmatvec_e4m3_fused2_raw(
13265 &self,
13266 b0: &CudaSlice<u8>,
13267 b1: &CudaSlice<u8>,
13268 x: &CudaSlice<f32>,
13269 in_f: usize,
13270 out0: usize,
13271 out1: usize,
13272 row_bytes: usize,
13273 ws0: f32,
13274 ws1: f32,
13275 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13276 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13277 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13278 }
13279
13280 #[allow(clippy::too_many_arguments)]
13281 pub fn qmatvec_e4m3_fused3_raw(
13282 &self,
13283 b0: &CudaSlice<u8>,
13284 b1: &CudaSlice<u8>,
13285 b2: &CudaSlice<u8>,
13286 x: &CudaSlice<f32>,
13287 in_f: usize,
13288 out0: usize,
13289 out1: usize,
13290 out2: usize,
13291 row_bytes: usize,
13292 ws0: f32,
13293 ws1: f32,
13294 ws2: f32,
13295 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13296 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13297 self.e4m3_fused3_core(
13298 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13299 )
13300 }
13301
13302 #[allow(clippy::too_many_arguments)]
13303 pub fn qmatvec_e4m3_fused2_t_raw(
13304 &self,
13305 b0: &CudaSlice<u8>,
13306 b1: &CudaSlice<u8>,
13307 x: &CudaSlice<f32>,
13308 m: usize,
13309 in_f: usize,
13310 out0: usize,
13311 out1: usize,
13312 row_bytes: usize,
13313 ws0: f32,
13314 ws1: f32,
13315 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13316 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13317 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13318 }
13319
13320 #[allow(clippy::too_many_arguments)]
13321 pub fn qmatvec_e4m3_fused3_t_raw(
13322 &self,
13323 b0: &CudaSlice<u8>,
13324 b1: &CudaSlice<u8>,
13325 b2: &CudaSlice<u8>,
13326 x: &CudaSlice<f32>,
13327 m: usize,
13328 in_f: usize,
13329 out0: usize,
13330 out1: usize,
13331 out2: usize,
13332 row_bytes: usize,
13333 ws0: f32,
13334 ws1: f32,
13335 ws2: f32,
13336 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13337 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13338 self.e4m3_fused3_t_core(
13339 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13340 )
13341 }
13342
13343 fn try_e4m3_blk_pre(
13354 &self,
13355 w: &crate::model::GpuTensor,
13356 aq: &CudaSlice<i8>,
13357 ad: &CudaSlice<f32>,
13358 m: usize,
13359 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13360 use crate::model::GpuTensor;
13361 if let GpuTensor::Quant {
13362 bytes,
13363 qtype,
13364 row_bytes,
13365 blk: Some(g),
13366 ..
13367 } = w
13368 {
13369 if *qtype == QT_F8_E4M3_BLK {
13370 if (2..=16).contains(&m)
13376 && std::env::var("MEMRA_NO_BATCHED").is_err()
13377 && (m <= 4 || Self::b8_enabled())
13378 {
13379 let mcols = Self::batched_mcols(m);
13380 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13381 bytes,
13382 aq,
13383 ad,
13384 &g.scales,
13385 m,
13386 w.in_features(),
13387 w.out_features(),
13388 *row_bytes,
13389 g.cols,
13390 mcols,
13391 )?));
13392 }
13393 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13394 bytes,
13395 aq,
13396 ad,
13397 &g.scales,
13398 m,
13399 w.in_features(),
13400 w.out_features(),
13401 *row_bytes,
13402 g.cols,
13403 )?));
13404 }
13405 }
13406 Ok(None)
13407 }
13408
13409 fn try_e4m3_blk_prefill(
13456 &self,
13457 w: &crate::model::GpuTensor,
13458 x: &CudaSlice<f32>,
13459 m: usize,
13460 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13461 use crate::model::GpuTensor;
13462 let GpuTensor::Quant {
13463 bytes,
13464 qtype,
13465 blk: Some(g),
13466 ..
13467 } = w
13468 else {
13469 return Ok(None);
13470 };
13471 if *qtype != QT_F8_E4M3_BLK {
13472 return Ok(None);
13473 }
13474 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
13479 return Ok(Some(y));
13480 }
13481 let (in_f, out_f) = (w.in_features(), w.out_features());
13482 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
13483 let tmp = GpuTensor::Quant {
13484 bytes: slab,
13485 qtype: QT_Q8_0,
13486 row_bytes: in_f / 32 * 34,
13487 ne: vec![in_f as u64, out_f as u64],
13488 scale: 1.0,
13489 rp: false,
13490 #[cfg(memra_cutlass)]
13491 cutlass: None,
13492 fp8: None,
13493 blk: None,
13494 f16: None,
13495 rp4: None,
13496 };
13497 Ok(Some(self.matmul(&tmp, x, m)?))
13499 }
13500
13501 pub fn matmul_pre_noscale(
13502 &self,
13503 w: &crate::model::GpuTensor,
13504 aq: &CudaSlice<i8>,
13505 ad: &CudaSlice<f32>,
13506 m: usize,
13507 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
13508 use crate::model::GpuTensor;
13509 if m == 1 {
13513 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
13514 return Ok(Some((y, 1.0)));
13515 }
13516 }
13517 if m != 1 || !self.uses_q8_1_fast(w) {
13519 return Ok(None);
13520 }
13521 let in_f = w.in_features();
13522 let out_f = w.out_features();
13523 let (bytes, qtype, row_bytes, scale, rp) = match w {
13524 GpuTensor::Quant {
13525 bytes,
13526 qtype,
13527 row_bytes,
13528 scale,
13529 rp,
13530 ..
13531 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13532 _ => return Ok(None),
13533 };
13534 if self.mmvq_supports(qtype) {
13536 let (mbytes, mrp) = match w {
13538 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
13539 _ => (bytes, rp),
13540 };
13541 let y = self.qmatvec_mmvq(
13542 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
13543 )?;
13544 return Ok(Some((y, scale)));
13545 }
13546 let name = match qtype {
13548 QT_Q8_0 => "qmatvec_q8_0_dp4a",
13549 QT_Q4_K => "qmatvec_q4_K_dp4a",
13550 QT_Q6_K => "qmatvec_q6_K_dp4a",
13551 QT_Q5_K => "qmatvec_q5_K_dp4a",
13552 QT_Q3_K => "qmatvec_q3_K_dp4a",
13553 QT_NVFP4 => {
13554 if rp {
13555 "qmatvec_nvfp4_dp4a_rp"
13556 } else {
13557 "qmatvec_nvfp4_dp4a"
13558 }
13559 }
13560 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
13561 _ => return Ok(None),
13562 };
13563 let f = self.func(name);
13564 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13565 let cfg = LaunchConfig {
13566 grid_dim: (out_f as u32, m as u32, 1),
13567 block_dim: (128, 1, 1),
13568 shared_mem_bytes: 0,
13569 };
13570 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13571 let __s_b = self.gpu.stream();
13572 let mut b = __s_b.launch_builder(&f);
13573 b.arg(bytes)
13574 .arg(aq)
13575 .arg(ad)
13576 .arg(&mut y)
13577 .arg(&inf)
13578 .arg(&outf)
13579 .arg(&mi)
13580 .arg(&rb);
13581 unsafe {
13582 b.launch(cfg)?;
13583 }
13584 Ok(Some((y, scale)))
13585 }
13586
13587 pub fn mmvq_supports(&self, qtype: i32) -> bool {
13590 if qtype == QT_F8_E4M3 {
13595 return true;
13596 }
13597 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13598 return false;
13599 }
13600 matches!(
13601 qtype,
13602 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
13603 )
13604 }
13605
13606 pub fn qmatvec_mmvq(
13611 &self,
13612 bytes: &CudaSlice<u8>,
13613 aq: &CudaSlice<i8>,
13614 ad: &CudaSlice<f32>,
13615 m: usize,
13616 in_f: usize,
13617 out_f: usize,
13618 qtype: i32,
13619 row_bytes: usize,
13620 scale: f32,
13621 rp: bool,
13622 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13623 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
13625 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
13626 )?;
13627 Ok(y)
13628 }
13629
13630 #[allow(clippy::too_many_arguments)]
13632 pub fn qmatvec_mmvq_into(
13633 &self,
13634 bytes: &CudaSlice<u8>,
13635 aq: &CudaSlice<i8>,
13636 ad: &CudaSlice<f32>,
13637 m: usize,
13638 in_f: usize,
13639 out_f: usize,
13640 qtype: i32,
13641 row_bytes: usize,
13642 scale: f32,
13643 rp: bool,
13644 y: &mut CudaSlice<f32>,
13645 ) -> Result<(), Box<dyn std::error::Error>> {
13646 debug_assert!(y.len() >= m * out_f);
13647 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
13653 && rp
13654 && m == 1
13655 && out_f >= 64
13656 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
13657 && {
13658 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13659 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
13660 }
13661 {
13662 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
13663 let cfg = LaunchConfig {
13664 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
13665 block_dim: (32, 2, 1),
13666 shared_mem_bytes: 0,
13667 };
13668 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
13669 let __s_b = self.gpu.stream();
13670 let mut b = __s_b.launch_builder(&f);
13671 b.arg(bytes)
13672 .arg(aq)
13673 .arg(ad)
13674 .arg(&mut *y)
13675 .arg(&inf)
13676 .arg(&outf)
13677 .arg(&mi)
13678 .arg(&rb);
13679 unsafe {
13680 b.launch(cfg)?;
13681 }
13682 if scale != 1.0 {
13683 self.scale_inplace(y, scale, out_f)?;
13684 }
13685 return Ok(());
13686 }
13687 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
13696 2
13697 } else {
13698 1
13699 };
13700 if m == 1 && qtype == QT_Q4_0 {
13705 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13706 mr = *Q40MR.get_or_init(|| {
13709 std::env::var("MEMRA_Q40_MR")
13710 .ok()
13711 .and_then(|v| v.parse().ok())
13712 .unwrap_or(1)
13713 });
13714 }
13715 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
13726 let q5_force = q5_mode.as_deref() == Some("2");
13727 let q5_il = qtype == QT_Q5_K
13730 && m == 1
13731 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
13732 if q5_il && !q5_force && out_f > 65536 {
13733 mr = 1;
13734 }
13735 if qtype == QT_Q4_0 && rp && mr != 1 {
13738 mr = 2;
13739 }
13740 if qtype == QT_Q8_0 && rp {
13744 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13745 mr = *Q80MR.get_or_init(|| {
13746 std::env::var("MEMRA_Q80_MR")
13747 .ok()
13748 .and_then(|v| v.parse().ok())
13749 .unwrap_or(1)
13750 });
13751 }
13752 let name = match (qtype, mr, rp) {
13753 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
13754 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
13755 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
13756 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
13757 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
13758 (QT_Q5_K, 2, _) => {
13759 if q5_il {
13760 "qmatvec_q5_K_mmvq_mr2_il"
13761 } else {
13762 "qmatvec_q5_K_mmvq_mr2"
13763 }
13764 }
13765 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
13766 (QT_Q8_0, _, true)
13771 if in_f % 1024 == 0 && {
13772 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13773 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
13774 } =>
13775 {
13776 "qmatvec_q8_0_mmvq_rpca"
13777 }
13778 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
13779 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
13780 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
13784 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
13785 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
13786 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
13787 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
13788 (QT_Q5_K, _, _) => {
13789 if q5_il {
13790 "qmatvec_q5_K_mmvq_il"
13791 } else {
13792 "qmatvec_q5_K_mmvq"
13793 }
13794 }
13795 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
13796 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
13797 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
13798 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
13799 };
13800 let f = self.func(name);
13801 let rows_per_block = ROWS_PER_BLOCK * mr;
13803 let cfg = LaunchConfig {
13804 grid_dim: (
13805 (out_f as u32 + rows_per_block - 1) / rows_per_block,
13806 m as u32,
13807 1,
13808 ),
13809 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
13812 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13813 let __s_b = self.gpu.stream();
13814 let mut b = __s_b.launch_builder(&f);
13815 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
13820 if Self::pdl_on()
13823 && Self::pdl_mmvq_on()
13824 && Self::pdl_nvfp4q8_on()
13825 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
13826 {
13827 use cudarc::driver::{DevicePtr, DevicePtrMut};
13828 let s = &self.gpu.stream();
13829 let (pw, _g0) = bytes.device_ptr(s);
13830 let (paq, _g1) = aq.device_ptr(s);
13831 let (pad, _g2) = ad.device_ptr(s);
13832 let (py, _g3) = y.device_ptr_mut(s);
13833 let mut ps = [
13834 &pw as *const _ as *mut std::ffi::c_void,
13835 &paq as *const _ as *mut _,
13836 &pad as *const _ as *mut _,
13837 &py as *const _ as *mut _,
13838 &inf as *const _ as *mut _,
13839 &outf as *const _ as *mut _,
13840 &mi as *const _ as *mut _,
13841 &rb as *const _ as *mut _,
13842 &scale as *const _ as *mut _,
13843 ];
13844 unsafe {
13845 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13846 }
13847 return Ok(());
13848 }
13849 b.arg(bytes)
13850 .arg(aq)
13851 .arg(ad)
13852 .arg(&mut *y)
13853 .arg(&inf)
13854 .arg(&outf)
13855 .arg(&mi)
13856 .arg(&rb)
13857 .arg(&scale);
13858 unsafe {
13859 b.launch(cfg)?;
13860 }
13861 } else if Self::pdl_on()
13862 && Self::pdl_mmvq_on()
13863 && (matches!(
13864 name,
13865 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
13866 ) || (Self::pdl_nvfp4q8_on()
13867 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
13868 {
13869 {
13873 use cudarc::driver::{DevicePtr, DevicePtrMut};
13874 let s = &self.gpu.stream();
13875 let (pw, _g0) = bytes.device_ptr(s);
13876 let (paq, _g1) = aq.device_ptr(s);
13877 let (pad, _g2) = ad.device_ptr(s);
13878 let (py, _g3) = y.device_ptr_mut(s);
13879 let mut ps = [
13880 &pw as *const _ as *mut std::ffi::c_void,
13881 &paq as *const _ as *mut _,
13882 &pad as *const _ as *mut _,
13883 &py as *const _ as *mut _,
13884 &inf as *const _ as *mut _,
13885 &outf as *const _ as *mut _,
13886 &mi as *const _ as *mut _,
13887 &rb as *const _ as *mut _,
13888 ];
13889 unsafe {
13890 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13891 }
13892 }
13893 if scale != 1.0 {
13894 self.scale_inplace(y, scale, m * out_f)?;
13895 }
13896 } else {
13897 b.arg(bytes)
13898 .arg(aq)
13899 .arg(ad)
13900 .arg(&mut *y)
13901 .arg(&inf)
13902 .arg(&outf)
13903 .arg(&mi)
13904 .arg(&rb);
13905 unsafe {
13906 b.launch(cfg)?;
13907 }
13908 if scale != 1.0 {
13909 self.scale_inplace(y, scale, m * out_f)?;
13910 }
13911 }
13912 Ok(())
13913 }
13914
13915 pub fn qmatvec_mmvq_raw(
13919 &self,
13920 bytes: &CudaSlice<u8>,
13921 x: &CudaSlice<f32>,
13922 m: usize,
13923 in_f: usize,
13924 out_f: usize,
13925 qtype: i32,
13926 row_bytes: usize,
13927 rp: bool,
13928 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13929 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13930 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
13931 }
13932
13933 pub fn batched_supports(&self, qtype: i32) -> bool {
13937 matches!(
13938 qtype,
13939 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
13940 )
13941 }
13942
13943 pub fn iq_fast_enabled() -> bool {
13951 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13952 *ON.get_or_init(|| {
13953 std::env::var("MEMRA_IQ_FAST")
13954 .map(|v| v != "0")
13955 .unwrap_or(true)
13956 })
13957 }
13958
13959 pub fn b8_enabled() -> bool {
13962 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13963 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
13964 }
13965
13966 pub fn batched_mcols(m: usize) -> usize {
13968 if m == 2 {
13969 2
13970 } else if m <= 4 {
13971 4
13972 } else if m <= 8 {
13973 8
13974 } else {
13975 16
13976 }
13977 }
13978
13979 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
13984 Some(match (qtype, mcols) {
13985 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
13986 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
13987 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
13988 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
13994 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
13995 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
13996 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
13997 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
14000 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
14001 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
14002 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
14003 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
14006 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
14007 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
14008 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
14009 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
14010 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
14011 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
14012 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
14013 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
14017 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
14018 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
14019 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
14020 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
14024 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
14025 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
14026 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
14027 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
14028 _ => return None,
14029 })
14030 }
14031
14032 pub fn sm_count(&self) -> i32 {
14067 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14068 *SMS.get_or_init(|| {
14069 use cudarc::driver::sys::CUdevice_attribute_enum as A;
14070 self.gpu
14071 .ctx
14072 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14073 .unwrap_or(82)
14074 })
14075 }
14076
14077 pub fn batched_variant(
14078 &self,
14079 _m: usize,
14080 in_f: usize,
14081 out_f: usize,
14082 qtype: i32,
14083 row_bytes: usize,
14084 mcols: usize,
14085 rp: bool,
14086 ) -> &'static str {
14087 if qtype == QT_Q8_0 {
14092 return if rp { "rp" } else { "base" };
14093 }
14094 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14095 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
14096 Ok("base") => "base",
14097 Ok("pf") => "pf",
14098 Ok("r2") => "r2",
14099 Ok("r2w8") => "r2w8",
14100 Ok("pfr2") => "pfr2",
14101 Ok("ca") => "ca",
14102 Ok("car2") => "car2",
14103 Ok("rp") => "rp",
14106 Ok("rpr2") => "rpr2",
14107 Ok("rpr2w8") => "rpr2w8",
14108 Ok("rpca") => "rpca",
14111 Ok("rpcar2") => "rpcar2",
14112 Ok("rpsc") => "rpsc",
14119 Ok("rpms") => "rpms",
14120 Ok("rpmsc") => "rpmsc",
14121 Ok("rpks") => "rpks",
14122 Ok("rpksc") => "rpksc",
14123 _ => "auto",
14124 });
14125 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
14129 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14134 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
14135 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
14136 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
14137 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14138 let sms = *SMS.get_or_init(|| {
14139 use cudarc::driver::sys::CUdevice_attribute_enum as A;
14140 self.gpu
14141 .ctx
14142 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14143 .unwrap_or(82)
14144 });
14145 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
14165 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14168 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
14169 Ok("base") => "base",
14170 Ok("r2") => "r2",
14171 Ok("r2w8") => "r2w8",
14172 _ => "auto",
14173 });
14174 let variant: &'static str = if qtype == QT_Q4_0 {
14175 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14179 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
14180 Ok("base") => "base",
14186 Ok("r2") => "r2",
14187 Ok("ms") => "ms",
14188 Ok("sm") => "sm",
14189 Ok("la") => "la",
14190 _ => "auto",
14191 });
14192 let v = if q40 != "auto" {
14193 q40
14194 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
14195 "r2"
14196 } else {
14197 "base"
14198 };
14199 if rp {
14204 match v {
14205 "ms" => "r2ms_rp",
14206 "sm" => "r2sm_rp",
14207 "la" => "r2la_rp",
14208 "r2" => "r2_rp",
14209 _ => "rp",
14210 }
14211 } else if matches!(v, "ms" | "sm" | "la") {
14212 "r2"
14213 } else {
14214 v
14215 }
14216 } else if qtype != QT_NVFP4 && !kq_r2 {
14217 "base"
14218 } else if kq_r2 && rp {
14219 "rp"
14223 } else if kq_r2 {
14224 if kq_bv != "auto" {
14227 if kq_bv == "r2w8" && mcols != 4 {
14228 "r2"
14229 } else {
14230 kq_bv
14231 }
14232 } else if bv != "auto" {
14233 match bv {
14234 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14235 "r2w8" | "rpr2w8" => {
14236 if mcols != 4 {
14237 "r2"
14238 } else {
14239 "r2w8"
14240 }
14241 }
14242 _ => "base", }
14244 } else {
14245 let blocks = (out_f + 7) / 8;
14246 let waves = blocks as f64 / (7 * sms as usize) as f64;
14247 let filled = blocks >= 4 * sms as usize;
14248 let use_r2 = if qtype == QT_Q4_K {
14249 filled
14250 } else {
14251 waves >= 2.0
14252 };
14253 if use_r2 { "r2" } else { "base" }
14254 }
14255 } else if bv != "auto" {
14256 let v = if bv == "r2w8" && mcols == 2 {
14261 "r2"
14262 } else if bv == "ca" && (!ca_ok || mcols == 8) {
14263 "pf"
14264 } else if bv == "car2" && (!ca_ok || mcols == 8) {
14265 "r2"
14266 } else if bv == "pfr2" && mcols == 8 {
14267 "r2"
14268 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14269 "rpr2"
14270 }
14271 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14273 if mcols == 8 { "rpr2w8" } else { "rpr2" }
14274 } else if bv == "rpcar2" && mcols == 2 {
14275 "rpca"
14276 }
14277 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14280 "rpr2"
14281 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14282 "rpr2"
14283 } else {
14284 bv
14285 };
14286 if rp {
14287 match v {
14288 "base" | "pf" | "ca" | "rp" => "rp",
14289 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14290 "r2w8" | "rpr2w8" => {
14291 if mcols == 2 {
14292 "rpr2"
14293 } else {
14294 "rpr2w8"
14295 }
14296 }
14297 other => other, }
14299 } else {
14300 v
14301 }
14302 } else if mcols == 8 {
14303 if rp {
14314 if sc_ok { "rpsc" } else { "rpr2w8" }
14315 } else {
14316 "r2w8"
14317 }
14318 } else if mcols >= 4 {
14319 let blocks = (out_f + 7) / 8;
14323 let r7 = 7 * sms as usize;
14324 let r8 = 8 * sms as usize;
14325 let waves = blocks as f64 / r7 as f64;
14326 let filled = blocks >= 4 * sms as usize;
14327 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14331 if rp { "rpr2w8" } else { "r2w8" }
14335 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14336 if rp { "rpr2" } else { "r2" }
14339 } else {
14340 if rp { "rp" } else { "pf" }
14344 }
14345 } else if in_f >= 6144 {
14346 if rp { "rpr2" } else { "r2" }
14350 } else if rp {
14351 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14356 if sc_ok && waves >= 0.9 && waves <= 1.1 {
14357 "rpsc"
14358 } else {
14359 "rp"
14360 }
14361 } else {
14362 "base"
14363 };
14364 variant
14365 }
14366
14367 pub fn qmatvec_mmvq_batched(
14368 &self,
14369 bytes: &CudaSlice<u8>,
14370 aq: &CudaSlice<i8>,
14371 ad: &CudaSlice<f32>,
14372 m: usize,
14373 in_f: usize,
14374 out_f: usize,
14375 qtype: i32,
14376 row_bytes: usize,
14377 mcols: usize,
14378 scale: f32,
14379 rp: bool,
14380 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14381 const ROWS_PER_BLOCK: u32 = 4;
14382 let forced: Option<&'static str> = {
14387 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14388 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14389 .as_deref()
14390 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14391 };
14392 let variant = match forced {
14393 Some(v) if !rp || v.contains("rp") => v,
14394 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14395 };
14396 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14397 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14398 })?;
14399 let variant = if mcols == 16 {
14403 if rp { "rp" } else { "base" }
14404 } else {
14405 variant
14406 };
14407 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14414 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14415 if b567
14416 && qtype == QT_NVFP4
14417 && rp
14418 && mcols == 8
14419 && (5..=7).contains(&m)
14420 && matches!(variant, "rpsc" | "rpr2w8")
14421 {
14422 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14423 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14425 let cfg = LaunchConfig {
14426 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14427 block_dim: (32, ROWS_PER_BLOCK, 1),
14428 shared_mem_bytes: 0,
14429 };
14430 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14431 let __s_b = self.gpu.stream();
14432 let mut b = __s_b.launch_builder(&f);
14433 b.arg(bytes)
14434 .arg(aq)
14435 .arg(ad)
14436 .arg(&mut y)
14437 .arg(&inf)
14438 .arg(&outf)
14439 .arg(&mi)
14440 .arg(&rb);
14441 unsafe {
14442 b.launch(cfg)?;
14443 }
14444 if scale != 1.0 {
14445 self.scale_inplace(&mut y, scale, m * out_f)?;
14446 }
14447 return Ok(y);
14448 }
14449 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14450 "base" => (base_name.into(), ROWS_PER_BLOCK),
14451 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14452 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14453 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14454 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14458 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14459 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14460 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14461 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14462 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14463 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14464 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
14466 debug_assert!(
14467 !rp || name.contains("_rp"),
14468 "rp weight dispatched to a GGUF-layout kernel"
14469 );
14470 let f = self.func(&name);
14471 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14472 let smem = if name.contains("_r2sm_rp") {
14474 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
14475 } else {
14476 0
14477 };
14478 let cfg = LaunchConfig {
14479 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14480 block_dim: (32, ROWS_PER_BLOCK, 1),
14481 shared_mem_bytes: smem,
14482 };
14483 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14484 let __s_b = self.gpu.stream();
14485 let mut b = __s_b.launch_builder(&f);
14486 b.arg(bytes)
14487 .arg(aq)
14488 .arg(ad)
14489 .arg(&mut y)
14490 .arg(&inf)
14491 .arg(&outf)
14492 .arg(&mi)
14493 .arg(&rb);
14494 unsafe {
14495 b.launch(cfg)?;
14496 }
14497 if scale != 1.0 {
14498 self.scale_inplace(&mut y, scale, m * out_f)?;
14499 }
14500 Ok(y)
14501 }
14502
14503 pub fn qmatvec_batched_raw(
14507 &self,
14508 bytes: &CudaSlice<u8>,
14509 x: &CudaSlice<f32>,
14510 m: usize,
14511 in_f: usize,
14512 out_f: usize,
14513 qtype: i32,
14514 row_bytes: usize,
14515 mcols: usize,
14516 rp: bool,
14517 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14518 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14519 self.qmatvec_mmvq_batched(
14520 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
14521 )
14522 }
14523
14524 pub fn qmatvec_nvfp4_batched_raw(
14526 &self,
14527 bytes: &CudaSlice<u8>,
14528 x: &CudaSlice<f32>,
14529 m: usize,
14530 in_f: usize,
14531 out_f: usize,
14532 row_bytes: usize,
14533 mcols: usize,
14534 rp: bool,
14535 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14536 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
14537 }
14538
14539 fn try_fp4_gemm(
14543 &self,
14544 w: &crate::model::GpuTensor,
14545 x: &CudaSlice<f32>,
14546 m: usize,
14547 in_f: usize,
14548 out_f: usize,
14549 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14550 use crate::model::GpuTensor;
14551 if cfg!(memra_portable_cuda) {
14552 return Ok(None);
14553 }
14554 if std::env::var("MEMRA_FP4").is_err() {
14555 return Ok(None);
14556 }
14557 #[cfg(memra_cutlass)]
14566 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
14567 if let GpuTensor::Quant {
14568 bytes,
14569 qtype,
14570 scale,
14571 row_bytes,
14572 cutlass,
14573 ..
14574 } = w
14575 {
14576 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
14577 if let Some(cw) = cutlass {
14578 let y = self.cutlass_fp4_gemm(
14580 &cw.b_packed,
14581 &cw.sfb_swizzled,
14582 x,
14583 *scale,
14584 m,
14585 out_f,
14586 in_f,
14587 )?;
14588 return Ok(Some(y));
14589 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
14590 let (b_packed, sfb_sw) =
14595 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
14596 let y =
14597 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
14598 return Ok(Some(y));
14599 }
14600 }
14601 }
14602 }
14603 if let GpuTensor::Quant {
14604 bytes,
14605 qtype,
14606 row_bytes,
14607 scale,
14608 rp,
14609 ..
14610 } = w
14611 {
14612 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
14615 let y =
14616 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
14617 return Ok(Some(y));
14618 }
14619 }
14620 Ok(None)
14621 }
14622
14623 pub fn rms_norm_f16out(
14627 &self,
14628 x: &CudaSlice<f32>,
14629 w: &CudaSlice<f32>,
14630 dst: &mut CudaSlice<f32>,
14631 dst16: &mut CudaSlice<u8>,
14632 ncols: usize,
14633 nrows: usize,
14634 eps: f32,
14635 ) -> Result<(), Box<dyn std::error::Error>> {
14636 let f = self.func("rms_norm_f16out_f32");
14637 let cfg = LaunchConfig {
14638 grid_dim: (nrows as u32, 1, 1),
14639 block_dim: (rms_block(), 1, 1),
14640 shared_mem_bytes: 0,
14641 };
14642 let (nc, e) = (ncols as i32, eps);
14643 let __s_b = self.gpu.stream();
14644 let mut b = __s_b.launch_builder(&f);
14645 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
14646 unsafe {
14647 b.launch(cfg)?;
14648 }
14649 Ok(())
14650 }
14651
14652 #[allow(clippy::too_many_arguments)]
14655 pub fn add_rms_norm_f16out(
14656 &self,
14657 a: &CudaSlice<f32>,
14658 b: &CudaSlice<f32>,
14659 w: &CudaSlice<f32>,
14660 res: &mut CudaSlice<f32>,
14661 dst: &mut CudaSlice<f32>,
14662 dst16: &mut CudaSlice<u8>,
14663 ncols: usize,
14664 nrows: usize,
14665 eps: f32,
14666 ) -> Result<(), Box<dyn std::error::Error>> {
14667 let f = self.func("add_rms_norm_f16out_f32");
14668 let cfg = LaunchConfig {
14669 grid_dim: (nrows as u32, 1, 1),
14670 block_dim: (rms_block(), 1, 1),
14671 shared_mem_bytes: 0,
14672 };
14673 let (nc, e) = (ncols as i32, eps);
14674 let __s_lb = self.gpu.stream();
14675 let mut lb = __s_lb.launch_builder(&f);
14676 lb.arg(a)
14677 .arg(b)
14678 .arg(w)
14679 .arg(res)
14680 .arg(dst)
14681 .arg(dst16)
14682 .arg(&nc)
14683 .arg(&e);
14684 unsafe {
14685 lb.launch(cfg)?;
14686 }
14687 Ok(())
14688 }
14689
14690 pub fn matmul_group_xh(
14693 &self,
14694 ws: &[&crate::model::GpuTensor],
14695 x: &CudaSlice<f32>,
14696 xh: &CudaSlice<u8>,
14697 m: usize,
14698 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14699 let mut out = Vec::with_capacity(ws.len());
14700 let in_f = ws[0].in_features();
14701 for w in ws {
14702 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
14703 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
14704 out.push(y);
14705 continue;
14706 }
14707 }
14708 out.push(self.matmul(w, x, m)?);
14709 }
14710 Ok(out)
14711 }
14712
14713 pub fn gdn_pad_mask(
14716 &self,
14717 beta: &mut CudaSlice<f32>,
14718 g_log: &mut CudaSlice<f32>,
14719 len_d: &CudaSlice<i32>,
14720 h: usize,
14721 t: usize,
14722 ) -> Result<(), Box<dyn std::error::Error>> {
14723 let f = self.func("gdn_pad_mask_f32");
14724 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
14725 let (hi, ti) = (h as i32, t as i32);
14726 let __s_b = self.gpu.stream();
14727 let mut b = __s_b.launch_builder(&f);
14728 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
14729 unsafe {
14730 b.launch(cfg)?;
14731 }
14732 Ok(())
14733 }
14734
14735 pub fn row_gather_dev(
14738 &self,
14739 src: &CudaSlice<f32>,
14740 dst: &mut CudaSlice<f32>,
14741 len_d: &CudaSlice<i32>,
14742 ncols: usize,
14743 ) -> Result<(), Box<dyn std::error::Error>> {
14744 let f = self.func("row_gather_dev_f32");
14745 let cfg = LaunchConfig::for_num_elems(ncols as u32);
14746 let nc = ncols as i32;
14747 let __s_b = self.gpu.stream();
14748 let mut b = __s_b.launch_builder(&f);
14749 b.arg(src).arg(dst).arg(len_d).arg(&nc);
14750 unsafe {
14751 b.launch(cfg)?;
14752 }
14753 Ok(())
14754 }
14755
14756 pub fn matmul_group(
14763 &self,
14764 ws: &[&crate::model::GpuTensor],
14765 x: &CudaSlice<f32>,
14766 m: usize,
14767 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14768 use crate::model::GpuTensor;
14769 let mut out = Vec::with_capacity(ws.len());
14770 let any_mirror = ws
14771 .iter()
14772 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
14773 if m >= 16 && any_mirror && !self.verify_exact_on() {
14774 let in_f = ws[0].in_features();
14775 let xh = self.f16_act(x, m * in_f, in_f)?;
14776 for w in ws {
14777 if w.in_features() == in_f {
14778 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
14779 out.push(y);
14780 continue;
14781 }
14782 }
14783 out.push(self.matmul(w, x, m)?);
14784 }
14785 return Ok(out);
14786 }
14787 for w in ws {
14788 out.push(self.matmul(w, x, m)?);
14789 }
14790 Ok(out)
14791 }
14792
14793 pub fn matmul_group_multi(
14800 &self,
14801 ws: &[&crate::model::GpuTensor],
14802 xs: &[&CudaSlice<f32>],
14803 ms: &[usize],
14804 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
14805 assert_eq!(xs.len(), ms.len());
14806 let in_f = ws[0].in_features();
14807 let total: usize = ms.iter().sum();
14808 let mut xcat = self.uninit(total * in_f)?;
14809 let mut off = 0usize;
14810 for (x, &m) in xs.iter().zip(ms) {
14811 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
14812 off += m;
14813 }
14814 let ys = self.matmul_group(ws, &xcat, total)?;
14815 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
14816 for (w, y) in ws.iter().zip(ys) {
14817 let out_f = w.out_features();
14818 let mut off = 0usize;
14819 for (s, &m) in ms.iter().enumerate() {
14820 let mut ys_s = self.uninit(m * out_f)?;
14821 let src = y.slice(off * out_f..(off + m) * out_f);
14822 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
14823 out[s].push(ys_s);
14824 off += m;
14825 }
14826 }
14827 Ok(out)
14828 }
14829
14830 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
14840 use crate::model::GpuTensor;
14841 if !legacy_quant_gemm_allowed(
14842 cfg!(memra_portable_cuda),
14843 cfg!(memra_hopper_mma),
14844 std::env::var_os("MEMRA_NO_GEMM").is_some(),
14845 ) {
14846 return false;
14847 }
14848 match w {
14849 GpuTensor::Quant { qtype, .. } => {
14850 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
14851 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
14852 }
14853 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
14854 }
14855 }
14856
14857 pub fn qmatvec_gemm(
14864 &self,
14865 w: &crate::model::GpuTensor,
14866 aq: &CudaSlice<i8>,
14867 ad: &CudaSlice<f32>,
14868 m: usize,
14869 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14870 use crate::model::GpuTensor;
14871 let in_f = w.in_features();
14872 let out_f = w.out_features();
14873 let (bytes, qtype, row_bytes, scale, rp) = match w {
14874 GpuTensor::Quant {
14875 bytes,
14876 qtype,
14877 row_bytes,
14878 scale,
14879 rp,
14880 ..
14881 } => (bytes, *qtype, *row_bytes, *scale, *rp),
14882 _ => unreachable!("gemm_supports guaranteed Quant"),
14883 };
14884 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
14890 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
14891 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
14892 if scale != 1.0 {
14893 self.scale_inplace(&mut y, scale, m * out_f)?;
14894 }
14895 return Ok(y);
14896 }
14897 }
14898 let name = match qtype {
14899 QT_Q8_0 => "qmatvec_gemm_q8_0",
14900 QT_Q4_K => "qmatvec_gemm_q4_K",
14901 QT_Q4_0 => {
14902 if rp {
14903 "qmatvec_gemm_q4_0_rp"
14904 } else {
14905 "qmatvec_gemm_q4_0"
14906 }
14907 }
14908 QT_Q5_K => "qmatvec_gemm_q5_K",
14909 QT_Q6_K => "qmatvec_gemm_q6_K",
14910 QT_NVFP4 => {
14911 if rp {
14912 "qmatvec_gemm_nvfp4_rp"
14913 } else {
14914 "qmatvec_gemm_nvfp4"
14915 }
14916 }
14917 _ => unreachable!(),
14918 };
14919 let f = self.func(name);
14920 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);
14925 let k1_tile = if is_k1 {
14927 k1_launch_override().unwrap_or((128, 128, 8))
14928 } else {
14929 (128, 128, 8)
14930 };
14931 let (bm, bn): (u32, u32) = if is_k1 {
14932 (k1_tile.0, k1_tile.1)
14933 } else {
14934 (64, 256)
14935 };
14936 let warps: u32 = if is_k1 {
14937 k1_tile.2
14938 } else {
14939 match qtype {
14940 QT_NVFP4 => 8,
14941 _ => 4,
14942 }
14943 };
14944 let cfg = LaunchConfig {
14945 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
14946 block_dim: (32, warps, 1),
14947 shared_mem_bytes: 0,
14948 };
14949 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14950 let __s_b = self.gpu.stream();
14951 let mut b = __s_b.launch_builder(&f);
14952 b.arg(bytes)
14953 .arg(aq)
14954 .arg(ad)
14955 .arg(&mut y)
14956 .arg(&inf)
14957 .arg(&outf)
14958 .arg(&mi)
14959 .arg(&rb);
14960 unsafe {
14961 b.launch(cfg)?;
14962 }
14963 if scale != 1.0 {
14964 self.scale_inplace(&mut y, scale, m * out_f)?;
14965 }
14966 Ok(y)
14967 }
14968
14969 pub fn qmatvec_gemm_raw(
14974 &self,
14975 bytes: &CudaSlice<u8>,
14976 x: &CudaSlice<f32>,
14977 m: usize,
14978 in_f: usize,
14979 out_f: usize,
14980 qtype: i32,
14981 row_bytes: usize,
14982 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14983 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14984 let name = match qtype {
14985 QT_Q8_0 => "qmatvec_gemm_q8_0",
14986 QT_Q4_K => "qmatvec_gemm_q4_K",
14987 QT_Q4_0 => "qmatvec_gemm_q4_0",
14988 QT_Q5_K => "qmatvec_gemm_q5_K",
14989 QT_Q6_K => "qmatvec_gemm_q6_K",
14990 QT_NVFP4 => "qmatvec_gemm_nvfp4",
14991 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
14992 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
14993 };
14994 let f = self.func(name);
14995 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);
14999 let k1_tile = if is_k1 {
15001 k1_launch_override().unwrap_or((128, 128, 8))
15002 } else {
15003 (128, 128, 8)
15004 };
15005 let (bm, bn): (u32, u32) = if is_k1 {
15006 (k1_tile.0, k1_tile.1)
15007 } else {
15008 (64, 256)
15009 };
15010 let warps: u32 = if is_k1 {
15011 k1_tile.2
15012 } else {
15013 match qtype {
15014 QT_NVFP4 | QT_NVFP4_RP => 8,
15015 _ => 4,
15016 }
15017 };
15018 let cfg = LaunchConfig {
15019 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15020 block_dim: (32, warps, 1),
15021 shared_mem_bytes: 0,
15022 };
15023 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15024 let __s_b = self.gpu.stream();
15025 let mut b = __s_b.launch_builder(&f);
15026 b.arg(bytes)
15027 .arg(&aq)
15028 .arg(&ad)
15029 .arg(&mut y)
15030 .arg(&inf)
15031 .arg(&outf)
15032 .arg(&mi)
15033 .arg(&rb);
15034 unsafe {
15035 b.launch(cfg)?;
15036 }
15037 Ok(y)
15038 }
15039
15040 pub fn qmatvec_gemm_q8_0_wgmma_raw(
15047 &self,
15048 rp4: &CudaSlice<u8>,
15049 aq: &CudaSlice<i8>,
15050 ad: &CudaSlice<f32>,
15051 m: usize,
15052 in_f: usize,
15053 out_f: usize,
15054 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15055 assert!(
15056 out_f % 64 == 0 && in_f % 32 == 0,
15057 "wgmma GEMM needs out_f%64==0, in_f%32==0"
15058 );
15059 let f = self.func("qmatvec_gemm_q8_0_wgmma");
15060 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
15062 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
15063 block_dim: (128, 1, 1),
15064 shared_mem_bytes: 0,
15065 };
15066 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
15067 let __s_b = self.gpu.stream();
15068 let mut b = __s_b.launch_builder(&f);
15069 b.arg(rp4)
15070 .arg(aq)
15071 .arg(ad)
15072 .arg(&mut y)
15073 .arg(&inf)
15074 .arg(&outf)
15075 .arg(&mi);
15076 unsafe {
15077 b.launch(cfg)?;
15078 }
15079 Ok(y)
15080 }
15081
15082 pub fn scale_inplace(
15084 &self,
15085 y: &mut CudaSlice<f32>,
15086 s: f32,
15087 n: usize,
15088 ) -> Result<(), Box<dyn std::error::Error>> {
15089 let f = self.func("scale_f32");
15090 let cfg = LaunchConfig::for_num_elems(n as u32);
15091 let (sf, ni) = (s, n as i32);
15092 let __s_b = self.gpu.stream();
15093 let mut b = __s_b.launch_builder(&f);
15094 b.arg(y).arg(&sf).arg(&ni);
15095 unsafe {
15096 b.launch(cfg)?;
15097 }
15098 Ok(())
15099 }
15100
15101 pub fn bf16_to_f32(
15106 &self,
15107 data: &cudarc::driver::CudaView<'_, u8>,
15108 n: usize,
15109 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15110 let mut out = self.alloc_uninit::<f32>(n)?;
15111 let f = self.func("bf16_to_f32");
15112 let cfg = LaunchConfig::for_num_elems(n as u32);
15113 let ni = n as i32;
15114 let __s_b = self.gpu.stream();
15115 let mut b = __s_b.launch_builder(&f);
15116 b.arg(data).arg(&mut out).arg(&ni);
15117 unsafe {
15118 b.launch(cfg)?;
15119 }
15120 Ok(out)
15121 }
15122
15123 fn linear_bf16_chunked(
15130 &self,
15131 x: &CudaSlice<f32>,
15132 data: &CudaSlice<u8>,
15133 m: usize,
15134 in_f: usize,
15135 out_f: usize,
15136 exact: bool,
15137 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15138 const CHUNK_BYTES: usize = 256 << 20;
15139 let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
15140 if chunk_rows >= out_f {
15141 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
15142 return if exact {
15143 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
15144 } else {
15145 self.linear(x, &wf32, m, in_f, out_f)
15146 };
15147 }
15148 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15149 let mut r0 = 0usize;
15150 while r0 < out_f {
15151 let rows = chunk_rows.min(out_f - r0);
15152 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
15153 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
15154 let yc = if exact {
15155 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
15156 } else {
15157 self.linear(x, &wf32, m, in_f, rows)?
15158 };
15159 for mi in 0..m {
15161 let src = yc.slice(mi * rows..(mi + 1) * rows);
15162 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
15163 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
15164 }
15165 r0 += rows;
15166 }
15167 Ok(y)
15168 }
15169
15170 pub fn linear_decode_exact(
15177 &self,
15178 x: &CudaSlice<f32>,
15179 w: &CudaSlice<f32>,
15180 m_tokens: usize,
15181 in_f: usize,
15182 out_f: usize,
15183 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15184 if m_tokens == 1 {
15185 return self.linear(x, w, 1, in_f, out_f);
15186 }
15187 let xv = self.view(x, m_tokens * in_f);
15188 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
15189 for t in 0..m_tokens {
15190 let row = xv.slice(t * in_f..(t + 1) * in_f);
15191 let mut xr = self.alloc_uninit::<f32>(in_f)?;
15192 self.copy_view_into(&mut xr, 0, &row, in_f)?;
15193 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
15194 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
15195 }
15196 Ok(y)
15197 }
15198
15199 pub fn linear(
15200 &self,
15201 x: &CudaSlice<f32>,
15202 w: &CudaSlice<f32>,
15203 m_tokens: usize,
15204 in_f: usize,
15205 out_f: usize,
15206 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15207 use cudarc::cublaslt::{Matmul, MatmulConfig};
15208 let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; let cfg = MatmulConfig {
15210 transa: true,
15211 transb: false,
15212 transc: false,
15213 m: out_f as u64,
15214 n: m_tokens as u64,
15215 k: in_f as u64,
15216 alpha: 1.0,
15217 lda: in_f as i64,
15218 ldb: in_f as i64,
15219 beta: 0.0,
15220 ldc: out_f as i64,
15221 stride_a: None,
15222 stride_b: None,
15223 stride_c: None,
15224 stride_bias: None,
15225 batch_size: None,
15226 };
15227 unsafe {
15228 self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15229 }
15230 Ok(c)
15231 }
15232
15233 pub fn sdpa_naive(
15235 &self,
15236 q: &CudaSlice<f32>,
15237 k: &CudaSlice<f32>,
15238 v: &CudaSlice<f32>,
15239 o: &mut CudaSlice<f32>,
15240 head_dim: usize,
15241 n_head: usize,
15242 n_head_kv: usize,
15243 t: usize,
15244 t_kv: usize,
15245 scale: f32,
15246 causal: bool,
15247 ) -> Result<(), Box<dyn std::error::Error>> {
15248 let f = self.func("sdpa_naive_f32");
15249 let cfg = LaunchConfig {
15250 grid_dim: (n_head as u32, t as u32, 1),
15251 block_dim: (128, 1, 1),
15252 shared_mem_bytes: (t_kv * 4) as u32,
15253 };
15254 let (hd, nh, nhkv, ti, tkvi, cz) = (
15255 head_dim as i32,
15256 n_head as i32,
15257 n_head_kv as i32,
15258 t as i32,
15259 t_kv as i32,
15260 causal as i32,
15261 );
15262 let __s_b = self.gpu.stream();
15263 let mut b = __s_b.launch_builder(&f);
15264 b.arg(q)
15265 .arg(k)
15266 .arg(v)
15267 .arg(o)
15268 .arg(&hd)
15269 .arg(&nh)
15270 .arg(&nhkv)
15271 .arg(&ti)
15272 .arg(&tkvi)
15273 .arg(&scale)
15274 .arg(&cz);
15275 unsafe {
15276 b.launch(cfg)?;
15277 }
15278 Ok(())
15279 }
15280
15281 #[allow(clippy::too_many_arguments)]
15286 pub fn sdpa_naive_island(
15287 &self,
15288 q: &CudaSlice<f32>,
15289 k: &CudaSlice<f32>,
15290 v: &CudaSlice<f32>,
15291 o: &mut CudaSlice<f32>,
15292 span_id: &CudaSlice<i32>,
15293 head_dim: usize,
15294 n_head: usize,
15295 n_head_kv: usize,
15296 t: usize,
15297 t_kv: usize,
15298 scale: f32,
15299 window: usize,
15300 ) -> Result<(), Box<dyn std::error::Error>> {
15301 let f = self.func("sdpa_naive_island_f32");
15302 let cfg = LaunchConfig {
15303 grid_dim: (n_head as u32, t as u32, 1),
15304 block_dim: (128, 1, 1),
15305 shared_mem_bytes: (t_kv * 4) as u32,
15306 };
15307 let (hd, nh, nhkv, ti, tkvi, wi) = (
15308 head_dim as i32,
15309 n_head as i32,
15310 n_head_kv as i32,
15311 t as i32,
15312 t_kv as i32,
15313 window as i32,
15314 );
15315 let __s_b = self.gpu.stream();
15316 let mut b = __s_b.launch_builder(&f);
15317 b.arg(q)
15318 .arg(k)
15319 .arg(v)
15320 .arg(o)
15321 .arg(span_id)
15322 .arg(&hd)
15323 .arg(&nh)
15324 .arg(&nhkv)
15325 .arg(&ti)
15326 .arg(&tkvi)
15327 .arg(&scale)
15328 .arg(&wi);
15329 unsafe {
15330 b.launch(cfg)?;
15331 }
15332 Ok(())
15333 }
15334
15335 #[allow(clippy::too_many_arguments)]
15337 pub fn sdpa_naive_w(
15338 &self,
15339 q: &CudaSlice<f32>,
15340 k: &CudaSlice<f32>,
15341 v: &CudaSlice<f32>,
15342 o: &mut CudaSlice<f32>,
15343 head_dim: usize,
15344 n_head: usize,
15345 n_head_kv: usize,
15346 t: usize,
15347 t_kv: usize,
15348 scale: f32,
15349 causal: bool,
15350 window: usize,
15351 ) -> Result<(), Box<dyn std::error::Error>> {
15352 let f = self.func("sdpa_naive_w_f32");
15353 let cfg = LaunchConfig {
15354 grid_dim: (n_head as u32, t as u32, 1),
15355 block_dim: (128, 1, 1),
15356 shared_mem_bytes: (t_kv * 4) as u32,
15357 };
15358 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15359 head_dim as i32,
15360 n_head as i32,
15361 n_head_kv as i32,
15362 t as i32,
15363 t_kv as i32,
15364 causal as i32,
15365 window as i32,
15366 );
15367 let __s_b = self.gpu.stream();
15368 let mut b = __s_b.launch_builder(&f);
15369 b.arg(q)
15370 .arg(k)
15371 .arg(v)
15372 .arg(o)
15373 .arg(&hd)
15374 .arg(&nh)
15375 .arg(&nhkv)
15376 .arg(&ti)
15377 .arg(&tkvi)
15378 .arg(&scale)
15379 .arg(&cz)
15380 .arg(&wi);
15381 unsafe {
15382 b.launch(cfg)?;
15383 }
15384 Ok(())
15385 }
15386
15387 pub fn sdpa_naive_view(
15389 &self,
15390 q: &CudaSlice<f32>,
15391 k: &cudarc::driver::CudaView<f32>,
15392 v: &cudarc::driver::CudaView<f32>,
15393 o: &mut CudaSlice<f32>,
15394 head_dim: usize,
15395 n_head: usize,
15396 n_head_kv: usize,
15397 t: usize,
15398 t_kv: usize,
15399 scale: f32,
15400 causal: bool,
15401 ) -> Result<(), Box<dyn std::error::Error>> {
15402 let f = self.func("sdpa_naive_f32");
15403 let cfg = LaunchConfig {
15404 grid_dim: (n_head as u32, t as u32, 1),
15405 block_dim: (128, 1, 1),
15406 shared_mem_bytes: (t_kv * 4) as u32,
15407 };
15408 let (hd, nh, nhkv, ti, tkvi, cz) = (
15409 head_dim as i32,
15410 n_head as i32,
15411 n_head_kv as i32,
15412 t as i32,
15413 t_kv as i32,
15414 causal as i32,
15415 );
15416 let __s_b = self.gpu.stream();
15417 let mut b = __s_b.launch_builder(&f);
15418 b.arg(q)
15419 .arg(k)
15420 .arg(v)
15421 .arg(o)
15422 .arg(&hd)
15423 .arg(&nh)
15424 .arg(&nhkv)
15425 .arg(&ti)
15426 .arg(&tkvi)
15427 .arg(&scale)
15428 .arg(&cz);
15429 unsafe {
15430 b.launch(cfg)?;
15431 }
15432 Ok(())
15433 }
15434
15435 #[allow(clippy::too_many_arguments)]
15443 pub fn fa_dequant_kv_view_f32(
15444 &self,
15445 k: &cudarc::driver::CudaView<u8>,
15446 v: &cudarc::driver::CudaView<u8>,
15447 kf: &mut CudaSlice<f32>,
15448 vf: &mut CudaSlice<f32>,
15449 kv_dim_k: usize,
15450 kv_dim_v: usize,
15451 t_kv: usize,
15452 k_tok_bytes: usize,
15453 v_tok_bytes: usize,
15454 g: bool,
15455 ) -> Result<(), Box<dyn std::error::Error>> {
15456 let f = if g {
15457 self.func_g("fa_dequant_kv_ws_f32")
15458 } else {
15459 self.func("fa_dequant_kv_ws_f32")
15460 };
15461 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
15462 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15463 let cfg = LaunchConfig {
15464 grid_dim: (nblk.max(1), 1, 1),
15465 block_dim: (256, 1, 1),
15466 shared_mem_bytes: 0,
15467 };
15468 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
15469 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15470 let __s_b = self.gpu.stream();
15471 let mut b = __s_b.launch_builder(&f);
15472 b.arg(k)
15473 .arg(v)
15474 .arg(&mut *kf)
15475 .arg(&mut *vf)
15476 .arg(&kdk)
15477 .arg(&kdv)
15478 .arg(&tkvi)
15479 .arg(&ktb)
15480 .arg(&vtb);
15481 unsafe {
15482 b.launch(cfg)?;
15483 }
15484 Ok(())
15485 }
15486
15487 #[allow(clippy::too_many_arguments)]
15488 pub fn sdpa_naive_quantized_view(
15489 &self,
15490 q: &CudaSlice<f32>,
15491 k: &cudarc::driver::CudaView<u8>,
15492 v: &cudarc::driver::CudaView<u8>,
15493 o: &mut CudaSlice<f32>,
15494 head_dim: usize,
15495 n_head: usize,
15496 n_head_kv: usize,
15497 t: usize,
15498 t_kv: usize,
15499 scale: f32,
15500 causal: bool,
15501 k_tok_bytes: usize,
15502 v_tok_bytes: usize,
15503 ) -> Result<(), Box<dyn std::error::Error>> {
15504 let kv_dim = n_head_kv * head_dim;
15505 let mut kf = self.uninit(t_kv * kv_dim)?;
15506 let mut vf = self.uninit(t_kv * kv_dim)?;
15507 let f = self.func("fa_dequant_kv_ws_f32");
15508 let total = (2 * t_kv * kv_dim) as u64;
15509 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15510 let cfg = LaunchConfig {
15511 grid_dim: (nblk.max(1), 1, 1),
15512 block_dim: (256, 1, 1),
15513 shared_mem_bytes: 0,
15514 };
15515 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15516 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15517 let __s_b = self.gpu.stream();
15518 let mut b = __s_b.launch_builder(&f);
15519 b.arg(k)
15520 .arg(v)
15521 .arg(&mut kf)
15522 .arg(&mut vf)
15523 .arg(&kv_dim_i)
15524 .arg(&kv_dim_i)
15525 .arg(&t_kv_i)
15526 .arg(&k_tok_bytes_i)
15527 .arg(&v_tok_bytes_i);
15528 unsafe { b.launch(cfg)? };
15529 self.sdpa_naive(
15530 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15531 )
15532 }
15533
15534 #[allow(clippy::too_many_arguments)]
15546 pub fn sdpa_naive_w_quantized_view(
15547 &self,
15548 q: &CudaSlice<f32>,
15549 k: &cudarc::driver::CudaView<u8>,
15550 v: &cudarc::driver::CudaView<u8>,
15551 o: &mut CudaSlice<f32>,
15552 head_dim: usize,
15553 n_head: usize,
15554 n_head_kv: usize,
15555 t: usize,
15556 t_kv: usize,
15557 scale: f32,
15558 causal: bool,
15559 window: usize,
15560 k_tok_bytes: usize,
15561 v_tok_bytes: usize,
15562 ) -> Result<(), Box<dyn std::error::Error>> {
15563 let kv_dim = n_head_kv * head_dim;
15564 let mut kf = self.uninit(t_kv * kv_dim)?;
15565 let mut vf = self.uninit(t_kv * kv_dim)?;
15566 let f = self.func("fa_dequant_kv_ws_f32");
15567 let total = (2 * t_kv * kv_dim) as u64;
15568 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15569 let cfg = LaunchConfig {
15570 grid_dim: (nblk.max(1), 1, 1),
15571 block_dim: (256, 1, 1),
15572 shared_mem_bytes: 0,
15573 };
15574 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15575 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15576 let __s_b = self.gpu.stream();
15577 let mut b = __s_b.launch_builder(&f);
15578 b.arg(k)
15579 .arg(v)
15580 .arg(&mut kf)
15581 .arg(&mut vf)
15582 .arg(&kv_dim_i)
15583 .arg(&kv_dim_i)
15584 .arg(&t_kv_i)
15585 .arg(&k_tok_bytes_i)
15586 .arg(&v_tok_bytes_i);
15587 unsafe { b.launch(cfg)? };
15588 self.sdpa_naive_w(
15589 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15590 )
15591 }
15592
15593 pub fn fa_prefill(
15597 &self,
15598 q: &CudaSlice<f32>,
15599 k: &CudaSlice<f32>,
15600 v: &CudaSlice<f32>,
15601 o: &mut CudaSlice<f32>,
15602 head_dim: usize,
15603 n_head: usize,
15604 n_head_kv: usize,
15605 t: usize,
15606 t_kv: usize,
15607 scale: f32,
15608 causal: bool,
15609 ) -> Result<(), Box<dyn std::error::Error>> {
15610 if portable_mma_gated() {
15611 return self.sdpa_naive(
15612 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15613 );
15614 }
15615 let fa3_on = head_dim == 256
15623 && causal
15624 && t == t_kv
15625 && match std::env::var("MEMRA_FA3").as_deref() {
15626 Ok("0") => false,
15627 Ok("1") => true,
15628 _ => cfg!(memra_hopper_mma),
15629 };
15630 if fa3_on {
15631 let n = t * n_head * head_dim;
15632 let nkv = t * n_head_kv * head_dim;
15633 let mut q16 = self.alloc_u8_uninit(n * 2)?;
15634 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
15635 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
15636 self.f32_to_bf16_into(q, &mut q16, n)?;
15637 self.f32_to_bf16_into(k, &mut k16, nkv)?;
15638 self.f32_to_bf16_into(v, &mut v16, nkv)?;
15639 let rc = {
15640 use cudarc::driver::{DevicePtr, DevicePtrMut};
15641 let stream = self.gpu.stream();
15642 let (qp, _g1) = q16.device_ptr(&stream);
15643 let (kp, _g2) = k16.device_ptr(&stream);
15644 let (vp, _g3) = v16.device_ptr(&stream);
15645 let (op, _g4) = o.device_ptr_mut(&stream);
15646 unsafe {
15647 memra_fa3_prefill(
15648 qp as *const core::ffi::c_void,
15649 kp as *const core::ffi::c_void,
15650 vp as *const core::ffi::c_void,
15651 op as *mut f32,
15652 t as i32,
15653 n_head as i32,
15654 n_head_kv as i32,
15655 head_dim as i32,
15656 scale,
15657 stream.cu_stream() as *mut core::ffi::c_void,
15658 )
15659 }
15660 };
15661 if rc != 0 {
15662 return Err(format!("memra_fa3_prefill rc={rc}").into());
15663 }
15664 return Ok(());
15665 }
15666 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15671 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
15672 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
15673 const BLOCK_Q: usize = 64;
15674 const BKX: usize = 32;
15675 let f = self.func("fa_prefill_bf16_p1");
15676 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
15677 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
15678 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15679 f.set_attribute(
15680 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15681 shmem as i32,
15682 )?;
15683 let cfg = LaunchConfig {
15684 grid_dim: (
15685 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15686 n_head as u32,
15687 1,
15688 ),
15689 block_dim: (32, 4, 1),
15690 shared_mem_bytes: shmem,
15691 };
15692 let (hd, nh, nhkv, ti, tkvi, cz) = (
15693 head_dim as i32,
15694 n_head as i32,
15695 n_head_kv as i32,
15696 t as i32,
15697 t_kv as i32,
15698 causal as i32,
15699 );
15700 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15701 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15702 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15703 let __s_b = self.gpu.stream();
15704 let mut b = __s_b.launch_builder(&f);
15705 b.arg(&qb)
15706 .arg(&kb)
15707 .arg(&vb)
15708 .arg(o)
15709 .arg(&hd)
15710 .arg(&nh)
15711 .arg(&nhkv)
15712 .arg(&ti)
15713 .arg(&tkvi)
15714 .arg(&scale)
15715 .arg(&cz);
15716 unsafe {
15717 b.launch(cfg)?;
15718 }
15719 return Ok(());
15720 }
15721 const BK: usize = 32;
15727 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
15730 let (block_q, warps, w2_sfx): (usize, u32, &str) =
15731 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
15732 let hd_sfx = fa_hd_suffix(head_dim)?;
15736 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15737 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
15742 let (kb16, vb16) = if bf16kv {
15743 let n = t_kv * n_head_kv * head_dim;
15744 let mut kb = self.alloc_u8_uninit(n * 2)?;
15745 let mut vb = self.alloc_u8_uninit(n * 2)?;
15746 let fcv = self.func("f32_to_bf16_bulk");
15747 let ni = n as i64;
15748 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
15749 let __s_b = self.gpu.stream();
15750 let mut b = __s_b.launch_builder(&fcv);
15751 b.arg(k).arg(&mut kb).arg(&ni);
15752 unsafe {
15753 b.launch(cfgc)?;
15754 }
15755 let __s_b = self.gpu.stream();
15756 let mut b = __s_b.launch_builder(&fcv);
15757 b.arg(v).arg(&mut vb).arg(&ni);
15758 unsafe {
15759 b.launch(cfgc)?;
15760 }
15761 (Some(kb), Some(vb))
15762 } else {
15763 (None, None)
15764 };
15765 let f = self.func(&if bf16kv {
15766 format!("fa_prefill_bf16kv_pp{hd_sfx}")
15767 } else {
15768 format!(
15769 "fa_prefill_f32{}{}{hd_sfx}",
15770 if floor { "" } else { "_pp" },
15771 if floor { "" } else { w2_sfx }
15772 )
15773 });
15774 let kv_stages = if bf16kv { 2 } else { 1 };
15777 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
15778 + 4 * (block_q * BK + 2 * block_q)) as u32;
15779 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15780 f.set_attribute(
15781 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15782 shmem as i32,
15783 )?;
15784 let cfg = LaunchConfig {
15785 grid_dim: (
15786 (t as u32 + block_q as u32 - 1) / block_q as u32,
15787 n_head as u32,
15788 1,
15789 ),
15790 block_dim: (32, warps, 1),
15791 shared_mem_bytes: shmem,
15792 };
15793 let (hd, nh, nhkv, ti, tkvi, cz) = (
15794 head_dim as i32,
15795 n_head as i32,
15796 n_head_kv as i32,
15797 t as i32,
15798 t_kv as i32,
15799 causal as i32,
15800 );
15801 let __s_b = self.gpu.stream();
15802 let mut b = __s_b.launch_builder(&f);
15803 b.arg(q);
15804 match (&kb16, &vb16) {
15805 (Some(kb), Some(vb)) => {
15806 b.arg(kb).arg(vb);
15807 }
15808 _ => {
15809 b.arg(k).arg(v);
15810 }
15811 }
15812 b.arg(o)
15813 .arg(&hd)
15814 .arg(&nh)
15815 .arg(&nhkv)
15816 .arg(&ti)
15817 .arg(&tkvi)
15818 .arg(&scale)
15819 .arg(&cz);
15820 unsafe {
15821 b.launch(cfg)?;
15822 }
15823 Ok(())
15824 }
15825
15826 #[allow(clippy::too_many_arguments)]
15830 pub fn fa_prefill_w(
15831 &self,
15832 q: &CudaSlice<f32>,
15833 k: &CudaSlice<f32>,
15834 v: &CudaSlice<f32>,
15835 o: &mut CudaSlice<f32>,
15836 head_dim: usize,
15837 n_head: usize,
15838 n_head_kv: usize,
15839 t: usize,
15840 t_kv: usize,
15841 scale: f32,
15842 causal: bool,
15843 window: usize,
15844 ) -> Result<(), Box<dyn std::error::Error>> {
15845 if portable_mma_gated() {
15848 return self.sdpa_naive_w(
15849 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15850 );
15851 }
15852 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15856 let faw_f32 =
15857 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
15858 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15859 self.fa_prefill_w_arm(
15860 q,
15861 k,
15862 v,
15863 o,
15864 head_dim,
15865 n_head,
15866 n_head_kv,
15867 t,
15868 t_kv,
15869 scale,
15870 causal,
15871 window,
15872 floor || faw_f32,
15873 floor,
15874 )
15875 }
15876
15877 #[allow(clippy::too_many_arguments)]
15880 pub fn fa_prefill_w_pre(
15881 &self,
15882 qb: &CudaSlice<u8>,
15883 kb: &CudaSlice<u8>,
15884 vb: &CudaSlice<u8>,
15885 o: &mut CudaSlice<f32>,
15886 head_dim: usize,
15887 n_head: usize,
15888 n_head_kv: usize,
15889 t: usize,
15890 t_kv: usize,
15891 scale: f32,
15892 causal: bool,
15893 window: usize,
15894 v_f16: bool,
15895 ) -> Result<(), Box<dyn std::error::Error>> {
15896 const BLOCK_Q: usize = 64;
15897 const BK: usize = 32;
15898 debug_assert_eq!(head_dim, 256);
15899 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
15900 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
15901 if hp {
15902 const BLOCK_QH: usize = 32;
15903 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
15906 let vh: &CudaSlice<u8> = if v_f16 {
15907 vb
15908 } else {
15909 let n = t_kv * n_head_kv * head_dim;
15910 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
15911 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
15912 }
15913 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
15914 vguard.as_ref().unwrap()
15915 };
15916 let f = self.func("fa_prefill_w_bf16_p1h2");
15917 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
15918 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15919 f.set_attribute(
15920 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15921 shmem as i32,
15922 )?;
15923 let cfg = LaunchConfig {
15924 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
15925 block_dim: (32, 4, 1),
15926 shared_mem_bytes: shmem,
15927 };
15928 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15929 head_dim as i32,
15930 n_head as i32,
15931 n_head_kv as i32,
15932 t as i32,
15933 t_kv as i32,
15934 causal as i32,
15935 window as i32,
15936 );
15937 let __s_b = self.gpu.stream();
15938 let mut b = __s_b.launch_builder(&f);
15939 b.arg(qb)
15940 .arg(kb)
15941 .arg(vh)
15942 .arg(o)
15943 .arg(&hd)
15944 .arg(&nh)
15945 .arg(&nhkv)
15946 .arg(&ti)
15947 .arg(&tkvi)
15948 .arg(&scale)
15949 .arg(&cz)
15950 .arg(&wi);
15951 unsafe {
15952 b.launch(cfg)?;
15953 }
15954 return Ok(());
15955 }
15956 let f = self.func("fa_prefill_w_bf16_p1");
15957 let shmem =
15958 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
15959 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15960 f.set_attribute(
15961 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15962 shmem as i32,
15963 )?;
15964 let cfg = LaunchConfig {
15965 grid_dim: (
15966 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15967 n_head as u32,
15968 1,
15969 ),
15970 block_dim: (32, 4, 1),
15971 shared_mem_bytes: shmem,
15972 };
15973 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15974 head_dim as i32,
15975 n_head as i32,
15976 n_head_kv as i32,
15977 t as i32,
15978 t_kv as i32,
15979 causal as i32,
15980 window as i32,
15981 );
15982 let __s_b = self.gpu.stream();
15983 let mut b = __s_b.launch_builder(&f);
15984 b.arg(qb)
15985 .arg(kb)
15986 .arg(vb)
15987 .arg(o)
15988 .arg(&hd)
15989 .arg(&nh)
15990 .arg(&nhkv)
15991 .arg(&ti)
15992 .arg(&tkvi)
15993 .arg(&scale)
15994 .arg(&cz)
15995 .arg(&wi);
15996 unsafe {
15997 b.launch(cfg)?;
15998 }
15999 Ok(())
16000 }
16001
16002 #[allow(clippy::too_many_arguments)]
16004 pub fn fa_prefill_w_arm(
16005 &self,
16006 q: &CudaSlice<f32>,
16007 k: &CudaSlice<f32>,
16008 v: &CudaSlice<f32>,
16009 o: &mut CudaSlice<f32>,
16010 head_dim: usize,
16011 n_head: usize,
16012 n_head_kv: usize,
16013 t: usize,
16014 t_kv: usize,
16015 scale: f32,
16016 causal: bool,
16017 window: usize,
16018 f32_stage: bool,
16019 floor: bool,
16020 ) -> Result<(), Box<dyn std::error::Error>> {
16021 const BLOCK_Q: usize = 64;
16022 const BK: usize = 32;
16023 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
16024 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16028 let p1 = !floor
16029 && !f32_stage
16030 && *P1_ON.get_or_init(|| {
16031 std::env::var("MEMRA_FAW_P1")
16032 .map(|v| v != "0")
16033 .unwrap_or(true)
16034 });
16035 let hp =
16036 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16037 if hp {
16038 const BLOCK_QH: usize = 32;
16039 let f = self.func("fa_prefill_w_bf16_p1h2");
16040 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16041 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16042 f.set_attribute(
16043 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16044 shmem as i32,
16045 )?;
16046 let cfg = LaunchConfig {
16047 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16048 block_dim: (32, 4, 1),
16049 shared_mem_bytes: shmem,
16050 };
16051 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16052 head_dim as i32,
16053 n_head as i32,
16054 n_head_kv as i32,
16055 t as i32,
16056 t_kv as i32,
16057 causal as i32,
16058 window as i32,
16059 );
16060 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16061 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16062 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
16063 let __s_b = self.gpu.stream();
16064 let mut b = __s_b.launch_builder(&f);
16065 b.arg(&qb)
16066 .arg(&kb)
16067 .arg(&vh)
16068 .arg(o)
16069 .arg(&hd)
16070 .arg(&nh)
16071 .arg(&nhkv)
16072 .arg(&ti)
16073 .arg(&tkvi)
16074 .arg(&scale)
16075 .arg(&cz)
16076 .arg(&wi);
16077 unsafe {
16078 b.launch(cfg)?;
16079 }
16080 return Ok(());
16081 }
16082 if p1 {
16083 let f = self.func("fa_prefill_w_bf16_p1");
16084 let shmem =
16085 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16086 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16087 f.set_attribute(
16088 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16089 shmem as i32,
16090 )?;
16091 let cfg = LaunchConfig {
16092 grid_dim: (
16093 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16094 n_head as u32,
16095 1,
16096 ),
16097 block_dim: (32, 4, 1),
16098 shared_mem_bytes: shmem,
16099 };
16100 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16101 head_dim as i32,
16102 n_head as i32,
16103 n_head_kv as i32,
16104 t as i32,
16105 t_kv as i32,
16106 causal as i32,
16107 window as i32,
16108 );
16109 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16110 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16111 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16112 let __s_b = self.gpu.stream();
16113 let mut b = __s_b.launch_builder(&f);
16114 b.arg(&qb)
16115 .arg(&kb)
16116 .arg(&vb)
16117 .arg(o)
16118 .arg(&hd)
16119 .arg(&nh)
16120 .arg(&nhkv)
16121 .arg(&ti)
16122 .arg(&tkvi)
16123 .arg(&scale)
16124 .arg(&cz)
16125 .arg(&wi);
16126 unsafe {
16127 b.launch(cfg)?;
16128 }
16129 return Ok(());
16130 }
16131 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16134 let g4 = !floor
16135 && !f32_stage
16136 && n_head_kv == 1
16137 && n_head % 4 == 0
16138 && *G4_ON.get_or_init(|| {
16139 std::env::var("MEMRA_FAW_G4")
16140 .map(|v| v != "0")
16141 .unwrap_or(true)
16142 });
16143 if g4 {
16144 const SP_M: usize = 16;
16145 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16148 let o2 = *O2_ON.get_or_init(|| {
16149 std::env::var("MEMRA_FAW_O2")
16150 .map(|v| v != "0")
16151 .unwrap_or(true)
16152 });
16153 let f = self.func(if o2 {
16154 "fa_prefill_w_bf16_g4o2"
16155 } else {
16156 "fa_prefill_w_bf16_g4"
16157 });
16158 let shmem = if o2 {
16159 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
16160 } else {
16161 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
16162 as u32
16163 };
16164 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16165 f.set_attribute(
16166 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16167 shmem as i32,
16168 )?;
16169 let cfg = LaunchConfig {
16170 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
16171 block_dim: (32, 4, 1),
16172 shared_mem_bytes: shmem,
16173 };
16174 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16175 head_dim as i32,
16176 n_head as i32,
16177 n_head_kv as i32,
16178 t as i32,
16179 t_kv as i32,
16180 causal as i32,
16181 window as i32,
16182 );
16183 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16184 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16185 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16186 let __s_b = self.gpu.stream();
16187 let mut b = __s_b.launch_builder(&f);
16188 b.arg(&qb)
16189 .arg(&kb)
16190 .arg(&vb)
16191 .arg(o)
16192 .arg(&hd)
16193 .arg(&nh)
16194 .arg(&nhkv)
16195 .arg(&ti)
16196 .arg(&tkvi)
16197 .arg(&scale)
16198 .arg(&cz)
16199 .arg(&wi);
16200 unsafe {
16201 b.launch(cfg)?;
16202 }
16203 return Ok(());
16204 }
16205 let f = self.func(if floor {
16206 "fa_prefill_w_f32"
16207 } else if f32_stage {
16208 "fa_prefill_w_f32_pp"
16209 } else {
16210 "fa_prefill_w_bf16_pp"
16211 });
16212 let shmem =
16213 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16214 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16215 f.set_attribute(
16216 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16217 shmem as i32,
16218 )?;
16219 let cfg = LaunchConfig {
16220 grid_dim: (
16221 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16222 n_head as u32,
16223 1,
16224 ),
16225 block_dim: (32, 4, 1),
16226 shared_mem_bytes: shmem,
16227 };
16228 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16229 head_dim as i32,
16230 n_head as i32,
16231 n_head_kv as i32,
16232 t as i32,
16233 t_kv as i32,
16234 causal as i32,
16235 window as i32,
16236 );
16237 if f32_stage {
16238 let __s_b = self.gpu.stream();
16239 let mut b = __s_b.launch_builder(&f);
16240 b.arg(q)
16241 .arg(k)
16242 .arg(v)
16243 .arg(o)
16244 .arg(&hd)
16245 .arg(&nh)
16246 .arg(&nhkv)
16247 .arg(&ti)
16248 .arg(&tkvi)
16249 .arg(&scale)
16250 .arg(&cz)
16251 .arg(&wi);
16252 unsafe {
16253 b.launch(cfg)?;
16254 }
16255 } else {
16256 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16257 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16258 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16259 let __s_b = self.gpu.stream();
16260 let mut b = __s_b.launch_builder(&f);
16261 b.arg(&qb)
16262 .arg(&kb)
16263 .arg(&vb)
16264 .arg(o)
16265 .arg(&hd)
16266 .arg(&nh)
16267 .arg(&nhkv)
16268 .arg(&ti)
16269 .arg(&tkvi)
16270 .arg(&scale)
16271 .arg(&cz)
16272 .arg(&wi);
16273 unsafe {
16274 b.launch(cfg)?;
16275 }
16276 }
16277 Ok(())
16278 }
16279
16280 #[allow(clippy::too_many_arguments)]
16284 pub fn fa_prefill_hd512(
16285 &self,
16286 q: &CudaSlice<f32>,
16287 k: &CudaSlice<f32>,
16288 v: &CudaSlice<f32>,
16289 o: &mut CudaSlice<f32>,
16290 head_dim: usize,
16291 n_head: usize,
16292 n_head_kv: usize,
16293 t: usize,
16294 t_kv: usize,
16295 scale: f32,
16296 causal: bool,
16297 ) -> Result<(), Box<dyn std::error::Error>> {
16298 if portable_mma_gated() {
16300 return self.sdpa_naive(
16301 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16302 );
16303 }
16304 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16310 let f32_stage =
16311 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16312 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16316 let sp = !f32_stage
16317 && *SP_ON.get_or_init(|| {
16318 std::env::var("MEMRA_FA512_SP")
16319 .map(|v| v != "0")
16320 .unwrap_or(true)
16321 });
16322 self.fa_prefill_hd512_arm(
16323 q,
16324 k,
16325 v,
16326 o,
16327 head_dim,
16328 n_head,
16329 n_head_kv,
16330 t,
16331 t_kv,
16332 scale,
16333 causal,
16334 f32_stage,
16335 sp,
16336 sp && fa_f16pv_on(),
16337 )
16338 }
16339
16340 #[allow(clippy::too_many_arguments)]
16342 pub fn fa_prefill_hd512_pre(
16343 &self,
16344 qb: &CudaSlice<u8>,
16345 kb: &CudaSlice<u8>,
16346 vb: &CudaSlice<u8>,
16347 o: &mut CudaSlice<f32>,
16348 head_dim: usize,
16349 n_head: usize,
16350 n_head_kv: usize,
16351 t: usize,
16352 t_kv: usize,
16353 scale: f32,
16354 causal: bool,
16355 v_f16: bool,
16356 ) -> Result<(), Box<dyn std::error::Error>> {
16357 debug_assert_eq!(head_dim, 512);
16358 const SP_M: usize = 16;
16359 const BKS: usize = 32;
16360 let f16pv = fa_f16pv_on();
16364 let nw = if f16pv { fa512_wide_warps() } else { 2 };
16365 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16366 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16367 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16368 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16369 let n = t_kv * n_head_kv * head_dim;
16371 let need = n * 2;
16372 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16373 *vguard = Some(self.alloc_uninit::<u8>(need)?);
16374 }
16375 let dst = vguard.as_mut().unwrap();
16376 self.bf16_to_f16_into(vb, n, dst)?;
16377 vguard.as_ref().unwrap()
16378 } else {
16379 vb
16380 };
16381 let f = self.func(if hp {
16382 "fa_prefill_bf16_hd512_sp16h2"
16383 } else {
16384 match (f16pv, nw) {
16385 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16386 (true, _) => "fa_prefill_bf16_hd512_sp16",
16387 _ => "fa_prefill_bf16_hd512_sp",
16388 }
16389 });
16390 let (nwarp, npart) = if hp {
16391 (4usize, 4usize)
16392 } else if nw > 2 {
16393 (nw, nw)
16394 } else {
16395 (2, 1)
16396 };
16397 let shmem = if hp {
16399 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
16400 as u32
16401 } else {
16402 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16403 + 4 * (npart * SP_M * BKS + SP_M)) as u32
16404 };
16405 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16406 f.set_attribute(
16407 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16408 shmem as i32,
16409 )?;
16410 let grid_y = if hp {
16411 (n_head / 2) as u32
16412 } else {
16413 n_head as u32
16414 };
16415 let cfg = LaunchConfig {
16416 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16417 block_dim: (32, nwarp as u32, 1),
16418 shared_mem_bytes: shmem,
16419 };
16420 let (hd, nh, nhkv, ti, tkvi, cz) = (
16421 head_dim as i32,
16422 n_head as i32,
16423 n_head_kv as i32,
16424 t as i32,
16425 t_kv as i32,
16426 causal as i32,
16427 );
16428 let __s_b = self.gpu.stream();
16429 let mut b = __s_b.launch_builder(&f);
16430 b.arg(qb)
16431 .arg(kb)
16432 .arg(vref)
16433 .arg(o)
16434 .arg(&hd)
16435 .arg(&nh)
16436 .arg(&nhkv)
16437 .arg(&ti)
16438 .arg(&tkvi)
16439 .arg(&scale)
16440 .arg(&cz);
16441 unsafe {
16442 b.launch(cfg)?;
16443 }
16444 Ok(())
16445 }
16446
16447 #[allow(clippy::too_many_arguments)]
16450 pub fn fa_prefill_hd512_arm(
16451 &self,
16452 q: &CudaSlice<f32>,
16453 k: &CudaSlice<f32>,
16454 v: &CudaSlice<f32>,
16455 o: &mut CudaSlice<f32>,
16456 head_dim: usize,
16457 n_head: usize,
16458 n_head_kv: usize,
16459 t: usize,
16460 t_kv: usize,
16461 scale: f32,
16462 causal: bool,
16463 f32_stage: bool,
16464 sp: bool,
16465 f16pv: bool,
16466 ) -> Result<(), Box<dyn std::error::Error>> {
16467 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
16468 if sp && !f32_stage {
16469 const SP_M: usize = 16;
16473 const BKS: usize = 32;
16474 let nw = if f16pv { fa512_wide_warps() } else { 2 };
16475 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16476 let f = self.func(if hp {
16477 "fa_prefill_bf16_hd512_sp16h2"
16478 } else {
16479 match (f16pv, nw) {
16480 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16481 (true, _) => "fa_prefill_bf16_hd512_sp16",
16482 _ => "fa_prefill_bf16_hd512_sp",
16483 }
16484 });
16485 let (nwarp, npart) = if hp {
16486 (4usize, 4usize)
16487 } else if nw > 2 {
16488 (nw, nw)
16489 } else {
16490 (2, 1)
16491 };
16492 let shmem = if hp {
16493 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
16494 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
16495 } else {
16496 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16497 + 4 * (npart * SP_M * BKS + SP_M)) as u32
16498 };
16499 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16500 f.set_attribute(
16501 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16502 shmem as i32,
16503 )?;
16504 let grid_y = if hp {
16505 (n_head / 2) as u32
16506 } else {
16507 n_head as u32
16508 };
16509 let cfg = LaunchConfig {
16510 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16511 block_dim: (32, nwarp as u32, 1),
16512 shared_mem_bytes: shmem,
16513 };
16514 let (hd, nh, nhkv, ti, tkvi, cz) = (
16515 head_dim as i32,
16516 n_head as i32,
16517 n_head_kv as i32,
16518 t as i32,
16519 t_kv as i32,
16520 causal as i32,
16521 );
16522 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16523 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16524 let vb = if f16pv {
16525 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
16526 } else {
16527 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
16528 };
16529 let __s_b = self.gpu.stream();
16530 let mut b = __s_b.launch_builder(&f);
16531 b.arg(&qb)
16532 .arg(&kb)
16533 .arg(&vb)
16534 .arg(o)
16535 .arg(&hd)
16536 .arg(&nh)
16537 .arg(&nhkv)
16538 .arg(&ti)
16539 .arg(&tkvi)
16540 .arg(&scale)
16541 .arg(&cz);
16542 unsafe {
16543 b.launch(cfg)?;
16544 }
16545 return Ok(());
16546 }
16547 const BLOCK_Q: usize = 32;
16548 const BK: usize = 32;
16549 const HALF: usize = 256;
16550 let f = self.func(if f32_stage {
16551 "fa_prefill_f32_hd512"
16552 } else {
16553 "fa_prefill_bf16_hd512"
16554 });
16555 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
16557 + 4 * BLOCK_Q) as u32;
16558 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16559 f.set_attribute(
16560 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16561 shmem as i32,
16562 )?;
16563 let cfg = LaunchConfig {
16564 grid_dim: (
16565 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16566 n_head as u32,
16567 2,
16568 ),
16569 block_dim: (32, 2, 1),
16570 shared_mem_bytes: shmem,
16571 };
16572 let (hd, nh, nhkv, ti, tkvi, cz) = (
16573 head_dim as i32,
16574 n_head as i32,
16575 n_head_kv as i32,
16576 t as i32,
16577 t_kv as i32,
16578 causal as i32,
16579 );
16580 if f32_stage {
16581 let __s_b = self.gpu.stream();
16582 let mut b = __s_b.launch_builder(&f);
16583 b.arg(q)
16584 .arg(k)
16585 .arg(v)
16586 .arg(o)
16587 .arg(&hd)
16588 .arg(&nh)
16589 .arg(&nhkv)
16590 .arg(&ti)
16591 .arg(&tkvi)
16592 .arg(&scale)
16593 .arg(&cz);
16594 unsafe {
16595 b.launch(cfg)?;
16596 }
16597 } else {
16598 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16599 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16600 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16601 let __s_b = self.gpu.stream();
16602 let mut b = __s_b.launch_builder(&f);
16603 b.arg(&qb)
16604 .arg(&kb)
16605 .arg(&vb)
16606 .arg(o)
16607 .arg(&hd)
16608 .arg(&nh)
16609 .arg(&nhkv)
16610 .arg(&ti)
16611 .arg(&tkvi)
16612 .arg(&scale)
16613 .arg(&cz);
16614 unsafe {
16615 b.launch(cfg)?;
16616 }
16617 }
16618 Ok(())
16619 }
16620
16621 #[allow(clippy::too_many_arguments)]
16625 pub fn rope_neox2_bf16e(
16626 &self,
16627 q: &mut CudaSlice<f32>,
16628 k: &mut CudaSlice<f32>,
16629 qb: &mut CudaSlice<u8>,
16630 kb: &mut CudaSlice<u8>,
16631 pos: &CudaSlice<i32>,
16632 head_dim: usize,
16633 n_dims: usize,
16634 nh_q: usize,
16635 nh_k: usize,
16636 n_tokens: usize,
16637 base: f32,
16638 freq_scale: f32,
16639 ff: Option<&CudaSlice<f32>>,
16640 ) -> Result<(), Box<dyn std::error::Error>> {
16641 let f = self.func("rope_neox2_bf16e_f32");
16642 let rows = ((nh_q + nh_k) * n_tokens) as u32;
16643 let cfg = LaunchConfig {
16644 grid_dim: (rows, 1, 1),
16645 block_dim: ((head_dim / 2) as u32, 1, 1),
16646 shared_mem_bytes: 0,
16647 };
16648 let theta_scale = base.powf(-2.0 / n_dims as f32);
16649 let (hd, nd, nhq, nhk, nt) = (
16650 head_dim as i32,
16651 n_dims as i32,
16652 nh_q as i32,
16653 nh_k as i32,
16654 n_tokens as i32,
16655 );
16656 let __s_b = self.gpu.stream();
16657 let mut b = __s_b.launch_builder(&f);
16658 match ff {
16659 Some(t) => {
16660 b.arg(&mut *q)
16661 .arg(&mut *k)
16662 .arg(&mut *qb)
16663 .arg(&mut *kb)
16664 .arg(pos)
16665 .arg(&hd)
16666 .arg(&nd)
16667 .arg(&nhq)
16668 .arg(&nhk)
16669 .arg(&nt)
16670 .arg(&theta_scale)
16671 .arg(&freq_scale)
16672 .arg(t);
16673 unsafe {
16674 b.launch(cfg)?;
16675 }
16676 }
16677 None => {
16678 let null: u64 = 0;
16679 b.arg(&mut *q)
16680 .arg(&mut *k)
16681 .arg(&mut *qb)
16682 .arg(&mut *kb)
16683 .arg(pos)
16684 .arg(&hd)
16685 .arg(&nd)
16686 .arg(&nhq)
16687 .arg(&nhk)
16688 .arg(&nt)
16689 .arg(&theta_scale)
16690 .arg(&freq_scale)
16691 .arg(&null);
16692 unsafe {
16693 b.launch(cfg)?;
16694 }
16695 }
16696 }
16697 Ok(())
16698 }
16699
16700 pub fn f32_to_bf16(
16703 &self,
16704 x: &CudaSlice<f32>,
16705 n: usize,
16706 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16707 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
16708 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16709 let f = self.func("f32_to_bf16_flat");
16710 let n_i = n as i64;
16711 let cfg = LaunchConfig {
16712 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16713 block_dim: (256, 1, 1),
16714 shared_mem_bytes: 0,
16715 };
16716 let __s_b = self.gpu.stream();
16717 let mut b = __s_b.launch_builder(&f);
16718 b.arg(x).arg(&mut y).arg(&n_i);
16719 unsafe {
16720 b.launch(cfg)?;
16721 }
16722 Ok(y)
16723 }
16724
16725 pub fn f32_to_f16(
16726 &self,
16727 x: &CudaSlice<f32>,
16728 n: usize,
16729 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16730 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
16731 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16732 let f = self.func("f32_to_f16_flat");
16733 let n_i = n as i64;
16734 let cfg = LaunchConfig {
16735 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16736 block_dim: (256, 1, 1),
16737 shared_mem_bytes: 0,
16738 };
16739 let __s_b = self.gpu.stream();
16740 let mut b = __s_b.launch_builder(&f);
16741 b.arg(x).arg(&mut y).arg(&n_i);
16742 unsafe {
16743 b.launch(cfg)?;
16744 }
16745 Ok(y)
16746 }
16747
16748 pub fn bf16_to_f16(
16750 &self,
16751 xb: &CudaSlice<u8>,
16752 n: usize,
16753 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16754 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16755 self.bf16_to_f16_into(xb, n, &mut y)?;
16756 Ok(y)
16757 }
16758
16759 pub fn bf16_to_f16_into(
16761 &self,
16762 xb: &CudaSlice<u8>,
16763 n: usize,
16764 y: &mut CudaSlice<u8>,
16765 ) -> Result<(), Box<dyn std::error::Error>> {
16766 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
16767 assert!(y.len() >= n * 2);
16768 let f = self.func("bf16_to_f16_flat");
16769 let n2 = (n / 2) as i64;
16770 let cfg = LaunchConfig {
16771 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
16772 block_dim: (256, 1, 1),
16773 shared_mem_bytes: 0,
16774 };
16775 let __s_b = self.gpu.stream();
16776 let mut b = __s_b.launch_builder(&f);
16777 b.arg(xb).arg(y).arg(&n2);
16778 unsafe {
16779 b.launch(cfg)?;
16780 }
16781 Ok(())
16782 }
16783
16784 #[allow(clippy::too_many_arguments)]
16789 pub fn fa_prefill_vl8(
16790 &self,
16791 seqs: &[FaSeqVl],
16792 head_dim: usize,
16793 n_head: usize,
16794 n_head_kv: usize,
16795 scale: f32,
16796 ) -> Result<(), Box<dyn std::error::Error>> {
16797 const BK: usize = 32;
16798 let b = seqs.len();
16799 assert!(b >= 1 && b <= 8);
16800 let mut packed = [FaSeqVl::default(); 8];
16801 packed[..b].copy_from_slice(seqs);
16802 let v = FaVl8(packed);
16803 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16804 let ept = (n_head_kv * head_dim) as i32;
16805 {
16806 let f = self.func("fa_mirror_vl");
16807 let max_n = (max_t as i64) * ept as i64;
16808 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
16809 for which in 0..2i32 {
16810 let cfg = LaunchConfig {
16811 grid_dim: (blocks, 1, b as u32),
16812 block_dim: (256, 1, 1),
16813 shared_mem_bytes: 0,
16814 };
16815 let __s_lb = self.gpu.stream();
16816 let mut lb = __s_lb.launch_builder(&f);
16817 lb.arg(&v).arg(&ept).arg(&which);
16818 unsafe {
16819 lb.launch(cfg)?;
16820 }
16821 }
16822 }
16823 let hd_sfx = fa_hd_suffix(head_dim)?;
16824 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
16825 let block_q = 64usize;
16826 let kv_stages = 2usize;
16827 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16828 + 4 * (block_q * BK + 2 * block_q)) as u32;
16829 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16830 f.set_attribute(
16831 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16832 shmem as i32,
16833 )?;
16834 let cfg = LaunchConfig {
16835 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
16836 block_dim: (32, 4, 1),
16837 shared_mem_bytes: shmem,
16838 };
16839 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16840 let __s_lb = self.gpu.stream();
16841 let mut lb = __s_lb.launch_builder(&f);
16842 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
16843 unsafe {
16844 lb.launch(cfg)?;
16845 }
16846 Ok(())
16847 }
16848
16849 #[allow(clippy::too_many_arguments)]
16853 pub fn attn_pre_vl8(
16854 &self,
16855 seqs: &[AttnPreVl],
16856 wq: &CudaSlice<f32>,
16857 wk: &CudaSlice<f32>,
16858 head_dim: usize,
16859 rope_dims: usize,
16860 n_head: usize,
16861 n_head_kv: usize,
16862 eps: f32,
16863 freq_base: f32,
16864 freq_scale: f32,
16865 kv_dim_k: usize,
16866 kv_dim_v: usize,
16867 k_tok_bytes: usize,
16868 v_tok_bytes: usize,
16869 ) -> Result<(), Box<dyn std::error::Error>> {
16870 let b = seqs.len();
16871 assert!(b >= 1 && b <= 8);
16872 let mut packed = [AttnPreVl::default(); 8];
16873 packed[..b].copy_from_slice(seqs);
16874 let v = AttnPreVl8(packed);
16875 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16876 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16877 {
16878 let f = self.func("q_gate_split_vl");
16879 let n = max_t * (n_head * head_dim) as u32;
16880 let cfg = LaunchConfig {
16881 grid_dim: (n.div_ceil(256), 1, b as u32),
16882 block_dim: (256, 1, 1),
16883 shared_mem_bytes: 0,
16884 };
16885 let __s_lb = self.gpu.stream();
16886 let mut lb = __s_lb.launch_builder(&f);
16887 lb.arg(&v).arg(&hd).arg(&nh);
16888 unsafe {
16889 lb.launch(cfg)?;
16890 }
16891 }
16892 {
16893 let f = self.func("attn_rms_vl");
16894 let cfg = LaunchConfig {
16895 grid_dim: (max_t * n_head as u32, 2, b as u32),
16896 block_dim: (rms_block(), 1, 1),
16897 shared_mem_bytes: 0,
16898 };
16899 let __s_lb = self.gpu.stream();
16900 let mut lb = __s_lb.launch_builder(&f);
16901 lb.arg(&v)
16902 .arg(wq)
16903 .arg(wk)
16904 .arg(&hd)
16905 .arg(&nh)
16906 .arg(&nhkv)
16907 .arg(&eps);
16908 unsafe {
16909 lb.launch(cfg)?;
16910 }
16911 }
16912 {
16913 let f = self.func("attn_rope_vl");
16914 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
16915 let nd = rope_dims as i32;
16916 let cfg = LaunchConfig {
16917 grid_dim: (max_t * n_head as u32, 2, b as u32),
16918 block_dim: ((head_dim / 2) as u32, 1, 1),
16919 shared_mem_bytes: 0,
16920 };
16921 let __s_lb = self.gpu.stream();
16922 let mut lb = __s_lb.launch_builder(&f);
16923 lb.arg(&v)
16924 .arg(&hd)
16925 .arg(&nd)
16926 .arg(&nh)
16927 .arg(&nhkv)
16928 .arg(&theta_scale)
16929 .arg(&freq_scale);
16930 unsafe {
16931 lb.launch(cfg)?;
16932 }
16933 }
16934 {
16935 let f = self.func("append_kv_vl");
16936 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
16937 let cfg = LaunchConfig {
16938 grid_dim: (nblk, max_t, b as u32),
16939 block_dim: (32, 1, 1),
16940 shared_mem_bytes: 0,
16941 };
16942 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
16943 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16944 let __s_lb = self.gpu.stream();
16945 let mut lb = __s_lb.launch_builder(&f);
16946 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
16947 unsafe {
16948 lb.launch(cfg)?;
16949 }
16950 }
16951 Ok(())
16952 }
16953
16954 pub fn fa_prefill_view(
16959 &self,
16960 q: &CudaSlice<f32>,
16961 k: &cudarc::driver::CudaView<u8>,
16962 v: &cudarc::driver::CudaView<u8>,
16963 o: &mut CudaSlice<f32>,
16964 head_dim: usize,
16965 n_head: usize,
16966 n_head_kv: usize,
16967 t: usize,
16968 t_kv: usize,
16969 scale: f32,
16970 causal: bool,
16971 k_tok_bytes: usize,
16972 v_tok_bytes: usize,
16973 g: bool,
16974 ) -> Result<(), Box<dyn std::error::Error>> {
16975 if portable_mma_gated() {
16976 return self.sdpa_naive_quantized_view(
16977 q,
16978 k,
16979 v,
16980 o,
16981 head_dim,
16982 n_head,
16983 n_head_kv,
16984 t,
16985 t_kv,
16986 scale,
16987 causal,
16988 k_tok_bytes,
16989 v_tok_bytes,
16990 );
16991 }
16992 const BLOCK_Q: usize = 64;
16993 const BK: usize = 32;
16994 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
16997 let f = if g {
16998 self.func_g(&name)
16999 } else {
17000 self.func(&name)
17001 };
17002 let shmem =
17003 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
17004 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17005 f.set_attribute(
17006 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17007 shmem as i32,
17008 )?;
17009 let cfg = LaunchConfig {
17010 grid_dim: (
17011 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17012 n_head as u32,
17013 1,
17014 ),
17015 block_dim: (32, 4, 1),
17016 shared_mem_bytes: shmem,
17017 };
17018 let (hd, nh, nhkv, ti, tkvi, cz) = (
17019 head_dim as i32,
17020 n_head as i32,
17021 n_head_kv as i32,
17022 t as i32,
17023 t_kv as i32,
17024 causal as i32,
17025 );
17026 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17027 let __s_b = self.gpu.stream();
17028 let mut b = __s_b.launch_builder(&f);
17029 b.arg(q)
17030 .arg(k)
17031 .arg(v)
17032 .arg(o)
17033 .arg(&hd)
17034 .arg(&nh)
17035 .arg(&nhkv)
17036 .arg(&ti)
17037 .arg(&tkvi)
17038 .arg(&scale)
17039 .arg(&cz)
17040 .arg(&ktb)
17041 .arg(&vtb);
17042 unsafe {
17043 b.launch(cfg)?;
17044 }
17045 Ok(())
17046 }
17047
17048 #[allow(clippy::too_many_arguments)]
17058 pub fn fa_prefill_view_ws(
17059 &self,
17060 q: &CudaSlice<f32>,
17061 k: &cudarc::driver::CudaView<u8>,
17062 v: &cudarc::driver::CudaView<u8>,
17063 o: &mut CudaSlice<f32>,
17064 head_dim: usize,
17065 n_head: usize,
17066 n_head_kv: usize,
17067 t: usize,
17068 t_kv: usize,
17069 scale: f32,
17070 causal: bool,
17071 k_tok_bytes: usize,
17072 v_tok_bytes: usize,
17073 g: bool,
17074 ) -> Result<(), Box<dyn std::error::Error>> {
17075 if portable_mma_gated() {
17076 return self.sdpa_naive_quantized_view(
17077 q,
17078 k,
17079 v,
17080 o,
17081 head_dim,
17082 n_head,
17083 n_head_kv,
17084 t,
17085 t_kv,
17086 scale,
17087 causal,
17088 k_tok_bytes,
17089 v_tok_bytes,
17090 );
17091 }
17092 const BLOCK_Q: usize = 64;
17093 const BK: usize = 32;
17094 let kv_dim_k = n_head_kv * head_dim;
17095 let kv_dim_v = n_head_kv * head_dim;
17096 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17098 let mut guard = self.prime_deqw_ws.lock().unwrap();
17100 let need_grow = match guard.as_ref() {
17101 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17102 None => true,
17103 };
17104 if need_grow {
17105 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17106 let (ck, cv) = guard
17107 .as_ref()
17108 .map(|(a, b)| (a.len(), b.len()))
17109 .unwrap_or((0, 0));
17110 *guard = Some((
17111 self.alloc_u8(grow(ck, k_ws_bytes))?,
17112 self.alloc_u8(grow(cv, v_ws_bytes))?,
17113 ));
17114 }
17115 let (kw, vw) = guard.as_mut().unwrap();
17116 {
17118 let f = if g {
17120 self.func_g("fa_dequant_kv_ws_bf16")
17121 } else {
17122 self.func("fa_dequant_kv_ws_bf16")
17123 };
17124 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17125 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17126 let cfg = LaunchConfig {
17127 grid_dim: (nblk.max(1), 1, 1),
17128 block_dim: (256, 1, 1),
17129 shared_mem_bytes: 0,
17130 };
17131 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17132 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17133 let __s_b = self.gpu.stream();
17134 let mut b = __s_b.launch_builder(&f);
17135 b.arg(k)
17136 .arg(v)
17137 .arg(&mut *kw)
17138 .arg(&mut *vw)
17139 .arg(&kdk)
17140 .arg(&kdv)
17141 .arg(&tkvi)
17142 .arg(&ktb)
17143 .arg(&vtb);
17144 unsafe {
17145 b.launch(cfg)?;
17146 }
17147 }
17148 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17156 .map(|v| v != "0")
17157 .unwrap_or(true);
17158 {
17159 let hd_sfx = fa_hd_suffix(head_dim)?;
17160 let f = self.func(&format!(
17161 "fa_prefill_qw{}{hd_sfx}",
17162 if db { "_db" } else { "" }
17163 ));
17164 let shmem = if db {
17165 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17167 } else {
17168 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17169 };
17170 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17171 f.set_attribute(
17172 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17173 shmem as i32,
17174 )?;
17175 let cfg = LaunchConfig {
17176 grid_dim: (
17177 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17178 n_head as u32,
17179 1,
17180 ),
17181 block_dim: (32, 4, 1),
17182 shared_mem_bytes: shmem,
17183 };
17184 let (hd, nh, nhkv, ti, tkvi, cz) = (
17185 head_dim as i32,
17186 n_head as i32,
17187 n_head_kv as i32,
17188 t as i32,
17189 t_kv as i32,
17190 causal as i32,
17191 );
17192 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17193 let __s_b = self.gpu.stream();
17194 let mut b = __s_b.launch_builder(&f);
17195 b.arg(q)
17196 .arg(&*kw)
17197 .arg(&*vw)
17198 .arg(o)
17199 .arg(&hd)
17200 .arg(&nh)
17201 .arg(&nhkv)
17202 .arg(&ti)
17203 .arg(&tkvi)
17204 .arg(&scale)
17205 .arg(&cz)
17206 .arg(&kdk)
17207 .arg(&kdv);
17208 unsafe {
17209 b.launch(cfg)?;
17210 }
17211 }
17212 Ok(())
17213 }
17214
17215 #[allow(clippy::too_many_arguments)]
17231 pub fn fa_prefill_view_ws_w_hd128(
17232 &self,
17233 q: &CudaSlice<f32>,
17234 k: &cudarc::driver::CudaView<u8>,
17235 v: &cudarc::driver::CudaView<u8>,
17236 o: &mut CudaSlice<f32>,
17237 head_dim: usize,
17238 n_head: usize,
17239 n_head_kv: usize,
17240 t: usize,
17241 t_kv: usize,
17242 scale: f32,
17243 causal: bool,
17244 window: usize,
17245 k_tok_bytes: usize,
17246 v_tok_bytes: usize,
17247 ) -> Result<(), Box<dyn std::error::Error>> {
17248 assert_eq!(
17249 head_dim, 128,
17250 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17251 );
17252 if portable_mma_gated() {
17253 return self.sdpa_naive_w_quantized_view(
17254 q,
17255 k,
17256 v,
17257 o,
17258 head_dim,
17259 n_head,
17260 n_head_kv,
17261 t,
17262 t_kv,
17263 scale,
17264 causal,
17265 window,
17266 k_tok_bytes,
17267 v_tok_bytes,
17268 );
17269 }
17270 const BLOCK_Q: usize = 64;
17271 const BK: usize = 32;
17272 let kv_dim_k = n_head_kv * head_dim;
17273 let kv_dim_v = n_head_kv * head_dim;
17274 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17276 let mut guard = self.prime_deqw_ws.lock().unwrap();
17277 let need_grow = match guard.as_ref() {
17278 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17279 None => true,
17280 };
17281 if need_grow {
17282 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17283 let (ck, cv) = guard
17284 .as_ref()
17285 .map(|(a, b)| (a.len(), b.len()))
17286 .unwrap_or((0, 0));
17287 *guard = Some((
17288 self.alloc_u8(grow(ck, k_ws_bytes))?,
17289 self.alloc_u8(grow(cv, v_ws_bytes))?,
17290 ));
17291 }
17292 let (kw, vw) = guard.as_mut().unwrap();
17293 {
17296 let f = self.func("fa_dequant_kv_ws_bf16");
17297 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17298 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17299 let cfg = LaunchConfig {
17300 grid_dim: (nblk.max(1), 1, 1),
17301 block_dim: (256, 1, 1),
17302 shared_mem_bytes: 0,
17303 };
17304 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17305 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17306 let __s_b = self.gpu.stream();
17307 let mut b = __s_b.launch_builder(&f);
17308 b.arg(k)
17309 .arg(v)
17310 .arg(&mut *kw)
17311 .arg(&mut *vw)
17312 .arg(&kdk)
17313 .arg(&kdv)
17314 .arg(&tkvi)
17315 .arg(&ktb)
17316 .arg(&vtb);
17317 unsafe {
17318 b.launch(cfg)?;
17319 }
17320 }
17321 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17323 .map(|v| v != "0")
17324 .unwrap_or(true);
17325 {
17326 let f = self.func(if db {
17327 "fa_prefill_qw_db_w_hd128"
17328 } else {
17329 "fa_prefill_qw_w_hd128"
17330 });
17331 let shmem = if db {
17332 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17333 } else {
17334 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17335 };
17336 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17337 f.set_attribute(
17338 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17339 shmem as i32,
17340 )?;
17341 let cfg = LaunchConfig {
17342 grid_dim: (
17343 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17344 n_head as u32,
17345 1,
17346 ),
17347 block_dim: (32, 4, 1),
17348 shared_mem_bytes: shmem,
17349 };
17350 let (hd, nh, nhkv, ti, tkvi, cz) = (
17351 head_dim as i32,
17352 n_head as i32,
17353 n_head_kv as i32,
17354 t as i32,
17355 t_kv as i32,
17356 causal as i32,
17357 );
17358 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17359 let __s_b = self.gpu.stream();
17360 let mut b = __s_b.launch_builder(&f);
17361 b.arg(q)
17362 .arg(&*kw)
17363 .arg(&*vw)
17364 .arg(o)
17365 .arg(&hd)
17366 .arg(&nh)
17367 .arg(&nhkv)
17368 .arg(&ti)
17369 .arg(&tkvi)
17370 .arg(&scale)
17371 .arg(&cz)
17372 .arg(&kdk)
17373 .arg(&kdv)
17374 .arg(&wnd);
17375 unsafe {
17376 b.launch(cfg)?;
17377 }
17378 }
17379 Ok(())
17380 }
17381
17382 pub fn fa_decode(
17386 &self,
17387 q: &CudaSlice<f32>,
17388 k: &cudarc::driver::CudaView<u8>,
17389 v: &cudarc::driver::CudaView<u8>,
17390 o: &mut CudaSlice<f32>,
17391 head_dim: usize,
17392 n_head: usize,
17393 n_head_kv: usize,
17394 t_kv: usize,
17395 scale: f32,
17396 k_tok_bytes: usize,
17397 v_tok_bytes: usize,
17398 ) -> Result<(), Box<dyn std::error::Error>> {
17399 self.fa_decode_kvmod(
17400 q,
17401 k,
17402 v,
17403 o,
17404 head_dim,
17405 n_head,
17406 n_head_kv,
17407 t_kv,
17408 scale,
17409 k_tok_bytes,
17410 v_tok_bytes,
17411 false,
17412 )
17413 }
17414
17415 #[allow(clippy::too_many_arguments)]
17419 #[allow(clippy::too_many_arguments)]
17423 #[allow(clippy::too_many_arguments)]
17424 fn fa_decode_scalar_unified(
17425 &self,
17426 q: &cudarc::driver::CudaView<f32>,
17427 k: &cudarc::driver::CudaView<u8>,
17428 v: &cudarc::driver::CudaView<u8>,
17429 o: &mut cudarc::driver::CudaViewMut<f32>,
17430 head_dim: usize,
17431 n_head: usize,
17432 n_head_kv: usize,
17433 t_kv_host: usize,
17434 t_kv_dev: Option<&CudaSlice<i32>>,
17435 scale: f32,
17436 n_splits: usize,
17437 split_keys: usize,
17438 k_tok_bytes: usize,
17439 v_tok_bytes: usize,
17440 g: bool,
17441 part_o: &mut CudaSlice<f32>,
17442 part_m: &mut CudaSlice<f32>,
17443 part_l: &mut CudaSlice<f32>,
17444 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
17445 ) -> Result<(), Box<dyn std::error::Error>> {
17446 let f = if g {
17447 self.func_g("fa_decode_f32")
17448 } else {
17449 self.fa_func("fa_decode_f32", head_dim)
17450 };
17451 let cfg = LaunchConfig {
17452 grid_dim: (n_head as u32, n_splits as u32, 1),
17453 block_dim: (head_dim as u32, 1, 1),
17454 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
17455 };
17456 let (hd, nh, nhkv, nsp) = (
17457 head_dim as i32,
17458 n_head as i32,
17459 n_head_kv as i32,
17460 n_splits as i32,
17461 );
17462 let (ktb, vtb, tkvi, ski) = (
17463 k_tok_bytes as i64,
17464 v_tok_bytes as i64,
17465 t_kv_host as i32,
17466 split_keys as i32,
17467 );
17468 let __s_b = self.gpu.stream();
17469 let mut b = __s_b.launch_builder(&f);
17470 match t_kv_dev {
17471 Some(d) => {
17472 b.arg(q)
17473 .arg(k)
17474 .arg(v)
17475 .arg(&mut *part_o)
17476 .arg(&mut *part_m)
17477 .arg(&mut *part_l)
17478 .arg(&hd)
17479 .arg(&nh)
17480 .arg(&nhkv)
17481 .arg(&tkvi)
17482 .arg(d)
17483 .arg(&scale)
17484 .arg(&nsp)
17485 .arg(&ski)
17486 .arg(&ktb)
17487 .arg(&vtb);
17488 unsafe {
17489 b.launch(cfg)?;
17490 }
17491 }
17492 None => {
17493 let null: u64 = 0;
17494 b.arg(q)
17495 .arg(k)
17496 .arg(v)
17497 .arg(&mut *part_o)
17498 .arg(&mut *part_m)
17499 .arg(&mut *part_l)
17500 .arg(&hd)
17501 .arg(&nh)
17502 .arg(&nhkv)
17503 .arg(&tkvi)
17504 .arg(&null)
17505 .arg(&scale)
17506 .arg(&nsp)
17507 .arg(&ski)
17508 .arg(&ktb)
17509 .arg(&vtb);
17510 unsafe {
17511 b.launch(cfg)?;
17512 }
17513 }
17514 }
17515 let cfg2 = LaunchConfig {
17516 grid_dim: (n_head as u32, 1, 1),
17517 block_dim: (head_dim as u32, 1, 1),
17518 shared_mem_bytes: 0,
17519 };
17520 if let Some((oq, od)) = q8_out {
17521 let fc = if g {
17523 self.func_g("fa_decode_combine_q8_1")
17524 } else {
17525 self.fa_func("fa_decode_combine_q8_1", head_dim)
17526 };
17527 let __s_b2 = self.gpu.stream();
17528 let mut b2 = __s_b2.launch_builder(&fc);
17529 b2.arg(&*part_o)
17530 .arg(&*part_m)
17531 .arg(&*part_l)
17532 .arg(oq)
17533 .arg(od)
17534 .arg(&hd)
17535 .arg(&nh)
17536 .arg(&nsp);
17537 unsafe {
17538 b2.launch(cfg2)?;
17539 }
17540 return Ok(());
17541 }
17542 let fc = if g {
17543 self.func_g("fa_decode_combine_f32")
17544 } else {
17545 self.fa_func("fa_decode_combine_f32", head_dim)
17546 };
17547 let __s_b2 = self.gpu.stream();
17548 let mut b2 = __s_b2.launch_builder(&fc);
17549 b2.arg(&*part_o)
17550 .arg(&*part_m)
17551 .arg(&*part_l)
17552 .arg(o)
17553 .arg(&hd)
17554 .arg(&nh)
17555 .arg(&nsp);
17556 unsafe {
17557 b2.launch(cfg2)?;
17558 }
17559 Ok(())
17560 }
17561
17562 pub fn fa_decode_kvmod(
17563 &self,
17564 q: &CudaSlice<f32>,
17565 k: &cudarc::driver::CudaView<u8>,
17566 v: &cudarc::driver::CudaView<u8>,
17567 o: &mut CudaSlice<f32>,
17568 head_dim: usize,
17569 n_head: usize,
17570 n_head_kv: usize,
17571 t_kv: usize,
17572 scale: f32,
17573 k_tok_bytes: usize,
17574 v_tok_bytes: usize,
17575 g: bool,
17576 ) -> Result<(), Box<dyn std::error::Error>> {
17577 let q_view = q.as_view();
17578 let mut o_view = o.as_view_mut();
17579 self.fa_decode_kvmod_view(
17580 &q_view,
17581 k,
17582 v,
17583 &mut o_view,
17584 head_dim,
17585 n_head,
17586 n_head_kv,
17587 t_kv,
17588 scale,
17589 k_tok_bytes,
17590 v_tok_bytes,
17591 g,
17592 )
17593 }
17594
17595 #[allow(clippy::too_many_arguments)]
17600 pub fn fa_decode_kvmod_view(
17601 &self,
17602 q: &cudarc::driver::CudaView<f32>,
17603 k: &cudarc::driver::CudaView<u8>,
17604 v: &cudarc::driver::CudaView<u8>,
17605 o: &mut cudarc::driver::CudaViewMut<f32>,
17606 head_dim: usize,
17607 n_head: usize,
17608 n_head_kv: usize,
17609 t_kv: usize,
17610 scale: f32,
17611 k_tok_bytes: usize,
17612 v_tok_bytes: usize,
17613 g: bool,
17614 ) -> Result<(), Box<dyn std::error::Error>> {
17615 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
17636 if g && head_dim == 256 && !fa_v4_at(t_kv) {
17640 fa_vec = false;
17641 }
17642 let sp = fa_split_keys(t_kv, n_head_kv);
17643 let n_splits = if fa_vec {
17644 ((t_kv + sp - 1) / sp).max(1)
17645 } else {
17646 ((t_kv + 255) / 256).max(1)
17647 };
17648 let o_len = n_head * n_splits * head_dim;
17649 let ml_len = n_head * n_splits;
17650 let mut part_guard = self.fa_part_pool.lock().unwrap();
17651 if part_guard
17652 .as_ref()
17653 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17654 .unwrap_or(true)
17655 {
17656 let old = part_guard.take();
17667 let (co, cm) = old
17668 .as_ref()
17669 .map(|pp| (pp.0.len(), pp.1.len()))
17670 .unwrap_or((0, 0));
17671 if let Some(old) = old {
17672 self.fa_part_retired.lock().unwrap().push(old);
17673 }
17674 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17675 eprintln!(
17676 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17677 co, o_len, cm, ml_len
17678 );
17679 }
17680 *part_guard = Some((
17681 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
17682 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17683 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17684 ));
17685 }
17686 let pg = part_guard.as_mut().unwrap();
17687 self.gpu
17688 .stream()
17689 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
17690 self.gpu
17691 .stream()
17692 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
17693 self.gpu
17694 .stream()
17695 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
17696 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
17697 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
17698 let (hd, nh, nhkv, tkvi, nsp) = (
17699 head_dim as i32,
17700 n_head as i32,
17701 n_head_kv as i32,
17702 t_kv as i32,
17703 n_splits as i32,
17704 );
17705 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17706 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
17710 let fa512_min = fa512_min_tkv();
17715 let deep = fa_vec
17718 && head_dim == 256
17719 && fa_v4_at(t_kv)
17720 && !g
17721 && fa_deep_at(t_kv)
17722 && !matches!(fa_v4_mode(), "noB3" | "stage");
17723 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
17724 let gqa = (n_head / n_head_kv).max(1) as u32;
17727 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
17728 (
17729 fv,
17730 LaunchConfig {
17731 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17732 block_dim: (32, gqa, 1),
17733 shared_mem_bytes: 0,
17734 },
17735 )
17736 } else if fa_vec && head_dim <= 256 {
17737 let gqa = (n_head / n_head_kv).max(1) as u32;
17738 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17749 let smem_tkv = *SMEM_TKV.get_or_init(|| {
17750 std::env::var("MEMRA_FA_SMEM_TKV")
17751 .ok()
17752 .and_then(|v| v.parse().ok())
17753 .unwrap_or_else(|| {
17754 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
17755 })
17756 });
17757 if fa_v4_at(t_kv) && head_dim == 256 {
17758 let v4name = match fa_v4_mode() {
17762 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
17765 _ => "fa_decode_vec_q_v4",
17766 };
17767 let fv = if g {
17768 self.func_g(v4name)
17769 } else {
17770 self.func(v4name)
17771 };
17772 let shmem = (if deep { 12160 } else { 11520 }
17775 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
17776 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17777 fv.set_attribute(
17778 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17779 shmem as i32,
17780 )?;
17781 (
17782 fv,
17783 LaunchConfig {
17784 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17785 block_dim: (32, gqa, 1),
17786 shared_mem_bytes: shmem,
17787 },
17788 )
17789 } else if fa_v3_active(head_dim) {
17790 let fv = if g {
17793 self.func_g("fa_decode_vec_q_v3")
17794 } else {
17795 self.func("fa_decode_vec_q_v3")
17796 };
17797 let shmem = (32 * head_dim * 2) as u32; (
17799 fv,
17800 LaunchConfig {
17801 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17802 block_dim: (32, gqa, 1),
17803 shared_mem_bytes: shmem,
17804 },
17805 )
17806 } else if fa_v2_on() {
17807 let fv = if g {
17811 self.func_g("fa_decode_vec_q_v2")
17812 } else {
17813 self.func("fa_decode_vec_q_v2")
17814 };
17815 let shmem = (2 * 32 * head_dim * 2) as u32; (
17817 fv,
17818 LaunchConfig {
17819 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17820 block_dim: (32, gqa, 1),
17821 shared_mem_bytes: shmem,
17822 },
17823 )
17824 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
17825 {
17826 let fv = if g {
17830 self.func_g("fa_decode_vec_q_smem")
17831 } else {
17832 self.func("fa_decode_vec_q_smem")
17833 };
17834 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
17836 fv.set_attribute(
17837 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17838 shmem as i32,
17839 )?;
17840 (
17841 fv,
17842 LaunchConfig {
17843 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17844 block_dim: (32, gqa, 1),
17845 shared_mem_bytes: shmem,
17846 },
17847 )
17848 } else {
17849 let fv = if g {
17852 self.func_g("fa_decode_vec_q")
17853 } else {
17854 self.func("fa_decode_vec_q")
17855 };
17856 (
17857 fv,
17858 LaunchConfig {
17859 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17860 block_dim: (32, gqa, 1),
17861 shared_mem_bytes: 0,
17862 },
17863 )
17864 }
17865 } else {
17866 return self.fa_decode_scalar_unified(
17869 q,
17870 k,
17871 v,
17872 o,
17873 head_dim,
17874 n_head,
17875 n_head_kv,
17876 t_kv,
17877 None,
17878 scale,
17879 n_splits,
17880 if fa_vec { sp } else { 256 },
17881 k_tok_bytes,
17882 v_tok_bytes,
17883 g,
17884 part_o,
17885 part_m,
17886 part_l,
17887 None,
17888 );
17889 };
17890 let __s_b = self.gpu.stream();
17891 let mut b = __s_b.launch_builder(&f);
17892 b.arg(q)
17893 .arg(k)
17894 .arg(v)
17895 .arg(&mut *part_o)
17896 .arg(&mut *part_m)
17897 .arg(&mut *part_l)
17898 .arg(&hd)
17899 .arg(&nh)
17900 .arg(&nhkv)
17901 .arg(&tkvi)
17902 .arg(&scale)
17903 .arg(&nsp)
17904 .arg(&ktb)
17905 .arg(&vtb);
17906 unsafe {
17907 b.launch(cfg)?;
17908 }
17909 let (fc, cfg2) = (
17912 if g {
17913 self.func_g("fa_decode_combine_f32")
17914 } else {
17915 self.fa_func("fa_decode_combine_f32", head_dim)
17916 },
17917 LaunchConfig {
17918 grid_dim: (n_head as u32, 1, 1),
17919 block_dim: (head_dim as u32, 1, 1),
17920 shared_mem_bytes: 0,
17921 },
17922 );
17923 let __s_b2 = self.gpu.stream();
17924 let mut b2 = __s_b2.launch_builder(&fc);
17925 b2.arg(&*part_o)
17926 .arg(&*part_m)
17927 .arg(&*part_l)
17928 .arg(o)
17929 .arg(&hd)
17930 .arg(&nh)
17931 .arg(&nsp);
17932 unsafe {
17933 b2.launch(cfg2)?;
17934 }
17935 Ok(())
17936 }
17937
17938 #[allow(clippy::too_many_arguments)]
17949 pub fn fa_decode_batch_seqs_v4(
17950 &self,
17951 q: &CudaSlice<f32>,
17952 kv_ptrs: &cudarc::driver::CudaView<u64>,
17953 pos_seq: &CudaSlice<i32>,
17954 o: &mut CudaSlice<f32>,
17955 head_dim: usize,
17956 n_head: usize,
17957 n_head_kv: usize,
17958 b_n: usize,
17959 t_kv_max: usize,
17960 scale: f32,
17961 split_keys: usize,
17962 k_tok_bytes: usize,
17963 v_tok_bytes: usize,
17964 ) -> Result<(), Box<dyn std::error::Error>> {
17965 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
17966 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
17967 let o_len = b_n * n_head * n_splits_max * head_dim;
17968 let ml_len = b_n * n_head * n_splits_max;
17969 let mut part_guard = self.fa_part_pool.lock().unwrap();
17970 if part_guard
17971 .as_ref()
17972 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17973 .unwrap_or(true)
17974 {
17975 let old = part_guard.take();
17986 let (co, cm) = old
17987 .as_ref()
17988 .map(|pp| (pp.0.len(), pp.1.len()))
17989 .unwrap_or((0, 0));
17990 if let Some(old) = old {
17991 self.fa_part_retired.lock().unwrap().push(old);
17992 }
17993 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17994 eprintln!(
17995 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17996 co, o_len, cm, ml_len
17997 );
17998 }
17999 *part_guard = Some((
18000 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18001 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18002 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18003 ));
18004 }
18005 let pg = part_guard.as_mut().unwrap();
18006 self.gpu
18007 .stream()
18008 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18009 self.gpu
18010 .stream()
18011 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18012 self.gpu
18013 .stream()
18014 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18015 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18016 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18017 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
18018 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18019 let gqa = (n_head / n_head_kv).max(1) as u32;
18020 let f = self.func("fa_decode_vec_q_seqs_v4");
18021 let shmem = (11520 + 32 * head_dim * 2) as u32;
18023 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18024 f.set_attribute(
18025 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18026 shmem as i32,
18027 )?;
18028 let cfg = LaunchConfig {
18029 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
18030 block_dim: (32, gqa, 1),
18031 shared_mem_bytes: shmem,
18032 };
18033 {
18034 let __s_b = self.gpu.stream();
18035 let mut b = __s_b.launch_builder(&f);
18036 b.arg(q)
18037 .arg(kv_ptrs)
18038 .arg(pos_seq)
18039 .arg(&mut *part_o)
18040 .arg(&mut *part_m)
18041 .arg(&mut *part_l)
18042 .arg(&hd)
18043 .arg(&nh)
18044 .arg(&nhkv)
18045 .arg(&scale)
18046 .arg(&nspm)
18047 .arg(&spk)
18048 .arg(&ktb)
18049 .arg(&vtb);
18050 unsafe {
18051 b.launch(cfg)?;
18052 }
18053 }
18054 let fc = self.func("fa_decode_combine_seqs");
18055 let cfg2 = LaunchConfig {
18056 grid_dim: (n_head as u32, b_n as u32, 1),
18057 block_dim: (head_dim as u32, 1, 1),
18058 shared_mem_bytes: 0,
18059 };
18060 let __s_b2 = self.gpu.stream();
18061 let mut b2 = __s_b2.launch_builder(&fc);
18062 b2.arg(&*part_o)
18063 .arg(&*part_m)
18064 .arg(&*part_l)
18065 .arg(o)
18066 .arg(&hd)
18067 .arg(&nh)
18068 .arg(pos_seq)
18069 .arg(&nspm)
18070 .arg(&spk);
18071 unsafe {
18072 b2.launch(cfg2)?;
18073 }
18074 Ok(())
18075 }
18076
18077 #[allow(clippy::too_many_arguments)]
18084 pub fn append_kv_quantized_seqs(
18085 &self,
18086 k_rows: &CudaSlice<f32>,
18087 v_rows: &CudaSlice<f32>,
18088 kv_ptrs: &cudarc::driver::CudaView<u64>,
18089 pos_seq: &CudaSlice<i32>,
18090 b_n: usize,
18091 kv_dim_k: usize,
18092 kv_dim_v: usize,
18093 k_tok_bytes: usize,
18094 v_tok_bytes: usize,
18095 ) -> Result<(), Box<dyn std::error::Error>> {
18096 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
18097 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
18098 let cfg = LaunchConfig {
18099 grid_dim: (nblk, b_n as u32, 1),
18100 block_dim: (32, 1, 1),
18101 shared_mem_bytes: 0,
18102 };
18103 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
18104 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18105 let __s_b = self.gpu.stream();
18106 let mut b = __s_b.launch_builder(&f);
18107 b.arg(k_rows)
18108 .arg(v_rows)
18109 .arg(kv_ptrs)
18110 .arg(pos_seq)
18111 .arg(&kdk)
18112 .arg(&kdv)
18113 .arg(&ktb)
18114 .arg(&vtb);
18115 unsafe {
18116 b.launch(cfg)?;
18117 }
18118 Ok(())
18119 }
18120
18121 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
18127 std::env::var("MEMRA_NO_FA_VEC").is_err()
18128 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
18129 && base_len + 1 >= fa_vec_min_tkv()
18130 && head_dim <= 256
18131 && head_dim % 32 == 0
18132 }
18133
18134 #[allow(clippy::too_many_arguments)]
18143 pub fn fa_decode_rows(
18144 &self,
18145 q: &CudaSlice<f32>,
18146 k: &cudarc::driver::CudaView<u8>,
18147 v: &cudarc::driver::CudaView<u8>,
18148 o: &mut CudaSlice<f32>,
18149 head_dim: usize,
18150 n_head: usize,
18151 n_head_kv: usize,
18152 base_len: usize,
18153 t: usize,
18154 scale: f32,
18155 k_tok_bytes: usize,
18156 v_tok_bytes: usize,
18157 base_dev: Option<(&CudaSlice<i32>, i32)>,
18161 kv_shared: bool,
18164 g: bool,
18168 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18171 ) -> Result<(), Box<dyn std::error::Error>> {
18172 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
18173 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
18180 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18181 let v = *SP512.get_or_init(|| {
18184 std::env::var("MEMRA_FA_SP512")
18185 .ok()
18186 .and_then(|x| x.parse().ok())
18187 .unwrap_or(0)
18188 });
18189 sp = if v >= 8 {
18190 v
18191 } else {
18192 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18193 };
18194 }
18195 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18196 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18197 let gqa = (n_head / n_head_kv).max(1) as u32;
18198 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
18209 groups.push((0, t, sp));
18210 } else {
18211 let mut r0 = 0usize;
18212 while r0 < t {
18213 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
18214 let mut r1 = r0 + 1;
18215 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18216 r1 += 1;
18217 }
18218 groups.push((r0, r1 - r0, sp_g));
18219 r0 = r1;
18220 }
18221 }
18222 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18226 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18227 std::env::var("MEMRA_FA_SMEM_TKV")
18228 .ok()
18229 .and_then(|v| v.parse().ok())
18230 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18231 });
18232 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18233 let v3 = fa_v3_active(head_dim);
18234 let smem_rows =
18235 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18236 let _ = kv_shared;
18241 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18244 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18258 let tb512 = head_dim == 512
18260 && sp <= 32
18261 && n_head / n_head_kv.max(1) <= 16
18262 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18263 let fname = if tb512 {
18264 "fa_decode_vec_q_rows_v4_512_tb"
18265 } else if i2 {
18266 "fa_decode_vec_q_rows_dpl16_i2"
18267 } else if head_dim == 512 {
18268 "fa_decode_vec_q_rows_dpl16"
18269 }
18270 else if v4 {
18272 "fa_decode_vec_q_rows_v4"
18273 } else if v3 {
18274 "fa_decode_vec_q_rows_v3"
18275 } else if fa_v2_on() {
18276 "fa_decode_vec_q_rows_v2"
18277 } else if smem_rows {
18278 "fa_decode_vec_q_rows_smem"
18279 } else {
18280 "fa_decode_vec_q_rows"
18281 };
18282 let f = if head_dim == 512 {
18283 self.fa_func(fname, head_dim)
18284 } else if g {
18285 self.func_g(if smem_rows {
18293 "fa_decode_vec_q_rows"
18294 } else {
18295 fname
18296 })
18297 } else {
18298 self.func(fname)
18299 };
18300 let shmem = if tb512 {
18301 let gk = Self::gkv_on();
18303 let sh =
18304 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18305 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18306 f.set_attribute(
18307 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18308 sh as i32,
18309 )?;
18310 sh
18311 } else if v4 || v3 || smem_rows || fa_v2_on() {
18312 let sh = (if v4 {
18314 11520 + 32 * head_dim * if g { 1 } else { 2 }
18315 } else if v3 {
18316 32 * head_dim * 2
18317 } else {
18318 2 * 32 * head_dim * 2
18319 }) as u32;
18320 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18321 f.set_attribute(
18322 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18323 sh as i32,
18324 )?;
18325 sh
18326 } else {
18327 0
18328 };
18329 for &(r0, t_g, sp_g) in &groups {
18333 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18334 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18335 let base_i = (base_len + r0) as i32;
18336 let o_len = t_g * n_head * n_splits_g * head_dim;
18337 let ml_len = t_g * n_head * n_splits_g;
18338 let mut part_guard = self.fa_part_pool.lock().unwrap();
18339 if part_guard
18340 .as_ref()
18341 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18342 .unwrap_or(true)
18343 {
18344 let old = part_guard.take();
18355 let (co, cm) = old
18356 .as_ref()
18357 .map(|pp| (pp.0.len(), pp.1.len()))
18358 .unwrap_or((0, 0));
18359 if let Some(old) = old {
18360 self.fa_part_retired.lock().unwrap().push(old);
18361 }
18362 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18363 eprintln!(
18364 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18365 co, o_len, cm, ml_len
18366 );
18367 }
18368 *part_guard = Some((
18369 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18370 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18371 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18372 ));
18373 }
18374 let pg = part_guard.as_mut().unwrap();
18375 self.gpu
18376 .stream()
18377 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18378 self.gpu
18379 .stream()
18380 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18381 self.gpu
18382 .stream()
18383 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18384 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18385 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18386 let qv = self.view(q, t * n_head * head_dim);
18387 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18388 let cfg = LaunchConfig {
18389 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
18390 block_dim: (32, gqa, 1),
18391 shared_mem_bytes: shmem,
18392 };
18393 {
18394 let __s_b = self.gpu.stream();
18395 let mut b = __s_b.launch_builder(&f);
18396 if tb512 {
18397 let (bd, plus) =
18399 base_dev.expect("hd512 rows twin requires a device base counter");
18400 let plus_g = plus + r0 as i32;
18401 let nr = t_g as i32;
18402 if Self::pdl_on() && Self::pdl_wb_on() {
18403 use cudarc::driver::{DevicePtr, DevicePtrMut};
18405 let s = &self.gpu.stream();
18406 let (pq, _b0) = q_g.device_ptr(s);
18407 let (pk, _b1) = k.device_ptr(s);
18408 let (pv, _b2) = v.device_ptr(s);
18409 let (po, _b3) = part_o.device_ptr_mut(s);
18410 let (pm, _b4) = part_m.device_ptr_mut(s);
18411 let (pl, _b5) = part_l.device_ptr_mut(s);
18412 let (pb, _b6) = bd.device_ptr(s);
18413 let mut ps = [
18414 &pq as *const _ as *mut std::ffi::c_void,
18415 &pk as *const _ as *mut _,
18416 &pv as *const _ as *mut _,
18417 &po as *const _ as *mut _,
18418 &pm as *const _ as *mut _,
18419 &pl as *const _ as *mut _,
18420 &hd as *const _ as *mut _,
18421 &nh as *const _ as *mut _,
18422 &nhkv as *const _ as *mut _,
18423 &pb as *const _ as *mut _,
18424 &plus_g as *const _ as *mut _,
18425 &scale as *const _ as *mut _,
18426 &nspm as *const _ as *mut _,
18427 &spk as *const _ as *mut _,
18428 &ktb as *const _ as *mut _,
18429 &vtb as *const _ as *mut _,
18430 &nr as *const _ as *mut _,
18431 ];
18432 unsafe {
18433 self.launch_pdl_flash(
18434 Self::gkv_on(),
18435 "fa_decode_vec_q_rows_v4_512_tb",
18436 (n_head_kv as u32, n_splits_g as u32, 1),
18437 (32, gqa, 1),
18438 shmem,
18439 &mut ps,
18440 )?;
18441 }
18442 } else {
18443 let cfg_tb = LaunchConfig {
18444 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
18445 block_dim: (32, gqa, 1),
18446 shared_mem_bytes: shmem,
18447 };
18448 b.arg(&q_g)
18449 .arg(k)
18450 .arg(v)
18451 .arg(&mut *part_o)
18452 .arg(&mut *part_m)
18453 .arg(&mut *part_l)
18454 .arg(&hd)
18455 .arg(&nh)
18456 .arg(&nhkv)
18457 .arg(bd)
18458 .arg(&plus_g)
18459 .arg(&scale)
18460 .arg(&nspm)
18461 .arg(&spk)
18462 .arg(&ktb)
18463 .arg(&vtb)
18464 .arg(&nr);
18465 unsafe {
18466 b.launch(cfg_tb)?;
18467 }
18468 }
18469 } else if head_dim == 512 {
18470 let (bd, plus) =
18471 base_dev.expect("hd512 rows twin requires a device base counter");
18472 let plus_g = plus + r0 as i32;
18473 b.arg(&q_g)
18474 .arg(k)
18475 .arg(v)
18476 .arg(&mut *part_o)
18477 .arg(&mut *part_m)
18478 .arg(&mut *part_l)
18479 .arg(&hd)
18480 .arg(&nh)
18481 .arg(&nhkv)
18482 .arg(bd)
18483 .arg(&plus_g)
18484 .arg(&scale)
18485 .arg(&nspm)
18486 .arg(&spk)
18487 .arg(&ktb)
18488 .arg(&vtb);
18489 unsafe {
18490 b.launch(cfg)?;
18491 }
18492 } else {
18493 b.arg(&q_g)
18494 .arg(k)
18495 .arg(v)
18496 .arg(&mut *part_o)
18497 .arg(&mut *part_m)
18498 .arg(&mut *part_l)
18499 .arg(&hd)
18500 .arg(&nh)
18501 .arg(&nhkv)
18502 .arg(&base_i)
18503 .arg(&scale)
18504 .arg(&nspm)
18505 .arg(&spk)
18506 .arg(&ktb)
18507 .arg(&vtb);
18508 unsafe {
18509 b.launch(cfg)?;
18510 }
18511 }
18512 }
18513 let cfg2 = LaunchConfig {
18514 grid_dim: (n_head as u32, t_g as u32, 1),
18515 block_dim: (head_dim as u32, 1, 1),
18516 shared_mem_bytes: 0,
18517 };
18518 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18519 if head_dim == 512 {
18520 let (bd, plus) = base_dev.unwrap();
18523 let plus_g = plus + r0 as i32;
18524 if let Some((oq, od)) = q8_out.as_mut() {
18525 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
18527 if Self::pdl_on() && Self::pdl_wb_on() {
18528 use cudarc::driver::{DevicePtr, DevicePtrMut};
18530 let s = &self.gpu.stream();
18531 let (po, _g0) = part_o.device_ptr(s);
18532 let (pm, _g1) = part_m.device_ptr(s);
18533 let (pl, _g2) = part_l.device_ptr(s);
18534 let (pq, _g3) = oq.device_ptr_mut(s);
18535 let (pd, _g4) = od.device_ptr_mut(s);
18536 let (pb, _g5) = bd.device_ptr(s);
18537 let mut ps = [
18538 &po as *const _ as *mut std::ffi::c_void,
18539 &pm as *const _ as *mut _,
18540 &pl as *const _ as *mut _,
18541 &pq as *const _ as *mut _,
18542 &pd as *const _ as *mut _,
18543 &hd as *const _ as *mut _,
18544 &nh as *const _ as *mut _,
18545 &pb as *const _ as *mut _,
18546 &plus_g as *const _ as *mut _,
18547 &nspm as *const _ as *mut _,
18548 &spk as *const _ as *mut _,
18549 ];
18550 unsafe {
18551 self.launch_pdl_flash(
18552 Self::gkv_on(),
18553 "fa_decode_combine_rows_dc_q8_1",
18554 cfg2.grid_dim,
18555 cfg2.block_dim,
18556 0,
18557 &mut ps,
18558 )?;
18559 }
18560 continue;
18561 }
18562 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
18563 let __s_b2 = self.gpu.stream();
18564 let mut b2 = __s_b2.launch_builder(&fc);
18565 b2.arg(&*part_o)
18566 .arg(&*part_m)
18567 .arg(&*part_l)
18568 .arg(&mut **oq)
18569 .arg(&mut **od)
18570 .arg(&hd)
18571 .arg(&nh)
18572 .arg(bd)
18573 .arg(&plus_g)
18574 .arg(&nspm)
18575 .arg(&spk);
18576 unsafe {
18577 b2.launch(cfg2)?;
18578 }
18579 continue;
18580 }
18581 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
18582 let __s_b2 = self.gpu.stream();
18583 let mut b2 = __s_b2.launch_builder(&fc);
18584 b2.arg(&*part_o)
18585 .arg(&*part_m)
18586 .arg(&*part_l)
18587 .arg(&mut o_g)
18588 .arg(&hd)
18589 .arg(&nh)
18590 .arg(bd)
18591 .arg(&plus_g)
18592 .arg(&nspm)
18593 .arg(&spk);
18594 unsafe {
18595 b2.launch(cfg2)?;
18596 }
18597 } else {
18598 assert!(
18601 q8_out.is_none(),
18602 "rows q8 emit requires the hd512 dc combine"
18603 );
18604 let fc = self.func("fa_decode_combine_rows");
18605 let __s_b2 = self.gpu.stream();
18606 let mut b2 = __s_b2.launch_builder(&fc);
18607 b2.arg(&*part_o)
18608 .arg(&*part_m)
18609 .arg(&*part_l)
18610 .arg(&mut o_g)
18611 .arg(&hd)
18612 .arg(&nh)
18613 .arg(&base_i)
18614 .arg(&nspm)
18615 .arg(&spk);
18616 unsafe {
18617 b2.launch(cfg2)?;
18618 }
18619 }
18620 }
18621 Ok(())
18622 }
18623
18624 #[allow(clippy::too_many_arguments)]
18628 pub fn fa_decode_rows_w(
18629 &self,
18630 q: &CudaSlice<f32>,
18631 k: &cudarc::driver::CudaView<u8>,
18632 v: &cudarc::driver::CudaView<u8>,
18633 o: &mut CudaSlice<f32>,
18634 head_dim: usize,
18635 n_head: usize,
18636 n_head_kv: usize,
18637 base_dev: &CudaSlice<i32>,
18638 base_plus: i32,
18639 t: usize,
18640 scale: f32,
18641 window: usize,
18642 k_tok_bytes: usize,
18643 v_tok_bytes: usize,
18644 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18645 ) -> Result<(), Box<dyn std::error::Error>> {
18646 debug_assert!(head_dim == 256);
18651 let sp = {
18659 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18660 let v = *SPW.get_or_init(|| {
18661 std::env::var("MEMRA_FA_SPW")
18662 .ok()
18663 .and_then(|x| x.parse().ok())
18664 .unwrap_or(0)
18665 });
18666 if v >= 8 {
18667 v
18668 } else {
18669 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18670 }
18671 };
18672 let n_splits_max = (window + sp - 1) / sp;
18673 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18674 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
18675 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18676 let gqa = (n_head / n_head_kv).max(1) as u32;
18677 let o_len = t * n_head * n_splits_max * head_dim;
18678 let ml_len = t * n_head * n_splits_max;
18679 let mut part_guard = self.fa_part_pool.lock().unwrap();
18680 if part_guard
18681 .as_ref()
18682 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18683 .unwrap_or(true)
18684 {
18685 let old = part_guard.take();
18696 let (co, cm) = old
18697 .as_ref()
18698 .map(|pp| (pp.0.len(), pp.1.len()))
18699 .unwrap_or((0, 0));
18700 if let Some(old) = old {
18701 self.fa_part_retired.lock().unwrap().push(old);
18702 }
18703 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18704 eprintln!(
18705 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18706 co, o_len, cm, ml_len
18707 );
18708 }
18709 *part_guard = Some((
18710 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18711 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18712 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18713 ));
18714 }
18715 let pg = part_guard.as_mut().unwrap();
18716 self.gpu
18717 .stream()
18718 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18719 self.gpu
18720 .stream()
18721 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18722 self.gpu
18723 .stream()
18724 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18725 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18726 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18732 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
18733 std::env::var("MEMRA_FA_SMEM_TKV")
18734 .ok()
18735 .and_then(|v| v.parse().ok())
18736 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18737 });
18738 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18744 let wg = Self::wkv_on();
18749 let sp2 =
18752 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
18753 if sp2 {
18754 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18755 if Self::pdl_on() && Self::pdl_wb_on() {
18756 use cudarc::driver::{DevicePtr, DevicePtrMut};
18758 let s = &self.gpu.stream();
18759 let (pq, _b0) = q.device_ptr(s);
18760 let (pk, _b1) = k.device_ptr(s);
18761 let (pv, _b2) = v.device_ptr(s);
18762 let (po, _b3) = part_o.device_ptr_mut(s);
18763 let (pm, _b4) = part_m.device_ptr_mut(s);
18764 let (pl, _b5) = part_l.device_ptr_mut(s);
18765 let (pb, _b6) = base_dev.device_ptr(s);
18766 let mut ps = [
18767 &pq as *const _ as *mut std::ffi::c_void,
18768 &pk as *const _ as *mut _,
18769 &pv as *const _ as *mut _,
18770 &po as *const _ as *mut _,
18771 &pm as *const _ as *mut _,
18772 &pl as *const _ as *mut _,
18773 &hd as *const _ as *mut _,
18774 &nh as *const _ as *mut _,
18775 &nhkv as *const _ as *mut _,
18776 &pb as *const _ as *mut _,
18777 &base_plus as *const _ as *mut _,
18778 &scale as *const _ as *mut _,
18779 &nspm as *const _ as *mut _,
18780 &spk as *const _ as *mut _,
18781 &ktb as *const _ as *mut _,
18782 &vtb as *const _ as *mut _,
18783 &wini as *const _ as *mut _,
18784 ];
18785 unsafe {
18786 self.launch_pdl_flash(
18787 wg,
18788 "fa_decode_vec_q_rows_v4_w_sp",
18789 (n_head_kv as u32, n_splits_max as u32, t as u32),
18790 (32, gqa + 1, 1),
18791 sh,
18792 &mut ps,
18793 )?;
18794 }
18795 } else {
18796 let f = if wg {
18797 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
18798 } else {
18799 self.func("fa_decode_vec_q_rows_v4_w_sp")
18800 };
18801 f.set_attribute(
18802 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18803 sh as i32,
18804 )?;
18805 let cfg = LaunchConfig {
18806 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18807 block_dim: (32, gqa + 1, 1),
18808 shared_mem_bytes: sh,
18809 };
18810 let __s_b = self.gpu.stream();
18811 let mut b = __s_b.launch_builder(&f);
18812 b.arg(q)
18813 .arg(k)
18814 .arg(v)
18815 .arg(&mut *part_o)
18816 .arg(&mut *part_m)
18817 .arg(&mut *part_l)
18818 .arg(&hd)
18819 .arg(&nh)
18820 .arg(&nhkv)
18821 .arg(base_dev)
18822 .arg(&base_plus)
18823 .arg(&scale)
18824 .arg(&nspm)
18825 .arg(&spk)
18826 .arg(&ktb)
18827 .arg(&vtb)
18828 .arg(&wini);
18829 unsafe {
18830 b.launch(cfg)?;
18831 }
18832 }
18833 } else {
18834 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
18835 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18837 use cudarc::driver::{DevicePtr, DevicePtrMut};
18838 let s = &self.gpu.stream();
18839 let (pq, _b0) = q.device_ptr(s);
18840 let (pk, _b1) = k.device_ptr(s);
18841 let (pv, _b2) = v.device_ptr(s);
18842 let (po, _b3) = part_o.device_ptr_mut(s);
18843 let (pm, _b4) = part_m.device_ptr_mut(s);
18844 let (pl, _b5) = part_l.device_ptr_mut(s);
18845 let (pb, _b6) = base_dev.device_ptr(s);
18846 let mut ps = [
18847 &pq as *const _ as *mut std::ffi::c_void,
18848 &pk as *const _ as *mut _,
18849 &pv as *const _ as *mut _,
18850 &po as *const _ as *mut _,
18851 &pm as *const _ as *mut _,
18852 &pl as *const _ as *mut _,
18853 &hd as *const _ as *mut _,
18854 &nh as *const _ as *mut _,
18855 &nhkv as *const _ as *mut _,
18856 &pb as *const _ as *mut _,
18857 &base_plus as *const _ as *mut _,
18858 &scale as *const _ as *mut _,
18859 &nspm as *const _ as *mut _,
18860 &spk as *const _ as *mut _,
18861 &ktb as *const _ as *mut _,
18862 &vtb as *const _ as *mut _,
18863 &wini as *const _ as *mut _,
18864 ];
18865 unsafe {
18866 self.launch_pdl_flash(
18867 wg,
18868 "fa_decode_vec_q_rows_v4_w",
18869 (n_head_kv as u32, n_splits_max as u32, t as u32),
18870 (32, gqa, 1),
18871 sh,
18872 &mut ps,
18873 )?;
18874 }
18875 } else {
18876 let pick = |name: &str| {
18877 if wg {
18878 self.func_g(name)
18879 } else {
18880 self.func(name)
18881 }
18882 };
18883 let (f, sh) = if fa_v4_at(window) {
18884 let f = pick("fa_decode_vec_q_rows_v4_w");
18885 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
18886 } else if smem_tkv > 0 && window >= smem_tkv {
18887 (
18890 pick("fa_decode_vec_q_rows_smem_w"),
18891 (2 * 32 * head_dim * 2) as u32,
18892 )
18893 } else {
18894 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
18895 };
18896 f.set_attribute(
18897 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18898 sh as i32,
18899 )?;
18900 let cfg = LaunchConfig {
18901 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18902 block_dim: (32, gqa, 1),
18903 shared_mem_bytes: sh,
18904 };
18905 let __s_b = self.gpu.stream();
18906 let mut b = __s_b.launch_builder(&f);
18907 b.arg(q)
18908 .arg(k)
18909 .arg(v)
18910 .arg(&mut *part_o)
18911 .arg(&mut *part_m)
18912 .arg(&mut *part_l)
18913 .arg(&hd)
18914 .arg(&nh)
18915 .arg(&nhkv)
18916 .arg(base_dev)
18917 .arg(&base_plus)
18918 .arg(&scale)
18919 .arg(&nspm)
18920 .arg(&spk)
18921 .arg(&ktb)
18922 .arg(&vtb)
18923 .arg(&wini);
18924 unsafe {
18925 b.launch(cfg)?;
18926 }
18927 }
18928 }
18929 let cfg2 = LaunchConfig {
18930 grid_dim: (n_head as u32, t as u32, 1),
18931 block_dim: (head_dim as u32, 1, 1),
18932 shared_mem_bytes: 0,
18933 };
18934 if let Some((oq, od)) = q8_out {
18935 if Self::pdl_on() && Self::pdl_wb_on() {
18938 use cudarc::driver::{DevicePtr, DevicePtrMut};
18940 let s = &self.gpu.stream();
18941 let (po, _g0) = part_o.device_ptr(s);
18942 let (pm, _g1) = part_m.device_ptr(s);
18943 let (pl, _g2) = part_l.device_ptr(s);
18944 let (pq, _g3) = oq.device_ptr_mut(s);
18945 let (pd, _g4) = od.device_ptr_mut(s);
18946 let mut ps = [
18947 &po as *const _ as *mut std::ffi::c_void,
18948 &pm as *const _ as *mut _,
18949 &pl as *const _ as *mut _,
18950 &pq as *const _ as *mut _,
18951 &pd as *const _ as *mut _,
18952 &hd as *const _ as *mut _,
18953 &nh as *const _ as *mut _,
18954 &nspm as *const _ as *mut _,
18955 &spk as *const _ as *mut _,
18956 &wini as *const _ as *mut _,
18957 ];
18958 unsafe {
18959 self.launch_pdl_flash(
18960 wg,
18961 "fa_decode_combine_rows_w_q8_1",
18962 cfg2.grid_dim,
18963 cfg2.block_dim,
18964 0,
18965 &mut ps,
18966 )?;
18967 }
18968 return Ok(());
18969 }
18970 let fc = if wg {
18971 self.func_g("fa_decode_combine_rows_w_q8_1")
18972 } else {
18973 self.func("fa_decode_combine_rows_w_q8_1")
18974 };
18975 let __s_b2 = self.gpu.stream();
18976 let mut b2 = __s_b2.launch_builder(&fc);
18977 b2.arg(&*part_o)
18978 .arg(&*part_m)
18979 .arg(&*part_l)
18980 .arg(oq)
18981 .arg(od)
18982 .arg(&hd)
18983 .arg(&nh)
18984 .arg(&nspm)
18985 .arg(&spk)
18986 .arg(&wini);
18987 unsafe {
18988 b2.launch(cfg2)?;
18989 }
18990 return Ok(());
18991 }
18992 let fc = if wg {
18993 self.func_g("fa_decode_combine_rows_w")
18994 } else {
18995 self.func("fa_decode_combine_rows_w")
18996 };
18997 let __s_b2 = self.gpu.stream();
18998 let mut b2 = __s_b2.launch_builder(&fc);
18999 b2.arg(&*part_o)
19000 .arg(&*part_m)
19001 .arg(&*part_l)
19002 .arg(o)
19003 .arg(&hd)
19004 .arg(&nh)
19005 .arg(&nspm)
19006 .arg(&spk)
19007 .arg(&wini);
19008 unsafe {
19009 b2.launch(cfg2)?;
19010 }
19011 Ok(())
19012 }
19013
19014 #[allow(clippy::too_many_arguments)]
19020 pub fn fa_decode_rows_dc(
19021 &self,
19022 q: &CudaSlice<f32>,
19023 k: &cudarc::driver::CudaView<u8>,
19024 v: &cudarc::driver::CudaView<u8>,
19025 o: &mut CudaSlice<f32>,
19026 head_dim: usize,
19027 n_head: usize,
19028 n_head_kv: usize,
19029 base_dev: &CudaSlice<i32>,
19030 t_kv_upper: usize,
19031 t: usize,
19032 scale: f32,
19033 k_tok_bytes: usize,
19034 v_tok_bytes: usize,
19035 base_plus: i32,
19036 g: bool,
19037 ) -> Result<(), Box<dyn std::error::Error>> {
19038 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
19039 assert!(
19040 v4 || fa_v3_active(head_dim),
19041 "stream fa rows requires the v3 or v4 lane"
19042 );
19043 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
19044 if v4 {
19045 let sp = fa_split_keys(t_kv_upper, n_head_kv);
19046 let n_splits_max = (t_kv_upper + sp - 1) / sp;
19047 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19048 let (nspm, spk) = (n_splits_max as i32, sp as i32);
19049 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19050 let gqa = (n_head / n_head_kv).max(1) as u32;
19051 let o_len = t * n_head * n_splits_max * head_dim;
19052 let ml_len = t * n_head * n_splits_max;
19053 let mut part_guard = self.fa_part_pool.lock().unwrap();
19054 if part_guard
19055 .as_ref()
19056 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19057 .unwrap_or(true)
19058 {
19059 let old = part_guard.take();
19070 let (co, cm) = old
19071 .as_ref()
19072 .map(|pp| (pp.0.len(), pp.1.len()))
19073 .unwrap_or((0, 0));
19074 if let Some(old) = old {
19075 self.fa_part_retired.lock().unwrap().push(old);
19076 }
19077 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19078 eprintln!(
19079 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19080 co, o_len, cm, ml_len
19081 );
19082 }
19083 *part_guard = Some((
19084 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19085 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19086 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19087 ));
19088 }
19089 let pg = part_guard.as_mut().unwrap();
19090 self.gpu
19091 .stream()
19092 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19093 self.gpu
19094 .stream()
19095 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19096 self.gpu
19097 .stream()
19098 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19099 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19100 let f = if g {
19101 self.func_g("fa_decode_vec_q_rows_v4_dc")
19102 } else {
19103 self.func("fa_decode_vec_q_rows_v4_dc")
19104 };
19105 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19106 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19107 f.set_attribute(
19108 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19109 sh as i32,
19110 )?;
19111 let cfg = LaunchConfig {
19112 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19113 block_dim: (32, gqa, 1),
19114 shared_mem_bytes: sh,
19115 };
19116 let __s_b = self.gpu.stream();
19117 let mut b = __s_b.launch_builder(&f);
19118 b.arg(q)
19119 .arg(k)
19120 .arg(v)
19121 .arg(&mut *part_o)
19122 .arg(&mut *part_m)
19123 .arg(&mut *part_l)
19124 .arg(&hd)
19125 .arg(&nh)
19126 .arg(&nhkv)
19127 .arg(base_dev)
19128 .arg(&base_plus)
19129 .arg(&scale)
19130 .arg(&nspm)
19131 .arg(&spk)
19132 .arg(&ktb)
19133 .arg(&vtb);
19134 unsafe {
19135 b.launch(cfg)?;
19136 }
19137 let fc = self.func("fa_decode_combine_rows_dc");
19138 let cfg2 = LaunchConfig {
19139 grid_dim: (n_head as u32, t as u32, 1),
19140 block_dim: (head_dim as u32, 1, 1),
19141 shared_mem_bytes: 0,
19142 };
19143 let __s_b2 = self.gpu.stream();
19144 let mut b2 = __s_b2.launch_builder(&fc);
19145 b2.arg(&*part_o)
19146 .arg(&*part_m)
19147 .arg(&*part_l)
19148 .arg(o)
19149 .arg(&hd)
19150 .arg(&nh)
19151 .arg(base_dev)
19152 .arg(&base_plus)
19153 .arg(&nspm)
19154 .arg(&spk);
19155 unsafe {
19156 b2.launch(cfg2)?;
19157 }
19158 return Ok(());
19159 }
19160 let sp = fa_split_keys(t_kv_upper, n_head_kv);
19161 let n_splits_max = (t_kv_upper + sp - 1) / sp;
19162 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19163 let (nspm, spk) = (n_splits_max as i32, sp as i32);
19164 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19165 let gqa = (n_head / n_head_kv).max(1) as u32;
19166 let o_len = t * n_head * n_splits_max * head_dim;
19167 let ml_len = t * n_head * n_splits_max;
19168 let mut part_guard = self.fa_part_pool.lock().unwrap();
19169 if part_guard
19170 .as_ref()
19171 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19172 .unwrap_or(true)
19173 {
19174 let old = part_guard.take();
19185 let (co, cm) = old
19186 .as_ref()
19187 .map(|pp| (pp.0.len(), pp.1.len()))
19188 .unwrap_or((0, 0));
19189 if let Some(old) = old {
19190 self.fa_part_retired.lock().unwrap().push(old);
19191 }
19192 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19193 eprintln!(
19194 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19195 co, o_len, cm, ml_len
19196 );
19197 }
19198 *part_guard = Some((
19199 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19200 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19201 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19202 ));
19203 }
19204 let pg = part_guard.as_mut().unwrap();
19205 self.gpu
19206 .stream()
19207 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19208 self.gpu
19209 .stream()
19210 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19211 self.gpu
19212 .stream()
19213 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19214 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19215 let f = self.func("fa_decode_vec_q_rows_v3_dc");
19216 let sh = (32 * head_dim * 2) as u32;
19217 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19218 f.set_attribute(
19219 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19220 sh as i32,
19221 )?;
19222 let cfg = LaunchConfig {
19223 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19224 block_dim: (32, gqa, 1),
19225 shared_mem_bytes: sh,
19226 };
19227 let __s_b = self.gpu.stream();
19228 let mut b = __s_b.launch_builder(&f);
19229 b.arg(q)
19230 .arg(k)
19231 .arg(v)
19232 .arg(&mut *part_o)
19233 .arg(&mut *part_m)
19234 .arg(&mut *part_l)
19235 .arg(&hd)
19236 .arg(&nh)
19237 .arg(&nhkv)
19238 .arg(base_dev)
19239 .arg(&scale)
19240 .arg(&nspm)
19241 .arg(&spk)
19242 .arg(&ktb)
19243 .arg(&vtb);
19244 unsafe {
19245 b.launch(cfg)?;
19246 }
19247 let fc = self.func("fa_decode_combine_rows_dc");
19248 let cfg2 = LaunchConfig {
19249 grid_dim: (n_head as u32, t as u32, 1),
19250 block_dim: (head_dim as u32, 1, 1),
19251 shared_mem_bytes: 0,
19252 };
19253 let plus0 = 0i32;
19254 let __s_b2 = self.gpu.stream();
19255 let mut b2 = __s_b2.launch_builder(&fc);
19256 b2.arg(&*part_o)
19257 .arg(&*part_m)
19258 .arg(&*part_l)
19259 .arg(o)
19260 .arg(&hd)
19261 .arg(&nh)
19262 .arg(base_dev)
19263 .arg(&plus0)
19264 .arg(&nspm)
19265 .arg(&spk);
19266 unsafe {
19267 b2.launch(cfg2)?;
19268 }
19269 Ok(())
19270 }
19271
19272 pub fn fa_decode_dc(
19283 &self,
19284 q: &CudaSlice<f32>,
19285 k: &cudarc::driver::CudaView<u8>,
19286 v: &cudarc::driver::CudaView<u8>,
19287 o: &mut CudaSlice<f32>,
19288 head_dim: usize,
19289 n_head: usize,
19290 n_head_kv: usize,
19291 t_kv_dev: &CudaSlice<i32>,
19292 bucket_max: usize,
19293 scale: f32,
19294 k_tok_bytes: usize,
19295 v_tok_bytes: usize,
19296 g: bool,
19297 ) -> Result<(), Box<dyn std::error::Error>> {
19298 self.fa_decode_dc_q8(
19299 q,
19300 k,
19301 v,
19302 o,
19303 head_dim,
19304 n_head,
19305 n_head_kv,
19306 t_kv_dev,
19307 bucket_max,
19308 scale,
19309 k_tok_bytes,
19310 v_tok_bytes,
19311 g,
19312 None,
19313 )
19314 }
19315
19316 #[allow(clippy::too_many_arguments)]
19319 pub fn fa_decode_dc_q8(
19320 &self,
19321 q: &CudaSlice<f32>,
19322 k: &cudarc::driver::CudaView<u8>,
19323 v: &cudarc::driver::CudaView<u8>,
19324 o: &mut CudaSlice<f32>,
19325 head_dim: usize,
19326 n_head: usize,
19327 n_head_kv: usize,
19328 t_kv_dev: &CudaSlice<i32>,
19329 bucket_max: usize,
19330 scale: f32,
19331 k_tok_bytes: usize,
19332 v_tok_bytes: usize,
19333 g: bool,
19334 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19335 ) -> Result<(), Box<dyn std::error::Error>> {
19336 let mut fa_vec =
19344 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19345 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19346 fa_vec = false;
19347 } let sp = fa_split_keys(bucket_max, n_head_kv);
19349 let n_splits = if fa_vec {
19350 ((bucket_max + sp - 1) / sp).max(1)
19351 } else {
19352 ((bucket_max + 255) / 256).max(1)
19353 };
19354 let o_len = n_head * n_splits * head_dim;
19355 let ml_len = n_head * n_splits;
19356 let mut part_guard = self.fa_part_pool.lock().unwrap();
19357 if part_guard
19358 .as_ref()
19359 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19360 .unwrap_or(true)
19361 {
19362 let old = part_guard.take();
19373 let (co, cm) = old
19374 .as_ref()
19375 .map(|pp| (pp.0.len(), pp.1.len()))
19376 .unwrap_or((0, 0));
19377 if let Some(old) = old {
19378 self.fa_part_retired.lock().unwrap().push(old);
19379 }
19380 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19381 eprintln!(
19382 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19383 co, o_len, cm, ml_len
19384 );
19385 }
19386 *part_guard = Some((
19387 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19388 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19389 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19390 ));
19391 }
19392 let pg = part_guard.as_mut().unwrap();
19393 self.gpu
19394 .stream()
19395 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19396 self.gpu
19397 .stream()
19398 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19399 self.gpu
19400 .stream()
19401 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19402 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19403 let (hd, nh, nhkv, nsp) = (
19404 head_dim as i32,
19405 n_head as i32,
19406 n_head_kv as i32,
19407 n_splits as i32,
19408 );
19409 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19410 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
19411 let deep = fa_vec
19414 && head_dim == 256
19415 && fa_v4_at(bucket_max)
19416 && !g
19417 && fa_deep_at(bucket_max)
19418 && !matches!(fa_v4_mode(), "noB3" | "stage");
19419 let (f, cfg) = if fa_vec
19420 && head_dim == 512
19421 && bucket_max >= {
19422 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19423 *FA512_MIN_DC.get_or_init(|| {
19424 std::env::var("MEMRA_FA512_MIN")
19425 .ok()
19426 .and_then(|v| v.parse().ok())
19427 .unwrap_or(512)
19428 })
19429 } {
19430 let gqa = (n_head / n_head_kv).max(1) as u32;
19432 (
19433 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
19434 LaunchConfig {
19435 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19436 block_dim: (32, gqa, 1),
19437 shared_mem_bytes: 0,
19438 },
19439 )
19440 } else if fa_vec && head_dim == 512 {
19441 let q_view = q.as_view();
19444 let mut o_view = o.as_view_mut();
19445 return self.fa_decode_scalar_unified(
19446 &q_view,
19447 k,
19448 v,
19449 &mut o_view,
19450 head_dim,
19451 n_head,
19452 n_head_kv,
19453 0,
19454 Some(t_kv_dev),
19455 scale,
19456 n_splits,
19457 sp,
19458 k_tok_bytes,
19459 v_tok_bytes,
19460 g,
19461 &mut *part_o,
19462 &mut *part_m,
19463 &mut *part_l,
19464 q8_out,
19465 );
19466 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
19467 let gqa = (n_head / n_head_kv).max(1) as u32;
19470 let fv = if g {
19471 self.func_g("fa_decode_vec_q_v4_dc")
19472 } else if deep {
19473 self.func("fa_decode_vec_q_v4_deep_dc")
19474 } else {
19475 self.func("fa_decode_vec_q_v4_dc")
19476 };
19477 let shmem =
19478 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19479 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19480 fv.set_attribute(
19481 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19482 shmem as i32,
19483 )?;
19484 (
19485 fv,
19486 LaunchConfig {
19487 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19488 block_dim: (32, gqa, 1),
19489 shared_mem_bytes: shmem,
19490 },
19491 )
19492 } else if fa_vec && fa_v3_active(head_dim) {
19493 let gqa = (n_head / n_head_kv).max(1) as u32;
19496 let fv = if g {
19497 self.func_g("fa_decode_vec_q_v3_dc")
19498 } else {
19499 self.func("fa_decode_vec_q_v3_dc")
19500 };
19501 let shmem = (32 * head_dim * 2) as u32; (
19503 fv,
19504 LaunchConfig {
19505 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19506 block_dim: (32, gqa, 1),
19507 shared_mem_bytes: shmem,
19508 },
19509 )
19510 } else if fa_vec && fa_v2_on() {
19511 let gqa = (n_head / n_head_kv).max(1) as u32;
19515 let fv = if g {
19516 self.func_g("fa_decode_vec_q_v2_dc")
19517 } else {
19518 self.func("fa_decode_vec_q_v2_dc")
19519 };
19520 let shmem = (2 * 32 * head_dim * 2) as u32; (
19522 fv,
19523 LaunchConfig {
19524 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19525 block_dim: (32, gqa, 1),
19526 shared_mem_bytes: shmem,
19527 },
19528 )
19529 } else if fa_vec {
19530 let gqa = (n_head / n_head_kv).max(1) as u32;
19531 let fv = if g {
19533 self.func_g("fa_decode_vec_q_dc")
19534 } else {
19535 self.func("fa_decode_vec_q_dc")
19536 };
19537 (
19538 fv,
19539 LaunchConfig {
19540 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19541 block_dim: (32, gqa, 1),
19542 shared_mem_bytes: 0,
19543 },
19544 )
19545 } else {
19546 let q_view = q.as_view();
19547 let mut o_view = o.as_view_mut();
19548 return self.fa_decode_scalar_unified(
19549 &q_view,
19550 k,
19551 v,
19552 &mut o_view,
19553 head_dim,
19554 n_head,
19555 n_head_kv,
19556 0,
19557 Some(t_kv_dev),
19558 scale,
19559 n_splits,
19560 if fa_vec { sp } else { 256 },
19561 k_tok_bytes,
19562 v_tok_bytes,
19563 g,
19564 &mut *part_o,
19565 &mut *part_m,
19566 &mut *part_l,
19567 q8_out,
19568 );
19569 };
19570 let ski = sp as i32; let __s_b = self.gpu.stream();
19572 let mut b = __s_b.launch_builder(&f);
19573 b.arg(q)
19574 .arg(k)
19575 .arg(v)
19576 .arg(&mut *part_o)
19577 .arg(&mut *part_m)
19578 .arg(&mut *part_l)
19579 .arg(&hd)
19580 .arg(&nh)
19581 .arg(&nhkv)
19582 .arg(t_kv_dev)
19583 .arg(&scale)
19584 .arg(&nsp)
19585 .arg(&ski)
19586 .arg(&ktb)
19587 .arg(&vtb);
19588 unsafe {
19589 b.launch(cfg)?;
19590 }
19591 let cfg2 = LaunchConfig {
19592 grid_dim: (n_head as u32, 1, 1),
19593 block_dim: (head_dim as u32, 1, 1),
19594 shared_mem_bytes: 0,
19595 };
19596 if let Some((oq, od)) = q8_out {
19597 let fc = if g {
19598 self.func_g("fa_decode_combine_q8_1")
19599 } else {
19600 self.fa_func("fa_decode_combine_q8_1", head_dim)
19601 };
19602 let __s_b2 = self.gpu.stream();
19603 let mut b2 = __s_b2.launch_builder(&fc);
19604 b2.arg(&*part_o)
19605 .arg(&*part_m)
19606 .arg(&*part_l)
19607 .arg(oq)
19608 .arg(od)
19609 .arg(&hd)
19610 .arg(&nh)
19611 .arg(&nsp);
19612 unsafe {
19613 b2.launch(cfg2)?;
19614 }
19615 return Ok(());
19616 }
19617 let fc = if g {
19618 self.func_g("fa_decode_combine_f32")
19619 } else {
19620 self.fa_func("fa_decode_combine_f32", head_dim)
19621 };
19622 let __s_b2 = self.gpu.stream();
19623 let mut b2 = __s_b2.launch_builder(&fc);
19624 b2.arg(&*part_o)
19625 .arg(&*part_m)
19626 .arg(&*part_l)
19627 .arg(o)
19628 .arg(&hd)
19629 .arg(&nh)
19630 .arg(&nsp);
19631 unsafe {
19632 b2.launch(cfg2)?;
19633 }
19634 Ok(())
19635 }
19636
19637 pub fn fa_geom_eager(
19643 &self,
19644 t_kv: usize,
19645 head_dim: usize,
19646 n_head_kv: usize,
19647 g: bool,
19648 ) -> (bool, usize) {
19649 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
19653 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
19659 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
19660 if g && head_dim == 256 && !fa_v4_at(t_kv) {
19666 fa_vec = false;
19667 }
19668 let sp = fa_split_keys(t_kv, n_head_kv);
19669 let n_splits = if fa_vec {
19670 ((t_kv + sp - 1) / sp).max(1)
19671 } else {
19672 ((t_kv + 255) / 256).max(1)
19673 };
19674 (fa_vec, n_splits)
19675 }
19676
19677 pub fn fa_bucket_key(
19683 &self,
19684 t_kv: usize,
19685 head_dim: usize,
19686 n_head_kv: usize,
19687 g: bool,
19688 ) -> (bool, usize) {
19689 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
19690 }
19691
19692 pub fn capture_graph_retained<F>(
19704 &self,
19705 step: F,
19706 ) -> Result<
19707 (
19708 cudarc::driver::CudaGraph,
19709 Vec<Box<dyn std::any::Any + Send>>,
19710 ),
19711 Box<dyn std::error::Error>,
19712 >
19713 where
19714 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19715 {
19716 use cudarc::driver::sys::CUgraphInstantiate_flags;
19717 self.capture_graph_retained_flags(
19718 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19719 step,
19720 )
19721 }
19722
19723 pub fn capture_graph_retained_flags<F>(
19728 &self,
19729 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
19730 mut step: F,
19731 ) -> Result<
19732 (
19733 cudarc::driver::CudaGraph,
19734 Vec<Box<dyn std::any::Any + Send>>,
19735 ),
19736 Box<dyn std::error::Error>,
19737 >
19738 where
19739 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19740 {
19741 use cudarc::driver::sys::CUstreamCaptureMode;
19742 self.capture_keep.lock().unwrap().clear();
19750 let was_tracking = self.gpu.ctx.is_event_tracking();
19751 if was_tracking {
19752 unsafe {
19753 self.gpu.ctx.disable_event_tracking();
19754 }
19755 }
19756 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19757 self.capture_keep_on
19758 .store(true, std::sync::atomic::Ordering::Relaxed);
19759 let w = (|| {
19760 step(self)?;
19761 step(self)
19762 })();
19763 self.capture_keep_on
19764 .store(false, std::sync::atomic::Ordering::Relaxed);
19765 w?;
19766 self.gpu.stream().synchronize()?;
19767 self.gpu
19768 .stream()
19769 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19770 let r = step(self);
19771 let g = self.gpu.stream().end_capture(flags);
19772 r?;
19773 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19774 graph.upload()?;
19775 Ok(graph)
19776 };
19777 let result = run();
19778 self.capture_keep_on
19779 .store(false, std::sync::atomic::Ordering::Relaxed);
19780 if was_tracking {
19781 unsafe {
19782 self.gpu.ctx.enable_event_tracking();
19783 }
19784 }
19785 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
19786 Ok((result?, keeper))
19787 }
19788
19789 pub fn capture_graph<F>(
19790 &self,
19791 mut step: F,
19792 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
19793 where
19794 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19795 {
19796 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
19797 let was_tracking = self.gpu.ctx.is_event_tracking();
19805 if was_tracking {
19806 unsafe {
19807 self.gpu.ctx.disable_event_tracking();
19808 }
19809 }
19810 let iflag = {
19817 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
19818 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
19819 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
19822 Ok("priority") => {
19823 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
19824 }
19825 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19826 })
19827 };
19828 let ct = {
19835 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19836 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
19837 };
19838 let warmups = {
19861 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19862 *W.get_or_init(|| {
19863 std::env::var("MEMRA_GRAPH_WARMUPS")
19864 .ok()
19865 .and_then(|v| v.parse().ok())
19866 .filter(|n| *n >= 1)
19867 .unwrap_or(1)
19868 })
19869 };
19870 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19871 let t_w = std::time::Instant::now();
19872 for _ in 0..warmups {
19874 step(self)?;
19875 }
19876 self.gpu.stream().synchronize()?;
19877 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
19878 let t_c = std::time::Instant::now();
19880 self.gpu
19881 .stream()
19882 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19883 let r = step(self);
19886 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
19887 let t_i = std::time::Instant::now();
19888 let g = self.gpu.stream().end_capture(iflag);
19889 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
19890 r?;
19891 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19892 let t_u = std::time::Instant::now();
19893 graph.upload()?;
19894 if ct {
19895 println!(
19896 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
19897 instantiate {ms_inst:.2} ms upload {:.2} ms",
19898 t_u.elapsed().as_secs_f64() * 1e3
19899 );
19900 }
19901 Ok(graph)
19902 };
19903 let result = run();
19904 if was_tracking {
19905 unsafe {
19906 self.gpu.ctx.enable_event_tracking();
19907 }
19908 }
19909 result
19910 }
19911
19912 pub fn gdn_scan_s128_view(
19914 &self,
19915 q: &CudaSlice<f32>,
19916 k: &CudaSlice<f32>,
19917 v: &CudaSlice<f32>,
19918 g: &CudaSlice<f32>,
19919 beta: &CudaSlice<f32>,
19920 state_in: &cudarc::driver::CudaView<f32>,
19921 state_out: &mut cudarc::driver::CudaViewMut<f32>,
19922 o: &mut CudaSlice<f32>,
19923 n_head: usize,
19924 t: usize,
19925 scale: f32,
19926 ) -> Result<(), Box<dyn std::error::Error>> {
19927 let f = self.func("gdn_scan_s128");
19928 const S_V: u32 = 128;
19929 const WARP: u32 = 32;
19930 const COLS: u32 = 4;
19931 let cfg = LaunchConfig {
19932 grid_dim: (n_head as u32, 1, S_V / COLS),
19933 block_dim: (WARP, COLS, 1),
19934 shared_mem_bytes: 0,
19935 };
19936 let (h, ti) = (n_head as i32, t as i32);
19937 let __s_b = self.gpu.stream();
19938 let mut b = __s_b.launch_builder(&f);
19939 b.arg(q)
19940 .arg(k)
19941 .arg(v)
19942 .arg(g)
19943 .arg(beta)
19944 .arg(state_in)
19945 .arg(state_out)
19946 .arg(o)
19947 .arg(&h)
19948 .arg(&ti)
19949 .arg(&scale);
19950 unsafe {
19951 b.launch(cfg)?;
19952 }
19953 Ok(())
19954 }
19955
19956 pub fn ssm_conv1d_view(
19958 &self,
19959 x: &cudarc::driver::CudaView<f32>,
19960 w: &CudaSlice<f32>,
19961 y: &mut CudaSlice<f32>,
19962 conv_dim: usize,
19963 t: usize,
19964 d_conv: usize,
19965 silu: bool,
19966 ) -> Result<(), Box<dyn std::error::Error>> {
19967 let f = self.func("ssm_conv1d_silu_f32");
19968 let cfg = LaunchConfig {
19970 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
19971 block_dim: (256, 1, 1),
19972 shared_mem_bytes: 0,
19973 };
19974 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
19975 let __s_b = self.gpu.stream();
19976 let mut b = __s_b.launch_builder(&f);
19977 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
19978 unsafe {
19979 b.launch(cfg)?;
19980 }
19981 Ok(())
19982 }
19983
19984 pub fn ssm_conv1d_tm(
19991 &self,
19992 qkv_tm: &CudaSlice<f32>,
19993 w: &CudaSlice<f32>,
19994 y: &mut CudaSlice<f32>,
19995 conv_dim: usize,
19996 t: usize,
19997 d_conv: usize,
19998 ) -> Result<(), Box<dyn std::error::Error>> {
19999 let f = self.func("ssm_conv1d_tm_f32");
20000 let cfg = LaunchConfig {
20001 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20002 block_dim: (256, 1, 1),
20003 shared_mem_bytes: 0,
20004 };
20005 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20006 let __s_b = self.gpu.stream();
20007 let mut b = __s_b.launch_builder(&f);
20008 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
20009 unsafe {
20010 b.launch(cfg)?;
20011 }
20012 Ok(())
20013 }
20014
20015 pub fn ssm_conv1d_tm_state(
20023 &self,
20024 qkv_tm: &CudaSlice<f32>,
20025 conv_state: &mut CudaSlice<f32>,
20026 w: &CudaSlice<f32>,
20027 y: &mut CudaSlice<f32>,
20028 conv_dim: usize,
20029 t: usize,
20030 d_conv: usize,
20031 ) -> Result<(), Box<dyn std::error::Error>> {
20032 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
20033 }
20034
20035 #[allow(clippy::too_many_arguments)]
20038 pub fn ssm_conv1d_tm_state_pad(
20039 &self,
20040 qkv_tm: &CudaSlice<f32>,
20041 conv_state: &mut CudaSlice<f32>,
20042 w: &CudaSlice<f32>,
20043 y: &mut CudaSlice<f32>,
20044 conv_dim: usize,
20045 t: usize,
20046 d_conv: usize,
20047 pad_len: Option<&CudaSlice<i32>>,
20048 ) -> Result<(), Box<dyn std::error::Error>> {
20049 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20050 let ring_old = if t < d_conv - 1 {
20054 Some(self.clone_dtod(conv_state)?)
20055 } else {
20056 None
20057 };
20058 {
20059 let f = self.func("ssm_conv1d_tm_state_f32");
20060 let cfg = LaunchConfig {
20061 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20062 block_dim: (256, 1, 1),
20063 shared_mem_bytes: 0,
20064 };
20065 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20066 let __s_b = self.gpu.stream();
20067 let mut b = __s_b.launch_builder(&f);
20068 b.arg(qkv_tm)
20069 .arg(&*conv_state)
20070 .arg(w)
20071 .arg(y)
20072 .arg(&cd)
20073 .arg(&ti)
20074 .arg(&dc);
20075 unsafe {
20076 b.launch(cfg)?;
20077 }
20078 }
20079 match (ring_old, pad_len) {
20080 (None, Some(len_d)) => {
20081 let f = self.func("ssm_conv_ring_update_dev_f32");
20082 let n = conv_dim * (d_conv - 1);
20083 let cfg = LaunchConfig::for_num_elems(n as u32);
20084 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20085 let __s_b = self.gpu.stream();
20086 let mut b = __s_b.launch_builder(&f);
20087 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20088 unsafe {
20089 b.launch(cfg)?;
20090 }
20091 }
20092 (None, None) => {
20093 let f = self.func("ssm_conv_ring_update_f32");
20094 let n = conv_dim * (d_conv - 1);
20095 let cfg = LaunchConfig::for_num_elems(n as u32);
20096 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20097 let __s_b = self.gpu.stream();
20098 let mut b = __s_b.launch_builder(&f);
20099 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20100 unsafe {
20101 b.launch(cfg)?;
20102 }
20103 }
20104 (Some(old), _) => {
20105 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
20106 }
20107 }
20108 Ok(())
20109 }
20110
20111 pub fn ssm_conv1d_tm_state_pad_v(
20113 &self,
20114 qkv_tm: &cudarc::driver::CudaView<f32>,
20115 conv_state: &mut CudaSlice<f32>,
20116 w: &CudaSlice<f32>,
20117 y: &mut CudaSlice<f32>,
20118 conv_dim: usize,
20119 t: usize,
20120 d_conv: usize,
20121 pad_len: Option<&CudaSlice<i32>>,
20122 ) -> Result<(), Box<dyn std::error::Error>> {
20123 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20124 let ring_old = if t < d_conv - 1 {
20128 Some(self.clone_dtod(conv_state)?)
20129 } else {
20130 None
20131 };
20132 {
20133 let f = self.func("ssm_conv1d_tm_state_f32");
20134 let cfg = LaunchConfig {
20135 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20136 block_dim: (256, 1, 1),
20137 shared_mem_bytes: 0,
20138 };
20139 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20140 let __s_b = self.gpu.stream();
20141 let mut b = __s_b.launch_builder(&f);
20142 b.arg(qkv_tm)
20143 .arg(&*conv_state)
20144 .arg(w)
20145 .arg(y)
20146 .arg(&cd)
20147 .arg(&ti)
20148 .arg(&dc);
20149 unsafe {
20150 b.launch(cfg)?;
20151 }
20152 }
20153 match (ring_old, pad_len) {
20154 (None, Some(len_d)) => {
20155 let f = self.func("ssm_conv_ring_update_dev_f32");
20156 let n = conv_dim * (d_conv - 1);
20157 let cfg = LaunchConfig::for_num_elems(n as u32);
20158 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20159 let __s_b = self.gpu.stream();
20160 let mut b = __s_b.launch_builder(&f);
20161 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20162 unsafe {
20163 b.launch(cfg)?;
20164 }
20165 }
20166 (None, None) => {
20167 let f = self.func("ssm_conv_ring_update_f32");
20168 let n = conv_dim * (d_conv - 1);
20169 let cfg = LaunchConfig::for_num_elems(n as u32);
20170 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20171 let __s_b = self.gpu.stream();
20172 let mut b = __s_b.launch_builder(&f);
20173 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20174 unsafe {
20175 b.launch(cfg)?;
20176 }
20177 }
20178 (Some(_), _) => unreachable!(
20179 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
20180 ),
20181 }
20182 Ok(())
20183 }
20184
20185 pub fn ssm_conv_ring_rebuild(
20190 &self,
20191 qkv_tm: &CudaSlice<f32>,
20192 ring_old: &CudaSlice<f32>,
20193 conv_state: &mut CudaSlice<f32>,
20194 conv_dim: usize,
20195 tc: usize,
20196 d_conv: usize,
20197 ) -> Result<(), Box<dyn std::error::Error>> {
20198 let f = self.func("ssm_conv_ring_rebuild_f32");
20199 let n = conv_dim * (d_conv - 1);
20200 let cfg = LaunchConfig::for_num_elems(n as u32);
20201 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
20202 let __s_b = self.gpu.stream();
20203 let mut b = __s_b.launch_builder(&f);
20204 b.arg(qkv_tm)
20205 .arg(ring_old)
20206 .arg(conv_state)
20207 .arg(&cd)
20208 .arg(&ti)
20209 .arg(&dc);
20210 unsafe {
20211 b.launch(cfg)?;
20212 }
20213 Ok(())
20214 }
20215
20216 #[allow(clippy::too_many_arguments)]
20221 pub fn gdn_prep_decode(
20222 &self,
20223 conv_out: &CudaSlice<f32>,
20224 beta_raw: &CudaSlice<f32>,
20225 alpha: &CudaSlice<f32>,
20226 dt_bias: &CudaSlice<f32>,
20227 a: &CudaSlice<f32>,
20228 q_l2: &mut CudaSlice<f32>,
20229 k_l2: &mut CudaSlice<f32>,
20230 v_g: &mut CudaSlice<f32>,
20231 beta: &mut CudaSlice<f32>,
20232 g_log: &mut CudaSlice<f32>,
20233 d_state: usize,
20234 num_v: usize,
20235 num_k: usize,
20236 key_dim: usize,
20237 eps: f32,
20238 ) -> Result<(), Box<dyn std::error::Error>> {
20239 let f = self.func("gdn_prep_decode_f32");
20240 let cfg = LaunchConfig {
20241 grid_dim: (num_v as u32, 1, 1),
20242 block_dim: (32, 4, 1),
20243 shared_mem_bytes: 0,
20244 };
20245 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20246 let __s_b = self.gpu.stream();
20247 let mut b = __s_b.launch_builder(&f);
20248 b.arg(conv_out)
20249 .arg(beta_raw)
20250 .arg(alpha)
20251 .arg(dt_bias)
20252 .arg(a)
20253 .arg(q_l2)
20254 .arg(k_l2)
20255 .arg(v_g)
20256 .arg(beta)
20257 .arg(g_log)
20258 .arg(&ds)
20259 .arg(&nv)
20260 .arg(&nk)
20261 .arg(&kd)
20262 .arg(&eps);
20263 unsafe {
20264 b.launch(cfg)?;
20265 }
20266 Ok(())
20267 }
20268
20269 #[allow(clippy::too_many_arguments)]
20273 pub fn ssm_conv1d_gdn(
20274 &self,
20275 qkv_tm: &CudaSlice<f32>,
20276 w: &CudaSlice<f32>,
20277 q_g: &mut CudaSlice<f32>,
20278 k_g: &mut CudaSlice<f32>,
20279 v_g: &mut CudaSlice<f32>,
20280 conv_dim: usize,
20281 t: usize,
20282 d_conv: usize,
20283 d_state: usize,
20284 num_v: usize,
20285 num_k: usize,
20286 key_dim: usize,
20287 ) -> Result<(), Box<dyn std::error::Error>> {
20288 let f = self.func("ssm_conv1d_gdn_f32");
20289 let cfg = LaunchConfig {
20290 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20291 block_dim: (256, 1, 1),
20292 shared_mem_bytes: 0,
20293 };
20294 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20295 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20296 let __s_b = self.gpu.stream();
20297 let mut b = __s_b.launch_builder(&f);
20298 b.arg(qkv_tm)
20299 .arg(w)
20300 .arg(q_g)
20301 .arg(k_g)
20302 .arg(v_g)
20303 .arg(&cd)
20304 .arg(&ti)
20305 .arg(&dc)
20306 .arg(&ds)
20307 .arg(&nv)
20308 .arg(&nk)
20309 .arg(&kd);
20310 unsafe {
20311 b.launch(cfg)?;
20312 }
20313 Ok(())
20314 }
20315
20316 pub fn ssm_conv1d(
20317 &self,
20318 x: &CudaSlice<f32>,
20319 w: &CudaSlice<f32>,
20320 y: &mut CudaSlice<f32>,
20321 conv_dim: usize,
20322 t: usize,
20323 d_conv: usize,
20324 silu: bool,
20325 ) -> Result<(), Box<dyn std::error::Error>> {
20326 let f = self.func("ssm_conv1d_silu_f32");
20327 let cfg = LaunchConfig {
20328 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20329 block_dim: (256, 1, 1),
20330 shared_mem_bytes: 0,
20331 };
20332 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20333 let __s_b = self.gpu.stream();
20334 let mut b = __s_b.launch_builder(&f);
20335 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20336 unsafe {
20337 b.launch(cfg)?;
20338 }
20339 Ok(())
20340 }
20341
20342 pub fn gdn_scan_s128(
20345 &self,
20346 q: &CudaSlice<f32>,
20347 k: &CudaSlice<f32>,
20348 v: &CudaSlice<f32>,
20349 g: &CudaSlice<f32>,
20350 beta: &CudaSlice<f32>,
20351 state_in: &CudaSlice<f32>,
20352 state_out: &mut CudaSlice<f32>,
20353 o: &mut CudaSlice<f32>,
20354 n_head: usize,
20355 t: usize,
20356 scale: f32,
20357 ) -> Result<(), Box<dyn std::error::Error>> {
20358 let f = self.func("gdn_scan_s128");
20359 const S_V: u32 = 128;
20360 const WARP: u32 = 32;
20361 const COLS_PER_BLOCK: u32 = 4;
20362 let cfg = LaunchConfig {
20363 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20364 block_dim: (WARP, COLS_PER_BLOCK, 1),
20365 shared_mem_bytes: 0,
20366 };
20367 let (h, ti) = (n_head as i32, t as i32);
20368 let __s_b = self.gpu.stream();
20369 let mut b = __s_b.launch_builder(&f);
20370 b.arg(q)
20371 .arg(k)
20372 .arg(v)
20373 .arg(g)
20374 .arg(beta)
20375 .arg(state_in)
20376 .arg(state_out)
20377 .arg(o)
20378 .arg(&h)
20379 .arg(&ti)
20380 .arg(&scale);
20381 unsafe {
20382 b.launch(cfg)?;
20383 }
20384 Ok(())
20385 }
20386
20387 #[allow(clippy::too_many_arguments)]
20392 pub fn ssm_conv1d_fused_decode_b(
20393 &self,
20394 qkv_cols: &CudaSlice<f32>,
20395 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20396 w: &CudaSlice<f32>,
20397 conv_outs: &mut CudaSlice<f32>,
20398 conv_dim: usize,
20399 d_conv: usize,
20400 b_n: usize,
20401 ) -> Result<(), Box<dyn std::error::Error>> {
20402 let f = self.func("ssm_conv1d_fused_decode_b_f32");
20403 let cfg = LaunchConfig {
20404 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20405 block_dim: (256, 1, 1),
20406 shared_mem_bytes: 0,
20407 };
20408 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20409 let __s_b = self.gpu.stream();
20410 let mut b = __s_b.launch_builder(&f);
20411 b.arg(qkv_cols)
20412 .arg(conv_state_ptrs)
20413 .arg(w)
20414 .arg(conv_outs)
20415 .arg(&cd)
20416 .arg(&dc);
20417 unsafe {
20418 b.launch(cfg)?;
20419 }
20420 Ok(())
20421 }
20422
20423 #[allow(clippy::too_many_arguments)]
20424 pub fn gdn_prep_decode_b(
20425 &self,
20426 conv_outs: &CudaSlice<f32>,
20427 beta_raws: &CudaSlice<f32>,
20428 alphas: &CudaSlice<f32>,
20429 dt_bias: &CudaSlice<f32>,
20430 a: &CudaSlice<f32>,
20431 q_l2: &mut CudaSlice<f32>,
20432 k_l2: &mut CudaSlice<f32>,
20433 v_g: &mut CudaSlice<f32>,
20434 beta: &mut CudaSlice<f32>,
20435 g_log: &mut CudaSlice<f32>,
20436 d_state: usize,
20437 num_v: usize,
20438 num_k: usize,
20439 key_dim: usize,
20440 eps: f32,
20441 conv_dim: usize,
20442 b_n: usize,
20443 ) -> Result<(), Box<dyn std::error::Error>> {
20444 let f = self.func("gdn_prep_decode_b_f32");
20445 let cfg = LaunchConfig {
20446 grid_dim: (num_v as u32, 1, b_n as u32),
20447 block_dim: (32, 4, 1),
20448 shared_mem_bytes: 0,
20449 };
20450 let (ds, nv, nk, kd, cd) = (
20451 d_state as i32,
20452 num_v as i32,
20453 num_k as i32,
20454 key_dim as i32,
20455 conv_dim as i32,
20456 );
20457 let __s_b = self.gpu.stream();
20458 let mut b = __s_b.launch_builder(&f);
20459 b.arg(conv_outs)
20460 .arg(beta_raws)
20461 .arg(alphas)
20462 .arg(dt_bias)
20463 .arg(a)
20464 .arg(q_l2)
20465 .arg(k_l2)
20466 .arg(v_g)
20467 .arg(beta)
20468 .arg(g_log)
20469 .arg(&ds)
20470 .arg(&nv)
20471 .arg(&nk)
20472 .arg(&kd)
20473 .arg(&eps)
20474 .arg(&cd);
20475 unsafe {
20476 b.launch(cfg)?;
20477 }
20478 Ok(())
20479 }
20480
20481 #[allow(clippy::too_many_arguments)]
20482 pub fn gdn_scan_s128_batched(
20483 &self,
20484 q: &CudaSlice<f32>,
20485 k: &CudaSlice<f32>,
20486 v: &CudaSlice<f32>,
20487 g: &CudaSlice<f32>,
20488 beta: &CudaSlice<f32>,
20489 state_in_ptrs: &cudarc::driver::CudaView<u64>,
20490 state_out_ptrs: &cudarc::driver::CudaView<u64>,
20491 o: &mut CudaSlice<f32>,
20492 n_head: usize,
20493 b_n: usize,
20494 scale: f32,
20495 ) -> Result<(), Box<dyn std::error::Error>> {
20496 let f = self.func("gdn_scan_s128_b");
20497 const S_V: u32 = 128;
20498 const WARP: u32 = 32;
20499 const COLS_PER_BLOCK: u32 = 4;
20500 let cfg = LaunchConfig {
20501 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20502 block_dim: (WARP, COLS_PER_BLOCK, 1),
20503 shared_mem_bytes: 0,
20504 };
20505 let h = n_head as i32;
20506 let __s_b = self.gpu.stream();
20507 let mut b = __s_b.launch_builder(&f);
20508 b.arg(q)
20509 .arg(k)
20510 .arg(v)
20511 .arg(g)
20512 .arg(beta)
20513 .arg(state_in_ptrs)
20514 .arg(state_out_ptrs)
20515 .arg(o)
20516 .arg(&h)
20517 .arg(&scale);
20518 unsafe {
20519 b.launch(cfg)?;
20520 }
20521 Ok(())
20522 }
20523
20524 #[allow(clippy::too_many_arguments)]
20530 pub fn ssm_conv1d_fused_decode_b_view(
20531 &self,
20532 qkv_cols: &cudarc::driver::CudaView<f32>,
20533 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20534 w: &CudaSlice<f32>,
20535 conv_outs: &mut CudaSlice<f32>,
20536 conv_dim: usize,
20537 d_conv: usize,
20538 b_n: usize,
20539 ) -> Result<(), Box<dyn std::error::Error>> {
20540 let f = self.func("ssm_conv1d_fused_decode_b_f32");
20541 let cfg = LaunchConfig {
20542 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20543 block_dim: (256, 1, 1),
20544 shared_mem_bytes: 0,
20545 };
20546 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20547 let __s_b = self.gpu.stream();
20548 let mut b = __s_b.launch_builder(&f);
20549 b.arg(qkv_cols)
20550 .arg(conv_state_ptrs)
20551 .arg(w)
20552 .arg(conv_outs)
20553 .arg(&cd)
20554 .arg(&dc);
20555 unsafe {
20556 b.launch(cfg)?;
20557 }
20558 Ok(())
20559 }
20560
20561 #[allow(clippy::too_many_arguments)]
20562 pub fn gdn_prep_decode_b_view(
20563 &self,
20564 conv_outs: &CudaSlice<f32>,
20565 beta_raws: &cudarc::driver::CudaView<f32>,
20566 alphas: &cudarc::driver::CudaView<f32>,
20567 dt_bias: &CudaSlice<f32>,
20568 a: &CudaSlice<f32>,
20569 q_l2: &mut CudaSlice<f32>,
20570 k_l2: &mut CudaSlice<f32>,
20571 v_g: &mut CudaSlice<f32>,
20572 beta: &mut CudaSlice<f32>,
20573 g_log: &mut CudaSlice<f32>,
20574 d_state: usize,
20575 num_v: usize,
20576 num_k: usize,
20577 key_dim: usize,
20578 eps: f32,
20579 conv_dim: usize,
20580 b_n: usize,
20581 ) -> Result<(), Box<dyn std::error::Error>> {
20582 let f = self.func("gdn_prep_decode_b_f32");
20583 let cfg = LaunchConfig {
20584 grid_dim: (num_v as u32, 1, b_n as u32),
20585 block_dim: (32, 4, 1),
20586 shared_mem_bytes: 0,
20587 };
20588 let (ds, nv, nk, kd, cd) = (
20589 d_state as i32,
20590 num_v as i32,
20591 num_k as i32,
20592 key_dim as i32,
20593 conv_dim as i32,
20594 );
20595 let __s_b = self.gpu.stream();
20596 let mut b = __s_b.launch_builder(&f);
20597 b.arg(conv_outs)
20598 .arg(beta_raws)
20599 .arg(alphas)
20600 .arg(dt_bias)
20601 .arg(a)
20602 .arg(q_l2)
20603 .arg(k_l2)
20604 .arg(v_g)
20605 .arg(beta)
20606 .arg(g_log)
20607 .arg(&ds)
20608 .arg(&nv)
20609 .arg(&nk)
20610 .arg(&kd)
20611 .arg(&eps)
20612 .arg(&cd);
20613 unsafe {
20614 b.launch(cfg)?;
20615 }
20616 Ok(())
20617 }
20618
20619 #[allow(clippy::too_many_arguments)]
20620 pub fn gdn_scan_s128_batched_view(
20621 &self,
20622 q: &CudaSlice<f32>,
20623 k: &CudaSlice<f32>,
20624 v: &CudaSlice<f32>,
20625 g: &CudaSlice<f32>,
20626 beta: &CudaSlice<f32>,
20627 state_in_ptrs: &cudarc::driver::CudaView<u64>,
20628 state_out_ptrs: &cudarc::driver::CudaView<u64>,
20629 o: &mut cudarc::driver::CudaViewMut<f32>,
20630 n_head: usize,
20631 b_n: usize,
20632 scale: f32,
20633 ) -> Result<(), Box<dyn std::error::Error>> {
20634 let f = self.func("gdn_scan_s128_b");
20635 const S_V: u32 = 128;
20636 const WARP: u32 = 32;
20637 const COLS_PER_BLOCK: u32 = 4;
20638 let cfg = LaunchConfig {
20639 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20640 block_dim: (WARP, COLS_PER_BLOCK, 1),
20641 shared_mem_bytes: 0,
20642 };
20643 let h = n_head as i32;
20644 let __s_b = self.gpu.stream();
20645 let mut b = __s_b.launch_builder(&f);
20646 b.arg(q)
20647 .arg(k)
20648 .arg(v)
20649 .arg(g)
20650 .arg(beta)
20651 .arg(state_in_ptrs)
20652 .arg(state_out_ptrs)
20653 .arg(o)
20654 .arg(&h)
20655 .arg(&scale);
20656 unsafe {
20657 b.launch(cfg)?;
20658 }
20659 Ok(())
20660 }
20661
20662 pub fn gdn_chunked_enabled() -> bool {
20671 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20672 *E.get_or_init(|| {
20673 std::env::var("MEMRA_GDN_CHUNKED")
20674 .map(|v| v != "0")
20675 .unwrap_or(true)
20676 })
20677 }
20678
20679 pub fn gdn_chunk_size() -> usize {
20684 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20685 *C.get_or_init(|| {
20686 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
20687 .ok()
20688 .and_then(|v| v.parse().ok())
20689 .unwrap_or(32);
20690 c.clamp(32, 128) / 32 * 32
20691 })
20692 }
20693
20694 #[allow(clippy::too_many_arguments)]
20699 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
20702 #[allow(clippy::too_many_arguments)]
20703 pub fn gdn_chunk_k123(
20704 &self,
20705 q: &CudaSlice<f32>,
20706 k: &CudaSlice<f32>,
20707 v: &CudaSlice<f32>,
20708 g: &CudaSlice<f32>,
20709 beta: &CudaSlice<f32>,
20710 wb16: Option<&mut CudaSlice<u8>>,
20711 n_head: usize,
20712 t: usize,
20713 c: usize,
20714 hk: usize,
20715 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
20716 ) -> Result<
20717 (
20718 CudaSlice<f32>,
20719 CudaSlice<f32>,
20720 CudaSlice<f32>,
20721 CudaSlice<f32>,
20722 ),
20723 Box<dyn std::error::Error>,
20724 > {
20725 const D: usize = 128;
20726 let h = n_head;
20727 let nc = (t + c - 1) / c;
20728 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
20729 let mut gcum = self.uninit(t * h)?;
20730 let mut a = self.uninit(nc * h * c * c)?;
20731 let mut p = self.uninit(nc * h * c * c)?;
20732 let mut u = self.uninit(nc * h * c * D)?;
20733 let mut w = self.uninit(nc * h * c * D)?;
20734 {
20735 let f = self.func("gdn_chunk_cumgate_f32");
20737 let cfg = LaunchConfig {
20738 grid_dim: (nc as u32, h as u32, 1),
20739 block_dim: (32, 1, 1),
20740 shared_mem_bytes: 0,
20741 };
20742 let __s_b = self.gpu.stream();
20743 let mut b = __s_b.launch_builder(&f);
20744 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
20745 unsafe {
20746 b.launch(cfg)?;
20747 }
20748 }
20749 if let Some((qb, kb, pb)) = k2w {
20750 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
20753 let f = self.func("gdn_k2_wgmma");
20754 let cfg = LaunchConfig {
20755 grid_dim: (nc as u32, h as u32, 1),
20756 block_dim: (128, 1, 1),
20757 shared_mem_bytes: 0,
20758 };
20759 let hki = hk as i32;
20760 let __s_b = self.gpu.stream();
20761 let mut b = __s_b.launch_builder(&f);
20762 b.arg(qb)
20763 .arg(kb)
20764 .arg(&gcum)
20765 .arg(beta)
20766 .arg(&mut a)
20767 .arg(&mut *pb)
20768 .arg(&hi)
20769 .arg(&ti)
20770 .arg(&ci)
20771 .arg(&hki);
20772 unsafe {
20773 b.launch(cfg)?;
20774 }
20775 } else if c <= 64 && !portable_mma_gated() {
20776 let f = self.func("gdn_chunk_attn_f32");
20778 let jt = ((c + 31) / 32) as u32;
20779 let cfg = LaunchConfig {
20780 grid_dim: (nc as u32, h as u32, jt),
20781 block_dim: (256, 1, 1),
20782 shared_mem_bytes: 0,
20783 };
20784 let hki = hk as i32;
20785 let __s_b = self.gpu.stream();
20786 let mut b = __s_b.launch_builder(&f);
20787 b.arg(q)
20788 .arg(k)
20789 .arg(&gcum)
20790 .arg(beta)
20791 .arg(&mut a)
20792 .arg(&mut p)
20793 .arg(&hi)
20794 .arg(&ti)
20795 .arg(&ci)
20796 .arg(&hki);
20797 unsafe {
20798 b.launch(cfg)?;
20799 }
20800 } else {
20801 assert!(
20803 hk == h,
20804 "generic K2 is broadcast-only (de-broadcast rides C==32)"
20805 );
20806 let f = self.func("gdn_chunk_attn_g_f32");
20807 let cfg = LaunchConfig {
20808 grid_dim: (nc as u32, h as u32, 1),
20809 block_dim: (32, 8, 1),
20810 shared_mem_bytes: 0,
20811 };
20812 let __s_b = self.gpu.stream();
20813 let mut b = __s_b.launch_builder(&f);
20814 b.arg(q)
20815 .arg(k)
20816 .arg(&gcum)
20817 .arg(beta)
20818 .arg(&mut a)
20819 .arg(&mut p)
20820 .arg(&hi)
20821 .arg(&ti)
20822 .arg(&ci);
20823 unsafe {
20824 b.launch(cfg)?;
20825 }
20826 }
20827 {
20828 let cfg = LaunchConfig {
20830 grid_dim: (nc as u32, h as u32, 1),
20831 block_dim: (256, 1, 1),
20832 shared_mem_bytes: 0,
20833 };
20834 match c {
20835 32 | 64 => {
20836 let f = self.func(if c == 32 {
20837 "gdn_chunk_solve32_f32"
20838 } else {
20839 "gdn_chunk_solve64_f32"
20840 });
20841 let wb: u64 = match wb16 {
20843 Some(d) => self.addr_u8(d),
20844 None => 0,
20845 };
20846 let hki = hk as i32;
20847 let __s_b = self.gpu.stream();
20848 let mut b = __s_b.launch_builder(&f);
20849 b.arg(v)
20850 .arg(k)
20851 .arg(&a)
20852 .arg(&gcum)
20853 .arg(&mut u)
20854 .arg(&mut w)
20855 .arg(&wb)
20856 .arg(&hi)
20857 .arg(&ti)
20858 .arg(&hki);
20859 unsafe {
20860 b.launch(cfg)?;
20861 }
20862 }
20863 _ => {
20864 assert!(hk == h, "generic K3 is broadcast-only");
20865 let f = self.func("gdn_chunk_solve_f32");
20866 let __s_b = self.gpu.stream();
20867 let mut b = __s_b.launch_builder(&f);
20868 b.arg(v)
20869 .arg(k)
20870 .arg(&a)
20871 .arg(&gcum)
20872 .arg(&mut u)
20873 .arg(&mut w)
20874 .arg(&hi)
20875 .arg(&ti)
20876 .arg(&ci);
20877 unsafe {
20878 b.launch(cfg)?;
20879 }
20880 }
20881 }
20882 }
20883 Ok((gcum, p, u, w))
20884 }
20885
20886 pub fn gdn_db_on() -> bool {
20890 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
20891 }
20892
20893 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
20902 !portable_mma_gated()
20903 && c == 32
20904 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
20905 Ok("1") => true,
20906 Ok("0") => false,
20907 _ => gdn_mma_default_on(),
20908 }
20909 }
20910
20911 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
20918 cfg!(memra_hopper_mma)
20919 && self.gdn_mma_enabled(c)
20920 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
20921 }
20922
20923 #[allow(clippy::too_many_arguments)]
20928 pub fn ssm_conv1d_gdn_state_pad(
20929 &self,
20930 qkv_tm: &cudarc::driver::CudaView<f32>,
20931 conv_state: &mut CudaSlice<f32>,
20932 w: &CudaSlice<f32>,
20933 q_g: &mut CudaSlice<f32>,
20934 k_g: &mut CudaSlice<f32>,
20935 v_g: &mut CudaSlice<f32>,
20936 conv_dim: usize,
20937 t: usize,
20938 d_conv: usize,
20939 d_state: usize,
20940 num_v: usize,
20941 num_k: usize,
20942 key_dim: usize,
20943 hk: usize,
20944 pad_len: Option<&CudaSlice<i32>>,
20945 ) -> Result<(), Box<dyn std::error::Error>> {
20946 assert!(
20947 t >= d_conv - 1,
20948 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
20949 );
20950 {
20951 let f = self.func("ssm_conv1d_gdn_state_f32");
20952 let cfg = LaunchConfig {
20953 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20954 block_dim: (256, 1, 1),
20955 shared_mem_bytes: 0,
20956 };
20957 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20958 let (ds, nv, nk, kd, hki) = (
20959 d_state as i32,
20960 num_v as i32,
20961 num_k as i32,
20962 key_dim as i32,
20963 hk as i32,
20964 );
20965 let __s_b = self.gpu.stream();
20966 let mut b = __s_b.launch_builder(&f);
20967 b.arg(qkv_tm)
20968 .arg(&*conv_state)
20969 .arg(w)
20970 .arg(q_g)
20971 .arg(k_g)
20972 .arg(v_g)
20973 .arg(&cd)
20974 .arg(&ti)
20975 .arg(&dc)
20976 .arg(&ds)
20977 .arg(&nv)
20978 .arg(&nk)
20979 .arg(&kd)
20980 .arg(&hki);
20981 unsafe {
20982 b.launch(cfg)?;
20983 }
20984 }
20985 match pad_len {
20986 Some(len_d) => {
20987 let f = self.func("ssm_conv_ring_update_dev_f32");
20988 let n = conv_dim * (d_conv - 1);
20989 let cfg = LaunchConfig::for_num_elems(n as u32);
20990 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20991 let __s_b = self.gpu.stream();
20992 let mut b = __s_b.launch_builder(&f);
20993 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20994 unsafe {
20995 b.launch(cfg)?;
20996 }
20997 }
20998 None => {
20999 let f = self.func("ssm_conv_ring_update_f32");
21000 let n = conv_dim * (d_conv - 1);
21001 let cfg = LaunchConfig::for_num_elems(n as u32);
21002 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21003 let __s_b = self.gpu.stream();
21004 let mut b = __s_b.launch_builder(&f);
21005 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
21006 unsafe {
21007 b.launch(cfg)?;
21008 }
21009 }
21010 }
21011 Ok(())
21012 }
21013
21014 pub fn gdn_chunk_alloc(
21018 &self,
21019 n_head: usize,
21020 t: usize,
21021 c: usize,
21022 hk: usize,
21023 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
21024 const D: usize = 128;
21025 assert!(
21026 c == 32,
21027 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
21028 );
21029 let h = n_head;
21030 let nc = (t + c - 1) / c;
21031 Ok(GdnChunkBufs {
21032 gcum: self.uninit(t * h)?,
21033 a: self.uninit(nc * h * c * c)?,
21034 p: self.uninit(nc * h * c * c)?,
21035 u: self.uninit(nc * h * c * D)?,
21036 w: self.uninit(nc * h * c * D)?,
21037 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21038 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21039 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21040 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
21041 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21042 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
21043 o: self.uninit(D * h * t)?,
21044 t,
21045 nc,
21046 })
21047 }
21048
21049 pub fn f32_to_bf16_v(
21051 &self,
21052 x: &cudarc::driver::CudaView<f32>,
21053 dst: &mut CudaSlice<u8>,
21054 n: usize,
21055 ) -> Result<(), Box<dyn std::error::Error>> {
21056 let f = self.func("f32_to_bf16_bulk");
21057 let ni = n as i64;
21058 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21059 let __s_b = self.gpu.stream();
21060 let mut b = __s_b.launch_builder(&f);
21061 b.arg(x).arg(dst).arg(&ni);
21062 unsafe {
21063 b.launch(cfg)?;
21064 }
21065 Ok(())
21066 }
21067
21068 pub fn f32_to_bf16_into(
21070 &self,
21071 x: &CudaSlice<f32>,
21072 dst: &mut CudaSlice<u8>,
21073 n: usize,
21074 ) -> Result<(), Box<dyn std::error::Error>> {
21075 let f = self.func("f32_to_bf16_bulk");
21076 let ni = n as i64;
21077 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21078 let __s_b = self.gpu.stream();
21079 let mut b = __s_b.launch_builder(&f);
21080 b.arg(x).arg(dst).arg(&ni);
21081 unsafe {
21082 b.launch(cfg)?;
21083 }
21084 Ok(())
21085 }
21086
21087 pub fn gdn_chunk_k123_vl8(
21090 &self,
21091 seqs: &[GdnSeqVl],
21092 n_head: usize,
21093 hk: usize,
21094 wq: Option<&GdnWVl8>,
21095 ) -> Result<(), Box<dyn std::error::Error>> {
21096 let b = seqs.len();
21097 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
21098 let mut packed = [GdnSeqVl::default(); 8];
21099 packed[..b].copy_from_slice(seqs);
21100 let v = GdnVl8(packed);
21101 let (hi, ci) = (n_head as i32, 32i32);
21102 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21103 {
21104 let f = self.func("gdn_chunk_cumgate_vl");
21105 let cfg = LaunchConfig {
21106 grid_dim: (max_nc, n_head as u32, b as u32),
21107 block_dim: (32, 1, 1),
21108 shared_mem_bytes: 0,
21109 };
21110 let __s_lb = self.gpu.stream();
21111 let mut lb = __s_lb.launch_builder(&f);
21112 lb.arg(&v).arg(&hi).arg(&ci);
21113 unsafe {
21114 lb.launch(cfg)?;
21115 }
21116 }
21117 let hki = hk as i32;
21118 if let Some(w) = wq {
21119 let f = self.func("gdn_k2_wgmma_vl");
21121 let cfg = LaunchConfig {
21122 grid_dim: (max_nc, n_head as u32, b as u32),
21123 block_dim: (128, 1, 1),
21124 shared_mem_bytes: 0,
21125 };
21126 let __s_lb = self.gpu.stream();
21127 let mut lb = __s_lb.launch_builder(&f);
21128 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
21129 unsafe {
21130 lb.launch(cfg)?;
21131 }
21132 } else {
21133 let f = self.func("gdn_chunk_attn_vl");
21134 let cfg = LaunchConfig {
21135 grid_dim: (max_nc, n_head as u32, b as u32),
21136 block_dim: (256, 1, 1),
21137 shared_mem_bytes: 0,
21138 };
21139 let __s_lb = self.gpu.stream();
21140 let mut lb = __s_lb.launch_builder(&f);
21141 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21142 unsafe {
21143 lb.launch(cfg)?;
21144 }
21145 }
21146 {
21147 let f = self.func("gdn_chunk_solve32_vl");
21148 let cfg = LaunchConfig {
21149 grid_dim: (max_nc, n_head as u32, b as u32),
21150 block_dim: (256, 1, 1),
21151 shared_mem_bytes: 0,
21152 };
21153 let __s_lb = self.gpu.stream();
21154 let mut lb = __s_lb.launch_builder(&f);
21155 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21156 unsafe {
21157 lb.launch(cfg)?;
21158 }
21159 }
21160 Ok(())
21161 }
21162
21163 #[allow(clippy::too_many_arguments)]
21167 pub fn gdn_prep_vl8(
21168 &self,
21169 seqs: &[GdnPrepVl],
21170 conv_w: &CudaSlice<f32>,
21171 dt_bias: &CudaSlice<f32>,
21172 a: &CudaSlice<f32>,
21173 conv_dim: usize,
21174 d_conv: usize,
21175 d_state: usize,
21176 num_v: usize,
21177 num_k: usize,
21178 key_dim: usize,
21179 hk: usize,
21180 eps: f32,
21181 ) -> Result<(), Box<dyn std::error::Error>> {
21182 let b = seqs.len();
21183 assert!(b >= 1 && b <= 8);
21184 let mut packed = [GdnPrepVl::default(); 8];
21185 packed[..b].copy_from_slice(seqs);
21186 let v = GdnPrepVl8(packed);
21187 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21188 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
21189 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
21190 assert!(
21191 conv_fuse || hk == num_v,
21192 "de-broadcast requires the fused conv"
21193 );
21194 if conv_fuse {
21195 let f = self.func("ssm_conv1d_gdn_state_vl");
21196 let cfg = LaunchConfig {
21197 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21198 block_dim: (256, 1, 1),
21199 shared_mem_bytes: 0,
21200 };
21201 let (dsi, nvi, nki, kdi, hki) = (
21202 d_state as i32,
21203 num_v as i32,
21204 num_k as i32,
21205 key_dim as i32,
21206 hk as i32,
21207 );
21208 let __s_lb = self.gpu.stream();
21209 let mut lb = __s_lb.launch_builder(&f);
21210 lb.arg(&v)
21211 .arg(conv_w)
21212 .arg(&cdi)
21213 .arg(&dci)
21214 .arg(&dsi)
21215 .arg(&nvi)
21216 .arg(&nki)
21217 .arg(&kdi)
21218 .arg(&hki);
21219 unsafe {
21220 lb.launch(cfg)?;
21221 }
21222 } else {
21223 let f = self.func("ssm_conv1d_tm_state_vl");
21224 let cfg = LaunchConfig {
21225 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21226 block_dim: (256, 1, 1),
21227 shared_mem_bytes: 0,
21228 };
21229 let __s_lb = self.gpu.stream();
21230 let mut lb = __s_lb.launch_builder(&f);
21231 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21232 unsafe {
21233 lb.launch(cfg)?;
21234 }
21235 }
21236 {
21237 let f = self.func("ssm_conv_ring_update_vl");
21238 let n = (conv_dim * (d_conv - 1)) as u32;
21239 let cfg = LaunchConfig {
21240 grid_dim: (n.div_ceil(256), 1, b as u32),
21241 block_dim: (256, 1, 1),
21242 shared_mem_bytes: 0,
21243 };
21244 let __s_lb = self.gpu.stream();
21245 let mut lb = __s_lb.launch_builder(&f);
21246 lb.arg(&v).arg(&cdi).arg(&dci);
21247 unsafe {
21248 lb.launch(cfg)?;
21249 }
21250 }
21251 if !conv_fuse {
21252 let f = self.func("qkv_to_gdn_repack_vl");
21253 let n = max_t * (num_v * d_state) as u32;
21254 let cfg = LaunchConfig {
21255 grid_dim: (n.div_ceil(256), 1, b as u32),
21256 block_dim: (256, 1, 1),
21257 shared_mem_bytes: 0,
21258 };
21259 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21260 let __s_lb = self.gpu.stream();
21261 let mut lb = __s_lb.launch_builder(&f);
21262 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21263 unsafe {
21264 lb.launch(cfg)?;
21265 }
21266 }
21267 if Self::l2_v2_on(d_state) {
21268 let f = self.func("gdn_l2_v2_vl");
21269 let cfg = LaunchConfig {
21270 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21271 block_dim: (256, 1, 1),
21272 shared_mem_bytes: 0,
21273 };
21274 let (dsi, nvi) = (d_state as i32, hk as i32);
21275 let __s_lb = self.gpu.stream();
21276 let mut lb = __s_lb.launch_builder(&f);
21277 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21278 unsafe {
21279 lb.launch(cfg)?;
21280 }
21281 } else {
21282 let f = self.func("gdn_l2_vl");
21283 let cfg = LaunchConfig {
21284 grid_dim: (max_t * hk as u32, 2, b as u32),
21285 block_dim: (256, 1, 1),
21286 shared_mem_bytes: 0,
21287 };
21288 let (dsi, nvi) = (d_state as i32, hk as i32);
21289 let __s_lb = self.gpu.stream();
21290 let mut lb = __s_lb.launch_builder(&f);
21291 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21292 unsafe {
21293 lb.launch(cfg)?;
21294 }
21295 }
21296 {
21297 let f = self.func("gdn_gate_prep_vl");
21298 let n = max_t * num_v as u32;
21299 let cfg = LaunchConfig {
21300 grid_dim: (n.div_ceil(256), 1, b as u32),
21301 block_dim: (256, 1, 1),
21302 shared_mem_bytes: 0,
21303 };
21304 let nvi = num_v as i32;
21305 let __s_lb = self.gpu.stream();
21306 let mut lb = __s_lb.launch_builder(&f);
21307 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21308 unsafe {
21309 lb.launch(cfg)?;
21310 }
21311 }
21312 Ok(())
21313 }
21314
21315 pub fn gdn_mirror_vl8(
21317 &self,
21318 seqs: &[GdnSeqVl],
21319 n_head: usize,
21320 which: i32,
21321 hk: usize,
21322 ) -> Result<(), Box<dyn std::error::Error>> {
21323 let b = seqs.len();
21324 assert!(b >= 1 && b <= 8);
21325 let mut packed = [GdnSeqVl::default(); 8];
21326 packed[..b].copy_from_slice(seqs);
21327 let v = GdnVl8(packed);
21328 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21329 let max_n = seqs
21330 .iter()
21331 .map(|s| {
21332 if which == 0 {
21333 s.t as i64 * ept as i64
21334 } else {
21335 s.nc as i64 * ept as i64 * 32
21336 }
21337 })
21338 .max()
21339 .unwrap();
21340 let f = self.func("gdn_mirror_vl");
21341 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21342 let cfg = LaunchConfig {
21343 grid_dim: (blocks, 1, b as u32),
21344 block_dim: (256, 1, 1),
21345 shared_mem_bytes: 0,
21346 };
21347 let __s_lb = self.gpu.stream();
21348 let mut lb = __s_lb.launch_builder(&f);
21349 lb.arg(&v).arg(&ept).arg(&which);
21350 unsafe {
21351 lb.launch(cfg)?;
21352 }
21353 Ok(())
21354 }
21355
21356 pub fn gdn_tail_vl8(
21358 &self,
21359 seqs: &[GdnPrepVl],
21360 norm_w: &CudaSlice<f32>,
21361 d_state: usize,
21362 num_v: usize,
21363 eps: f32,
21364 ) -> Result<(), Box<dyn std::error::Error>> {
21365 let b = seqs.len();
21366 assert!(b >= 1 && b <= 8);
21367 let mut packed = [GdnPrepVl::default(); 8];
21368 packed[..b].copy_from_slice(seqs);
21369 let v = GdnPrepVl8(packed);
21370 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21371 let f = self.func("gated_rmsnorm_f16out_vl");
21372 let cfg = LaunchConfig {
21374 grid_dim: (max_t * num_v as u32, 1, b as u32),
21375 block_dim: (128, 1, 1),
21376 shared_mem_bytes: 0,
21377 };
21378 let (dsi, nvi) = (d_state as i32, num_v as i32);
21379 let __s_lb = self.gpu.stream();
21380 let mut lb = __s_lb.launch_builder(&f);
21381 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21382 unsafe {
21383 lb.launch(cfg)?;
21384 }
21385 Ok(())
21386 }
21387
21388 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
21391 use cudarc::driver::DevicePtr;
21392 let s = self.gpu.stream();
21393 let (p, _g) = x.device_ptr(&s);
21394 p as u64
21395 }
21396 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
21397 use cudarc::driver::DevicePtrMut;
21398 let s = self.gpu.stream();
21399 let (p, _g) = x.device_ptr_mut(&s);
21400 p as u64
21401 }
21402 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
21403 use cudarc::driver::DevicePtr;
21404 let s = self.gpu.stream();
21405 let (p, _g) = x.device_ptr(&s);
21406 p as u64
21407 }
21408 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
21409 use cudarc::driver::DevicePtr;
21410 let s = self.gpu.stream();
21411 let (p, _g) = x.device_ptr(&s);
21412 p as u64
21413 }
21414
21415 pub fn gdn_chunk_vl8(
21419 &self,
21420 seqs: &[GdnSeqVl],
21421 n_head: usize,
21422 scale: f32,
21423 hk: usize,
21424 wq: Option<&GdnWVl8>,
21425 ) -> Result<(), Box<dyn std::error::Error>> {
21426 const NSPLIT: u32 = 4;
21427 let b = seqs.len();
21428 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
21429 let mut packed = [GdnSeqVl::default(); 8];
21430 packed[..b].copy_from_slice(seqs);
21431 let v = GdnVl8(packed);
21432 let (hi, ci) = (n_head as i32, 32i32);
21433 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21434 let hki = hk as i32;
21435 if let Some(w) = wq {
21436 let f = self.func("gdn_k45_wgmma_vl");
21438 let cfg = LaunchConfig {
21439 grid_dim: (n_head as u32, NSPLIT, b as u32),
21440 block_dim: (256, 1, 1),
21441 shared_mem_bytes: 0,
21442 };
21443 let __s_lb = self.gpu.stream();
21444 let mut lb = __s_lb.launch_builder(&f);
21445 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
21446 unsafe {
21447 lb.launch(cfg)?;
21448 }
21449 let _ = max_nc;
21450 return Ok(());
21451 }
21452 {
21453 let f = self.func("gdn_chunk_state_mma_vl");
21454 let cfg = LaunchConfig {
21455 grid_dim: (n_head as u32, NSPLIT, b as u32),
21456 block_dim: (256, 1, 1),
21457 shared_mem_bytes: 0,
21458 };
21459 let __s_lb = self.gpu.stream();
21460 let mut lb = __s_lb.launch_builder(&f);
21461 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21462 unsafe {
21463 lb.launch(cfg)?;
21464 }
21465 }
21466 {
21467 let f = self.func("gdn_chunk_output_mma_vl");
21468 let cfg = LaunchConfig {
21469 grid_dim: (max_nc, n_head as u32, b as u32),
21470 block_dim: (256, 1, 1),
21471 shared_mem_bytes: 0,
21472 };
21473 let __s_lb = self.gpu.stream();
21474 let mut lb = __s_lb.launch_builder(&f);
21475 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
21476 unsafe {
21477 lb.launch(cfg)?;
21478 }
21479 }
21480 Ok(())
21481 }
21482 pub fn gdn_scan_chunked(
21483 &self,
21484 q: &CudaSlice<f32>,
21485 k: &CudaSlice<f32>,
21486 v: &CudaSlice<f32>,
21487 g: &CudaSlice<f32>,
21488 beta: &CudaSlice<f32>,
21489 kb16_pre: Option<&CudaSlice<u8>>,
21490 qb16_pre: Option<&CudaSlice<u8>>,
21491 state_in: &CudaSlice<f32>,
21492 state_out: &mut CudaSlice<f32>,
21493 o: &mut CudaSlice<f32>,
21494 n_head: usize,
21495 t: usize,
21496 scale: f32,
21497 c: usize,
21498 hk: usize,
21499 ) -> Result<(), Box<dyn std::error::Error>> {
21500 const D: usize = 128;
21501 const NSPLIT: u32 = 4;
21502 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
21503 let h = n_head;
21504 let nc = (t + c - 1) / c;
21505 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21506 let gdn_mma_pre = !portable_mma_gated()
21511 && c == 32
21512 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21513 Ok("1") => true,
21514 Ok("0") => false,
21515 _ => gdn_mma_default_on(),
21516 };
21517 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
21518 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
21519 } else {
21520 None
21521 };
21522 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
21527 && gdn_mma_pre
21528 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
21529 let nk = t * hk * D;
21530 let mut kb16_local: Option<CudaSlice<u8>> = None;
21531 if gdn_mma_pre && kb16_pre.is_none() {
21532 let mut kb = self.alloc_u8_uninit(nk * 2)?;
21533 let f = self.func("f32_to_bf16_bulk");
21534 let n2 = nk as i64;
21535 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21536 let __s_b = self.gpu.stream();
21537 let mut b = __s_b.launch_builder(&f);
21538 b.arg(k).arg(&mut kb).arg(&n2);
21539 unsafe {
21540 b.launch(cfg2)?;
21541 }
21542 kb16_local = Some(kb);
21543 }
21544 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
21545 if let Some(kb) = kb16_pre {
21546 assert!(kb.len() >= nk * 2, "kb16_pre too small");
21547 }
21548 let mut qb16: Option<CudaSlice<u8>> = None;
21549 let mut pb16: Option<CudaSlice<u8>> = None;
21550 if gdn_wgmma_pre {
21551 if qb16_pre.is_none() {
21554 let mut qb = self.alloc_u8_uninit(nk * 2)?;
21555 let f = self.func("f32_to_bf16_bulk");
21556 let n2 = nk as i64;
21557 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21558 let __s_b = self.gpu.stream();
21559 let mut b = __s_b.launch_builder(&f);
21560 b.arg(q).arg(&mut qb).arg(&n2);
21561 unsafe {
21562 b.launch(cfg2)?;
21563 }
21564 qb16 = Some(qb);
21565 } else if let Some(qb) = qb16_pre {
21566 assert!(qb.len() >= nk * 2, "qb16_pre too small");
21567 }
21568 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
21569 }
21570 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
21571 let k2w = if gdn_wgmma_pre {
21572 Some((
21573 *qb16_ref0.as_ref().unwrap(),
21574 *kb16_ref0.as_ref().unwrap(),
21575 pb16.as_mut().unwrap(),
21576 ))
21577 } else {
21578 None
21579 };
21580 let (gcum, p, u, w) =
21581 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
21582 let _ = &w;
21583 let mut y = self.uninit(nc * h * c * D)?;
21584 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
21598 && c == 32
21599 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21600 Ok("1") => true,
21601 Ok("0") => false,
21602 _ => gdn_mma_default_on(),
21603 };
21604 if gdn_mma {
21605 let wb16 = wb16_pre
21606 .take()
21607 .expect("mma path pre-allocates wb16 (K3 store fold)");
21608 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
21609 if gdn_wgmma_pre {
21621 let qb16 = qb16_ref0.unwrap();
21623 let pb16 = pb16.as_ref().unwrap();
21624 {
21625 let f = self.func("gdn_k45_wgmma");
21626 let cfg = LaunchConfig {
21627 grid_dim: (h as u32, 4, 1),
21628 block_dim: (256, 1, 1),
21629 shared_mem_bytes: 0,
21630 };
21631 let hki = hk as i32;
21632 let __s_b = self.gpu.stream();
21633 let mut b = __s_b.launch_builder(&f);
21634 b.arg(kb16_ref)
21635 .arg(&gcum)
21636 .arg(beta)
21637 .arg(&u)
21638 .arg(&wb16)
21639 .arg(qb16)
21640 .arg(pb16)
21641 .arg(o)
21642 .arg(&scale)
21643 .arg(state_in)
21644 .arg(&mut *state_out)
21645 .arg(&hi)
21646 .arg(&ti)
21647 .arg(&ci)
21648 .arg(&hki);
21649 unsafe {
21650 b.launch(cfg)?;
21651 }
21652 }
21653 return Ok(());
21654 }
21655 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
21659 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
21660 {
21661 let f = self.func("gdn_chunk_state_mma");
21662 let cfg = LaunchConfig {
21663 grid_dim: (h as u32, NSPLIT, 1),
21664 block_dim: (256, 1, 1),
21665 shared_mem_bytes: 0,
21666 };
21667 let hki = hk as i32;
21668 let __s_b = self.gpu.stream();
21669 let mut b = __s_b.launch_builder(&f);
21670 b.arg(kb16_ref)
21671 .arg(&gcum)
21672 .arg(beta)
21673 .arg(&u)
21674 .arg(&wb16)
21675 .arg(&mut y16)
21676 .arg(&mut ssnap16)
21677 .arg(state_in)
21678 .arg(&mut *state_out)
21679 .arg(&hi)
21680 .arg(&ti)
21681 .arg(&ci)
21682 .arg(&hki);
21683 unsafe {
21684 b.launch(cfg)?;
21685 }
21686 }
21687 {
21688 let f = self.func("gdn_chunk_output_mma");
21690 let jt = ((c + 31) / 32) as u32;
21691 let cfg = LaunchConfig {
21692 grid_dim: (nc as u32, h as u32, jt),
21693 block_dim: (256, 1, 1),
21694 shared_mem_bytes: 0,
21695 };
21696 let hki = hk as i32;
21697 let __s_b = self.gpu.stream();
21698 let mut b = __s_b.launch_builder(&f);
21699 b.arg(q)
21700 .arg(&gcum)
21701 .arg(&p)
21702 .arg(&y16)
21703 .arg(&ssnap16)
21704 .arg(o)
21705 .arg(&hi)
21706 .arg(&ti)
21707 .arg(&ci)
21708 .arg(&scale)
21709 .arg(&hki);
21710 unsafe {
21711 b.launch(cfg)?;
21712 }
21713 }
21714 return Ok(());
21715 }
21716 {
21717 let f = self.func("gdn_chunk_state_f32");
21719 let cfg = LaunchConfig {
21720 grid_dim: (h as u32, NSPLIT, 1),
21721 block_dim: (256, 1, 1),
21722 shared_mem_bytes: 0,
21723 };
21724 let __s_b = self.gpu.stream();
21725 let mut b = __s_b.launch_builder(&f);
21726 b.arg(k)
21727 .arg(&gcum)
21728 .arg(beta)
21729 .arg(&u)
21730 .arg(&w)
21731 .arg(&mut y)
21732 .arg(&mut ssnap)
21733 .arg(state_in)
21734 .arg(&mut *state_out)
21735 .arg(&hi)
21736 .arg(&ti)
21737 .arg(&ci);
21738 unsafe {
21739 b.launch(cfg)?;
21740 }
21741 }
21742 {
21743 let f = self.func("gdn_chunk_output_f32");
21745 let jt = ((c + 31) / 32) as u32;
21746 let cfg = LaunchConfig {
21747 grid_dim: (nc as u32, h as u32, jt),
21748 block_dim: (256, 1, 1),
21749 shared_mem_bytes: 0,
21750 };
21751 let __s_b = self.gpu.stream();
21752 let mut b = __s_b.launch_builder(&f);
21753 b.arg(q)
21754 .arg(&gcum)
21755 .arg(&p)
21756 .arg(&y)
21757 .arg(&ssnap)
21758 .arg(o)
21759 .arg(&hi)
21760 .arg(&ti)
21761 .arg(&ci)
21762 .arg(&scale);
21763 unsafe {
21764 b.launch(cfg)?;
21765 }
21766 }
21767 Ok(())
21768 }
21769
21770 #[allow(clippy::too_many_arguments)]
21779 #[allow(clippy::too_many_arguments)]
21780 pub fn gdn_scan_prefill(
21781 &self,
21782 q: &CudaSlice<f32>,
21783 k: &CudaSlice<f32>,
21784 v: &CudaSlice<f32>,
21785 g: &CudaSlice<f32>,
21786 beta: &CudaSlice<f32>,
21787 kb16_pre: Option<&CudaSlice<u8>>,
21788 qb16_pre: Option<&CudaSlice<u8>>,
21789 state_in: &CudaSlice<f32>,
21790 state_out: &mut CudaSlice<f32>,
21791 o: &mut CudaSlice<f32>,
21792 n_head: usize,
21793 t: usize,
21794 scale: f32,
21795 hk: usize,
21796 ) -> Result<(), Box<dyn std::error::Error>> {
21797 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
21798 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
21799 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
21800 }
21801 if Self::gdn_chunked_enabled() && t >= 16 {
21802 self.gdn_scan_chunked(
21803 q,
21804 k,
21805 v,
21806 g,
21807 beta,
21808 kb16_pre,
21809 qb16_pre,
21810 state_in,
21811 state_out,
21812 o,
21813 n_head,
21814 t,
21815 scale,
21816 Self::gdn_chunk_size(),
21817 hk,
21818 )
21819 } else {
21820 assert!(
21821 hk == n_head,
21822 "s128 scan is broadcast-only (prep guarantees by predicate)"
21823 );
21824 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
21825 }
21826 }
21827
21828 #[allow(clippy::too_many_arguments)]
21830 fn gdn_scan_diff(
21831 &self,
21832 q: &CudaSlice<f32>,
21833 k: &CudaSlice<f32>,
21834 v: &CudaSlice<f32>,
21835 g: &CudaSlice<f32>,
21836 beta: &CudaSlice<f32>,
21837 state_in: &CudaSlice<f32>,
21838 state_out: &mut CudaSlice<f32>,
21839 o: &mut CudaSlice<f32>,
21840 n_head: usize,
21841 t: usize,
21842 scale: f32,
21843 ) -> Result<(), Box<dyn std::error::Error>> {
21844 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
21845 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
21846 let mut o_c = self.uninit(o.len())?;
21847 let mut st_c = self.uninit(state_out.len())?;
21848 self.gdn_scan_chunked(
21849 q,
21850 k,
21851 v,
21852 g,
21853 beta,
21854 None,
21855 None,
21856 state_in,
21857 &mut st_c,
21858 &mut o_c,
21859 n_head,
21860 t,
21861 scale,
21862 Self::gdn_chunk_size(),
21863 n_head,
21864 )?;
21865 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
21866 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
21867 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
21868 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
21869 let mut max_abs = 0f32;
21870 let mut max_rel = 0f32;
21871 let mut sum_rel = 0f64;
21872 for (x, y) in a.iter().zip(b) {
21873 let ad = (x - y).abs();
21874 let rel = ad / x.abs().max(y.abs()).max(1e-3);
21875 if ad > max_abs {
21876 max_abs = ad;
21877 }
21878 if rel > max_rel {
21879 max_rel = rel;
21880 }
21881 sum_rel += rel as f64;
21882 }
21883 (max_abs, max_rel, sum_rel / a.len() as f64)
21884 };
21885 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
21886 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
21887 println!(
21888 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
21889 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
21890 Self::gdn_chunk_size()
21891 );
21892 Ok(())
21893 }
21894
21895 pub fn gdn_glog(
21897 &self,
21898 alpha: &CudaSlice<f32>,
21899 dt_bias: &CudaSlice<f32>,
21900 a: &CudaSlice<f32>,
21901 g_log: &mut CudaSlice<f32>,
21902 n_head: usize,
21903 t: usize,
21904 ) -> Result<(), Box<dyn std::error::Error>> {
21905 let f = self.func("gdn_glog_f32");
21906 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
21907 let (h, ti) = (n_head as i32, t as i32);
21908 let __s_b = self.gpu.stream();
21909 let mut b = __s_b.launch_builder(&f);
21910 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
21911 unsafe {
21912 b.launch(cfg)?;
21913 }
21914 Ok(())
21915 }
21916
21917 pub fn sigmoid_v(
21920 &self,
21921 x: &cudarc::driver::CudaView<f32>,
21922 y: &mut CudaSlice<f32>,
21923 n: usize,
21924 ) -> Result<(), Box<dyn std::error::Error>> {
21925 let f = self.func("sigmoid_f32");
21926 let cfg = LaunchConfig::for_num_elems(n as u32);
21927 let ni = n as i32;
21928 let __s_b = self.gpu.stream();
21929 let mut b = __s_b.launch_builder(&f);
21930 b.arg(x).arg(y).arg(&ni);
21931 unsafe {
21932 b.launch(cfg)?;
21933 }
21934 Ok(())
21935 }
21936
21937 pub fn gdn_glog_v(
21938 &self,
21939 alpha: &cudarc::driver::CudaView<f32>,
21940 dt_bias: &CudaSlice<f32>,
21941 a: &CudaSlice<f32>,
21942 g_log: &mut CudaSlice<f32>,
21943 n_head: usize,
21944 t: usize,
21945 ) -> Result<(), Box<dyn std::error::Error>> {
21946 let f = self.func("gdn_glog_f32");
21947 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
21948 let (h, ti) = (n_head as i32, t as i32);
21949 let __s_b = self.gpu.stream();
21950 let mut b = __s_b.launch_builder(&f);
21951 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
21952 unsafe {
21953 b.launch(cfg)?;
21954 }
21955 Ok(())
21956 }
21957
21958 pub fn sigmoid(
21959 &self,
21960 x: &CudaSlice<f32>,
21961 y: &mut CudaSlice<f32>,
21962 n: usize,
21963 ) -> Result<(), Box<dyn std::error::Error>> {
21964 let f = self.func("sigmoid_f32");
21965 let cfg = LaunchConfig::for_num_elems(n as u32);
21966 let ni = n as i32;
21967 let __s_b = self.gpu.stream();
21968 let mut b = __s_b.launch_builder(&f);
21969 b.arg(x).arg(y).arg(&ni);
21970 unsafe {
21971 b.launch(cfg)?;
21972 }
21973 Ok(())
21974 }
21975
21976 pub fn sig_mul_f16out(
21979 &self,
21980 a: &CudaSlice<f32>,
21981 g: &CudaSlice<f32>,
21982 dst: &mut CudaSlice<f32>,
21983 dst16: &mut CudaSlice<u8>,
21984 n: usize,
21985 ) -> Result<(), Box<dyn std::error::Error>> {
21986 let f = self.func("sig_mul_f16out_f32");
21987 let cfg = LaunchConfig::for_num_elems(n as u32);
21988 let ni = n as i32;
21989 let __s_b = self.gpu.stream();
21990 let mut b = __s_b.launch_builder(&f);
21991 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
21992 unsafe {
21993 b.launch(cfg)?;
21994 }
21995 Ok(())
21996 }
21997
21998 #[allow(clippy::too_many_arguments)]
22007 pub fn attn_head_gate(
22008 &self,
22009 a: &CudaSlice<f32>,
22010 g: &CudaSlice<f32>,
22011 dst: &mut CudaSlice<f32>,
22012 dst16: Option<&mut CudaSlice<u8>>,
22013 head_dim: usize,
22014 n_head: usize,
22015 t: usize,
22016 ) -> Result<(), Box<dyn std::error::Error>> {
22017 let f = self.func("attn_head_gate_f32");
22018 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22019 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22020 let d16: u64 = match dst16 {
22022 Some(d) => self.addr_u8(d),
22023 None => 0,
22024 };
22025 let __s_b = self.gpu.stream();
22026 let mut b = __s_b.launch_builder(&f);
22027 b.arg(a)
22028 .arg(g)
22029 .arg(dst)
22030 .arg(&d16)
22031 .arg(&hd)
22032 .arg(&nh)
22033 .arg(&ti);
22034 unsafe {
22035 b.launch(cfg)?;
22036 }
22037 Ok(())
22038 }
22039
22040 #[allow(clippy::too_many_arguments)]
22049 pub fn swiglu_clamped_mul_scaled(
22050 &self,
22051 gate: &CudaSlice<f32>,
22052 up: &CudaSlice<f32>,
22053 gs: f32,
22054 us: f32,
22055 limit: f32,
22056 dst: &mut CudaSlice<f32>,
22057 n: usize,
22058 ) -> Result<(), Box<dyn std::error::Error>> {
22059 debug_assert!(
22060 limit > 1e-6,
22061 "swiglu_clamped needs a live limit; use silu_mul_scaled"
22062 );
22063 let f = self.func("swiglu_clamped_mul_scaled_f32");
22064 let cfg = LaunchConfig::for_num_elems(n as u32);
22065 let ni = n as i32;
22066 let __s_b = self.gpu.stream();
22067 let mut b = __s_b.launch_builder(&f);
22068 b.arg(gate)
22069 .arg(up)
22070 .arg(&gs)
22071 .arg(&us)
22072 .arg(&limit)
22073 .arg(dst)
22074 .arg(&ni);
22075 unsafe {
22076 b.launch(cfg)?;
22077 }
22078 Ok(())
22079 }
22080
22081 pub fn gated_rmsnorm(
22083 &self,
22084 o: &CudaSlice<f32>,
22085 w: &CudaSlice<f32>,
22086 z: &CudaSlice<f32>,
22087 dst: &mut CudaSlice<f32>,
22088 ncols: usize,
22089 nrows: usize,
22090 eps: f32,
22091 ) -> Result<(), Box<dyn std::error::Error>> {
22092 let f = self.func("gated_rmsnorm_f32");
22093 let cfg = LaunchConfig {
22094 grid_dim: (nrows as u32, 1, 1),
22095 block_dim: (128, 1, 1),
22096 shared_mem_bytes: 0,
22097 };
22098 let (nc, e) = (ncols as i32, eps);
22099 let __s_b = self.gpu.stream();
22100 let mut b = __s_b.launch_builder(&f);
22101 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22102 unsafe {
22103 b.launch(cfg)?;
22104 }
22105 Ok(())
22106 }
22107
22108 pub fn gated_rmsnorm_f16out(
22111 &self,
22112 o: &CudaSlice<f32>,
22113 w: &CudaSlice<f32>,
22114 z: &CudaSlice<f32>,
22115 dst: &mut CudaSlice<f32>,
22116 dst16: &mut CudaSlice<u8>,
22117 ncols: usize,
22118 nrows: usize,
22119 eps: f32,
22120 ) -> Result<(), Box<dyn std::error::Error>> {
22121 let f = self.func("gated_rmsnorm_f16out_f32");
22122 let cfg = LaunchConfig {
22124 grid_dim: (nrows as u32, 1, 1),
22125 block_dim: (128, 1, 1),
22126 shared_mem_bytes: 0,
22127 };
22128 let (nc, e) = (ncols as i32, eps);
22129 let __s_b = self.gpu.stream();
22130 let mut b = __s_b.launch_builder(&f);
22131 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22132 unsafe {
22133 b.launch(cfg)?;
22134 }
22135 Ok(())
22136 }
22137
22138 #[allow(clippy::too_many_arguments)]
22142 pub fn add_rms_norm_zq8(
22143 &self,
22144 a: &CudaSlice<f32>,
22145 b_in: &CudaSlice<f32>,
22146 w: &CudaSlice<f32>,
22147 res: &mut CudaSlice<f32>,
22148 z: &mut CudaSlice<f32>,
22149 ncols: usize,
22150 nrows: usize,
22151 eps: f32,
22152 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22153 assert!(ncols % 32 == 0);
22154 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
22155 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22156 let f = self.func("add_rms_norm_zq8");
22157 let cfg = LaunchConfig {
22158 grid_dim: (nrows as u32, 1, 1),
22159 block_dim: (1024, 1, 1),
22160 shared_mem_bytes: 0,
22161 };
22162 let (nc, ep) = (ncols as i32, eps);
22163 let __s_b = self.gpu.stream();
22164 let mut b = __s_b.launch_builder(&f);
22165 b.arg(a)
22166 .arg(b_in)
22167 .arg(w)
22168 .arg(res)
22169 .arg(z)
22170 .arg(&mut q)
22171 .arg(&mut d)
22172 .arg(&nc)
22173 .arg(&ep);
22174 unsafe {
22175 b.launch(cfg)?;
22176 }
22177 Ok((q, d))
22178 }
22179
22180 pub fn gated_rmsnorm_zv(
22185 &self,
22186 o: &CudaSlice<f32>,
22187 w: &CudaSlice<f32>,
22188 z: &cudarc::driver::CudaView<f32>,
22189 dst: &mut CudaSlice<f32>,
22190 ncols: usize,
22191 nrows: usize,
22192 eps: f32,
22193 ) -> Result<(), Box<dyn std::error::Error>> {
22194 let f = self.func("gated_rmsnorm_f32");
22195 let cfg = LaunchConfig {
22196 grid_dim: (nrows as u32, 1, 1),
22197 block_dim: (128, 1, 1),
22198 shared_mem_bytes: 0,
22199 };
22200 let (nc, e) = (ncols as i32, eps);
22201 let __s_b = self.gpu.stream();
22202 let mut b = __s_b.launch_builder(&f);
22203 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22204 unsafe {
22205 b.launch(cfg)?;
22206 }
22207 Ok(())
22208 }
22209
22210 pub fn gated_rmsnorm_f16out_zv(
22211 &self,
22212 o: &CudaSlice<f32>,
22213 w: &CudaSlice<f32>,
22214 z: &cudarc::driver::CudaView<f32>,
22215 dst: &mut CudaSlice<f32>,
22216 dst16: &mut CudaSlice<u8>,
22217 ncols: usize,
22218 nrows: usize,
22219 eps: f32,
22220 ) -> Result<(), Box<dyn std::error::Error>> {
22221 let f = self.func("gated_rmsnorm_f16out_f32");
22222 let cfg = LaunchConfig {
22224 grid_dim: (nrows as u32, 1, 1),
22225 block_dim: (128, 1, 1),
22226 shared_mem_bytes: 0,
22227 };
22228 let (nc, e) = (ncols as i32, eps);
22229 let __s_b = self.gpu.stream();
22230 let mut b = __s_b.launch_builder(&f);
22231 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22232 unsafe {
22233 b.launch(cfg)?;
22234 }
22235 Ok(())
22236 }
22237
22238 pub fn gated_rmsnorm_q8_1(
22239 &self,
22240 o: &CudaSlice<f32>,
22241 w: &CudaSlice<f32>,
22242 z: &CudaSlice<f32>,
22243 ncols: usize,
22244 nrows: usize,
22245 eps: f32,
22246 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22247 assert!(ncols % 32 == 0);
22248 let f = self.func("gated_rmsnorm_q8_1");
22249 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22250 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22251 let cfg = LaunchConfig {
22252 grid_dim: (nrows as u32, 1, 1),
22253 block_dim: (128, 1, 1),
22254 shared_mem_bytes: 0,
22255 };
22256 let (nc, ep) = (ncols as i32, eps);
22257 let __s_b = self.gpu.stream();
22258 let mut b = __s_b.launch_builder(&f);
22259 b.arg(o)
22260 .arg(w)
22261 .arg(z)
22262 .arg(&mut out_q)
22263 .arg(&mut out_d)
22264 .arg(&nc)
22265 .arg(&ep);
22266 unsafe {
22267 b.launch(cfg)?;
22268 }
22269 Ok((out_q, out_d))
22270 }
22271
22272 pub fn transpose(
22274 &self,
22275 inp: &CudaSlice<f32>,
22276 rows: usize,
22277 cols: usize,
22278 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22279 let f = self.func("transpose_f32");
22280 let mut out = self.zeros(rows * cols)?;
22281 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22282 let (r, c) = (rows as i32, cols as i32);
22283 let __s_b = self.gpu.stream();
22284 let mut b = __s_b.launch_builder(&f);
22285 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22286 unsafe {
22287 b.launch(cfg)?;
22288 }
22289 Ok(out)
22290 }
22291
22292 pub fn repeat_heads(
22294 &self,
22295 inp: &CudaSlice<f32>,
22296 out: &mut CudaSlice<f32>,
22297 head_dim: usize,
22298 n_in: usize,
22299 n_out: usize,
22300 t: usize,
22301 ) -> Result<(), Box<dyn std::error::Error>> {
22302 let f = self.func("repeat_heads_f32");
22303 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22304 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22305 let __s_b = self.gpu.stream();
22306 let mut b = __s_b.launch_builder(&f);
22307 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22308 unsafe {
22309 b.launch(cfg)?;
22310 }
22311 Ok(())
22312 }
22313
22314 pub fn q_gate_split(
22321 &self,
22322 qf: &CudaSlice<f32>,
22323 q_out: &mut CudaSlice<f32>,
22324 gate_out: &mut CudaSlice<f32>,
22325 head_dim: usize,
22326 n_head: usize,
22327 t: usize,
22328 ) -> Result<(), Box<dyn std::error::Error>> {
22329 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
22330 let out_need = head_dim * n_head * t;
22331 if q_out.len() < out_need || gate_out.len() < out_need {
22332 return Err(format!(
22333 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
22334 q_out.len(),
22335 gate_out.len()
22336 )
22337 .into());
22338 }
22339 let f = self.func("q_gate_split_f32");
22340 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22341 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22342 let __s_b = self.gpu.stream();
22343 let mut b = __s_b.launch_builder(&f);
22344 b.arg(qf)
22345 .arg(q_out)
22346 .arg(gate_out)
22347 .arg(&hd)
22348 .arg(&nh)
22349 .arg(&ti);
22350 unsafe {
22351 b.launch(cfg)?;
22352 }
22353 Ok(())
22354 }
22355
22356 pub fn qkv_to_gdn_repack(
22360 &self,
22361 conv_out: &CudaSlice<f32>,
22362 q_g: &mut CudaSlice<f32>,
22363 k_g: &mut CudaSlice<f32>,
22364 v_g: &mut CudaSlice<f32>,
22365 d_state: usize,
22366 num_v: usize,
22367 num_k: usize,
22368 key_dim: usize,
22369 t: usize,
22370 ) -> Result<(), Box<dyn std::error::Error>> {
22371 let f = self.func("qkv_to_gdn_repack_f32");
22372 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22373 let (ds, nv, nk, kd, ti) = (
22374 d_state as i32,
22375 num_v as i32,
22376 num_k as i32,
22377 key_dim as i32,
22378 t as i32,
22379 );
22380 let __s_b = self.gpu.stream();
22381 let mut b = __s_b.launch_builder(&f);
22382 b.arg(conv_out)
22383 .arg(q_g)
22384 .arg(k_g)
22385 .arg(v_g)
22386 .arg(&ds)
22387 .arg(&nv)
22388 .arg(&nk)
22389 .arg(&kd)
22390 .arg(&ti);
22391 unsafe {
22392 b.launch(cfg)?;
22393 }
22394 Ok(())
22395 }
22396
22397 pub fn conv_left_pad(
22400 &self,
22401 src: &CudaSlice<f32>,
22402 dst: &mut CudaSlice<f32>,
22403 conv_dim: usize,
22404 t: usize,
22405 pad: usize,
22406 ) -> Result<(), Box<dyn std::error::Error>> {
22407 let f = self.func("conv_left_pad_f32");
22408 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
22409 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
22410 let __s_b = self.gpu.stream();
22411 let mut b = __s_b.launch_builder(&f);
22412 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
22413 unsafe {
22414 b.launch(cfg)?;
22415 }
22416 Ok(())
22417 }
22418
22419 pub fn conv_assemble_and_roll(
22423 &self,
22424 qkv_col: &CudaSlice<f32>,
22425 conv_state: &mut CudaSlice<f32>,
22426 conv_in: &mut CudaSlice<f32>,
22427 conv_dim: usize,
22428 pad: usize,
22429 ) -> Result<(), Box<dyn std::error::Error>> {
22430 let f = self.func("conv_assemble_and_roll_f32");
22431 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22432 let (cd, p) = (conv_dim as i32, pad as i32);
22433 let __s_b = self.gpu.stream();
22434 let mut b = __s_b.launch_builder(&f);
22435 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
22436 unsafe {
22437 b.launch(cfg)?;
22438 }
22439 Ok(())
22440 }
22441
22442 pub fn ssm_conv1d_fused_decode(
22448 &self,
22449 qkv_col: &CudaSlice<f32>,
22450 conv_state: &mut CudaSlice<f32>,
22451 w: &CudaSlice<f32>,
22452 conv_out: &mut CudaSlice<f32>,
22453 conv_dim: usize,
22454 d_conv: usize,
22455 ) -> Result<(), Box<dyn std::error::Error>> {
22456 let f = self.func("ssm_conv1d_fused_decode_f32");
22457 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22458 let (cd, dc) = (conv_dim as i32, d_conv as i32);
22459 let __s_b = self.gpu.stream();
22460 let mut b = __s_b.launch_builder(&f);
22461 b.arg(qkv_col)
22462 .arg(conv_state)
22463 .arg(w)
22464 .arg(conv_out)
22465 .arg(&cd)
22466 .arg(&dc);
22467 unsafe {
22468 b.launch(cfg)?;
22469 }
22470 Ok(())
22471 }
22472
22473 pub fn slice_range(
22476 &self,
22477 src: &CudaSlice<f32>,
22478 start: usize,
22479 len: usize,
22480 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22481 let host = self.gpu.stream().clone_dtoh(src)?;
22482 self.gpu.stream().synchronize()?;
22483 Ok(self.htod(&host[start..start + len])?)
22484 }
22485}
22486
22487#[cfg(test)]
22488mod target_dispatch_tests {
22489 use super::legacy_quant_gemm_allowed;
22490
22491 #[test]
22492 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
22493 assert!(legacy_quant_gemm_allowed(false, false, false));
22495 assert!(!legacy_quant_gemm_allowed(false, false, true));
22496 assert!(!legacy_quant_gemm_allowed(true, false, false));
22498 assert!(!legacy_quant_gemm_allowed(true, false, true));
22499 assert!(legacy_quant_gemm_allowed(true, true, false));
22501 assert!(!legacy_quant_gemm_allowed(true, true, true));
22502 }
22503
22504 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
22505 #[test]
22506 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
22507 assert!(!legacy_quant_gemm_allowed(
22508 cfg!(memra_portable_cuda),
22509 cfg!(memra_hopper_mma),
22510 false
22511 ));
22512 }
22513
22514 #[cfg(memra_hopper_mma)]
22515 #[test]
22516 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
22517 assert!(legacy_quant_gemm_allowed(
22518 cfg!(memra_portable_cuda),
22519 cfg!(memra_hopper_mma),
22520 false
22521 ));
22522 assert!(super::portable_mma_gated() == false);
22523 }
22524}
22525
22526impl memra_kv::KvDev for Engine {
22529 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22530 Engine::zeros(self, n)
22531 }
22532 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22533 Engine::uninit(self, n)
22534 }
22535 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
22536 Engine::alloc_u8(self, n)
22537 }
22538 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
22539 Engine::htod_i32(self, v)
22540 }
22541 fn clone_dtod(
22542 &self,
22543 src: &CudaSlice<f32>,
22544 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22545 Engine::clone_dtod(self, src)
22546 }
22547 fn copy_into(
22548 &self,
22549 dst: &mut CudaSlice<f32>,
22550 off: usize,
22551 src: &CudaSlice<f32>,
22552 len: usize,
22553 ) -> Result<(), Box<dyn std::error::Error>> {
22554 Engine::copy_into(self, dst, off, src, len)
22555 }
22556 fn set_i32_one(
22557 &self,
22558 d: &mut CudaSlice<i32>,
22559 v: i32,
22560 ) -> Result<(), Box<dyn std::error::Error>> {
22561 Engine::set_i32_one(self, d, v)
22562 }
22563}
22564
22565#[cfg(test)]
22566mod fused_gate_bounds_tests {
22567 use super::*;
22568
22569 #[test]
22582 #[ignore = "requires a CUDA GPU"]
22583 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
22584 let e = Engine::new(0).unwrap();
22585 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
22586 let fused = 2 * head_dim * n_head * t;
22587 let out_n = head_dim * n_head * t;
22588
22589 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
22591 let mut q = e.uninit(out_n).unwrap();
22592 let mut gate = e.uninit(out_n).unwrap();
22593 let err = e
22594 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
22595 .expect_err("half-width wq must be refused, not read past")
22596 .to_string();
22597 assert!(err.contains("NO fused gate"), "{err}");
22598 assert!(err.contains(&format!("{fused}")), "{err}");
22599
22600 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
22603 let wide = e.htod(&host).unwrap();
22604 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
22605 .expect("full-width wq splits");
22606 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
22607 for tok in 0..t {
22608 for hh in 0..n_head {
22609 for d in 0..head_dim {
22610 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
22611 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
22612 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
22613 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
22614 }
22615 }
22616 }
22617
22618 let mut small = e.uninit(out_n - 1).unwrap();
22620 assert!(
22621 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
22622 .is_err()
22623 );
22624 }
22625}
22626
22627#[cfg(test)]
22631mod fused_rope_width_tests {
22632 use super::Engine;
22633
22634 #[test]
22637 fn full_width_is_accepted() {
22638 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
22639 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
22640 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
22641 }
22642
22643 #[test]
22656 fn gemma4_official_artifact_widths_pass() {
22657 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
22658 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
22659 }
22660
22661 #[test]
22664 fn partial_rotary_is_refused_with_the_geometry_named() {
22665 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
22667 .expect_err("partial rotary must refuse");
22668 let msg = err.to_string();
22669 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
22670 assert!(msg.contains("n_rot 64"), "{msg}");
22671 assert!(msg.contains("head_dim 256"), "{msg}");
22672 assert!(
22673 msg.contains("64..256"),
22674 "names the band it would corrupt: {msg}"
22675 );
22676 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
22678 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
22680 }
22681}