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 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1463 let coop_on =
1464 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1465 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1466 if coop_on && 16 * nrow <= self.sm_count() as usize {
1467 let f = self.func("filter_stats_coop_f32");
1468 let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1469 let cfg = LaunchConfig {
1470 grid_dim: (16, nrow as u32, 1),
1471 block_dim: (512, 1, 1),
1472 shared_mem_bytes: 0,
1473 };
1474 let __s_b = self.gpu.stream();
1475 let mut b = __s_b.launch_builder(&f);
1476 b.arg(x)
1477 .arg(&rs)
1478 .arg(rows)
1479 .arg(&mut *out_th)
1480 .arg(&mut *out_z)
1481 .arg(&mut *out_max)
1482 .arg(&mut ws)
1483 .arg(&ni)
1484 .arg(&nr)
1485 .arg(&temp)
1486 .arg(&top_k)
1487 .arg(&top_p)
1488 .arg(&min_p);
1489 unsafe {
1490 b.launch_cooperative(cfg)?;
1491 }
1492 return Ok(());
1493 }
1494 let f = self.func("filter_stats_f32");
1495 let cfg = LaunchConfig {
1496 grid_dim: (nrow as u32, 1, 1),
1497 block_dim: (1024, 1, 1),
1498 shared_mem_bytes: 0,
1499 };
1500 let __s_b = self.gpu.stream();
1501 let mut b = __s_b.launch_builder(&f);
1502 b.arg(x)
1503 .arg(&rs)
1504 .arg(rows)
1505 .arg(&mut *out_th)
1506 .arg(&mut *out_z)
1507 .arg(&mut *out_max)
1508 .arg(&ni)
1509 .arg(&nr)
1510 .arg(&temp)
1511 .arg(&top_k)
1512 .arg(&top_p)
1513 .arg(&min_p);
1514 unsafe {
1515 b.launch(cfg)?;
1516 }
1517 Ok(())
1518 }
1519
1520 #[allow(clippy::too_many_arguments)]
1522 pub fn softmax_gather_filtered(
1523 &self,
1524 x: &CudaSlice<f32>,
1525 row_stride: usize,
1526 ids: &CudaSlice<u32>,
1527 rows: &CudaSlice<i32>,
1528 th: &CudaSlice<f32>,
1529 z: &CudaSlice<f32>,
1530 out: &mut CudaSlice<f32>,
1531 n: usize,
1532 npair: usize,
1533 temp: f32,
1534 ) -> Result<(), Box<dyn std::error::Error>> {
1535 let f = self.func("softmax_gather_filtered_f32");
1536 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1537 let cfg = LaunchConfig {
1538 grid_dim: (npair as u32, 1, 1),
1539 block_dim: (256, 1, 1),
1540 shared_mem_bytes: 0,
1541 };
1542 let __s_b = self.gpu.stream();
1543 let mut b = __s_b.launch_builder(&f);
1544 b.arg(x)
1545 .arg(&rs)
1546 .arg(ids)
1547 .arg(rows)
1548 .arg(th)
1549 .arg(z)
1550 .arg(&mut *out)
1551 .arg(&ni)
1552 .arg(&np)
1553 .arg(&temp);
1554 unsafe {
1555 b.launch(cfg)?;
1556 }
1557 Ok(())
1558 }
1559
1560 #[allow(clippy::too_many_arguments)]
1562 pub fn residual_sample_filtered(
1563 &self,
1564 p: &CudaSlice<f32>,
1565 q: Option<&CudaSlice<f32>>,
1566 n: usize,
1567 temp: f32,
1568 seed: u64,
1569 stream_pos: u32,
1570 p_stats: (f32, f32, f32),
1571 q_stats: (f32, f32, f32),
1572 out_tok: &mut CudaSlice<u32>,
1573 ) -> Result<(), Box<dyn std::error::Error>> {
1574 let f = self.func("residual_sample_filtered_f32");
1575 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1576 let has_q: i32 = q.is_some() as i32;
1577 let qbuf = q.unwrap_or(p);
1578 let (pm, pth, pz) = p_stats;
1579 let (qm, qth, qz) = q_stats;
1580 let cfg = LaunchConfig {
1581 grid_dim: (1, 1, 1),
1582 block_dim: (1024, 1, 1),
1583 shared_mem_bytes: 0,
1584 };
1585 let __s_b = self.gpu.stream();
1586 let mut b = __s_b.launch_builder(&f);
1587 b.arg(p)
1588 .arg(qbuf)
1589 .arg(&has_q)
1590 .arg(&ni)
1591 .arg(&temp)
1592 .arg(&slo)
1593 .arg(&shi)
1594 .arg(&stream_pos)
1595 .arg(&pm)
1596 .arg(&pth)
1597 .arg(&pz)
1598 .arg(&qm)
1599 .arg(&qth)
1600 .arg(&qz)
1601 .arg(&mut *out_tok);
1602 unsafe {
1603 b.launch(cfg)?;
1604 }
1605 Ok(())
1606 }
1607
1608 #[allow(clippy::too_many_arguments)]
1614 pub fn residual_sample_sparse_q(
1615 &self,
1616 p: &CudaSlice<f32>,
1617 cand_ids: &CudaSlice<u32>,
1618 q_probs: &CudaSlice<f32>,
1619 n_cand: usize,
1620 n: usize,
1621 temp: f32,
1622 seed: u64,
1623 stream_pos: u32,
1624 p_stats: (f32, f32, f32),
1625 out_tok: &mut CudaSlice<u32>,
1626 ) -> Result<(), Box<dyn std::error::Error>> {
1627 assert!(
1628 n_cand >= 1 && n_cand <= 32,
1629 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1630 );
1631 let f = self.func("residual_sample_sparse_q_f32");
1632 let (ni, nc) = (n as i32, n_cand as i32);
1633 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1634 let (pm, pth, pz) = p_stats;
1635 let cfg = LaunchConfig {
1636 grid_dim: (1, 1, 1),
1637 block_dim: (1024, 1, 1),
1638 shared_mem_bytes: 0,
1639 };
1640 let __s_b = self.gpu.stream();
1641 let mut b = __s_b.launch_builder(&f);
1642 b.arg(p)
1643 .arg(cand_ids)
1644 .arg(q_probs)
1645 .arg(&nc)
1646 .arg(&ni)
1647 .arg(&temp)
1648 .arg(&slo)
1649 .arg(&shi)
1650 .arg(&stream_pos)
1651 .arg(&pm)
1652 .arg(&pth)
1653 .arg(&pz)
1654 .arg(&mut *out_tok);
1655 unsafe {
1656 b.launch(cfg)?;
1657 }
1658 Ok(())
1659 }
1660
1661 #[allow(clippy::too_many_arguments)]
1663 pub fn gumbel_perturb_filtered(
1664 &self,
1665 x: &CudaSlice<f32>,
1666 y: &mut CudaSlice<f32>,
1667 n: usize,
1668 seed: u64,
1669 stream_pos: u32,
1670 temp: f32,
1671 row_max: f32,
1672 th: f32,
1673 ) -> Result<(), Box<dyn std::error::Error>> {
1674 let f = self.func("gumbel_perturb_filtered_f32");
1675 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1676 let cfg = LaunchConfig {
1677 grid_dim: (n.div_ceil(256) as u32, 1, 1),
1678 block_dim: (256, 1, 1),
1679 shared_mem_bytes: 0,
1680 };
1681 let __s_b = self.gpu.stream();
1682 let mut b = __s_b.launch_builder(&f);
1683 b.arg(x)
1684 .arg(&mut *y)
1685 .arg(&ni)
1686 .arg(&slo)
1687 .arg(&shi)
1688 .arg(&stream_pos)
1689 .arg(&temp)
1690 .arg(&row_max)
1691 .arg(&th);
1692 unsafe {
1693 b.launch(cfg)?;
1694 }
1695 Ok(())
1696 }
1697
1698 #[allow(clippy::too_many_arguments)]
1702 pub fn penalize_logits(
1703 &self,
1704 x: &mut CudaSlice<f32>,
1705 hist: &CudaSlice<u32>,
1706 n_hist: usize,
1707 rep: f32,
1708 freq: f32,
1709 present: f32,
1710 n: usize,
1711 ) -> Result<(), Box<dyn std::error::Error>> {
1712 if n_hist == 0 {
1713 return Ok(());
1714 }
1715 let f = self.func("penalize_logits_f32");
1716 let (nh, ni) = (n_hist as i32, n as i32);
1717 let cfg = LaunchConfig {
1718 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1719 block_dim: (128, 1, 1),
1720 shared_mem_bytes: 0,
1721 };
1722 let __s_b = self.gpu.stream();
1723 let mut b = __s_b.launch_builder(&f);
1724 b.arg(&mut *x)
1725 .arg(hist)
1726 .arg(&nh)
1727 .arg(&rep)
1728 .arg(&freq)
1729 .arg(&present)
1730 .arg(&ni);
1731 unsafe {
1732 b.launch(cfg)?;
1733 }
1734 Ok(())
1735 }
1736
1737 #[allow(clippy::too_many_arguments)]
1739 pub fn penalize_logits_rows(
1740 &self,
1741 x: &mut CudaSlice<f32>,
1742 hist: &CudaSlice<u32>,
1743 n_hist: usize,
1744 rep: f32,
1745 freq: f32,
1746 present: f32,
1747 n: usize,
1748 nrow: usize,
1749 ) -> Result<(), Box<dyn std::error::Error>> {
1750 if n_hist == 0 || nrow == 0 {
1751 return Ok(());
1752 }
1753 let f = self.func("penalize_logits_rows_f32");
1754 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1755 let cfg = LaunchConfig {
1756 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1757 block_dim: (128, 1, 1),
1758 shared_mem_bytes: 0,
1759 };
1760 let __s_b = self.gpu.stream();
1761 let mut b = __s_b.launch_builder(&f);
1762 b.arg(&mut *x)
1763 .arg(hist)
1764 .arg(&nh)
1765 .arg(&rep)
1766 .arg(&freq)
1767 .arg(&present)
1768 .arg(&ni)
1769 .arg(&nr);
1770 unsafe {
1771 b.launch(cfg)?;
1772 }
1773 Ok(())
1774 }
1775
1776 #[allow(clippy::too_many_arguments)]
1784 pub fn penalize_logits_rows_inc(
1785 &self,
1786 x: &mut CudaSlice<f32>,
1787 hist: &CudaSlice<u32>,
1788 n_hist0: usize,
1789 rep: f32,
1790 freq: f32,
1791 present: f32,
1792 n: usize,
1793 nrow: usize,
1794 win: usize,
1795 ) -> Result<(), Box<dyn std::error::Error>> {
1796 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
1797 return Ok(());
1798 }
1799 debug_assert!(
1800 hist.len() >= n_hist0 + nrow - 1,
1801 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
1802 );
1803 let f = self.func("penalize_logits_rows_inc_f32");
1804 let max_len = win.min(n_hist0 + nrow - 1).max(1);
1805 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
1806 let cfg = LaunchConfig {
1807 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
1808 block_dim: (128, 1, 1),
1809 shared_mem_bytes: 0,
1810 };
1811 let __s_b = self.gpu.stream();
1812 let mut b = __s_b.launch_builder(&f);
1813 b.arg(&mut *x)
1814 .arg(hist)
1815 .arg(&nh)
1816 .arg(&rep)
1817 .arg(&freq)
1818 .arg(&present)
1819 .arg(&ni)
1820 .arg(&nr)
1821 .arg(&wi);
1822 unsafe {
1823 b.launch(cfg)?;
1824 }
1825 Ok(())
1826 }
1827
1828 pub fn wpf_level() -> u32 {
1836 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1837 *ON.get_or_init(|| {
1838 std::env::var("MEMRA_WPF")
1839 .ok()
1840 .and_then(|v| v.parse().ok())
1841 .unwrap_or(1)
1842 })
1843 }
1844
1845 pub fn set_verify_exact(&self, on: bool) {
1857 self.verify_exact
1858 .store(on, std::sync::atomic::Ordering::Relaxed);
1859 }
1860 pub(crate) fn verify_exact_on(&self) -> bool {
1861 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1862 }
1863
1864 pub fn qkv_append_on() -> bool {
1867 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1868 *ON.get_or_init(|| {
1869 std::env::var("MEMRA_QKV_APPEND")
1870 .map(|v| v != "0")
1871 .unwrap_or(true)
1872 })
1873 }
1874
1875 pub fn pdl_wb_on() -> bool {
1878 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1879 *ON.get_or_init(|| {
1880 std::env::var("MEMRA_PDL_WB")
1881 .map(|v| v != "0")
1882 .unwrap_or(true)
1883 })
1884 }
1885
1886 pub fn norm_ilp_on() -> bool {
1894 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1895 *ON.get_or_init(|| {
1896 std::env::var("MEMRA_NORM_ILP")
1897 .map(|v| v != "0")
1898 .unwrap_or(true)
1899 })
1900 }
1901
1902 pub fn tk_ffn_dual_on() -> bool {
1910 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1911 *ON.get_or_init(|| {
1912 std::env::var("MEMRA_TK_FFN_DUAL")
1913 .map(|v| v != "0")
1914 .unwrap_or(true)
1915 })
1916 }
1917
1918 pub fn pdl_mmvq_on() -> bool {
1922 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1923 *ON.get_or_init(|| {
1924 std::env::var("MEMRA_PDL_MMVQ")
1925 .map(|v| v != "0")
1926 .unwrap_or(true)
1927 })
1928 }
1929
1930 pub fn pdl_on() -> bool {
1931 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1932 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1933 }
1934
1935 pub fn pdl_nvfp4q8_on() -> bool {
1941 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1942 *ON.get_or_init(|| {
1943 std::env::var("MEMRA_PDL_NVFP4")
1944 .map(|v| v != "0")
1945 .unwrap_or(true)
1946 })
1947 }
1948
1949 fn q40_mr1_on() -> bool {
1955 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1956 match *Q40MR.get_or_init(|| {
1957 std::env::var("MEMRA_Q40_MR")
1958 .ok()
1959 .and_then(|v| v.parse().ok())
1960 }) {
1961 Some(v) => v == 1,
1962 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1963 }
1964 }
1965
1966 fn pdl_func_flash(
1971 &self,
1972 g: bool,
1973 name: &'static str,
1974 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1975 use cudarc::driver::sys as cu;
1976 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1983 std::sync::Mutex::new(None);
1984 static FNS: std::sync::Mutex<
1985 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1986 > = std::sync::Mutex::new(None);
1987 let ctx_key = self.ctx().cu_ctx() as usize;
1988 if let Some(&f) = FNS
1989 .lock()
1990 .unwrap()
1991 .get_or_insert_with(Default::default)
1992 .get(&(ctx_key, g, name))
1993 {
1994 return Ok(f as cu::CUfunction);
1995 }
1996 let module = {
1997 let mut mods = MODS.lock().unwrap();
1998 let map = mods.get_or_insert_with(Default::default);
1999 match map.get(&(ctx_key, g)) {
2000 Some(&m) => m,
2001 None => {
2002 let m = self.pdl_load_module_in_ctx(if g {
2003 FLASH_FATBIN_KF8VF8
2004 } else {
2005 FLASH_FATBIN
2006 })?;
2007 map.insert((ctx_key, g), m);
2008 m
2009 }
2010 }
2011 };
2012 let cname = std::ffi::CString::new(name)?;
2013 let mut f: cu::CUfunction = std::ptr::null_mut();
2014 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2015 if r != cu::CUresult::CUDA_SUCCESS {
2016 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2017 }
2018 FNS.lock()
2019 .unwrap()
2020 .get_or_insert_with(Default::default)
2021 .insert((ctx_key, g, name), f as usize);
2022 Ok(f)
2023 }
2024
2025 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2030 use cudarc::driver::sys as cu;
2031 let mut prev: cu::CUcontext = std::ptr::null_mut();
2032 unsafe {
2033 cu::cuCtxGetCurrent(&mut prev).result()?;
2034 }
2035 self.ctx().bind_to_thread()?;
2036 let mut m: cu::CUmodule = std::ptr::null_mut();
2037 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2038 let restore = if prev.is_null() {
2039 cu::CUresult::CUDA_SUCCESS
2040 } else {
2041 unsafe { cu::cuCtxSetCurrent(prev) }
2042 };
2043 if r != cu::CUresult::CUDA_SUCCESS {
2044 return Err(format!("pdl module load: {r:?}").into());
2045 }
2046 if restore != cu::CUresult::CUDA_SUCCESS {
2047 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2048 }
2049 Ok(m as usize)
2050 }
2051
2052 fn pdl_func(
2053 &self,
2054 name: &'static str,
2055 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2056 use cudarc::driver::sys as cu;
2057 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2060 std::sync::Mutex::new(None);
2061 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2064 std::sync::Mutex::new(None);
2065 static FNS: std::sync::Mutex<
2066 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2067 > = std::sync::Mutex::new(None);
2068 let ctx_key = self.ctx().cu_ctx() as usize;
2069 if let Some(&f) = FNS
2070 .lock()
2071 .unwrap()
2072 .get_or_insert_with(Default::default)
2073 .get(&(ctx_key, name))
2074 {
2075 return Ok(f as cu::CUfunction);
2076 }
2077 let module = {
2078 let mut mods = MODULES.lock().unwrap();
2079 let map = mods.get_or_insert_with(Default::default);
2080 match map.get(&ctx_key) {
2081 Some(&m) => m,
2082 None => {
2083 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2084 map.insert(ctx_key, m);
2085 m
2086 }
2087 }
2088 };
2089 let cname = std::ffi::CString::new(name)?;
2090 let mut f: cu::CUfunction = std::ptr::null_mut();
2091 let mut r =
2092 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2093 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2094 let qmodule = {
2095 let mut mods = QMODULES.lock().unwrap();
2096 let map = mods.get_or_insert_with(Default::default);
2097 match map.get(&ctx_key) {
2098 Some(&m) => m,
2099 None => {
2100 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2101 map.insert(ctx_key, m);
2102 m
2103 }
2104 }
2105 };
2106 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2107 }
2108 if r != cu::CUresult::CUDA_SUCCESS {
2109 return Err(format!("pdl_func {name}: {r:?}").into());
2110 }
2111 FNS.lock()
2112 .unwrap()
2113 .get_or_insert_with(Default::default)
2114 .insert((ctx_key, name), f as usize);
2115 Ok(f)
2116 }
2117
2118 unsafe fn launch_pdl_flash(
2130 &self,
2131 g: bool,
2132 name: &'static str,
2133 grid: (u32, u32, u32),
2134 block: (u32, u32, u32),
2135 smem: u32,
2136 params: &mut [*mut std::ffi::c_void],
2137 ) -> Result<(), Box<dyn std::error::Error>> {
2138 use cudarc::driver::sys as cu;
2139 let f = self.pdl_func_flash(g, name)?;
2140 if smem > 0 {
2141 let r =
2143 unsafe {
2144 cu::cuFuncSetAttribute(f,
2145 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2146 smem as i32)
2147 };
2148 if r != cu::CUresult::CUDA_SUCCESS {
2149 return Err(format!("pdl smem attr {name}: {r:?}").into());
2150 }
2151 }
2152 let mut attr = cu::CUlaunchAttribute {
2153 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2154 pad: [0; 4],
2155 value: cu::CUlaunchAttributeValue {
2156 programmaticStreamSerializationAllowed: 1,
2157 },
2158 };
2159 let cfg = cu::CUlaunchConfig {
2160 gridDimX: grid.0,
2161 gridDimY: grid.1,
2162 gridDimZ: grid.2,
2163 blockDimX: block.0,
2164 blockDimY: block.1,
2165 blockDimZ: block.2,
2166 sharedMemBytes: smem,
2167 hStream: self.gpu.stream().cu_stream(),
2168 attrs: &mut attr,
2169 numAttrs: 1,
2170 };
2171 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2172 if r != cu::CUresult::CUDA_SUCCESS {
2173 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2174 }
2175 Ok(())
2176 }
2177
2178 unsafe fn launch_pdl(
2179 &self,
2180 name: &'static str,
2181 grid: (u32, u32, u32),
2182 block: (u32, u32, u32),
2183 params: &mut [*mut std::ffi::c_void],
2184 ) -> Result<(), Box<dyn std::error::Error>> {
2185 use cudarc::driver::sys as cu;
2186 let f = self.pdl_func(name)?;
2187 let mut attr = cu::CUlaunchAttribute {
2188 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2189 pad: [0; 4],
2190 value: cu::CUlaunchAttributeValue {
2191 programmaticStreamSerializationAllowed: 1,
2192 },
2193 };
2194 let cfg = cu::CUlaunchConfig {
2195 gridDimX: grid.0,
2196 gridDimY: grid.1,
2197 gridDimZ: grid.2,
2198 blockDimX: block.0,
2199 blockDimY: block.1,
2200 blockDimZ: block.2,
2201 sharedMemBytes: 0,
2202 hStream: self.gpu.stream().cu_stream(),
2203 attrs: &mut attr,
2204 numAttrs: 1,
2205 };
2206 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2207 if r != cu::CUresult::CUDA_SUCCESS {
2208 return Err(format!("launch_pdl {name}: {r:?}").into());
2209 }
2210 Ok(())
2211 }
2212
2213 pub fn prefetch_weight_l2(
2216 &self,
2217 w: &crate::model::GpuTensor,
2218 ) -> Result<(), Box<dyn std::error::Error>> {
2219 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2220 let p = rp4.as_ref().unwrap_or(bytes);
2221 self.prefetch_l2(p, p.len())?;
2222 }
2223 Ok(())
2224 }
2225
2226 pub fn gather_row_bf16(
2229 &self,
2230 table: &CudaSlice<u8>,
2231 tok: &CudaSlice<u32>,
2232 idx: usize,
2233 dst: &mut CudaSlice<f32>,
2234 ncols: usize,
2235 ) -> Result<(), Box<dyn std::error::Error>> {
2236 let f = self.func("gather_row_bf16_f32");
2237 let cfg = LaunchConfig {
2238 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2239 block_dim: (256, 1, 1),
2240 shared_mem_bytes: 0,
2241 };
2242 let (nc, ix) = (ncols as i32, idx as i32);
2243 let __s_b = self.gpu.stream();
2244 let mut b = __s_b.launch_builder(&f);
2245 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2246 unsafe {
2247 b.launch(cfg)?;
2248 }
2249 Ok(())
2250 }
2251
2252 #[allow(clippy::too_many_arguments)]
2258 pub fn dflash2_dynconv(
2259 &self,
2260 x: &CudaSlice<f32>,
2261 dyn_: &CudaSlice<f32>,
2262 base: &CudaSlice<f32>,
2263 out: &mut CudaSlice<f32>,
2264 rows: usize,
2265 hidden: usize,
2266 group_size: usize,
2267 ksize: usize,
2268 half: usize,
2269 ) -> Result<(), Box<dyn std::error::Error>> {
2270 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2271 let f = self.func("dflash2_dynconv_f32");
2272 let n = rows * hidden;
2273 let cfg = LaunchConfig {
2274 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2275 block_dim: (256, 1, 1),
2276 shared_mem_bytes: 0,
2277 };
2278 let (ri, hi, gi, ki, hf) = (
2279 rows as i32,
2280 hidden as i32,
2281 group_size as i32,
2282 ksize as i32,
2283 half as i32,
2284 );
2285 let __s_b = self.gpu.stream();
2286 let mut b = __s_b.launch_builder(&f);
2287 b.arg(x)
2288 .arg(dyn_)
2289 .arg(base)
2290 .arg(out)
2291 .arg(&ri)
2292 .arg(&hi)
2293 .arg(&gi)
2294 .arg(&ki)
2295 .arg(&hf);
2296 unsafe {
2297 b.launch(cfg)?;
2298 }
2299 Ok(())
2300 }
2301
2302 pub fn topk_rows(
2306 &self,
2307 logits: &CudaSlice<f32>,
2308 n_rows: usize,
2309 n_cols: usize,
2310 k: usize,
2311 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2312 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2313 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2314 let f = self.func("topk_rows_f32");
2315 let nth = 256usize;
2316 let mut vals = self.uninit(n_rows * k)?;
2317 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2318 let cfg = LaunchConfig {
2319 grid_dim: (n_rows as u32, 1, 1),
2320 block_dim: (nth as u32, 1, 1),
2321 shared_mem_bytes: (nth * k * 8) as u32,
2322 };
2323 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2324 let __s_b = self.gpu.stream();
2325 let mut b = __s_b.launch_builder(&f);
2326 b.arg(logits)
2327 .arg(&nr)
2328 .arg(&nc)
2329 .arg(&ki)
2330 .arg(&mut vals)
2331 .arg(&mut idxs);
2332 unsafe {
2333 b.launch(cfg)?;
2334 }
2335 Ok((vals, idxs))
2336 }
2337
2338 pub fn add_row_inplace(
2340 &self,
2341 logits: &mut CudaSlice<f32>,
2342 bias: &CudaSlice<f32>,
2343 n: usize,
2344 row_off: usize,
2345 ) -> Result<(), Box<dyn std::error::Error>> {
2346 let f = self.func("add_row_inplace_f32");
2347 let cfg = LaunchConfig {
2348 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2349 block_dim: (256, 1, 1),
2350 shared_mem_bytes: 0,
2351 };
2352 let (ni, off) = (n as i32, row_off as i64);
2353 let __s_b = self.gpu.stream();
2354 let mut b = __s_b.launch_builder(&f);
2355 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2356 unsafe {
2357 b.launch(cfg)?;
2358 }
2359 Ok(())
2360 }
2361
2362 pub fn prefetch_l2(
2364 &self,
2365 p: &CudaSlice<u8>,
2366 n: usize,
2367 ) -> Result<(), Box<dyn std::error::Error>> {
2368 let f = self.func("prefetch_l2_bytes");
2369 let lines = n.div_ceil(128);
2370 let ni = n as i64;
2371 let cfg = LaunchConfig {
2372 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2373 block_dim: (256, 1, 1),
2374 shared_mem_bytes: 0,
2375 };
2376 let __s_b = self.gpu.stream();
2377 let mut b = __s_b.launch_builder(&f);
2378 b.arg(p).arg(&ni);
2379 unsafe {
2380 b.launch(cfg)?;
2381 }
2382 Ok(())
2383 }
2384
2385 pub fn router_gemv(
2388 &self,
2389 w: &CudaSlice<f32>,
2390 x: &CudaSlice<f32>,
2391 n_embd: usize,
2392 n_experts: usize,
2393 t: usize,
2394 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2395 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2401 Ok("0") => false,
2402 Ok(_) => true,
2403 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2404 };
2405 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2415 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2416 }
2417
2418 pub fn router_gemv_form(
2421 &self,
2422 w: &CudaSlice<f32>,
2423 x: &CudaSlice<f32>,
2424 n_embd: usize,
2425 n_experts: usize,
2426 t: usize,
2427 w8: bool,
2428 batch: bool,
2429 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2430 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2431 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2432 let f = if batch {
2433 self.func("router_gemv_f32_w8_batch")
2434 } else if w8 {
2435 self.func("router_gemv_f32_w8")
2436 } else {
2437 self.func("router_gemv_f32")
2438 };
2439 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2440 let cfg = if batch {
2441 LaunchConfig {
2442 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2443 block_dim: (32, 8, 1),
2444 shared_mem_bytes: 0,
2445 }
2446 } else {
2447 LaunchConfig {
2448 grid_dim: (n_experts as u32, t as u32, 1),
2449 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2450 shared_mem_bytes: 0,
2451 }
2452 };
2453 let __s_b = self.gpu.stream();
2454 let mut b = __s_b.launch_builder(&f);
2455 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2456 unsafe {
2457 b.launch(cfg)?;
2458 }
2459 Ok(y)
2460 }
2461
2462 pub fn rows_permute(
2464 &self,
2465 src: &CudaSlice<f32>,
2466 idx: &CudaSlice<i32>,
2467 nrows: usize,
2468 ncols: usize,
2469 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2470 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2471 let f = self.func("rows_permute_f32");
2472 let (nc, nr) = (ncols as i32, nrows as i32);
2473 let cfg = LaunchConfig {
2474 grid_dim: (nrows as u32, 1, 1),
2475 block_dim: (256, 1, 1),
2476 shared_mem_bytes: 0,
2477 };
2478 let __s_b = self.gpu.stream();
2479 let mut b = __s_b.launch_builder(&f);
2480 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2481 unsafe {
2482 b.launch(cfg)?;
2483 }
2484 Ok(dst)
2485 }
2486
2487 pub fn sigmoid_dot_rows(
2492 &self,
2493 x: &CudaSlice<f32>,
2494 w: &CudaSlice<f32>,
2495 n_embd: usize,
2496 t: usize,
2497 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2498 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2501 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2502 let gs = self.linear(x, w, t, n_embd, 1)?;
2503 let mut g = self.uninit(t)?;
2504 self.sigmoid(&gs, &mut g, t)?;
2505 return Ok(g);
2506 }
2507 let mut g = self.alloc_uninit::<f32>(t)?;
2513 let f = self.func("sigmoid_dot_rows_f32");
2514 let (ne, ti) = (n_embd as i32, t as i32);
2515 let cfg = LaunchConfig {
2516 grid_dim: (t as u32, 1, 1),
2517 block_dim: (32, 8, 1),
2518 shared_mem_bytes: 0,
2519 };
2520 let __s_b = self.gpu.stream();
2521 let mut b = __s_b.launch_builder(&f);
2522 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2523 unsafe {
2524 b.launch(cfg)?;
2525 }
2526 Ok(g)
2527 }
2528
2529 pub fn spec_rollback_stream(
2531 &self,
2532 len_ptrs: &CudaSlice<u64>,
2533 pos_start: &CudaSlice<i32>,
2534 acc: &CudaSlice<u32>,
2535 base: usize,
2536 n_rows: usize,
2537 ) -> Result<(), Box<dyn std::error::Error>> {
2538 let f = self.func("spec_rollback_stream");
2539 let (b, nr) = (base as i32, n_rows as i32);
2540 let cfg = LaunchConfig {
2541 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2542 block_dim: (64, 1, 1),
2543 shared_mem_bytes: 0,
2544 };
2545 let __s_bl = self.gpu.stream();
2546 let mut bl = __s_bl.launch_builder(&f);
2547 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2548 unsafe {
2549 bl.launch(cfg)?;
2550 }
2551 Ok(())
2552 }
2553
2554 pub fn plain_tok_ring(
2556 &self,
2557 vam: &CudaSlice<u32>,
2558 pos_start: &CudaSlice<i32>,
2559 base: usize,
2560 ring: &mut CudaSlice<u32>,
2561 ) -> Result<(), Box<dyn std::error::Error>> {
2562 let f = self.func("plain_tok_ring");
2563 let (b, cap) = (base as i32, ring.len() as i32);
2564 let cfg = LaunchConfig {
2565 grid_dim: (1, 1, 1),
2566 block_dim: (32, 1, 1),
2567 shared_mem_bytes: 0,
2568 };
2569 let __s_bl = self.gpu.stream();
2570 let mut bl = __s_bl.launch_builder(&f);
2571 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2572 unsafe {
2573 bl.launch(cfg)?;
2574 }
2575 Ok(())
2576 }
2577
2578 pub fn spec_ring_commit(
2580 &self,
2581 vtok: &CudaSlice<u32>,
2582 acc: &CudaSlice<u32>,
2583 brk: &CudaSlice<u32>,
2584 ring: &mut CudaSlice<u32>,
2585 pend: &mut CudaSlice<u32>,
2586 ) -> Result<(), Box<dyn std::error::Error>> {
2587 let f = self.func("spec_ring_commit");
2588 let cfg = LaunchConfig {
2589 grid_dim: (1, 1, 1),
2590 block_dim: (32, 1, 1),
2591 shared_mem_bytes: 0,
2592 };
2593 let __s_b = self.gpu.stream();
2594 let mut b = __s_b.launch_builder(&f);
2595 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2596 unsafe {
2597 b.launch(cfg)?;
2598 }
2599 Ok(())
2600 }
2601 pub fn i32_copy_add(
2602 &self,
2603 src: &CudaSlice<i32>,
2604 dst: &mut CudaSlice<i32>,
2605 delta: i32,
2606 ) -> Result<(), Box<dyn std::error::Error>> {
2607 let f = self.func("i32_copy_add");
2608 let cfg = LaunchConfig {
2609 grid_dim: (1, 1, 1),
2610 block_dim: (32, 1, 1),
2611 shared_mem_bytes: 0,
2612 };
2613 let __s_b = self.gpu.stream();
2614 let mut b = __s_b.launch_builder(&f);
2615 b.arg(src).arg(dst).arg(&delta);
2616 unsafe {
2617 b.launch(cfg)?;
2618 }
2619 Ok(())
2620 }
2621 pub fn u32_copy(
2622 &self,
2623 src: &CudaSlice<u32>,
2624 dst: &mut CudaSlice<u32>,
2625 ) -> Result<(), Box<dyn std::error::Error>> {
2626 let f = self.func("u32_copy");
2627 let cfg = LaunchConfig {
2628 grid_dim: (1, 1, 1),
2629 block_dim: (32, 1, 1),
2630 shared_mem_bytes: 0,
2631 };
2632 let __s_b = self.gpu.stream();
2633 let mut b = __s_b.launch_builder(&f);
2634 b.arg(src).arg(dst);
2635 unsafe {
2636 b.launch(cfg)?;
2637 }
2638 Ok(())
2639 }
2640
2641 pub fn spec_adapt_k(
2645 &self,
2646 acc: &CudaSlice<u32>,
2647 brk: &mut CudaSlice<u32>,
2648 floor: usize,
2649 cap: usize,
2650 ) -> Result<(), Box<dyn std::error::Error>> {
2651 let f = self.func("spec_adapt_k");
2652 let (fl, cp) = (floor as i32, cap as i32);
2653 let cfg = LaunchConfig {
2654 grid_dim: (1, 1, 1),
2655 block_dim: (32, 1, 1),
2656 shared_mem_bytes: 0,
2657 };
2658 let __s_b = self.gpu.stream();
2659 let mut b = __s_b.launch_builder(&f);
2660 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2661 unsafe {
2662 b.launch(cfg)?;
2663 }
2664 Ok(())
2665 }
2666
2667 pub fn spec_accept_greedy_dc(
2669 &self,
2670 preds: &CudaSlice<u32>,
2671 vtok: &CudaSlice<u32>,
2672 last_pred: &CudaSlice<u32>,
2673 brk: &CudaSlice<u32>,
2674 out: &mut CudaSlice<u32>,
2675 ) -> Result<(), Box<dyn std::error::Error>> {
2676 let f = self.func("spec_accept_greedy_dc");
2677 let cfg = LaunchConfig {
2678 grid_dim: (1, 1, 1),
2679 block_dim: (32, 1, 1),
2680 shared_mem_bytes: 0,
2681 };
2682 let __s_b = self.gpu.stream();
2683 let mut b = __s_b.launch_builder(&f);
2684 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2685 unsafe {
2686 b.launch(cfg)?;
2687 }
2688 Ok(())
2689 }
2690
2691 pub fn pos_iota(
2693 &self,
2694 pos0: &CudaSlice<i32>,
2695 out: &mut CudaSlice<i32>,
2696 t: usize,
2697 ) -> Result<(), Box<dyn std::error::Error>> {
2698 let f = self.func("pos_iota_i32");
2699 let ti = t as i32;
2700 let cfg = LaunchConfig {
2701 grid_dim: (1, 1, 1),
2702 block_dim: (t.max(1) as u32, 1, 1),
2703 shared_mem_bytes: 0,
2704 };
2705 let __s_b = self.gpu.stream();
2706 let mut b = __s_b.launch_builder(&f);
2707 b.arg(pos0).arg(out).arg(&ti);
2708 unsafe {
2709 b.launch(cfg)?;
2710 }
2711 Ok(())
2712 }
2713 #[allow(clippy::too_many_arguments)]
2714 pub fn append_kv_quantized_rows_dc(
2715 &self,
2716 k_rows: &CudaSlice<f32>,
2717 v_rows: &CudaSlice<f32>,
2718 kc: &mut CudaSlice<u8>,
2719 vc: &mut CudaSlice<u8>,
2720 t0_dev: &CudaSlice<i32>,
2721 t: usize,
2722 kv_dim_k: usize,
2723 kv_dim_v: usize,
2724 k_tok_bytes: usize,
2725 v_tok_bytes: usize,
2726 g: bool,
2727 ) -> Result<(), Box<dyn std::error::Error>> {
2728 let f = if g {
2729 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2730 } else {
2731 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2732 };
2733 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2734 let cfg = LaunchConfig {
2735 grid_dim: (nblk, t as u32, 1),
2736 block_dim: (32, 1, 1),
2737 shared_mem_bytes: 0,
2738 };
2739 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2740 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2741 let __s_b = self.gpu.stream();
2742 let mut b = __s_b.launch_builder(&f);
2743 b.arg(k_rows)
2744 .arg(v_rows)
2745 .arg(kc)
2746 .arg(vc)
2747 .arg(t0_dev)
2748 .arg(&kdk)
2749 .arg(&kdv)
2750 .arg(&ktb)
2751 .arg(&vtb);
2752 unsafe {
2753 b.launch(cfg)?;
2754 }
2755 Ok(())
2756 }
2757
2758 #[allow(clippy::too_many_arguments)]
2761 pub fn append_kv_quantized_row_dc_inc(
2762 &self,
2763 k_row: &CudaSlice<f32>,
2764 v_row: &CudaSlice<f32>,
2765 kc: &mut CudaSlice<u8>,
2766 vc: &mut CudaSlice<u8>,
2767 t0_dev: &mut CudaSlice<i32>,
2768 kv_dim_k: usize,
2769 kv_dim_v: usize,
2770 k_tok_bytes: usize,
2771 v_tok_bytes: usize,
2772 g: bool,
2773 ) -> Result<(), Box<dyn std::error::Error>> {
2774 let f = if g {
2775 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2776 } else {
2777 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2778 };
2779 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2780 let cfg = LaunchConfig {
2781 grid_dim: (1, 1, 1),
2782 block_dim: (nthreads, 1, 1),
2783 shared_mem_bytes: 0,
2784 };
2785 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2786 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2787 let __s_b = self.gpu.stream();
2788 let mut b = __s_b.launch_builder(&f);
2789 b.arg(k_row)
2790 .arg(v_row)
2791 .arg(kc)
2792 .arg(vc)
2793 .arg(t0_dev)
2794 .arg(&kdk)
2795 .arg(&kdv)
2796 .arg(&ktb)
2797 .arg(&vtb);
2798 unsafe {
2799 b.launch(cfg)?;
2800 }
2801 Ok(())
2802 }
2803
2804 pub fn pack_tok_p(
2806 &self,
2807 tok: &CudaSlice<u32>,
2808 p: &CudaSlice<f32>,
2809 out: &mut CudaSlice<u32>,
2810 slot: usize,
2811 ) -> Result<(), Box<dyn std::error::Error>> {
2812 let f = self.func("pack_tok_p");
2813 let sl = slot as i32;
2814 let cfg = LaunchConfig {
2815 grid_dim: (1, 1, 1),
2816 block_dim: (32, 1, 1),
2817 shared_mem_bytes: 0,
2818 };
2819 let __s_b = self.gpu.stream();
2820 let mut b = __s_b.launch_builder(&f);
2821 b.arg(tok).arg(p).arg(out).arg(&sl);
2822 unsafe {
2823 b.launch(cfg)?;
2824 }
2825 Ok(())
2826 }
2827 pub fn tok_map_u32(
2828 &self,
2829 tok: &mut CudaSlice<u32>,
2830 map: &CudaSlice<u32>,
2831 ) -> Result<(), Box<dyn std::error::Error>> {
2832 let f = self.func("tok_map_u32");
2833 let cfg = LaunchConfig {
2834 grid_dim: (1, 1, 1),
2835 block_dim: (32, 1, 1),
2836 shared_mem_bytes: 0,
2837 };
2838 let __s_b = self.gpu.stream();
2839 let mut b = __s_b.launch_builder(&f);
2840 b.arg(tok).arg(map);
2841 unsafe {
2842 b.launch(cfg)?;
2843 }
2844 Ok(())
2845 }
2846
2847 #[allow(clippy::too_many_arguments)]
2849 pub fn spec_assemble_verify(
2850 &self,
2851 tokp: &CudaSlice<u32>,
2852 pend: &CudaSlice<u32>,
2853 d2t: Option<&CudaSlice<u32>>,
2854 vtok: &mut CudaSlice<u32>,
2855 brk: &mut CudaSlice<u32>,
2856 p_min: f32,
2857 k: usize,
2858 pmin0: bool,
2859 ) -> Result<(), Box<dyn std::error::Error>> {
2860 let f = self.func("spec_assemble_verify");
2861 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2862 let cfg = LaunchConfig {
2863 grid_dim: (1, 1, 1),
2864 block_dim: (32, 1, 1),
2865 shared_mem_bytes: 0,
2866 };
2867 let __s_b = self.gpu.stream();
2868 let mut b = __s_b.launch_builder(&f);
2869 match d2t {
2870 Some(m) => {
2871 b.arg(tokp)
2872 .arg(pend)
2873 .arg(m)
2874 .arg(vtok)
2875 .arg(brk)
2876 .arg(&p_min)
2877 .arg(&ki)
2878 .arg(&pm);
2879 unsafe {
2880 b.launch(cfg)?;
2881 }
2882 }
2883 None => {
2884 let null: u64 = 0;
2885 b.arg(tokp)
2886 .arg(pend)
2887 .arg(&null)
2888 .arg(vtok)
2889 .arg(brk)
2890 .arg(&p_min)
2891 .arg(&ki)
2892 .arg(&pm);
2893 unsafe {
2894 b.launch(cfg)?;
2895 }
2896 }
2897 }
2898 Ok(())
2899 }
2900
2901 #[allow(clippy::too_many_arguments)]
2903 pub fn ssm_conv_ring_rebuild_dc(
2904 &self,
2905 qkv_tm: &CudaSlice<f32>,
2906 ring_old: &CudaSlice<f32>,
2907 conv_state: &mut CudaSlice<f32>,
2908 conv_dim: usize,
2909 acc: &CudaSlice<u32>,
2910 base: usize,
2911 t_v: usize,
2912 d_conv: usize,
2913 ) -> Result<(), Box<dyn std::error::Error>> {
2914 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2915 let n = conv_dim * (d_conv - 1);
2916 let cfg = LaunchConfig::for_num_elems(n as u32);
2917 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2918 let __s_b = self.gpu.stream();
2919 let mut b = __s_b.launch_builder(&f);
2920 b.arg(qkv_tm)
2921 .arg(ring_old)
2922 .arg(conv_state)
2923 .arg(&cd)
2924 .arg(acc)
2925 .arg(&b0)
2926 .arg(&tv)
2927 .arg(&dc);
2928 unsafe {
2929 b.launch(cfg)?;
2930 }
2931 Ok(())
2932 }
2933 #[allow(clippy::too_many_arguments)]
2934 pub fn gdn_scan_s128_dc(
2935 &self,
2936 q: &CudaSlice<f32>,
2937 k: &CudaSlice<f32>,
2938 v: &CudaSlice<f32>,
2939 g: &CudaSlice<f32>,
2940 beta: &CudaSlice<f32>,
2941 state_in: &CudaSlice<f32>,
2942 state_out: &mut CudaSlice<f32>,
2943 o: &mut CudaSlice<f32>,
2944 n_head: usize,
2945 acc: &CudaSlice<u32>,
2946 base: usize,
2947 t_v: usize,
2948 scale: f32,
2949 ) -> Result<(), Box<dyn std::error::Error>> {
2950 let f = self.func("gdn_scan_s128_dc");
2951 const S_V: u32 = 128;
2952 const WARP: u32 = 32;
2953 const COLS_PER_BLOCK: u32 = 4;
2954 let cfg = LaunchConfig {
2955 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2956 block_dim: (WARP, COLS_PER_BLOCK, 1),
2957 shared_mem_bytes: 0,
2958 };
2959 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2960 let __s_b = self.gpu.stream();
2961 let mut b = __s_b.launch_builder(&f);
2962 b.arg(q)
2963 .arg(k)
2964 .arg(v)
2965 .arg(g)
2966 .arg(beta)
2967 .arg(state_in)
2968 .arg(state_out)
2969 .arg(o)
2970 .arg(&h)
2971 .arg(acc)
2972 .arg(&b0)
2973 .arg(&tv)
2974 .arg(&scale);
2975 unsafe {
2976 b.launch(cfg)?;
2977 }
2978 Ok(())
2979 }
2980
2981 pub fn spec_rollback_kv(
2983 &self,
2984 len_ptrs: &CudaSlice<u64>,
2985 saved: &CudaSlice<i32>,
2986 acc: &CudaSlice<u32>,
2987 base: usize,
2988 n_layer: usize,
2989 ) -> Result<(), Box<dyn std::error::Error>> {
2990 let f = self.func("spec_rollback_kv");
2991 let (b, nl) = (base as i32, n_layer as i32);
2992 let cfg = LaunchConfig {
2993 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2994 block_dim: (64, 1, 1),
2995 shared_mem_bytes: 0,
2996 };
2997 let __s_bl = self.gpu.stream();
2998 let mut bl = __s_bl.launch_builder(&f);
2999 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3000 unsafe {
3001 bl.launch(cfg)?;
3002 }
3003 Ok(())
3004 }
3005
3006 pub fn spec_fork_valid(
3008 &self,
3009 acc: &CudaSlice<u32>,
3010 optimistic_pending: u32,
3011 valid: &mut CudaSlice<u32>,
3012 ) -> Result<(), Box<dyn std::error::Error>> {
3013 let f = self.func("spec_fork_valid");
3014 let cfg = LaunchConfig {
3015 grid_dim: (1, 1, 1),
3016 block_dim: (1, 1, 1),
3017 shared_mem_bytes: 0,
3018 };
3019 let __s_bl = self.gpu.stream();
3020 let mut bl = __s_bl.launch_builder(&f);
3021 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3022 unsafe {
3023 bl.launch(cfg)?;
3024 }
3025 Ok(())
3026 }
3027
3028 pub fn spec_fork_reconcile_kv(
3030 &self,
3031 len_ptrs: &CudaSlice<u64>,
3032 saved: &CudaSlice<i32>,
3033 acc: &CudaSlice<u32>,
3034 valid: &CudaSlice<u32>,
3035 base: usize,
3036 n_layer: usize,
3037 ) -> Result<(), Box<dyn std::error::Error>> {
3038 let f = self.func("spec_fork_reconcile_kv");
3039 let (b, nl) = (base as i32, n_layer as i32);
3040 let cfg = LaunchConfig {
3041 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3042 block_dim: (64, 1, 1),
3043 shared_mem_bytes: 0,
3044 };
3045 let __s_bl = self.gpu.stream();
3046 let mut bl = __s_bl.launch_builder(&f);
3047 bl.arg(len_ptrs)
3048 .arg(saved)
3049 .arg(acc)
3050 .arg(valid)
3051 .arg(&b)
3052 .arg(&nl);
3053 unsafe {
3054 bl.launch(cfg)?;
3055 }
3056 Ok(())
3057 }
3058
3059 pub fn spec_fork_restore_f32(
3061 &self,
3062 snapshot: &CudaSlice<f32>,
3063 state: &mut CudaSlice<f32>,
3064 valid: &CudaSlice<u32>,
3065 ) -> Result<(), Box<dyn std::error::Error>> {
3066 assert_eq!(
3067 snapshot.len(),
3068 state.len(),
3069 "fork recurrent snapshot shape mismatch"
3070 );
3071 let f = self.func("spec_fork_restore_f32");
3072 let n = state.len() as i32;
3073 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3074 let cfg = LaunchConfig {
3075 grid_dim: (blocks, 1, 1),
3076 block_dim: (256, 1, 1),
3077 shared_mem_bytes: 0,
3078 };
3079 let __s_bl = self.gpu.stream();
3080 let mut bl = __s_bl.launch_builder(&f);
3081 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3082 unsafe {
3083 bl.launch(cfg)?;
3084 }
3085 Ok(())
3086 }
3087
3088 pub fn spec_seed_gather(
3091 &self,
3092 vx: &CudaSlice<f32>,
3093 fill_prev: &CudaSlice<f32>,
3094 acc: &CudaSlice<u32>,
3095 h_seed: &mut CudaSlice<f32>,
3096 base: usize,
3097 n_embd: usize,
3098 ) -> Result<(), Box<dyn std::error::Error>> {
3099 let f = self.func("spec_seed_gather");
3100 let (b, ne) = (base as i32, n_embd as i32);
3101 let cfg = LaunchConfig {
3102 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3103 block_dim: (256, 1, 1),
3104 shared_mem_bytes: 0,
3105 };
3106 let __s_bl = self.gpu.stream();
3107 let mut bl = __s_bl.launch_builder(&f);
3108 bl.arg(vx)
3109 .arg(fill_prev)
3110 .arg(acc)
3111 .arg(h_seed)
3112 .arg(&b)
3113 .arg(&ne);
3114 unsafe {
3115 bl.launch(cfg)?;
3116 }
3117 Ok(())
3118 }
3119
3120 pub fn spec_accept_greedy(
3122 &self,
3123 preds: &CudaSlice<u32>,
3124 draft: &CudaSlice<u32>,
3125 last_pred: u32,
3126 base: usize,
3127 k_round: usize,
3128 out: &mut CudaSlice<u32>,
3129 ) -> Result<(), Box<dyn std::error::Error>> {
3130 let f = self.func("spec_accept_greedy");
3131 let (b, k) = (base as i32, k_round as i32);
3132 let cfg = LaunchConfig {
3133 grid_dim: (1, 1, 1),
3134 block_dim: (32, 1, 1),
3135 shared_mem_bytes: 0,
3136 };
3137 let __s_bl = self.gpu.stream();
3138 let mut bl = __s_bl.launch_builder(&f);
3139 bl.arg(preds)
3140 .arg(draft)
3141 .arg(&last_pred)
3142 .arg(&b)
3143 .arg(&k)
3144 .arg(out);
3145 unsafe {
3146 bl.launch(cfg)?;
3147 }
3148 Ok(())
3149 }
3150
3151 pub fn gumbel_perturb(
3158 &self,
3159 x: &CudaSlice<f32>,
3160 y: &mut CudaSlice<f32>,
3161 n: usize,
3162 seed: u64,
3163 stream_pos: u32,
3164 temp: f32,
3165 ) -> Result<(), Box<dyn std::error::Error>> {
3166 let f = self.func("gumbel_perturb_f32");
3167 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3168 let cfg = LaunchConfig {
3169 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3170 block_dim: (256, 1, 1),
3171 shared_mem_bytes: 0,
3172 };
3173 let __s_b = self.gpu.stream();
3174 let mut b = __s_b.launch_builder(&f);
3175 b.arg(x)
3176 .arg(&mut *y)
3177 .arg(&ni)
3178 .arg(&slo)
3179 .arg(&shi)
3180 .arg(&stream_pos)
3181 .arg(&temp);
3182 unsafe {
3183 b.launch(cfg)?;
3184 }
3185 Ok(())
3186 }
3187
3188 pub fn mask_logits_col(
3196 &self,
3197 logits: &mut CudaSlice<f32>,
3198 mask: &CudaSlice<u32>,
3199 col: usize,
3200 n: usize,
3201 mask_words: usize,
3202 ) -> Result<(), Box<dyn std::error::Error>> {
3203 let f = self.func("mask_logits_f32");
3204 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3205 let cfg = LaunchConfig {
3206 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3207 block_dim: (256, 1, 1),
3208 shared_mem_bytes: 0,
3209 };
3210 let __s_b = self.gpu.stream();
3211 let mut b = __s_b.launch_builder(&f);
3212 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3213 unsafe {
3214 b.launch(cfg)?;
3215 }
3216 Ok(())
3217 }
3218
3219 pub fn gumbel_perturb_col(
3226 &self,
3227 x: &CudaSlice<f32>,
3228 col: usize,
3229 y: &mut CudaSlice<f32>,
3230 n: usize,
3231 seed: u64,
3232 stream_pos: u32,
3233 temp: f32,
3234 ) -> Result<(), Box<dyn std::error::Error>> {
3235 let f = self.func("gumbel_perturb_f32");
3236 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3237 let col_view = x.slice(col * n..(col + 1) * n);
3238 let cfg = LaunchConfig {
3239 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3240 block_dim: (256, 1, 1),
3241 shared_mem_bytes: 0,
3242 };
3243 let __s_b = self.gpu.stream();
3244 let mut b = __s_b.launch_builder(&f);
3245 b.arg(&col_view)
3246 .arg(&mut *y)
3247 .arg(&ni)
3248 .arg(&slo)
3249 .arg(&shi)
3250 .arg(&stream_pos)
3251 .arg(&temp);
3252 unsafe {
3253 b.launch(cfg)?;
3254 }
3255 Ok(())
3256 }
3257
3258 #[allow(clippy::too_many_arguments)]
3264 pub fn gumbel_perturb_filtered_col(
3265 &self,
3266 x: &CudaSlice<f32>,
3267 col: usize,
3268 y: &mut CudaSlice<f32>,
3269 n: usize,
3270 seed: u64,
3271 stream_pos: u32,
3272 temp: f32,
3273 stat_max: &CudaSlice<f32>,
3274 stat_th: &CudaSlice<f32>,
3275 stat_idx: usize,
3276 ) -> Result<(), Box<dyn std::error::Error>> {
3277 let f = self.func("gumbel_perturb_filtered_col_f32");
3278 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3279 let (ci, si) = (col as i32, stat_idx as i32);
3280 let cfg = LaunchConfig {
3281 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3282 block_dim: (256, 1, 1),
3283 shared_mem_bytes: 0,
3284 };
3285 let __s_b = self.gpu.stream();
3286 let mut b = __s_b.launch_builder(&f);
3287 b.arg(x)
3288 .arg(&ci)
3289 .arg(&mut *y)
3290 .arg(&ni)
3291 .arg(&slo)
3292 .arg(&shi)
3293 .arg(&stream_pos)
3294 .arg(&temp)
3295 .arg(stat_max)
3296 .arg(stat_th)
3297 .arg(&si);
3298 unsafe {
3299 b.launch(cfg)?;
3300 }
3301 Ok(())
3302 }
3303
3304 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3309 let f = self.func("memra_sctr_inc");
3310 let cfg = LaunchConfig {
3311 grid_dim: (1, 1, 1),
3312 block_dim: (1, 1, 1),
3313 shared_mem_bytes: 0,
3314 };
3315 let __s_b = self.gpu.stream();
3316 let mut b = __s_b.launch_builder(&f);
3317 b.arg(&mut *ctr);
3318 unsafe {
3319 b.launch(cfg)?;
3320 }
3321 Ok(())
3322 }
3323
3324 pub fn gumbel_perturb_ctr(
3329 &self,
3330 x: &CudaSlice<f32>,
3331 y: &mut CudaSlice<f32>,
3332 n: usize,
3333 seed: u64,
3334 ctr: &CudaSlice<u32>,
3335 temp: f32,
3336 ) -> Result<(), Box<dyn std::error::Error>> {
3337 let f = self.func("gumbel_perturb_ctr_f32");
3338 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3339 let cfg = LaunchConfig {
3340 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3341 block_dim: (256, 1, 1),
3342 shared_mem_bytes: 0,
3343 };
3344 let __s_b = self.gpu.stream();
3345 let mut b = __s_b.launch_builder(&f);
3346 b.arg(x)
3347 .arg(&mut *y)
3348 .arg(&ni)
3349 .arg(&slo)
3350 .arg(&shi)
3351 .arg(ctr)
3352 .arg(&temp);
3353 unsafe {
3354 b.launch(cfg)?;
3355 }
3356 Ok(())
3357 }
3358
3359 pub fn softmax_gather(
3363 &self,
3364 x: &CudaSlice<f32>,
3365 row_stride: usize,
3366 ids: &CudaSlice<u32>,
3367 rows: &CudaSlice<i32>,
3368 out: &mut CudaSlice<f32>,
3369 n: usize,
3370 npair: usize,
3371 temp: f32,
3372 ) -> Result<(), Box<dyn std::error::Error>> {
3373 let f = self.func("softmax_gather_f32");
3374 let (ni, rs) = (n as i32, row_stride as i64);
3375 let np = npair as i32;
3376 let cfg = LaunchConfig {
3377 grid_dim: (npair as u32, 1, 1),
3378 block_dim: (256, 1, 1),
3379 shared_mem_bytes: 0,
3380 };
3381 let __s_b = self.gpu.stream();
3382 let mut b = __s_b.launch_builder(&f);
3383 b.arg(x)
3384 .arg(&rs)
3385 .arg(ids)
3386 .arg(rows)
3387 .arg(&mut *out)
3388 .arg(&ni)
3389 .arg(&np)
3390 .arg(&temp);
3391 unsafe {
3392 b.launch(cfg)?;
3393 }
3394 Ok(())
3395 }
3396
3397 pub fn residual_sample(
3401 &self,
3402 p: &CudaSlice<f32>,
3403 q: Option<&CudaSlice<f32>>,
3404 n: usize,
3405 temp: f32,
3406 seed: u64,
3407 stream_pos: u32,
3408 out_tok: &mut CudaSlice<u32>,
3409 ) -> Result<(), Box<dyn std::error::Error>> {
3410 let f = self.func("residual_sample_f32");
3411 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3412 let nth = 1024u32;
3413 let cfg = LaunchConfig {
3414 grid_dim: (1, 1, 1),
3415 block_dim: (nth, 1, 1),
3416 shared_mem_bytes: 0,
3417 };
3418 let has_q: i32 = q.is_some() as i32;
3419 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3421 let mut b = __s_b.launch_builder(&f);
3422 b.arg(p)
3423 .arg(qbuf)
3424 .arg(&has_q)
3425 .arg(&ni)
3426 .arg(&temp)
3427 .arg(&slo)
3428 .arg(&shi)
3429 .arg(&stream_pos)
3430 .arg(&mut *out_tok);
3431 unsafe {
3432 b.launch(cfg)?;
3433 }
3434 Ok(())
3435 }
3436
3437 pub fn with_moe_cache<R>(
3442 &self,
3443 max_block_bytes: usize,
3444 f: impl FnOnce(
3445 &mut crate::moe_cache::MoeSlotCache,
3446 &Engine,
3447 ) -> Result<R, Box<dyn std::error::Error>>,
3448 ) -> Result<R, Box<dyn std::error::Error>> {
3449 let mut guard = self.moe_cache.lock().unwrap();
3450 if guard.is_none() {
3451 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3452 }
3453 let cache = guard.as_mut().unwrap();
3454 f(cache, self)
3455 }
3456
3457 pub fn freeze_moe_cache(&self) {
3460 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3461 cache.freeze();
3462 }
3463 }
3464
3465 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3468 self.moe_cache
3469 .lock()
3470 .unwrap()
3471 .as_ref()
3472 .map(crate::moe_cache::MoeSlotCache::export_residency)
3473 }
3474
3475 pub(crate) fn moe_cache_frozen(&self) -> bool {
3476 self.moe_cache
3477 .lock()
3478 .unwrap()
3479 .as_ref()
3480 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3481 }
3482
3483 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3490 crate::cpu_experts::configured()
3491 && self.moe_cache_frozen()
3492 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3493 }
3494
3495 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3497 assert!(
3498 self.moe_cache.lock().unwrap().is_none(),
3499 "MoE cache layout configured after cache construction"
3500 );
3501 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3502 }
3503
3504 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3505 self.moe_cache_layout.lock().unwrap().clone()
3506 }
3507
3508 pub fn moe_cache_enabled() -> bool {
3510 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3511 }
3512
3513 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3516 let guard = self.moe_cache.lock().unwrap();
3517 guard
3518 .as_ref()
3519 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3520 }
3521
3522 pub fn cpu_expert_stats(
3526 &self,
3527 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3528 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3529 }
3530
3531 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3534 crate::cpu_experts::predictor_stats()
3535 }
3536
3537 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3538 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3539 }
3540
3541 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3544 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3545 }
3546
3547 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3550 let guard = self.moe_cache.lock().unwrap();
3551 guard
3552 .as_ref()
3553 .and_then(|cache| cache.pread_stats())
3554 .map(|stats| {
3555 (
3556 stats.reads,
3557 stats.bytes,
3558 stats.read_errors,
3559 stats.short_reads,
3560 stats.fallbacks,
3561 stats.buffer_waits,
3562 stats.ring_full,
3563 )
3564 })
3565 }
3566
3567 pub fn spill_config_fallbacks(&self) -> u64 {
3569 crate::spill_pread::config_fallbacks()
3570 }
3571
3572 pub fn moe_cache_reset_counters(&self) {
3574 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3575 c.reset_counters();
3576 }
3577 }
3578
3579 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3580 Ok(self.gpu.stream().clone_htod(v)?)
3581 }
3582
3583 pub fn htod_bytes_padded(
3587 &self,
3588 v: &[u8],
3589 pad: usize,
3590 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3591 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3592 {
3593 let mut view = d.slice_mut(0..v.len());
3594 self.gpu.stream().memcpy_htod(v, &mut view)?;
3595 }
3596 Ok(d)
3597 }
3598
3599 pub fn copy_into(
3601 &self,
3602 dst: &mut CudaSlice<f32>,
3603 off: usize,
3604 src: &CudaSlice<f32>,
3605 len: usize,
3606 ) -> Result<(), Box<dyn std::error::Error>> {
3607 let mut view = dst.slice_mut(off..off + len);
3608 self.gpu
3609 .stream()
3610 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3611 Ok(())
3612 }
3613
3614 pub fn copy_u8_into(
3617 &self,
3618 dst: &mut CudaSlice<u8>,
3619 off: usize,
3620 src: &CudaSlice<u8>,
3621 len: usize,
3622 ) -> Result<(), Box<dyn std::error::Error>> {
3623 let mut view = dst.slice_mut(off..off + len);
3624 self.gpu
3625 .stream()
3626 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3627 Ok(())
3628 }
3629
3630 pub fn copy_u8_range_into(
3632 &self,
3633 dst: &mut CudaSlice<u8>,
3634 dst_off: usize,
3635 src: &CudaSlice<u8>,
3636 src_off: usize,
3637 len: usize,
3638 ) -> Result<(), Box<dyn std::error::Error>> {
3639 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3640 self.gpu
3641 .stream()
3642 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3643 Ok(())
3644 }
3645
3646 pub fn prepare_kv_append(
3650 &self,
3651 kv: &mut crate::cache::KvLayer,
3652 retain_from: usize,
3653 append_rows: usize,
3654 ) -> Result<usize, Box<dyn std::error::Error>> {
3655 let Some(plan) = kv
3656 .ring
3657 .as_ref()
3658 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3659 .transpose()?
3660 else {
3661 return Ok(kv.len);
3662 };
3663 match plan {
3664 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3665 crate::cache::KvRingAppend::Rebase {
3666 src_row,
3667 keep_rows,
3668 new_base,
3669 write_row,
3670 } => {
3671 if keep_rows > 0 {
3672 let k_len = keep_rows * kv.k_tok_bytes;
3673 let v_len = keep_rows * kv.v_tok_bytes;
3674 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3675 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3676 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3677 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3678 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3679 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3680 }
3681 kv.ring.as_mut().unwrap().apply_rebase(new_base);
3682 Ok(write_row)
3683 }
3684 }
3685 }
3686
3687 pub fn htod_u8_into(
3690 &self,
3691 dst: &mut CudaSlice<u8>,
3692 off: usize,
3693 src: &[u8],
3694 ) -> Result<(), Box<dyn std::error::Error>> {
3695 let mut view = dst.slice_mut(off..off + src.len());
3696 self.gpu.stream().memcpy_htod(src, &mut view)?;
3697 Ok(())
3698 }
3699
3700 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3701 b.slice(0..len)
3702 }
3703
3704 pub fn view_u8_range<'a>(
3707 &self,
3708 b: &'a CudaSlice<u8>,
3709 start: usize,
3710 end: usize,
3711 ) -> cudarc::driver::CudaView<'a, u8> {
3712 b.slice(start..end)
3713 }
3714 pub fn view_u8<'a>(
3715 &self,
3716 b: &'a CudaSlice<u8>,
3717 len: usize,
3718 ) -> cudarc::driver::CudaView<'a, u8> {
3719 b.slice(0..len)
3720 }
3721
3722 pub fn append_kv_quantized(
3726 &self,
3727 k_row: &CudaSlice<f32>,
3728 v_row: &CudaSlice<f32>,
3729 kc: &mut CudaSlice<u8>,
3730 vc: &mut CudaSlice<u8>,
3731 t: usize,
3732 kv_dim_k: usize,
3733 kv_dim_v: usize,
3734 k_tok_bytes: usize,
3735 v_tok_bytes: usize,
3736 g: bool,
3737 ) -> Result<(), Box<dyn std::error::Error>> {
3738 let f = if g {
3739 self.func_g("append_quantize_kv_q8_0_q5_1")
3740 } else {
3741 self.func("append_quantize_kv_q8_0_q5_1")
3742 };
3743 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3744 let cfg = LaunchConfig {
3745 grid_dim: (nblk, 1, 1),
3746 block_dim: (32, 1, 1),
3747 shared_mem_bytes: 0,
3748 };
3749 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3750 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3751 let __s_b = self.gpu.stream();
3752 let mut b = __s_b.launch_builder(&f);
3753 b.arg(k_row)
3754 .arg(v_row)
3755 .arg(kc)
3756 .arg(vc)
3757 .arg(&ti)
3758 .arg(&kdk)
3759 .arg(&kdv)
3760 .arg(&ktb)
3761 .arg(&vtb);
3762 unsafe {
3763 b.launch(cfg)?;
3764 }
3765 Ok(())
3766 }
3767
3768 pub fn append_kv_quantized_dc(
3772 &self,
3773 k_row: &CudaSlice<f32>,
3774 v_row: &CudaSlice<f32>,
3775 kc: &mut CudaSlice<u8>,
3776 vc: &mut CudaSlice<u8>,
3777 t_dev: &CudaSlice<i32>,
3778 kv_dim_k: usize,
3779 kv_dim_v: usize,
3780 k_tok_bytes: usize,
3781 v_tok_bytes: usize,
3782 g: bool,
3783 ) -> Result<(), Box<dyn std::error::Error>> {
3784 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3785 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3786 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3787 if Self::pdl_on() && Self::pdl_wb_on() {
3789 use cudarc::driver::{DevicePtr, DevicePtrMut};
3790 let s = &self.gpu.stream();
3791 let (pk, _g0) = k_row.device_ptr(s);
3792 let (pv, _g1) = v_row.device_ptr(s);
3793 let (pkc, _g2) = kc.device_ptr_mut(s);
3794 let (pvc, _g3) = vc.device_ptr_mut(s);
3795 let (pt, _g4) = t_dev.device_ptr(s);
3796 let mut ps = [
3797 &pk as *const _ as *mut std::ffi::c_void,
3798 &pv as *const _ as *mut _,
3799 &pkc as *const _ as *mut _,
3800 &pvc as *const _ as *mut _,
3801 &pt as *const _ as *mut _,
3802 &kdk as *const _ as *mut _,
3803 &kdv as *const _ as *mut _,
3804 &ktb as *const _ as *mut _,
3805 &vtb as *const _ as *mut _,
3806 ];
3807 unsafe {
3808 self.launch_pdl_flash(
3809 g,
3810 "append_quantize_kv_q8_0_q5_1_dc",
3811 (nblk, 1, 1),
3812 (32, 1, 1),
3813 0,
3814 &mut ps,
3815 )?;
3816 }
3817 return Ok(());
3818 }
3819 let f = if g {
3820 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3821 } else {
3822 self.func("append_quantize_kv_q8_0_q5_1_dc")
3823 };
3824 let cfg = LaunchConfig {
3825 grid_dim: (nblk, 1, 1),
3826 block_dim: (32, 1, 1),
3827 shared_mem_bytes: 0,
3828 };
3829 let __s_b = self.gpu.stream();
3830 let mut b = __s_b.launch_builder(&f);
3831 b.arg(k_row)
3832 .arg(v_row)
3833 .arg(kc)
3834 .arg(vc)
3835 .arg(t_dev)
3836 .arg(&kdk)
3837 .arg(&kdv)
3838 .arg(&ktb)
3839 .arg(&vtb);
3840 unsafe {
3841 b.launch(cfg)?;
3842 }
3843 Ok(())
3844 }
3845
3846 #[allow(clippy::too_many_arguments)]
3853 pub fn append_kv_quantized_rows(
3854 &self,
3855 k_rows: &CudaSlice<f32>,
3856 v_rows: &CudaSlice<f32>,
3857 kc: &mut CudaSlice<u8>,
3858 vc: &mut CudaSlice<u8>,
3859 t0: usize,
3860 t: usize,
3861 kv_dim_k: usize,
3862 kv_dim_v: usize,
3863 k_tok_bytes: usize,
3864 v_tok_bytes: usize,
3865 g: bool,
3866 ) -> Result<(), Box<dyn std::error::Error>> {
3867 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3868 for i in 0..t {
3869 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3870 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3871 self.append_kv_quantized_view(
3872 &k_row,
3873 &v_row,
3874 kc,
3875 vc,
3876 t0 + i,
3877 kv_dim_k,
3878 kv_dim_v,
3879 k_tok_bytes,
3880 v_tok_bytes,
3881 g,
3882 )?;
3883 }
3884 return Ok(());
3885 }
3886 let f = if g {
3887 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3888 } else {
3889 self.func("append_quantize_kv_q8_0_q5_1_rows")
3890 };
3891 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3892 let cfg = LaunchConfig {
3893 grid_dim: (nblk, t as u32, 1),
3894 block_dim: (32, 1, 1),
3895 shared_mem_bytes: 0,
3896 };
3897 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3898 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3899 let __s_b = self.gpu.stream();
3900 let mut b = __s_b.launch_builder(&f);
3901 b.arg(k_rows)
3902 .arg(v_rows)
3903 .arg(kc)
3904 .arg(vc)
3905 .arg(&t0i)
3906 .arg(&kdk)
3907 .arg(&kdv)
3908 .arg(&ktb)
3909 .arg(&vtb);
3910 unsafe {
3911 b.launch(cfg)?;
3912 }
3913 Ok(())
3914 }
3915
3916 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3920 let f = self.func("inc_i32");
3921 let cfg = LaunchConfig {
3922 grid_dim: (1, 1, 1),
3923 block_dim: (1, 1, 1),
3924 shared_mem_bytes: 0,
3925 };
3926 let __s_b = self.gpu.stream();
3927 let mut b = __s_b.launch_builder(&f);
3928 b.arg(p);
3929 unsafe {
3930 b.launch(cfg)?;
3931 }
3932 Ok(())
3933 }
3934
3935 pub fn append_kv_quantized_view(
3938 &self,
3939 k_row: &cudarc::driver::CudaView<f32>,
3940 v_row: &cudarc::driver::CudaView<f32>,
3941 kc: &mut CudaSlice<u8>,
3942 vc: &mut CudaSlice<u8>,
3943 t: usize,
3944 kv_dim_k: usize,
3945 kv_dim_v: usize,
3946 k_tok_bytes: usize,
3947 v_tok_bytes: usize,
3948 g: bool,
3949 ) -> Result<(), Box<dyn std::error::Error>> {
3950 let f = if g {
3951 self.func_g("append_quantize_kv_q8_0_q5_1")
3952 } else {
3953 self.func("append_quantize_kv_q8_0_q5_1")
3954 };
3955 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3956 let cfg = LaunchConfig {
3957 grid_dim: (nblk, 1, 1),
3958 block_dim: (32, 1, 1),
3959 shared_mem_bytes: 0,
3960 };
3961 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3962 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3963 let __s_b = self.gpu.stream();
3964 let mut b = __s_b.launch_builder(&f);
3965 b.arg(k_row)
3966 .arg(v_row)
3967 .arg(kc)
3968 .arg(vc)
3969 .arg(&ti)
3970 .arg(&kdk)
3971 .arg(&kdv)
3972 .arg(&ktb)
3973 .arg(&vtb);
3974 unsafe {
3975 b.launch(cfg)?;
3976 }
3977 Ok(())
3978 }
3979
3980 pub fn copy_view_into(
3983 &self,
3984 dst: &mut CudaSlice<f32>,
3985 off: usize,
3986 src: &cudarc::driver::CudaView<f32>,
3987 len: usize,
3988 ) -> Result<(), Box<dyn std::error::Error>> {
3989 let mut view = dst.slice_mut(off..off + len);
3990 self.gpu
3991 .stream()
3992 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3993 Ok(())
3994 }
3995
3996 pub fn clone_dtod(
4000 &self,
4001 src: &CudaSlice<f32>,
4002 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4003 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4004 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4005 Ok(dst)
4006 }
4007
4008 pub fn dtod_copy_view(
4011 &self,
4012 src: &cudarc::driver::CudaView<f32>,
4013 dst: &mut CudaSlice<f32>,
4014 ) -> Result<(), Box<dyn std::error::Error>> {
4015 self.gpu.stream().memcpy_dtod(src, dst)?;
4016 Ok(())
4017 }
4018
4019 pub fn dtod_copy_view_i8(
4021 &self,
4022 src: &cudarc::driver::CudaView<i8>,
4023 dst: &mut CudaSlice<i8>,
4024 ) -> Result<(), Box<dyn std::error::Error>> {
4025 self.gpu.stream().memcpy_dtod(src, dst)?;
4026 Ok(())
4027 }
4028
4029 pub fn dtod_copy_into(
4031 &self,
4032 src: &CudaSlice<f32>,
4033 dst: &mut CudaSlice<f32>,
4034 offset: usize,
4035 ) -> Result<(), Box<dyn std::error::Error>> {
4036 let n = src.len();
4037 let mut dv = dst.slice_mut(offset..offset + n);
4038 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4039 Ok(())
4040 }
4041
4042 pub fn copy_batch_uniform_f32(
4048 &self,
4049 table: &CudaSlice<u64>,
4050 n: usize,
4051 words: usize,
4052 ) -> Result<(), Box<dyn std::error::Error>> {
4053 if n == 0 || words == 0 {
4054 return Ok(());
4055 }
4056 debug_assert!(
4057 table.len() >= 2 * n,
4058 "pointer table must hold n srcs + n dsts"
4059 );
4060 let f = self.func("copy_batch_uniform_f32");
4061 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4064 let (ni, wi) = (n as i32, words as i32);
4065 let cfg = LaunchConfig {
4066 grid_dim: (chunks, n as u32, 1),
4067 block_dim: (256, 1, 1),
4068 shared_mem_bytes: 0,
4069 };
4070 let __s = self.gpu.stream();
4071 let mut b = __s.launch_builder(&f);
4072 b.arg(table).arg(&ni).arg(&wi);
4073 unsafe {
4074 b.launch(cfg)?;
4075 }
4076 Ok(())
4077 }
4078
4079 pub fn htod_u64_into(
4082 &self,
4083 v: &[u64],
4084 dst: &mut CudaSlice<u64>,
4085 ) -> Result<(), Box<dyn std::error::Error>> {
4086 let mut view = dst.slice_mut(0..v.len());
4087 self.gpu.stream().memcpy_htod(v, &mut view)?;
4088 Ok(())
4089 }
4090
4091 pub fn copy_indirect_src_f32(
4096 &self,
4097 src_entry: &cudarc::driver::CudaView<u64>,
4098 dst: &mut CudaSlice<f32>,
4099 dst_off: usize,
4100 words: usize,
4101 ) -> Result<(), Box<dyn std::error::Error>> {
4102 let f = self.func("copy_indirect_src_f32");
4103 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4104 let wi = words as i32;
4105 let cfg = LaunchConfig {
4106 grid_dim: (chunks, 1, 1),
4107 block_dim: (256, 1, 1),
4108 shared_mem_bytes: 0,
4109 };
4110 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4111 let __s = self.gpu.stream();
4112 let mut b = __s.launch_builder(&f);
4113 b.arg(src_entry).arg(&mut dv).arg(&wi);
4114 unsafe {
4115 b.launch(cfg)?;
4116 }
4117 Ok(())
4118 }
4119
4120 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4122 self.alloc_uninit::<i8>(n)
4123 }
4124
4125 pub fn qmatvec(
4127 &self,
4128 w: &CudaSlice<u8>,
4129 x: &CudaSlice<f32>,
4130 m: usize,
4131 in_f: usize,
4132 out_f: usize,
4133 qtype: i32,
4134 row_bytes: usize,
4135 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4136 let f = self.func("qmatvec_f32");
4137 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4139 grid_dim: (out_f as u32, m as u32, 1),
4140 block_dim: (256, 1, 1),
4141 shared_mem_bytes: 0,
4142 };
4143 let (inf, outf, mi, qt, rb) =
4144 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4145 let __s_b = self.gpu.stream();
4146 let mut b = __s_b.launch_builder(&f);
4147 b.arg(w)
4148 .arg(x)
4149 .arg(&mut y)
4150 .arg(&inf)
4151 .arg(&outf)
4152 .arg(&mi)
4153 .arg(&qt)
4154 .arg(&rb);
4155 unsafe {
4156 b.launch(cfg)?;
4157 }
4158 Ok(y)
4159 }
4160
4161 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4163 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4164 self.keep_if_capturing(&s);
4165 Ok(s)
4166 }
4167
4168 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4172 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4173 self.keep_if_capturing(&s);
4174 Ok(s)
4175 }
4176
4177 pub fn memset_zeros_view(
4180 &self,
4181 dst: &mut cudarc::driver::CudaViewMut<f32>,
4182 ) -> Result<(), Box<dyn std::error::Error>> {
4183 self.gpu.stream().memset_zeros(dst)?;
4184 Ok(())
4185 }
4186
4187 pub fn stage_expert(
4193 &self,
4194 host_bytes: &[u8],
4195 scratch: &mut CudaSlice<u8>,
4196 off: usize,
4197 ) -> Result<(), Box<dyn std::error::Error>> {
4198 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4201 }
4202
4203 pub fn moe_router_topk(
4209 &self,
4210 logits: &CudaSlice<f32>,
4211 t: usize,
4212 n_expert: usize,
4213 n_used: usize,
4214 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4215 let f = self.func("moe_router_topk_f32");
4216 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 {
4219 grid_dim: (t as u32, 1, 1),
4220 block_dim: (n_expert as u32, 1, 1),
4221 shared_mem_bytes: 0,
4222 };
4223 let (ne, nu) = (n_expert as i32, n_used as i32);
4224 let __s_b = self.gpu.stream();
4225 let mut b = __s_b.launch_builder(&f);
4226 b.arg(logits)
4227 .arg(&mut sel_idx)
4228 .arg(&mut sel_w)
4229 .arg(&ne)
4230 .arg(&nu);
4231 unsafe {
4232 b.launch(cfg)?;
4233 }
4234 Ok((sel_idx, sel_w))
4235 }
4236
4237 pub fn moe_router_topk_scaled(
4240 &self,
4241 logits: &CudaSlice<f32>,
4242 t: usize,
4243 n_expert: usize,
4244 n_used: usize,
4245 ex_scale: &CudaSlice<f32>,
4246 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4247 let f = self.func("moe_router_topk_scaled_f32");
4252 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4253 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4254 let cfg = LaunchConfig {
4255 grid_dim: (t as u32, 1, 1),
4256 block_dim: (n_expert as u32, 1, 1),
4257 shared_mem_bytes: 0,
4258 };
4259 let (ne, nu) = (n_expert as i32, n_used as i32);
4260 let __s_b = self.gpu.stream();
4261 let mut b = __s_b.launch_builder(&f);
4262 b.arg(logits)
4263 .arg(&mut sel_idx)
4264 .arg(&mut sel_w)
4265 .arg(&ne)
4266 .arg(&nu)
4267 .arg(ex_scale);
4268 unsafe {
4269 b.launch(cfg)?;
4270 }
4271 Ok((sel_idx, sel_w))
4272 }
4273
4274 pub fn moe_router_topk_host(
4282 &self,
4283 logits: &CudaSlice<f32>,
4284 t: usize,
4285 n_expert: usize,
4286 n_used: usize,
4287 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4288 let f = self.func("moe_router_topk_f32");
4289 let n = t * n_used;
4290 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4291 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4292 let cfg = LaunchConfig {
4293 grid_dim: (t as u32, 1, 1),
4294 block_dim: (n_expert as u32, 1, 1),
4295 shared_mem_bytes: 0,
4296 };
4297 let (ne, nu) = (n_expert as i32, n_used as i32);
4298 let __s_b = self.gpu.stream();
4299 let mut b = __s_b.launch_builder(&f);
4300 b.arg(logits)
4301 .arg(&mut sel_idx)
4302 .arg(&mut sel_w)
4303 .arg(&ne)
4304 .arg(&nu);
4305 unsafe {
4306 b.launch(cfg)?;
4307 }
4308 let bytes = n * 8;
4310 let mut guard = self.router_stage.lock().unwrap();
4311 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4312 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4313 }
4314 let stage = guard.as_mut().unwrap();
4315 let (si, sw) = unsafe {
4316 (
4317 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4318 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4319 )
4320 };
4321 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()))
4325 }
4326
4327 #[allow(clippy::too_many_arguments)]
4331 pub fn moe_router_sigmoid_topk(
4332 &self,
4333 logits: &CudaSlice<f32>,
4334 t: usize,
4335 n_expert: usize,
4336 n_used: usize,
4337 active_count: usize,
4338 correction_bias: &CudaSlice<f32>,
4339 active: &CudaSlice<u8>,
4340 scaling_factor: f32,
4341 route_norm: bool,
4342 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4343 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4344 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4345 return Err(format!(
4346 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4347 )
4348 .into());
4349 }
4350 if logits.len() < t * n_expert
4351 || correction_bias.len() != n_expert
4352 || active.len() != n_expert
4353 {
4354 return Err(format!(
4355 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4356 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4357 ).into());
4358 }
4359 let f = self.func("moe_router_sigmoid_topk_f32");
4360 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4361 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4362 let threads = n_expert.div_ceil(32) * 32;
4363 let cfg = LaunchConfig {
4364 grid_dim: (t as u32, 1, 1),
4365 block_dim: (threads as u32, 1, 1),
4366 shared_mem_bytes: 0,
4367 };
4368 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4369 let __s_b = self.gpu.stream();
4370 let mut b = __s_b.launch_builder(&f);
4371 b.arg(logits)
4372 .arg(correction_bias)
4373 .arg(active)
4374 .arg(&mut sel_idx)
4375 .arg(&mut sel_w)
4376 .arg(&ne)
4377 .arg(&nu)
4378 .arg(&scaling_factor)
4379 .arg(&rn);
4380 unsafe {
4381 b.launch(cfg)?;
4382 }
4383 Ok((sel_idx, sel_w))
4384 }
4385
4386 #[allow(clippy::too_many_arguments)]
4389 pub fn moe_router_sigmoid_topk_host(
4390 &self,
4391 logits: &CudaSlice<f32>,
4392 t: usize,
4393 n_expert: usize,
4394 n_used: usize,
4395 active_count: usize,
4396 correction_bias: &CudaSlice<f32>,
4397 active: &CudaSlice<u8>,
4398 scaling_factor: f32,
4399 route_norm: bool,
4400 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4401 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4402 logits,
4403 t,
4404 n_expert,
4405 n_used,
4406 active_count,
4407 correction_bias,
4408 active,
4409 scaling_factor,
4410 route_norm,
4411 )?;
4412 let n = t * n_used;
4413 let bytes = n * 8;
4414 let mut guard = self.router_stage.lock().unwrap();
4415 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4416 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4417 }
4418 let stage = guard.as_mut().unwrap();
4419 let (si, sw) = unsafe {
4420 (
4421 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4422 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4423 )
4424 };
4425 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4426 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4427 self.gpu.stream().synchronize()?;
4428 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4429 }
4430
4431 pub fn stage_expert_async(
4435 &self,
4436 host_bytes: &[u8],
4437 scratch: &mut CudaSlice<u8>,
4438 off: usize,
4439 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4440 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4441 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4442 Ok(self.copy_stream.record_event(None)?)
4443 }
4444
4445 pub fn compute_wait(
4447 &self,
4448 ev: &cudarc::driver::CudaEvent,
4449 ) -> Result<(), Box<dyn std::error::Error>> {
4450 self.gpu.stream().wait(ev)?;
4451 Ok(())
4452 }
4453
4454 pub fn qmatvec_view(
4459 &self,
4460 w: &CudaSlice<u8>,
4461 range: std::ops::Range<usize>,
4462 x: &cudarc::driver::CudaView<f32>,
4463 m: usize,
4464 in_f: usize,
4465 out_f: usize,
4466 qtype: i32,
4467 row_bytes: usize,
4468 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4469 let f = self.func("qmatvec_f32");
4470 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4473 grid_dim: (out_f as u32, m as u32, 1),
4474 block_dim: (256, 1, 1),
4475 shared_mem_bytes: 0,
4476 };
4477 let (inf, outf, mi, qt, rb) =
4478 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4479 let __s_b = self.gpu.stream();
4480 let mut b = __s_b.launch_builder(&f);
4481 b.arg(&wv)
4482 .arg(x)
4483 .arg(&mut y)
4484 .arg(&inf)
4485 .arg(&outf)
4486 .arg(&mi)
4487 .arg(&qt)
4488 .arg(&rb);
4489 unsafe {
4490 b.launch(cfg)?;
4491 }
4492 Ok(y)
4493 }
4494
4495 #[allow(clippy::too_many_arguments)]
4502 pub fn moe_gate_up_silu8_q8(
4506 &self,
4507 gp: WPtr8,
4508 up: WPtr8,
4509 aq: &CudaSlice<i8>,
4510 ad: &CudaSlice<f32>,
4511 in_f: usize,
4512 n_ff: usize,
4513 n_used: usize,
4514 qt_g: i32,
4515 qt_u: i32,
4516 rb_g: usize,
4517 rb_u: usize,
4518 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4519 let f = self.func("moe_gate_up_silu8_q8");
4520 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4521 let cfg = LaunchConfig {
4522 grid_dim: (n_ff as u32, n_used as u32, 1),
4523 block_dim: (32, 1, 1),
4524 shared_mem_bytes: 0,
4525 };
4526 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4527 let __s_b = self.gpu.stream();
4528 let mut b = __s_b.launch_builder(&f);
4529 b.arg(&gp)
4530 .arg(&up)
4531 .arg(aq)
4532 .arg(ad)
4533 .arg(&mut act)
4534 .arg(&inf)
4535 .arg(&nff)
4536 .arg(&qt_g)
4537 .arg(&qt_u)
4538 .arg(&rbg)
4539 .arg(&rbu);
4540 unsafe {
4541 b.launch(cfg)?;
4542 }
4543 Ok(act)
4544 }
4545
4546 #[allow(clippy::too_many_arguments)]
4547 pub fn moe_down8_fma_q8(
4548 &self,
4549 dp: WPtr8,
4550 w: F32x8,
4551 aq2: &CudaSlice<i8>,
4552 ad2: &CudaSlice<f32>,
4553 dst: &mut cudarc::driver::CudaViewMut<f32>,
4554 in_f: usize,
4555 out_f: usize,
4556 n_used: usize,
4557 qt: i32,
4558 rb: usize,
4559 ) -> Result<(), Box<dyn std::error::Error>> {
4560 let f = self.func("moe_down8_fma_q8");
4561 let cfg = LaunchConfig {
4562 grid_dim: (out_f as u32, 1, 1),
4563 block_dim: (32, 1, 1),
4564 shared_mem_bytes: 0,
4565 };
4566 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4567 let __s_b = self.gpu.stream();
4568 let mut b = __s_b.launch_builder(&f);
4569 b.arg(&dp)
4570 .arg(&w)
4571 .arg(aq2)
4572 .arg(ad2)
4573 .arg(dst)
4574 .arg(&inf)
4575 .arg(&outf)
4576 .arg(&nu)
4577 .arg(&qt)
4578 .arg(&rbi);
4579 unsafe {
4580 b.launch(cfg)?;
4581 }
4582 Ok(())
4583 }
4584
4585 pub fn qmatvec_expert_q8(
4587 &self,
4588 w: &CudaSlice<u8>,
4589 range: std::ops::Range<usize>,
4590 aq: &CudaSlice<i8>,
4591 ad: &CudaSlice<f32>,
4592 m: usize,
4593 in_f: usize,
4594 out_f: usize,
4595 qtype: i32,
4596 row_bytes: usize,
4597 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4598 let f = self.func("qmatvec_expert_q8");
4599 let wv = w.slice(range);
4600 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4601 const ROWS: u32 = 4; let cfg = LaunchConfig {
4603 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4604 block_dim: (32, ROWS, 1),
4605 shared_mem_bytes: 0,
4606 };
4607 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4608 let __s_b = self.gpu.stream();
4609 let mut b = __s_b.launch_builder(&f);
4610 b.arg(&wv)
4611 .arg(aq)
4612 .arg(ad)
4613 .arg(&mut y)
4614 .arg(&inf)
4615 .arg(&outf)
4616 .arg(&mi)
4617 .arg(&qtype)
4618 .arg(&rbi);
4619 unsafe {
4620 b.launch(cfg)?;
4621 }
4622 Ok(y)
4623 }
4624
4625 pub fn moe_gate_up_silu8(
4626 &self,
4627 gp: WPtr8,
4628 up: WPtr8,
4629 x: &cudarc::driver::CudaView<f32>,
4630 in_f: usize,
4631 n_ff: usize,
4632 n_used: usize,
4633 qt_g: i32,
4634 qt_u: i32,
4635 rb_g: usize,
4636 rb_u: usize,
4637 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4638 let f = self.func("moe_gate_up_silu8_f32");
4639 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
4641 grid_dim: (n_ff as u32, n_used as u32, 1),
4642 block_dim: (256, 1, 1),
4643 shared_mem_bytes: 0,
4644 };
4645 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4646 let __s_b = self.gpu.stream();
4647 let mut b = __s_b.launch_builder(&f);
4648 b.arg(&gp)
4649 .arg(&up)
4650 .arg(x)
4651 .arg(&mut act)
4652 .arg(&inf)
4653 .arg(&nff)
4654 .arg(&qt_g)
4655 .arg(&qt_u)
4656 .arg(&rbg)
4657 .arg(&rbu);
4658 unsafe {
4659 b.launch(cfg)?;
4660 }
4661 Ok(act)
4662 }
4663
4664 #[allow(clippy::too_many_arguments)]
4670 pub fn moe_down8_fma_into(
4671 &self,
4672 dp: WPtr8,
4673 w: F32x8,
4674 act: &CudaSlice<f32>,
4675 dst: &mut cudarc::driver::CudaViewMut<f32>,
4676 in_f: usize,
4677 out_f: usize,
4678 n_used: usize,
4679 qt: i32,
4680 rb: usize,
4681 ) -> Result<(), Box<dyn std::error::Error>> {
4682 let f = self.func("moe_down8_fma_f32");
4683 let cfg = LaunchConfig {
4684 grid_dim: (out_f as u32, 1, 1),
4685 block_dim: (256, 1, 1),
4686 shared_mem_bytes: 0,
4687 };
4688 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4689 let __s_b = self.gpu.stream();
4690 let mut b = __s_b.launch_builder(&f);
4691 b.arg(&dp)
4692 .arg(&w)
4693 .arg(act)
4694 .arg(dst)
4695 .arg(&inf)
4696 .arg(&outf)
4697 .arg(&nu)
4698 .arg(&qt)
4699 .arg(&rbv);
4700 unsafe {
4701 b.launch(cfg)?;
4702 }
4703 Ok(())
4704 }
4705
4706 #[allow(clippy::too_many_arguments)]
4711 #[allow(clippy::too_many_arguments)]
4726 #[allow(clippy::too_many_arguments)]
4728 pub fn moe_pairs_matvec_q8(
4729 &self,
4730 table: &CudaSlice<u64>,
4731 proj: i32,
4732 pair_tok: &CudaSlice<i32>,
4733 pair_ex: &CudaSlice<i32>,
4734 aq: &CudaSlice<i8>,
4735 ad: &CudaSlice<f32>,
4736 in_f: usize,
4737 out_f: usize,
4738 n_expert: usize,
4739 n_pairs: usize,
4740 qtype: i32,
4741 row_bytes: usize,
4742 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4743 let f = self.func("moe_pairs_matvec_q8");
4744 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4745 const ROWS: u32 = 4;
4746 let cfg = LaunchConfig {
4747 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4748 block_dim: (32, ROWS, 1),
4749 shared_mem_bytes: 0,
4750 };
4751 let (inf, outf, ne, np, rbi) = (
4752 in_f as i32,
4753 out_f as i32,
4754 n_expert as i32,
4755 n_pairs as i32,
4756 row_bytes as i64,
4757 );
4758 let __s_b = self.gpu.stream();
4759 let mut b = __s_b.launch_builder(&f);
4760 b.arg(table)
4761 .arg(&proj)
4762 .arg(pair_tok)
4763 .arg(pair_ex)
4764 .arg(aq)
4765 .arg(ad)
4766 .arg(&mut y)
4767 .arg(&inf)
4768 .arg(&outf)
4769 .arg(&ne)
4770 .arg(&np)
4771 .arg(&qtype)
4772 .arg(&rbi);
4773 unsafe {
4774 b.launch(cfg)?;
4775 }
4776 Ok(y)
4777 }
4778
4779 #[allow(clippy::too_many_arguments)]
4781 pub fn moe_pairs_matvec_q8_em(
4782 &self,
4783 table: &CudaSlice<u64>,
4784 proj: i32,
4785 ex_ids: &CudaSlice<i32>,
4786 ex_off: &CudaSlice<i32>,
4787 ex_pairs: &CudaSlice<i32>,
4788 pair_tok: &CudaSlice<i32>,
4789 aq: &CudaSlice<i8>,
4790 ad: &CudaSlice<f32>,
4791 in_f: usize,
4792 out_f: usize,
4793 n_expert: usize,
4794 n_active: usize,
4795 n_pairs: usize,
4796 qtype: i32,
4797 row_bytes: usize,
4798 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4799 let f = self.func("moe_pairs_matvec_q8_em");
4800 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4801 const ROWS: u32 = 4;
4802 let cfg = LaunchConfig {
4803 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4804 block_dim: (32, ROWS, 1),
4805 shared_mem_bytes: 0,
4806 };
4807 let (inf, outf, ne, na, rbi) = (
4808 in_f as i32,
4809 out_f as i32,
4810 n_expert as i32,
4811 n_active as i32,
4812 row_bytes as i64,
4813 );
4814 let __s_b = self.gpu.stream();
4815 let mut b = __s_b.launch_builder(&f);
4816 b.arg(table)
4817 .arg(&proj)
4818 .arg(ex_ids)
4819 .arg(ex_off)
4820 .arg(ex_pairs)
4821 .arg(pair_tok)
4822 .arg(aq)
4823 .arg(ad)
4824 .arg(&mut y)
4825 .arg(&inf)
4826 .arg(&outf)
4827 .arg(&ne)
4828 .arg(&na)
4829 .arg(&qtype)
4830 .arg(&rbi);
4831 unsafe {
4832 b.launch(cfg)?;
4833 }
4834 Ok(y)
4835 }
4836
4837 #[allow(clippy::too_many_arguments)]
4840 pub fn moe_pairs_matvec_q8_dec(
4841 &self,
4842 table: &CudaSlice<u64>,
4843 proj: i32,
4844 ex_ids: &CudaSlice<i32>,
4845 ex_off: &CudaSlice<i32>,
4846 ex_pairs: &CudaSlice<i32>,
4847 pair_tok: &CudaSlice<i32>,
4848 aq: &CudaSlice<i8>,
4849 ad: &CudaSlice<f32>,
4850 in_f: usize,
4851 out_f: usize,
4852 n_expert: usize,
4853 n_active: usize,
4854 n_pairs: usize,
4855 qtype: i32,
4856 row_bytes: usize,
4857 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4858 let f = self.func("moe_pairs_matvec_q8_dec");
4859 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4860 const ROWS: u32 = 4;
4861 let cfg = LaunchConfig {
4862 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4863 block_dim: (32, ROWS, 1),
4864 shared_mem_bytes: 0,
4865 };
4866 let (inf, outf, ne, na, rbi) = (
4867 in_f as i32,
4868 out_f as i32,
4869 n_expert as i32,
4870 n_active as i32,
4871 row_bytes as i64,
4872 );
4873 let __s_b = self.gpu.stream();
4874 let mut b = __s_b.launch_builder(&f);
4875 b.arg(table)
4876 .arg(&proj)
4877 .arg(ex_ids)
4878 .arg(ex_off)
4879 .arg(ex_pairs)
4880 .arg(pair_tok)
4881 .arg(aq)
4882 .arg(ad)
4883 .arg(&mut y)
4884 .arg(&inf)
4885 .arg(&outf)
4886 .arg(&ne)
4887 .arg(&na)
4888 .arg(&qtype)
4889 .arg(&rbi);
4890 unsafe {
4891 b.launch(cfg)?;
4892 }
4893 Ok(y)
4894 }
4895
4896 pub fn moe_pairs_gelu_mul(
4897 &self,
4898 gate: &CudaSlice<f32>,
4899 up: &CudaSlice<f32>,
4900 n: usize,
4901 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4902 let f = self.func("moe_pairs_gelu_mul");
4903 let mut act = self.alloc_uninit::<f32>(n)?;
4904 let cfg = LaunchConfig::for_num_elems(n as u32);
4905 let nl = n as i64;
4906 let __s_b = self.gpu.stream();
4907 let mut b = __s_b.launch_builder(&f);
4908 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4909 unsafe {
4910 b.launch(cfg)?;
4911 }
4912 Ok(act)
4913 }
4914
4915 pub fn moe_pairs_silu_mul(
4916 &self,
4917 gate: &CudaSlice<f32>,
4918 up: &CudaSlice<f32>,
4919 n: usize,
4920 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4921 let f = self.func("moe_pairs_silu_mul");
4922 let mut act = self.alloc_uninit::<f32>(n)?;
4923 let cfg = LaunchConfig::for_num_elems(n as u32);
4924 let nl = n as i64;
4925 let __s_b = self.gpu.stream();
4926 let mut b = __s_b.launch_builder(&f);
4927 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4928 unsafe {
4929 b.launch(cfg)?;
4930 }
4931 Ok(act)
4932 }
4933
4934 #[allow(clippy::too_many_arguments)]
4935 pub fn moe_pairs_scatter(
4936 &self,
4937 y_down: &CudaSlice<f32>,
4938 pair_w: &CudaSlice<f32>,
4939 tok_pair_off: &CudaSlice<i32>,
4940 tok_pair_ids: &CudaSlice<i32>,
4941 moe_out: &mut CudaSlice<f32>,
4942 t: usize,
4943 n_embd: usize,
4944 ) -> Result<(), Box<dyn std::error::Error>> {
4945 let f = self.func("moe_pairs_scatter");
4946 let cfg = LaunchConfig {
4947 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4948 block_dim: (256, 1, 1),
4949 shared_mem_bytes: 0,
4950 };
4951 let ne = n_embd as i32;
4952 let __s_b = self.gpu.stream();
4953 let mut b = __s_b.launch_builder(&f);
4954 b.arg(y_down)
4955 .arg(pair_w)
4956 .arg(tok_pair_off)
4957 .arg(tok_pair_ids)
4958 .arg(moe_out)
4959 .arg(&ne);
4960 unsafe {
4961 b.launch(cfg)?;
4962 }
4963 Ok(())
4964 }
4965
4966 #[allow(clippy::too_many_arguments)]
4970 pub fn moe_gate_up_gelu8_dev_q8(
4971 &self,
4972 table: &CudaSlice<u64>,
4973 sel: &cudarc::driver::CudaView<i32>,
4974 aq: &CudaSlice<i8>,
4975 ad: &CudaSlice<f32>,
4976 in_f: usize,
4977 n_ff: usize,
4978 n_used: usize,
4979 n_expert: usize,
4980 qt_g: i32,
4981 qt_u: i32,
4982 rb_g: usize,
4983 rb_u: usize,
4984 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4985 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4986 let (inf, nff, ne, rbg, rbu) = (
4987 in_f as i32,
4988 n_ff as i32,
4989 n_expert as i32,
4990 rb_g as i64,
4991 rb_u as i64,
4992 );
4993 let f = self.func("moe_gate_up_gelu8_dev_q8");
4994 let cfg = LaunchConfig {
4995 grid_dim: (n_ff as u32, n_used as u32, 1),
4996 block_dim: (32, 1, 1),
4997 shared_mem_bytes: 0,
4998 };
4999 let __s_b = self.gpu.stream();
5000 let mut b = __s_b.launch_builder(&f);
5001 b.arg(table)
5002 .arg(sel)
5003 .arg(aq)
5004 .arg(ad)
5005 .arg(&mut act)
5006 .arg(&inf)
5007 .arg(&nff)
5008 .arg(&ne)
5009 .arg(&qt_g)
5010 .arg(&qt_u)
5011 .arg(&rbg)
5012 .arg(&rbu);
5013 unsafe {
5014 b.launch(cfg)?;
5015 }
5016 Ok(act)
5017 }
5018
5019 #[allow(clippy::too_many_arguments)]
5021 pub fn moe_gate_up_gelu8_dev_q8_rows(
5022 &self,
5023 table: &CudaSlice<u64>,
5024 sel: &CudaSlice<i32>,
5025 aq: &CudaSlice<i8>,
5026 ad: &CudaSlice<f32>,
5027 t: usize,
5028 in_f: usize,
5029 n_ff: usize,
5030 n_used: usize,
5031 n_expert: usize,
5032 qt_g: i32,
5033 qt_u: i32,
5034 rb_g: usize,
5035 rb_u: usize,
5036 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5037 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5038 let (inf, nff, ne, rbg, rbu, nu) = (
5039 in_f as i32,
5040 n_ff as i32,
5041 n_expert as i32,
5042 rb_g as i64,
5043 rb_u as i64,
5044 n_used as i32,
5045 );
5046 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5047 let cfg = LaunchConfig {
5048 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5049 block_dim: (32, 1, 1),
5050 shared_mem_bytes: 0,
5051 };
5052 let __s_b = self.gpu.stream();
5053 let mut b = __s_b.launch_builder(&f);
5054 b.arg(table)
5055 .arg(sel)
5056 .arg(aq)
5057 .arg(ad)
5058 .arg(&mut act)
5059 .arg(&inf)
5060 .arg(&nff)
5061 .arg(&ne)
5062 .arg(&qt_g)
5063 .arg(&qt_u)
5064 .arg(&rbg)
5065 .arg(&rbu)
5066 .arg(&nu);
5067 unsafe {
5068 b.launch(cfg)?;
5069 }
5070 Ok(act)
5071 }
5072
5073 #[allow(clippy::too_many_arguments)]
5075 pub fn moe_gate_up_gelu8_dev_q8_csr(
5076 &self,
5077 table: &CudaSlice<u64>,
5078 sel: &CudaSlice<i32>,
5079 aq: &CudaSlice<i8>,
5080 ad: &CudaSlice<f32>,
5081 n_pairs: usize,
5082 in_f: usize,
5083 n_ff: usize,
5084 n_used: usize,
5085 n_expert: usize,
5086 qt_g: i32,
5087 qt_u: i32,
5088 rb_g: usize,
5089 rb_u: usize,
5090 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5091 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5092 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5093 in_f as i32,
5094 n_ff as i32,
5095 n_expert as i32,
5096 rb_g as i64,
5097 rb_u as i64,
5098 n_used as i32,
5099 n_pairs as i32,
5100 );
5101 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5102 let cfg = LaunchConfig {
5103 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5104 block_dim: (32, 1, 1),
5105 shared_mem_bytes: 0,
5106 };
5107 let __s_b = self.gpu.stream();
5108 let mut b = __s_b.launch_builder(&f);
5109 b.arg(table)
5110 .arg(sel)
5111 .arg(aq)
5112 .arg(ad)
5113 .arg(&mut act)
5114 .arg(&inf)
5115 .arg(&nff)
5116 .arg(&ne)
5117 .arg(&qt_g)
5118 .arg(&qt_u)
5119 .arg(&rbg)
5120 .arg(&rbu)
5121 .arg(&nu)
5122 .arg(&npi);
5123 unsafe {
5124 b.launch(cfg)?;
5125 }
5126 Ok(act)
5127 }
5128
5129 #[allow(clippy::too_many_arguments)]
5131 pub fn moe_down8_fma_dev_q8_rows_g(
5132 &self,
5133 table: &CudaSlice<u64>,
5134 sel: &CudaSlice<i32>,
5135 w: &CudaSlice<f32>,
5136 aq2: &CudaSlice<i8>,
5137 ad2: &CudaSlice<f32>,
5138 dst: &mut CudaSlice<f32>,
5139 t: usize,
5140 in_f: usize,
5141 out_f: usize,
5142 n_used: usize,
5143 n_expert: usize,
5144 qt: i32,
5145 rb: usize,
5146 ) -> Result<(), Box<dyn std::error::Error>> {
5147 let (inf, outf, nu, ne, rbi) = (
5148 in_f as i32,
5149 out_f as i32,
5150 n_used as i32,
5151 n_expert as i32,
5152 rb as i64,
5153 );
5154 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5158 let f = self.func(if step_b1_w8 {
5159 "moe_down8_fma_dev_q8_rows_w8"
5160 } else {
5161 "moe_down8_fma_dev_q8_rows_g"
5162 });
5163 let cfg = LaunchConfig {
5164 grid_dim: (out_f as u32, 1, t as u32),
5165 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5166 shared_mem_bytes: 0,
5167 };
5168 let __s_b = self.gpu.stream();
5169 let mut b = __s_b.launch_builder(&f);
5170 b.arg(table)
5171 .arg(sel)
5172 .arg(w)
5173 .arg(aq2)
5174 .arg(ad2)
5175 .arg(dst)
5176 .arg(&inf)
5177 .arg(&outf)
5178 .arg(&nu)
5179 .arg(&ne)
5180 .arg(&qt)
5181 .arg(&rbi);
5182 unsafe {
5183 b.launch(cfg)?;
5184 }
5185 Ok(())
5186 }
5187
5188 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5192 let (out_f, in_f) = (2048usize, 2816usize);
5193 let nblk = in_f / 32;
5194 let mut seed = 0x9E3779B97F4A7C15u64;
5195 let mut rng = move || {
5196 seed = seed
5197 .wrapping_mul(6364136223846793005)
5198 .wrapping_add(1442695040888963407);
5199 (seed >> 33) as u8
5200 };
5201 let mut w = vec![0u8; out_f * nblk * 18];
5202 for b in w.iter_mut() {
5203 *b = rng();
5204 }
5205 for r in 0..out_f {
5206 for g in 0..nblk {
5207 let off = (r * nblk + g) * 18;
5208 w[off] = 0x00;
5209 w[off + 1] = 0x2C; }
5211 }
5212 let qplane = out_f * nblk * 16;
5213 let mut wrp = vec![0u8; w.len()];
5214 for r in 0..out_f {
5215 for g in 0..nblk {
5216 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5217 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5218 .copy_from_slice(&src[0..2]);
5219 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5220 }
5221 }
5222 let w_d = self.htod_bytes(&w)?;
5223 let wrp_d = self.htod_bytes(&wrp)?;
5224 let mut aq = vec![0i8; m * in_f];
5225 for v in aq.iter_mut() {
5226 *v = rng() as i8;
5227 }
5228 let aq_d = self.htod_i8(&aq)?;
5229 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5230 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5231 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5232 const RPB: u32 = 4;
5233 let cfg = LaunchConfig {
5234 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5235 block_dim: (32, RPB, 1),
5236 shared_mem_bytes: 0,
5237 };
5238 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5239 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5240 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5241 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5242 {
5243 let __s_b = self.gpu.stream();
5244 let mut b = __s_b.launch_builder(&fb);
5245 b.arg(&w_d)
5246 .arg(&aq_d)
5247 .arg(&ad_d)
5248 .arg(&mut y0)
5249 .arg(&inf)
5250 .arg(&outf)
5251 .arg(&mi)
5252 .arg(&rb);
5253 unsafe {
5254 b.launch(cfg)?;
5255 }
5256 let __s_b = self.gpu.stream();
5257 let mut b = __s_b.launch_builder(&fr);
5258 b.arg(&wrp_d)
5259 .arg(&aq_d)
5260 .arg(&ad_d)
5261 .arg(&mut y1)
5262 .arg(&inf)
5263 .arg(&outf)
5264 .arg(&mi)
5265 .arg(&qp);
5266 unsafe {
5267 b.launch(cfg)?;
5268 }
5269 }
5270 self.gpu.stream().synchronize()?;
5271 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5272 let nd = h0
5273 .iter()
5274 .zip(&h1)
5275 .filter(|(a, b)| a.to_bits() != b.to_bits())
5276 .count();
5277 if nd != 0 {
5278 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5279 }
5280 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5281 self.gpu.stream().synchronize()?;
5282 let t0 = std::time::Instant::now();
5283 for _ in 0..500 {
5284 if rp {
5285 let __s_b = self.gpu.stream();
5286 let mut b = __s_b.launch_builder(&fr);
5287 b.arg(&wrp_d)
5288 .arg(&aq_d)
5289 .arg(&ad_d)
5290 .arg(&mut y1)
5291 .arg(&inf)
5292 .arg(&outf)
5293 .arg(&mi)
5294 .arg(&qp);
5295 unsafe {
5296 b.launch(cfg)?;
5297 }
5298 } else {
5299 let __s_b = self.gpu.stream();
5300 let mut b = __s_b.launch_builder(&fb);
5301 b.arg(&w_d)
5302 .arg(&aq_d)
5303 .arg(&ad_d)
5304 .arg(&mut y0)
5305 .arg(&inf)
5306 .arg(&outf)
5307 .arg(&mi)
5308 .arg(&rb);
5309 unsafe {
5310 b.launch(cfg)?;
5311 }
5312 }
5313 }
5314 self.gpu.stream().synchronize()?;
5315 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5316 };
5317 let _ = time(false)?;
5318 let _ = time(true)?; Ok((time(false)?, time(true)?))
5320 }
5321
5322 pub fn build_q4_rp4(
5327 &self,
5328 t: &mut crate::model::GpuTensor,
5329 ) -> Result<(), Box<dyn std::error::Error>> {
5330 use crate::model::GpuTensor;
5331 let GpuTensor::Quant {
5332 bytes,
5333 qtype,
5334 row_bytes,
5335 ne,
5336 rp4,
5337 ..
5338 } = t
5339 else {
5340 return Ok(());
5341 };
5342 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5343 return Ok(());
5344 }
5345 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5346 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5347 return Ok(());
5348 }
5349 let nblk = in_f / 32;
5350 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5351 let f = self.func("q4_0_split_rp_build");
5352 let n = (out_f * nblk) as i32;
5353 let cfg = LaunchConfig {
5354 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5355 block_dim: (256, 1, 1),
5356 shared_mem_bytes: 0,
5357 };
5358 let (of, nb) = (out_f as i32, nblk as i32);
5359 let _ = n;
5360 let __s_b = self.gpu.stream();
5361 let mut b = __s_b.launch_builder(&f);
5362 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5363 unsafe {
5364 b.launch(cfg)?;
5365 }
5366 *rp4 = Some(dst);
5367 Ok(())
5368 }
5369
5370 pub fn build_q8_rp4(
5375 &self,
5376 t: &mut crate::model::GpuTensor,
5377 ) -> Result<(), Box<dyn std::error::Error>> {
5378 use crate::model::GpuTensor;
5379 let GpuTensor::Quant {
5380 bytes,
5381 qtype,
5382 row_bytes,
5383 ne,
5384 rp4,
5385 ..
5386 } = t
5387 else {
5388 return Ok(());
5389 };
5390 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5391 return Ok(());
5392 }
5393 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5394 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5395 return Ok(());
5396 }
5397 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5398 Ok(())
5399 }
5400
5401 pub fn build_q8_rp4_raw(
5404 &self,
5405 bytes: &CudaSlice<u8>,
5406 in_f: usize,
5407 out_f: usize,
5408 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5409 assert!(in_f % 32 == 0);
5410 let nblk = in_f / 32;
5411 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5412 let f = self.func("q8_0_split_rp_build");
5413 let cfg = LaunchConfig {
5414 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5415 block_dim: (256, 1, 1),
5416 shared_mem_bytes: 0,
5417 };
5418 let (of, nb) = (out_f as i32, nblk as i32);
5419 let __s_b = self.gpu.stream();
5420 let mut b = __s_b.launch_builder(&f);
5421 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5422 unsafe {
5423 b.launch(cfg)?;
5424 }
5425 Ok(dst)
5426 }
5427
5428 pub fn build_q4k_rp4(
5436 &self,
5437 t: &mut crate::model::GpuTensor,
5438 ) -> Result<(), Box<dyn std::error::Error>> {
5439 use crate::model::GpuTensor;
5440 let GpuTensor::Quant {
5441 bytes,
5442 qtype,
5443 row_bytes,
5444 ne,
5445 rp4,
5446 ..
5447 } = t
5448 else {
5449 return Ok(());
5450 };
5451 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5452 return Ok(());
5453 }
5454 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5455 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5456 return Ok(());
5457 }
5458 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5459 Ok(())
5460 }
5461
5462 pub fn build_q6k_rp4(
5463 &self,
5464 t: &mut crate::model::GpuTensor,
5465 ) -> Result<(), Box<dyn std::error::Error>> {
5466 use crate::model::GpuTensor;
5467 let GpuTensor::Quant {
5468 bytes,
5469 qtype,
5470 row_bytes,
5471 ne,
5472 rp4,
5473 ..
5474 } = t
5475 else {
5476 return Ok(());
5477 };
5478 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5479 return Ok(());
5480 }
5481 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5482 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5483 return Ok(());
5484 }
5485 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5486 Ok(())
5487 }
5488
5489 pub fn build_kq_rp4_raw(
5491 &self,
5492 bytes: &CudaSlice<u8>,
5493 in_f: usize,
5494 out_f: usize,
5495 qtype: i32,
5496 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5497 assert!(in_f % 256 == 0);
5498 let nsbk = in_f / 256;
5499 let (sb_bytes, kname) = match qtype {
5500 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5501 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5502 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5503 };
5504 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5505 let f = self.func(kname);
5506 let cfg = LaunchConfig {
5507 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5508 block_dim: (256, 1, 1),
5509 shared_mem_bytes: 0,
5510 };
5511 let (of, nb) = (out_f as i32, nsbk as i32);
5512 let __s_b = self.gpu.stream();
5513 let mut b = __s_b.launch_builder(&f);
5514 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5515 unsafe {
5516 b.launch(cfg)?;
5517 }
5518 Ok(dst)
5519 }
5520
5521 pub fn kqrp_enabled() -> bool {
5525 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5526 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5527 Ok("0") => false,
5528 Ok(_) => true,
5529 Err(_) => cfg!(memra_hopper_mma),
5530 })
5531 }
5532
5533 pub fn build_q4_rp_swap(
5539 &self,
5540 t: &mut crate::model::GpuTensor,
5541 ) -> Result<bool, Box<dyn std::error::Error>> {
5542 use crate::model::GpuTensor;
5543 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5553 return Ok(false);
5554 }
5555 self.build_q4_rp4(t)?;
5556 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5558 return Ok(false);
5559 };
5560 match rp4.take() {
5561 Some(split) => {
5562 *bytes = split; *rp = true;
5564 Ok(true)
5565 }
5566 None => Ok(false),
5567 }
5568 }
5569
5570 pub fn q4rp_enabled() -> bool {
5572 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5573 *ON.get_or_init(|| {
5574 std::env::var("MEMRA_Q4RP")
5575 .map(|v| v != "0")
5576 .unwrap_or(true)
5577 })
5578 }
5579
5580 pub fn copy_rows_strided(
5583 &self,
5584 src: &CudaSlice<f32>,
5585 dst: &mut CudaSlice<f32>,
5586 row_elems: usize,
5587 n_rows: usize,
5588 src_stride: usize,
5589 src_off: usize,
5590 ) -> Result<(), Box<dyn std::error::Error>> {
5591 let f = self.func("copy_rows_strided_f32");
5592 let cfg = LaunchConfig {
5593 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5594 block_dim: (256, 1, 1),
5595 shared_mem_bytes: 0,
5596 };
5597 let (re, nr) = (row_elems as i32, n_rows as i32);
5598 let (st, off) = (src_stride as i64, src_off as i64);
5599 let __s_b = self.gpu.stream();
5600 let mut b = __s_b.launch_builder(&f);
5601 b.arg(src)
5602 .arg(&mut *dst)
5603 .arg(&re)
5604 .arg(&nr)
5605 .arg(&st)
5606 .arg(&off);
5607 unsafe {
5608 b.launch(cfg)?;
5609 }
5610 Ok(())
5611 }
5612
5613 pub fn u32_set_k(
5615 &self,
5616 dst: &mut CudaSlice<u32>,
5617 v: u32,
5618 idx: usize,
5619 ) -> Result<(), Box<dyn std::error::Error>> {
5620 let f = self.func("u32_set_k");
5621 let cfg = LaunchConfig {
5622 grid_dim: (1, 1, 1),
5623 block_dim: (1, 1, 1),
5624 shared_mem_bytes: 0,
5625 };
5626 let ii = idx as i32;
5627 let __s_b = self.gpu.stream();
5628 let mut b = __s_b.launch_builder(&f);
5629 b.arg(dst).arg(&v).arg(&ii);
5630 unsafe {
5631 b.launch(cfg)?;
5632 }
5633 Ok(())
5634 }
5635
5636 pub fn i32_add_k(
5638 &self,
5639 d: &mut CudaSlice<i32>,
5640 v: i32,
5641 ) -> Result<(), Box<dyn std::error::Error>> {
5642 let f = self.func("i32_add_k");
5643 let cfg = LaunchConfig {
5644 grid_dim: (1, 1, 1),
5645 block_dim: (32, 1, 1),
5646 shared_mem_bytes: 0,
5647 };
5648 let __s_b = self.gpu.stream();
5649 let mut b = __s_b.launch_builder(&f);
5650 b.arg(d).arg(&v);
5651 unsafe {
5652 b.launch(cfg)?;
5653 }
5654 Ok(())
5655 }
5656
5657 pub fn i32_iota_from(
5659 &self,
5660 ctr: &CudaSlice<i32>,
5661 dst: &mut CudaSlice<i32>,
5662 n: usize,
5663 ) -> Result<(), Box<dyn std::error::Error>> {
5664 let f = self.func("i32_iota_from");
5665 let cfg = LaunchConfig::for_num_elems(n as u32);
5666 let ni = n as i32;
5667 let __s_b = self.gpu.stream();
5668 let mut b = __s_b.launch_builder(&f);
5669 b.arg(ctr).arg(dst).arg(&ni);
5670 unsafe {
5671 b.launch(cfg)?;
5672 }
5673 Ok(())
5674 }
5675
5676 pub fn u32_map_k(
5678 &self,
5679 buf: &mut CudaSlice<u32>,
5680 map: &CudaSlice<u32>,
5681 idx: usize,
5682 ) -> Result<(), Box<dyn std::error::Error>> {
5683 let f = self.func("u32_map_k");
5684 let cfg = LaunchConfig {
5685 grid_dim: (1, 1, 1),
5686 block_dim: (1, 1, 1),
5687 shared_mem_bytes: 0,
5688 };
5689 let ii = idx as i32;
5690 let __s_b = self.gpu.stream();
5691 let mut b = __s_b.launch_builder(&f);
5692 b.arg(buf).arg(map).arg(&ii);
5693 unsafe {
5694 b.launch(cfg)?;
5695 }
5696 Ok(())
5697 }
5698
5699 #[allow(clippy::too_many_arguments)]
5701 pub fn u32_pack2(
5702 &self,
5703 a: &CudaSlice<u32>,
5704 off_a: usize,
5705 n1: usize,
5706 b_in: &CudaSlice<u32>,
5707 n2: usize,
5708 out: &mut CudaSlice<u32>,
5709 ) -> Result<(), Box<dyn std::error::Error>> {
5710 let f = self.func("u32_pack2");
5711 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5712 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5713 let __s_b = self.gpu.stream();
5714 let mut b = __s_b.launch_builder(&f);
5715 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5716 unsafe {
5717 b.launch(cfg)?;
5718 }
5719 Ok(())
5720 }
5721
5722 pub fn moe_w_exscale(
5724 &self,
5725 w: &mut CudaSlice<f32>,
5726 sel: &CudaSlice<i32>,
5727 s: &CudaSlice<f32>,
5728 n: usize,
5729 ) -> Result<(), Box<dyn std::error::Error>> {
5730 let f = self.func("moe_w_exscale");
5731 let cfg = LaunchConfig::for_num_elems(n as u32);
5732 let ni = n as i32;
5733 let __s_b = self.gpu.stream();
5734 let mut b = __s_b.launch_builder(&f);
5735 b.arg(w).arg(sel).arg(s).arg(&ni);
5736 unsafe {
5737 b.launch(cfg)?;
5738 }
5739 Ok(())
5740 }
5741
5742 pub fn moe_w_scale_by_expert(
5745 &self,
5746 w: &mut CudaSlice<f32>,
5747 sel: &CudaSlice<i32>,
5748 macros: &CudaSlice<f32>,
5749 n_expert: usize,
5750 n: usize,
5751 ) -> Result<(), Box<dyn std::error::Error>> {
5752 let f = self.func("moe_w_scale_by_expert");
5753 let cfg = LaunchConfig {
5754 grid_dim: (n.div_ceil(64) as u32, 1, 1),
5755 block_dim: (64, 1, 1),
5756 shared_mem_bytes: 0,
5757 };
5758 let (ne, nn) = (n_expert as i32, n as i32);
5759 let __s_b = self.gpu.stream();
5760 let mut b = __s_b.launch_builder(&f);
5761 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5762 unsafe {
5763 b.launch(cfg)?;
5764 }
5765 Ok(())
5766 }
5767
5768 pub fn moe_gate_up_silu8_dev_q8(
5769 &self,
5770 table: &CudaSlice<u64>,
5771 sel: &cudarc::driver::CudaView<i32>,
5772 aq: &CudaSlice<i8>,
5773 ad: &CudaSlice<f32>,
5774 in_f: usize,
5775 n_ff: usize,
5776 n_used: usize,
5777 n_expert: usize,
5778 qt_g: i32,
5779 qt_u: i32,
5780 rb_g: usize,
5781 rb_u: usize,
5782 macros: &CudaSlice<f32>,
5783 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5784 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5785 let (mode, wpb) = GU.get_or_init(|| {
5786 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5787 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5788 .ok()
5789 .and_then(|v| v.parse().ok())
5790 .unwrap_or(4u32)
5791 .clamp(1, 16);
5792 (mode, wpb)
5793 });
5794 let (mode, wpb) = (mode.as_str(), *wpb);
5795 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5796 let (inf, nff, ne, rbg, rbu) = (
5797 in_f as i32,
5798 n_ff as i32,
5799 n_expert as i32,
5800 rb_g as i64,
5801 rb_u as i64,
5802 );
5803 let (f, cfg) = match mode {
5804 "1" | "2" | "4" => {
5805 let rpw: u32 = mode.parse().unwrap();
5806 let f = self.func(match rpw {
5807 1 => "moe_gate_up_silu8_dev_q8_r1",
5808 2 => "moe_gate_up_silu8_dev_q8_r2",
5809 _ => "moe_gate_up_silu8_dev_q8_r4",
5810 });
5811 let rows_per_block = (rpw * wpb) as usize;
5812 let gx = n_ff.div_ceil(rows_per_block) as u32;
5813 (
5814 f,
5815 LaunchConfig {
5816 grid_dim: (gx, n_used as u32, 1),
5817 block_dim: (32, wpb, 1),
5818 shared_mem_bytes: 0,
5819 },
5820 )
5821 }
5822 "j8" if n_used <= 32 => (
5823 self.func("moe_gate_up_silu8_dev_q8_j8"),
5824 LaunchConfig {
5825 grid_dim: (n_ff as u32, 1, 1),
5826 block_dim: (32, n_used as u32, 1),
5827 shared_mem_bytes: 0,
5828 },
5829 ),
5830 "vsm2" => {
5832 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5833 let sh = (rb_g + rb_u) as u32;
5834 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5835 f.set_attribute(
5836 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5837 sh as i32,
5838 )?;
5839 (
5840 f,
5841 LaunchConfig {
5842 grid_dim: (n_ff as u32, n_used as u32, 1),
5843 block_dim: (32, 1, 1),
5844 shared_mem_bytes: sh,
5845 },
5846 )
5847 }
5848 "vsm" => {
5849 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5850 let sh = (rb_g + rb_u) as u32;
5851 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5852 f.set_attribute(
5853 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5854 sh as i32,
5855 )?;
5856 (
5857 f,
5858 LaunchConfig {
5859 grid_dim: (n_ff as u32, n_used as u32, 1),
5860 block_dim: (32, 1, 1),
5861 shared_mem_bytes: sh,
5862 },
5863 )
5864 }
5865 "sg" => (
5866 self.func("moe_gate_up_silu8_dev_q8_sg"),
5867 LaunchConfig {
5868 grid_dim: (n_ff as u32, n_used as u32, 1),
5869 block_dim: (32, 1, 1),
5870 shared_mem_bytes: 0,
5871 },
5872 ),
5873 "j8sg" if n_used <= 32 => (
5874 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5875 LaunchConfig {
5876 grid_dim: (n_ff as u32, 1, 1),
5877 block_dim: (32, n_used as u32, 1),
5878 shared_mem_bytes: 0,
5879 },
5880 ),
5881 "u64" if in_f == 2048 => (
5882 self.func("moe_gate_up_silu8_dev_q8_u64"),
5883 LaunchConfig {
5884 grid_dim: (n_ff as u32, n_used as u32, 1),
5885 block_dim: (32, 1, 1),
5886 shared_mem_bytes: 0,
5887 },
5888 ),
5889 "gs4" if in_f == 2048 => (
5890 self.func("moe_gate_up_silu8_dev_q8_gs4"),
5891 LaunchConfig {
5892 grid_dim: (n_ff as u32, n_used as u32, 1),
5893 block_dim: (32, 4, 1),
5894 shared_mem_bytes: 0,
5895 },
5896 ),
5897 "v" | "" => (
5899 self.func("moe_gate_up_silu8_dev_q8_v"),
5900 LaunchConfig {
5901 grid_dim: (n_ff as u32, n_used as u32, 1),
5902 block_dim: (32, 1, 1),
5903 shared_mem_bytes: 0,
5904 },
5905 ),
5906 "s2" => (
5907 self.func("moe_gate_up_silu8_dev_q8_s2"),
5908 LaunchConfig {
5909 grid_dim: (n_ff as u32, n_used as u32, 1),
5910 block_dim: (32, 2, 1),
5911 shared_mem_bytes: 0,
5912 },
5913 ),
5914 "s2z" => {
5915 let rz = wpb.min(16); (
5917 self.func("moe_gate_up_silu8_dev_q8_s2z"),
5918 LaunchConfig {
5919 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5920 block_dim: (32, 2, rz),
5921 shared_mem_bytes: 0,
5922 },
5923 )
5924 }
5925 _ => (
5926 self.func("moe_gate_up_silu8_dev_q8"),
5927 LaunchConfig {
5928 grid_dim: (n_ff as u32, n_used as u32, 1),
5929 block_dim: (32, 1, 1),
5930 shared_mem_bytes: 0,
5931 },
5932 ),
5933 };
5934 let __s_b = self.gpu.stream();
5935 let mut b = __s_b.launch_builder(&f);
5936 b.arg(table)
5937 .arg(sel)
5938 .arg(aq)
5939 .arg(ad)
5940 .arg(&mut act)
5941 .arg(&inf)
5942 .arg(&nff)
5943 .arg(&ne)
5944 .arg(&qt_g)
5945 .arg(&qt_u)
5946 .arg(&rbg)
5947 .arg(&rbu)
5948 .arg(macros);
5949 unsafe {
5950 b.launch(cfg)?;
5951 }
5952 Ok(act)
5953 }
5954
5955 #[allow(clippy::too_many_arguments)]
5956 pub fn moe_down8_fma_dev_q8(
5957 &self,
5958 table: &CudaSlice<u64>,
5959 sel: &cudarc::driver::CudaView<i32>,
5960 w: &cudarc::driver::CudaView<f32>,
5961 aq2: &CudaSlice<i8>,
5962 ad2: &CudaSlice<f32>,
5963 dst: &mut cudarc::driver::CudaViewMut<f32>,
5964 in_f: usize,
5965 out_f: usize,
5966 n_used: usize,
5967 n_expert: usize,
5968 qt: i32,
5969 rb: usize,
5970 ) -> Result<(), Box<dyn std::error::Error>> {
5971 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5972 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5973 let (inf, outf, nu, ne, rbi) = (
5974 in_f as i32,
5975 out_f as i32,
5976 n_used as i32,
5977 n_expert as i32,
5978 rb as i64,
5979 );
5980 let (f, cfg) = match mode.as_str() {
5983 m @ ("1" | "2" | "4") if n_used <= 8 => {
5984 let rpw: usize = m.parse().unwrap();
5985 let f = self.func(match rpw {
5986 1 => "moe_down8_fma_dev_q8_w8r1",
5987 2 => "moe_down8_fma_dev_q8_w8r2",
5988 _ => "moe_down8_fma_dev_q8_w8r4",
5989 });
5990 (
5991 f,
5992 LaunchConfig {
5993 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5994 block_dim: (32, n_used as u32, 1),
5995 shared_mem_bytes: 0,
5996 },
5997 )
5998 }
5999 "h2" if in_f == 512 => (
6000 self.func("moe_down8_fma_dev_q8_h2"),
6001 LaunchConfig {
6002 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6003 block_dim: (32, 1, 1),
6004 shared_mem_bytes: 0,
6005 },
6006 ),
6007 "" if in_f == 704 && n_used <= 8 => (
6010 self.func("moe_down8_fma_dev_q8_w8r2"),
6011 LaunchConfig {
6012 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6013 block_dim: (32, n_used as u32, 1),
6014 shared_mem_bytes: 0,
6015 },
6016 ),
6017 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6021 self.func("moe_down8_fma_dev_q8_w8h2v"),
6022 LaunchConfig {
6023 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6024 block_dim: (32, n_used as u32, 1),
6025 shared_mem_bytes: 0,
6026 },
6027 ),
6028 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6029 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6030 LaunchConfig {
6031 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6032 block_dim: (32, n_used as u32, 1),
6033 shared_mem_bytes: 0,
6034 },
6035 ),
6036 "w8h2r2" if in_f == 512 && n_used <= 8 => (
6037 self.func("moe_down8_fma_dev_q8_w8h2r2"),
6038 LaunchConfig {
6039 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6040 block_dim: (32, n_used as u32, 1),
6041 shared_mem_bytes: 0,
6042 },
6043 ),
6044 "w8h2" if in_f == 512 && n_used <= 8 => (
6045 self.func("moe_down8_fma_dev_q8_w8h2"),
6046 LaunchConfig {
6047 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6048 block_dim: (32, n_used as u32, 1),
6049 shared_mem_bytes: 0,
6050 },
6051 ),
6052 _ => (
6053 self.func("moe_down8_fma_dev_q8"),
6054 LaunchConfig {
6055 grid_dim: (out_f as u32, 1, 1),
6056 block_dim: (32, 1, 1),
6057 shared_mem_bytes: 0,
6058 },
6059 ),
6060 };
6061 let __s_b = self.gpu.stream();
6062 let mut b = __s_b.launch_builder(&f);
6063 b.arg(table)
6064 .arg(sel)
6065 .arg(w)
6066 .arg(aq2)
6067 .arg(ad2)
6068 .arg(dst)
6069 .arg(&inf)
6070 .arg(&outf)
6071 .arg(&nu)
6072 .arg(&ne)
6073 .arg(&qt)
6074 .arg(&rbi);
6075 unsafe {
6076 b.launch(cfg)?;
6077 }
6078 Ok(())
6079 }
6080
6081 #[allow(clippy::too_many_arguments)]
6088 pub fn moe_gate_up_silu8_dev_q8_rows(
6089 &self,
6090 table: &CudaSlice<u64>,
6091 sel: &CudaSlice<i32>,
6092 aq: &CudaSlice<i8>,
6093 ad: &CudaSlice<f32>,
6094 t: usize,
6095 in_f: usize,
6096 n_ff: usize,
6097 n_used: usize,
6098 n_expert: usize,
6099 qt_g: i32,
6100 qt_u: i32,
6101 rb_g: usize,
6102 rb_u: usize,
6103 macros: &CudaSlice<f32>,
6104 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6105 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6106 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6107 let cfg = LaunchConfig {
6108 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6109 block_dim: (32, 1, 1),
6110 shared_mem_bytes: 0,
6111 };
6112 let (inf, nff, ne, nu, rbg, rbu) = (
6113 in_f as i32,
6114 n_ff as i32,
6115 n_expert as i32,
6116 n_used as i32,
6117 rb_g as i64,
6118 rb_u as i64,
6119 );
6120 let __s_b = self.gpu.stream();
6121 let mut b = __s_b.launch_builder(&f);
6122 b.arg(table)
6123 .arg(sel)
6124 .arg(aq)
6125 .arg(ad)
6126 .arg(&mut act)
6127 .arg(&inf)
6128 .arg(&nff)
6129 .arg(&ne)
6130 .arg(&qt_g)
6131 .arg(&qt_u)
6132 .arg(&rbg)
6133 .arg(&rbu)
6134 .arg(&nu)
6135 .arg(macros);
6136 unsafe {
6137 b.launch(cfg)?;
6138 }
6139 Ok(act)
6140 }
6141
6142 #[allow(clippy::too_many_arguments)]
6147 pub fn moe_down8_fma_dev_q8_rows(
6148 &self,
6149 table: &CudaSlice<u64>,
6150 sel: &CudaSlice<i32>,
6151 w: &CudaSlice<f32>,
6152 aq2: &CudaSlice<i8>,
6153 ad2: &CudaSlice<f32>,
6154 dst: &mut CudaSlice<f32>,
6155 t: usize,
6156 in_f: usize,
6157 out_f: usize,
6158 n_used: usize,
6159 n_expert: usize,
6160 qt: i32,
6161 rb: usize,
6162 ) -> Result<(), Box<dyn std::error::Error>> {
6163 assert!(
6164 in_f == 512 && n_used <= 8,
6165 "down rows twin is w8h2v shape-gated"
6166 );
6167 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6168 let cfg = LaunchConfig {
6169 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6170 block_dim: (32, n_used as u32, 1),
6171 shared_mem_bytes: 0,
6172 };
6173 let (inf, outf, nu, ne, rbi) = (
6174 in_f as i32,
6175 out_f as i32,
6176 n_used as i32,
6177 n_expert as i32,
6178 rb as i64,
6179 );
6180 let __s_b = self.gpu.stream();
6181 let mut b = __s_b.launch_builder(&f);
6182 b.arg(table)
6183 .arg(sel)
6184 .arg(w)
6185 .arg(aq2)
6186 .arg(ad2)
6187 .arg(dst)
6188 .arg(&inf)
6189 .arg(&outf)
6190 .arg(&nu)
6191 .arg(&ne)
6192 .arg(&qt)
6193 .arg(&rbi);
6194 unsafe {
6195 b.launch(cfg)?;
6196 }
6197 Ok(())
6198 }
6199
6200 #[allow(clippy::too_many_arguments)]
6204 pub fn moe_gate_up_silu8_dev_q8_csr(
6205 &self,
6206 table: &CudaSlice<u64>,
6207 sel: &CudaSlice<i32>,
6208 aq: &CudaSlice<i8>,
6209 ad: &CudaSlice<f32>,
6210 n_pairs: usize,
6211 in_f: usize,
6212 n_ff: usize,
6213 n_used: usize,
6214 n_expert: usize,
6215 qt_g: i32,
6216 qt_u: i32,
6217 rb_g: usize,
6218 rb_u: usize,
6219 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6220 let f = if qt_g == crate::QT_NVFP4 {
6223 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6224 } else {
6225 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6226 };
6227 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6228 let cfg = LaunchConfig {
6229 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6230 block_dim: (32, 1, 1),
6231 shared_mem_bytes: 0,
6232 };
6233 let (inf, nff, ne, nu, npi, rbg, rbu) = (
6234 in_f as i32,
6235 n_ff as i32,
6236 n_expert as i32,
6237 n_used as i32,
6238 n_pairs as i32,
6239 rb_g as i64,
6240 rb_u as i64,
6241 );
6242 let __s_b = self.gpu.stream();
6243 let mut b = __s_b.launch_builder(&f);
6244 b.arg(table)
6245 .arg(sel)
6246 .arg(aq)
6247 .arg(ad)
6248 .arg(&mut act)
6249 .arg(&inf)
6250 .arg(&nff)
6251 .arg(&ne)
6252 .arg(&qt_g)
6253 .arg(&qt_u)
6254 .arg(&rbg)
6255 .arg(&rbu)
6256 .arg(&nu)
6257 .arg(&npi);
6258 unsafe {
6259 b.launch(cfg)?;
6260 }
6261 Ok(act)
6262 }
6263
6264 #[allow(clippy::too_many_arguments)]
6268 pub fn moe_down8_fma_dev_q8_variant(
6269 &self,
6270 variant: &str,
6271 table: &CudaSlice<u64>,
6272 sel: &cudarc::driver::CudaView<i32>,
6273 w: &cudarc::driver::CudaView<f32>,
6274 aq2: &CudaSlice<i8>,
6275 ad2: &CudaSlice<f32>,
6276 dst: &mut cudarc::driver::CudaViewMut<f32>,
6277 in_f: usize,
6278 out_f: usize,
6279 n_used: usize,
6280 n_expert: usize,
6281 qt: i32,
6282 rb: usize,
6283 ) -> Result<(), Box<dyn std::error::Error>> {
6284 let (inf, outf, nu, ne, rbi) = (
6285 in_f as i32,
6286 out_f as i32,
6287 n_used as i32,
6288 n_expert as i32,
6289 rb as i64,
6290 );
6291 let (f, cfg) = match variant {
6292 "w8h2" | "w8h2v" => (
6293 self.func(if variant == "w8h2" {
6294 "moe_down8_fma_dev_q8_w8h2"
6295 } else {
6296 "moe_down8_fma_dev_q8_w8h2v"
6297 }),
6298 LaunchConfig {
6299 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6300 block_dim: (32, n_used as u32, 1),
6301 shared_mem_bytes: 0,
6302 },
6303 ),
6304 "w8h2r2" | "w8h2r2v" => (
6305 self.func(if variant == "w8h2r2" {
6306 "moe_down8_fma_dev_q8_w8h2r2"
6307 } else {
6308 "moe_down8_fma_dev_q8_w8h2r2v"
6309 }),
6310 LaunchConfig {
6311 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6312 block_dim: (32, n_used as u32, 1),
6313 shared_mem_bytes: 0,
6314 },
6315 ),
6316 _ => (
6317 self.func("moe_down8_fma_dev_q8"),
6318 LaunchConfig {
6319 grid_dim: (out_f as u32, 1, 1),
6320 block_dim: (32, 1, 1),
6321 shared_mem_bytes: 0,
6322 },
6323 ),
6324 };
6325 let __s_b = self.gpu.stream();
6326 let mut b = __s_b.launch_builder(&f);
6327 b.arg(table)
6328 .arg(sel)
6329 .arg(w)
6330 .arg(aq2)
6331 .arg(ad2)
6332 .arg(dst)
6333 .arg(&inf)
6334 .arg(&outf)
6335 .arg(&nu)
6336 .arg(&ne)
6337 .arg(&qt)
6338 .arg(&rbi);
6339 unsafe {
6340 b.launch(cfg)?;
6341 }
6342 Ok(())
6343 }
6344
6345 #[allow(clippy::too_many_arguments)]
6347 pub fn moe_gate_up_silu8_dev_q8_variant(
6348 &self,
6349 variant: &str,
6350 table: &CudaSlice<u64>,
6351 sel: &cudarc::driver::CudaView<i32>,
6352 aq: &CudaSlice<i8>,
6353 ad: &CudaSlice<f32>,
6354 in_f: usize,
6355 n_ff: usize,
6356 n_used: usize,
6357 n_expert: usize,
6358 qt_g: i32,
6359 qt_u: i32,
6360 rb_g: usize,
6361 rb_u: usize,
6362 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6363 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6364 let (inf, nff, ne, rbg, rbu) = (
6365 in_f as i32,
6366 n_ff as i32,
6367 n_expert as i32,
6368 rb_g as i64,
6369 rb_u as i64,
6370 );
6371 let f = self.func(if variant == "v" {
6372 "moe_gate_up_silu8_dev_q8_v"
6373 } else {
6374 "moe_gate_up_silu8_dev_q8"
6375 });
6376 let cfg = LaunchConfig {
6377 grid_dim: (n_ff as u32, n_used as u32, 1),
6378 block_dim: (32, 1, 1),
6379 shared_mem_bytes: 0,
6380 };
6381 let __s_b = self.gpu.stream();
6382 let mut b = __s_b.launch_builder(&f);
6383 b.arg(table)
6384 .arg(sel)
6385 .arg(aq)
6386 .arg(ad)
6387 .arg(&mut act)
6388 .arg(&inf)
6389 .arg(&nff)
6390 .arg(&ne)
6391 .arg(&qt_g)
6392 .arg(&qt_u)
6393 .arg(&rbg)
6394 .arg(&rbu);
6395 unsafe {
6396 b.launch(cfg)?;
6397 }
6398 Ok(act)
6399 }
6400
6401 pub fn moe_gate_up_silu8_dev(
6402 &self,
6403 table: &CudaSlice<u64>,
6404 sel: &cudarc::driver::CudaView<i32>,
6405 x: &cudarc::driver::CudaView<f32>,
6406 in_f: usize,
6407 n_ff: usize,
6408 n_used: usize,
6409 n_expert: usize,
6410 qt_g: i32,
6411 qt_u: i32,
6412 rb_g: usize,
6413 rb_u: usize,
6414 macros: &CudaSlice<f32>,
6415 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6416 let f = self.func("moe_gate_up_silu8_dev");
6417 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6419 grid_dim: (n_ff as u32, n_used as u32, 1),
6420 block_dim: (256, 1, 1),
6421 shared_mem_bytes: 0,
6422 };
6423 let (inf, nff, ne, rbg, rbu) = (
6424 in_f as i32,
6425 n_ff as i32,
6426 n_expert as i32,
6427 rb_g as i64,
6428 rb_u as i64,
6429 );
6430 let __s_b = self.gpu.stream();
6431 let mut b = __s_b.launch_builder(&f);
6432 b.arg(table)
6433 .arg(sel)
6434 .arg(x)
6435 .arg(&mut act)
6436 .arg(&inf)
6437 .arg(&nff)
6438 .arg(&ne)
6439 .arg(&qt_g)
6440 .arg(&qt_u)
6441 .arg(&rbg)
6442 .arg(&rbu)
6443 .arg(macros);
6444 unsafe {
6445 b.launch(cfg)?;
6446 }
6447 Ok(act)
6448 }
6449
6450 #[allow(clippy::too_many_arguments)]
6453 pub fn moe_down8_fma_dev(
6454 &self,
6455 table: &CudaSlice<u64>,
6456 sel: &cudarc::driver::CudaView<i32>,
6457 w: &cudarc::driver::CudaView<f32>,
6458 act: &CudaSlice<f32>,
6459 dst: &mut cudarc::driver::CudaViewMut<f32>,
6460 in_f: usize,
6461 out_f: usize,
6462 n_used: usize,
6463 n_expert: usize,
6464 qt: i32,
6465 rb: usize,
6466 ) -> Result<(), Box<dyn std::error::Error>> {
6467 let f = self.func("moe_down8_fma_dev");
6468 let cfg = LaunchConfig {
6469 grid_dim: (out_f as u32, 1, 1),
6470 block_dim: (256, 1, 1),
6471 shared_mem_bytes: 0,
6472 };
6473 let (inf, outf, nu, ne, rbv) = (
6474 in_f as i32,
6475 out_f as i32,
6476 n_used as i32,
6477 n_expert as i32,
6478 rb as i64,
6479 );
6480 let __s_b = self.gpu.stream();
6481 let mut b = __s_b.launch_builder(&f);
6482 b.arg(table)
6483 .arg(sel)
6484 .arg(w)
6485 .arg(act)
6486 .arg(dst)
6487 .arg(&inf)
6488 .arg(&outf)
6489 .arg(&nu)
6490 .arg(&ne)
6491 .arg(&qt)
6492 .arg(&rbv);
6493 unsafe {
6494 b.launch(cfg)?;
6495 }
6496 Ok(())
6497 }
6498
6499 pub fn axpy_into(
6501 &self,
6502 src: &CudaSlice<f32>,
6503 alpha: f32,
6504 dst: &mut cudarc::driver::CudaViewMut<f32>,
6505 n: usize,
6506 ) -> Result<(), Box<dyn std::error::Error>> {
6507 let f = self.func("axpy_f32");
6508 let cfg = LaunchConfig::for_num_elems(n as u32);
6509 let (a, ni) = (alpha, n as i32);
6510 let __s_b = self.gpu.stream();
6511 let mut b = __s_b.launch_builder(&f);
6512 b.arg(src).arg(dst).arg(&a).arg(&ni);
6513 unsafe {
6514 b.launch(cfg)?;
6515 }
6516 Ok(())
6517 }
6518
6519 pub fn add_scaled_rows(
6521 &self,
6522 src: &CudaSlice<f32>,
6523 scale: &CudaSlice<f32>,
6524 dst: &mut CudaSlice<f32>,
6525 ncols: usize,
6526 nrows: usize,
6527 ) -> Result<(), Box<dyn std::error::Error>> {
6528 let f = self.func("add_scaled_rows_f32");
6529 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6530 let (nc, nr) = (ncols as i32, nrows as i32);
6531 let __s_b = self.gpu.stream();
6532 let mut b = __s_b.launch_builder(&f);
6533 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6534 unsafe {
6535 b.launch(cfg)?;
6536 }
6537 Ok(())
6538 }
6539
6540 pub fn gather_rows(
6544 &self,
6545 src: &CudaSlice<f32>,
6546 idx: &CudaSlice<i32>,
6547 dst: &mut CudaSlice<f32>,
6548 ncols: usize,
6549 m_e: usize,
6550 ) -> Result<(), Box<dyn std::error::Error>> {
6551 let f = self.func("gather_rows_f32");
6552 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6553 let (nc, me) = (ncols as i32, m_e as i32);
6554 let __s_b = self.gpu.stream();
6555 let mut b = __s_b.launch_builder(&f);
6556 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6557 unsafe {
6558 b.launch(cfg)?;
6559 }
6560 Ok(())
6561 }
6562
6563 pub fn scatter_slot(
6568 &self,
6569 src: &CudaSlice<f32>,
6570 tok_idx: &CudaSlice<i32>,
6571 slot_idx: &CudaSlice<i32>,
6572 weight: &CudaSlice<f32>,
6573 dst: &mut CudaSlice<f32>,
6574 wbuf: &mut CudaSlice<f32>,
6575 ncols: usize,
6576 n_used: usize,
6577 m_e: usize,
6578 ) -> Result<(), Box<dyn std::error::Error>> {
6579 let f = self.func("scatter_add_slot_f32");
6580 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6581 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6582 let __s_b = self.gpu.stream();
6583 let mut b = __s_b.launch_builder(&f);
6584 b.arg(src)
6585 .arg(tok_idx)
6586 .arg(slot_idx)
6587 .arg(weight)
6588 .arg(dst)
6589 .arg(wbuf)
6590 .arg(&nc)
6591 .arg(&nu)
6592 .arg(&me);
6593 unsafe {
6594 b.launch(cfg)?;
6595 }
6596 Ok(())
6597 }
6598
6599 pub fn reduce_slots(
6603 &self,
6604 slots: &CudaSlice<f32>,
6605 wbuf: &CudaSlice<f32>,
6606 dst: &mut CudaSlice<f32>,
6607 ncols: usize,
6608 n_used: usize,
6609 t: usize,
6610 ) -> Result<(), Box<dyn std::error::Error>> {
6611 let f = self.func("reduce_slots_f32");
6612 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6613 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6614 let __s_b = self.gpu.stream();
6615 let mut b = __s_b.launch_builder(&f);
6616 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6617 unsafe {
6618 b.launch(cfg)?;
6619 }
6620 Ok(())
6621 }
6622
6623 pub fn quantize_q8_1_view(
6630 &self,
6631 x: &cudarc::driver::CudaView<f32>,
6632 m: usize,
6633 in_f: usize,
6634 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6635 let f = self.func("quantize_q8_1");
6636 let nblk = in_f / 32;
6637 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6638 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6639 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6640 let (inf, mi) = (in_f as i32, m as i32);
6641 let __s_b = self.gpu.stream();
6642 let mut b = __s_b.launch_builder(&f);
6643 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6644 unsafe {
6645 b.launch(cfg)?;
6646 }
6647 Ok((q, d))
6648 }
6649
6650 pub fn quantize_q8_1(
6651 &self,
6652 x: &CudaSlice<f32>,
6653 m: usize,
6654 in_f: usize,
6655 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6656 let nblk = in_f / 32;
6657 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);
6661 let (inf, mi) = (in_f as i32, m as i32);
6662 if Self::pdl_on() && Self::pdl_wb_on() {
6663 {
6664 use cudarc::driver::{DevicePtr, DevicePtrMut};
6665 let s = &self.gpu.stream();
6666 let (px, _g0) = x.device_ptr(s);
6667 let (pq, _g1) = q.device_ptr_mut(s);
6668 let (pd, _g2) = d.device_ptr_mut(s);
6669 let mut ps = [
6670 &px as *const _ as *mut std::ffi::c_void,
6671 &pq as *const _ as *mut _,
6672 &pd as *const _ as *mut _,
6673 &inf as *const _ as *mut _,
6674 &mi as *const _ as *mut _,
6675 ];
6676 unsafe {
6677 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6678 }
6679 }
6680 return Ok((q, d));
6681 }
6682 let f = self.func("quantize_q8_1");
6683 let __s_b = self.gpu.stream();
6684 let mut b = __s_b.launch_builder(&f);
6685 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6686 unsafe {
6687 b.launch(cfg)?;
6688 }
6689 Ok((q, d))
6690 }
6691
6692 pub fn quantize_fp4_act(
6696 &self,
6697 x: &CudaSlice<f32>,
6698 m: usize,
6699 in_f: usize,
6700 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6701 let f = self.func("quantize_fp4_act");
6702 let nb16 = in_f / 16;
6703 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);
6706 let (inf, mi) = (in_f as i32, m as i32);
6707 let __s_b = self.gpu.stream();
6708 let mut b = __s_b.launch_builder(&f);
6709 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6710 unsafe {
6711 b.launch(cfg)?;
6712 }
6713 Ok((aq4, ad4))
6714 }
6715
6716 pub fn qmatvec_gemm_nvfp4_fp4(
6721 &self,
6722 bytes: &CudaSlice<u8>,
6723 x: &CudaSlice<f32>,
6724 m: usize,
6725 in_f: usize,
6726 out_f: usize,
6727 row_bytes: usize,
6728 scale: f32,
6729 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6730 assert!(
6731 in_f % 64 == 0,
6732 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6733 );
6734 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6735 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6736 if scale != 1.0 {
6737 self.scale_inplace(&mut y, scale, m * out_f)?;
6738 }
6739 Ok(y)
6740 }
6741
6742 fn fp4_gemm_launch(
6745 &self,
6746 bytes: &CudaSlice<u8>,
6747 aq4: &CudaSlice<u32>,
6748 ad4: &CudaSlice<u8>,
6749 m: usize,
6750 in_f: usize,
6751 out_f: usize,
6752 row_bytes: usize,
6753 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6754 let f = self.func("qmatvec_gemm_nvfp4_fp4");
6755 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
6757 const BN: u32 = 256;
6758 let cfg = LaunchConfig {
6759 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6760 block_dim: (32, 4, 1),
6761 shared_mem_bytes: 0,
6762 };
6763 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6764 let __s_b = self.gpu.stream();
6765 let mut b = __s_b.launch_builder(&f);
6766 b.arg(bytes)
6767 .arg(aq4)
6768 .arg(ad4)
6769 .arg(&mut y)
6770 .arg(&inf)
6771 .arg(&outf)
6772 .arg(&mi)
6773 .arg(&rb);
6774 unsafe {
6775 b.launch(cfg)?;
6776 }
6777 Ok(y)
6778 }
6779
6780 pub fn qmatvec_gemm_nvfp4_fp4_raw(
6782 &self,
6783 bytes: &CudaSlice<u8>,
6784 x: &CudaSlice<f32>,
6785 m: usize,
6786 in_f: usize,
6787 out_f: usize,
6788 row_bytes: usize,
6789 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6790 assert!(
6791 in_f % 64 == 0,
6792 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6793 );
6794 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6795 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6796 }
6797
6798 pub fn qmatvec_q8_0_fast(
6800 &self,
6801 w: &CudaSlice<u8>,
6802 x: &CudaSlice<f32>,
6803 m: usize,
6804 in_f: usize,
6805 out_f: usize,
6806 row_bytes: usize,
6807 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6808 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6809 let f = self.func("qmatvec_q8_0_dp4a");
6810 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6812 grid_dim: (out_f as u32, m as u32, 1),
6813 block_dim: (128, 1, 1),
6814 shared_mem_bytes: 0,
6815 };
6816 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6817 let __s_b = self.gpu.stream();
6818 let mut b = __s_b.launch_builder(&f);
6819 b.arg(w)
6820 .arg(&aq)
6821 .arg(&ad)
6822 .arg(&mut y)
6823 .arg(&inf)
6824 .arg(&outf)
6825 .arg(&mi)
6826 .arg(&rb);
6827 unsafe {
6828 b.launch(cfg)?;
6829 }
6830 Ok(y)
6831 }
6832
6833 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
6836 &self,
6837 w: &CudaSlice<u8>,
6838 x: &CudaSlice<f32>,
6839 m: usize,
6840 in_f: usize,
6841 out_f: usize,
6842 row_bytes: usize,
6843 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6844 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6845 let f = self.func("qmatvec_q4_K_dp4a");
6846 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6848 grid_dim: (out_f as u32, m as u32, 1),
6849 block_dim: (128, 1, 1),
6850 shared_mem_bytes: 0,
6851 };
6852 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6853 let __s_b = self.gpu.stream();
6854 let mut b = __s_b.launch_builder(&f);
6855 b.arg(w)
6856 .arg(&aq)
6857 .arg(&ad)
6858 .arg(&mut y)
6859 .arg(&inf)
6860 .arg(&outf)
6861 .arg(&mi)
6862 .arg(&rb);
6863 unsafe {
6864 b.launch(cfg)?;
6865 }
6866 Ok(y)
6867 }
6868
6869 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
6872 &self,
6873 w: &CudaSlice<u8>,
6874 x: &CudaSlice<f32>,
6875 m: usize,
6876 in_f: usize,
6877 out_f: usize,
6878 row_bytes: usize,
6879 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6880 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6881 let f = self.func("qmatvec_q6_K_dp4a");
6882 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6884 grid_dim: (out_f as u32, m as u32, 1),
6885 block_dim: (128, 1, 1),
6886 shared_mem_bytes: 0,
6887 };
6888 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6889 let __s_b = self.gpu.stream();
6890 let mut b = __s_b.launch_builder(&f);
6891 b.arg(w)
6892 .arg(&aq)
6893 .arg(&ad)
6894 .arg(&mut y)
6895 .arg(&inf)
6896 .arg(&outf)
6897 .arg(&mi)
6898 .arg(&rb);
6899 unsafe {
6900 b.launch(cfg)?;
6901 }
6902 Ok(y)
6903 }
6904
6905 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
6908 &self,
6909 w: &CudaSlice<u8>,
6910 x: &CudaSlice<f32>,
6911 m: usize,
6912 in_f: usize,
6913 out_f: usize,
6914 row_bytes: usize,
6915 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6916 self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6917 }
6918 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
6921 &self,
6922 w: &CudaSlice<u8>,
6923 x: &CudaSlice<f32>,
6924 m: usize,
6925 in_f: usize,
6926 out_f: usize,
6927 row_bytes: usize,
6928 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6929 self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6930 }
6931 pub fn qmatvec_nvfp4_fast_rp(
6933 &self,
6934 w: &CudaSlice<u8>,
6935 x: &CudaSlice<f32>,
6936 m: usize,
6937 in_f: usize,
6938 out_f: usize,
6939 row_bytes: usize,
6940 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6941 assert!(
6942 in_f % 64 == 0,
6943 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6944 );
6945 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6946 }
6947 pub fn qmatvec_nvfp4_fast(
6949 &self,
6950 w: &CudaSlice<u8>,
6951 x: &CudaSlice<f32>,
6952 m: usize,
6953 in_f: usize,
6954 out_f: usize,
6955 row_bytes: usize,
6956 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6957 assert!(
6960 in_f % 64 == 0,
6961 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6962 );
6963 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6964 }
6965 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
6968 &self,
6969 w: &CudaSlice<u8>,
6970 x: &CudaSlice<f32>,
6971 m: usize,
6972 in_f: usize,
6973 out_f: usize,
6974 row_bytes: usize,
6975 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6976 self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6977 }
6978
6979 fn qmatvec_dp4a_named(
6981 &self,
6982 name: &str,
6983 w: &CudaSlice<u8>,
6984 x: &CudaSlice<f32>,
6985 m: usize,
6986 in_f: usize,
6987 out_f: usize,
6988 row_bytes: usize,
6989 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6990 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6991 let f = self.func(name);
6992 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6994 grid_dim: (out_f as u32, m as u32, 1),
6995 block_dim: (128, 1, 1),
6996 shared_mem_bytes: 0,
6997 };
6998 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6999 let __s_b = self.gpu.stream();
7000 let mut b = __s_b.launch_builder(&f);
7001 b.arg(w)
7002 .arg(&aq)
7003 .arg(&ad)
7004 .arg(&mut y)
7005 .arg(&inf)
7006 .arg(&outf)
7007 .arg(&mi)
7008 .arg(&rb);
7009 unsafe {
7010 b.launch(cfg)?;
7011 }
7012 Ok(y)
7013 }
7014
7015 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7016 Ok(self.gpu.stream().clone_htod(v)?)
7017 }
7018 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
7019 Ok(self.gpu.stream().clone_htod(v)?)
7020 }
7021 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7023 Ok(self.gpu.stream().clone_htod(v)?)
7024 }
7025 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
7026 Ok(self.gpu.stream().clone_htod(v)?)
7027 }
7028 pub fn dtoh_view(
7030 &self,
7031 d: &cudarc::driver::CudaView<f32>,
7032 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7033 let v = self.gpu.stream().clone_dtoh(d)?;
7034 self.gpu.stream().synchronize()?;
7035 Ok(v)
7036 }
7037 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7038 let v = self.gpu.stream().clone_dtoh(d)?;
7039 self.gpu.stream().synchronize()?;
7040 Ok(v)
7041 }
7042 pub fn dtoh_pair(
7046 &self,
7047 a: &CudaSlice<f32>,
7048 b: &CudaSlice<f32>,
7049 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7050 let av = self.gpu.stream().clone_dtoh(a)?;
7051 let bv = self.gpu.stream().clone_dtoh(b)?;
7052 self.gpu.stream().synchronize()?;
7053 Ok((av, bv))
7054 }
7055 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
7057 let v = self.gpu.stream().clone_dtoh(d)?;
7058 self.gpu.stream().synchronize()?;
7059 Ok(v)
7060 }
7061 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7063 let v = self.gpu.stream().clone_dtoh(d)?;
7064 self.gpu.stream().synchronize()?;
7065 Ok(v)
7066 }
7067 pub fn dtoh_u8_view(
7068 &self,
7069 d: &cudarc::driver::CudaView<u8>,
7070 ) -> Result<Vec<u8>, 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 zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7076 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
7077 self.keep_if_capturing(&s);
7078 Ok(s)
7079 }
7080
7081 pub fn prob_of_token_device(
7090 &self,
7091 logits: &CudaSlice<f32>,
7092 tok: &CudaSlice<u32>,
7093 n_vocab: usize,
7094 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7095 let nb = ARGMAX_NB;
7096 let mut part = self.alloc_uninit::<f32>(nb)?;
7097 let mut p = self.alloc_uninit::<f32>(1)?;
7098 let f1 = self.func("prob_of_token_partial_f32");
7099 let cfg1 = LaunchConfig {
7100 grid_dim: (nb as u32, 1, 1),
7101 block_dim: (256, 1, 1),
7102 shared_mem_bytes: 0,
7103 };
7104 let nv = n_vocab as i32;
7105 let __s_b1 = self.gpu.stream();
7106 let mut b1 = __s_b1.launch_builder(&f1);
7107 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7108 unsafe {
7109 b1.launch(cfg1)?;
7110 }
7111 let f2 = self.func("prob_of_token_final_f32");
7112 let cfg2 = LaunchConfig {
7113 grid_dim: (1, 1, 1),
7114 block_dim: (256, 1, 1),
7115 shared_mem_bytes: 0,
7116 };
7117 let nbi = nb as i32;
7118 let __s_b2 = self.gpu.stream();
7119 let mut b2 = __s_b2.launch_builder(&f2);
7120 b2.arg(&part).arg(&mut p).arg(&nbi);
7121 unsafe {
7122 b2.launch(cfg2)?;
7123 }
7124 Ok(p)
7125 }
7126
7127 pub fn prob_of_token_device_col(
7134 &self,
7135 logits: &CudaSlice<f32>,
7136 tok_all: &CudaSlice<u32>,
7137 tok_idx: usize,
7138 p_out: &mut CudaSlice<f32>,
7139 p_idx: usize,
7140 n_vocab: usize,
7141 ) -> Result<(), Box<dyn std::error::Error>> {
7142 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
7143 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
7144 let nb = ARGMAX_NB;
7145 let mut part = self.alloc_uninit::<f32>(nb)?;
7146 let f1 = self.func("prob_of_token_partial_f32");
7147 let cfg1 = LaunchConfig {
7148 grid_dim: (nb as u32, 1, 1),
7149 block_dim: (256, 1, 1),
7150 shared_mem_bytes: 0,
7151 };
7152 let nv = n_vocab as i32;
7153 let __s_b1 = self.gpu.stream();
7154 let mut b1 = __s_b1.launch_builder(&f1);
7155 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
7156 unsafe {
7157 b1.launch(cfg1)?;
7158 }
7159 let f2 = self.func("prob_of_token_final_f32");
7160 let cfg2 = LaunchConfig {
7161 grid_dim: (1, 1, 1),
7162 block_dim: (256, 1, 1),
7163 shared_mem_bytes: 0,
7164 };
7165 let nbi = nb as i32;
7166 let __s_b2 = self.gpu.stream();
7167 let mut b2 = __s_b2.launch_builder(&f2);
7168 b2.arg(&part).arg(&mut p_v).arg(&nbi);
7169 unsafe {
7170 b2.launch(cfg2)?;
7171 }
7172 Ok(())
7173 }
7174
7175 pub fn prob_of_token_device_into(
7176 &self,
7177 logits: &CudaSlice<f32>,
7178 tok: &CudaSlice<u32>,
7179 p_out: &mut CudaSlice<f32>,
7180 n_vocab: usize,
7181 ) -> Result<(), Box<dyn std::error::Error>> {
7182 let nb = ARGMAX_NB;
7183 let mut part = self.alloc_uninit::<f32>(nb)?;
7184 let f1 = self.func("prob_of_token_partial_f32");
7185 let cfg1 = LaunchConfig {
7186 grid_dim: (nb as u32, 1, 1),
7187 block_dim: (256, 1, 1),
7188 shared_mem_bytes: 0,
7189 };
7190 let nv = n_vocab as i32;
7191 let __s_b1 = self.gpu.stream();
7192 let mut b1 = __s_b1.launch_builder(&f1);
7193 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7194 unsafe {
7195 b1.launch(cfg1)?;
7196 }
7197 let f2 = self.func("prob_of_token_final_f32");
7198 let cfg2 = LaunchConfig {
7199 grid_dim: (1, 1, 1),
7200 block_dim: (256, 1, 1),
7201 shared_mem_bytes: 0,
7202 };
7203 let nbi = nb as i32;
7204 let __s_b2 = self.gpu.stream();
7205 let mut b2 = __s_b2.launch_builder(&f2);
7206 b2.arg(&part).arg(p_out).arg(&nbi);
7207 unsafe {
7208 b2.launch(cfg2)?;
7209 }
7210 Ok(())
7211 }
7212
7213 pub fn argmax_token_device(
7214 &self,
7215 logits: &CudaSlice<f32>,
7216 n_vocab: usize,
7217 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7218 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
7219 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
7220 Ok(tok)
7221 }
7222 pub fn argmax_token_device_into(
7229 &self,
7230 logits: &CudaSlice<f32>,
7231 tok: &mut CudaSlice<u32>,
7232 n_vocab: usize,
7233 ) -> Result<(), Box<dyn std::error::Error>> {
7234 let nb = ARGMAX_NB;
7235 let f1 = self.func("argmax_partial_f32");
7236 let f2 = self.func("argmax_final_f32");
7237 let mut guard = self.argmax_partials.lock().unwrap();
7238 if guard.is_none() {
7239 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7242 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7243 *guard = Some((pv, pi));
7244 }
7245 let (part_v, part_i) = guard.as_mut().unwrap();
7246 let nv = n_vocab as i32;
7247 let nbi = nb as i32;
7248 let cfg1 = LaunchConfig {
7250 grid_dim: (nb as u32, 1, 1),
7251 block_dim: (256, 1, 1),
7252 shared_mem_bytes: 0,
7253 };
7254 let __s_b1 = self.gpu.stream();
7255 let mut b1 = __s_b1.launch_builder(&f1);
7256 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
7257 unsafe {
7258 b1.launch(cfg1)?;
7259 }
7260 let cfg2 = LaunchConfig {
7262 grid_dim: (1, 1, 1),
7263 block_dim: (256, 1, 1),
7264 shared_mem_bytes: 0,
7265 };
7266 let __s_b2 = self.gpu.stream();
7267 let mut b2 = __s_b2.launch_builder(&f2);
7268 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
7269 unsafe {
7270 b2.launch(cfg2)?;
7271 }
7272 Ok(())
7273 }
7274 pub fn argmax_token_device_col(
7280 &self,
7281 logits: &CudaSlice<f32>,
7282 col: usize,
7283 n_vocab: usize,
7284 toks: &mut CudaSlice<u32>,
7285 out_idx: usize,
7286 ) -> Result<(), Box<dyn std::error::Error>> {
7287 let nb = ARGMAX_NB;
7288 let f1 = self.func("argmax_partial_f32");
7289 let f2 = self.func("argmax_final_f32");
7290 let mut guard = self.argmax_partials.lock().unwrap();
7291 if guard.is_none() {
7292 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7293 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7294 *guard = Some((pv, pi));
7295 }
7296 let (part_v, part_i) = guard.as_mut().unwrap();
7297 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
7298 let nv = n_vocab as i32;
7299 let nbi = nb as i32;
7300 let cfg1 = LaunchConfig {
7301 grid_dim: (nb as u32, 1, 1),
7302 block_dim: (256, 1, 1),
7303 shared_mem_bytes: 0,
7304 };
7305 let __s_b1 = self.gpu.stream();
7306 let mut b1 = __s_b1.launch_builder(&f1);
7307 b1.arg(&col_view)
7308 .arg(&mut *part_v)
7309 .arg(&mut *part_i)
7310 .arg(&nv);
7311 unsafe {
7312 b1.launch(cfg1)?;
7313 }
7314 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
7315 let cfg2 = LaunchConfig {
7316 grid_dim: (1, 1, 1),
7317 block_dim: (256, 1, 1),
7318 shared_mem_bytes: 0,
7319 };
7320 let __s_b2 = self.gpu.stream();
7321 let mut b2 = __s_b2.launch_builder(&f2);
7322 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
7323 unsafe {
7324 b2.launch(cfg2)?;
7325 }
7326 Ok(())
7327 }
7328 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7330 Ok(self.gpu.stream().clone_htod(v)?)
7331 }
7332 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7333 let v = self.gpu.stream().clone_dtoh(d)?;
7334 self.gpu.stream().synchronize()?;
7335 Ok(v)
7336 }
7337 pub fn htod_u32_into(
7341 &self,
7342 dst: &mut CudaSlice<u32>,
7343 src: &[u32],
7344 ) -> Result<(), Box<dyn std::error::Error>> {
7345 let mut view = dst.slice_mut(0..src.len());
7346 self.gpu.stream().memcpy_htod(src, &mut view)?;
7347 Ok(())
7348 }
7349
7350 pub fn htod_i32_into(
7353 &self,
7354 dst: &mut CudaSlice<i32>,
7355 src: &[i32],
7356 ) -> Result<(), Box<dyn std::error::Error>> {
7357 let mut view = dst.slice_mut(0..src.len());
7358 self.gpu.stream().memcpy_htod(src, &mut view)?;
7359 Ok(())
7360 }
7361
7362 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7363 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
7364 self.keep_if_capturing(&s);
7365 Ok(s)
7366 }
7367 pub fn embed_gather_device_into(
7370 &self,
7371 embd: &CudaSlice<u8>,
7372 token_d: &CudaSlice<u32>,
7373 x_out: &mut CudaSlice<f32>,
7374 n_embd: usize,
7375 qtype: i32,
7376 row_bytes: usize,
7377 ) -> Result<(), Box<dyn std::error::Error>> {
7378 let f = self.func("embed_gather_u32");
7379 let cfg = LaunchConfig {
7380 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7381 block_dim: (256, 1, 1),
7382 shared_mem_bytes: 0,
7383 };
7384 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7385 let __s_b = self.gpu.stream();
7386 let mut b = __s_b.launch_builder(&f);
7387 b.arg(embd)
7388 .arg(token_d)
7389 .arg(x_out)
7390 .arg(&ne)
7391 .arg(&qt)
7392 .arg(&rb);
7393 unsafe {
7394 b.launch(cfg)?;
7395 }
7396 Ok(())
7397 }
7398 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7400 let v = self.gpu.stream().clone_dtoh(d)?;
7401 self.gpu.stream().synchronize()?;
7402 Ok(v[0])
7403 }
7404 pub fn i32_set_k(
7411 &self,
7412 dst: &mut CudaSlice<i32>,
7413 v: i32,
7414 ) -> Result<(), Box<dyn std::error::Error>> {
7415 let f = self.func("i32_set_k");
7416 let cfg = LaunchConfig {
7417 grid_dim: (1, 1, 1),
7418 block_dim: (1, 1, 1),
7419 shared_mem_bytes: 0,
7420 };
7421 let idx = 0i32;
7422 let __s_b = self.gpu.stream();
7423 let mut b = __s_b.launch_builder(&f);
7424 b.arg(dst).arg(&v).arg(&idx);
7425 unsafe {
7426 b.launch(cfg)?;
7427 }
7428 Ok(())
7429 }
7430
7431 pub fn set_i32_one(
7432 &self,
7433 d: &mut CudaSlice<i32>,
7434 v: i32,
7435 ) -> Result<(), Box<dyn std::error::Error>> {
7436 self.gpu.stream().memcpy_htod(&[v], d)?;
7437 Ok(())
7438 }
7439 pub fn set_u32_one(
7442 &self,
7443 d: &mut CudaSlice<u32>,
7444 v: u32,
7445 ) -> Result<(), Box<dyn std::error::Error>> {
7446 self.gpu.stream().memcpy_htod(&[v], d)?;
7447 Ok(())
7448 }
7449 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7451 let v = self.gpu.stream().clone_dtoh(d)?;
7452 self.gpu.stream().synchronize()?;
7453 Ok(v[0])
7454 }
7455 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7457 Ok(self.gpu.stream().clone_htod(bytes)?)
7458 }
7459 pub fn embed_gather_device(
7463 &self,
7464 embd: &CudaSlice<u8>,
7465 token_d: &CudaSlice<u32>,
7466 n_embd: usize,
7467 qtype: i32,
7468 row_bytes: usize,
7469 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7470 let f = self.func("embed_gather_u32");
7471 let mut x = self.alloc_uninit::<f32>(n_embd)?;
7472 let cfg = LaunchConfig {
7473 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7474 block_dim: (256, 1, 1),
7475 shared_mem_bytes: 0,
7476 };
7477 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7478 let __s_b = self.gpu.stream();
7479 let mut b = __s_b.launch_builder(&f);
7480 b.arg(embd)
7481 .arg(token_d)
7482 .arg(&mut x)
7483 .arg(&ne)
7484 .arg(&qt)
7485 .arg(&rb);
7486 unsafe {
7487 b.launch(cfg)?;
7488 }
7489 Ok(x)
7490 }
7491
7492 pub fn embed_gather_device_t(
7496 &self,
7497 embd: &CudaSlice<u8>,
7498 tokens: &[u32],
7499 n_embd: usize,
7500 qtype: i32,
7501 row_bytes: usize,
7502 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7503 let t = tokens.len();
7504 let tok_d = self.gpu.stream().clone_htod(tokens)?;
7505 let f = self.func("embed_gather_u32_t");
7506 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7507 let cfg = LaunchConfig {
7508 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7509 block_dim: (256, 1, 1),
7510 shared_mem_bytes: 0,
7511 };
7512 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7513 let __s_b = self.gpu.stream();
7514 let mut b = __s_b.launch_builder(&f);
7515 b.arg(embd)
7516 .arg(&tok_d)
7517 .arg(&mut x)
7518 .arg(&ne)
7519 .arg(&qt)
7520 .arg(&rb)
7521 .arg(&ti);
7522 unsafe {
7523 b.launch(cfg)?;
7524 }
7525 Ok(x)
7526 }
7527
7528 pub fn embed_gather_device_tv(
7533 &self,
7534 embd: &CudaSlice<u8>,
7535 tok_v: &cudarc::driver::CudaView<u32>,
7536 t: usize,
7537 n_embd: usize,
7538 qtype: i32,
7539 row_bytes: usize,
7540 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7541 let f = self.func("embed_gather_u32_t");
7542 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7543 let cfg = LaunchConfig {
7544 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7545 block_dim: (256, 1, 1),
7546 shared_mem_bytes: 0,
7547 };
7548 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7549 let __s_b = self.gpu.stream();
7550 let mut b = __s_b.launch_builder(&f);
7551 b.arg(embd)
7552 .arg(tok_v)
7553 .arg(&mut x)
7554 .arg(&ne)
7555 .arg(&qt)
7556 .arg(&rb)
7557 .arg(&ti);
7558 unsafe {
7559 b.launch(cfg)?;
7560 }
7561 Ok(x)
7562 }
7563
7564 pub fn embed_gather_device_td(
7565 &self,
7566 embd: &CudaSlice<u8>,
7567 tok_d: &CudaSlice<u32>,
7568 t: usize,
7569 n_embd: usize,
7570 qtype: i32,
7571 row_bytes: usize,
7572 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7573 let f = self.func("embed_gather_u32_t");
7574 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7575 let cfg = LaunchConfig {
7576 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7577 block_dim: (256, 1, 1),
7578 shared_mem_bytes: 0,
7579 };
7580 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7581 let __s_b = self.gpu.stream();
7582 let mut b = __s_b.launch_builder(&f);
7583 b.arg(embd)
7584 .arg(tok_d)
7585 .arg(&mut x)
7586 .arg(&ne)
7587 .arg(&qt)
7588 .arg(&rb)
7589 .arg(&ti);
7590 unsafe {
7591 b.launch(cfg)?;
7592 }
7593 Ok(x)
7594 }
7595
7596 #[inline]
7602 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7604 if self
7605 .capture_keep_on
7606 .load(std::sync::atomic::Ordering::Relaxed)
7607 {
7608 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7609 }
7610 }
7611
7612 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7613 &self,
7614 n: usize,
7615 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7616 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7617 {
7621 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7622 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7623 use cudarc::driver::DevicePtrMut;
7625 let n_bytes = s.len() * std::mem::size_of::<T>();
7626 let stream = self.gpu.stream();
7627 let (p_, _g) = s.device_ptr_mut(&stream);
7628 unsafe {
7629 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7630 .result()?;
7631 }
7632 }
7633 }
7634 self.keep_if_capturing(&s);
7635 Ok(s)
7636 }
7637
7638 pub fn uninit_q8_pair(
7643 &self,
7644 n: usize,
7645 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7646 Ok((
7647 self.alloc_uninit::<i8>(n)?,
7648 self.alloc_uninit::<f32>(n / 32)?,
7649 ))
7650 }
7651
7652 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7653 self.alloc_uninit::<f32>(n)
7654 }
7655
7656 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7658 self.alloc_uninit::<i8>(n)
7659 }
7660
7661 #[allow(clippy::too_many_arguments)]
7665 pub fn rms_norm3(
7666 &self,
7667 x: &CudaSlice<f32>,
7668 w0: &CudaSlice<f32>,
7669 w1: &CudaSlice<f32>,
7670 w2: &CudaSlice<f32>,
7671 d0: &mut CudaSlice<f32>,
7672 d1: &mut CudaSlice<f32>,
7673 d2: &mut CudaSlice<f32>,
7674 ncols: usize,
7675 nrows: usize,
7676 eps: f32,
7677 ) -> Result<(), Box<dyn std::error::Error>> {
7678 let f = self.func("rms_norm3_f32");
7679 let cfg = LaunchConfig {
7680 grid_dim: (nrows as u32, 1, 1),
7681 block_dim: (rms_block(), 1, 1),
7682 shared_mem_bytes: 0,
7683 };
7684 let (nc, e) = (ncols as i32, eps);
7685 let __s_b = self.gpu.stream();
7686 let mut b = __s_b.launch_builder(&f);
7687 b.arg(x)
7688 .arg(w0)
7689 .arg(w1)
7690 .arg(w2)
7691 .arg(d0)
7692 .arg(d1)
7693 .arg(d2)
7694 .arg(&nc)
7695 .arg(&e);
7696 unsafe {
7697 b.launch(cfg)?;
7698 }
7699 Ok(())
7700 }
7701
7702 #[allow(clippy::too_many_arguments)]
7704 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7707 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7708 *WARP_ON.get_or_init(|| {
7709 std::env::var("MEMRA_QKVNORM_W")
7710 .map(|v| v != "0")
7711 .unwrap_or(true)
7712 }) && ncols % 4 == 0
7713 && rows >= 64
7714 }
7715
7716 #[allow(clippy::too_many_arguments)]
7719 pub fn rms_norm_qkv_w4b(
7720 &self,
7721 q: &CudaSlice<f32>,
7722 k: &CudaSlice<f32>,
7723 v: &CudaSlice<f32>,
7724 wq: &CudaSlice<f32>,
7725 wk: &CudaSlice<f32>,
7726 wv: &CudaSlice<f32>,
7727 dq: &mut CudaSlice<f32>,
7728 dk: &mut CudaSlice<f32>,
7729 dv: &mut CudaSlice<f32>,
7730 dvb: &mut CudaSlice<u8>,
7731 ncols: usize,
7732 rq: usize,
7733 rk: usize,
7734 eps: f32,
7735 vf16: bool,
7736 ) -> Result<(), Box<dyn std::error::Error>> {
7737 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7738 let f = self.func("rms_norm_qkv_w4b_f32");
7739 let rows = (rq + 2 * rk) as u32;
7740 let cfg = LaunchConfig {
7741 grid_dim: (rows.div_ceil(8), 1, 1),
7742 block_dim: (256, 1, 1),
7743 shared_mem_bytes: 0,
7744 };
7745 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7746 let vf = vf16 as i32;
7747 let __s_b = self.gpu.stream();
7748 let mut b = __s_b.launch_builder(&f);
7749 b.arg(q)
7750 .arg(k)
7751 .arg(v)
7752 .arg(wq)
7753 .arg(wk)
7754 .arg(wv)
7755 .arg(dq)
7756 .arg(dk)
7757 .arg(dv)
7758 .arg(&mut *dvb)
7759 .arg(&nc)
7760 .arg(&rqi)
7761 .arg(&rki)
7762 .arg(&rvi)
7763 .arg(&e)
7764 .arg(&vf);
7765 unsafe {
7766 b.launch(cfg)?;
7767 }
7768 Ok(())
7769 }
7770
7771 pub fn rms_norm_qkv(
7772 &self,
7773 q: &CudaSlice<f32>,
7774 k: &CudaSlice<f32>,
7775 v: &CudaSlice<f32>,
7776 wq: &CudaSlice<f32>,
7777 wk: &CudaSlice<f32>,
7778 wv: &CudaSlice<f32>,
7779 dq: &mut CudaSlice<f32>,
7780 dk: &mut CudaSlice<f32>,
7781 dv: &mut CudaSlice<f32>,
7782 ncols: usize,
7783 rq: usize,
7784 rk: usize,
7785 eps: f32,
7786 ) -> Result<(), Box<dyn std::error::Error>> {
7787 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7791 let warp_on = *WARP_ON.get_or_init(|| {
7792 std::env::var("MEMRA_QKVNORM_W")
7793 .map(|v| v != "0")
7794 .unwrap_or(true)
7795 });
7796 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7799 let f = self.func("rms_norm_qkv_w4_f32");
7800 let rows = (rq + 2 * rk) as u32;
7801 let cfg = LaunchConfig {
7802 grid_dim: (rows.div_ceil(8), 1, 1),
7803 block_dim: (256, 1, 1),
7804 shared_mem_bytes: 0,
7805 };
7806 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7807 let __s_b = self.gpu.stream();
7808 let mut b = __s_b.launch_builder(&f);
7809 b.arg(q)
7810 .arg(k)
7811 .arg(v)
7812 .arg(wq)
7813 .arg(wk)
7814 .arg(wv)
7815 .arg(dq)
7816 .arg(dk)
7817 .arg(dv)
7818 .arg(&nc)
7819 .arg(&rqi)
7820 .arg(&rki)
7821 .arg(&rvi)
7822 .arg(&e);
7823 unsafe {
7824 b.launch(cfg)?;
7825 }
7826 return Ok(());
7827 }
7828 let f = self.func("rms_norm_qkv_f32");
7829 let grid = (rq + 2 * rk) as u32;
7830 let cfg = LaunchConfig {
7831 grid_dim: (grid, 1, 1),
7832 block_dim: (rms_block(), 1, 1),
7833 shared_mem_bytes: 0,
7834 };
7835 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7836 let __s_b = self.gpu.stream();
7837 let mut b = __s_b.launch_builder(&f);
7838 b.arg(q)
7839 .arg(k)
7840 .arg(v)
7841 .arg(wq)
7842 .arg(wk)
7843 .arg(wv)
7844 .arg(dq)
7845 .arg(dk)
7846 .arg(dv)
7847 .arg(&nc)
7848 .arg(&rqi)
7849 .arg(&rki)
7850 .arg(&e);
7851 unsafe {
7852 b.launch(cfg)?;
7853 }
7854 Ok(())
7855 }
7856
7857 #[allow(clippy::too_many_arguments)]
7859 pub fn rms_norm2x(
7860 &self,
7861 a: &CudaSlice<f32>,
7862 bb: &CudaSlice<f32>,
7863 wa: &CudaSlice<f32>,
7864 wb: &CudaSlice<f32>,
7865 da: &mut CudaSlice<f32>,
7866 db: &mut CudaSlice<f32>,
7867 ncols: usize,
7868 nrows: usize,
7869 eps: f32,
7870 ) -> Result<(), Box<dyn std::error::Error>> {
7871 let f = self.func("rms_norm2x_f32");
7872 let cfg = LaunchConfig {
7873 grid_dim: (2 * nrows as u32, 1, 1),
7874 block_dim: (rms_block(), 1, 1),
7875 shared_mem_bytes: 0,
7876 };
7877 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7878 let __s_b = self.gpu.stream();
7879 let mut b = __s_b.launch_builder(&f);
7880 b.arg(a)
7881 .arg(bb)
7882 .arg(wa)
7883 .arg(wb)
7884 .arg(da)
7885 .arg(db)
7886 .arg(&nc)
7887 .arg(&nr)
7888 .arg(&e);
7889 unsafe {
7890 b.launch(cfg)?;
7891 }
7892 Ok(())
7893 }
7894
7895 pub fn softcap(
7897 &self,
7898 y: &mut CudaSlice<f32>,
7899 cap: f32,
7900 n: usize,
7901 ) -> Result<(), Box<dyn std::error::Error>> {
7902 let f = self.func("softcap_f32");
7903 let cfg = LaunchConfig::for_num_elems(n as u32);
7904 let ni = n as i32;
7905 let __s_b = self.gpu.stream();
7906 let mut b = __s_b.launch_builder(&f);
7907 b.arg(y).arg(&cap).arg(&ni);
7908 unsafe {
7909 b.launch(cfg)?;
7910 }
7911 Ok(())
7912 }
7913
7914 pub fn mask_ids_rows(
7917 &self,
7918 y: &mut CudaSlice<f32>,
7919 ids: &CudaSlice<i32>,
7920 n_ids: usize,
7921 n_vocab: usize,
7922 t: usize,
7923 ) -> Result<(), Box<dyn std::error::Error>> {
7924 let f = self.func("mask_ids_rows_f32");
7925 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7926 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7927 let __s_b = self.gpu.stream();
7928 let mut b = __s_b.launch_builder(&f);
7929 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7930 unsafe {
7931 b.launch(cfg)?;
7932 }
7933 Ok(())
7934 }
7935
7936 #[allow(clippy::too_many_arguments)]
7938 pub fn add_scale_rms_norm(
7939 &self,
7940 a: &CudaSlice<f32>,
7941 b_in: &CudaSlice<f32>,
7942 c: f32,
7943 w: &CudaSlice<f32>,
7944 res: &mut CudaSlice<f32>,
7945 dst: &mut CudaSlice<f32>,
7946 ncols: usize,
7947 nrows: usize,
7948 eps: f32,
7949 ) -> Result<(), Box<dyn std::error::Error>> {
7950 let f = self.func("add_scale_rms_norm_f32");
7951 let cfg = LaunchConfig {
7952 grid_dim: (nrows as u32, 1, 1),
7953 block_dim: (rms_block(), 1, 1),
7954 shared_mem_bytes: 0,
7955 };
7956 let (nc, e2) = (ncols as i32, eps);
7957 let __s_b = self.gpu.stream();
7958 let mut b = __s_b.launch_builder(&f);
7959 b.arg(a)
7960 .arg(b_in)
7961 .arg(&c)
7962 .arg(w)
7963 .arg(res)
7964 .arg(dst)
7965 .arg(&nc)
7966 .arg(&e2);
7967 unsafe {
7968 b.launch(cfg)?;
7969 }
7970 Ok(())
7971 }
7972
7973 #[allow(clippy::too_many_arguments)]
7976 pub fn add_scale_rms_norm_q8_1(
7977 &self,
7978 a: &CudaSlice<f32>,
7979 b_in: &CudaSlice<f32>,
7980 c: f32,
7981 w: &CudaSlice<f32>,
7982 res: &mut CudaSlice<f32>,
7983 ncols: usize,
7984 nrows: usize,
7985 eps: f32,
7986 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7987 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7988 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7989 let (nc, e2) = (ncols as i32, eps);
7990 if Self::pdl_on() && Self::pdl_wb_on() {
7991 {
7992 use cudarc::driver::{DevicePtr, DevicePtrMut};
7993 let s = &self.gpu.stream();
7994 let (pa, _g0) = a.device_ptr(s);
7995 let (pb, _g1) = b_in.device_ptr(s);
7996 let (pw, _g2) = w.device_ptr(s);
7997 let (pr, _g3) = res.device_ptr_mut(s);
7998 let (pq, _g4) = out_q.device_ptr_mut(s);
7999 let (pd, _g5) = out_d.device_ptr_mut(s);
8000 let mut ps = [
8001 &pa as *const _ as *mut std::ffi::c_void,
8002 &pb as *const _ as *mut _,
8003 &c as *const _ as *mut _,
8004 &pw as *const _ as *mut _,
8005 &pr as *const _ as *mut _,
8006 &pq as *const _ as *mut _,
8007 &pd as *const _ as *mut _,
8008 &nc as *const _ as *mut _,
8009 &e2 as *const _ as *mut _,
8010 ];
8011 unsafe {
8012 self.launch_pdl(
8013 "add_scale_rms_norm_q8_1",
8014 (nrows as u32, 1, 1),
8015 (rms_block(), 1, 1),
8016 &mut ps,
8017 )?;
8018 }
8019 }
8020 return Ok((out_q, out_d));
8021 }
8022 let f = self.func("add_scale_rms_norm_q8_1");
8023 let cfg = LaunchConfig {
8024 grid_dim: (nrows as u32, 1, 1),
8025 block_dim: (rms_block(), 1, 1),
8026 shared_mem_bytes: 0,
8027 };
8028 let __s_b = self.gpu.stream();
8029 let mut b = __s_b.launch_builder(&f);
8030 b.arg(a)
8031 .arg(b_in)
8032 .arg(&c)
8033 .arg(w)
8034 .arg(res)
8035 .arg(&mut out_q)
8036 .arg(&mut out_d)
8037 .arg(&nc)
8038 .arg(&e2);
8039 unsafe {
8040 b.launch(cfg)?;
8041 }
8042 Ok((out_q, out_d))
8043 }
8044
8045 #[allow(clippy::too_many_arguments)]
8047 pub fn add_scale_rms_norm_q8_1_into(
8048 &self,
8049 a: &CudaSlice<f32>,
8050 b_in: &CudaSlice<f32>,
8051 c: f32,
8052 w: &CudaSlice<f32>,
8053 res: &mut CudaSlice<f32>,
8054 ncols: usize,
8055 nrows: usize,
8056 eps: f32,
8057 out_q: &mut CudaSlice<i8>,
8058 out_d: &mut CudaSlice<f32>,
8059 ) -> Result<(), Box<dyn std::error::Error>> {
8060 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8061 let (nc, e2) = (ncols as i32, eps);
8062 if Self::pdl_on() && Self::pdl_wb_on() {
8063 use cudarc::driver::{DevicePtr, DevicePtrMut};
8064 let s = &self.gpu.stream();
8065 let (pa, _g0) = a.device_ptr(s);
8066 let (pb, _g1) = b_in.device_ptr(s);
8067 let (pw, _g2) = w.device_ptr(s);
8068 let (pr, _g3) = res.device_ptr_mut(s);
8069 let (pq, _g4) = out_q.device_ptr_mut(s);
8070 let (pd, _g5) = out_d.device_ptr_mut(s);
8071 let mut ps = [
8072 &pa as *const _ as *mut std::ffi::c_void,
8073 &pb as *const _ as *mut _,
8074 &c as *const _ as *mut _,
8075 &pw as *const _ as *mut _,
8076 &pr as *const _ as *mut _,
8077 &pq as *const _ as *mut _,
8078 &pd as *const _ as *mut _,
8079 &nc as *const _ as *mut _,
8080 &e2 as *const _ as *mut _,
8081 ];
8082 unsafe {
8083 self.launch_pdl(
8084 "add_scale_rms_norm_q8_1",
8085 (nrows as u32, 1, 1),
8086 (rms_block(), 1, 1),
8087 &mut ps,
8088 )?;
8089 }
8090 return Ok(());
8091 }
8092 let f = self.func("add_scale_rms_norm_q8_1");
8093 let cfg = LaunchConfig {
8094 grid_dim: (nrows as u32, 1, 1),
8095 block_dim: (rms_block(), 1, 1),
8096 shared_mem_bytes: 0,
8097 };
8098 let __s_b = self.gpu.stream();
8099 let mut b = __s_b.launch_builder(&f);
8100 b.arg(a)
8101 .arg(b_in)
8102 .arg(&c)
8103 .arg(w)
8104 .arg(res)
8105 .arg(&mut *out_q)
8106 .arg(&mut *out_d)
8107 .arg(&nc)
8108 .arg(&e2);
8109 unsafe {
8110 b.launch(cfg)?;
8111 }
8112 Ok(())
8113 }
8114
8115 #[allow(clippy::too_many_arguments)]
8118 pub fn rms_pre_add_scale_rms_norm_q8_1(
8119 &self,
8120 a: &CudaSlice<f32>,
8121 wa: &CudaSlice<f32>,
8122 b_in: &CudaSlice<f32>,
8123 c: f32,
8124 w: &CudaSlice<f32>,
8125 res: &mut CudaSlice<f32>,
8126 ncols: usize,
8127 nrows: usize,
8128 eps: f32,
8129 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8130 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8131 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8132 let (nc, e2) = (ncols as i32, eps);
8133 if Self::pdl_on() {
8134 {
8135 use cudarc::driver::{DevicePtr, DevicePtrMut};
8136 let s = &self.gpu.stream();
8137 let (pa, _g0) = a.device_ptr(s);
8138 let (pwa, _g1) = wa.device_ptr(s);
8139 let (pb, _g2) = b_in.device_ptr(s);
8140 let (pw, _g3) = w.device_ptr(s);
8141 let (pr, _g4) = res.device_ptr_mut(s);
8142 let (pq, _g5) = out_q.device_ptr_mut(s);
8143 let (pd, _g6) = out_d.device_ptr_mut(s);
8144 let mut ps = [
8145 &pa as *const _ as *mut std::ffi::c_void,
8146 &pwa as *const _ as *mut _,
8147 &pb as *const _ as *mut _,
8148 &c as *const _ as *mut _,
8149 &pw as *const _ as *mut _,
8150 &pr as *const _ as *mut _,
8151 &pq as *const _ as *mut _,
8152 &pd as *const _ as *mut _,
8153 &nc as *const _ as *mut _,
8154 &e2 as *const _ as *mut _,
8155 ];
8156 unsafe {
8157 self.launch_pdl(
8158 "rms_pre_add_scale_rms_norm_q8_1",
8159 (nrows as u32, 1, 1),
8160 (rms_block(), 1, 1),
8161 &mut ps,
8162 )?;
8163 }
8164 }
8165 return Ok((out_q, out_d));
8166 }
8167 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8168 let cfg = LaunchConfig {
8169 grid_dim: (nrows as u32, 1, 1),
8170 block_dim: (rms_block(), 1, 1),
8171 shared_mem_bytes: 0,
8172 };
8173 let __s_b = self.gpu.stream();
8174 let mut b = __s_b.launch_builder(&f);
8175 b.arg(a)
8176 .arg(wa)
8177 .arg(b_in)
8178 .arg(&c)
8179 .arg(w)
8180 .arg(res)
8181 .arg(&mut out_q)
8182 .arg(&mut out_d)
8183 .arg(&nc)
8184 .arg(&e2);
8185 unsafe {
8186 b.launch(cfg)?;
8187 }
8188 Ok((out_q, out_d))
8189 }
8190
8191 pub fn gelu_tanh_mul_q8_1(
8194 &self,
8195 gate: &CudaSlice<f32>,
8196 up: &cudarc::driver::CudaView<f32>,
8197 act: &mut CudaSlice<f32>,
8198 ncols: usize,
8199 nrows: usize,
8200 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8201 debug_assert!(ncols % 128 == 0);
8202 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8203 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8204 let nc = ncols as i32;
8205 if Self::pdl_on() {
8206 {
8207 use cudarc::driver::{DevicePtr, DevicePtrMut};
8208 let s = &self.gpu.stream();
8209 let (pg, _g0) = gate.device_ptr(s);
8210 let (pu, _g1) = up.device_ptr(s);
8211 let (pact, _g2) = act.device_ptr_mut(s);
8212 let (pq, _g3) = out_q.device_ptr_mut(s);
8213 let (pd, _g4) = out_d.device_ptr_mut(s);
8214 let mut ps = [
8215 &pg as *const _ as *mut std::ffi::c_void,
8216 &pu as *const _ as *mut _,
8217 &pact as *const _ as *mut _,
8218 &pq as *const _ as *mut _,
8219 &pd as *const _ as *mut _,
8220 &nc as *const _ as *mut _,
8221 ];
8222 unsafe {
8223 self.launch_pdl(
8224 "gelu_tanh_mul_q8_1",
8225 (nrows as u32, 1, 1),
8226 (rms_block(), 1, 1),
8227 &mut ps,
8228 )?;
8229 }
8230 }
8231 return Ok((out_q, out_d));
8232 }
8233 let f = self.func("gelu_tanh_mul_q8_1");
8234 let cfg = LaunchConfig {
8235 grid_dim: (nrows as u32, 1, 1),
8236 block_dim: (rms_block(), 1, 1),
8237 shared_mem_bytes: 0,
8238 };
8239 let __s_b = self.gpu.stream();
8240 let mut b = __s_b.launch_builder(&f);
8241 b.arg(gate)
8242 .arg(up)
8243 .arg(act)
8244 .arg(&mut out_q)
8245 .arg(&mut out_d)
8246 .arg(&nc);
8247 unsafe {
8248 b.launch(cfg)?;
8249 }
8250 Ok((out_q, out_d))
8251 }
8252
8253 #[allow(clippy::too_many_arguments)]
8255 pub fn gelu_tanh_mul_q8_1_into(
8256 &self,
8257 gate: &CudaSlice<f32>,
8258 up: &cudarc::driver::CudaView<f32>,
8259 act: &mut CudaSlice<f32>,
8260 ncols: usize,
8261 nrows: usize,
8262 out_q: &mut CudaSlice<i8>,
8263 out_d: &mut CudaSlice<f32>,
8264 ) -> Result<(), Box<dyn std::error::Error>> {
8265 debug_assert!(ncols % 128 == 0);
8266 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8267 let nc = ncols as i32;
8268 if Self::pdl_on() {
8269 use cudarc::driver::{DevicePtr, DevicePtrMut};
8270 let s = &self.gpu.stream();
8271 let (pg, _g0) = gate.device_ptr(s);
8272 let (pu, _g1) = up.device_ptr(s);
8273 let (pact, _g2) = act.device_ptr_mut(s);
8274 let (pq, _g3) = out_q.device_ptr_mut(s);
8275 let (pd, _g4) = out_d.device_ptr_mut(s);
8276 let mut ps = [
8277 &pg as *const _ as *mut std::ffi::c_void,
8278 &pu as *const _ as *mut _,
8279 &pact as *const _ as *mut _,
8280 &pq as *const _ as *mut _,
8281 &pd as *const _ as *mut _,
8282 &nc as *const _ as *mut _,
8283 ];
8284 unsafe {
8285 self.launch_pdl(
8286 "gelu_tanh_mul_q8_1",
8287 (nrows as u32, 1, 1),
8288 (rms_block(), 1, 1),
8289 &mut ps,
8290 )?;
8291 }
8292 return Ok(());
8293 }
8294 let f = self.func("gelu_tanh_mul_q8_1");
8295 let cfg = LaunchConfig {
8296 grid_dim: (nrows as u32, 1, 1),
8297 block_dim: (rms_block(), 1, 1),
8298 shared_mem_bytes: 0,
8299 };
8300 let __s_b = self.gpu.stream();
8301 let mut b = __s_b.launch_builder(&f);
8302 b.arg(gate)
8303 .arg(up)
8304 .arg(&mut *act)
8305 .arg(&mut *out_q)
8306 .arg(&mut *out_d)
8307 .arg(&nc);
8308 unsafe {
8309 b.launch(cfg)?;
8310 }
8311 Ok(())
8312 }
8313
8314 #[allow(clippy::too_many_arguments)]
8316 pub fn add_rms_norm3_q8z(
8317 &self,
8318 a: &CudaSlice<f32>,
8319 b_in: &CudaSlice<f32>,
8320 w0: &CudaSlice<f32>,
8321 w1: &CudaSlice<f32>,
8322 w2: &CudaSlice<f32>,
8323 res: &mut CudaSlice<f32>,
8324 out1: &mut CudaSlice<f32>,
8325 ncols: usize,
8326 nrows: usize,
8327 eps: f32,
8328 ) -> Result<
8329 (
8330 (CudaSlice<i8>, CudaSlice<f32>),
8331 (CudaSlice<i8>, CudaSlice<f32>),
8332 ),
8333 Box<dyn std::error::Error>,
8334 > {
8335 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
8336 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8337 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
8338 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8339 let f = self.func("add_rms_norm3_q8z_f32");
8340 let cfg = LaunchConfig {
8341 grid_dim: (nrows as u32, 1, 1),
8342 block_dim: (rms_block(), 1, 1),
8343 shared_mem_bytes: 0,
8344 };
8345 let (nc, e2) = (ncols as i32, eps);
8346 let __s_b = self.gpu.stream();
8347 let mut b = __s_b.launch_builder(&f);
8348 b.arg(a)
8349 .arg(b_in)
8350 .arg(w0)
8351 .arg(w1)
8352 .arg(w2)
8353 .arg(res)
8354 .arg(&mut q0)
8355 .arg(&mut d0)
8356 .arg(out1)
8357 .arg(&mut q2)
8358 .arg(&mut d2)
8359 .arg(&nc)
8360 .arg(&e2);
8361 unsafe {
8362 b.launch(cfg)?;
8363 }
8364 Ok(((q0, d0), (q2, d2)))
8365 }
8366
8367 #[allow(clippy::too_many_arguments)]
8369 pub fn add_rms_norm3(
8370 &self,
8371 a: &CudaSlice<f32>,
8372 b_in: &CudaSlice<f32>,
8373 w0: &CudaSlice<f32>,
8374 w1: &CudaSlice<f32>,
8375 w2: &CudaSlice<f32>,
8376 res: &mut CudaSlice<f32>,
8377 d0: &mut CudaSlice<f32>,
8378 d1: &mut CudaSlice<f32>,
8379 d2: &mut CudaSlice<f32>,
8380 ncols: usize,
8381 nrows: usize,
8382 eps: f32,
8383 ) -> Result<(), Box<dyn std::error::Error>> {
8384 let f = self.func("add_rms_norm3_f32");
8385 let cfg = LaunchConfig {
8386 grid_dim: (nrows as u32, 1, 1),
8387 block_dim: (rms_block(), 1, 1),
8388 shared_mem_bytes: 0,
8389 };
8390 let (nc, e2) = (ncols as i32, eps);
8391 let __s_b = self.gpu.stream();
8392 let mut b = __s_b.launch_builder(&f);
8393 b.arg(a)
8394 .arg(b_in)
8395 .arg(w0)
8396 .arg(w1)
8397 .arg(w2)
8398 .arg(res)
8399 .arg(d0)
8400 .arg(d1)
8401 .arg(d2)
8402 .arg(&nc)
8403 .arg(&e2);
8404 unsafe {
8405 b.launch(cfg)?;
8406 }
8407 Ok(())
8408 }
8409
8410 pub fn add_scale(
8412 &self,
8413 a: &CudaSlice<f32>,
8414 b_in: &CudaSlice<f32>,
8415 c: f32,
8416 dst: &mut CudaSlice<f32>,
8417 n: usize,
8418 ) -> Result<(), Box<dyn std::error::Error>> {
8419 let f = self.func("add_scale_f32");
8420 let cfg = LaunchConfig::for_num_elems(n as u32);
8421 let ni = n as i32;
8422 let __s_b = self.gpu.stream();
8423 let mut b = __s_b.launch_builder(&f);
8424 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8425 unsafe {
8426 b.launch(cfg)?;
8427 }
8428 Ok(())
8429 }
8430
8431 pub fn layer_norm_bias(
8433 &self,
8434 x: &CudaSlice<f32>,
8435 w: &CudaSlice<f32>,
8436 b: &CudaSlice<f32>,
8437 dst: &mut CudaSlice<f32>,
8438 ncols: usize,
8439 nrows: usize,
8440 eps: f32,
8441 ) -> Result<(), Box<dyn std::error::Error>> {
8442 let f = self.func("layer_norm_bias_f32");
8443 let (nc, e) = (ncols as i32, eps);
8444 let cfg = LaunchConfig {
8445 grid_dim: (nrows as u32, 1, 1),
8446 block_dim: (256, 1, 1),
8447 shared_mem_bytes: 0,
8448 };
8449 let __s_b = self.gpu.stream();
8450 let mut lb = __s_b.launch_builder(&f);
8451 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8452 unsafe {
8453 lb.launch(cfg)?;
8454 }
8455 Ok(())
8456 }
8457
8458 pub fn gelu_tanh(
8460 &self,
8461 x: &CudaSlice<f32>,
8462 dst: &mut CudaSlice<f32>,
8463 n: usize,
8464 ) -> Result<(), Box<dyn std::error::Error>> {
8465 let f = self.func("gelu_tanh_f32");
8466 let ni = n as i64;
8467 let cfg = LaunchConfig {
8468 grid_dim: (n.div_ceil(256) as u32, 1, 1),
8469 block_dim: (256, 1, 1),
8470 shared_mem_bytes: 0,
8471 };
8472 let __s_b = self.gpu.stream();
8473 let mut lb = __s_b.launch_builder(&f);
8474 lb.arg(x).arg(&mut *dst).arg(&ni);
8475 unsafe {
8476 lb.launch(cfg)?;
8477 }
8478 Ok(())
8479 }
8480
8481 pub fn row_softmax(
8483 &self,
8484 x: &mut CudaSlice<f32>,
8485 ncols: usize,
8486 nrows: usize,
8487 ) -> Result<(), Box<dyn std::error::Error>> {
8488 let f = self.func("row_softmax_f32");
8489 let nc = ncols as i32;
8490 let cfg = LaunchConfig {
8491 grid_dim: (nrows as u32, 1, 1),
8492 block_dim: (256, 1, 1),
8493 shared_mem_bytes: 0,
8494 };
8495 let __s_b = self.gpu.stream();
8496 let mut lb = __s_b.launch_builder(&f);
8497 lb.arg(&mut *x).arg(&nc);
8498 unsafe {
8499 lb.launch(cfg)?;
8500 }
8501 Ok(())
8502 }
8503
8504 pub fn rms_norm(
8505 &self,
8506 x: &CudaSlice<f32>,
8507 w: &CudaSlice<f32>,
8508 dst: &mut CudaSlice<f32>,
8509 ncols: usize,
8510 nrows: usize,
8511 eps: f32,
8512 ) -> Result<(), Box<dyn std::error::Error>> {
8513 let (nc, e) = (ncols as i32, eps);
8514 let kname = if Self::norm_ilp_on() {
8515 "rms_norm_f32_v2"
8516 } else {
8517 "rms_norm_f32"
8518 };
8519 if Self::pdl_on() && Self::pdl_wb_on() {
8520 use cudarc::driver::{DevicePtr, DevicePtrMut};
8521 let s = &self.gpu.stream();
8522 let (px, _g0) = x.device_ptr(s);
8523 let (pw, _g1) = w.device_ptr(s);
8524 let (pd, _g2) = dst.device_ptr_mut(s);
8525 let mut ps = [
8526 &px as *const _ as *mut std::ffi::c_void,
8527 &pw as *const _ as *mut _,
8528 &pd as *const _ as *mut _,
8529 &nc as *const _ as *mut _,
8530 &e as *const _ as *mut _,
8531 ];
8532 unsafe {
8533 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
8534 }
8535 return Ok(());
8536 }
8537 let f = self.func(kname);
8538 let cfg = LaunchConfig {
8539 grid_dim: (nrows as u32, 1, 1),
8540 block_dim: (rms_block(), 1, 1),
8541 shared_mem_bytes: 0,
8542 };
8543 let __s_b = self.gpu.stream();
8544 let mut b = __s_b.launch_builder(&f);
8545 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8546 unsafe {
8547 b.launch(cfg)?;
8548 }
8549 Ok(())
8550 }
8551
8552 pub fn rms_norm_decode(
8560 &self,
8561 x: &CudaSlice<f32>,
8562 w: &CudaSlice<f32>,
8563 dst: &mut CudaSlice<f32>,
8564 ncols: usize,
8565 nrows: usize,
8566 eps: f32,
8567 ) -> Result<(), Box<dyn std::error::Error>> {
8568 let f = self.func(if Self::norm_ilp_on() {
8569 "rms_norm_f32_v2"
8570 } else {
8571 "rms_norm_f32"
8572 });
8573 let cfg = LaunchConfig {
8574 grid_dim: (nrows as u32, 1, 1),
8575 block_dim: (1024, 1, 1),
8576 shared_mem_bytes: 0,
8577 };
8578 let (nc, e) = (ncols as i32, eps);
8579 let __s_b = self.gpu.stream();
8580 let mut b = __s_b.launch_builder(&f);
8581 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8582 unsafe {
8583 b.launch(cfg)?;
8584 }
8585 Ok(())
8586 }
8587
8588 pub fn rms_norm_q8_1(
8592 &self,
8593 x: &CudaSlice<f32>,
8594 w: &CudaSlice<f32>,
8595 ncols: usize,
8596 nrows: usize,
8597 eps: f32,
8598 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8599 let nblk = ncols / 32;
8600 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8601 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8602 let (nc, e) = (ncols as i32, eps);
8603 if Self::pdl_on() {
8604 {
8605 use cudarc::driver::{DevicePtr, DevicePtrMut};
8606 let s = &self.gpu.stream();
8607 let (px, _g0) = x.device_ptr(s);
8608 let (pw, _g1) = w.device_ptr(s);
8609 let (pq, _g2) = q.device_ptr_mut(s);
8610 let (pd, _g3) = d.device_ptr_mut(s);
8611 let mut ps = [
8612 &px as *const _ as *mut std::ffi::c_void,
8613 &pw as *const _ as *mut _,
8614 &pq as *const _ as *mut _,
8615 &pd as *const _ as *mut _,
8616 &nc as *const _ as *mut _,
8617 &e as *const _ as *mut _,
8618 ];
8619 unsafe {
8620 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8621 }
8622 }
8623 return Ok((q, d));
8624 }
8625 let f = self.func("rms_norm_q8_1");
8626 let cfg = LaunchConfig {
8629 grid_dim: (nrows as u32, 1, 1),
8630 block_dim: (1024, 1, 1),
8631 shared_mem_bytes: 0,
8632 };
8633 let __s_b = self.gpu.stream();
8634 let mut b = __s_b.launch_builder(&f);
8635 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8636 unsafe {
8637 b.launch(cfg)?;
8638 }
8639 Ok((q, d))
8640 }
8641
8642 pub fn rms_norm_q8_1_into(
8645 &self,
8646 x: &CudaSlice<f32>,
8647 w: &CudaSlice<f32>,
8648 ncols: usize,
8649 nrows: usize,
8650 eps: f32,
8651 q: &mut CudaSlice<i8>,
8652 d: &mut CudaSlice<f32>,
8653 ) -> Result<(), Box<dyn std::error::Error>> {
8654 let nblk = ncols / 32;
8655 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8656 let (nc, e) = (ncols as i32, eps);
8657 if Self::pdl_on() {
8658 use cudarc::driver::{DevicePtr, DevicePtrMut};
8659 let s = &self.gpu.stream();
8660 let (px, _g0) = x.device_ptr(s);
8661 let (pw, _g1) = w.device_ptr(s);
8662 let (pq, _g2) = q.device_ptr_mut(s);
8663 let (pd, _g3) = d.device_ptr_mut(s);
8664 let mut ps = [
8665 &px as *const _ as *mut std::ffi::c_void,
8666 &pw as *const _ as *mut _,
8667 &pq as *const _ as *mut _,
8668 &pd as *const _ as *mut _,
8669 &nc as *const _ as *mut _,
8670 &e as *const _ as *mut _,
8671 ];
8672 unsafe {
8673 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8674 }
8675 return Ok(());
8676 }
8677 let f = self.func("rms_norm_q8_1");
8678 let cfg = LaunchConfig {
8679 grid_dim: (nrows as u32, 1, 1),
8680 block_dim: (1024, 1, 1),
8681 shared_mem_bytes: 0,
8682 };
8683 let __s_b = self.gpu.stream();
8684 let mut b = __s_b.launch_builder(&f);
8685 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8686 unsafe {
8687 b.launch(cfg)?;
8688 }
8689 Ok(())
8690 }
8691
8692 pub fn quantize_q8_1_into(
8694 &self,
8695 x: &CudaSlice<f32>,
8696 m: usize,
8697 in_f: usize,
8698 q: &mut CudaSlice<i8>,
8699 d: &mut CudaSlice<f32>,
8700 ) -> Result<(), Box<dyn std::error::Error>> {
8701 let nblk = in_f / 32;
8702 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8703 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8704 let (inf, mi) = (in_f as i32, m as i32);
8705 if Self::pdl_on() && Self::pdl_wb_on() {
8706 use cudarc::driver::{DevicePtr, DevicePtrMut};
8707 let s = &self.gpu.stream();
8708 let (px, _g0) = x.device_ptr(s);
8709 let (pq, _g1) = q.device_ptr_mut(s);
8710 let (pd, _g2) = d.device_ptr_mut(s);
8711 let mut ps = [
8712 &px as *const _ as *mut std::ffi::c_void,
8713 &pq as *const _ as *mut _,
8714 &pd as *const _ as *mut _,
8715 &inf as *const _ as *mut _,
8716 &mi as *const _ as *mut _,
8717 ];
8718 unsafe {
8719 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8720 }
8721 return Ok(());
8722 }
8723 let f = self.func("quantize_q8_1");
8724 let __s_b = self.gpu.stream();
8725 let mut b = __s_b.launch_builder(&f);
8726 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8727 unsafe {
8728 b.launch(cfg)?;
8729 }
8730 Ok(())
8731 }
8732
8733 pub fn add_rms_norm_q8_1(
8737 &self,
8738 a: &CudaSlice<f32>,
8739 b_in: &CudaSlice<f32>,
8740 w: &CudaSlice<f32>,
8741 res: &mut CudaSlice<f32>,
8742 ncols: usize,
8743 nrows: usize,
8744 eps: f32,
8745 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8746 let nblk = ncols / 32;
8747 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8748 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8749 let f = self.func("add_rms_norm_q8_1");
8750 let cfg = LaunchConfig {
8752 grid_dim: (nrows as u32, 1, 1),
8753 block_dim: (1024, 1, 1),
8754 shared_mem_bytes: 0,
8755 };
8756 let (nc, e) = (ncols as i32, eps);
8757 let __s_bld = self.gpu.stream();
8758 let mut bld = __s_bld.launch_builder(&f);
8759 bld.arg(a)
8760 .arg(b_in)
8761 .arg(w)
8762 .arg(res)
8763 .arg(&mut q)
8764 .arg(&mut d)
8765 .arg(&nc)
8766 .arg(&e);
8767 unsafe {
8768 bld.launch(cfg)?;
8769 }
8770 Ok((q, d))
8771 }
8772
8773 pub fn add_rms_norm(
8777 &self,
8778 a: &CudaSlice<f32>,
8779 b: &CudaSlice<f32>,
8780 w: &CudaSlice<f32>,
8781 res: &mut CudaSlice<f32>,
8782 dst: &mut CudaSlice<f32>,
8783 ncols: usize,
8784 nrows: usize,
8785 eps: f32,
8786 ) -> Result<(), Box<dyn std::error::Error>> {
8787 let (nc, e) = (ncols as i32, eps);
8788 let kname = if Self::norm_ilp_on() {
8789 "add_rms_norm_f32_v2"
8790 } else {
8791 "add_rms_norm_f32"
8792 };
8793 if Self::pdl_on() && Self::pdl_wb_on() {
8794 use cudarc::driver::{DevicePtr, DevicePtrMut};
8795 let s = &self.gpu.stream();
8796 let (pa, _g0) = a.device_ptr(s);
8797 let (pb, _g1) = b.device_ptr(s);
8798 let (pw, _g2) = w.device_ptr(s);
8799 let (pr, _g3) = res.device_ptr_mut(s);
8800 let (pd, _g4) = dst.device_ptr_mut(s);
8801 let mut ps = [
8802 &pa as *const _ as *mut std::ffi::c_void,
8803 &pb as *const _ as *mut _,
8804 &pw as *const _ as *mut _,
8805 &pr as *const _ as *mut _,
8806 &pd as *const _ as *mut _,
8807 &nc as *const _ as *mut _,
8808 &e as *const _ as *mut _,
8809 ];
8810 unsafe {
8811 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
8812 }
8813 return Ok(());
8814 }
8815 let f = self.func(kname);
8816 let cfg = LaunchConfig {
8817 grid_dim: (nrows as u32, 1, 1),
8818 block_dim: (rms_block(), 1, 1),
8819 shared_mem_bytes: 0,
8820 };
8821 let __s_b2 = self.gpu.stream();
8822 let mut b2 = __s_b2.launch_builder(&f);
8823 b2.arg(a)
8824 .arg(b)
8825 .arg(w)
8826 .arg(&mut *res)
8827 .arg(&mut *dst)
8828 .arg(&nc)
8829 .arg(&e);
8830 unsafe {
8831 b2.launch(cfg)?;
8832 }
8833 Ok(())
8834 }
8835
8836 #[allow(clippy::too_many_arguments)]
8839 pub fn rms_pre_add_rms_norm(
8840 &self,
8841 a: &CudaSlice<f32>,
8842 wa: &CudaSlice<f32>,
8843 b: &CudaSlice<f32>,
8844 w: &CudaSlice<f32>,
8845 res: &mut CudaSlice<f32>,
8846 dst: &mut CudaSlice<f32>,
8847 ncols: usize,
8848 nrows: usize,
8849 eps: f32,
8850 ) -> Result<(), Box<dyn std::error::Error>> {
8851 let f = self.func("rms_pre_add_rms_norm_f32");
8852 let cfg = LaunchConfig {
8853 grid_dim: (nrows as u32, 1, 1),
8854 block_dim: (rms_block(), 1, 1),
8855 shared_mem_bytes: 0,
8856 };
8857 let (nc, e) = (ncols as i32, eps);
8858 let __s_b2 = self.gpu.stream();
8859 let mut b2 = __s_b2.launch_builder(&f);
8860 b2.arg(a)
8861 .arg(wa)
8862 .arg(b)
8863 .arg(w)
8864 .arg(&mut *res)
8865 .arg(&mut *dst)
8866 .arg(&nc)
8867 .arg(&e);
8868 unsafe {
8869 b2.launch(cfg)?;
8870 }
8871 Ok(())
8872 }
8873
8874 #[allow(clippy::too_many_arguments)]
8876 pub fn rms_pre_add_rms_norm_q8z(
8877 &self,
8878 a: &CudaSlice<f32>,
8879 wa: &CudaSlice<f32>,
8880 b: &CudaSlice<f32>,
8881 w: &CudaSlice<f32>,
8882 res: &mut CudaSlice<f32>,
8883 dst: &mut CudaSlice<f32>,
8884 ncols: usize,
8885 nrows: usize,
8886 eps: f32,
8887 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8888 debug_assert!(ncols % 128 == 0);
8889 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8890 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8891 let (nc, e) = (ncols as i32, eps);
8892 if Self::pdl_on() {
8893 {
8894 use cudarc::driver::{DevicePtr, DevicePtrMut};
8895 let s = &self.gpu.stream();
8896 let (pa, _g0) = a.device_ptr(s);
8897 let (pwa, _g1) = wa.device_ptr(s);
8898 let (pb, _g2) = b.device_ptr(s);
8899 let (pw, _g3) = w.device_ptr(s);
8900 let (pr, _g4) = res.device_ptr_mut(s);
8901 let (pdst, _g5) = dst.device_ptr_mut(s);
8902 let (pq, _g6) = out_q.device_ptr_mut(s);
8903 let (pd, _g7) = out_d.device_ptr_mut(s);
8904 let mut ps = [
8905 &pa as *const _ as *mut std::ffi::c_void,
8906 &pwa as *const _ as *mut _,
8907 &pb as *const _ as *mut _,
8908 &pw as *const _ as *mut _,
8909 &pr as *const _ as *mut _,
8910 &pdst as *const _ as *mut _,
8911 &pq as *const _ as *mut _,
8912 &pd as *const _ as *mut _,
8913 &nc as *const _ as *mut _,
8914 &e as *const _ as *mut _,
8915 ];
8916 unsafe {
8917 self.launch_pdl(
8918 "rms_pre_add_rms_norm_q8z_f32",
8919 (nrows as u32, 1, 1),
8920 (rms_block(), 1, 1),
8921 &mut ps,
8922 )?;
8923 }
8924 }
8925 return Ok((out_q, out_d));
8926 }
8927 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8928 let cfg = LaunchConfig {
8929 grid_dim: (nrows as u32, 1, 1),
8930 block_dim: (rms_block(), 1, 1),
8931 shared_mem_bytes: 0,
8932 };
8933 let __s_b2 = self.gpu.stream();
8934 let mut b2 = __s_b2.launch_builder(&f);
8935 b2.arg(a)
8936 .arg(wa)
8937 .arg(b)
8938 .arg(w)
8939 .arg(&mut *res)
8940 .arg(&mut *dst)
8941 .arg(&mut out_q)
8942 .arg(&mut out_d)
8943 .arg(&nc)
8944 .arg(&e);
8945 unsafe {
8946 b2.launch(cfg)?;
8947 }
8948 Ok((out_q, out_d))
8949 }
8950
8951 #[allow(clippy::too_many_arguments)]
8955 pub fn rms_pre_add_rms_norm_q8z_into(
8956 &self,
8957 a: &CudaSlice<f32>,
8958 wa: &CudaSlice<f32>,
8959 b: &CudaSlice<f32>,
8960 w: &CudaSlice<f32>,
8961 res: &mut CudaSlice<f32>,
8962 dst: &mut CudaSlice<f32>,
8963 ncols: usize,
8964 nrows: usize,
8965 eps: f32,
8966 out_q: &mut CudaSlice<i8>,
8967 out_d: &mut CudaSlice<f32>,
8968 ) -> Result<(), Box<dyn std::error::Error>> {
8969 debug_assert!(ncols % 128 == 0);
8970 let (nc, e) = (ncols as i32, eps);
8971 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8972 let cfg = LaunchConfig {
8973 grid_dim: (nrows as u32, 1, 1),
8974 block_dim: (rms_block(), 1, 1),
8975 shared_mem_bytes: 0,
8976 };
8977 let __s_b = self.gpu.stream();
8978 let mut b2 = __s_b.launch_builder(&f);
8979 b2.arg(a)
8980 .arg(wa)
8981 .arg(b)
8982 .arg(w)
8983 .arg(&mut *res)
8984 .arg(&mut *dst)
8985 .arg(&mut *out_q)
8986 .arg(&mut *out_d)
8987 .arg(&nc)
8988 .arg(&e);
8989 unsafe {
8990 b2.launch(cfg)?;
8991 }
8992 Ok(())
8993 }
8994
8995 #[allow(clippy::too_many_arguments)]
8998 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8999 &self,
9000 a: &CudaSlice<f32>,
9001 wa: &CudaSlice<f32>,
9002 b_in: &CudaSlice<f32>,
9003 c: f32,
9004 w: &CudaSlice<f32>,
9005 res: &mut CudaSlice<f32>,
9006 ncols: usize,
9007 nrows: usize,
9008 eps: f32,
9009 out_q: &mut CudaSlice<i8>,
9010 out_d: &mut CudaSlice<f32>,
9011 ) -> Result<(), Box<dyn std::error::Error>> {
9012 debug_assert!(ncols % 128 == 0);
9013 let (nc, e2) = (ncols as i32, eps);
9014 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9015 let cfg = LaunchConfig {
9016 grid_dim: (nrows as u32, 1, 1),
9017 block_dim: (rms_block(), 1, 1),
9018 shared_mem_bytes: 0,
9019 };
9020 let __s_b = self.gpu.stream();
9021 let mut b2 = __s_b.launch_builder(&f);
9022 b2.arg(a)
9023 .arg(wa)
9024 .arg(b_in)
9025 .arg(&c)
9026 .arg(w)
9027 .arg(&mut *res)
9028 .arg(&mut *out_q)
9029 .arg(&mut *out_d)
9030 .arg(&nc)
9031 .arg(&e2);
9032 unsafe {
9033 b2.launch(cfg)?;
9034 }
9035 Ok(())
9036 }
9037
9038 pub fn g4_pnfold_on() -> bool {
9046 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9047 *ON.get_or_init(|| {
9048 std::env::var("MEMRA_G4_PNFOLD")
9049 .map(|v| v != "0")
9050 .unwrap_or(true)
9051 })
9052 }
9053
9054 pub fn build_q4_out_concat3(
9058 &self,
9059 w0: &crate::model::GpuTensor,
9060 w1: &crate::model::GpuTensor,
9061 w2: &crate::model::GpuTensor,
9062 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
9063 use crate::model::GpuTensor;
9064 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
9065 match w {
9066 GpuTensor::Quant {
9067 qtype,
9068 row_bytes,
9069 rp,
9070 ..
9071 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
9072 _ => None,
9073 }
9074 };
9075 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
9076 else {
9077 return Ok(None);
9078 };
9079 if rb0 != rb1
9080 || rb0 != rb2
9081 || w0.in_features() != w1.in_features()
9082 || w0.in_features() != w2.in_features()
9083 {
9084 return Ok(None);
9085 }
9086 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
9087 match w {
9088 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
9089 _ => unreachable!(),
9090 }
9091 }
9092 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
9093 let total = rb0 * (o0 + o1 + o2);
9094 let mut cat = self.alloc_u8(total)?;
9095 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
9096 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
9097 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
9098 Ok(Some(GpuTensor::Quant {
9099 bytes: cat,
9100 qtype: QT_Q4_0,
9101 row_bytes: rb0,
9102 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
9103 scale: 1.0,
9104 rp: false,
9105 #[cfg(memra_cutlass)]
9106 cutlass: None,
9107 fp8: None,
9108 blk: None,
9109 rp4: None,
9110 f16: None,
9111 }))
9112 }
9113
9114 fn full_width_rope_only(
9132 kernel: &str,
9133 n_rot: usize,
9134 head_dim: usize,
9135 ) -> Result<(), Box<dyn std::error::Error>> {
9136 if n_rot == head_dim {
9137 return Ok(());
9138 }
9139 Err(format!(
9140 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
9141 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
9142 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
9143 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
9144 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
9145 )
9146 .into())
9147 }
9148
9149 #[allow(clippy::too_many_arguments)]
9153 pub fn rms_norm_qkv_rope_cat(
9154 &self,
9155 qkv: &CudaSlice<f32>,
9156 wq: &CudaSlice<f32>,
9157 wk: &CudaSlice<f32>,
9158 wv: &CudaSlice<f32>,
9159 q: &mut CudaSlice<f32>,
9160 k: &mut CudaSlice<f32>,
9161 v: &mut CudaSlice<f32>,
9162 head_dim: usize,
9163 n_rot: usize,
9164 rq: usize,
9165 rk: usize,
9166 pos: &CudaSlice<i32>,
9167 nh_q: usize,
9168 nh_k: usize,
9169 base: f32,
9170 freq_scale: f32,
9171 ff: Option<&CudaSlice<f32>>,
9172 eps: f32,
9173 ) -> Result<(), Box<dyn std::error::Error>> {
9174 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
9175 let rows = rq + rk + rk;
9176 let theta_scale = base.powf(-2.0 / head_dim as f32);
9177 let (nc, rqi, rki, nhq, nhk) = (
9178 head_dim as i32,
9179 rq as i32,
9180 rk as i32,
9181 nh_q as i32,
9182 nh_k as i32,
9183 );
9184 if Self::pdl_on() {
9185 use cudarc::driver::{DevicePtr, DevicePtrMut};
9186 let s = &self.gpu.stream();
9187 let (pqkv, _g0) = qkv.device_ptr(s);
9188 let (pwq, _g1) = wq.device_ptr(s);
9189 let (pwk, _g2) = wk.device_ptr(s);
9190 let (pwv, _g3) = wv.device_ptr(s);
9191 let (pq, _g4) = q.device_ptr_mut(s);
9192 let (pk, _g5) = k.device_ptr_mut(s);
9193 let (pv, _g6) = v.device_ptr_mut(s);
9194 let (ppos, _g7) = pos.device_ptr(s);
9195 let (pff, _g8) = match ff {
9196 Some(t) => {
9197 let (p, g) = t.device_ptr(s);
9198 (p, Some(g))
9199 }
9200 None => (0, None),
9201 };
9202 let mut ps = [
9203 &pqkv as *const _ as *mut std::ffi::c_void,
9204 &pwq as *const _ as *mut _,
9205 &pwk as *const _ as *mut _,
9206 &pwv as *const _ as *mut _,
9207 &pq as *const _ as *mut _,
9208 &pk as *const _ as *mut _,
9209 &pv as *const _ as *mut _,
9210 &nc as *const _ as *mut _,
9211 &rqi as *const _ as *mut _,
9212 &rki as *const _ as *mut _,
9213 &ppos as *const _ as *mut _,
9214 &nhq as *const _ as *mut _,
9215 &nhk as *const _ as *mut _,
9216 &theta_scale as *const _ as *mut _,
9217 &freq_scale as *const _ as *mut _,
9218 &pff as *const _ as *mut _,
9219 &eps as *const _ as *mut _,
9220 ];
9221 unsafe {
9222 self.launch_pdl(
9223 "rms_norm_qkv_rope_cat_f32",
9224 (rows as u32, 1, 1),
9225 (rms_block(), 1, 1),
9226 &mut ps,
9227 )?;
9228 }
9229 return Ok(());
9230 }
9231 let f = self.func("rms_norm_qkv_rope_cat_f32");
9232 let cfg = LaunchConfig {
9233 grid_dim: (rows as u32, 1, 1),
9234 block_dim: (rms_block(), 1, 1),
9235 shared_mem_bytes: 0,
9236 };
9237 let __s_b = self.gpu.stream();
9238 let mut b = __s_b.launch_builder(&f);
9239 match ff {
9240 Some(t) => {
9241 b.arg(qkv)
9242 .arg(wq)
9243 .arg(wk)
9244 .arg(wv)
9245 .arg(&mut *q)
9246 .arg(&mut *k)
9247 .arg(&mut *v)
9248 .arg(&nc)
9249 .arg(&rqi)
9250 .arg(&rki)
9251 .arg(pos)
9252 .arg(&nhq)
9253 .arg(&nhk)
9254 .arg(&theta_scale)
9255 .arg(&freq_scale)
9256 .arg(t)
9257 .arg(&eps);
9258 unsafe {
9259 b.launch(cfg)?;
9260 }
9261 }
9262 None => {
9263 let null: u64 = 0;
9264 b.arg(qkv)
9265 .arg(wq)
9266 .arg(wk)
9267 .arg(wv)
9268 .arg(&mut *q)
9269 .arg(&mut *k)
9270 .arg(&mut *v)
9271 .arg(&nc)
9272 .arg(&rqi)
9273 .arg(&rki)
9274 .arg(pos)
9275 .arg(&nhq)
9276 .arg(&nhk)
9277 .arg(&theta_scale)
9278 .arg(&freq_scale)
9279 .arg(&null)
9280 .arg(&eps);
9281 unsafe {
9282 b.launch(cfg)?;
9283 }
9284 }
9285 }
9286 Ok(())
9287 }
9288
9289 #[allow(clippy::too_many_arguments)]
9293 pub fn rms_norm_qkv_rope(
9294 &self,
9295 q0: &CudaSlice<f32>,
9296 k0: &CudaSlice<f32>,
9297 v0: &CudaSlice<f32>,
9298 wq: &CudaSlice<f32>,
9299 wk: &CudaSlice<f32>,
9300 wv: &CudaSlice<f32>,
9301 q: &mut CudaSlice<f32>,
9302 k: &mut CudaSlice<f32>,
9303 v: &mut CudaSlice<f32>,
9304 head_dim: usize,
9305 n_rot: usize,
9306 rq: usize,
9307 rk: usize,
9308 pos: &CudaSlice<i32>,
9309 nh_q: usize,
9310 nh_k: usize,
9311 base: f32,
9312 freq_scale: f32,
9313 ff: Option<&CudaSlice<f32>>,
9314 eps: f32,
9315 ) -> Result<(), Box<dyn std::error::Error>> {
9316 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
9317 let f = self.func("rms_norm_qkv_rope_f32");
9318 let rows = rq + rk + rk; let cfg = LaunchConfig {
9320 grid_dim: (rows as u32, 1, 1),
9321 block_dim: (rms_block(), 1, 1),
9322 shared_mem_bytes: 0,
9323 };
9324 let theta_scale = base.powf(-2.0 / head_dim as f32);
9325 let (nc, rqi, rki, nhq, nhk) = (
9326 head_dim as i32,
9327 rq as i32,
9328 rk as i32,
9329 nh_q as i32,
9330 nh_k as i32,
9331 );
9332 let __s_b = self.gpu.stream();
9333 let mut b = __s_b.launch_builder(&f);
9334 match ff {
9335 Some(t) => {
9336 b.arg(q0)
9337 .arg(k0)
9338 .arg(v0)
9339 .arg(wq)
9340 .arg(wk)
9341 .arg(wv)
9342 .arg(&mut *q)
9343 .arg(&mut *k)
9344 .arg(&mut *v)
9345 .arg(&nc)
9346 .arg(&rqi)
9347 .arg(&rki)
9348 .arg(pos)
9349 .arg(&nhq)
9350 .arg(&nhk)
9351 .arg(&theta_scale)
9352 .arg(&freq_scale)
9353 .arg(t)
9354 .arg(&eps);
9355 unsafe {
9356 b.launch(cfg)?;
9357 }
9358 }
9359 None => {
9360 let null: u64 = 0;
9361 b.arg(q0)
9362 .arg(k0)
9363 .arg(v0)
9364 .arg(wq)
9365 .arg(wk)
9366 .arg(wv)
9367 .arg(&mut *q)
9368 .arg(&mut *k)
9369 .arg(&mut *v)
9370 .arg(&nc)
9371 .arg(&rqi)
9372 .arg(&rki)
9373 .arg(pos)
9374 .arg(&nhq)
9375 .arg(&nhk)
9376 .arg(&theta_scale)
9377 .arg(&freq_scale)
9378 .arg(&null)
9379 .arg(&eps);
9380 unsafe {
9381 b.launch(cfg)?;
9382 }
9383 }
9384 }
9385 Ok(())
9386 }
9387
9388 #[allow(clippy::too_many_arguments)]
9394 pub fn rms_norm_qkv_rope_append_dc(
9395 &self,
9396 q0: &CudaSlice<f32>,
9397 k0: &CudaSlice<f32>,
9398 v0: &CudaSlice<f32>,
9399 wq: &CudaSlice<f32>,
9400 wk: &CudaSlice<f32>,
9401 wv: &CudaSlice<f32>,
9402 q: &mut CudaSlice<f32>,
9403 k: &mut CudaSlice<f32>,
9404 v: &mut CudaSlice<f32>,
9405 head_dim: usize,
9406 n_rot: usize,
9407 rq: usize,
9408 rk: usize,
9409 pos: &CudaSlice<i32>,
9410 nh_q: usize,
9411 nh_k: usize,
9412 base: f32,
9413 freq_scale: f32,
9414 ff: Option<&CudaSlice<f32>>,
9415 eps: f32,
9416 kc: &mut CudaSlice<u8>,
9417 vc: &mut CudaSlice<u8>,
9418 t_dev: &CudaSlice<i32>,
9419 k_tok_bytes: usize,
9420 v_tok_bytes: usize,
9421 g: bool,
9422 ) -> Result<(), Box<dyn std::error::Error>> {
9423 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
9424 let rows = rq + rk + rk;
9425 let theta_scale = base.powf(-2.0 / head_dim as f32);
9426 let (nc, rqi, rki, nhq, nhk) = (
9427 head_dim as i32,
9428 rq as i32,
9429 rk as i32,
9430 nh_q as i32,
9431 nh_k as i32,
9432 );
9433 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9434 if Self::pdl_on() && Self::pdl_wb_on() {
9435 use cudarc::driver::{DevicePtr, DevicePtrMut};
9436 let s = &self.gpu.stream();
9437 let (p0, _a0) = q0.device_ptr(s);
9438 let (p1, _a1) = k0.device_ptr(s);
9439 let (p2, _a2) = v0.device_ptr(s);
9440 let (pwq, _a3) = wq.device_ptr(s);
9441 let (pwk, _a4) = wk.device_ptr(s);
9442 let (pwv, _a5) = wv.device_ptr(s);
9443 let (pq, _a6) = q.device_ptr_mut(s);
9444 let (pk, _a7) = k.device_ptr_mut(s);
9445 let (pv, _a8) = v.device_ptr_mut(s);
9446 let (pp, _a9) = pos.device_ptr(s);
9447 let pff: u64 = match ff {
9448 Some(t) => {
9449 let (p, _gg) = t.device_ptr(s);
9450 p as u64
9451 }
9452 None => 0,
9453 };
9454 let (pkc, _a10) = kc.device_ptr_mut(s);
9455 let (pvc, _a11) = vc.device_ptr_mut(s);
9456 let (pt, _a12) = t_dev.device_ptr(s);
9457 let mut ps = [
9458 &p0 as *const _ as *mut std::ffi::c_void,
9459 &p1 as *const _ as *mut _,
9460 &p2 as *const _ as *mut _,
9461 &pwq as *const _ as *mut _,
9462 &pwk as *const _ as *mut _,
9463 &pwv as *const _ as *mut _,
9464 &pq as *const _ as *mut _,
9465 &pk as *const _ as *mut _,
9466 &pv as *const _ as *mut _,
9467 &nc as *const _ as *mut _,
9468 &rqi as *const _ as *mut _,
9469 &rki as *const _ as *mut _,
9470 &pp as *const _ as *mut _,
9471 &nhq as *const _ as *mut _,
9472 &nhk as *const _ as *mut _,
9473 &theta_scale as *const _ as *mut _,
9474 &freq_scale as *const _ as *mut _,
9475 &pff as *const _ as *mut _,
9476 &eps as *const _ as *mut _,
9477 &pkc as *const _ as *mut _,
9478 &pvc as *const _ as *mut _,
9479 &pt as *const _ as *mut _,
9480 &ktb as *const _ as *mut _,
9481 &vtb as *const _ as *mut _,
9482 ];
9483 unsafe {
9484 self.launch_pdl_flash(
9485 g,
9486 "rms_norm_qkv_rope_append_dc_f32",
9487 (rows as u32, 1, 1),
9488 (rms_block(), 1, 1),
9489 0,
9490 &mut ps,
9491 )?;
9492 }
9493 return Ok(());
9494 }
9495 let f = if g {
9496 self.func_g("rms_norm_qkv_rope_append_dc_f32")
9497 } else {
9498 self.func("rms_norm_qkv_rope_append_dc_f32")
9499 };
9500 let cfg = LaunchConfig {
9501 grid_dim: (rows as u32, 1, 1),
9502 block_dim: (rms_block(), 1, 1),
9503 shared_mem_bytes: 0,
9504 };
9505 let __s_b = self.gpu.stream();
9506 let mut b = __s_b.launch_builder(&f);
9507 match ff {
9508 Some(t) => {
9509 b.arg(q0)
9510 .arg(k0)
9511 .arg(v0)
9512 .arg(wq)
9513 .arg(wk)
9514 .arg(wv)
9515 .arg(&mut *q)
9516 .arg(&mut *k)
9517 .arg(&mut *v)
9518 .arg(&nc)
9519 .arg(&rqi)
9520 .arg(&rki)
9521 .arg(pos)
9522 .arg(&nhq)
9523 .arg(&nhk)
9524 .arg(&theta_scale)
9525 .arg(&freq_scale)
9526 .arg(t)
9527 .arg(&eps)
9528 .arg(&mut *kc)
9529 .arg(&mut *vc)
9530 .arg(t_dev)
9531 .arg(&ktb)
9532 .arg(&vtb);
9533 unsafe {
9534 b.launch(cfg)?;
9535 }
9536 }
9537 None => {
9538 let null: u64 = 0;
9539 b.arg(q0)
9540 .arg(k0)
9541 .arg(v0)
9542 .arg(wq)
9543 .arg(wk)
9544 .arg(wv)
9545 .arg(&mut *q)
9546 .arg(&mut *k)
9547 .arg(&mut *v)
9548 .arg(&nc)
9549 .arg(&rqi)
9550 .arg(&rki)
9551 .arg(pos)
9552 .arg(&nhq)
9553 .arg(&nhk)
9554 .arg(&theta_scale)
9555 .arg(&freq_scale)
9556 .arg(&null)
9557 .arg(&eps)
9558 .arg(&mut *kc)
9559 .arg(&mut *vc)
9560 .arg(t_dev)
9561 .arg(&ktb)
9562 .arg(&vtb);
9563 unsafe {
9564 b.launch(cfg)?;
9565 }
9566 }
9567 }
9568 Ok(())
9569 }
9570
9571 #[allow(clippy::too_many_arguments)]
9579 pub fn rms_norm_qkv_rope_append(
9580 &self,
9581 q0: &CudaSlice<f32>,
9582 k0: &CudaSlice<f32>,
9583 v0: &CudaSlice<f32>,
9584 wq: &CudaSlice<f32>,
9585 wk: &CudaSlice<f32>,
9586 wv: &CudaSlice<f32>,
9587 q: &mut CudaSlice<f32>,
9588 k: &mut CudaSlice<f32>,
9589 v: &mut CudaSlice<f32>,
9590 head_dim: usize,
9591 n_rot: usize,
9592 rq: usize,
9593 rk: usize,
9594 pos: &CudaSlice<i32>,
9595 nh_q: usize,
9596 nh_k: usize,
9597 base: f32,
9598 freq_scale: f32,
9599 ff: Option<&CudaSlice<f32>>,
9600 eps: f32,
9601 kc: &mut CudaSlice<u8>,
9602 vc: &mut CudaSlice<u8>,
9603 t: usize,
9604 k_tok_bytes: usize,
9605 v_tok_bytes: usize,
9606 g: bool,
9607 ) -> Result<(), Box<dyn std::error::Error>> {
9608 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
9609 let rows = rq + rk + rk;
9610 let theta_scale = base.powf(-2.0 / head_dim as f32);
9611 let (nc, rqi, rki, nhq, nhk) = (
9612 head_dim as i32,
9613 rq as i32,
9614 rk as i32,
9615 nh_q as i32,
9616 nh_k as i32,
9617 );
9618 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9619 let ti = t as i32;
9620 if Self::pdl_on() && Self::pdl_wb_on() {
9621 use cudarc::driver::{DevicePtr, DevicePtrMut};
9622 let s = &self.gpu.stream();
9623 let (p0, _a0) = q0.device_ptr(s);
9624 let (p1, _a1) = k0.device_ptr(s);
9625 let (p2, _a2) = v0.device_ptr(s);
9626 let (pwq, _a3) = wq.device_ptr(s);
9627 let (pwk, _a4) = wk.device_ptr(s);
9628 let (pwv, _a5) = wv.device_ptr(s);
9629 let (pq, _a6) = q.device_ptr_mut(s);
9630 let (pk, _a7) = k.device_ptr_mut(s);
9631 let (pv, _a8) = v.device_ptr_mut(s);
9632 let (pp, _a9) = pos.device_ptr(s);
9633 let pff: u64 = match ff {
9634 Some(t) => {
9635 let (p, _gg) = t.device_ptr(s);
9636 p as u64
9637 }
9638 None => 0,
9639 };
9640 let (pkc, _a10) = kc.device_ptr_mut(s);
9641 let (pvc, _a11) = vc.device_ptr_mut(s);
9642 let mut ps = [
9643 &p0 as *const _ as *mut std::ffi::c_void,
9644 &p1 as *const _ as *mut _,
9645 &p2 as *const _ as *mut _,
9646 &pwq as *const _ as *mut _,
9647 &pwk as *const _ as *mut _,
9648 &pwv as *const _ as *mut _,
9649 &pq as *const _ as *mut _,
9650 &pk as *const _ as *mut _,
9651 &pv as *const _ as *mut _,
9652 &nc as *const _ as *mut _,
9653 &rqi as *const _ as *mut _,
9654 &rki as *const _ as *mut _,
9655 &pp as *const _ as *mut _,
9656 &nhq as *const _ as *mut _,
9657 &nhk as *const _ as *mut _,
9658 &theta_scale as *const _ as *mut _,
9659 &freq_scale as *const _ as *mut _,
9660 &pff as *const _ as *mut _,
9661 &eps as *const _ as *mut _,
9662 &pkc as *const _ as *mut _,
9663 &pvc as *const _ as *mut _,
9664 &ti as *const _ as *mut _,
9665 &ktb as *const _ as *mut _,
9666 &vtb as *const _ as *mut _,
9667 ];
9668 unsafe {
9669 self.launch_pdl_flash(
9670 g,
9671 "rms_norm_qkv_rope_append_f32",
9672 (rows as u32, 1, 1),
9673 (rms_block(), 1, 1),
9674 0,
9675 &mut ps,
9676 )?;
9677 }
9678 return Ok(());
9679 }
9680 let f = if g {
9681 self.func_g("rms_norm_qkv_rope_append_f32")
9682 } else {
9683 self.func("rms_norm_qkv_rope_append_f32")
9684 };
9685 let cfg = LaunchConfig {
9686 grid_dim: (rows as u32, 1, 1),
9687 block_dim: (rms_block(), 1, 1),
9688 shared_mem_bytes: 0,
9689 };
9690 let __s_b = self.gpu.stream();
9691 let mut b = __s_b.launch_builder(&f);
9692 let null: u64 = 0;
9693 b.arg(q0)
9694 .arg(k0)
9695 .arg(v0)
9696 .arg(wq)
9697 .arg(wk)
9698 .arg(wv)
9699 .arg(&mut *q)
9700 .arg(&mut *k)
9701 .arg(&mut *v)
9702 .arg(&nc)
9703 .arg(&rqi)
9704 .arg(&rki)
9705 .arg(pos)
9706 .arg(&nhq)
9707 .arg(&nhk)
9708 .arg(&theta_scale)
9709 .arg(&freq_scale);
9710 match ff {
9711 Some(t) => {
9712 b.arg(t);
9713 }
9714 None => {
9715 b.arg(&null);
9716 }
9717 }
9718 b.arg(&eps)
9719 .arg(&mut *kc)
9720 .arg(&mut *vc)
9721 .arg(&ti)
9722 .arg(&ktb)
9723 .arg(&vtb);
9724 unsafe {
9725 b.launch(cfg)?;
9726 }
9727 Ok(())
9728 }
9729
9730 pub fn add_q8_1(
9731 &self,
9732 a: &CudaSlice<f32>,
9733 b: &CudaSlice<f32>,
9734 res: &mut CudaSlice<f32>,
9735 ncols: usize,
9736 nrows: usize,
9737 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9738 debug_assert!(ncols % 128 == 0);
9739 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9740 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9741 let f = self.func("add_q8_1_f32");
9742 let cfg = LaunchConfig {
9743 grid_dim: (nrows as u32, 1, 1),
9744 block_dim: (rms_block(), 1, 1),
9745 shared_mem_bytes: 0,
9746 };
9747 let nc = ncols as i32;
9748 let __s_b2 = self.gpu.stream();
9749 let mut b2 = __s_b2.launch_builder(&f);
9750 b2.arg(a)
9751 .arg(b)
9752 .arg(&mut *res)
9753 .arg(&mut out_q)
9754 .arg(&mut out_d)
9755 .arg(&nc);
9756 unsafe {
9757 b2.launch(cfg)?;
9758 }
9759 Ok((out_q, out_d))
9760 }
9761
9762 pub fn rms_pre_add_q8_1(
9766 &self,
9767 a: &CudaSlice<f32>,
9768 wa: &CudaSlice<f32>,
9769 b: &CudaSlice<f32>,
9770 res: &mut CudaSlice<f32>,
9771 ncols: usize,
9772 nrows: usize,
9773 eps: f32,
9774 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9775 debug_assert!(ncols % 128 == 0);
9776 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9777 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9778 let f = self.func("rms_pre_add_q8_1_f32");
9779 let cfg = LaunchConfig {
9780 grid_dim: (nrows as u32, 1, 1),
9781 block_dim: (rms_block(), 1, 1),
9782 shared_mem_bytes: 0,
9783 };
9784 let (nc, ep) = (ncols as i32, eps);
9785 let __s_b2 = self.gpu.stream();
9786 let mut b2 = __s_b2.launch_builder(&f);
9787 b2.arg(a)
9788 .arg(wa)
9789 .arg(b)
9790 .arg(&mut *res)
9791 .arg(&mut out_q)
9792 .arg(&mut out_d)
9793 .arg(&nc)
9794 .arg(&ep);
9795 unsafe {
9796 b2.launch(cfg)?;
9797 }
9798 Ok((out_q, out_d))
9799 }
9800
9801 pub fn l2_v2_on(ncols: usize) -> bool {
9805 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9806 }
9807
9808 pub fn l2_norm_pp(
9809 &self,
9810 x: &CudaSlice<f32>,
9811 dst: &mut CudaSlice<f32>,
9812 dst16: Option<&mut CudaSlice<u8>>,
9813 ncols: usize,
9814 nrows: usize,
9815 eps: f32,
9816 ) -> Result<(), Box<dyn std::error::Error>> {
9817 if Self::l2_v2_on(ncols) {
9818 let f = self.func("l2_norm_pp_v2_f32");
9819 let rows_per_block = 8u32; let cfg = LaunchConfig {
9821 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9822 block_dim: (256, 1, 1),
9823 shared_mem_bytes: 0,
9824 };
9825 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9826 let d16: u64 = match dst16 {
9828 Some(d) => self.addr_u8(d),
9829 None => 0,
9830 };
9831 let __s_b = self.gpu.stream();
9832 let mut b = __s_b.launch_builder(&f);
9833 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9834 unsafe {
9835 b.launch(cfg)?;
9836 }
9837 return Ok(());
9838 }
9839 self.l2_norm(x, dst, ncols, nrows, eps)
9840 }
9841
9842 pub fn l2_norm(
9843 &self,
9844 x: &CudaSlice<f32>,
9845 dst: &mut CudaSlice<f32>,
9846 ncols: usize,
9847 nrows: usize,
9848 eps: f32,
9849 ) -> Result<(), Box<dyn std::error::Error>> {
9850 let f = self.func("l2_norm_f32");
9851 let cfg = LaunchConfig {
9852 grid_dim: (nrows as u32, 1, 1),
9853 block_dim: (256, 1, 1),
9854 shared_mem_bytes: 0,
9855 };
9856 let (nc, e) = (ncols as i32, eps);
9857 let __s_b = self.gpu.stream();
9858 let mut b = __s_b.launch_builder(&f);
9859 b.arg(x).arg(dst).arg(&nc).arg(&e);
9860 unsafe {
9861 b.launch(cfg)?;
9862 }
9863 Ok(())
9864 }
9865
9866 pub fn l2_norm_decode(
9872 &self,
9873 x: &CudaSlice<f32>,
9874 dst: &mut CudaSlice<f32>,
9875 ncols: usize,
9876 nrows: usize,
9877 eps: f32,
9878 ) -> Result<(), Box<dyn std::error::Error>> {
9879 let f = self.func("l2_norm_f32");
9880 let cfg = LaunchConfig {
9881 grid_dim: (nrows as u32, 1, 1),
9882 block_dim: (32, 1, 1),
9883 shared_mem_bytes: 0,
9884 };
9885 let (nc, e) = (ncols as i32, eps);
9886 let __s_b = self.gpu.stream();
9887 let mut b = __s_b.launch_builder(&f);
9888 b.arg(x).arg(dst).arg(&nc).arg(&e);
9889 unsafe {
9890 b.launch(cfg)?;
9891 }
9892 Ok(())
9893 }
9894
9895 pub fn rope_neox(
9897 &self,
9898 x: &mut CudaSlice<f32>,
9899 pos: &CudaSlice<i32>,
9900 head_dim: usize,
9901 n_dims: usize,
9902 n_heads: usize,
9903 n_tokens: usize,
9904 freq_base: f32,
9905 freq_scale: f32,
9906 ) -> Result<(), Box<dyn std::error::Error>> {
9907 let f = self.func("rope_neox_f32");
9908 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9909 let grid = (n_heads * n_tokens) as u32;
9910 let cfg = LaunchConfig {
9911 grid_dim: (grid, 1, 1),
9912 block_dim: ((head_dim / 2) as u32, 1, 1),
9913 shared_mem_bytes: 0,
9914 };
9915 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9916 let __s_b = self.gpu.stream();
9917 let mut b = __s_b.launch_builder(&f);
9918 b.arg(x)
9919 .arg(pos)
9920 .arg(&hd)
9921 .arg(&nd)
9922 .arg(&nh)
9923 .arg(&theta_scale)
9924 .arg(&freq_scale);
9925 unsafe {
9926 b.launch(cfg)?;
9927 }
9928 Ok(())
9929 }
9930
9931 pub fn rope_neox_ff(
9933 &self,
9934 x: &mut CudaSlice<f32>,
9935 pos: &CudaSlice<i32>,
9936 head_dim: usize,
9937 n_dims: usize,
9938 n_heads: usize,
9939 n_tokens: usize,
9940 freq_base: f32,
9941 freq_scale: f32,
9942 ff: &CudaSlice<f32>,
9943 ) -> Result<(), Box<dyn std::error::Error>> {
9944 let f = self.func("rope_neox_ff_f32");
9945 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9946 let grid = (n_heads * n_tokens) as u32;
9947 let cfg = LaunchConfig {
9948 grid_dim: (grid, 1, 1),
9949 block_dim: ((head_dim / 2) as u32, 1, 1),
9950 shared_mem_bytes: 0,
9951 };
9952 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9953 let __s_b = self.gpu.stream();
9954 let mut b = __s_b.launch_builder(&f);
9955 b.arg(x)
9956 .arg(pos)
9957 .arg(&hd)
9958 .arg(&nd)
9959 .arg(&nh)
9960 .arg(&theta_scale)
9961 .arg(&freq_scale)
9962 .arg(ff);
9963 unsafe {
9964 b.launch(cfg)?;
9965 }
9966 Ok(())
9967 }
9968
9969 #[allow(clippy::too_many_arguments)]
9971 pub fn rope_neox2(
9972 &self,
9973 q: &mut CudaSlice<f32>,
9974 k: &mut CudaSlice<f32>,
9975 pos: &CudaSlice<i32>,
9976 head_dim: usize,
9977 n_dims: usize,
9978 nh_q: usize,
9979 nh_k: usize,
9980 n_tokens: usize,
9981 freq_base: f32,
9982 freq_scale: f32,
9983 ff: Option<&CudaSlice<f32>>,
9984 ) -> Result<(), Box<dyn std::error::Error>> {
9985 let f = self.func("rope_neox2_f32");
9986 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9987 let grid = ((nh_q + nh_k) * n_tokens) as u32;
9988 let cfg = LaunchConfig {
9989 grid_dim: (grid, 1, 1),
9990 block_dim: ((head_dim / 2) as u32, 1, 1),
9991 shared_mem_bytes: 0,
9992 };
9993 let (hd, nd, nq, nk, nt) = (
9994 head_dim as i32,
9995 n_dims as i32,
9996 nh_q as i32,
9997 nh_k as i32,
9998 n_tokens as i32,
9999 );
10000 let __s_b = self.gpu.stream();
10001 let mut b = __s_b.launch_builder(&f);
10002 b.arg(q)
10003 .arg(k)
10004 .arg(pos)
10005 .arg(&hd)
10006 .arg(&nd)
10007 .arg(&nq)
10008 .arg(&nk)
10009 .arg(&nt)
10010 .arg(&theta_scale)
10011 .arg(&freq_scale);
10012 match ff {
10013 Some(ffv) => {
10014 b.arg(ffv);
10015 unsafe {
10016 b.launch(cfg)?;
10017 }
10018 }
10019 None => {
10020 let null: u64 = 0;
10021 b.arg(&null);
10022 unsafe {
10023 b.launch(cfg)?;
10024 }
10025 }
10026 }
10027 Ok(())
10028 }
10029
10030 pub fn gelu_tanh_mul(
10032 &self,
10033 gate: &CudaSlice<f32>,
10034 up: &CudaSlice<f32>,
10035 dst: &mut CudaSlice<f32>,
10036 n: usize,
10037 ) -> Result<(), Box<dyn std::error::Error>> {
10038 let f = self.func("gelu_tanh_mul_f32");
10039 let cfg = LaunchConfig::for_num_elems(n as u32);
10040 let ni = n as i32;
10041 let __s_b = self.gpu.stream();
10042 let mut b = __s_b.launch_builder(&f);
10043 b.arg(gate).arg(up).arg(dst).arg(&ni);
10044 unsafe {
10045 b.launch(cfg)?;
10046 }
10047 Ok(())
10048 }
10049
10050 pub fn silu_mul(
10051 &self,
10052 gate: &CudaSlice<f32>,
10053 up: &CudaSlice<f32>,
10054 dst: &mut CudaSlice<f32>,
10055 n: usize,
10056 ) -> Result<(), Box<dyn std::error::Error>> {
10057 let f = self.func("silu_mul_f32");
10058 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10060 let ni = n as i32;
10061 let __s_b = self.gpu.stream();
10062 let mut b = __s_b.launch_builder(&f);
10063 b.arg(gate).arg(up).arg(dst).arg(&ni);
10064 unsafe {
10065 b.launch(cfg)?;
10066 }
10067 Ok(())
10068 }
10069
10070 pub fn silu_mul_f16out(
10073 &self,
10074 gate: &CudaSlice<f32>,
10075 up: &CudaSlice<f32>,
10076 dst: &mut CudaSlice<f32>,
10077 dst16: &mut CudaSlice<u8>,
10078 n: usize,
10079 ) -> Result<(), Box<dyn std::error::Error>> {
10080 let f = self.func("silu_mul_f16out_f32");
10081 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10082 let ni = n as i32;
10083 let __s_b = self.gpu.stream();
10084 let mut b = __s_b.launch_builder(&f);
10085 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
10086 unsafe {
10087 b.launch(cfg)?;
10088 }
10089 Ok(())
10090 }
10091
10092 pub fn silu_mul_scaled(
10099 &self,
10100 gate: &CudaSlice<f32>,
10101 up: &CudaSlice<f32>,
10102 gs: f32,
10103 us: f32,
10104 dst: &mut CudaSlice<f32>,
10105 n: usize,
10106 ) -> Result<(), Box<dyn std::error::Error>> {
10107 let f = self.func("silu_mul_scaled_f32");
10108 let cfg = LaunchConfig::for_num_elems(n as u32);
10109 let ni = n as i32;
10110 let (gsf, usf) = (gs, us);
10111 let __s_b = self.gpu.stream();
10112 let mut b = __s_b.launch_builder(&f);
10113 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
10114 unsafe {
10115 b.launch(cfg)?;
10116 }
10117 Ok(())
10118 }
10119
10120 #[allow(clippy::too_many_arguments)]
10124 pub fn swigluoai_mul_scaled(
10125 &self,
10126 gate: &CudaSlice<f32>,
10127 up: &CudaSlice<f32>,
10128 gs: f32,
10129 us: f32,
10130 alpha: f32,
10131 limit: f32,
10132 dst: &mut CudaSlice<f32>,
10133 n: usize,
10134 ) -> Result<(), Box<dyn std::error::Error>> {
10135 let f = self.func("swigluoai_mul_scaled_f32");
10136 let cfg = LaunchConfig::for_num_elems(n as u32);
10137 let ni = n as i32;
10138 let __s_b = self.gpu.stream();
10139 let mut b = __s_b.launch_builder(&f);
10140 b.arg(gate)
10141 .arg(up)
10142 .arg(&gs)
10143 .arg(&us)
10144 .arg(&alpha)
10145 .arg(&limit)
10146 .arg(dst)
10147 .arg(&ni);
10148 unsafe {
10149 b.launch(cfg)?;
10150 }
10151 Ok(())
10152 }
10153
10154 pub fn silu_mul_scaled_q8_1(
10162 &self,
10163 gate: &CudaSlice<f32>,
10164 up: &CudaSlice<f32>,
10165 gs: f32,
10166 us: f32,
10167 n: usize,
10168 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10169 let f = self.func("silu_mul_scaled_q8_1");
10170 let nblk = n / 32;
10171 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);
10175 let (gsf, usf, ni) = (gs, us, n as i32);
10176 let __s_b = self.gpu.stream();
10177 let mut b = __s_b.launch_builder(&f);
10178 b.arg(gate)
10179 .arg(up)
10180 .arg(&gsf)
10181 .arg(&usf)
10182 .arg(&mut aq)
10183 .arg(&mut ad)
10184 .arg(&ni);
10185 unsafe {
10186 b.launch(cfg)?;
10187 }
10188 Ok((aq, ad))
10189 }
10190
10191 pub fn add(
10192 &self,
10193 a: &CudaSlice<f32>,
10194 b_in: &CudaSlice<f32>,
10195 dst: &mut CudaSlice<f32>,
10196 n: usize,
10197 ) -> Result<(), Box<dyn std::error::Error>> {
10198 let f = self.func("add_f32");
10199 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10201 let ni = n as i32;
10202 let __s_bld = self.gpu.stream();
10203 let mut bld = __s_bld.launch_builder(&f);
10204 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
10205 unsafe {
10206 bld.launch(cfg)?;
10207 }
10208 Ok(())
10209 }
10210
10211 pub fn mul(
10212 &self,
10213 a: &CudaSlice<f32>,
10214 b_in: &CudaSlice<f32>,
10215 dst: &mut CudaSlice<f32>,
10216 n: usize,
10217 ) -> Result<(), Box<dyn std::error::Error>> {
10218 let f = self.func("mul_f32");
10219 let cfg = LaunchConfig::for_num_elems(n as u32);
10220 let ni = n as i32;
10221 let __s_bld = self.gpu.stream();
10222 let mut bld = __s_bld.launch_builder(&f);
10223 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
10224 unsafe {
10225 bld.launch(cfg)?;
10226 }
10227 Ok(())
10228 }
10229
10230 pub fn matmul(
10233 &self,
10234 w: &crate::model::GpuTensor,
10235 x: &CudaSlice<f32>,
10236 m: usize,
10237 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10238 use crate::model::GpuTensor;
10239 let in_f = w.in_features();
10240 let out_f = w.out_features();
10241 #[allow(non_snake_case)]
10249 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
10252 usize::MAX
10253 } else {
10254 16usize
10255 };
10256
10257 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
10282 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
10283 return Ok(y);
10284 }
10285 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
10292 return Ok(y);
10293 }
10294 if let Some(y) = self.try_f16_gemm(w, x, m)? {
10297 return Ok(y);
10298 }
10299 }
10300 if let GpuTensor::Quant { qtype, .. } = w {
10315 if *qtype == QT_F8_E4M3_BLK {
10316 if m >= GEMM_M_THRESHOLD {
10317 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
10318 return Ok(y);
10319 }
10320 }
10321 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10322 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10323 return Ok(y);
10324 }
10325 }
10326 }
10327 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
10328 return self.qmatvec_mmq(w, x, m);
10329 }
10330 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
10331 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10332 return self.qmatvec_gemm(w, &aq, &ad, m);
10333 }
10334 if m >= GEMM_M_THRESHOLD {
10337 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
10338 return Ok(y);
10339 }
10340 }
10341 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
10345 if m == 1 && fast {
10350 if let GpuTensor::Quant {
10351 bytes,
10352 qtype,
10353 row_bytes,
10354 rp,
10355 rp4,
10356 scale,
10357 ..
10358 } = w
10359 {
10360 if self.mmvq_supports(*qtype) {
10361 let (bytes, rp) = match rp4 {
10365 Some(m4) => (m4, true),
10366 None => (bytes, *rp),
10367 };
10368 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10369 return self.qmatvec_mmvq(
10370 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
10371 );
10372 }
10373 }
10374 }
10375 if (2..=16).contains(&m)
10391 && fast
10392 && std::env::var("MEMRA_NO_BATCHED").is_err()
10393 && (m <= 4 || Self::b8_enabled())
10394 {
10395 let m_ok = m <= 8
10405 || matches!(w, GpuTensor::Quant { qtype, .. }
10406 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
10407 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
10408 if m_ok {
10409 if let GpuTensor::Quant {
10410 bytes,
10411 qtype,
10412 row_bytes,
10413 rp,
10414 rp4,
10415 ..
10416 } = w
10417 {
10418 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
10419 let (bytes, rp) = match rp4 {
10420 Some(m4) => (m4, true),
10421 None => (bytes, *rp),
10422 };
10423 let mcols = Self::batched_mcols(m);
10424 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10425 let mut y = self.qmatvec_mmvq_batched(
10426 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
10427 )?;
10428 if let GpuTensor::Quant { scale, .. } = w {
10429 if *scale != 1.0 {
10430 self.scale_inplace(&mut y, *scale, m * out_f)?;
10431 }
10432 }
10433 return Ok(y);
10434 }
10435 }
10436 }
10437 }
10438 if fast {
10444 if let GpuTensor::Quant {
10445 bytes,
10446 qtype,
10447 row_bytes,
10448 scale,
10449 ..
10450 } = w
10451 {
10452 if *qtype == QT_F8_E4M3 {
10453 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10454 return self.qmatvec_mmvq(
10455 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10456 );
10457 }
10458 }
10459 }
10460 let mut y = match w {
10461 GpuTensor::Quant {
10462 bytes,
10463 qtype,
10464 row_bytes,
10465 ..
10466 } if fast && *qtype == QT_Q8_0 => {
10467 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10468 }
10469 GpuTensor::Quant {
10470 bytes,
10471 qtype,
10472 row_bytes,
10473 ..
10474 } if fast && *qtype == QT_Q4_K => {
10475 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10476 }
10477 GpuTensor::Quant {
10478 bytes,
10479 qtype,
10480 row_bytes,
10481 ..
10482 } if fast && *qtype == QT_Q6_K => {
10483 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10484 }
10485 GpuTensor::Quant {
10486 bytes,
10487 qtype,
10488 row_bytes,
10489 ..
10490 } if fast && *qtype == QT_Q5_K => {
10491 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10492 }
10493 GpuTensor::Quant {
10494 bytes,
10495 qtype,
10496 row_bytes,
10497 ..
10498 } if fast && *qtype == QT_Q3_K => {
10499 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10500 }
10501 GpuTensor::Quant {
10502 bytes,
10503 qtype,
10504 row_bytes,
10505 rp,
10506 ..
10507 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10508 if *rp {
10509 "qmatvec_nvfp4_dp4a_rp"
10510 } else {
10511 "qmatvec_nvfp4_dp4a"
10512 },
10513 bytes,
10514 x,
10515 m,
10516 in_f,
10517 out_f,
10518 *row_bytes,
10519 )?,
10520 GpuTensor::Quant {
10524 bytes,
10525 qtype,
10526 row_bytes,
10527 ..
10528 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10529 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10530 }
10531 GpuTensor::Quant {
10536 bytes,
10537 qtype,
10538 row_bytes,
10539 rp,
10540 ..
10541 } =>
10542 {
10545 self.qmatvec(
10546 bytes,
10547 x,
10548 m,
10549 in_f,
10550 out_f,
10551 if *rp && *qtype == QT_NVFP4 {
10552 QT_NVFP4_RP
10553 } else {
10554 *qtype
10555 },
10556 *row_bytes,
10557 )?
10558 }
10559 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10560 GpuTensor::FloatBf16 { data, .. } => {
10563 self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10564 }
10565 };
10566 if let GpuTensor::Quant { scale, .. } = w {
10568 if *scale != 1.0 {
10569 self.scale_inplace(&mut y, *scale, m * out_f)?;
10570 }
10571 }
10572 Ok(y)
10573 }
10574
10575 pub fn stage_a_raw_needed() -> bool {
10585 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10586 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
10587 }
10588
10589 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10592 use crate::model::GpuTensor;
10593 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10594 return false;
10595 }
10596 match w {
10597 GpuTensor::Quant { qtype, .. } => {
10604 matches!(
10605 *qtype,
10606 QT_Q8_0
10607 | QT_Q4_K
10608 | QT_Q6_K
10609 | QT_Q5_K
10610 | QT_Q3_K
10611 | QT_NVFP4
10612 | QT_F8_E4M3
10613 | QT_F8_E4M3_BLK
10614 | QT_Q4_0
10615 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10616 }
10617 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10618 }
10619 }
10620
10621 pub fn matmul_pre(
10626 &self,
10627 w: &crate::model::GpuTensor,
10628 aq: &CudaSlice<i8>,
10629 ad: &CudaSlice<f32>,
10630 x_fallback: &CudaSlice<f32>,
10631 m: usize,
10632 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10633 use crate::model::GpuTensor;
10634 let x_raw_ok = x_fallback.len() >= m * w.in_features();
10640 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10643 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10644 return Ok(y);
10645 }
10646 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10649 return Ok(y);
10650 }
10651 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10653 return Ok(y);
10654 }
10655 }
10656 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10662 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10663 return Ok(y);
10664 }
10665 }
10666 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10667 return Ok(y);
10668 }
10669 if m >= 16
10674 && w.out_features() >= 128
10675 && self.mmq_supports(w)
10676 && !self.verify_exact_on()
10677 && x_raw_ok
10678 {
10679 return self.qmatvec_mmq(w, x_fallback, m);
10680 }
10681 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10684 if let Some(y) =
10685 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10686 {
10687 return Ok(y);
10688 }
10689 }
10690 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10693 return self.qmatvec_gemm(w, aq, ad, m);
10694 }
10695 if !self.uses_q8_1_fast(w) {
10714 if !x_raw_ok {
10715 return Err(format!(
10716 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
10717 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
10718 dropped the f32, so there is nothing to fall back to — pass the real f32 \
10719 activation (see Engine::rms_norm_decode, which is bit-identical to \
10720 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
10721 x_fallback.len(),
10722 m,
10723 w.in_features(),
10724 m * w.in_features()
10725 )
10726 .into());
10727 }
10728 return self.matmul(w, x_fallback, m);
10729 }
10730 let in_f = w.in_features();
10731 let out_f = w.out_features();
10732 let (bytes, qtype, row_bytes, scale, rp) = match w {
10733 GpuTensor::Quant {
10734 bytes,
10735 qtype,
10736 row_bytes,
10737 scale,
10738 rp,
10739 ..
10740 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10741 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10742 };
10743 let (mbytes, mrp) = match w {
10746 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10747 _ => (bytes, rp),
10748 };
10749 if m == 1 && self.mmvq_supports(qtype) {
10753 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10754 }
10755 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10768 && std::env::var("MEMRA_NO_BATCHED").is_err()
10769 && (m <= 4 || Self::b8_enabled())
10770 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10774 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10775 {
10776 let mcols = Self::batched_mcols(m);
10777 return self.qmatvec_mmvq_batched(
10778 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10779 );
10780 }
10781 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10787 let (b2, r2) = if qtype == QT_Q4_0 {
10788 (mbytes, mrp)
10789 } else {
10790 (bytes, rp)
10791 };
10792 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10793 }
10794 let name = match qtype {
10795 QT_Q8_0 => "qmatvec_q8_0_dp4a",
10796 QT_Q4_K => "qmatvec_q4_K_dp4a",
10797 QT_Q6_K => "qmatvec_q6_K_dp4a",
10798 QT_Q5_K => "qmatvec_q5_K_dp4a",
10799 QT_Q3_K => "qmatvec_q3_K_dp4a",
10800 QT_NVFP4 => {
10801 if rp {
10802 "qmatvec_nvfp4_dp4a_rp"
10803 } else {
10804 "qmatvec_nvfp4_dp4a"
10805 }
10806 }
10807 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10808 _ => unreachable!(),
10809 };
10810 let f = self.func(name);
10811 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10813 grid_dim: (out_f as u32, m as u32, 1),
10814 block_dim: (128, 1, 1),
10815 shared_mem_bytes: 0,
10816 };
10817 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10818 let __s_b = self.gpu.stream();
10819 let mut b = __s_b.launch_builder(&f);
10820 b.arg(bytes)
10821 .arg(aq)
10822 .arg(ad)
10823 .arg(&mut y)
10824 .arg(&inf)
10825 .arg(&outf)
10826 .arg(&mi)
10827 .arg(&rb);
10828 unsafe {
10829 b.launch(cfg)?;
10830 }
10831 if scale != 1.0 {
10832 self.scale_inplace(&mut y, scale, m * out_f)?;
10833 }
10834 Ok(y)
10835 }
10836
10837 pub fn matmul_decode_exact(
10845 &self,
10846 w: &crate::model::GpuTensor,
10847 x: &CudaSlice<f32>,
10848 m: usize,
10849 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10850 use crate::model::GpuTensor;
10851 if let GpuTensor::Float { data, .. } = w {
10859 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10860 }
10861 if let GpuTensor::FloatBf16 { data, .. } = w {
10864 let (in_f, out_f) = (w.in_features(), w.out_features());
10865 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10866 }
10867 if !self.uses_q8_1_fast(w) {
10868 return self.matmul(w, x, m);
10869 }
10870 let in_f = w.in_features();
10871 let out_f = w.out_features();
10872 let (bytes, qtype, row_bytes, scale, rp) = match w {
10873 GpuTensor::Quant {
10874 bytes,
10875 qtype,
10876 row_bytes,
10877 scale,
10878 rp,
10879 ..
10880 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10881 _ => return self.matmul(w, x, m),
10882 };
10883 let (bytes, rp) = match w {
10886 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10887 _ => (bytes, rp),
10888 };
10889 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10890 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10894 return Ok(y);
10895 }
10896 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10905 && std::env::var("MEMRA_NO_BATCHED").is_err()
10906 && (m <= 4 || Self::b8_enabled())
10907 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10910 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10911 {
10912 let mcols = Self::batched_mcols(m);
10913 return self.qmatvec_mmvq_batched(
10914 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10915 );
10916 }
10917 if self.mmvq_supports(qtype) {
10918 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10921 }
10922 self.matmul_pre(w, &aq, &ad, x, m)
10925 }
10926
10927 pub fn matmul_decode_exact_pre(
10937 &self,
10938 w: &crate::model::GpuTensor,
10939 aq: &CudaSlice<i8>,
10940 ad: &CudaSlice<f32>,
10941 m: usize,
10942 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10943 use crate::model::GpuTensor;
10944 debug_assert!(
10945 self.uses_q8_1_fast(w),
10946 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10947 );
10948 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10950 return Ok(y);
10951 }
10952 let in_f = w.in_features();
10953 let out_f = w.out_features();
10954 let (bytes, qtype, row_bytes, scale, rp) = match w {
10955 GpuTensor::Quant {
10956 bytes,
10957 qtype,
10958 row_bytes,
10959 scale,
10960 rp,
10961 ..
10962 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10963 _ => {
10964 return Err(
10965 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10966 );
10967 }
10968 };
10969 let (bytes, rp) = match w {
10971 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10972 _ => (bytes, rp),
10973 };
10974 if (2..=16).contains(&m)
10976 && self.batched_supports(qtype)
10977 && self.mmvq_supports(qtype)
10978 && std::env::var("MEMRA_NO_BATCHED").is_err()
10979 && (m <= 4 || Self::b8_enabled())
10980 && (m <= 8
10981 || qtype == QT_Q4_0
10982 || qtype == QT_Q6_K
10983 || qtype == QT_F8_E4M3
10984 || qtype == QT_NVFP4
10985 || qtype == QT_Q4_K
10986 || qtype == QT_Q5_K
10987 || qtype == QT_Q8_0)
10988 {
10989 let mcols = Self::batched_mcols(m);
10990 return self.qmatvec_mmvq_batched(
10991 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10992 );
10993 }
10994 if self.mmvq_supports(qtype) {
10995 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10996 }
10997 let x0 = self.zeros(0)?;
11000 self.matmul_pre(w, aq, ad, &x0, m)
11001 }
11002
11003 pub fn matmul_decode_exact_dual_pre(
11012 &self,
11013 w0: &crate::model::GpuTensor,
11014 w1: &crate::model::GpuTensor,
11015 aq: &CudaSlice<i8>,
11016 ad: &CudaSlice<f32>,
11017 m: usize,
11018 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11019 {
11020 use crate::model::GpuTensor;
11021 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11022 let on = *ON.get_or_init(|| {
11023 std::env::var("MEMRA_SPEC_DUAL_T")
11024 .map(|v| v != "0")
11025 .unwrap_or(true)
11026 });
11027 if !on
11028 || !(2..=7).contains(&m)
11029 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11030 || !self.uses_q8_1_fast(w0)
11031 || !self.uses_q8_1_fast(w1)
11032 {
11033 return Ok(None);
11034 }
11035 if !self.mmvq_supports(QT_NVFP4) {
11040 return Ok(None);
11041 }
11042 let (in_f, out_f) = (w0.in_features(), w0.out_features());
11043 if w1.in_features() != in_f || w1.out_features() != out_f {
11044 return Ok(None);
11045 }
11046 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
11047 (
11048 GpuTensor::Quant {
11049 bytes: b0,
11050 qtype: q0,
11051 row_bytes: rb0,
11052 scale: s0,
11053 rp: rp0,
11054 rp4: None,
11055 ..
11056 },
11057 GpuTensor::Quant {
11058 bytes: b1,
11059 qtype: q1,
11060 row_bytes: rb1,
11061 scale: s1,
11062 rp: rp1,
11063 rp4: None,
11064 ..
11065 },
11066 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11067 (b0, b1, *rb0, *s0, *s1, *rp0)
11068 }
11069 _ => return Ok(None),
11070 };
11071 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
11074 {
11075 return Ok(None);
11076 }
11077 let (y0, y1) =
11078 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
11079 Ok(Some(((y0, s0), (y1, s1))))
11080 }
11081
11082 pub fn matmul_decode_exact_group4_pre(
11092 &self,
11093 ws: [&crate::model::GpuTensor; 4],
11094 aq: &CudaSlice<i8>,
11095 ad: &CudaSlice<f32>,
11096 m: usize,
11097 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11098 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11099 let on = *ON.get_or_init(|| {
11100 std::env::var("MEMRA_TK_GDN_GROUP")
11101 .map(|v| v != "0")
11102 .unwrap_or(true)
11103 });
11104 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
11105 }
11106
11107 pub fn matmul_decode_exact_group3_pre(
11112 &self,
11113 ws: [&crate::model::GpuTensor; 3],
11114 aq: &CudaSlice<i8>,
11115 ad: &CudaSlice<f32>,
11116 m: usize,
11117 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11118 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11119 let on = *ON.get_or_init(|| {
11120 std::env::var("MEMRA_TK_FA_GROUP")
11121 .map(|v| v != "0")
11122 .unwrap_or(true)
11123 });
11124 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
11125 }
11126
11127 fn matmul_decode_exact_group_pre(
11131 &self,
11132 ws: &[&crate::model::GpuTensor],
11133 aq: &CudaSlice<i8>,
11134 ad: &CudaSlice<f32>,
11135 m: usize,
11136 on: bool,
11137 tag: &'static str,
11138 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11139 use crate::model::GpuTensor;
11140 if !on
11141 || !(2..=8).contains(&m)
11142 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11143 || (m > 4 && !Self::b8_enabled())
11144 || !self.mmvq_supports(QT_NVFP4)
11145 || !self.batched_supports(QT_NVFP4)
11146 {
11147 return Ok(None);
11148 }
11149 let in_f = ws[0].in_features();
11150 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
11151 for w in ws {
11152 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
11153 return Ok(None);
11154 }
11155 match w {
11156 GpuTensor::Quant {
11157 bytes,
11158 qtype,
11159 scale,
11160 rp: true,
11161 rp4: None,
11162 ..
11163 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
11164 parts.push((bytes, w.out_features(), *scale));
11165 }
11166 _ => return Ok(None),
11167 }
11168 }
11169 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11171 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
11172 let mcols = if (5..=7).contains(&m) && b567 {
11173 m
11174 } else {
11175 Self::batched_mcols(m)
11176 };
11177 let kname: &'static str = match mcols {
11178 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
11179 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
11180 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
11181 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
11182 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
11183 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
11184 _ => return Ok(None),
11185 };
11186 if std::env::var("MEMRA_DEBUG").is_ok() {
11189 use std::sync::Mutex;
11190 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
11191 let mut seen = SEEN.lock().unwrap();
11192 if !seen.contains(&tag) {
11193 seen.push(tag);
11194 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
11195 }
11196 }
11197 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
11200 let three = parts.len() == 3;
11201 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
11202 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
11203 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
11204 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
11207 let cfg = LaunchConfig {
11208 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
11209 block_dim: (32, ROWS_PER_BLOCK, 1),
11210 shared_mem_bytes: 0,
11211 };
11212 let (inf, mi) = (in_f as i32, m as i32);
11213 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
11214 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
11215 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
11216 let s3 = if three { 1.0f32 } else { parts[3].2 };
11217 let w3 = if three { parts[0].0 } else { parts[3].0 };
11218 let f = self.func(kname);
11219 let __s_b = self.gpu.stream();
11220 let mut b = __s_b.launch_builder(&f);
11221 b.arg(parts[0].0)
11222 .arg(parts[1].0)
11223 .arg(parts[2].0)
11224 .arg(w3)
11225 .arg(aq)
11226 .arg(ad)
11227 .arg(&mut y0)
11228 .arg(&mut y1)
11229 .arg(&mut y2)
11230 .arg(&mut y3)
11231 .arg(&inf)
11232 .arg(&n0)
11233 .arg(&n1)
11234 .arg(&n2)
11235 .arg(&n3)
11236 .arg(&mi)
11237 .arg(&s0)
11238 .arg(&s1)
11239 .arg(&s2)
11240 .arg(&s3);
11241 unsafe {
11242 b.launch(cfg)?;
11243 }
11244 Ok(Some(if three {
11245 vec![y0, y1, y2]
11246 } else {
11247 vec![y0, y1, y2, y3]
11248 }))
11249 }
11250
11251 pub fn matmul_decode_exact_dual(
11267 &self,
11268 w0: &crate::model::GpuTensor,
11269 w1: &crate::model::GpuTensor,
11270 x: &CudaSlice<f32>,
11271 m: usize,
11272 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11273 use crate::model::GpuTensor;
11274 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11275 let on = *ON.get_or_init(|| {
11276 std::env::var("MEMRA_SPEC_DUAL_T")
11277 .map(|v| v != "0")
11278 .unwrap_or(true)
11279 });
11280 if !on
11281 || !(2..=4).contains(&m)
11282 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11283 || !self.uses_q8_1_fast(w0)
11284 || !self.uses_q8_1_fast(w1)
11285 {
11286 return Ok(None);
11287 }
11288 if !self.mmvq_supports(QT_NVFP4) {
11293 return Ok(None);
11294 }
11295 let (in_f, out_f) = (w0.in_features(), w0.out_features());
11296 if w1.in_features() != in_f || w1.out_features() != out_f {
11297 return Ok(None);
11298 }
11299 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
11300 (
11301 GpuTensor::Quant {
11302 bytes: b0,
11303 qtype: q0,
11304 row_bytes: rb0,
11305 scale: s0,
11306 rp: rp0,
11307 rp4: None,
11308 ..
11309 },
11310 GpuTensor::Quant {
11311 bytes: b1,
11312 qtype: q1,
11313 row_bytes: rb1,
11314 scale: s1,
11315 rp: rp1,
11316 rp4: None,
11317 ..
11318 },
11319 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11320 (b0, b1, *rb0, *s0, *s1, *rp0)
11321 }
11322 _ => return Ok(None),
11323 };
11324 if std::env::var("MEMRA_DEBUG").is_ok() {
11327 static ONCE: std::sync::Once = std::sync::Once::new();
11328 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
11329 }
11330 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11331 let (y0, y1) =
11332 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
11333 let mut y0 = y0;
11334 let mut y1 = y1;
11335 if s0 != 1.0 {
11336 self.scale_inplace(&mut y0, s0, m * out_f)?;
11337 }
11338 if s1 != 1.0 {
11339 self.scale_inplace(&mut y1, s1, m * out_f)?;
11340 }
11341 Ok(Some((y0, y1)))
11342 }
11343
11344 #[allow(clippy::too_many_arguments)]
11349 pub fn qmatvec_batched_dual_raw(
11350 &self,
11351 b0: &CudaSlice<u8>,
11352 b1: &CudaSlice<u8>,
11353 aq: &CudaSlice<i8>,
11354 ad: &CudaSlice<f32>,
11355 m: usize,
11356 in_f: usize,
11357 out_f: usize,
11358 row_bytes: usize,
11359 rp: bool,
11360 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11361 const ROWS_PER_BLOCK: u32 = 4;
11362 let mcols = Self::batched_mcols(m);
11363 let tiny_rp1 = rp
11366 && mcols == 4
11367 && out_f <= 128
11368 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
11369 let (name, rows_per_block) = if tiny_rp1 {
11370 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
11371 } else {
11372 match (mcols, rp, m) {
11373 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
11374 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
11375 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
11376 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
11377 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
11378 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
11379 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
11380 _ => {
11381 return Err(
11382 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
11383 );
11384 }
11385 }
11386 };
11387 let f = self.func(name);
11388 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
11389 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
11390 let cfg = LaunchConfig {
11391 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11392 block_dim: (32, ROWS_PER_BLOCK, 1),
11393 shared_mem_bytes: 0,
11394 };
11395 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11396 let __s_b = self.gpu.stream();
11397 let mut b = __s_b.launch_builder(&f);
11398 b.arg(b0)
11399 .arg(b1)
11400 .arg(aq)
11401 .arg(ad)
11402 .arg(&mut y0)
11403 .arg(&mut y1)
11404 .arg(&inf)
11405 .arg(&outf)
11406 .arg(&mi)
11407 .arg(&rb);
11408 unsafe {
11409 b.launch(cfg)?;
11410 }
11411 Ok((y0, y1))
11412 }
11413
11414 pub fn matmul_pre_dual_noscale(
11426 &self,
11427 w0: &crate::model::GpuTensor,
11428 w1: &crate::model::GpuTensor,
11429 aq: &CudaSlice<i8>,
11430 ad: &CudaSlice<f32>,
11431 m: usize,
11432 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11433 {
11434 use crate::model::GpuTensor;
11435 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11436 return Ok(None);
11437 }
11438 if !self.mmvq_supports(QT_NVFP4) {
11448 return Ok(None);
11449 }
11450 let (in_f, out_f) = (w0.in_features(), w0.out_features());
11451 if w1.in_features() != in_f || w1.out_features() != out_f {
11452 return Ok(None);
11453 }
11454 let no_mirror =
11467 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
11468 if self.q8_ffn_fuse2_on()
11469 && no_mirror(w0)
11470 && no_mirror(w1)
11471 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
11472 {
11473 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
11474 return Ok(Some(((y0, 1.0), (y1, 1.0))));
11475 }
11476 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11486 let (y0, y1) =
11487 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
11488 return Ok(Some(((y0, p0.3), (y1, p1.3))));
11489 }
11490 let (b0, q0, rb0, s0, rp0) = match w0 {
11491 GpuTensor::Quant {
11492 bytes,
11493 qtype,
11494 row_bytes,
11495 scale,
11496 rp,
11497 ..
11498 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11499 _ => return Ok(None),
11500 };
11501 let (b1, q1, rb1, s1, rp1) = match w1 {
11502 GpuTensor::Quant {
11503 bytes,
11504 qtype,
11505 row_bytes,
11506 scale,
11507 rp,
11508 ..
11509 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11510 _ => return Ok(None),
11511 };
11512 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
11513 return Ok(None);
11514 }
11515 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11517 let rows_per_block = ROWS_PER_BLOCK * RPW;
11518 let f = self.func(if rp0 {
11519 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
11520 } else {
11521 "qmatvec_nvfp4_mmvq_dual_mr2"
11522 });
11523 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
11524 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
11525 let cfg = LaunchConfig {
11526 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11527 block_dim: (32, ROWS_PER_BLOCK, 1),
11528 shared_mem_bytes: 0,
11529 };
11530 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
11531 let one = 1.0f32;
11534 let __s_b = self.gpu.stream();
11535 let mut b = __s_b.launch_builder(&f);
11536 b.arg(b0)
11537 .arg(b1)
11538 .arg(aq)
11539 .arg(ad)
11540 .arg(&mut y0)
11541 .arg(&mut y1)
11542 .arg(&inf)
11543 .arg(&outf)
11544 .arg(&mi)
11545 .arg(&rb)
11546 .arg(&one)
11547 .arg(&one);
11548 unsafe {
11549 b.launch(cfg)?;
11550 }
11551 Ok(Some(((y0, s0), (y1, s1))))
11552 }
11553
11554 #[allow(clippy::too_many_arguments)]
11562 pub fn matmul_nvfp4_fused3(
11563 &self,
11564 w0: &crate::model::GpuTensor,
11565 w1: &crate::model::GpuTensor,
11566 w2: &crate::model::GpuTensor,
11567 aq: &CudaSlice<i8>,
11568 ad: &CudaSlice<f32>,
11569 m: usize,
11570 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11571 {
11572 use crate::model::GpuTensor;
11573 if !(1..=8).contains(&m)
11580 || !self.mmvq_supports(QT_NVFP4)
11581 || !self.uses_q8_1_fast(w0)
11582 || !self.uses_q8_1_fast(w1)
11583 || !self.uses_q8_1_fast(w2)
11584 {
11585 return Ok(None);
11586 }
11587 if m > 1 {
11588 let in_f = w0.in_features();
11589 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
11590 || !self.batched_supports(QT_NVFP4)
11591 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11592 || (m > 4 && !Self::b8_enabled())
11593 || in_f % 512 != 0
11594 || in_f / 64 > 272
11595 {
11596 return Ok(None);
11597 }
11598 }
11599 let unpack = |w: &crate::model::GpuTensor| match w {
11600 GpuTensor::Quant {
11601 bytes,
11602 qtype,
11603 scale,
11604 rp,
11605 ..
11606 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11607 _ => None,
11608 };
11609 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
11610 return Ok(None);
11611 };
11612 let in_f = w0.in_features();
11613 if w1.in_features() != in_f || w2.in_features() != in_f {
11614 return Ok(None);
11615 }
11616 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
11617 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11619 let rows_pb = ROWS_PER_BLOCK * RPW;
11620 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11621 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11622 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11623 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11624 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
11625 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
11628 if m > 1 {
11629 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
11631 return Ok(None);
11632 }
11633 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
11634 let cfg = LaunchConfig {
11635 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
11636 block_dim: (32, ROWS_PER_BLOCK, 1),
11637 shared_mem_bytes: 0,
11638 };
11639 let __s_b = self.gpu.stream();
11640 let mut b = __s_b.launch_builder(&f);
11641 b.arg(b0)
11642 .arg(b1)
11643 .arg(b2)
11644 .arg(aq)
11645 .arg(ad)
11646 .arg(&mut y0)
11647 .arg(&mut y1)
11648 .arg(&mut y2)
11649 .arg(&inf)
11650 .arg(&oi0)
11651 .arg(&oi1)
11652 .arg(&oi2)
11653 .arg(&mi);
11654 unsafe {
11655 b.launch(cfg)?;
11656 }
11657 return Ok(Some((y0, y1, y2)));
11658 }
11659 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
11660 let cfg = LaunchConfig {
11661 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
11662 block_dim: (32, ROWS_PER_BLOCK, 1),
11663 shared_mem_bytes: 0,
11664 };
11665 let __s_b = self.gpu.stream();
11666 let mut b = __s_b.launch_builder(&f);
11667 b.arg(b0)
11668 .arg(b1)
11669 .arg(b2)
11670 .arg(aq)
11671 .arg(ad)
11672 .arg(&mut y0)
11673 .arg(&mut y1)
11674 .arg(&mut y2)
11675 .arg(&inf)
11676 .arg(&oi0)
11677 .arg(&oi1)
11678 .arg(&oi2)
11679 .arg(&mi)
11680 .arg(&p0.1)
11681 .arg(&p1.1)
11682 .arg(&p2.1);
11683 unsafe {
11684 b.launch(cfg)?;
11685 }
11686 Ok(Some((y0, y1, y2)))
11687 }
11688
11689 pub fn matmul_nvfp4_fused2(
11698 &self,
11699 w0: &crate::model::GpuTensor,
11700 w1: &crate::model::GpuTensor,
11701 aq: &CudaSlice<i8>,
11702 ad: &CudaSlice<f32>,
11703 m: usize,
11704 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11705 use crate::model::GpuTensor;
11706 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11707 let off =
11708 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11709 if off
11712 || m != 1
11713 || !self.mmvq_supports(QT_NVFP4)
11714 || !self.uses_q8_1_fast(w0)
11715 || !self.uses_q8_1_fast(w1)
11716 {
11717 return Ok(None);
11718 }
11719 let unpack = |w: &crate::model::GpuTensor| match w {
11720 GpuTensor::Quant {
11721 bytes,
11722 qtype,
11723 scale,
11724 rp,
11725 ..
11726 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11727 _ => None,
11728 };
11729 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11730 return Ok(None);
11731 };
11732 let in_f = w0.in_features();
11733 if w1.in_features() != in_f {
11734 return Ok(None);
11735 }
11736 let (o0, o1) = (w0.out_features(), w1.out_features());
11737 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11739 let rows_pb = ROWS_PER_BLOCK * RPW;
11740 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11741 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11742 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11743 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11744 let cfg = LaunchConfig {
11745 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11746 block_dim: (32, ROWS_PER_BLOCK, 1),
11747 shared_mem_bytes: 0,
11748 };
11749 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11750 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11753 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11756 {
11757 use cudarc::driver::{DevicePtr, DevicePtrMut};
11758 let s = &self.gpu.stream();
11759 let (pw0, _g0) = b0.device_ptr(s);
11760 let (pw1, _g1) = b1.device_ptr(s);
11761 let (paq, _g2) = aq.device_ptr(s);
11762 let (pad, _g3) = ad.device_ptr(s);
11763 let (py0, _g4) = y0.device_ptr_mut(s);
11764 let (py1, _g5) = y1.device_ptr_mut(s);
11765 let (s0, s1) = (p0.1, p1.1);
11766 let mut ps = [
11767 &pw0 as *const _ as *mut std::ffi::c_void,
11768 &pw1 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 &inf as *const _ as *mut _,
11774 &oi0 as *const _ as *mut _,
11775 &oi1 as *const _ as *mut _,
11776 &mi as *const _ as *mut _,
11777 &s0 as *const _ as *mut _,
11778 &s1 as *const _ as *mut _,
11779 ];
11780 unsafe {
11781 self.launch_pdl(
11782 "qmatvec_nvfp4_mmvq_fused2_rp",
11783 cfg.grid_dim,
11784 cfg.block_dim,
11785 &mut ps,
11786 )?;
11787 }
11788 }
11789 return Ok(Some((y0, y1)));
11790 }
11791 let __s_b = self.gpu.stream();
11792 let mut b = __s_b.launch_builder(&f);
11793 b.arg(b0)
11794 .arg(b1)
11795 .arg(aq)
11796 .arg(ad)
11797 .arg(&mut y0)
11798 .arg(&mut y1)
11799 .arg(&inf)
11800 .arg(&oi0)
11801 .arg(&oi1)
11802 .arg(&mi)
11803 .arg(&p0.1)
11804 .arg(&p1.1);
11805 unsafe {
11806 b.launch(cfg)?;
11807 }
11808 Ok(Some((y0, y1)))
11809 }
11810
11811 pub fn matmul_nvfp4_fused2_into(
11816 &self,
11817 w0: &crate::model::GpuTensor,
11818 w1: &crate::model::GpuTensor,
11819 aq: &CudaSlice<i8>,
11820 ad: &CudaSlice<f32>,
11821 y0: &mut CudaSlice<f32>,
11822 y1: &mut CudaSlice<f32>,
11823 ) -> Result<bool, Box<dyn std::error::Error>> {
11824 use crate::model::GpuTensor;
11825 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11826 let off =
11827 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11828 if off
11829 || !self.mmvq_supports(QT_NVFP4)
11830 || !self.uses_q8_1_fast(w0)
11831 || !self.uses_q8_1_fast(w1)
11832 {
11833 return Ok(false);
11834 }
11835 let unpack = |w: &crate::model::GpuTensor| match w {
11836 GpuTensor::Quant {
11837 bytes,
11838 qtype,
11839 scale,
11840 rp,
11841 ..
11842 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11843 _ => None,
11844 };
11845 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11846 return Ok(false);
11847 };
11848 let in_f = w0.in_features();
11849 if w1.in_features() != in_f {
11850 return Ok(false);
11851 }
11852 let (o0, o1) = (w0.out_features(), w1.out_features());
11853 if y0.len() < o0 || y1.len() < o1 {
11854 return Ok(false);
11855 }
11856 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11858 let rows_pb = ROWS_PER_BLOCK * RPW;
11859 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11860 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11861 let cfg = LaunchConfig {
11862 grid_dim: (nb(o0) + nb(o1), 1, 1),
11863 block_dim: (32, ROWS_PER_BLOCK, 1),
11864 shared_mem_bytes: 0,
11865 };
11866 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11867 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11870 let __s_b = self.gpu.stream();
11871 let mut b = __s_b.launch_builder(&f);
11872 b.arg(b0)
11873 .arg(b1)
11874 .arg(aq)
11875 .arg(ad)
11876 .arg(&mut *y0)
11877 .arg(&mut *y1)
11878 .arg(&inf)
11879 .arg(&oi0)
11880 .arg(&oi1)
11881 .arg(&mi)
11882 .arg(&p0.1)
11883 .arg(&p1.1);
11884 unsafe {
11885 b.launch(cfg)?;
11886 }
11887 Ok(true)
11888 }
11889
11890 #[allow(clippy::type_complexity)]
11895 pub fn matmul_nvfp4_fused4(
11896 &self,
11897 w0: &crate::model::GpuTensor,
11898 w1: &crate::model::GpuTensor,
11899 w2: &crate::model::GpuTensor,
11900 w3: &crate::model::GpuTensor,
11901 aq: &CudaSlice<i8>,
11902 ad: &CudaSlice<f32>,
11903 m: usize,
11904 ) -> Result<
11905 Option<(
11906 CudaSlice<f32>,
11907 CudaSlice<f32>,
11908 CudaSlice<f32>,
11909 CudaSlice<f32>,
11910 )>,
11911 Box<dyn std::error::Error>,
11912 > {
11913 use crate::model::GpuTensor;
11914 if !(1..=8).contains(&m)
11921 || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11922 || !self.mmvq_supports(QT_NVFP4)
11923 || !self.uses_q8_1_fast(w0)
11924 || !self.uses_q8_1_fast(w1)
11925 || !self.uses_q8_1_fast(w2)
11926 || !self.uses_q8_1_fast(w3)
11927 {
11928 return Ok(None);
11929 }
11930 if m > 1 {
11931 let in_f = w0.in_features();
11934 if !self.batched_supports(QT_NVFP4)
11935 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11936 || (m > 4 && !Self::b8_enabled())
11937 || in_f % 512 != 0
11938 || in_f / 64 > 272
11939 {
11940 return Ok(None);
11941 }
11942 }
11943 let unpack = |w: &crate::model::GpuTensor| match w {
11944 GpuTensor::Quant {
11945 bytes,
11946 qtype,
11947 scale,
11948 rp,
11949 ..
11950 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11951 _ => None,
11952 };
11953 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11954 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11955 else {
11956 return Ok(None);
11957 };
11958 let in_f = w0.in_features();
11959 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11960 return Ok(None);
11961 }
11962 let (o0, o1, o2, o3) = (
11963 w0.out_features(),
11964 w1.out_features(),
11965 w2.out_features(),
11966 w3.out_features(),
11967 );
11968 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11970 let rows_pb = ROWS_PER_BLOCK * RPW;
11971 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11972 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11973 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11974 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11975 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11976 let (inf, oi0, oi1, oi2, oi3, mi) = (
11977 in_f as i32,
11978 o0 as i32,
11979 o1 as i32,
11980 o2 as i32,
11981 o3 as i32,
11982 m as i32,
11983 );
11984 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11987 if m > 1 {
11988 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
11991 return Ok(None);
11992 }
11993 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
11994 let cfg = LaunchConfig {
11995 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
11996 block_dim: (32, ROWS_PER_BLOCK, 1),
11997 shared_mem_bytes: 0,
11998 };
11999 let __s_b = self.gpu.stream();
12000 let mut b = __s_b.launch_builder(&f);
12001 b.arg(b0)
12002 .arg(b1)
12003 .arg(b2)
12004 .arg(b3)
12005 .arg(aq)
12006 .arg(ad)
12007 .arg(&mut y0)
12008 .arg(&mut y1)
12009 .arg(&mut y2)
12010 .arg(&mut y3)
12011 .arg(&inf)
12012 .arg(&oi0)
12013 .arg(&oi1)
12014 .arg(&oi2)
12015 .arg(&oi3)
12016 .arg(&mi);
12017 unsafe {
12018 b.launch(cfg)?;
12019 }
12020 return Ok(Some((y0, y1, y2, y3)));
12021 }
12022 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
12023 let cfg = LaunchConfig {
12024 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
12025 block_dim: (32, ROWS_PER_BLOCK, 1),
12026 shared_mem_bytes: 0,
12027 };
12028 let __s_b = self.gpu.stream();
12029 let mut b = __s_b.launch_builder(&f);
12030 b.arg(b0)
12031 .arg(b1)
12032 .arg(b2)
12033 .arg(b3)
12034 .arg(aq)
12035 .arg(ad)
12036 .arg(&mut y0)
12037 .arg(&mut y1)
12038 .arg(&mut y2)
12039 .arg(&mut y3)
12040 .arg(&inf)
12041 .arg(&oi0)
12042 .arg(&oi1)
12043 .arg(&oi2)
12044 .arg(&oi3)
12045 .arg(&mi)
12046 .arg(&p0.1)
12047 .arg(&p1.1)
12048 .arg(&p2.1)
12049 .arg(&p3.1);
12050 unsafe {
12051 b.launch(cfg)?;
12052 }
12053 Ok(Some((y0, y1, y2, y3)))
12054 }
12055
12056 pub fn matmul_q8_fused2(
12064 &self,
12065 w0: &crate::model::GpuTensor,
12066 w1: &crate::model::GpuTensor,
12067 aq: &CudaSlice<i8>,
12068 ad: &CudaSlice<f32>,
12069 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12070 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12076 return Ok(Some(self.e4m3_fused2_core(
12077 p0.0,
12078 p1.0,
12079 aq,
12080 ad,
12081 w0.in_features(),
12082 p0.1,
12083 p1.1,
12084 p0.2,
12085 p0.3,
12086 p1.3,
12087 )?));
12088 }
12089 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12090 return Ok(None);
12091 };
12092 Ok(Some(self.q8_fused2_core(
12093 p0.0,
12094 p1.0,
12095 aq,
12096 ad,
12097 w0.in_features(),
12098 p0.1,
12099 p1.1,
12100 p0.2,
12101 )?))
12102 }
12103
12104 #[allow(clippy::too_many_arguments)]
12105 fn q8_fused2_core(
12106 &self,
12107 b0: &CudaSlice<u8>,
12108 b1: &CudaSlice<u8>,
12109 aq: &CudaSlice<i8>,
12110 ad: &CudaSlice<f32>,
12111 in_f: usize,
12112 out0: usize,
12113 out1: usize,
12114 row_bytes: usize,
12115 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12116 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12118 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12119 let f = self.func("qmatvec_q8_0_mmvq_fused2");
12120 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12121 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12122 let cfg = LaunchConfig {
12123 grid_dim: (nb0 + nb1, 1, 1),
12124 block_dim: (32, ROWS_PER_BLOCK, 1),
12125 shared_mem_bytes: 0,
12126 };
12127 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12128 let __s_b = self.gpu.stream();
12129 let mut b = __s_b.launch_builder(&f);
12130 b.arg(b0)
12131 .arg(b1)
12132 .arg(aq)
12133 .arg(ad)
12134 .arg(&mut y0)
12135 .arg(&mut y1)
12136 .arg(&inf)
12137 .arg(&o0)
12138 .arg(&o1)
12139 .arg(&rbl);
12140 unsafe {
12141 b.launch(cfg)?;
12142 }
12143 Ok((y0, y1))
12144 }
12145
12146 pub fn matmul_q8_fused2_x(
12152 &self,
12153 w0: &crate::model::GpuTensor,
12154 w1: &crate::model::GpuTensor,
12155 x: &CudaSlice<f32>,
12156 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12157 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12158 return Ok(None);
12159 }
12160 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12161 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
12162 return Ok(Some(self.e4m3_fused2_core(
12163 p0.0,
12164 p1.0,
12165 &aq,
12166 &ad,
12167 w0.in_features(),
12168 p0.1,
12169 p1.1,
12170 p0.2,
12171 p0.3,
12172 p1.3,
12173 )?));
12174 }
12175 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12176 return Ok(None);
12177 };
12178 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
12179 Ok(Some(self.q8_fused2_core(
12180 p0.0,
12181 p1.0,
12182 &aq,
12183 &ad,
12184 w0.in_features(),
12185 p0.1,
12186 p1.1,
12187 p0.2,
12188 )?))
12189 }
12190
12191 #[allow(clippy::too_many_arguments)]
12194 pub fn qmatvec_q8_fused2_raw(
12195 &self,
12196 b0: &CudaSlice<u8>,
12197 b1: &CudaSlice<u8>,
12198 x: &CudaSlice<f32>,
12199 in_f: usize,
12200 out0: usize,
12201 out1: usize,
12202 row_bytes: usize,
12203 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12204 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12205 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
12206 }
12207
12208 pub fn matmul_q4_fused3(
12214 &self,
12215 w0: &crate::model::GpuTensor,
12216 w1: &crate::model::GpuTensor,
12217 w2: &crate::model::GpuTensor,
12218 aq: &CudaSlice<i8>,
12219 ad: &CudaSlice<f32>,
12220 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12221 {
12222 use crate::model::GpuTensor;
12223 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12224 match w {
12225 GpuTensor::Quant {
12226 qtype, row_bytes, ..
12227 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12228 _ => None,
12229 }
12230 };
12231 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
12232 return Ok(None);
12233 };
12234 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12235 return Ok(None);
12236 }
12237 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12241 match w {
12242 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12243 Some(m) => (m, true),
12244 None => (bytes, *rp),
12245 },
12246 _ => unreachable!(),
12247 }
12248 }
12249 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12250 if rp0 != rp1 || rp1 != rp2 {
12251 return Ok(None);
12252 }
12253 let rp = rp0;
12254 let rpb: u32 = 4;
12255 let mr1 = rp && Self::q40_mr1_on();
12259 let nb = |o: usize| {
12260 if mr1 {
12261 (o as u32).div_ceil(rpb)
12262 } else {
12263 (o as u32).div_ceil(2).div_ceil(rpb)
12264 }
12265 };
12266 let grid = nb(o0) + nb(o1) + nb(o2);
12267 let mut y0 = self.alloc_uninit::<f32>(o0)?;
12268 let mut y1 = self.alloc_uninit::<f32>(o1)?;
12269 let mut y2 = self.alloc_uninit::<f32>(o2)?;
12270 let f = self.func(if mr1 {
12271 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12272 } else if rp {
12273 "qmatvec_q4_0_mmvq_fused3_rp"
12274 } else {
12275 "qmatvec_q4_0_mmvq_fused3"
12276 });
12277 let cfg = LaunchConfig {
12278 grid_dim: (grid, 1, 1),
12279 block_dim: (32, rpb, 1),
12280 shared_mem_bytes: 0,
12281 };
12282 let inf = w0.in_features() as i32;
12283 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12284 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12285 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12288 {
12289 use cudarc::driver::{DevicePtr, DevicePtrMut};
12290 let s = &self.gpu.stream();
12291 let (p0, _g0) = b0.device_ptr(s);
12292 let (p1, _g1) = b1.device_ptr(s);
12293 let (p2, _g2) = b2.device_ptr(s);
12294 let (paq, _g3) = aq.device_ptr(s);
12295 let (pad, _g4) = ad.device_ptr(s);
12296 let (py0, _g5) = y0.device_ptr_mut(s);
12297 let (py1, _g6) = y1.device_ptr_mut(s);
12298 let (py2, _g7) = y2.device_ptr_mut(s);
12299 let mut ps = [
12300 &p0 as *const _ as *mut std::ffi::c_void,
12301 &p1 as *const _ as *mut _,
12302 &p2 as *const _ as *mut _,
12303 &paq as *const _ as *mut _,
12304 &pad as *const _ as *mut _,
12305 &py0 as *const _ as *mut _,
12306 &py1 as *const _ as *mut _,
12307 &py2 as *const _ as *mut _,
12308 &inf as *const _ as *mut _,
12309 &oo0 as *const _ as *mut _,
12310 &oo1 as *const _ as *mut _,
12311 &oo2 as *const _ as *mut _,
12312 &r0 as *const _ as *mut _,
12313 &r1 as *const _ as *mut _,
12314 &r2 as *const _ as *mut _,
12315 ];
12316 unsafe {
12317 self.launch_pdl(
12318 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12319 (grid, 1, 1),
12320 (32, rpb, 1),
12321 &mut ps,
12322 )?;
12323 }
12324 }
12325 return Ok(Some((y0, y1, y2)));
12326 }
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(&r0)
12342 .arg(&r1)
12343 .arg(&r2);
12344 unsafe {
12345 b.launch(cfg)?;
12346 }
12347 Ok(Some((y0, y1, y2)))
12348 }
12349
12350 #[allow(clippy::too_many_arguments)]
12353 pub fn matmul_q4_fused3_into(
12354 &self,
12355 w0: &crate::model::GpuTensor,
12356 w1: &crate::model::GpuTensor,
12357 w2: &crate::model::GpuTensor,
12358 aq: &CudaSlice<i8>,
12359 ad: &CudaSlice<f32>,
12360 y0: &mut CudaSlice<f32>,
12361 y1: &mut CudaSlice<f32>,
12362 y2: &mut CudaSlice<f32>,
12363 ) -> Result<bool, Box<dyn std::error::Error>> {
12364 use crate::model::GpuTensor;
12365 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12366 match w {
12367 GpuTensor::Quant {
12368 qtype, row_bytes, ..
12369 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12370 _ => None,
12371 }
12372 };
12373 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
12374 return Ok(false);
12375 };
12376 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12377 return Ok(false);
12378 }
12379 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12380 match w {
12381 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12382 Some(m) => (m, true),
12383 None => (bytes, *rp),
12384 },
12385 _ => unreachable!(),
12386 }
12387 }
12388 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12389 if rp0 != rp1 || rp1 != rp2 {
12390 return Ok(false);
12391 }
12392 let rp = rp0;
12393 let rpb: u32 = 4;
12394 let mr1 = rp && Self::q40_mr1_on();
12395 let nb = |o: usize| {
12396 if mr1 {
12397 (o as u32).div_ceil(rpb)
12398 } else {
12399 (o as u32).div_ceil(2).div_ceil(rpb)
12400 }
12401 };
12402 let grid = nb(o0) + nb(o1) + nb(o2);
12403 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
12404 let f = self.func(if mr1 {
12405 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12406 } else if rp {
12407 "qmatvec_q4_0_mmvq_fused3_rp"
12408 } else {
12409 "qmatvec_q4_0_mmvq_fused3"
12410 });
12411 let cfg = LaunchConfig {
12412 grid_dim: (grid, 1, 1),
12413 block_dim: (32, rpb, 1),
12414 shared_mem_bytes: 0,
12415 };
12416 let inf = w0.in_features() as i32;
12417 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12418 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12419 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12421 use cudarc::driver::{DevicePtr, DevicePtrMut};
12422 let s = &self.gpu.stream();
12423 let (p0, _g0) = b0.device_ptr(s);
12424 let (p1, _g1) = b1.device_ptr(s);
12425 let (p2, _g2) = b2.device_ptr(s);
12426 let (paq, _g3) = aq.device_ptr(s);
12427 let (pad, _g4) = ad.device_ptr(s);
12428 let (py0, _g5) = y0.device_ptr_mut(s);
12429 let (py1, _g6) = y1.device_ptr_mut(s);
12430 let (py2, _g7) = y2.device_ptr_mut(s);
12431 let mut ps = [
12432 &p0 as *const _ as *mut std::ffi::c_void,
12433 &p1 as *const _ as *mut _,
12434 &p2 as *const _ as *mut _,
12435 &paq as *const _ as *mut _,
12436 &pad as *const _ as *mut _,
12437 &py0 as *const _ as *mut _,
12438 &py1 as *const _ as *mut _,
12439 &py2 as *const _ as *mut _,
12440 &inf as *const _ as *mut _,
12441 &oo0 as *const _ as *mut _,
12442 &oo1 as *const _ as *mut _,
12443 &oo2 as *const _ as *mut _,
12444 &r0 as *const _ as *mut _,
12445 &r1 as *const _ as *mut _,
12446 &r2 as *const _ as *mut _,
12447 ];
12448 unsafe {
12449 self.launch_pdl(
12450 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12451 (grid, 1, 1),
12452 (32, rpb, 1),
12453 &mut ps,
12454 )?;
12455 }
12456 return Ok(true);
12457 }
12458 let __s_b = self.gpu.stream();
12459 let mut b = __s_b.launch_builder(&f);
12460 b.arg(b0)
12461 .arg(b1)
12462 .arg(b2)
12463 .arg(aq)
12464 .arg(ad)
12465 .arg(&mut *y0)
12466 .arg(&mut *y1)
12467 .arg(&mut *y2)
12468 .arg(&inf)
12469 .arg(&oo0)
12470 .arg(&oo1)
12471 .arg(&oo2)
12472 .arg(&r0)
12473 .arg(&r1)
12474 .arg(&r2);
12475 unsafe {
12476 b.launch(cfg)?;
12477 }
12478 Ok(true)
12479 }
12480
12481 pub fn matmul_q4_fused2(
12483 &self,
12484 w0: &crate::model::GpuTensor,
12485 w1: &crate::model::GpuTensor,
12486 aq: &CudaSlice<i8>,
12487 ad: &CudaSlice<f32>,
12488 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12489 use crate::model::GpuTensor;
12490 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12491 match w {
12492 GpuTensor::Quant {
12493 qtype, row_bytes, ..
12494 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12495 _ => None,
12496 }
12497 };
12498 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12499 return Ok(None);
12500 };
12501 if w0.in_features() != w1.in_features() {
12502 return Ok(None);
12503 }
12504 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12506 match w {
12507 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12508 Some(m) => (m, true),
12509 None => (bytes, *rp),
12510 },
12511 _ => unreachable!(),
12512 }
12513 }
12514 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12515 if rp0 != rp1 {
12516 return Ok(None);
12517 }
12518 let rp = rp0;
12519 let rpb: u32 = 4;
12520 let mr1 = rp && Self::q40_mr1_on();
12522 let nb = |o: usize| {
12523 if mr1 {
12524 (o as u32).div_ceil(rpb)
12525 } else {
12526 (o as u32).div_ceil(2).div_ceil(rpb)
12527 }
12528 };
12529 let grid = nb(o0) + nb(o1);
12530 let mut y0 = self.alloc_uninit::<f32>(o0)?;
12531 let mut y1 = self.alloc_uninit::<f32>(o1)?;
12532 let f = self.func(if mr1 {
12533 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12534 } else if rp {
12535 "qmatvec_q4_0_mmvq_fused2_rp"
12536 } else {
12537 "qmatvec_q4_0_mmvq_fused2"
12538 });
12539 let cfg = LaunchConfig {
12540 grid_dim: (grid, 1, 1),
12541 block_dim: (32, rpb, 1),
12542 shared_mem_bytes: 0,
12543 };
12544 let inf = w0.in_features() as i32;
12545 let (oo0, oo1) = (o0 as i32, o1 as i32);
12546 let (r0, r1) = (rb0 as i64, rb1 as i64);
12547 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12549 {
12550 use cudarc::driver::{DevicePtr, DevicePtrMut};
12551 let s = &self.gpu.stream();
12552 let (p0, _g0) = b0.device_ptr(s);
12553 let (p1, _g1) = b1.device_ptr(s);
12554 let (paq, _g2) = aq.device_ptr(s);
12555 let (pad, _g3) = ad.device_ptr(s);
12556 let (py0, _g4) = y0.device_ptr_mut(s);
12557 let (py1, _g5) = y1.device_ptr_mut(s);
12558 let mut ps = [
12559 &p0 as *const _ as *mut std::ffi::c_void,
12560 &p1 as *const _ as *mut _,
12561 &paq as *const _ as *mut _,
12562 &pad as *const _ as *mut _,
12563 &py0 as *const _ as *mut _,
12564 &py1 as *const _ as *mut _,
12565 &inf as *const _ as *mut _,
12566 &oo0 as *const _ as *mut _,
12567 &oo1 as *const _ as *mut _,
12568 &r0 as *const _ as *mut _,
12569 &r1 as *const _ as *mut _,
12570 ];
12571 unsafe {
12572 self.launch_pdl(
12573 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12574 (grid, 1, 1),
12575 (32, rpb, 1),
12576 &mut ps,
12577 )?;
12578 }
12579 }
12580 return Ok(Some((y0, y1)));
12581 }
12582 let __s_b = self.gpu.stream();
12583 let mut b = __s_b.launch_builder(&f);
12584 b.arg(b0)
12585 .arg(b1)
12586 .arg(aq)
12587 .arg(ad)
12588 .arg(&mut y0)
12589 .arg(&mut y1)
12590 .arg(&inf)
12591 .arg(&oo0)
12592 .arg(&oo1)
12593 .arg(&r0)
12594 .arg(&r1);
12595 unsafe {
12596 b.launch(cfg)?;
12597 }
12598 Ok(Some((y0, y1)))
12599 }
12600
12601 pub fn matmul_q4_fused2_into(
12603 &self,
12604 w0: &crate::model::GpuTensor,
12605 w1: &crate::model::GpuTensor,
12606 aq: &CudaSlice<i8>,
12607 ad: &CudaSlice<f32>,
12608 y0: &mut CudaSlice<f32>,
12609 y1: &mut CudaSlice<f32>,
12610 ) -> Result<bool, Box<dyn std::error::Error>> {
12611 use crate::model::GpuTensor;
12612 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12613 match w {
12614 GpuTensor::Quant {
12615 qtype, row_bytes, ..
12616 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12617 _ => None,
12618 }
12619 };
12620 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12621 return Ok(false);
12622 };
12623 if w0.in_features() != w1.in_features() {
12624 return Ok(false);
12625 }
12626 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12627 match w {
12628 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12629 Some(m) => (m, true),
12630 None => (bytes, *rp),
12631 },
12632 _ => unreachable!(),
12633 }
12634 }
12635 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12636 if rp0 != rp1 {
12637 return Ok(false);
12638 }
12639 let rp = rp0;
12640 let rpb: u32 = 4;
12641 let mr1 = rp && Self::q40_mr1_on();
12642 let nb = |o: usize| {
12643 if mr1 {
12644 (o as u32).div_ceil(rpb)
12645 } else {
12646 (o as u32).div_ceil(2).div_ceil(rpb)
12647 }
12648 };
12649 let grid = nb(o0) + nb(o1);
12650 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
12651 let f = self.func(if mr1 {
12652 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12653 } else if rp {
12654 "qmatvec_q4_0_mmvq_fused2_rp"
12655 } else {
12656 "qmatvec_q4_0_mmvq_fused2"
12657 });
12658 let cfg = LaunchConfig {
12659 grid_dim: (grid, 1, 1),
12660 block_dim: (32, rpb, 1),
12661 shared_mem_bytes: 0,
12662 };
12663 let inf = w0.in_features() as i32;
12664 let (oo0, oo1) = (o0 as i32, o1 as i32);
12665 let (r0, r1) = (rb0 as i64, rb1 as i64);
12666 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12668 use cudarc::driver::{DevicePtr, DevicePtrMut};
12669 let s = &self.gpu.stream();
12670 let (p0, _g0) = b0.device_ptr(s);
12671 let (p1, _g1) = b1.device_ptr(s);
12672 let (paq, _g2) = aq.device_ptr(s);
12673 let (pad, _g3) = ad.device_ptr(s);
12674 let (py0, _g4) = y0.device_ptr_mut(s);
12675 let (py1, _g5) = y1.device_ptr_mut(s);
12676 let mut ps = [
12677 &p0 as *const _ as *mut std::ffi::c_void,
12678 &p1 as *const _ as *mut _,
12679 &paq as *const _ as *mut _,
12680 &pad as *const _ as *mut _,
12681 &py0 as *const _ as *mut _,
12682 &py1 as *const _ as *mut _,
12683 &inf as *const _ as *mut _,
12684 &oo0 as *const _ as *mut _,
12685 &oo1 as *const _ as *mut _,
12686 &r0 as *const _ as *mut _,
12687 &r1 as *const _ as *mut _,
12688 ];
12689 unsafe {
12690 self.launch_pdl(
12691 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12692 (grid, 1, 1),
12693 (32, rpb, 1),
12694 &mut ps,
12695 )?;
12696 }
12697 return Ok(true);
12698 }
12699 let __s_b = self.gpu.stream();
12700 let mut b = __s_b.launch_builder(&f);
12701 b.arg(b0)
12702 .arg(b1)
12703 .arg(aq)
12704 .arg(ad)
12705 .arg(&mut *y0)
12706 .arg(&mut *y1)
12707 .arg(&inf)
12708 .arg(&oo0)
12709 .arg(&oo1)
12710 .arg(&r0)
12711 .arg(&r1);
12712 unsafe {
12713 b.launch(cfg)?;
12714 }
12715 Ok(true)
12716 }
12717
12718 pub fn matmul_q4_fused2_batched(
12723 &self,
12724 w0: &crate::model::GpuTensor,
12725 w1: &crate::model::GpuTensor,
12726 aq: &CudaSlice<i8>,
12727 ad: &CudaSlice<f32>,
12728 m: usize,
12729 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12730 use crate::model::GpuTensor;
12731 if m < 2 || m > 8 {
12732 return Ok(None);
12733 }
12734 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12735 match w {
12736 GpuTensor::Quant {
12737 qtype, row_bytes, ..
12738 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12739 _ => None,
12740 }
12741 };
12742 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
12743 return Ok(None);
12744 };
12745 if w0.in_features() != w1.in_features() {
12746 return Ok(None);
12747 }
12748 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12749 match w {
12750 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12751 Some(mr) => (mr, true),
12752 None => (bytes, *rp),
12753 },
12754 _ => unreachable!(),
12755 }
12756 }
12757 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12758 if !rp0 || !rp1 {
12759 return Ok(None);
12760 }
12761 let mcols = Self::batched_mcols(m);
12762 let rpb: u32 = 4;
12763 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12764 let grid = nb(o0) + nb(o1);
12765 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12766 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12767 let f = self.func(match mcols {
12768 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12769 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12770 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12771 });
12772 let cfg = LaunchConfig {
12773 grid_dim: (grid, 1, 1),
12774 block_dim: (32, rpb, 1),
12775 shared_mem_bytes: 0,
12776 };
12777 let inf = w0.in_features() as i32;
12778 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12779 let rb = rb0 as i64;
12780 let __s_b = self.gpu.stream();
12781 let mut b = __s_b.launch_builder(&f);
12782 b.arg(b0)
12783 .arg(b1)
12784 .arg(aq)
12785 .arg(ad)
12786 .arg(&mut y0)
12787 .arg(&mut y1)
12788 .arg(&inf)
12789 .arg(&oo0)
12790 .arg(&oo1)
12791 .arg(&mi)
12792 .arg(&rb);
12793 unsafe {
12794 b.launch(cfg)?;
12795 }
12796 Ok(Some((y0, y1)))
12797 }
12798
12799 #[allow(clippy::too_many_arguments)]
12802 pub fn matmul_q4_fused3_batched(
12803 &self,
12804 w0: &crate::model::GpuTensor,
12805 w1: &crate::model::GpuTensor,
12806 w2: &crate::model::GpuTensor,
12807 aq: &CudaSlice<i8>,
12808 ad: &CudaSlice<f32>,
12809 m: usize,
12810 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12811 {
12812 use crate::model::GpuTensor;
12813 if m < 2 || m > 8 {
12814 return Ok(None);
12815 }
12816 let q4 = |w: &GpuTensor| -> Option<usize> {
12817 match w {
12818 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12819 _ => None,
12820 }
12821 };
12822 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12823 return Ok(None);
12824 };
12825 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12826 return Ok(None);
12827 }
12828 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12829 match w {
12830 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12831 Some(mr) => (mr, true),
12832 None => (bytes, *rp),
12833 },
12834 _ => unreachable!(),
12835 }
12836 }
12837 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12838 if !rp0 || !rp1 || !rp2 {
12839 return Ok(None);
12840 }
12841 let mcols = Self::batched_mcols(m);
12842 let rpb: u32 = 4;
12843 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12844 let grid = nb(o0) + nb(o1) + nb(o2);
12845 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12846 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12847 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12848 let f = self.func(match mcols {
12849 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12850 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12851 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12852 });
12853 let cfg = LaunchConfig {
12854 grid_dim: (grid, 1, 1),
12855 block_dim: (32, rpb, 1),
12856 shared_mem_bytes: 0,
12857 };
12858 let inf = w0.in_features() as i32;
12859 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12860 let rb = 0i64;
12861 let __s_b = self.gpu.stream();
12862 let mut b = __s_b.launch_builder(&f);
12863 b.arg(b0)
12864 .arg(b1)
12865 .arg(b2)
12866 .arg(aq)
12867 .arg(ad)
12868 .arg(&mut y0)
12869 .arg(&mut y1)
12870 .arg(&mut y2)
12871 .arg(&inf)
12872 .arg(&oo0)
12873 .arg(&oo1)
12874 .arg(&oo2)
12875 .arg(&mi)
12876 .arg(&rb);
12877 unsafe {
12878 b.launch(cfg)?;
12879 }
12880 Ok(Some((y0, y1, y2)))
12881 }
12882
12883 pub fn matmul_q8_fused3(
12884 &self,
12885 w0: &crate::model::GpuTensor,
12886 w1: &crate::model::GpuTensor,
12887 w2: &crate::model::GpuTensor,
12888 aq: &CudaSlice<i8>,
12889 ad: &CudaSlice<f32>,
12890 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12891 {
12892 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12895 return Ok(Some(self.e4m3_fused3_core(
12896 p0.0,
12897 p1.0,
12898 p2.0,
12899 aq,
12900 ad,
12901 w0.in_features(),
12902 p0.1,
12903 p1.1,
12904 p2.1,
12905 p0.2,
12906 p0.3,
12907 p1.3,
12908 p2.3,
12909 )?));
12910 }
12911 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12912 return Ok(None);
12913 };
12914 Ok(Some(self.q8_fused3_core(
12915 p0.0,
12916 p1.0,
12917 p2.0,
12918 aq,
12919 ad,
12920 w0.in_features(),
12921 p0.1,
12922 p1.1,
12923 p2.1,
12924 p0.2,
12925 )?))
12926 }
12927
12928 #[allow(clippy::too_many_arguments)]
12929 fn q8_fused3_core(
12930 &self,
12931 b0: &CudaSlice<u8>,
12932 b1: &CudaSlice<u8>,
12933 b2: &CudaSlice<u8>,
12934 aq: &CudaSlice<i8>,
12935 ad: &CudaSlice<f32>,
12936 in_f: usize,
12937 out0: usize,
12938 out1: usize,
12939 out2: usize,
12940 row_bytes: usize,
12941 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12942 const ROWS_PER_BLOCK: u32 = 4;
12943 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12944 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12945 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12946 let f = self.func("qmatvec_q8_0_mmvq_fused3");
12947 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12948 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12949 let mut y2 = self.alloc_uninit::<f32>(out2)?;
12950 let cfg = LaunchConfig {
12951 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12952 block_dim: (32, ROWS_PER_BLOCK, 1),
12953 shared_mem_bytes: 0,
12954 };
12955 let (inf, o0, o1, o2, rbl) = (
12956 in_f as i32,
12957 out0 as i32,
12958 out1 as i32,
12959 out2 as i32,
12960 row_bytes as i64,
12961 );
12962 let __s_b = self.gpu.stream();
12963 let mut b = __s_b.launch_builder(&f);
12964 b.arg(b0)
12965 .arg(b1)
12966 .arg(b2)
12967 .arg(aq)
12968 .arg(ad)
12969 .arg(&mut y0)
12970 .arg(&mut y1)
12971 .arg(&mut y2)
12972 .arg(&inf)
12973 .arg(&o0)
12974 .arg(&o1)
12975 .arg(&o2)
12976 .arg(&rbl);
12977 unsafe {
12978 b.launch(cfg)?;
12979 }
12980 Ok((y0, y1, y2))
12981 }
12982
12983 #[allow(clippy::too_many_arguments)]
12985 pub fn qmatvec_q8_fused3_raw(
12986 &self,
12987 b0: &CudaSlice<u8>,
12988 b1: &CudaSlice<u8>,
12989 b2: &CudaSlice<u8>,
12990 x: &CudaSlice<f32>,
12991 in_f: usize,
12992 out0: usize,
12993 out1: usize,
12994 out2: usize,
12995 row_bytes: usize,
12996 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12997 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12998 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12999 }
13000
13001 pub fn matmul_q8_fused2_t(
13012 &self,
13013 w0: &crate::model::GpuTensor,
13014 w1: &crate::model::GpuTensor,
13015 aq: &CudaSlice<i8>,
13016 ad: &CudaSlice<f32>,
13017 m: usize,
13018 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13019 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
13023 return Ok(None);
13024 }
13025 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13028 if m > 4 && !Self::b8_enabled() {
13029 return Ok(None);
13030 }
13031 return Ok(Some(self.e4m3_fused2_t_core(
13032 p0.0,
13033 p1.0,
13034 aq,
13035 ad,
13036 m,
13037 w0.in_features(),
13038 p0.1,
13039 p1.1,
13040 p0.2,
13041 p0.3,
13042 p1.3,
13043 )?));
13044 }
13045 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13046 return Ok(None);
13047 };
13048 Ok(Some(self.q8_fused2_t_core(
13049 p0.0,
13050 p1.0,
13051 aq,
13052 ad,
13053 m,
13054 w0.in_features(),
13055 p0.1,
13056 p1.1,
13057 p0.2,
13058 )?))
13059 }
13060
13061 #[allow(clippy::too_many_arguments)]
13062 fn q8_fused2_t_core(
13063 &self,
13064 b0: &CudaSlice<u8>,
13065 b1: &CudaSlice<u8>,
13066 aq: &CudaSlice<i8>,
13067 ad: &CudaSlice<f32>,
13068 m: usize,
13069 in_f: usize,
13070 out0: usize,
13071 out1: usize,
13072 row_bytes: usize,
13073 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13074 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13076 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13077 let f = self.func(match Self::batched_mcols(m) {
13078 2 => "qmatvec_q8_0_mmvq_fused2_b2",
13079 4 => "qmatvec_q8_0_mmvq_fused2_b4",
13080 _ => "qmatvec_q8_0_mmvq_fused2_b8",
13082 });
13083 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13084 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13085 let cfg = LaunchConfig {
13086 grid_dim: (nb0 + nb1, 1, 1),
13087 block_dim: (32, ROWS_PER_BLOCK, 1),
13088 shared_mem_bytes: 0,
13089 };
13090 let (inf, o0, o1, mi, rbl) = (
13091 in_f as i32,
13092 out0 as i32,
13093 out1 as i32,
13094 m as i32,
13095 row_bytes as i64,
13096 );
13097 let __s_b = self.gpu.stream();
13098 let mut b = __s_b.launch_builder(&f);
13099 b.arg(b0)
13100 .arg(b1)
13101 .arg(aq)
13102 .arg(ad)
13103 .arg(&mut y0)
13104 .arg(&mut y1)
13105 .arg(&inf)
13106 .arg(&o0)
13107 .arg(&o1)
13108 .arg(&mi)
13109 .arg(&rbl);
13110 unsafe {
13111 b.launch(cfg)?;
13112 }
13113 Ok((y0, y1))
13114 }
13115
13116 #[allow(clippy::too_many_arguments)]
13119 pub fn qmatvec_q8_fused2_t_raw(
13120 &self,
13121 b0: &CudaSlice<u8>,
13122 b1: &CudaSlice<u8>,
13123 x: &CudaSlice<f32>,
13124 m: usize,
13125 in_f: usize,
13126 out0: usize,
13127 out1: usize,
13128 row_bytes: usize,
13129 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13130 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13131 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
13132 }
13133
13134 #[allow(clippy::too_many_arguments)]
13137 pub fn matmul_q8_fused3_t(
13138 &self,
13139 w0: &crate::model::GpuTensor,
13140 w1: &crate::model::GpuTensor,
13141 w2: &crate::model::GpuTensor,
13142 aq: &CudaSlice<i8>,
13143 ad: &CudaSlice<f32>,
13144 m: usize,
13145 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13146 {
13147 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
13148 return Ok(None);
13149 }
13150 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
13151 return Ok(Some(self.e4m3_fused3_t_core(
13152 p0.0,
13153 p1.0,
13154 p2.0,
13155 aq,
13156 ad,
13157 m,
13158 w0.in_features(),
13159 p0.1,
13160 p1.1,
13161 p2.1,
13162 p0.2,
13163 p0.3,
13164 p1.3,
13165 p2.3,
13166 )?));
13167 }
13168 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
13169 return Ok(None);
13170 };
13171 Ok(Some(self.q8_fused3_t_core(
13172 p0.0,
13173 p1.0,
13174 p2.0,
13175 aq,
13176 ad,
13177 m,
13178 w0.in_features(),
13179 p0.1,
13180 p1.1,
13181 p2.1,
13182 p0.2,
13183 )?))
13184 }
13185
13186 #[allow(clippy::too_many_arguments)]
13187 fn q8_fused3_t_core(
13188 &self,
13189 b0: &CudaSlice<u8>,
13190 b1: &CudaSlice<u8>,
13191 b2: &CudaSlice<u8>,
13192 aq: &CudaSlice<i8>,
13193 ad: &CudaSlice<f32>,
13194 m: usize,
13195 in_f: usize,
13196 out0: usize,
13197 out1: usize,
13198 out2: usize,
13199 row_bytes: usize,
13200 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13201 const ROWS_PER_BLOCK: u32 = 4;
13202 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13203 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13204 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13205 let f = self.func(if Self::batched_mcols(m) == 2 {
13206 "qmatvec_q8_0_mmvq_fused3_b2"
13207 } else {
13208 "qmatvec_q8_0_mmvq_fused3_b4"
13209 });
13210 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13211 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13212 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13213 let cfg = LaunchConfig {
13214 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13215 block_dim: (32, ROWS_PER_BLOCK, 1),
13216 shared_mem_bytes: 0,
13217 };
13218 let (inf, o0, o1, o2, mi, rbl) = (
13219 in_f as i32,
13220 out0 as i32,
13221 out1 as i32,
13222 out2 as i32,
13223 m as i32,
13224 row_bytes as i64,
13225 );
13226 let __s_b = self.gpu.stream();
13227 let mut b = __s_b.launch_builder(&f);
13228 b.arg(b0)
13229 .arg(b1)
13230 .arg(b2)
13231 .arg(aq)
13232 .arg(ad)
13233 .arg(&mut y0)
13234 .arg(&mut y1)
13235 .arg(&mut y2)
13236 .arg(&inf)
13237 .arg(&o0)
13238 .arg(&o1)
13239 .arg(&o2)
13240 .arg(&mi)
13241 .arg(&rbl);
13242 unsafe {
13243 b.launch(cfg)?;
13244 }
13245 Ok((y0, y1, y2))
13246 }
13247
13248 #[allow(clippy::too_many_arguments)]
13250 pub fn qmatvec_q8_fused3_t_raw(
13251 &self,
13252 b0: &CudaSlice<u8>,
13253 b1: &CudaSlice<u8>,
13254 b2: &CudaSlice<u8>,
13255 x: &CudaSlice<f32>,
13256 m: usize,
13257 in_f: usize,
13258 out0: usize,
13259 out1: usize,
13260 out2: usize,
13261 row_bytes: usize,
13262 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13263 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13264 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
13265 }
13266
13267 pub fn q8_ffn_fuse2_on(&self) -> bool {
13271 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13272 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
13273 }
13274
13275 #[allow(clippy::type_complexity)]
13281 fn q8_fused_params<'w, const N: usize>(
13282 &self,
13283 ws: &[&'w crate::model::GpuTensor; N],
13284 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
13285 use crate::model::GpuTensor;
13286 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13287 return None;
13288 }
13289 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
13290 return None;
13291 }
13292 let in_f = ws[0].in_features();
13293 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
13294 for (i, w) in ws.iter().enumerate() {
13295 match w {
13296 GpuTensor::Quant {
13297 bytes,
13298 qtype,
13299 row_bytes,
13300 scale,
13301 ..
13302 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
13303 out[i] = Some((bytes, w.out_features(), *row_bytes))
13304 }
13305 _ => return None,
13306 }
13307 }
13308 Some(out.map(|o| o.unwrap()))
13309 }
13310
13311 pub fn e4m3_dual_on(&self) -> bool {
13314 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13315 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
13316 }
13317
13318 #[allow(clippy::type_complexity)]
13330 fn e4m3_fused_params<'w, const N: usize>(
13331 &self,
13332 ws: &[&'w crate::model::GpuTensor; N],
13333 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
13334 use crate::model::GpuTensor;
13335 if !self.e4m3_dual_on() {
13336 return None;
13337 }
13338 let in_f = ws[0].in_features();
13339 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
13340 for (i, w) in ws.iter().enumerate() {
13341 match w {
13342 GpuTensor::Quant {
13343 bytes,
13344 qtype,
13345 row_bytes,
13346 scale,
13347 rp,
13348 rp4,
13349 ..
13350 } if *qtype == QT_F8_E4M3
13351 && w.in_features() == in_f
13352 && *row_bytes == in_f
13353 && !*rp
13354 && rp4.is_none() =>
13355 {
13356 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
13357 }
13358 _ => return None,
13359 }
13360 }
13361 Some(out.map(|o| o.unwrap()))
13362 }
13363
13364 #[allow(clippy::too_many_arguments)]
13368 fn e4m3_fused2_core(
13369 &self,
13370 b0: &CudaSlice<u8>,
13371 b1: &CudaSlice<u8>,
13372 aq: &CudaSlice<i8>,
13373 ad: &CudaSlice<f32>,
13374 in_f: usize,
13375 out0: usize,
13376 out1: usize,
13377 row_bytes: usize,
13378 ws0: f32,
13379 ws1: f32,
13380 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13381 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13383 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13384 let f = self.func("qmatvec_e4m3_mmvq_fused2");
13385 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13386 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13387 let cfg = LaunchConfig {
13388 grid_dim: (nb0 + nb1, 1, 1),
13389 block_dim: (32, ROWS_PER_BLOCK, 1),
13390 shared_mem_bytes: 0,
13391 };
13392 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13393 let __s_b = self.gpu.stream();
13394 let mut b = __s_b.launch_builder(&f);
13395 b.arg(b0)
13396 .arg(b1)
13397 .arg(aq)
13398 .arg(ad)
13399 .arg(&mut y0)
13400 .arg(&mut y1)
13401 .arg(&inf)
13402 .arg(&o0)
13403 .arg(&o1)
13404 .arg(&rbl)
13405 .arg(&ws0)
13406 .arg(&ws1);
13407 unsafe {
13408 b.launch(cfg)?;
13409 }
13410 Ok((y0, y1))
13411 }
13412
13413 #[allow(clippy::too_many_arguments)]
13415 fn e4m3_fused3_core(
13416 &self,
13417 b0: &CudaSlice<u8>,
13418 b1: &CudaSlice<u8>,
13419 b2: &CudaSlice<u8>,
13420 aq: &CudaSlice<i8>,
13421 ad: &CudaSlice<f32>,
13422 in_f: usize,
13423 out0: usize,
13424 out1: usize,
13425 out2: usize,
13426 row_bytes: usize,
13427 ws0: f32,
13428 ws1: f32,
13429 ws2: f32,
13430 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13431 const ROWS_PER_BLOCK: u32 = 4;
13432 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13433 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13434 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13435 let f = self.func("qmatvec_e4m3_mmvq_fused3");
13436 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13437 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13438 let mut y2 = self.alloc_uninit::<f32>(out2)?;
13439 let cfg = LaunchConfig {
13440 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13441 block_dim: (32, ROWS_PER_BLOCK, 1),
13442 shared_mem_bytes: 0,
13443 };
13444 let (inf, o0, o1, o2, rbl) = (
13445 in_f as i32,
13446 out0 as i32,
13447 out1 as i32,
13448 out2 as i32,
13449 row_bytes as i64,
13450 );
13451 let __s_b = self.gpu.stream();
13452 let mut b = __s_b.launch_builder(&f);
13453 b.arg(b0)
13454 .arg(b1)
13455 .arg(b2)
13456 .arg(aq)
13457 .arg(ad)
13458 .arg(&mut y0)
13459 .arg(&mut y1)
13460 .arg(&mut y2)
13461 .arg(&inf)
13462 .arg(&o0)
13463 .arg(&o1)
13464 .arg(&o2)
13465 .arg(&rbl)
13466 .arg(&ws0)
13467 .arg(&ws1)
13468 .arg(&ws2);
13469 unsafe {
13470 b.launch(cfg)?;
13471 }
13472 Ok((y0, y1, y2))
13473 }
13474
13475 #[allow(clippy::too_many_arguments)]
13479 fn e4m3_fused2_t_core(
13480 &self,
13481 b0: &CudaSlice<u8>,
13482 b1: &CudaSlice<u8>,
13483 aq: &CudaSlice<i8>,
13484 ad: &CudaSlice<f32>,
13485 m: usize,
13486 in_f: usize,
13487 out0: usize,
13488 out1: usize,
13489 row_bytes: usize,
13490 ws0: f32,
13491 ws1: f32,
13492 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13493 const ROWS_PER_BLOCK: u32 = 4;
13494 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13495 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13496 let f = self.func(match Self::batched_mcols(m) {
13497 2 => "qmatvec_e4m3_mmvq_fused2_b2",
13498 4 => "qmatvec_e4m3_mmvq_fused2_b4",
13499 _ => "qmatvec_e4m3_mmvq_fused2_b8",
13500 });
13501 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13502 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13503 let cfg = LaunchConfig {
13504 grid_dim: (nb0 + nb1, 1, 1),
13505 block_dim: (32, ROWS_PER_BLOCK, 1),
13506 shared_mem_bytes: 0,
13507 };
13508 let (inf, o0, o1, mi, rbl) = (
13509 in_f as i32,
13510 out0 as i32,
13511 out1 as i32,
13512 m as i32,
13513 row_bytes as i64,
13514 );
13515 let __s_b = self.gpu.stream();
13516 let mut b = __s_b.launch_builder(&f);
13517 b.arg(b0)
13518 .arg(b1)
13519 .arg(aq)
13520 .arg(ad)
13521 .arg(&mut y0)
13522 .arg(&mut y1)
13523 .arg(&inf)
13524 .arg(&o0)
13525 .arg(&o1)
13526 .arg(&mi)
13527 .arg(&rbl);
13528 unsafe {
13529 b.launch(cfg)?;
13530 }
13531 if ws0 != 1.0 {
13532 self.scale_inplace(&mut y0, ws0, m * out0)?;
13533 }
13534 if ws1 != 1.0 {
13535 self.scale_inplace(&mut y1, ws1, m * out1)?;
13536 }
13537 Ok((y0, y1))
13538 }
13539
13540 #[allow(clippy::too_many_arguments)]
13542 fn e4m3_fused3_t_core(
13543 &self,
13544 b0: &CudaSlice<u8>,
13545 b1: &CudaSlice<u8>,
13546 b2: &CudaSlice<u8>,
13547 aq: &CudaSlice<i8>,
13548 ad: &CudaSlice<f32>,
13549 m: usize,
13550 in_f: usize,
13551 out0: usize,
13552 out1: usize,
13553 out2: usize,
13554 row_bytes: usize,
13555 ws0: f32,
13556 ws1: f32,
13557 ws2: f32,
13558 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13559 const ROWS_PER_BLOCK: u32 = 4;
13560 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13561 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13562 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13563 let f = self.func(if Self::batched_mcols(m) == 2 {
13564 "qmatvec_e4m3_mmvq_fused3_b2"
13565 } else {
13566 "qmatvec_e4m3_mmvq_fused3_b4"
13567 });
13568 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13569 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13570 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13571 let cfg = LaunchConfig {
13572 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13573 block_dim: (32, ROWS_PER_BLOCK, 1),
13574 shared_mem_bytes: 0,
13575 };
13576 let (inf, o0, o1, o2, mi, rbl) = (
13577 in_f as i32,
13578 out0 as i32,
13579 out1 as i32,
13580 out2 as i32,
13581 m as i32,
13582 row_bytes as i64,
13583 );
13584 let __s_b = self.gpu.stream();
13585 let mut b = __s_b.launch_builder(&f);
13586 b.arg(b0)
13587 .arg(b1)
13588 .arg(b2)
13589 .arg(aq)
13590 .arg(ad)
13591 .arg(&mut y0)
13592 .arg(&mut y1)
13593 .arg(&mut y2)
13594 .arg(&inf)
13595 .arg(&o0)
13596 .arg(&o1)
13597 .arg(&o2)
13598 .arg(&mi)
13599 .arg(&rbl);
13600 unsafe {
13601 b.launch(cfg)?;
13602 }
13603 if ws0 != 1.0 {
13604 self.scale_inplace(&mut y0, ws0, m * out0)?;
13605 }
13606 if ws1 != 1.0 {
13607 self.scale_inplace(&mut y1, ws1, m * out1)?;
13608 }
13609 if ws2 != 1.0 {
13610 self.scale_inplace(&mut y2, ws2, m * out2)?;
13611 }
13612 Ok((y0, y1, y2))
13613 }
13614
13615 pub fn qmatvec_e4m3_blk_mmvq(
13625 &self,
13626 bytes: &CudaSlice<u8>,
13627 aq: &CudaSlice<i8>,
13628 ad: &CudaSlice<f32>,
13629 scales: &CudaSlice<f32>,
13630 m: usize,
13631 in_f: usize,
13632 out_f: usize,
13633 row_bytes: usize,
13634 scale_cols: usize,
13635 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13636 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
13638 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
13639 )?;
13640 Ok(y)
13641 }
13642
13643 #[allow(clippy::too_many_arguments)]
13645 pub fn qmatvec_e4m3_blk_mmvq_into(
13646 &self,
13647 bytes: &CudaSlice<u8>,
13648 aq: &CudaSlice<i8>,
13649 ad: &CudaSlice<f32>,
13650 scales: &CudaSlice<f32>,
13651 m: usize,
13652 in_f: usize,
13653 out_f: usize,
13654 row_bytes: usize,
13655 scale_cols: usize,
13656 y: &mut CudaSlice<f32>,
13657 ) -> Result<(), Box<dyn std::error::Error>> {
13658 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
13660 let cfg = LaunchConfig {
13661 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
13662 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
13665 let (inf, outf, mi, rb, sc) = (
13666 in_f as i32,
13667 out_f as i32,
13668 m as i32,
13669 row_bytes as i64,
13670 scale_cols as i32,
13671 );
13672 let __s_b = self.gpu.stream();
13673 let mut b = __s_b.launch_builder(&f);
13674 b.arg(bytes)
13675 .arg(aq)
13676 .arg(ad)
13677 .arg(scales)
13678 .arg(&mut *y)
13679 .arg(&inf)
13680 .arg(&outf)
13681 .arg(&mi)
13682 .arg(&rb)
13683 .arg(&sc);
13684 unsafe {
13685 b.launch(cfg)?;
13686 }
13687 Ok(())
13688 }
13689
13690 #[allow(clippy::too_many_arguments)]
13696 pub fn qmatvec_e4m3_blk_mmvq_batched(
13697 &self,
13698 bytes: &CudaSlice<u8>,
13699 aq: &CudaSlice<i8>,
13700 ad: &CudaSlice<f32>,
13701 scales: &CudaSlice<f32>,
13702 m: usize,
13703 in_f: usize,
13704 out_f: usize,
13705 row_bytes: usize,
13706 scale_cols: usize,
13707 mcols: usize,
13708 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13709 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
13711 let name = match mcols {
13712 2 => "qmatvec_e4m3_blk_mmvq_b2",
13713 4 => "qmatvec_e4m3_blk_mmvq_b4",
13714 8 => "qmatvec_e4m3_blk_mmvq_b8",
13715 16 => "qmatvec_e4m3_blk_mmvq_b16",
13716 _ => {
13717 return Err(
13718 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
13719 );
13720 }
13721 };
13722 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13723 let f = self.func(name);
13724 let cfg = LaunchConfig {
13725 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
13726 block_dim: (32, ROWS_PER_BLOCK, 1),
13727 shared_mem_bytes: 0,
13728 };
13729 let (inf, outf, mi, rb, sc) = (
13730 in_f as i32,
13731 out_f as i32,
13732 m as i32,
13733 row_bytes as i64,
13734 scale_cols as i32,
13735 );
13736 let __s_b = self.gpu.stream();
13737 let mut b = __s_b.launch_builder(&f);
13738 b.arg(bytes)
13739 .arg(aq)
13740 .arg(ad)
13741 .arg(scales)
13742 .arg(&mut y)
13743 .arg(&inf)
13744 .arg(&outf)
13745 .arg(&mi)
13746 .arg(&rb)
13747 .arg(&sc);
13748 unsafe {
13749 b.launch(cfg)?;
13750 }
13751 Ok(y)
13752 }
13753
13754 #[allow(clippy::too_many_arguments)]
13757 pub fn qmatvec_e4m3_blk_batched_raw(
13758 &self,
13759 bytes: &CudaSlice<u8>,
13760 x: &CudaSlice<f32>,
13761 scales: &CudaSlice<f32>,
13762 m: usize,
13763 in_f: usize,
13764 out_f: usize,
13765 row_bytes: usize,
13766 scale_cols: usize,
13767 mcols: usize,
13768 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13769 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13770 self.qmatvec_e4m3_blk_mmvq_batched(
13771 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13772 )
13773 }
13774
13775 #[allow(clippy::too_many_arguments)]
13778 pub fn qmatvec_e4m3_blk_mmvq_raw(
13779 &self,
13780 bytes: &CudaSlice<u8>,
13781 x: &CudaSlice<f32>,
13782 scales: &CudaSlice<f32>,
13783 m: usize,
13784 in_f: usize,
13785 out_f: usize,
13786 row_bytes: usize,
13787 scale_cols: usize,
13788 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13789 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13790 self.qmatvec_e4m3_blk_mmvq(
13791 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13792 )
13793 }
13794
13795 #[allow(clippy::too_many_arguments)]
13798 pub fn qmatvec_e4m3_fused2_raw(
13799 &self,
13800 b0: &CudaSlice<u8>,
13801 b1: &CudaSlice<u8>,
13802 x: &CudaSlice<f32>,
13803 in_f: usize,
13804 out0: usize,
13805 out1: usize,
13806 row_bytes: usize,
13807 ws0: f32,
13808 ws1: f32,
13809 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13810 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13811 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13812 }
13813
13814 #[allow(clippy::too_many_arguments)]
13815 pub fn qmatvec_e4m3_fused3_raw(
13816 &self,
13817 b0: &CudaSlice<u8>,
13818 b1: &CudaSlice<u8>,
13819 b2: &CudaSlice<u8>,
13820 x: &CudaSlice<f32>,
13821 in_f: usize,
13822 out0: usize,
13823 out1: usize,
13824 out2: usize,
13825 row_bytes: usize,
13826 ws0: f32,
13827 ws1: f32,
13828 ws2: f32,
13829 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13830 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13831 self.e4m3_fused3_core(
13832 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13833 )
13834 }
13835
13836 #[allow(clippy::too_many_arguments)]
13837 pub fn qmatvec_e4m3_fused2_t_raw(
13838 &self,
13839 b0: &CudaSlice<u8>,
13840 b1: &CudaSlice<u8>,
13841 x: &CudaSlice<f32>,
13842 m: usize,
13843 in_f: usize,
13844 out0: usize,
13845 out1: usize,
13846 row_bytes: usize,
13847 ws0: f32,
13848 ws1: f32,
13849 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13850 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13851 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13852 }
13853
13854 #[allow(clippy::too_many_arguments)]
13855 pub fn qmatvec_e4m3_fused3_t_raw(
13856 &self,
13857 b0: &CudaSlice<u8>,
13858 b1: &CudaSlice<u8>,
13859 b2: &CudaSlice<u8>,
13860 x: &CudaSlice<f32>,
13861 m: usize,
13862 in_f: usize,
13863 out0: usize,
13864 out1: usize,
13865 out2: usize,
13866 row_bytes: usize,
13867 ws0: f32,
13868 ws1: f32,
13869 ws2: f32,
13870 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13871 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13872 self.e4m3_fused3_t_core(
13873 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13874 )
13875 }
13876
13877 fn try_e4m3_blk_pre(
13888 &self,
13889 w: &crate::model::GpuTensor,
13890 aq: &CudaSlice<i8>,
13891 ad: &CudaSlice<f32>,
13892 m: usize,
13893 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13894 use crate::model::GpuTensor;
13895 if let GpuTensor::Quant {
13896 bytes,
13897 qtype,
13898 row_bytes,
13899 blk: Some(g),
13900 ..
13901 } = w
13902 {
13903 if *qtype == QT_F8_E4M3_BLK {
13904 if (2..=16).contains(&m)
13910 && std::env::var("MEMRA_NO_BATCHED").is_err()
13911 && (m <= 4 || Self::b8_enabled())
13912 {
13913 let mcols = Self::batched_mcols(m);
13914 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13915 bytes,
13916 aq,
13917 ad,
13918 &g.scales,
13919 m,
13920 w.in_features(),
13921 w.out_features(),
13922 *row_bytes,
13923 g.cols,
13924 mcols,
13925 )?));
13926 }
13927 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13928 bytes,
13929 aq,
13930 ad,
13931 &g.scales,
13932 m,
13933 w.in_features(),
13934 w.out_features(),
13935 *row_bytes,
13936 g.cols,
13937 )?));
13938 }
13939 }
13940 Ok(None)
13941 }
13942
13943 fn try_e4m3_blk_prefill(
13990 &self,
13991 w: &crate::model::GpuTensor,
13992 x: &CudaSlice<f32>,
13993 m: usize,
13994 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13995 use crate::model::GpuTensor;
13996 let GpuTensor::Quant {
13997 bytes,
13998 qtype,
13999 blk: Some(g),
14000 ..
14001 } = w
14002 else {
14003 return Ok(None);
14004 };
14005 if *qtype != QT_F8_E4M3_BLK {
14006 return Ok(None);
14007 }
14008 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
14013 return Ok(Some(y));
14014 }
14015 let (in_f, out_f) = (w.in_features(), w.out_features());
14016 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
14017 let tmp = GpuTensor::Quant {
14018 bytes: slab,
14019 qtype: QT_Q8_0,
14020 row_bytes: in_f / 32 * 34,
14021 ne: vec![in_f as u64, out_f as u64],
14022 scale: 1.0,
14023 rp: false,
14024 #[cfg(memra_cutlass)]
14025 cutlass: None,
14026 fp8: None,
14027 blk: None,
14028 f16: None,
14029 rp4: None,
14030 };
14031 Ok(Some(self.matmul(&tmp, x, m)?))
14033 }
14034
14035 pub fn matmul_pre_noscale(
14036 &self,
14037 w: &crate::model::GpuTensor,
14038 aq: &CudaSlice<i8>,
14039 ad: &CudaSlice<f32>,
14040 m: usize,
14041 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
14042 use crate::model::GpuTensor;
14043 if m == 1 {
14047 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
14048 return Ok(Some((y, 1.0)));
14049 }
14050 }
14051 if m != 1 || !self.uses_q8_1_fast(w) {
14053 return Ok(None);
14054 }
14055 let in_f = w.in_features();
14056 let out_f = w.out_features();
14057 let (bytes, qtype, row_bytes, scale, rp) = match w {
14058 GpuTensor::Quant {
14059 bytes,
14060 qtype,
14061 row_bytes,
14062 scale,
14063 rp,
14064 ..
14065 } => (bytes, *qtype, *row_bytes, *scale, *rp),
14066 _ => return Ok(None),
14067 };
14068 if self.mmvq_supports(qtype) {
14070 let (mbytes, mrp) = match w {
14072 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
14073 _ => (bytes, rp),
14074 };
14075 let y = self.qmatvec_mmvq(
14076 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
14077 )?;
14078 return Ok(Some((y, scale)));
14079 }
14080 let name = match qtype {
14082 QT_Q8_0 => "qmatvec_q8_0_dp4a",
14083 QT_Q4_K => "qmatvec_q4_K_dp4a",
14084 QT_Q6_K => "qmatvec_q6_K_dp4a",
14085 QT_Q5_K => "qmatvec_q5_K_dp4a",
14086 QT_Q3_K => "qmatvec_q3_K_dp4a",
14087 QT_NVFP4 => {
14088 if rp {
14089 "qmatvec_nvfp4_dp4a_rp"
14090 } else {
14091 "qmatvec_nvfp4_dp4a"
14092 }
14093 }
14094 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
14095 _ => return Ok(None),
14096 };
14097 let f = self.func(name);
14098 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14099 let cfg = LaunchConfig {
14100 grid_dim: (out_f as u32, m as u32, 1),
14101 block_dim: (128, 1, 1),
14102 shared_mem_bytes: 0,
14103 };
14104 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14105 let __s_b = self.gpu.stream();
14106 let mut b = __s_b.launch_builder(&f);
14107 b.arg(bytes)
14108 .arg(aq)
14109 .arg(ad)
14110 .arg(&mut y)
14111 .arg(&inf)
14112 .arg(&outf)
14113 .arg(&mi)
14114 .arg(&rb);
14115 unsafe {
14116 b.launch(cfg)?;
14117 }
14118 Ok(Some((y, scale)))
14119 }
14120
14121 pub fn mmvq_supports(&self, qtype: i32) -> bool {
14124 if qtype == QT_F8_E4M3 {
14129 return true;
14130 }
14131 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14132 return false;
14133 }
14134 matches!(
14135 qtype,
14136 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
14137 )
14138 }
14139
14140 pub fn qmatvec_mmvq(
14145 &self,
14146 bytes: &CudaSlice<u8>,
14147 aq: &CudaSlice<i8>,
14148 ad: &CudaSlice<f32>,
14149 m: usize,
14150 in_f: usize,
14151 out_f: usize,
14152 qtype: i32,
14153 row_bytes: usize,
14154 scale: f32,
14155 rp: bool,
14156 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14157 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
14159 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
14160 )?;
14161 Ok(y)
14162 }
14163
14164 #[allow(clippy::too_many_arguments)]
14166 pub fn qmatvec_mmvq_into(
14167 &self,
14168 bytes: &CudaSlice<u8>,
14169 aq: &CudaSlice<i8>,
14170 ad: &CudaSlice<f32>,
14171 m: usize,
14172 in_f: usize,
14173 out_f: usize,
14174 qtype: i32,
14175 row_bytes: usize,
14176 scale: f32,
14177 rp: bool,
14178 y: &mut CudaSlice<f32>,
14179 ) -> Result<(), Box<dyn std::error::Error>> {
14180 debug_assert!(y.len() >= m * out_f);
14181 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
14187 && rp
14188 && m == 1
14189 && out_f >= 64
14190 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
14191 && {
14192 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14193 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
14194 }
14195 {
14196 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
14197 let cfg = LaunchConfig {
14198 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
14199 block_dim: (32, 2, 1),
14200 shared_mem_bytes: 0,
14201 };
14202 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
14203 let __s_b = self.gpu.stream();
14204 let mut b = __s_b.launch_builder(&f);
14205 b.arg(bytes)
14206 .arg(aq)
14207 .arg(ad)
14208 .arg(&mut *y)
14209 .arg(&inf)
14210 .arg(&outf)
14211 .arg(&mi)
14212 .arg(&rb);
14213 unsafe {
14214 b.launch(cfg)?;
14215 }
14216 if scale != 1.0 {
14217 self.scale_inplace(y, scale, out_f)?;
14218 }
14219 return Ok(());
14220 }
14221 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
14230 2
14231 } else {
14232 1
14233 };
14234 if m == 1 && qtype == QT_Q4_0 {
14239 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
14240 mr = *Q40MR.get_or_init(|| {
14243 std::env::var("MEMRA_Q40_MR")
14244 .ok()
14245 .and_then(|v| v.parse().ok())
14246 .unwrap_or(1)
14247 });
14248 }
14249 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
14260 let q5_force = q5_mode.as_deref() == Some("2");
14261 let q5_il = qtype == QT_Q5_K
14264 && m == 1
14265 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
14266 if q5_il && !q5_force && out_f > 65536 {
14267 mr = 1;
14268 }
14269 if qtype == QT_Q4_0 && rp && mr != 1 {
14272 mr = 2;
14273 }
14274 if qtype == QT_Q8_0 && rp {
14278 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
14279 mr = *Q80MR.get_or_init(|| {
14280 std::env::var("MEMRA_Q80_MR")
14281 .ok()
14282 .and_then(|v| v.parse().ok())
14283 .unwrap_or(1)
14284 });
14285 }
14286 let name = match (qtype, mr, rp) {
14287 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
14288 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
14289 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
14290 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
14291 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
14292 (QT_Q5_K, 2, _) => {
14293 if q5_il {
14294 "qmatvec_q5_K_mmvq_mr2_il"
14295 } else {
14296 "qmatvec_q5_K_mmvq_mr2"
14297 }
14298 }
14299 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
14300 (QT_Q8_0, _, true)
14305 if in_f % 1024 == 0 && {
14306 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14307 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
14308 } =>
14309 {
14310 "qmatvec_q8_0_mmvq_rpca"
14311 }
14312 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
14313 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
14314 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
14318 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
14319 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
14320 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
14321 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
14322 (QT_Q5_K, _, _) => {
14323 if q5_il {
14324 "qmatvec_q5_K_mmvq_il"
14325 } else {
14326 "qmatvec_q5_K_mmvq"
14327 }
14328 }
14329 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
14330 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
14331 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
14332 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
14333 };
14334 let f = self.func(name);
14335 let rows_per_block = ROWS_PER_BLOCK * mr;
14337 let cfg = LaunchConfig {
14338 grid_dim: (
14339 (out_f as u32 + rows_per_block - 1) / rows_per_block,
14340 m as u32,
14341 1,
14342 ),
14343 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
14346 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14347 let __s_b = self.gpu.stream();
14348 let mut b = __s_b.launch_builder(&f);
14349 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
14354 if Self::pdl_on()
14357 && Self::pdl_mmvq_on()
14358 && Self::pdl_nvfp4q8_on()
14359 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
14360 {
14361 use cudarc::driver::{DevicePtr, DevicePtrMut};
14362 let s = &self.gpu.stream();
14363 let (pw, _g0) = bytes.device_ptr(s);
14364 let (paq, _g1) = aq.device_ptr(s);
14365 let (pad, _g2) = ad.device_ptr(s);
14366 let (py, _g3) = y.device_ptr_mut(s);
14367 let mut ps = [
14368 &pw as *const _ as *mut std::ffi::c_void,
14369 &paq as *const _ as *mut _,
14370 &pad as *const _ as *mut _,
14371 &py as *const _ as *mut _,
14372 &inf as *const _ as *mut _,
14373 &outf as *const _ as *mut _,
14374 &mi as *const _ as *mut _,
14375 &rb as *const _ as *mut _,
14376 &scale as *const _ as *mut _,
14377 ];
14378 unsafe {
14379 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14380 }
14381 return Ok(());
14382 }
14383 b.arg(bytes)
14384 .arg(aq)
14385 .arg(ad)
14386 .arg(&mut *y)
14387 .arg(&inf)
14388 .arg(&outf)
14389 .arg(&mi)
14390 .arg(&rb)
14391 .arg(&scale);
14392 unsafe {
14393 b.launch(cfg)?;
14394 }
14395 } else if Self::pdl_on()
14396 && Self::pdl_mmvq_on()
14397 && (matches!(
14398 name,
14399 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
14400 ) || (Self::pdl_nvfp4q8_on()
14401 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
14402 {
14403 {
14407 use cudarc::driver::{DevicePtr, DevicePtrMut};
14408 let s = &self.gpu.stream();
14409 let (pw, _g0) = bytes.device_ptr(s);
14410 let (paq, _g1) = aq.device_ptr(s);
14411 let (pad, _g2) = ad.device_ptr(s);
14412 let (py, _g3) = y.device_ptr_mut(s);
14413 let mut ps = [
14414 &pw as *const _ as *mut std::ffi::c_void,
14415 &paq as *const _ as *mut _,
14416 &pad as *const _ as *mut _,
14417 &py as *const _ as *mut _,
14418 &inf as *const _ as *mut _,
14419 &outf as *const _ as *mut _,
14420 &mi as *const _ as *mut _,
14421 &rb as *const _ as *mut _,
14422 ];
14423 unsafe {
14424 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14425 }
14426 }
14427 if scale != 1.0 {
14428 self.scale_inplace(y, scale, m * out_f)?;
14429 }
14430 } else {
14431 b.arg(bytes)
14432 .arg(aq)
14433 .arg(ad)
14434 .arg(&mut *y)
14435 .arg(&inf)
14436 .arg(&outf)
14437 .arg(&mi)
14438 .arg(&rb);
14439 unsafe {
14440 b.launch(cfg)?;
14441 }
14442 if scale != 1.0 {
14443 self.scale_inplace(y, scale, m * out_f)?;
14444 }
14445 }
14446 Ok(())
14447 }
14448
14449 pub fn qmatvec_mmvq_raw(
14453 &self,
14454 bytes: &CudaSlice<u8>,
14455 x: &CudaSlice<f32>,
14456 m: usize,
14457 in_f: usize,
14458 out_f: usize,
14459 qtype: i32,
14460 row_bytes: usize,
14461 rp: bool,
14462 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14463 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14464 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
14465 }
14466
14467 pub fn batched_supports(&self, qtype: i32) -> bool {
14471 matches!(
14472 qtype,
14473 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
14474 )
14475 }
14476
14477 pub fn iq_fast_enabled() -> bool {
14485 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14486 *ON.get_or_init(|| {
14487 std::env::var("MEMRA_IQ_FAST")
14488 .map(|v| v != "0")
14489 .unwrap_or(true)
14490 })
14491 }
14492
14493 pub fn b8_enabled() -> bool {
14496 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14497 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
14498 }
14499
14500 pub fn batched_mcols(m: usize) -> usize {
14502 if m == 2 {
14503 2
14504 } else if m <= 4 {
14505 4
14506 } else if m <= 8 {
14507 8
14508 } else {
14509 16
14510 }
14511 }
14512
14513 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
14518 Some(match (qtype, mcols) {
14519 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
14520 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
14521 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
14522 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
14528 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
14529 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
14530 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
14531 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
14534 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
14535 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
14536 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
14537 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
14540 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
14541 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
14542 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
14543 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
14544 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
14545 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
14546 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
14547 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
14551 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
14552 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
14553 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
14554 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
14558 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
14559 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
14560 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
14561 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
14562 _ => return None,
14563 })
14564 }
14565
14566 pub fn sm_count(&self) -> i32 {
14601 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14602 *SMS.get_or_init(|| {
14603 use cudarc::driver::sys::CUdevice_attribute_enum as A;
14604 self.gpu
14605 .ctx
14606 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14607 .unwrap_or(82)
14608 })
14609 }
14610
14611 pub fn batched_variant(
14612 &self,
14613 _m: usize,
14614 in_f: usize,
14615 out_f: usize,
14616 qtype: i32,
14617 row_bytes: usize,
14618 mcols: usize,
14619 rp: bool,
14620 ) -> &'static str {
14621 if qtype == QT_Q8_0 {
14626 return if rp { "rp" } else { "base" };
14627 }
14628 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14629 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
14630 Ok("base") => "base",
14631 Ok("pf") => "pf",
14632 Ok("r2") => "r2",
14633 Ok("r2w8") => "r2w8",
14634 Ok("pfr2") => "pfr2",
14635 Ok("ca") => "ca",
14636 Ok("car2") => "car2",
14637 Ok("rp") => "rp",
14640 Ok("rpr2") => "rpr2",
14641 Ok("rpr2w8") => "rpr2w8",
14642 Ok("rpca") => "rpca",
14645 Ok("rpcar2") => "rpcar2",
14646 Ok("rpsc") => "rpsc",
14653 Ok("rpms") => "rpms",
14654 Ok("rpmsc") => "rpmsc",
14655 Ok("rpks") => "rpks",
14656 Ok("rpksc") => "rpksc",
14657 _ => "auto",
14658 });
14659 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
14663 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14668 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
14669 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
14670 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
14671 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14672 let sms = *SMS.get_or_init(|| {
14673 use cudarc::driver::sys::CUdevice_attribute_enum as A;
14674 self.gpu
14675 .ctx
14676 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14677 .unwrap_or(82)
14678 });
14679 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
14699 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14702 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
14703 Ok("base") => "base",
14704 Ok("r2") => "r2",
14705 Ok("r2w8") => "r2w8",
14706 _ => "auto",
14707 });
14708 let variant: &'static str = if qtype == QT_Q4_0 {
14709 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14713 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
14714 Ok("base") => "base",
14720 Ok("r2") => "r2",
14721 Ok("ms") => "ms",
14722 Ok("sm") => "sm",
14723 Ok("la") => "la",
14724 _ => "auto",
14725 });
14726 let v = if q40 != "auto" {
14727 q40
14728 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
14729 "r2"
14730 } else {
14731 "base"
14732 };
14733 if rp {
14738 match v {
14739 "ms" => "r2ms_rp",
14740 "sm" => "r2sm_rp",
14741 "la" => "r2la_rp",
14742 "r2" => "r2_rp",
14743 _ => "rp",
14744 }
14745 } else if matches!(v, "ms" | "sm" | "la") {
14746 "r2"
14747 } else {
14748 v
14749 }
14750 } else if qtype != QT_NVFP4 && !kq_r2 {
14751 "base"
14752 } else if kq_r2 && rp {
14753 "rp"
14757 } else if kq_r2 {
14758 if kq_bv != "auto" {
14761 if kq_bv == "r2w8" && mcols != 4 {
14762 "r2"
14763 } else {
14764 kq_bv
14765 }
14766 } else if bv != "auto" {
14767 match bv {
14768 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14769 "r2w8" | "rpr2w8" => {
14770 if mcols != 4 {
14771 "r2"
14772 } else {
14773 "r2w8"
14774 }
14775 }
14776 _ => "base", }
14778 } else {
14779 let blocks = (out_f + 7) / 8;
14780 let waves = blocks as f64 / (7 * sms as usize) as f64;
14781 let filled = blocks >= 4 * sms as usize;
14782 let use_r2 = if qtype == QT_Q4_K {
14783 filled
14784 } else {
14785 waves >= 2.0
14786 };
14787 if use_r2 { "r2" } else { "base" }
14788 }
14789 } else if bv != "auto" {
14790 let v = if bv == "r2w8" && mcols == 2 {
14795 "r2"
14796 } else if bv == "ca" && (!ca_ok || mcols == 8) {
14797 "pf"
14798 } else if bv == "car2" && (!ca_ok || mcols == 8) {
14799 "r2"
14800 } else if bv == "pfr2" && mcols == 8 {
14801 "r2"
14802 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14803 "rpr2"
14804 }
14805 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14807 if mcols == 8 { "rpr2w8" } else { "rpr2" }
14808 } else if bv == "rpcar2" && mcols == 2 {
14809 "rpca"
14810 }
14811 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14814 "rpr2"
14815 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14816 "rpr2"
14817 } else {
14818 bv
14819 };
14820 if rp {
14821 match v {
14822 "base" | "pf" | "ca" | "rp" => "rp",
14823 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14824 "r2w8" | "rpr2w8" => {
14825 if mcols == 2 {
14826 "rpr2"
14827 } else {
14828 "rpr2w8"
14829 }
14830 }
14831 other => other, }
14833 } else {
14834 v
14835 }
14836 } else if mcols == 8 {
14837 if rp {
14848 if sc_ok { "rpsc" } else { "rpr2w8" }
14849 } else {
14850 "r2w8"
14851 }
14852 } else if mcols >= 4 {
14853 let blocks = (out_f + 7) / 8;
14857 let r7 = 7 * sms as usize;
14858 let r8 = 8 * sms as usize;
14859 let waves = blocks as f64 / r7 as f64;
14860 let filled = blocks >= 4 * sms as usize;
14861 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14865 if rp { "rpr2w8" } else { "r2w8" }
14869 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14870 if rp { "rpr2" } else { "r2" }
14873 } else {
14874 if rp { "rp" } else { "pf" }
14878 }
14879 } else if in_f >= 6144 {
14880 if rp { "rpr2" } else { "r2" }
14884 } else if rp {
14885 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14890 if sc_ok && waves >= 0.9 && waves <= 1.1 {
14891 "rpsc"
14892 } else {
14893 "rp"
14894 }
14895 } else {
14896 "base"
14897 };
14898 variant
14899 }
14900
14901 pub fn qmatvec_mmvq_batched(
14902 &self,
14903 bytes: &CudaSlice<u8>,
14904 aq: &CudaSlice<i8>,
14905 ad: &CudaSlice<f32>,
14906 m: usize,
14907 in_f: usize,
14908 out_f: usize,
14909 qtype: i32,
14910 row_bytes: usize,
14911 mcols: usize,
14912 scale: f32,
14913 rp: bool,
14914 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14915 const ROWS_PER_BLOCK: u32 = 4;
14916 let forced: Option<&'static str> = {
14921 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14922 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14923 .as_deref()
14924 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14925 };
14926 let variant = match forced {
14927 Some(v) if !rp || v.contains("rp") => v,
14928 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14929 };
14930 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14931 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14932 })?;
14933 let variant = if mcols == 16 {
14937 if rp { "rp" } else { "base" }
14938 } else {
14939 variant
14940 };
14941 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14948 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14949 if b567
14950 && qtype == QT_NVFP4
14951 && rp
14952 && mcols == 8
14953 && (5..=7).contains(&m)
14954 && matches!(variant, "rpsc" | "rpr2w8")
14955 {
14956 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14957 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14959 let cfg = LaunchConfig {
14960 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14961 block_dim: (32, ROWS_PER_BLOCK, 1),
14962 shared_mem_bytes: 0,
14963 };
14964 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14965 let __s_b = self.gpu.stream();
14966 let mut b = __s_b.launch_builder(&f);
14967 b.arg(bytes)
14968 .arg(aq)
14969 .arg(ad)
14970 .arg(&mut y)
14971 .arg(&inf)
14972 .arg(&outf)
14973 .arg(&mi)
14974 .arg(&rb);
14975 unsafe {
14976 b.launch(cfg)?;
14977 }
14978 if scale != 1.0 {
14979 self.scale_inplace(&mut y, scale, m * out_f)?;
14980 }
14981 return Ok(y);
14982 }
14983 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14984 "base" => (base_name.into(), ROWS_PER_BLOCK),
14985 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14986 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14987 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14988 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14992 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14993 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14994 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14995 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14996 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14997 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14998 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
15000 debug_assert!(
15001 !rp || name.contains("_rp"),
15002 "rp weight dispatched to a GGUF-layout kernel"
15003 );
15004 let f = self.func(&name);
15005 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15006 let smem = if name.contains("_r2sm_rp") {
15008 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
15009 } else {
15010 0
15011 };
15012 let cfg = LaunchConfig {
15013 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
15014 block_dim: (32, ROWS_PER_BLOCK, 1),
15015 shared_mem_bytes: smem,
15016 };
15017 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15018 let __s_b = self.gpu.stream();
15019 let mut b = __s_b.launch_builder(&f);
15020 b.arg(bytes)
15021 .arg(aq)
15022 .arg(ad)
15023 .arg(&mut y)
15024 .arg(&inf)
15025 .arg(&outf)
15026 .arg(&mi)
15027 .arg(&rb);
15028 unsafe {
15029 b.launch(cfg)?;
15030 }
15031 if scale != 1.0 {
15032 self.scale_inplace(&mut y, scale, m * out_f)?;
15033 }
15034 Ok(y)
15035 }
15036
15037 pub fn qmatvec_batched_raw(
15041 &self,
15042 bytes: &CudaSlice<u8>,
15043 x: &CudaSlice<f32>,
15044 m: usize,
15045 in_f: usize,
15046 out_f: usize,
15047 qtype: i32,
15048 row_bytes: usize,
15049 mcols: usize,
15050 rp: bool,
15051 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15052 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15053 self.qmatvec_mmvq_batched(
15054 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
15055 )
15056 }
15057
15058 pub fn qmatvec_nvfp4_batched_raw(
15060 &self,
15061 bytes: &CudaSlice<u8>,
15062 x: &CudaSlice<f32>,
15063 m: usize,
15064 in_f: usize,
15065 out_f: usize,
15066 row_bytes: usize,
15067 mcols: usize,
15068 rp: bool,
15069 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15070 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
15071 }
15072
15073 fn try_fp4_gemm(
15077 &self,
15078 w: &crate::model::GpuTensor,
15079 x: &CudaSlice<f32>,
15080 m: usize,
15081 in_f: usize,
15082 out_f: usize,
15083 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15084 use crate::model::GpuTensor;
15085 if cfg!(memra_portable_cuda) {
15086 return Ok(None);
15087 }
15088 if std::env::var("MEMRA_FP4").is_err() {
15089 return Ok(None);
15090 }
15091 #[cfg(memra_cutlass)]
15100 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
15101 if let GpuTensor::Quant {
15102 bytes,
15103 qtype,
15104 scale,
15105 row_bytes,
15106 cutlass,
15107 ..
15108 } = w
15109 {
15110 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
15111 if let Some(cw) = cutlass {
15112 let y = self.cutlass_fp4_gemm(
15114 &cw.b_packed,
15115 &cw.sfb_swizzled,
15116 x,
15117 *scale,
15118 m,
15119 out_f,
15120 in_f,
15121 )?;
15122 return Ok(Some(y));
15123 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
15124 let (b_packed, sfb_sw) =
15129 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
15130 let y =
15131 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
15132 return Ok(Some(y));
15133 }
15134 }
15135 }
15136 }
15137 if let GpuTensor::Quant {
15138 bytes,
15139 qtype,
15140 row_bytes,
15141 scale,
15142 rp,
15143 ..
15144 } = w
15145 {
15146 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
15149 let y =
15150 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
15151 return Ok(Some(y));
15152 }
15153 }
15154 Ok(None)
15155 }
15156
15157 pub fn rms_norm_f16out(
15161 &self,
15162 x: &CudaSlice<f32>,
15163 w: &CudaSlice<f32>,
15164 dst: &mut CudaSlice<f32>,
15165 dst16: &mut CudaSlice<u8>,
15166 ncols: usize,
15167 nrows: usize,
15168 eps: f32,
15169 ) -> Result<(), Box<dyn std::error::Error>> {
15170 let f = self.func("rms_norm_f16out_f32");
15171 let cfg = LaunchConfig {
15172 grid_dim: (nrows as u32, 1, 1),
15173 block_dim: (rms_block(), 1, 1),
15174 shared_mem_bytes: 0,
15175 };
15176 let (nc, e) = (ncols as i32, eps);
15177 let __s_b = self.gpu.stream();
15178 let mut b = __s_b.launch_builder(&f);
15179 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
15180 unsafe {
15181 b.launch(cfg)?;
15182 }
15183 Ok(())
15184 }
15185
15186 #[allow(clippy::too_many_arguments)]
15189 pub fn add_rms_norm_f16out(
15190 &self,
15191 a: &CudaSlice<f32>,
15192 b: &CudaSlice<f32>,
15193 w: &CudaSlice<f32>,
15194 res: &mut CudaSlice<f32>,
15195 dst: &mut CudaSlice<f32>,
15196 dst16: &mut CudaSlice<u8>,
15197 ncols: usize,
15198 nrows: usize,
15199 eps: f32,
15200 ) -> Result<(), Box<dyn std::error::Error>> {
15201 let f = self.func("add_rms_norm_f16out_f32");
15202 let cfg = LaunchConfig {
15203 grid_dim: (nrows as u32, 1, 1),
15204 block_dim: (rms_block(), 1, 1),
15205 shared_mem_bytes: 0,
15206 };
15207 let (nc, e) = (ncols as i32, eps);
15208 let __s_lb = self.gpu.stream();
15209 let mut lb = __s_lb.launch_builder(&f);
15210 lb.arg(a)
15211 .arg(b)
15212 .arg(w)
15213 .arg(res)
15214 .arg(dst)
15215 .arg(dst16)
15216 .arg(&nc)
15217 .arg(&e);
15218 unsafe {
15219 lb.launch(cfg)?;
15220 }
15221 Ok(())
15222 }
15223
15224 pub fn matmul_group_xh(
15227 &self,
15228 ws: &[&crate::model::GpuTensor],
15229 x: &CudaSlice<f32>,
15230 xh: &CudaSlice<u8>,
15231 m: usize,
15232 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15233 let mut out = Vec::with_capacity(ws.len());
15234 let in_f = ws[0].in_features();
15235 for w in ws {
15236 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
15237 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
15238 out.push(y);
15239 continue;
15240 }
15241 }
15242 out.push(self.matmul(w, x, m)?);
15243 }
15244 Ok(out)
15245 }
15246
15247 pub fn gdn_pad_mask(
15250 &self,
15251 beta: &mut CudaSlice<f32>,
15252 g_log: &mut CudaSlice<f32>,
15253 len_d: &CudaSlice<i32>,
15254 h: usize,
15255 t: usize,
15256 ) -> Result<(), Box<dyn std::error::Error>> {
15257 let f = self.func("gdn_pad_mask_f32");
15258 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
15259 let (hi, ti) = (h as i32, t as i32);
15260 let __s_b = self.gpu.stream();
15261 let mut b = __s_b.launch_builder(&f);
15262 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
15263 unsafe {
15264 b.launch(cfg)?;
15265 }
15266 Ok(())
15267 }
15268
15269 pub fn row_gather_dev(
15272 &self,
15273 src: &CudaSlice<f32>,
15274 dst: &mut CudaSlice<f32>,
15275 len_d: &CudaSlice<i32>,
15276 ncols: usize,
15277 ) -> Result<(), Box<dyn std::error::Error>> {
15278 let f = self.func("row_gather_dev_f32");
15279 let cfg = LaunchConfig::for_num_elems(ncols as u32);
15280 let nc = ncols as i32;
15281 let __s_b = self.gpu.stream();
15282 let mut b = __s_b.launch_builder(&f);
15283 b.arg(src).arg(dst).arg(len_d).arg(&nc);
15284 unsafe {
15285 b.launch(cfg)?;
15286 }
15287 Ok(())
15288 }
15289
15290 pub fn matmul_group(
15297 &self,
15298 ws: &[&crate::model::GpuTensor],
15299 x: &CudaSlice<f32>,
15300 m: usize,
15301 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15302 use crate::model::GpuTensor;
15303 let mut out = Vec::with_capacity(ws.len());
15304 let any_mirror = ws
15305 .iter()
15306 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
15307 if m >= 16 && any_mirror && !self.verify_exact_on() {
15308 let in_f = ws[0].in_features();
15309 let xh = self.f16_act(x, m * in_f, in_f)?;
15310 for w in ws {
15311 if w.in_features() == in_f {
15312 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
15313 out.push(y);
15314 continue;
15315 }
15316 }
15317 out.push(self.matmul(w, x, m)?);
15318 }
15319 return Ok(out);
15320 }
15321 for w in ws {
15322 out.push(self.matmul(w, x, m)?);
15323 }
15324 Ok(out)
15325 }
15326
15327 pub fn matmul_group_multi(
15334 &self,
15335 ws: &[&crate::model::GpuTensor],
15336 xs: &[&CudaSlice<f32>],
15337 ms: &[usize],
15338 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
15339 assert_eq!(xs.len(), ms.len());
15340 let in_f = ws[0].in_features();
15341 let total: usize = ms.iter().sum();
15342 let mut xcat = self.uninit(total * in_f)?;
15343 let mut off = 0usize;
15344 for (x, &m) in xs.iter().zip(ms) {
15345 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
15346 off += m;
15347 }
15348 let ys = self.matmul_group(ws, &xcat, total)?;
15349 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
15350 for (w, y) in ws.iter().zip(ys) {
15351 let out_f = w.out_features();
15352 let mut off = 0usize;
15353 for (s, &m) in ms.iter().enumerate() {
15354 let mut ys_s = self.uninit(m * out_f)?;
15355 let src = y.slice(off * out_f..(off + m) * out_f);
15356 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
15357 out[s].push(ys_s);
15358 off += m;
15359 }
15360 }
15361 Ok(out)
15362 }
15363
15364 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
15374 use crate::model::GpuTensor;
15375 if !legacy_quant_gemm_allowed(
15376 cfg!(memra_portable_cuda),
15377 cfg!(memra_hopper_mma),
15378 std::env::var_os("MEMRA_NO_GEMM").is_some(),
15379 ) {
15380 return false;
15381 }
15382 match w {
15383 GpuTensor::Quant { qtype, .. } => {
15384 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
15385 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
15386 }
15387 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
15388 }
15389 }
15390
15391 pub fn qmatvec_gemm(
15398 &self,
15399 w: &crate::model::GpuTensor,
15400 aq: &CudaSlice<i8>,
15401 ad: &CudaSlice<f32>,
15402 m: usize,
15403 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15404 use crate::model::GpuTensor;
15405 let in_f = w.in_features();
15406 let out_f = w.out_features();
15407 let (bytes, qtype, row_bytes, scale, rp) = match w {
15408 GpuTensor::Quant {
15409 bytes,
15410 qtype,
15411 row_bytes,
15412 scale,
15413 rp,
15414 ..
15415 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15416 _ => unreachable!("gemm_supports guaranteed Quant"),
15417 };
15418 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
15424 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
15425 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
15426 if scale != 1.0 {
15427 self.scale_inplace(&mut y, scale, m * out_f)?;
15428 }
15429 return Ok(y);
15430 }
15431 }
15432 let name = match qtype {
15433 QT_Q8_0 => "qmatvec_gemm_q8_0",
15434 QT_Q4_K => "qmatvec_gemm_q4_K",
15435 QT_Q4_0 => {
15436 if rp {
15437 "qmatvec_gemm_q4_0_rp"
15438 } else {
15439 "qmatvec_gemm_q4_0"
15440 }
15441 }
15442 QT_Q5_K => "qmatvec_gemm_q5_K",
15443 QT_Q6_K => "qmatvec_gemm_q6_K",
15444 QT_NVFP4 => {
15445 if rp {
15446 "qmatvec_gemm_nvfp4_rp"
15447 } else {
15448 "qmatvec_gemm_nvfp4"
15449 }
15450 }
15451 _ => unreachable!(),
15452 };
15453 let f = self.func(name);
15454 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);
15459 let k1_tile = if is_k1 {
15461 k1_launch_override().unwrap_or((128, 128, 8))
15462 } else {
15463 (128, 128, 8)
15464 };
15465 let (bm, bn): (u32, u32) = if is_k1 {
15466 (k1_tile.0, k1_tile.1)
15467 } else {
15468 (64, 256)
15469 };
15470 let warps: u32 = if is_k1 {
15471 k1_tile.2
15472 } else {
15473 match qtype {
15474 QT_NVFP4 => 8,
15475 _ => 4,
15476 }
15477 };
15478 let cfg = LaunchConfig {
15479 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15480 block_dim: (32, warps, 1),
15481 shared_mem_bytes: 0,
15482 };
15483 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15484 let __s_b = self.gpu.stream();
15485 let mut b = __s_b.launch_builder(&f);
15486 b.arg(bytes)
15487 .arg(aq)
15488 .arg(ad)
15489 .arg(&mut y)
15490 .arg(&inf)
15491 .arg(&outf)
15492 .arg(&mi)
15493 .arg(&rb);
15494 unsafe {
15495 b.launch(cfg)?;
15496 }
15497 if scale != 1.0 {
15498 self.scale_inplace(&mut y, scale, m * out_f)?;
15499 }
15500 Ok(y)
15501 }
15502
15503 pub fn qmatvec_gemm_raw(
15508 &self,
15509 bytes: &CudaSlice<u8>,
15510 x: &CudaSlice<f32>,
15511 m: usize,
15512 in_f: usize,
15513 out_f: usize,
15514 qtype: i32,
15515 row_bytes: usize,
15516 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15517 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15518 let name = match qtype {
15519 QT_Q8_0 => "qmatvec_gemm_q8_0",
15520 QT_Q4_K => "qmatvec_gemm_q4_K",
15521 QT_Q4_0 => "qmatvec_gemm_q4_0",
15522 QT_Q5_K => "qmatvec_gemm_q5_K",
15523 QT_Q6_K => "qmatvec_gemm_q6_K",
15524 QT_NVFP4 => "qmatvec_gemm_nvfp4",
15525 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
15526 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
15527 };
15528 let f = self.func(name);
15529 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);
15533 let k1_tile = if is_k1 {
15535 k1_launch_override().unwrap_or((128, 128, 8))
15536 } else {
15537 (128, 128, 8)
15538 };
15539 let (bm, bn): (u32, u32) = if is_k1 {
15540 (k1_tile.0, k1_tile.1)
15541 } else {
15542 (64, 256)
15543 };
15544 let warps: u32 = if is_k1 {
15545 k1_tile.2
15546 } else {
15547 match qtype {
15548 QT_NVFP4 | QT_NVFP4_RP => 8,
15549 _ => 4,
15550 }
15551 };
15552 let cfg = LaunchConfig {
15553 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15554 block_dim: (32, warps, 1),
15555 shared_mem_bytes: 0,
15556 };
15557 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15558 let __s_b = self.gpu.stream();
15559 let mut b = __s_b.launch_builder(&f);
15560 b.arg(bytes)
15561 .arg(&aq)
15562 .arg(&ad)
15563 .arg(&mut y)
15564 .arg(&inf)
15565 .arg(&outf)
15566 .arg(&mi)
15567 .arg(&rb);
15568 unsafe {
15569 b.launch(cfg)?;
15570 }
15571 Ok(y)
15572 }
15573
15574 pub fn qmatvec_gemm_q8_0_wgmma_raw(
15581 &self,
15582 rp4: &CudaSlice<u8>,
15583 aq: &CudaSlice<i8>,
15584 ad: &CudaSlice<f32>,
15585 m: usize,
15586 in_f: usize,
15587 out_f: usize,
15588 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15589 assert!(
15590 out_f % 64 == 0 && in_f % 32 == 0,
15591 "wgmma GEMM needs out_f%64==0, in_f%32==0"
15592 );
15593 let f = self.func("qmatvec_gemm_q8_0_wgmma");
15594 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
15596 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
15597 block_dim: (128, 1, 1),
15598 shared_mem_bytes: 0,
15599 };
15600 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
15601 let __s_b = self.gpu.stream();
15602 let mut b = __s_b.launch_builder(&f);
15603 b.arg(rp4)
15604 .arg(aq)
15605 .arg(ad)
15606 .arg(&mut y)
15607 .arg(&inf)
15608 .arg(&outf)
15609 .arg(&mi);
15610 unsafe {
15611 b.launch(cfg)?;
15612 }
15613 Ok(y)
15614 }
15615
15616 pub fn scale_inplace(
15618 &self,
15619 y: &mut CudaSlice<f32>,
15620 s: f32,
15621 n: usize,
15622 ) -> Result<(), Box<dyn std::error::Error>> {
15623 let f = self.func("scale_f32");
15624 let cfg = LaunchConfig::for_num_elems(n as u32);
15625 let (sf, ni) = (s, n as i32);
15626 let __s_b = self.gpu.stream();
15627 let mut b = __s_b.launch_builder(&f);
15628 b.arg(y).arg(&sf).arg(&ni);
15629 unsafe {
15630 b.launch(cfg)?;
15631 }
15632 Ok(())
15633 }
15634
15635 pub fn bf16_to_f32(
15640 &self,
15641 data: &cudarc::driver::CudaView<'_, u8>,
15642 n: usize,
15643 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15644 let mut out = self.alloc_uninit::<f32>(n)?;
15645 let f = self.func("bf16_to_f32");
15646 let cfg = LaunchConfig::for_num_elems(n as u32);
15647 let ni = n as i32;
15648 let __s_b = self.gpu.stream();
15649 let mut b = __s_b.launch_builder(&f);
15650 b.arg(data).arg(&mut out).arg(&ni);
15651 unsafe {
15652 b.launch(cfg)?;
15653 }
15654 Ok(out)
15655 }
15656
15657 fn linear_bf16_chunked(
15664 &self,
15665 x: &CudaSlice<f32>,
15666 data: &CudaSlice<u8>,
15667 m: usize,
15668 in_f: usize,
15669 out_f: usize,
15670 exact: bool,
15671 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15672 const CHUNK_BYTES: usize = 256 << 20;
15673 let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
15674 if chunk_rows >= out_f {
15675 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
15676 return if exact {
15677 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
15678 } else {
15679 self.linear(x, &wf32, m, in_f, out_f)
15680 };
15681 }
15682 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15683 let mut r0 = 0usize;
15684 while r0 < out_f {
15685 let rows = chunk_rows.min(out_f - r0);
15686 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
15687 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
15688 let yc = if exact {
15689 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
15690 } else {
15691 self.linear(x, &wf32, m, in_f, rows)?
15692 };
15693 for mi in 0..m {
15695 let src = yc.slice(mi * rows..(mi + 1) * rows);
15696 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
15697 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
15698 }
15699 r0 += rows;
15700 }
15701 Ok(y)
15702 }
15703
15704 pub fn linear_decode_exact(
15711 &self,
15712 x: &CudaSlice<f32>,
15713 w: &CudaSlice<f32>,
15714 m_tokens: usize,
15715 in_f: usize,
15716 out_f: usize,
15717 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15718 if m_tokens == 1 {
15719 return self.linear(x, w, 1, in_f, out_f);
15720 }
15721 let xv = self.view(x, m_tokens * in_f);
15722 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
15723 for t in 0..m_tokens {
15724 let row = xv.slice(t * in_f..(t + 1) * in_f);
15725 let mut xr = self.alloc_uninit::<f32>(in_f)?;
15726 self.copy_view_into(&mut xr, 0, &row, in_f)?;
15727 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
15728 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
15729 }
15730 Ok(y)
15731 }
15732
15733 pub fn linear(
15734 &self,
15735 x: &CudaSlice<f32>,
15736 w: &CudaSlice<f32>,
15737 m_tokens: usize,
15738 in_f: usize,
15739 out_f: usize,
15740 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15741 use cudarc::cublaslt::{Matmul, MatmulConfig};
15742 let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; let cfg = MatmulConfig {
15744 transa: true,
15745 transb: false,
15746 transc: false,
15747 m: out_f as u64,
15748 n: m_tokens as u64,
15749 k: in_f as u64,
15750 alpha: 1.0,
15751 lda: in_f as i64,
15752 ldb: in_f as i64,
15753 beta: 0.0,
15754 ldc: out_f as i64,
15755 stride_a: None,
15756 stride_b: None,
15757 stride_c: None,
15758 stride_bias: None,
15759 batch_size: None,
15760 };
15761 unsafe {
15762 self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15763 }
15764 Ok(c)
15765 }
15766
15767 pub fn sdpa_naive(
15769 &self,
15770 q: &CudaSlice<f32>,
15771 k: &CudaSlice<f32>,
15772 v: &CudaSlice<f32>,
15773 o: &mut CudaSlice<f32>,
15774 head_dim: usize,
15775 n_head: usize,
15776 n_head_kv: usize,
15777 t: usize,
15778 t_kv: usize,
15779 scale: f32,
15780 causal: bool,
15781 ) -> Result<(), Box<dyn std::error::Error>> {
15782 let f = self.func("sdpa_naive_f32");
15783 let cfg = LaunchConfig {
15784 grid_dim: (n_head as u32, t as u32, 1),
15785 block_dim: (128, 1, 1),
15786 shared_mem_bytes: (t_kv * 4) as u32,
15787 };
15788 let (hd, nh, nhkv, ti, tkvi, cz) = (
15789 head_dim as i32,
15790 n_head as i32,
15791 n_head_kv as i32,
15792 t as i32,
15793 t_kv as i32,
15794 causal as i32,
15795 );
15796 let __s_b = self.gpu.stream();
15797 let mut b = __s_b.launch_builder(&f);
15798 b.arg(q)
15799 .arg(k)
15800 .arg(v)
15801 .arg(o)
15802 .arg(&hd)
15803 .arg(&nh)
15804 .arg(&nhkv)
15805 .arg(&ti)
15806 .arg(&tkvi)
15807 .arg(&scale)
15808 .arg(&cz);
15809 unsafe {
15810 b.launch(cfg)?;
15811 }
15812 Ok(())
15813 }
15814
15815 #[allow(clippy::too_many_arguments)]
15820 pub fn sdpa_naive_island(
15821 &self,
15822 q: &CudaSlice<f32>,
15823 k: &CudaSlice<f32>,
15824 v: &CudaSlice<f32>,
15825 o: &mut CudaSlice<f32>,
15826 span_id: &CudaSlice<i32>,
15827 head_dim: usize,
15828 n_head: usize,
15829 n_head_kv: usize,
15830 t: usize,
15831 t_kv: usize,
15832 scale: f32,
15833 window: usize,
15834 ) -> Result<(), Box<dyn std::error::Error>> {
15835 let f = self.func("sdpa_naive_island_f32");
15836 let cfg = LaunchConfig {
15837 grid_dim: (n_head as u32, t as u32, 1),
15838 block_dim: (128, 1, 1),
15839 shared_mem_bytes: (t_kv * 4) as u32,
15840 };
15841 let (hd, nh, nhkv, ti, tkvi, wi) = (
15842 head_dim as i32,
15843 n_head as i32,
15844 n_head_kv as i32,
15845 t as i32,
15846 t_kv as i32,
15847 window as i32,
15848 );
15849 let __s_b = self.gpu.stream();
15850 let mut b = __s_b.launch_builder(&f);
15851 b.arg(q)
15852 .arg(k)
15853 .arg(v)
15854 .arg(o)
15855 .arg(span_id)
15856 .arg(&hd)
15857 .arg(&nh)
15858 .arg(&nhkv)
15859 .arg(&ti)
15860 .arg(&tkvi)
15861 .arg(&scale)
15862 .arg(&wi);
15863 unsafe {
15864 b.launch(cfg)?;
15865 }
15866 Ok(())
15867 }
15868
15869 #[allow(clippy::too_many_arguments)]
15871 pub fn sdpa_naive_w(
15872 &self,
15873 q: &CudaSlice<f32>,
15874 k: &CudaSlice<f32>,
15875 v: &CudaSlice<f32>,
15876 o: &mut CudaSlice<f32>,
15877 head_dim: usize,
15878 n_head: usize,
15879 n_head_kv: usize,
15880 t: usize,
15881 t_kv: usize,
15882 scale: f32,
15883 causal: bool,
15884 window: usize,
15885 ) -> Result<(), Box<dyn std::error::Error>> {
15886 let f = self.func("sdpa_naive_w_f32");
15887 let cfg = LaunchConfig {
15888 grid_dim: (n_head as u32, t as u32, 1),
15889 block_dim: (128, 1, 1),
15890 shared_mem_bytes: (t_kv * 4) as u32,
15891 };
15892 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15893 head_dim as i32,
15894 n_head as i32,
15895 n_head_kv as i32,
15896 t as i32,
15897 t_kv as i32,
15898 causal as i32,
15899 window as i32,
15900 );
15901 let __s_b = self.gpu.stream();
15902 let mut b = __s_b.launch_builder(&f);
15903 b.arg(q)
15904 .arg(k)
15905 .arg(v)
15906 .arg(o)
15907 .arg(&hd)
15908 .arg(&nh)
15909 .arg(&nhkv)
15910 .arg(&ti)
15911 .arg(&tkvi)
15912 .arg(&scale)
15913 .arg(&cz)
15914 .arg(&wi);
15915 unsafe {
15916 b.launch(cfg)?;
15917 }
15918 Ok(())
15919 }
15920
15921 #[allow(clippy::too_many_arguments)]
15931 pub fn sdpa_naive_w_lo(
15932 &self,
15933 q: &CudaSlice<f32>,
15934 k: &CudaSlice<f32>,
15935 v: &CudaSlice<f32>,
15936 o: &mut CudaSlice<f32>,
15937 head_dim: usize,
15938 n_head: usize,
15939 n_head_kv: usize,
15940 t: usize,
15941 t_kv: usize,
15942 scale: f32,
15943 causal: bool,
15944 window: usize,
15945 ) -> Result<(), Box<dyn std::error::Error>> {
15946 let kv_lo = if window > 0 {
15947 (t_kv - t + 1).saturating_sub(window)
15948 } else {
15949 0
15950 };
15951 let smem = (t_kv - kv_lo) * 4;
15952 if smem > 48 * 1024 {
15953 return Err(format!(
15954 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
15955 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
15956 a window this wide needs the multi-pass long-ctx kernel"
15957 )
15958 .into());
15959 }
15960 let f = self.func("sdpa_naive_w_lo_f32");
15961 let cfg = LaunchConfig {
15962 grid_dim: (n_head as u32, t as u32, 1),
15963 block_dim: (128, 1, 1),
15964 shared_mem_bytes: smem as u32,
15965 };
15966 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
15967 head_dim as i32,
15968 n_head as i32,
15969 n_head_kv as i32,
15970 t as i32,
15971 t_kv as i32,
15972 causal as i32,
15973 window as i32,
15974 kv_lo as i32,
15975 );
15976 let __s_b = self.gpu.stream();
15977 let mut b = __s_b.launch_builder(&f);
15978 b.arg(q)
15979 .arg(k)
15980 .arg(v)
15981 .arg(o)
15982 .arg(&hd)
15983 .arg(&nh)
15984 .arg(&nhkv)
15985 .arg(&ti)
15986 .arg(&tkvi)
15987 .arg(&scale)
15988 .arg(&cz)
15989 .arg(&wi)
15990 .arg(&lo);
15991 unsafe {
15992 b.launch(cfg)?;
15993 }
15994 Ok(())
15995 }
15996
15997 pub fn sdpa_naive_view(
15999 &self,
16000 q: &CudaSlice<f32>,
16001 k: &cudarc::driver::CudaView<f32>,
16002 v: &cudarc::driver::CudaView<f32>,
16003 o: &mut CudaSlice<f32>,
16004 head_dim: usize,
16005 n_head: usize,
16006 n_head_kv: usize,
16007 t: usize,
16008 t_kv: usize,
16009 scale: f32,
16010 causal: bool,
16011 ) -> Result<(), Box<dyn std::error::Error>> {
16012 let f = self.func("sdpa_naive_f32");
16013 let cfg = LaunchConfig {
16014 grid_dim: (n_head as u32, t as u32, 1),
16015 block_dim: (128, 1, 1),
16016 shared_mem_bytes: (t_kv * 4) as u32,
16017 };
16018 let (hd, nh, nhkv, ti, tkvi, cz) = (
16019 head_dim as i32,
16020 n_head as i32,
16021 n_head_kv as i32,
16022 t as i32,
16023 t_kv as i32,
16024 causal as i32,
16025 );
16026 let __s_b = self.gpu.stream();
16027 let mut b = __s_b.launch_builder(&f);
16028 b.arg(q)
16029 .arg(k)
16030 .arg(v)
16031 .arg(o)
16032 .arg(&hd)
16033 .arg(&nh)
16034 .arg(&nhkv)
16035 .arg(&ti)
16036 .arg(&tkvi)
16037 .arg(&scale)
16038 .arg(&cz);
16039 unsafe {
16040 b.launch(cfg)?;
16041 }
16042 Ok(())
16043 }
16044
16045 #[allow(clippy::too_many_arguments)]
16053 pub fn fa_dequant_kv_view_f32(
16054 &self,
16055 k: &cudarc::driver::CudaView<u8>,
16056 v: &cudarc::driver::CudaView<u8>,
16057 kf: &mut CudaSlice<f32>,
16058 vf: &mut CudaSlice<f32>,
16059 kv_dim_k: usize,
16060 kv_dim_v: usize,
16061 t_kv: usize,
16062 k_tok_bytes: usize,
16063 v_tok_bytes: usize,
16064 g: bool,
16065 ) -> Result<(), Box<dyn std::error::Error>> {
16066 let f = if g {
16067 self.func_g("fa_dequant_kv_ws_f32")
16068 } else {
16069 self.func("fa_dequant_kv_ws_f32")
16070 };
16071 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
16072 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16073 let cfg = LaunchConfig {
16074 grid_dim: (nblk.max(1), 1, 1),
16075 block_dim: (256, 1, 1),
16076 shared_mem_bytes: 0,
16077 };
16078 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
16079 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16080 let __s_b = self.gpu.stream();
16081 let mut b = __s_b.launch_builder(&f);
16082 b.arg(k)
16083 .arg(v)
16084 .arg(&mut *kf)
16085 .arg(&mut *vf)
16086 .arg(&kdk)
16087 .arg(&kdv)
16088 .arg(&tkvi)
16089 .arg(&ktb)
16090 .arg(&vtb);
16091 unsafe {
16092 b.launch(cfg)?;
16093 }
16094 Ok(())
16095 }
16096
16097 #[allow(clippy::too_many_arguments)]
16098 pub fn sdpa_naive_quantized_view(
16099 &self,
16100 q: &CudaSlice<f32>,
16101 k: &cudarc::driver::CudaView<u8>,
16102 v: &cudarc::driver::CudaView<u8>,
16103 o: &mut CudaSlice<f32>,
16104 head_dim: usize,
16105 n_head: usize,
16106 n_head_kv: usize,
16107 t: usize,
16108 t_kv: usize,
16109 scale: f32,
16110 causal: bool,
16111 k_tok_bytes: usize,
16112 v_tok_bytes: usize,
16113 ) -> Result<(), Box<dyn std::error::Error>> {
16114 let kv_dim = n_head_kv * head_dim;
16115 let mut kf = self.uninit(t_kv * kv_dim)?;
16116 let mut vf = self.uninit(t_kv * kv_dim)?;
16117 let f = self.func("fa_dequant_kv_ws_f32");
16118 let total = (2 * t_kv * kv_dim) as u64;
16119 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16120 let cfg = LaunchConfig {
16121 grid_dim: (nblk.max(1), 1, 1),
16122 block_dim: (256, 1, 1),
16123 shared_mem_bytes: 0,
16124 };
16125 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
16126 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
16127 let __s_b = self.gpu.stream();
16128 let mut b = __s_b.launch_builder(&f);
16129 b.arg(k)
16130 .arg(v)
16131 .arg(&mut kf)
16132 .arg(&mut vf)
16133 .arg(&kv_dim_i)
16134 .arg(&kv_dim_i)
16135 .arg(&t_kv_i)
16136 .arg(&k_tok_bytes_i)
16137 .arg(&v_tok_bytes_i);
16138 unsafe { b.launch(cfg)? };
16139 self.sdpa_naive(
16140 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16141 )
16142 }
16143
16144 #[allow(clippy::too_many_arguments)]
16156 pub fn sdpa_naive_w_quantized_view(
16157 &self,
16158 q: &CudaSlice<f32>,
16159 k: &cudarc::driver::CudaView<u8>,
16160 v: &cudarc::driver::CudaView<u8>,
16161 o: &mut CudaSlice<f32>,
16162 head_dim: usize,
16163 n_head: usize,
16164 n_head_kv: usize,
16165 t: usize,
16166 t_kv: usize,
16167 scale: f32,
16168 causal: bool,
16169 window: usize,
16170 k_tok_bytes: usize,
16171 v_tok_bytes: usize,
16172 ) -> Result<(), Box<dyn std::error::Error>> {
16173 let kv_dim = n_head_kv * head_dim;
16174 let mut kf = self.uninit(t_kv * kv_dim)?;
16175 let mut vf = self.uninit(t_kv * kv_dim)?;
16176 let f = self.func("fa_dequant_kv_ws_f32");
16177 let total = (2 * t_kv * kv_dim) as u64;
16178 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16179 let cfg = LaunchConfig {
16180 grid_dim: (nblk.max(1), 1, 1),
16181 block_dim: (256, 1, 1),
16182 shared_mem_bytes: 0,
16183 };
16184 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
16185 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
16186 let __s_b = self.gpu.stream();
16187 let mut b = __s_b.launch_builder(&f);
16188 b.arg(k)
16189 .arg(v)
16190 .arg(&mut kf)
16191 .arg(&mut vf)
16192 .arg(&kv_dim_i)
16193 .arg(&kv_dim_i)
16194 .arg(&t_kv_i)
16195 .arg(&k_tok_bytes_i)
16196 .arg(&v_tok_bytes_i);
16197 unsafe { b.launch(cfg)? };
16198 self.sdpa_naive_w(
16199 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
16200 )
16201 }
16202
16203 pub fn fa_prefill(
16207 &self,
16208 q: &CudaSlice<f32>,
16209 k: &CudaSlice<f32>,
16210 v: &CudaSlice<f32>,
16211 o: &mut CudaSlice<f32>,
16212 head_dim: usize,
16213 n_head: usize,
16214 n_head_kv: usize,
16215 t: usize,
16216 t_kv: usize,
16217 scale: f32,
16218 causal: bool,
16219 ) -> Result<(), Box<dyn std::error::Error>> {
16220 if portable_mma_gated() {
16221 return self.sdpa_naive(
16222 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16223 );
16224 }
16225 let fa3_on = head_dim == 256
16233 && causal
16234 && t == t_kv
16235 && match std::env::var("MEMRA_FA3").as_deref() {
16236 Ok("0") => false,
16237 Ok("1") => true,
16238 _ => cfg!(memra_hopper_mma),
16239 };
16240 if fa3_on {
16241 let n = t * n_head * head_dim;
16242 let nkv = t * n_head_kv * head_dim;
16243 let mut q16 = self.alloc_u8_uninit(n * 2)?;
16244 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
16245 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
16246 self.f32_to_bf16_into(q, &mut q16, n)?;
16247 self.f32_to_bf16_into(k, &mut k16, nkv)?;
16248 self.f32_to_bf16_into(v, &mut v16, nkv)?;
16249 let rc = {
16250 use cudarc::driver::{DevicePtr, DevicePtrMut};
16251 let stream = self.gpu.stream();
16252 let (qp, _g1) = q16.device_ptr(&stream);
16253 let (kp, _g2) = k16.device_ptr(&stream);
16254 let (vp, _g3) = v16.device_ptr(&stream);
16255 let (op, _g4) = o.device_ptr_mut(&stream);
16256 unsafe {
16257 memra_fa3_prefill(
16258 qp as *const core::ffi::c_void,
16259 kp as *const core::ffi::c_void,
16260 vp as *const core::ffi::c_void,
16261 op as *mut f32,
16262 t as i32,
16263 n_head as i32,
16264 n_head_kv as i32,
16265 head_dim as i32,
16266 scale,
16267 stream.cu_stream() as *mut core::ffi::c_void,
16268 )
16269 }
16270 };
16271 if rc != 0 {
16272 return Err(format!("memra_fa3_prefill rc={rc}").into());
16273 }
16274 return Ok(());
16275 }
16276 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16281 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
16282 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
16283 const BLOCK_Q: usize = 64;
16284 const BKX: usize = 32;
16285 let f = self.func("fa_prefill_bf16_p1");
16286 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
16287 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
16288 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16289 f.set_attribute(
16290 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16291 shmem as i32,
16292 )?;
16293 let cfg = LaunchConfig {
16294 grid_dim: (
16295 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16296 n_head as u32,
16297 1,
16298 ),
16299 block_dim: (32, 4, 1),
16300 shared_mem_bytes: shmem,
16301 };
16302 let (hd, nh, nhkv, ti, tkvi, cz) = (
16303 head_dim as i32,
16304 n_head as i32,
16305 n_head_kv as i32,
16306 t as i32,
16307 t_kv as i32,
16308 causal as i32,
16309 );
16310 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16311 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16312 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16313 let __s_b = self.gpu.stream();
16314 let mut b = __s_b.launch_builder(&f);
16315 b.arg(&qb)
16316 .arg(&kb)
16317 .arg(&vb)
16318 .arg(o)
16319 .arg(&hd)
16320 .arg(&nh)
16321 .arg(&nhkv)
16322 .arg(&ti)
16323 .arg(&tkvi)
16324 .arg(&scale)
16325 .arg(&cz);
16326 unsafe {
16327 b.launch(cfg)?;
16328 }
16329 return Ok(());
16330 }
16331 const BK: usize = 32;
16337 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
16340 let (block_q, warps, w2_sfx): (usize, u32, &str) =
16341 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
16342 let hd_sfx = fa_hd_suffix(head_dim)?;
16346 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
16347 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
16352 let (kb16, vb16) = if bf16kv {
16353 let n = t_kv * n_head_kv * head_dim;
16354 let mut kb = self.alloc_u8_uninit(n * 2)?;
16355 let mut vb = self.alloc_u8_uninit(n * 2)?;
16356 let fcv = self.func("f32_to_bf16_bulk");
16357 let ni = n as i64;
16358 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16359 let __s_b = self.gpu.stream();
16360 let mut b = __s_b.launch_builder(&fcv);
16361 b.arg(k).arg(&mut kb).arg(&ni);
16362 unsafe {
16363 b.launch(cfgc)?;
16364 }
16365 let __s_b = self.gpu.stream();
16366 let mut b = __s_b.launch_builder(&fcv);
16367 b.arg(v).arg(&mut vb).arg(&ni);
16368 unsafe {
16369 b.launch(cfgc)?;
16370 }
16371 (Some(kb), Some(vb))
16372 } else {
16373 (None, None)
16374 };
16375 let f = self.func(&if bf16kv {
16376 format!("fa_prefill_bf16kv_pp{hd_sfx}")
16377 } else {
16378 format!(
16379 "fa_prefill_f32{}{}{hd_sfx}",
16380 if floor { "" } else { "_pp" },
16381 if floor { "" } else { w2_sfx }
16382 )
16383 });
16384 let kv_stages = if bf16kv { 2 } else { 1 };
16387 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16388 + 4 * (block_q * BK + 2 * block_q)) as u32;
16389 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16390 f.set_attribute(
16391 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16392 shmem as i32,
16393 )?;
16394 let cfg = LaunchConfig {
16395 grid_dim: (
16396 (t as u32 + block_q as u32 - 1) / block_q as u32,
16397 n_head as u32,
16398 1,
16399 ),
16400 block_dim: (32, warps, 1),
16401 shared_mem_bytes: shmem,
16402 };
16403 let (hd, nh, nhkv, ti, tkvi, cz) = (
16404 head_dim as i32,
16405 n_head as i32,
16406 n_head_kv as i32,
16407 t as i32,
16408 t_kv as i32,
16409 causal as i32,
16410 );
16411 let __s_b = self.gpu.stream();
16412 let mut b = __s_b.launch_builder(&f);
16413 b.arg(q);
16414 match (&kb16, &vb16) {
16415 (Some(kb), Some(vb)) => {
16416 b.arg(kb).arg(vb);
16417 }
16418 _ => {
16419 b.arg(k).arg(v);
16420 }
16421 }
16422 b.arg(o)
16423 .arg(&hd)
16424 .arg(&nh)
16425 .arg(&nhkv)
16426 .arg(&ti)
16427 .arg(&tkvi)
16428 .arg(&scale)
16429 .arg(&cz);
16430 unsafe {
16431 b.launch(cfg)?;
16432 }
16433 Ok(())
16434 }
16435
16436 #[allow(clippy::too_many_arguments)]
16440 pub fn fa_prefill_w(
16441 &self,
16442 q: &CudaSlice<f32>,
16443 k: &CudaSlice<f32>,
16444 v: &CudaSlice<f32>,
16445 o: &mut CudaSlice<f32>,
16446 head_dim: usize,
16447 n_head: usize,
16448 n_head_kv: usize,
16449 t: usize,
16450 t_kv: usize,
16451 scale: f32,
16452 causal: bool,
16453 window: usize,
16454 ) -> Result<(), Box<dyn std::error::Error>> {
16455 if portable_mma_gated() {
16458 return self.sdpa_naive_w(
16459 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
16460 );
16461 }
16462 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16466 let faw_f32 =
16467 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
16468 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
16469 self.fa_prefill_w_arm(
16470 q,
16471 k,
16472 v,
16473 o,
16474 head_dim,
16475 n_head,
16476 n_head_kv,
16477 t,
16478 t_kv,
16479 scale,
16480 causal,
16481 window,
16482 floor || faw_f32,
16483 floor,
16484 )
16485 }
16486
16487 #[allow(clippy::too_many_arguments)]
16490 pub fn fa_prefill_w_pre(
16491 &self,
16492 qb: &CudaSlice<u8>,
16493 kb: &CudaSlice<u8>,
16494 vb: &CudaSlice<u8>,
16495 o: &mut CudaSlice<f32>,
16496 head_dim: usize,
16497 n_head: usize,
16498 n_head_kv: usize,
16499 t: usize,
16500 t_kv: usize,
16501 scale: f32,
16502 causal: bool,
16503 window: usize,
16504 v_f16: bool,
16505 ) -> Result<(), Box<dyn std::error::Error>> {
16506 const BLOCK_Q: usize = 64;
16507 const BK: usize = 32;
16508 debug_assert_eq!(head_dim, 256);
16509 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16510 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
16511 if hp {
16512 const BLOCK_QH: usize = 32;
16513 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16516 let vh: &CudaSlice<u8> = if v_f16 {
16517 vb
16518 } else {
16519 let n = t_kv * n_head_kv * head_dim;
16520 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
16521 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
16522 }
16523 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
16524 vguard.as_ref().unwrap()
16525 };
16526 let f = self.func("fa_prefill_w_bf16_p1h2");
16527 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16528 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16529 f.set_attribute(
16530 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16531 shmem as i32,
16532 )?;
16533 let cfg = LaunchConfig {
16534 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16535 block_dim: (32, 4, 1),
16536 shared_mem_bytes: shmem,
16537 };
16538 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16539 head_dim as i32,
16540 n_head as i32,
16541 n_head_kv as i32,
16542 t as i32,
16543 t_kv as i32,
16544 causal as i32,
16545 window as i32,
16546 );
16547 let __s_b = self.gpu.stream();
16548 let mut b = __s_b.launch_builder(&f);
16549 b.arg(qb)
16550 .arg(kb)
16551 .arg(vh)
16552 .arg(o)
16553 .arg(&hd)
16554 .arg(&nh)
16555 .arg(&nhkv)
16556 .arg(&ti)
16557 .arg(&tkvi)
16558 .arg(&scale)
16559 .arg(&cz)
16560 .arg(&wi);
16561 unsafe {
16562 b.launch(cfg)?;
16563 }
16564 return Ok(());
16565 }
16566 let f = self.func("fa_prefill_w_bf16_p1");
16567 let shmem =
16568 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16569 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16570 f.set_attribute(
16571 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16572 shmem as i32,
16573 )?;
16574 let cfg = LaunchConfig {
16575 grid_dim: (
16576 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16577 n_head as u32,
16578 1,
16579 ),
16580 block_dim: (32, 4, 1),
16581 shared_mem_bytes: shmem,
16582 };
16583 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16584 head_dim as i32,
16585 n_head as i32,
16586 n_head_kv as i32,
16587 t as i32,
16588 t_kv as i32,
16589 causal as i32,
16590 window as i32,
16591 );
16592 let __s_b = self.gpu.stream();
16593 let mut b = __s_b.launch_builder(&f);
16594 b.arg(qb)
16595 .arg(kb)
16596 .arg(vb)
16597 .arg(o)
16598 .arg(&hd)
16599 .arg(&nh)
16600 .arg(&nhkv)
16601 .arg(&ti)
16602 .arg(&tkvi)
16603 .arg(&scale)
16604 .arg(&cz)
16605 .arg(&wi);
16606 unsafe {
16607 b.launch(cfg)?;
16608 }
16609 Ok(())
16610 }
16611
16612 #[allow(clippy::too_many_arguments)]
16614 pub fn fa_prefill_w_arm(
16615 &self,
16616 q: &CudaSlice<f32>,
16617 k: &CudaSlice<f32>,
16618 v: &CudaSlice<f32>,
16619 o: &mut CudaSlice<f32>,
16620 head_dim: usize,
16621 n_head: usize,
16622 n_head_kv: usize,
16623 t: usize,
16624 t_kv: usize,
16625 scale: f32,
16626 causal: bool,
16627 window: usize,
16628 f32_stage: bool,
16629 floor: bool,
16630 ) -> Result<(), Box<dyn std::error::Error>> {
16631 const BLOCK_Q: usize = 64;
16632 const BK: usize = 32;
16633 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
16634 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16638 let p1 = !floor
16639 && !f32_stage
16640 && *P1_ON.get_or_init(|| {
16641 std::env::var("MEMRA_FAW_P1")
16642 .map(|v| v != "0")
16643 .unwrap_or(true)
16644 });
16645 let hp =
16646 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16647 if hp {
16648 const BLOCK_QH: usize = 32;
16649 let f = self.func("fa_prefill_w_bf16_p1h2");
16650 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16651 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16652 f.set_attribute(
16653 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16654 shmem as i32,
16655 )?;
16656 let cfg = LaunchConfig {
16657 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16658 block_dim: (32, 4, 1),
16659 shared_mem_bytes: shmem,
16660 };
16661 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16662 head_dim as i32,
16663 n_head as i32,
16664 n_head_kv as i32,
16665 t as i32,
16666 t_kv as i32,
16667 causal as i32,
16668 window as i32,
16669 );
16670 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16671 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16672 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
16673 let __s_b = self.gpu.stream();
16674 let mut b = __s_b.launch_builder(&f);
16675 b.arg(&qb)
16676 .arg(&kb)
16677 .arg(&vh)
16678 .arg(o)
16679 .arg(&hd)
16680 .arg(&nh)
16681 .arg(&nhkv)
16682 .arg(&ti)
16683 .arg(&tkvi)
16684 .arg(&scale)
16685 .arg(&cz)
16686 .arg(&wi);
16687 unsafe {
16688 b.launch(cfg)?;
16689 }
16690 return Ok(());
16691 }
16692 if p1 {
16693 let f = self.func("fa_prefill_w_bf16_p1");
16694 let shmem =
16695 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16696 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16697 f.set_attribute(
16698 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16699 shmem as i32,
16700 )?;
16701 let cfg = LaunchConfig {
16702 grid_dim: (
16703 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16704 n_head as u32,
16705 1,
16706 ),
16707 block_dim: (32, 4, 1),
16708 shared_mem_bytes: shmem,
16709 };
16710 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16711 head_dim as i32,
16712 n_head as i32,
16713 n_head_kv as i32,
16714 t as i32,
16715 t_kv as i32,
16716 causal as i32,
16717 window as i32,
16718 );
16719 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16720 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16721 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16722 let __s_b = self.gpu.stream();
16723 let mut b = __s_b.launch_builder(&f);
16724 b.arg(&qb)
16725 .arg(&kb)
16726 .arg(&vb)
16727 .arg(o)
16728 .arg(&hd)
16729 .arg(&nh)
16730 .arg(&nhkv)
16731 .arg(&ti)
16732 .arg(&tkvi)
16733 .arg(&scale)
16734 .arg(&cz)
16735 .arg(&wi);
16736 unsafe {
16737 b.launch(cfg)?;
16738 }
16739 return Ok(());
16740 }
16741 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16744 let g4 = !floor
16745 && !f32_stage
16746 && n_head_kv == 1
16747 && n_head % 4 == 0
16748 && *G4_ON.get_or_init(|| {
16749 std::env::var("MEMRA_FAW_G4")
16750 .map(|v| v != "0")
16751 .unwrap_or(true)
16752 });
16753 if g4 {
16754 const SP_M: usize = 16;
16755 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16758 let o2 = *O2_ON.get_or_init(|| {
16759 std::env::var("MEMRA_FAW_O2")
16760 .map(|v| v != "0")
16761 .unwrap_or(true)
16762 });
16763 let f = self.func(if o2 {
16764 "fa_prefill_w_bf16_g4o2"
16765 } else {
16766 "fa_prefill_w_bf16_g4"
16767 });
16768 let shmem = if o2 {
16769 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
16770 } else {
16771 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
16772 as u32
16773 };
16774 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16775 f.set_attribute(
16776 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16777 shmem as i32,
16778 )?;
16779 let cfg = LaunchConfig {
16780 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
16781 block_dim: (32, 4, 1),
16782 shared_mem_bytes: shmem,
16783 };
16784 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16785 head_dim as i32,
16786 n_head as i32,
16787 n_head_kv as i32,
16788 t as i32,
16789 t_kv as i32,
16790 causal as i32,
16791 window as i32,
16792 );
16793 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16794 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16795 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16796 let __s_b = self.gpu.stream();
16797 let mut b = __s_b.launch_builder(&f);
16798 b.arg(&qb)
16799 .arg(&kb)
16800 .arg(&vb)
16801 .arg(o)
16802 .arg(&hd)
16803 .arg(&nh)
16804 .arg(&nhkv)
16805 .arg(&ti)
16806 .arg(&tkvi)
16807 .arg(&scale)
16808 .arg(&cz)
16809 .arg(&wi);
16810 unsafe {
16811 b.launch(cfg)?;
16812 }
16813 return Ok(());
16814 }
16815 let f = self.func(if floor {
16816 "fa_prefill_w_f32"
16817 } else if f32_stage {
16818 "fa_prefill_w_f32_pp"
16819 } else {
16820 "fa_prefill_w_bf16_pp"
16821 });
16822 let shmem =
16823 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16824 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16825 f.set_attribute(
16826 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16827 shmem as i32,
16828 )?;
16829 let cfg = LaunchConfig {
16830 grid_dim: (
16831 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16832 n_head as u32,
16833 1,
16834 ),
16835 block_dim: (32, 4, 1),
16836 shared_mem_bytes: shmem,
16837 };
16838 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16839 head_dim as i32,
16840 n_head as i32,
16841 n_head_kv as i32,
16842 t as i32,
16843 t_kv as i32,
16844 causal as i32,
16845 window as i32,
16846 );
16847 if f32_stage {
16848 let __s_b = self.gpu.stream();
16849 let mut b = __s_b.launch_builder(&f);
16850 b.arg(q)
16851 .arg(k)
16852 .arg(v)
16853 .arg(o)
16854 .arg(&hd)
16855 .arg(&nh)
16856 .arg(&nhkv)
16857 .arg(&ti)
16858 .arg(&tkvi)
16859 .arg(&scale)
16860 .arg(&cz)
16861 .arg(&wi);
16862 unsafe {
16863 b.launch(cfg)?;
16864 }
16865 } else {
16866 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16867 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16868 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16869 let __s_b = self.gpu.stream();
16870 let mut b = __s_b.launch_builder(&f);
16871 b.arg(&qb)
16872 .arg(&kb)
16873 .arg(&vb)
16874 .arg(o)
16875 .arg(&hd)
16876 .arg(&nh)
16877 .arg(&nhkv)
16878 .arg(&ti)
16879 .arg(&tkvi)
16880 .arg(&scale)
16881 .arg(&cz)
16882 .arg(&wi);
16883 unsafe {
16884 b.launch(cfg)?;
16885 }
16886 }
16887 Ok(())
16888 }
16889
16890 #[allow(clippy::too_many_arguments)]
16894 pub fn fa_prefill_hd512(
16895 &self,
16896 q: &CudaSlice<f32>,
16897 k: &CudaSlice<f32>,
16898 v: &CudaSlice<f32>,
16899 o: &mut CudaSlice<f32>,
16900 head_dim: usize,
16901 n_head: usize,
16902 n_head_kv: usize,
16903 t: usize,
16904 t_kv: usize,
16905 scale: f32,
16906 causal: bool,
16907 ) -> Result<(), Box<dyn std::error::Error>> {
16908 if portable_mma_gated() {
16910 return self.sdpa_naive(
16911 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16912 );
16913 }
16914 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16920 let f32_stage =
16921 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16922 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16926 let sp = !f32_stage
16927 && *SP_ON.get_or_init(|| {
16928 std::env::var("MEMRA_FA512_SP")
16929 .map(|v| v != "0")
16930 .unwrap_or(true)
16931 });
16932 self.fa_prefill_hd512_arm(
16933 q,
16934 k,
16935 v,
16936 o,
16937 head_dim,
16938 n_head,
16939 n_head_kv,
16940 t,
16941 t_kv,
16942 scale,
16943 causal,
16944 f32_stage,
16945 sp,
16946 sp && fa_f16pv_on(),
16947 )
16948 }
16949
16950 #[allow(clippy::too_many_arguments)]
16952 pub fn fa_prefill_hd512_pre(
16953 &self,
16954 qb: &CudaSlice<u8>,
16955 kb: &CudaSlice<u8>,
16956 vb: &CudaSlice<u8>,
16957 o: &mut CudaSlice<f32>,
16958 head_dim: usize,
16959 n_head: usize,
16960 n_head_kv: usize,
16961 t: usize,
16962 t_kv: usize,
16963 scale: f32,
16964 causal: bool,
16965 v_f16: bool,
16966 ) -> Result<(), Box<dyn std::error::Error>> {
16967 debug_assert_eq!(head_dim, 512);
16968 const SP_M: usize = 16;
16969 const BKS: usize = 32;
16970 let f16pv = fa_f16pv_on();
16974 let nw = if f16pv { fa512_wide_warps() } else { 2 };
16975 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16976 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16977 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16978 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16979 let n = t_kv * n_head_kv * head_dim;
16981 let need = n * 2;
16982 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16983 *vguard = Some(self.alloc_uninit::<u8>(need)?);
16984 }
16985 let dst = vguard.as_mut().unwrap();
16986 self.bf16_to_f16_into(vb, n, dst)?;
16987 vguard.as_ref().unwrap()
16988 } else {
16989 vb
16990 };
16991 let f = self.func(if hp {
16992 "fa_prefill_bf16_hd512_sp16h2"
16993 } else {
16994 match (f16pv, nw) {
16995 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16996 (true, _) => "fa_prefill_bf16_hd512_sp16",
16997 _ => "fa_prefill_bf16_hd512_sp",
16998 }
16999 });
17000 let (nwarp, npart) = if hp {
17001 (4usize, 4usize)
17002 } else if nw > 2 {
17003 (nw, nw)
17004 } else {
17005 (2, 1)
17006 };
17007 let shmem = if hp {
17009 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
17010 as u32
17011 } else {
17012 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
17013 + 4 * (npart * SP_M * BKS + SP_M)) as u32
17014 };
17015 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17016 f.set_attribute(
17017 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17018 shmem as i32,
17019 )?;
17020 let grid_y = if hp {
17021 (n_head / 2) as u32
17022 } else {
17023 n_head as u32
17024 };
17025 let cfg = LaunchConfig {
17026 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
17027 block_dim: (32, nwarp as u32, 1),
17028 shared_mem_bytes: shmem,
17029 };
17030 let (hd, nh, nhkv, ti, tkvi, cz) = (
17031 head_dim as i32,
17032 n_head as i32,
17033 n_head_kv as i32,
17034 t as i32,
17035 t_kv as i32,
17036 causal as i32,
17037 );
17038 let __s_b = self.gpu.stream();
17039 let mut b = __s_b.launch_builder(&f);
17040 b.arg(qb)
17041 .arg(kb)
17042 .arg(vref)
17043 .arg(o)
17044 .arg(&hd)
17045 .arg(&nh)
17046 .arg(&nhkv)
17047 .arg(&ti)
17048 .arg(&tkvi)
17049 .arg(&scale)
17050 .arg(&cz);
17051 unsafe {
17052 b.launch(cfg)?;
17053 }
17054 Ok(())
17055 }
17056
17057 #[allow(clippy::too_many_arguments)]
17060 pub fn fa_prefill_hd512_arm(
17061 &self,
17062 q: &CudaSlice<f32>,
17063 k: &CudaSlice<f32>,
17064 v: &CudaSlice<f32>,
17065 o: &mut CudaSlice<f32>,
17066 head_dim: usize,
17067 n_head: usize,
17068 n_head_kv: usize,
17069 t: usize,
17070 t_kv: usize,
17071 scale: f32,
17072 causal: bool,
17073 f32_stage: bool,
17074 sp: bool,
17075 f16pv: bool,
17076 ) -> Result<(), Box<dyn std::error::Error>> {
17077 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
17078 if sp && !f32_stage {
17079 const SP_M: usize = 16;
17083 const BKS: usize = 32;
17084 let nw = if f16pv { fa512_wide_warps() } else { 2 };
17085 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
17086 let f = self.func(if hp {
17087 "fa_prefill_bf16_hd512_sp16h2"
17088 } else {
17089 match (f16pv, nw) {
17090 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
17091 (true, _) => "fa_prefill_bf16_hd512_sp16",
17092 _ => "fa_prefill_bf16_hd512_sp",
17093 }
17094 });
17095 let (nwarp, npart) = if hp {
17096 (4usize, 4usize)
17097 } else if nw > 2 {
17098 (nw, nw)
17099 } else {
17100 (2, 1)
17101 };
17102 let shmem = if hp {
17103 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
17104 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
17105 } else {
17106 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
17107 + 4 * (npart * SP_M * BKS + SP_M)) as u32
17108 };
17109 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17110 f.set_attribute(
17111 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17112 shmem as i32,
17113 )?;
17114 let grid_y = if hp {
17115 (n_head / 2) as u32
17116 } else {
17117 n_head as u32
17118 };
17119 let cfg = LaunchConfig {
17120 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
17121 block_dim: (32, nwarp as u32, 1),
17122 shared_mem_bytes: shmem,
17123 };
17124 let (hd, nh, nhkv, ti, tkvi, cz) = (
17125 head_dim as i32,
17126 n_head as i32,
17127 n_head_kv as i32,
17128 t as i32,
17129 t_kv as i32,
17130 causal as i32,
17131 );
17132 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
17133 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
17134 let vb = if f16pv {
17135 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
17136 } else {
17137 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
17138 };
17139 let __s_b = self.gpu.stream();
17140 let mut b = __s_b.launch_builder(&f);
17141 b.arg(&qb)
17142 .arg(&kb)
17143 .arg(&vb)
17144 .arg(o)
17145 .arg(&hd)
17146 .arg(&nh)
17147 .arg(&nhkv)
17148 .arg(&ti)
17149 .arg(&tkvi)
17150 .arg(&scale)
17151 .arg(&cz);
17152 unsafe {
17153 b.launch(cfg)?;
17154 }
17155 return Ok(());
17156 }
17157 const BLOCK_Q: usize = 32;
17158 const BK: usize = 32;
17159 const HALF: usize = 256;
17160 let f = self.func(if f32_stage {
17161 "fa_prefill_f32_hd512"
17162 } else {
17163 "fa_prefill_bf16_hd512"
17164 });
17165 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
17167 + 4 * BLOCK_Q) as u32;
17168 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17169 f.set_attribute(
17170 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17171 shmem as i32,
17172 )?;
17173 let cfg = LaunchConfig {
17174 grid_dim: (
17175 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17176 n_head as u32,
17177 2,
17178 ),
17179 block_dim: (32, 2, 1),
17180 shared_mem_bytes: shmem,
17181 };
17182 let (hd, nh, nhkv, ti, tkvi, cz) = (
17183 head_dim as i32,
17184 n_head as i32,
17185 n_head_kv as i32,
17186 t as i32,
17187 t_kv as i32,
17188 causal as i32,
17189 );
17190 if f32_stage {
17191 let __s_b = self.gpu.stream();
17192 let mut b = __s_b.launch_builder(&f);
17193 b.arg(q)
17194 .arg(k)
17195 .arg(v)
17196 .arg(o)
17197 .arg(&hd)
17198 .arg(&nh)
17199 .arg(&nhkv)
17200 .arg(&ti)
17201 .arg(&tkvi)
17202 .arg(&scale)
17203 .arg(&cz);
17204 unsafe {
17205 b.launch(cfg)?;
17206 }
17207 } else {
17208 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
17209 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
17210 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
17211 let __s_b = self.gpu.stream();
17212 let mut b = __s_b.launch_builder(&f);
17213 b.arg(&qb)
17214 .arg(&kb)
17215 .arg(&vb)
17216 .arg(o)
17217 .arg(&hd)
17218 .arg(&nh)
17219 .arg(&nhkv)
17220 .arg(&ti)
17221 .arg(&tkvi)
17222 .arg(&scale)
17223 .arg(&cz);
17224 unsafe {
17225 b.launch(cfg)?;
17226 }
17227 }
17228 Ok(())
17229 }
17230
17231 #[allow(clippy::too_many_arguments)]
17235 pub fn rope_neox2_bf16e(
17236 &self,
17237 q: &mut CudaSlice<f32>,
17238 k: &mut CudaSlice<f32>,
17239 qb: &mut CudaSlice<u8>,
17240 kb: &mut CudaSlice<u8>,
17241 pos: &CudaSlice<i32>,
17242 head_dim: usize,
17243 n_dims: usize,
17244 nh_q: usize,
17245 nh_k: usize,
17246 n_tokens: usize,
17247 base: f32,
17248 freq_scale: f32,
17249 ff: Option<&CudaSlice<f32>>,
17250 ) -> Result<(), Box<dyn std::error::Error>> {
17251 let f = self.func("rope_neox2_bf16e_f32");
17252 let rows = ((nh_q + nh_k) * n_tokens) as u32;
17253 let cfg = LaunchConfig {
17254 grid_dim: (rows, 1, 1),
17255 block_dim: ((head_dim / 2) as u32, 1, 1),
17256 shared_mem_bytes: 0,
17257 };
17258 let theta_scale = base.powf(-2.0 / n_dims as f32);
17259 let (hd, nd, nhq, nhk, nt) = (
17260 head_dim as i32,
17261 n_dims as i32,
17262 nh_q as i32,
17263 nh_k as i32,
17264 n_tokens as i32,
17265 );
17266 let __s_b = self.gpu.stream();
17267 let mut b = __s_b.launch_builder(&f);
17268 match ff {
17269 Some(t) => {
17270 b.arg(&mut *q)
17271 .arg(&mut *k)
17272 .arg(&mut *qb)
17273 .arg(&mut *kb)
17274 .arg(pos)
17275 .arg(&hd)
17276 .arg(&nd)
17277 .arg(&nhq)
17278 .arg(&nhk)
17279 .arg(&nt)
17280 .arg(&theta_scale)
17281 .arg(&freq_scale)
17282 .arg(t);
17283 unsafe {
17284 b.launch(cfg)?;
17285 }
17286 }
17287 None => {
17288 let null: u64 = 0;
17289 b.arg(&mut *q)
17290 .arg(&mut *k)
17291 .arg(&mut *qb)
17292 .arg(&mut *kb)
17293 .arg(pos)
17294 .arg(&hd)
17295 .arg(&nd)
17296 .arg(&nhq)
17297 .arg(&nhk)
17298 .arg(&nt)
17299 .arg(&theta_scale)
17300 .arg(&freq_scale)
17301 .arg(&null);
17302 unsafe {
17303 b.launch(cfg)?;
17304 }
17305 }
17306 }
17307 Ok(())
17308 }
17309
17310 pub fn f32_to_bf16(
17313 &self,
17314 x: &CudaSlice<f32>,
17315 n: usize,
17316 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17317 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
17318 let mut y = self.alloc_uninit::<u8>(n * 2)?;
17319 let f = self.func("f32_to_bf16_flat");
17320 let n_i = n as i64;
17321 let cfg = LaunchConfig {
17322 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
17323 block_dim: (256, 1, 1),
17324 shared_mem_bytes: 0,
17325 };
17326 let __s_b = self.gpu.stream();
17327 let mut b = __s_b.launch_builder(&f);
17328 b.arg(x).arg(&mut y).arg(&n_i);
17329 unsafe {
17330 b.launch(cfg)?;
17331 }
17332 Ok(y)
17333 }
17334
17335 pub fn f32_to_f16(
17336 &self,
17337 x: &CudaSlice<f32>,
17338 n: usize,
17339 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17340 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
17341 let mut y = self.alloc_uninit::<u8>(n * 2)?;
17342 let f = self.func("f32_to_f16_flat");
17343 let n_i = n as i64;
17344 let cfg = LaunchConfig {
17345 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
17346 block_dim: (256, 1, 1),
17347 shared_mem_bytes: 0,
17348 };
17349 let __s_b = self.gpu.stream();
17350 let mut b = __s_b.launch_builder(&f);
17351 b.arg(x).arg(&mut y).arg(&n_i);
17352 unsafe {
17353 b.launch(cfg)?;
17354 }
17355 Ok(y)
17356 }
17357
17358 pub fn bf16_to_f16(
17360 &self,
17361 xb: &CudaSlice<u8>,
17362 n: usize,
17363 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17364 let mut y = self.alloc_uninit::<u8>(n * 2)?;
17365 self.bf16_to_f16_into(xb, n, &mut y)?;
17366 Ok(y)
17367 }
17368
17369 pub fn bf16_to_f16_into(
17371 &self,
17372 xb: &CudaSlice<u8>,
17373 n: usize,
17374 y: &mut CudaSlice<u8>,
17375 ) -> Result<(), Box<dyn std::error::Error>> {
17376 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
17377 assert!(y.len() >= n * 2);
17378 let f = self.func("bf16_to_f16_flat");
17379 let n2 = (n / 2) as i64;
17380 let cfg = LaunchConfig {
17381 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
17382 block_dim: (256, 1, 1),
17383 shared_mem_bytes: 0,
17384 };
17385 let __s_b = self.gpu.stream();
17386 let mut b = __s_b.launch_builder(&f);
17387 b.arg(xb).arg(y).arg(&n2);
17388 unsafe {
17389 b.launch(cfg)?;
17390 }
17391 Ok(())
17392 }
17393
17394 #[allow(clippy::too_many_arguments)]
17399 pub fn fa_prefill_vl8(
17400 &self,
17401 seqs: &[FaSeqVl],
17402 head_dim: usize,
17403 n_head: usize,
17404 n_head_kv: usize,
17405 scale: f32,
17406 ) -> Result<(), Box<dyn std::error::Error>> {
17407 const BK: usize = 32;
17408 let b = seqs.len();
17409 assert!(b >= 1 && b <= 8);
17410 let mut packed = [FaSeqVl::default(); 8];
17411 packed[..b].copy_from_slice(seqs);
17412 let v = FaVl8(packed);
17413 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17414 let ept = (n_head_kv * head_dim) as i32;
17415 {
17416 let f = self.func("fa_mirror_vl");
17417 let max_n = (max_t as i64) * ept as i64;
17418 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
17419 for which in 0..2i32 {
17420 let cfg = LaunchConfig {
17421 grid_dim: (blocks, 1, b as u32),
17422 block_dim: (256, 1, 1),
17423 shared_mem_bytes: 0,
17424 };
17425 let __s_lb = self.gpu.stream();
17426 let mut lb = __s_lb.launch_builder(&f);
17427 lb.arg(&v).arg(&ept).arg(&which);
17428 unsafe {
17429 lb.launch(cfg)?;
17430 }
17431 }
17432 }
17433 let hd_sfx = fa_hd_suffix(head_dim)?;
17434 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
17435 let block_q = 64usize;
17436 let kv_stages = 2usize;
17437 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
17438 + 4 * (block_q * BK + 2 * block_q)) as u32;
17439 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17440 f.set_attribute(
17441 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17442 shmem as i32,
17443 )?;
17444 let cfg = LaunchConfig {
17445 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
17446 block_dim: (32, 4, 1),
17447 shared_mem_bytes: shmem,
17448 };
17449 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17450 let __s_lb = self.gpu.stream();
17451 let mut lb = __s_lb.launch_builder(&f);
17452 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
17453 unsafe {
17454 lb.launch(cfg)?;
17455 }
17456 Ok(())
17457 }
17458
17459 #[allow(clippy::too_many_arguments)]
17463 pub fn attn_pre_vl8(
17464 &self,
17465 seqs: &[AttnPreVl],
17466 wq: &CudaSlice<f32>,
17467 wk: &CudaSlice<f32>,
17468 head_dim: usize,
17469 rope_dims: usize,
17470 n_head: usize,
17471 n_head_kv: usize,
17472 eps: f32,
17473 freq_base: f32,
17474 freq_scale: f32,
17475 kv_dim_k: usize,
17476 kv_dim_v: usize,
17477 k_tok_bytes: usize,
17478 v_tok_bytes: usize,
17479 ) -> Result<(), Box<dyn std::error::Error>> {
17480 let b = seqs.len();
17481 assert!(b >= 1 && b <= 8);
17482 let mut packed = [AttnPreVl::default(); 8];
17483 packed[..b].copy_from_slice(seqs);
17484 let v = AttnPreVl8(packed);
17485 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17486 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17487 {
17488 let f = self.func("q_gate_split_vl");
17489 let n = max_t * (n_head * head_dim) as u32;
17490 let cfg = LaunchConfig {
17491 grid_dim: (n.div_ceil(256), 1, b as u32),
17492 block_dim: (256, 1, 1),
17493 shared_mem_bytes: 0,
17494 };
17495 let __s_lb = self.gpu.stream();
17496 let mut lb = __s_lb.launch_builder(&f);
17497 lb.arg(&v).arg(&hd).arg(&nh);
17498 unsafe {
17499 lb.launch(cfg)?;
17500 }
17501 }
17502 {
17503 let f = self.func("attn_rms_vl");
17504 let cfg = LaunchConfig {
17505 grid_dim: (max_t * n_head as u32, 2, b as u32),
17506 block_dim: (rms_block(), 1, 1),
17507 shared_mem_bytes: 0,
17508 };
17509 let __s_lb = self.gpu.stream();
17510 let mut lb = __s_lb.launch_builder(&f);
17511 lb.arg(&v)
17512 .arg(wq)
17513 .arg(wk)
17514 .arg(&hd)
17515 .arg(&nh)
17516 .arg(&nhkv)
17517 .arg(&eps);
17518 unsafe {
17519 lb.launch(cfg)?;
17520 }
17521 }
17522 {
17523 let f = self.func("attn_rope_vl");
17524 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
17525 let nd = rope_dims as i32;
17526 let cfg = LaunchConfig {
17527 grid_dim: (max_t * n_head as u32, 2, b as u32),
17528 block_dim: ((head_dim / 2) as u32, 1, 1),
17529 shared_mem_bytes: 0,
17530 };
17531 let __s_lb = self.gpu.stream();
17532 let mut lb = __s_lb.launch_builder(&f);
17533 lb.arg(&v)
17534 .arg(&hd)
17535 .arg(&nd)
17536 .arg(&nh)
17537 .arg(&nhkv)
17538 .arg(&theta_scale)
17539 .arg(&freq_scale);
17540 unsafe {
17541 lb.launch(cfg)?;
17542 }
17543 }
17544 {
17545 let f = self.func("append_kv_vl");
17546 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
17547 let cfg = LaunchConfig {
17548 grid_dim: (nblk, max_t, b as u32),
17549 block_dim: (32, 1, 1),
17550 shared_mem_bytes: 0,
17551 };
17552 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17553 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17554 let __s_lb = self.gpu.stream();
17555 let mut lb = __s_lb.launch_builder(&f);
17556 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
17557 unsafe {
17558 lb.launch(cfg)?;
17559 }
17560 }
17561 Ok(())
17562 }
17563
17564 pub fn fa_prefill_view(
17569 &self,
17570 q: &CudaSlice<f32>,
17571 k: &cudarc::driver::CudaView<u8>,
17572 v: &cudarc::driver::CudaView<u8>,
17573 o: &mut CudaSlice<f32>,
17574 head_dim: usize,
17575 n_head: usize,
17576 n_head_kv: usize,
17577 t: usize,
17578 t_kv: usize,
17579 scale: f32,
17580 causal: bool,
17581 k_tok_bytes: usize,
17582 v_tok_bytes: usize,
17583 g: bool,
17584 ) -> Result<(), Box<dyn std::error::Error>> {
17585 if portable_mma_gated() {
17586 return self.sdpa_naive_quantized_view(
17587 q,
17588 k,
17589 v,
17590 o,
17591 head_dim,
17592 n_head,
17593 n_head_kv,
17594 t,
17595 t_kv,
17596 scale,
17597 causal,
17598 k_tok_bytes,
17599 v_tok_bytes,
17600 );
17601 }
17602 const BLOCK_Q: usize = 64;
17603 const BK: usize = 32;
17604 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
17607 let f = if g {
17608 self.func_g(&name)
17609 } else {
17610 self.func(&name)
17611 };
17612 let shmem =
17613 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
17614 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17615 f.set_attribute(
17616 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17617 shmem as i32,
17618 )?;
17619 let cfg = LaunchConfig {
17620 grid_dim: (
17621 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17622 n_head as u32,
17623 1,
17624 ),
17625 block_dim: (32, 4, 1),
17626 shared_mem_bytes: shmem,
17627 };
17628 let (hd, nh, nhkv, ti, tkvi, cz) = (
17629 head_dim as i32,
17630 n_head as i32,
17631 n_head_kv as i32,
17632 t as i32,
17633 t_kv as i32,
17634 causal as i32,
17635 );
17636 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17637 let __s_b = self.gpu.stream();
17638 let mut b = __s_b.launch_builder(&f);
17639 b.arg(q)
17640 .arg(k)
17641 .arg(v)
17642 .arg(o)
17643 .arg(&hd)
17644 .arg(&nh)
17645 .arg(&nhkv)
17646 .arg(&ti)
17647 .arg(&tkvi)
17648 .arg(&scale)
17649 .arg(&cz)
17650 .arg(&ktb)
17651 .arg(&vtb);
17652 unsafe {
17653 b.launch(cfg)?;
17654 }
17655 Ok(())
17656 }
17657
17658 #[allow(clippy::too_many_arguments)]
17668 pub fn fa_prefill_view_ws(
17669 &self,
17670 q: &CudaSlice<f32>,
17671 k: &cudarc::driver::CudaView<u8>,
17672 v: &cudarc::driver::CudaView<u8>,
17673 o: &mut CudaSlice<f32>,
17674 head_dim: usize,
17675 n_head: usize,
17676 n_head_kv: usize,
17677 t: usize,
17678 t_kv: usize,
17679 scale: f32,
17680 causal: bool,
17681 k_tok_bytes: usize,
17682 v_tok_bytes: usize,
17683 g: bool,
17684 ) -> Result<(), Box<dyn std::error::Error>> {
17685 if portable_mma_gated() {
17686 return self.sdpa_naive_quantized_view(
17687 q,
17688 k,
17689 v,
17690 o,
17691 head_dim,
17692 n_head,
17693 n_head_kv,
17694 t,
17695 t_kv,
17696 scale,
17697 causal,
17698 k_tok_bytes,
17699 v_tok_bytes,
17700 );
17701 }
17702 const BLOCK_Q: usize = 64;
17703 const BK: usize = 32;
17704 let kv_dim_k = n_head_kv * head_dim;
17705 let kv_dim_v = n_head_kv * head_dim;
17706 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17708 let mut guard = self.prime_deqw_ws.lock().unwrap();
17710 let need_grow = match guard.as_ref() {
17711 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17712 None => true,
17713 };
17714 if need_grow {
17715 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17716 let (ck, cv) = guard
17717 .as_ref()
17718 .map(|(a, b)| (a.len(), b.len()))
17719 .unwrap_or((0, 0));
17720 *guard = Some((
17721 self.alloc_u8(grow(ck, k_ws_bytes))?,
17722 self.alloc_u8(grow(cv, v_ws_bytes))?,
17723 ));
17724 }
17725 let (kw, vw) = guard.as_mut().unwrap();
17726 {
17728 let f = if g {
17730 self.func_g("fa_dequant_kv_ws_bf16")
17731 } else {
17732 self.func("fa_dequant_kv_ws_bf16")
17733 };
17734 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17735 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17736 let cfg = LaunchConfig {
17737 grid_dim: (nblk.max(1), 1, 1),
17738 block_dim: (256, 1, 1),
17739 shared_mem_bytes: 0,
17740 };
17741 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17742 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17743 let __s_b = self.gpu.stream();
17744 let mut b = __s_b.launch_builder(&f);
17745 b.arg(k)
17746 .arg(v)
17747 .arg(&mut *kw)
17748 .arg(&mut *vw)
17749 .arg(&kdk)
17750 .arg(&kdv)
17751 .arg(&tkvi)
17752 .arg(&ktb)
17753 .arg(&vtb);
17754 unsafe {
17755 b.launch(cfg)?;
17756 }
17757 }
17758 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17766 .map(|v| v != "0")
17767 .unwrap_or(true);
17768 {
17769 let hd_sfx = fa_hd_suffix(head_dim)?;
17770 let f = self.func(&format!(
17771 "fa_prefill_qw{}{hd_sfx}",
17772 if db { "_db" } else { "" }
17773 ));
17774 let shmem = if db {
17775 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17777 } else {
17778 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17779 };
17780 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17781 f.set_attribute(
17782 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17783 shmem as i32,
17784 )?;
17785 let cfg = LaunchConfig {
17786 grid_dim: (
17787 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17788 n_head as u32,
17789 1,
17790 ),
17791 block_dim: (32, 4, 1),
17792 shared_mem_bytes: shmem,
17793 };
17794 let (hd, nh, nhkv, ti, tkvi, cz) = (
17795 head_dim as i32,
17796 n_head as i32,
17797 n_head_kv as i32,
17798 t as i32,
17799 t_kv as i32,
17800 causal as i32,
17801 );
17802 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17803 let __s_b = self.gpu.stream();
17804 let mut b = __s_b.launch_builder(&f);
17805 b.arg(q)
17806 .arg(&*kw)
17807 .arg(&*vw)
17808 .arg(o)
17809 .arg(&hd)
17810 .arg(&nh)
17811 .arg(&nhkv)
17812 .arg(&ti)
17813 .arg(&tkvi)
17814 .arg(&scale)
17815 .arg(&cz)
17816 .arg(&kdk)
17817 .arg(&kdv);
17818 unsafe {
17819 b.launch(cfg)?;
17820 }
17821 }
17822 Ok(())
17823 }
17824
17825 #[allow(clippy::too_many_arguments)]
17841 pub fn fa_prefill_view_ws_w_hd128(
17842 &self,
17843 q: &CudaSlice<f32>,
17844 k: &cudarc::driver::CudaView<u8>,
17845 v: &cudarc::driver::CudaView<u8>,
17846 o: &mut CudaSlice<f32>,
17847 head_dim: usize,
17848 n_head: usize,
17849 n_head_kv: usize,
17850 t: usize,
17851 t_kv: usize,
17852 scale: f32,
17853 causal: bool,
17854 window: usize,
17855 k_tok_bytes: usize,
17856 v_tok_bytes: usize,
17857 ) -> Result<(), Box<dyn std::error::Error>> {
17858 assert_eq!(
17859 head_dim, 128,
17860 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17861 );
17862 if portable_mma_gated() {
17863 return self.sdpa_naive_w_quantized_view(
17864 q,
17865 k,
17866 v,
17867 o,
17868 head_dim,
17869 n_head,
17870 n_head_kv,
17871 t,
17872 t_kv,
17873 scale,
17874 causal,
17875 window,
17876 k_tok_bytes,
17877 v_tok_bytes,
17878 );
17879 }
17880 const BLOCK_Q: usize = 64;
17881 const BK: usize = 32;
17882 let kv_dim_k = n_head_kv * head_dim;
17883 let kv_dim_v = n_head_kv * head_dim;
17884 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17886 let mut guard = self.prime_deqw_ws.lock().unwrap();
17887 let need_grow = match guard.as_ref() {
17888 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17889 None => true,
17890 };
17891 if need_grow {
17892 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17893 let (ck, cv) = guard
17894 .as_ref()
17895 .map(|(a, b)| (a.len(), b.len()))
17896 .unwrap_or((0, 0));
17897 *guard = Some((
17898 self.alloc_u8(grow(ck, k_ws_bytes))?,
17899 self.alloc_u8(grow(cv, v_ws_bytes))?,
17900 ));
17901 }
17902 let (kw, vw) = guard.as_mut().unwrap();
17903 {
17906 let f = self.func("fa_dequant_kv_ws_bf16");
17907 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17908 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17909 let cfg = LaunchConfig {
17910 grid_dim: (nblk.max(1), 1, 1),
17911 block_dim: (256, 1, 1),
17912 shared_mem_bytes: 0,
17913 };
17914 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17915 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17916 let __s_b = self.gpu.stream();
17917 let mut b = __s_b.launch_builder(&f);
17918 b.arg(k)
17919 .arg(v)
17920 .arg(&mut *kw)
17921 .arg(&mut *vw)
17922 .arg(&kdk)
17923 .arg(&kdv)
17924 .arg(&tkvi)
17925 .arg(&ktb)
17926 .arg(&vtb);
17927 unsafe {
17928 b.launch(cfg)?;
17929 }
17930 }
17931 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17933 .map(|v| v != "0")
17934 .unwrap_or(true);
17935 {
17936 let f = self.func(if db {
17937 "fa_prefill_qw_db_w_hd128"
17938 } else {
17939 "fa_prefill_qw_w_hd128"
17940 });
17941 let shmem = if db {
17942 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17943 } else {
17944 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17945 };
17946 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17947 f.set_attribute(
17948 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17949 shmem as i32,
17950 )?;
17951 let cfg = LaunchConfig {
17952 grid_dim: (
17953 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17954 n_head as u32,
17955 1,
17956 ),
17957 block_dim: (32, 4, 1),
17958 shared_mem_bytes: shmem,
17959 };
17960 let (hd, nh, nhkv, ti, tkvi, cz) = (
17961 head_dim as i32,
17962 n_head as i32,
17963 n_head_kv as i32,
17964 t as i32,
17965 t_kv as i32,
17966 causal as i32,
17967 );
17968 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17969 let __s_b = self.gpu.stream();
17970 let mut b = __s_b.launch_builder(&f);
17971 b.arg(q)
17972 .arg(&*kw)
17973 .arg(&*vw)
17974 .arg(o)
17975 .arg(&hd)
17976 .arg(&nh)
17977 .arg(&nhkv)
17978 .arg(&ti)
17979 .arg(&tkvi)
17980 .arg(&scale)
17981 .arg(&cz)
17982 .arg(&kdk)
17983 .arg(&kdv)
17984 .arg(&wnd);
17985 unsafe {
17986 b.launch(cfg)?;
17987 }
17988 }
17989 Ok(())
17990 }
17991
17992 pub fn fa_decode(
17996 &self,
17997 q: &CudaSlice<f32>,
17998 k: &cudarc::driver::CudaView<u8>,
17999 v: &cudarc::driver::CudaView<u8>,
18000 o: &mut CudaSlice<f32>,
18001 head_dim: usize,
18002 n_head: usize,
18003 n_head_kv: usize,
18004 t_kv: usize,
18005 scale: f32,
18006 k_tok_bytes: usize,
18007 v_tok_bytes: usize,
18008 ) -> Result<(), Box<dyn std::error::Error>> {
18009 self.fa_decode_kvmod(
18010 q,
18011 k,
18012 v,
18013 o,
18014 head_dim,
18015 n_head,
18016 n_head_kv,
18017 t_kv,
18018 scale,
18019 k_tok_bytes,
18020 v_tok_bytes,
18021 false,
18022 )
18023 }
18024
18025 #[allow(clippy::too_many_arguments)]
18029 #[allow(clippy::too_many_arguments)]
18033 #[allow(clippy::too_many_arguments)]
18034 fn fa_decode_scalar_unified(
18035 &self,
18036 q: &cudarc::driver::CudaView<f32>,
18037 k: &cudarc::driver::CudaView<u8>,
18038 v: &cudarc::driver::CudaView<u8>,
18039 o: &mut cudarc::driver::CudaViewMut<f32>,
18040 head_dim: usize,
18041 n_head: usize,
18042 n_head_kv: usize,
18043 t_kv_host: usize,
18044 t_kv_dev: Option<&CudaSlice<i32>>,
18045 scale: f32,
18046 n_splits: usize,
18047 split_keys: usize,
18048 k_tok_bytes: usize,
18049 v_tok_bytes: usize,
18050 g: bool,
18051 part_o: &mut CudaSlice<f32>,
18052 part_m: &mut CudaSlice<f32>,
18053 part_l: &mut CudaSlice<f32>,
18054 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18055 ) -> Result<(), Box<dyn std::error::Error>> {
18056 let f = if g {
18057 self.func_g("fa_decode_f32")
18058 } else {
18059 self.fa_func("fa_decode_f32", head_dim)
18060 };
18061 let cfg = LaunchConfig {
18062 grid_dim: (n_head as u32, n_splits as u32, 1),
18063 block_dim: (head_dim as u32, 1, 1),
18064 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
18065 };
18066 let (hd, nh, nhkv, nsp) = (
18067 head_dim as i32,
18068 n_head as i32,
18069 n_head_kv as i32,
18070 n_splits as i32,
18071 );
18072 let (ktb, vtb, tkvi, ski) = (
18073 k_tok_bytes as i64,
18074 v_tok_bytes as i64,
18075 t_kv_host as i32,
18076 split_keys as i32,
18077 );
18078 let __s_b = self.gpu.stream();
18079 let mut b = __s_b.launch_builder(&f);
18080 match t_kv_dev {
18081 Some(d) => {
18082 b.arg(q)
18083 .arg(k)
18084 .arg(v)
18085 .arg(&mut *part_o)
18086 .arg(&mut *part_m)
18087 .arg(&mut *part_l)
18088 .arg(&hd)
18089 .arg(&nh)
18090 .arg(&nhkv)
18091 .arg(&tkvi)
18092 .arg(d)
18093 .arg(&scale)
18094 .arg(&nsp)
18095 .arg(&ski)
18096 .arg(&ktb)
18097 .arg(&vtb);
18098 unsafe {
18099 b.launch(cfg)?;
18100 }
18101 }
18102 None => {
18103 let null: u64 = 0;
18104 b.arg(q)
18105 .arg(k)
18106 .arg(v)
18107 .arg(&mut *part_o)
18108 .arg(&mut *part_m)
18109 .arg(&mut *part_l)
18110 .arg(&hd)
18111 .arg(&nh)
18112 .arg(&nhkv)
18113 .arg(&tkvi)
18114 .arg(&null)
18115 .arg(&scale)
18116 .arg(&nsp)
18117 .arg(&ski)
18118 .arg(&ktb)
18119 .arg(&vtb);
18120 unsafe {
18121 b.launch(cfg)?;
18122 }
18123 }
18124 }
18125 let cfg2 = LaunchConfig {
18126 grid_dim: (n_head as u32, 1, 1),
18127 block_dim: (head_dim as u32, 1, 1),
18128 shared_mem_bytes: 0,
18129 };
18130 if let Some((oq, od)) = q8_out {
18131 let fc = if g {
18133 self.func_g("fa_decode_combine_q8_1")
18134 } else {
18135 self.fa_func("fa_decode_combine_q8_1", head_dim)
18136 };
18137 let __s_b2 = self.gpu.stream();
18138 let mut b2 = __s_b2.launch_builder(&fc);
18139 b2.arg(&*part_o)
18140 .arg(&*part_m)
18141 .arg(&*part_l)
18142 .arg(oq)
18143 .arg(od)
18144 .arg(&hd)
18145 .arg(&nh)
18146 .arg(&nsp);
18147 unsafe {
18148 b2.launch(cfg2)?;
18149 }
18150 return Ok(());
18151 }
18152 let fc = if g {
18153 self.func_g("fa_decode_combine_f32")
18154 } else {
18155 self.fa_func("fa_decode_combine_f32", head_dim)
18156 };
18157 let __s_b2 = self.gpu.stream();
18158 let mut b2 = __s_b2.launch_builder(&fc);
18159 b2.arg(&*part_o)
18160 .arg(&*part_m)
18161 .arg(&*part_l)
18162 .arg(o)
18163 .arg(&hd)
18164 .arg(&nh)
18165 .arg(&nsp);
18166 unsafe {
18167 b2.launch(cfg2)?;
18168 }
18169 Ok(())
18170 }
18171
18172 pub fn fa_decode_kvmod(
18173 &self,
18174 q: &CudaSlice<f32>,
18175 k: &cudarc::driver::CudaView<u8>,
18176 v: &cudarc::driver::CudaView<u8>,
18177 o: &mut CudaSlice<f32>,
18178 head_dim: usize,
18179 n_head: usize,
18180 n_head_kv: usize,
18181 t_kv: usize,
18182 scale: f32,
18183 k_tok_bytes: usize,
18184 v_tok_bytes: usize,
18185 g: bool,
18186 ) -> Result<(), Box<dyn std::error::Error>> {
18187 let q_view = q.as_view();
18188 let mut o_view = o.as_view_mut();
18189 self.fa_decode_kvmod_view(
18190 &q_view,
18191 k,
18192 v,
18193 &mut o_view,
18194 head_dim,
18195 n_head,
18196 n_head_kv,
18197 t_kv,
18198 scale,
18199 k_tok_bytes,
18200 v_tok_bytes,
18201 g,
18202 )
18203 }
18204
18205 #[allow(clippy::too_many_arguments)]
18210 pub fn fa_decode_kvmod_view(
18211 &self,
18212 q: &cudarc::driver::CudaView<f32>,
18213 k: &cudarc::driver::CudaView<u8>,
18214 v: &cudarc::driver::CudaView<u8>,
18215 o: &mut cudarc::driver::CudaViewMut<f32>,
18216 head_dim: usize,
18217 n_head: usize,
18218 n_head_kv: usize,
18219 t_kv: usize,
18220 scale: f32,
18221 k_tok_bytes: usize,
18222 v_tok_bytes: usize,
18223 g: bool,
18224 ) -> Result<(), Box<dyn std::error::Error>> {
18225 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
18246 if g && head_dim == 256 && !fa_v4_at(t_kv) {
18250 fa_vec = false;
18251 }
18252 let sp = fa_split_keys(t_kv, n_head_kv);
18253 let n_splits = if fa_vec {
18254 ((t_kv + sp - 1) / sp).max(1)
18255 } else {
18256 ((t_kv + 255) / 256).max(1)
18257 };
18258 let o_len = n_head * n_splits * head_dim;
18259 let ml_len = n_head * n_splits;
18260 let mut part_guard = self.fa_part_pool.lock().unwrap();
18261 if part_guard
18262 .as_ref()
18263 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18264 .unwrap_or(true)
18265 {
18266 let old = part_guard.take();
18277 let (co, cm) = old
18278 .as_ref()
18279 .map(|pp| (pp.0.len(), pp.1.len()))
18280 .unwrap_or((0, 0));
18281 if let Some(old) = old {
18282 self.fa_part_retired.lock().unwrap().push(old);
18283 }
18284 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18285 eprintln!(
18286 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18287 co, o_len, cm, ml_len
18288 );
18289 }
18290 *part_guard = Some((
18291 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18292 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18293 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18294 ));
18295 }
18296 let pg = part_guard.as_mut().unwrap();
18297 self.gpu
18298 .stream()
18299 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18300 self.gpu
18301 .stream()
18302 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18303 self.gpu
18304 .stream()
18305 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18306 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18307 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18308 let (hd, nh, nhkv, tkvi, nsp) = (
18309 head_dim as i32,
18310 n_head as i32,
18311 n_head_kv as i32,
18312 t_kv as i32,
18313 n_splits as i32,
18314 );
18315 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18316 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
18320 let fa512_min = fa512_min_tkv();
18325 let deep = fa_vec
18328 && head_dim == 256
18329 && fa_v4_at(t_kv)
18330 && !g
18331 && fa_deep_at(t_kv)
18332 && !matches!(fa_v4_mode(), "noB3" | "stage");
18333 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
18334 let gqa = (n_head / n_head_kv).max(1) as u32;
18337 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
18338 (
18339 fv,
18340 LaunchConfig {
18341 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18342 block_dim: (32, gqa, 1),
18343 shared_mem_bytes: 0,
18344 },
18345 )
18346 } else if fa_vec && head_dim <= 256 {
18347 let gqa = (n_head / n_head_kv).max(1) as u32;
18348 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18359 let smem_tkv = *SMEM_TKV.get_or_init(|| {
18360 std::env::var("MEMRA_FA_SMEM_TKV")
18361 .ok()
18362 .and_then(|v| v.parse().ok())
18363 .unwrap_or_else(|| {
18364 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18365 })
18366 });
18367 if fa_v4_at(t_kv) && head_dim == 256 {
18368 let v4name = match fa_v4_mode() {
18372 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
18375 _ => "fa_decode_vec_q_v4",
18376 };
18377 let fv = if g {
18378 self.func_g(v4name)
18379 } else {
18380 self.func(v4name)
18381 };
18382 let shmem = (if deep { 12160 } else { 11520 }
18385 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
18386 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18387 fv.set_attribute(
18388 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18389 shmem as i32,
18390 )?;
18391 (
18392 fv,
18393 LaunchConfig {
18394 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18395 block_dim: (32, gqa, 1),
18396 shared_mem_bytes: shmem,
18397 },
18398 )
18399 } else if fa_v3_active(head_dim) {
18400 let fv = if g {
18403 self.func_g("fa_decode_vec_q_v3")
18404 } else {
18405 self.func("fa_decode_vec_q_v3")
18406 };
18407 let shmem = (32 * head_dim * 2) as u32; (
18409 fv,
18410 LaunchConfig {
18411 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18412 block_dim: (32, gqa, 1),
18413 shared_mem_bytes: shmem,
18414 },
18415 )
18416 } else if fa_v2_on() {
18417 let fv = if g {
18421 self.func_g("fa_decode_vec_q_v2")
18422 } else {
18423 self.func("fa_decode_vec_q_v2")
18424 };
18425 let shmem = (2 * 32 * head_dim * 2) as u32; (
18427 fv,
18428 LaunchConfig {
18429 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18430 block_dim: (32, gqa, 1),
18431 shared_mem_bytes: shmem,
18432 },
18433 )
18434 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
18435 {
18436 let fv = if g {
18440 self.func_g("fa_decode_vec_q_smem")
18441 } else {
18442 self.func("fa_decode_vec_q_smem")
18443 };
18444 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
18446 fv.set_attribute(
18447 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18448 shmem as i32,
18449 )?;
18450 (
18451 fv,
18452 LaunchConfig {
18453 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18454 block_dim: (32, gqa, 1),
18455 shared_mem_bytes: shmem,
18456 },
18457 )
18458 } else {
18459 let fv = if g {
18462 self.func_g("fa_decode_vec_q")
18463 } else {
18464 self.func("fa_decode_vec_q")
18465 };
18466 (
18467 fv,
18468 LaunchConfig {
18469 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18470 block_dim: (32, gqa, 1),
18471 shared_mem_bytes: 0,
18472 },
18473 )
18474 }
18475 } else {
18476 return self.fa_decode_scalar_unified(
18479 q,
18480 k,
18481 v,
18482 o,
18483 head_dim,
18484 n_head,
18485 n_head_kv,
18486 t_kv,
18487 None,
18488 scale,
18489 n_splits,
18490 if fa_vec { sp } else { 256 },
18491 k_tok_bytes,
18492 v_tok_bytes,
18493 g,
18494 part_o,
18495 part_m,
18496 part_l,
18497 None,
18498 );
18499 };
18500 let __s_b = self.gpu.stream();
18501 let mut b = __s_b.launch_builder(&f);
18502 b.arg(q)
18503 .arg(k)
18504 .arg(v)
18505 .arg(&mut *part_o)
18506 .arg(&mut *part_m)
18507 .arg(&mut *part_l)
18508 .arg(&hd)
18509 .arg(&nh)
18510 .arg(&nhkv)
18511 .arg(&tkvi)
18512 .arg(&scale)
18513 .arg(&nsp)
18514 .arg(&ktb)
18515 .arg(&vtb);
18516 unsafe {
18517 b.launch(cfg)?;
18518 }
18519 let (fc, cfg2) = (
18522 if g {
18523 self.func_g("fa_decode_combine_f32")
18524 } else {
18525 self.fa_func("fa_decode_combine_f32", head_dim)
18526 },
18527 LaunchConfig {
18528 grid_dim: (n_head as u32, 1, 1),
18529 block_dim: (head_dim as u32, 1, 1),
18530 shared_mem_bytes: 0,
18531 },
18532 );
18533 let __s_b2 = self.gpu.stream();
18534 let mut b2 = __s_b2.launch_builder(&fc);
18535 b2.arg(&*part_o)
18536 .arg(&*part_m)
18537 .arg(&*part_l)
18538 .arg(o)
18539 .arg(&hd)
18540 .arg(&nh)
18541 .arg(&nsp);
18542 unsafe {
18543 b2.launch(cfg2)?;
18544 }
18545 Ok(())
18546 }
18547
18548 #[allow(clippy::too_many_arguments)]
18559 pub fn fa_decode_batch_seqs_v4(
18560 &self,
18561 q: &CudaSlice<f32>,
18562 kv_ptrs: &cudarc::driver::CudaView<u64>,
18563 pos_seq: &CudaSlice<i32>,
18564 o: &mut CudaSlice<f32>,
18565 head_dim: usize,
18566 n_head: usize,
18567 n_head_kv: usize,
18568 b_n: usize,
18569 t_kv_max: usize,
18570 scale: f32,
18571 split_keys: usize,
18572 k_tok_bytes: usize,
18573 v_tok_bytes: usize,
18574 ) -> Result<(), Box<dyn std::error::Error>> {
18575 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
18576 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
18577 let o_len = b_n * n_head * n_splits_max * head_dim;
18578 let ml_len = b_n * n_head * n_splits_max;
18579 let mut part_guard = self.fa_part_pool.lock().unwrap();
18580 if part_guard
18581 .as_ref()
18582 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18583 .unwrap_or(true)
18584 {
18585 let old = part_guard.take();
18596 let (co, cm) = old
18597 .as_ref()
18598 .map(|pp| (pp.0.len(), pp.1.len()))
18599 .unwrap_or((0, 0));
18600 if let Some(old) = old {
18601 self.fa_part_retired.lock().unwrap().push(old);
18602 }
18603 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18604 eprintln!(
18605 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18606 co, o_len, cm, ml_len
18607 );
18608 }
18609 *part_guard = Some((
18610 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18611 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18612 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18613 ));
18614 }
18615 let pg = part_guard.as_mut().unwrap();
18616 self.gpu
18617 .stream()
18618 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18619 self.gpu
18620 .stream()
18621 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18622 self.gpu
18623 .stream()
18624 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18625 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18626 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18627 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
18628 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18629 let gqa = (n_head / n_head_kv).max(1) as u32;
18630 let f = self.func("fa_decode_vec_q_seqs_v4");
18631 let shmem = (11520 + 32 * head_dim * 2) as u32;
18633 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18634 f.set_attribute(
18635 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18636 shmem as i32,
18637 )?;
18638 let cfg = LaunchConfig {
18639 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
18640 block_dim: (32, gqa, 1),
18641 shared_mem_bytes: shmem,
18642 };
18643 {
18644 let __s_b = self.gpu.stream();
18645 let mut b = __s_b.launch_builder(&f);
18646 b.arg(q)
18647 .arg(kv_ptrs)
18648 .arg(pos_seq)
18649 .arg(&mut *part_o)
18650 .arg(&mut *part_m)
18651 .arg(&mut *part_l)
18652 .arg(&hd)
18653 .arg(&nh)
18654 .arg(&nhkv)
18655 .arg(&scale)
18656 .arg(&nspm)
18657 .arg(&spk)
18658 .arg(&ktb)
18659 .arg(&vtb);
18660 unsafe {
18661 b.launch(cfg)?;
18662 }
18663 }
18664 let fc = self.func("fa_decode_combine_seqs");
18665 let cfg2 = LaunchConfig {
18666 grid_dim: (n_head as u32, b_n as u32, 1),
18667 block_dim: (head_dim as u32, 1, 1),
18668 shared_mem_bytes: 0,
18669 };
18670 let __s_b2 = self.gpu.stream();
18671 let mut b2 = __s_b2.launch_builder(&fc);
18672 b2.arg(&*part_o)
18673 .arg(&*part_m)
18674 .arg(&*part_l)
18675 .arg(o)
18676 .arg(&hd)
18677 .arg(&nh)
18678 .arg(pos_seq)
18679 .arg(&nspm)
18680 .arg(&spk);
18681 unsafe {
18682 b2.launch(cfg2)?;
18683 }
18684 Ok(())
18685 }
18686
18687 #[allow(clippy::too_many_arguments)]
18694 pub fn append_kv_quantized_seqs(
18695 &self,
18696 k_rows: &CudaSlice<f32>,
18697 v_rows: &CudaSlice<f32>,
18698 kv_ptrs: &cudarc::driver::CudaView<u64>,
18699 pos_seq: &CudaSlice<i32>,
18700 b_n: usize,
18701 kv_dim_k: usize,
18702 kv_dim_v: usize,
18703 k_tok_bytes: usize,
18704 v_tok_bytes: usize,
18705 ) -> Result<(), Box<dyn std::error::Error>> {
18706 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
18707 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
18708 let cfg = LaunchConfig {
18709 grid_dim: (nblk, b_n as u32, 1),
18710 block_dim: (32, 1, 1),
18711 shared_mem_bytes: 0,
18712 };
18713 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
18714 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18715 let __s_b = self.gpu.stream();
18716 let mut b = __s_b.launch_builder(&f);
18717 b.arg(k_rows)
18718 .arg(v_rows)
18719 .arg(kv_ptrs)
18720 .arg(pos_seq)
18721 .arg(&kdk)
18722 .arg(&kdv)
18723 .arg(&ktb)
18724 .arg(&vtb);
18725 unsafe {
18726 b.launch(cfg)?;
18727 }
18728 Ok(())
18729 }
18730
18731 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
18737 std::env::var("MEMRA_NO_FA_VEC").is_err()
18738 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
18739 && base_len + 1 >= fa_vec_min_tkv()
18740 && head_dim <= 256
18741 && head_dim % 32 == 0
18742 }
18743
18744 #[allow(clippy::too_many_arguments)]
18753 pub fn fa_decode_rows(
18754 &self,
18755 q: &CudaSlice<f32>,
18756 k: &cudarc::driver::CudaView<u8>,
18757 v: &cudarc::driver::CudaView<u8>,
18758 o: &mut CudaSlice<f32>,
18759 head_dim: usize,
18760 n_head: usize,
18761 n_head_kv: usize,
18762 base_len: usize,
18763 t: usize,
18764 scale: f32,
18765 k_tok_bytes: usize,
18766 v_tok_bytes: usize,
18767 base_dev: Option<(&CudaSlice<i32>, i32)>,
18771 kv_shared: bool,
18774 g: bool,
18778 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18781 ) -> Result<(), Box<dyn std::error::Error>> {
18782 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
18783 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
18790 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18791 let v = *SP512.get_or_init(|| {
18794 std::env::var("MEMRA_FA_SP512")
18795 .ok()
18796 .and_then(|x| x.parse().ok())
18797 .unwrap_or(0)
18798 });
18799 sp = if v >= 8 {
18800 v
18801 } else {
18802 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18803 };
18804 }
18805 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18806 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18807 let gqa = (n_head / n_head_kv).max(1) as u32;
18808 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
18819 groups.push((0, t, sp));
18820 } else {
18821 let mut r0 = 0usize;
18822 while r0 < t {
18823 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
18824 let mut r1 = r0 + 1;
18825 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18826 r1 += 1;
18827 }
18828 groups.push((r0, r1 - r0, sp_g));
18829 r0 = r1;
18830 }
18831 }
18832 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18836 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18837 std::env::var("MEMRA_FA_SMEM_TKV")
18838 .ok()
18839 .and_then(|v| v.parse().ok())
18840 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18841 });
18842 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18843 let v3 = fa_v3_active(head_dim);
18844 let smem_rows =
18845 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18846 let _ = kv_shared;
18851 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18854 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18868 let tb512 = head_dim == 512
18870 && sp <= 32
18871 && n_head / n_head_kv.max(1) <= 16
18872 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18873 let fname = if tb512 {
18874 "fa_decode_vec_q_rows_v4_512_tb"
18875 } else if i2 {
18876 "fa_decode_vec_q_rows_dpl16_i2"
18877 } else if head_dim == 512 {
18878 "fa_decode_vec_q_rows_dpl16"
18879 }
18880 else if v4 {
18882 "fa_decode_vec_q_rows_v4"
18883 } else if v3 {
18884 "fa_decode_vec_q_rows_v3"
18885 } else if fa_v2_on() {
18886 "fa_decode_vec_q_rows_v2"
18887 } else if smem_rows {
18888 "fa_decode_vec_q_rows_smem"
18889 } else {
18890 "fa_decode_vec_q_rows"
18891 };
18892 let f = if head_dim == 512 {
18893 self.fa_func(fname, head_dim)
18894 } else if g {
18895 self.func_g(if smem_rows {
18903 "fa_decode_vec_q_rows"
18904 } else {
18905 fname
18906 })
18907 } else {
18908 self.func(fname)
18909 };
18910 let shmem = if tb512 {
18911 let gk = Self::gkv_on();
18913 let sh =
18914 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18915 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18916 f.set_attribute(
18917 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18918 sh as i32,
18919 )?;
18920 sh
18921 } else if v4 || v3 || smem_rows || fa_v2_on() {
18922 let sh = (if v4 {
18924 11520 + 32 * head_dim * if g { 1 } else { 2 }
18925 } else if v3 {
18926 32 * head_dim * 2
18927 } else {
18928 2 * 32 * head_dim * 2
18929 }) as u32;
18930 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18931 f.set_attribute(
18932 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18933 sh as i32,
18934 )?;
18935 sh
18936 } else {
18937 0
18938 };
18939 for &(r0, t_g, sp_g) in &groups {
18943 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18944 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18945 let base_i = (base_len + r0) as i32;
18946 let o_len = t_g * n_head * n_splits_g * head_dim;
18947 let ml_len = t_g * n_head * n_splits_g;
18948 let mut part_guard = self.fa_part_pool.lock().unwrap();
18949 if part_guard
18950 .as_ref()
18951 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18952 .unwrap_or(true)
18953 {
18954 let old = part_guard.take();
18965 let (co, cm) = old
18966 .as_ref()
18967 .map(|pp| (pp.0.len(), pp.1.len()))
18968 .unwrap_or((0, 0));
18969 if let Some(old) = old {
18970 self.fa_part_retired.lock().unwrap().push(old);
18971 }
18972 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18973 eprintln!(
18974 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18975 co, o_len, cm, ml_len
18976 );
18977 }
18978 *part_guard = Some((
18979 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18980 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18981 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18982 ));
18983 }
18984 let pg = part_guard.as_mut().unwrap();
18985 self.gpu
18986 .stream()
18987 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18988 self.gpu
18989 .stream()
18990 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18991 self.gpu
18992 .stream()
18993 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18994 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18995 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18996 let qv = self.view(q, t * n_head * head_dim);
18997 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18998 let cfg = LaunchConfig {
18999 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
19000 block_dim: (32, gqa, 1),
19001 shared_mem_bytes: shmem,
19002 };
19003 {
19004 let __s_b = self.gpu.stream();
19005 let mut b = __s_b.launch_builder(&f);
19006 if tb512 {
19007 let (bd, plus) =
19009 base_dev.expect("hd512 rows twin requires a device base counter");
19010 let plus_g = plus + r0 as i32;
19011 let nr = t_g as i32;
19012 if Self::pdl_on() && Self::pdl_wb_on() {
19013 use cudarc::driver::{DevicePtr, DevicePtrMut};
19015 let s = &self.gpu.stream();
19016 let (pq, _b0) = q_g.device_ptr(s);
19017 let (pk, _b1) = k.device_ptr(s);
19018 let (pv, _b2) = v.device_ptr(s);
19019 let (po, _b3) = part_o.device_ptr_mut(s);
19020 let (pm, _b4) = part_m.device_ptr_mut(s);
19021 let (pl, _b5) = part_l.device_ptr_mut(s);
19022 let (pb, _b6) = bd.device_ptr(s);
19023 let mut ps = [
19024 &pq as *const _ as *mut std::ffi::c_void,
19025 &pk as *const _ as *mut _,
19026 &pv as *const _ as *mut _,
19027 &po as *const _ as *mut _,
19028 &pm as *const _ as *mut _,
19029 &pl as *const _ as *mut _,
19030 &hd as *const _ as *mut _,
19031 &nh as *const _ as *mut _,
19032 &nhkv as *const _ as *mut _,
19033 &pb as *const _ as *mut _,
19034 &plus_g as *const _ as *mut _,
19035 &scale as *const _ as *mut _,
19036 &nspm as *const _ as *mut _,
19037 &spk as *const _ as *mut _,
19038 &ktb as *const _ as *mut _,
19039 &vtb as *const _ as *mut _,
19040 &nr as *const _ as *mut _,
19041 ];
19042 unsafe {
19043 self.launch_pdl_flash(
19044 Self::gkv_on(),
19045 "fa_decode_vec_q_rows_v4_512_tb",
19046 (n_head_kv as u32, n_splits_g as u32, 1),
19047 (32, gqa, 1),
19048 shmem,
19049 &mut ps,
19050 )?;
19051 }
19052 } else {
19053 let cfg_tb = LaunchConfig {
19054 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
19055 block_dim: (32, gqa, 1),
19056 shared_mem_bytes: shmem,
19057 };
19058 b.arg(&q_g)
19059 .arg(k)
19060 .arg(v)
19061 .arg(&mut *part_o)
19062 .arg(&mut *part_m)
19063 .arg(&mut *part_l)
19064 .arg(&hd)
19065 .arg(&nh)
19066 .arg(&nhkv)
19067 .arg(bd)
19068 .arg(&plus_g)
19069 .arg(&scale)
19070 .arg(&nspm)
19071 .arg(&spk)
19072 .arg(&ktb)
19073 .arg(&vtb)
19074 .arg(&nr);
19075 unsafe {
19076 b.launch(cfg_tb)?;
19077 }
19078 }
19079 } else if head_dim == 512 {
19080 let (bd, plus) =
19081 base_dev.expect("hd512 rows twin requires a device base counter");
19082 let plus_g = plus + r0 as i32;
19083 b.arg(&q_g)
19084 .arg(k)
19085 .arg(v)
19086 .arg(&mut *part_o)
19087 .arg(&mut *part_m)
19088 .arg(&mut *part_l)
19089 .arg(&hd)
19090 .arg(&nh)
19091 .arg(&nhkv)
19092 .arg(bd)
19093 .arg(&plus_g)
19094 .arg(&scale)
19095 .arg(&nspm)
19096 .arg(&spk)
19097 .arg(&ktb)
19098 .arg(&vtb);
19099 unsafe {
19100 b.launch(cfg)?;
19101 }
19102 } else {
19103 b.arg(&q_g)
19104 .arg(k)
19105 .arg(v)
19106 .arg(&mut *part_o)
19107 .arg(&mut *part_m)
19108 .arg(&mut *part_l)
19109 .arg(&hd)
19110 .arg(&nh)
19111 .arg(&nhkv)
19112 .arg(&base_i)
19113 .arg(&scale)
19114 .arg(&nspm)
19115 .arg(&spk)
19116 .arg(&ktb)
19117 .arg(&vtb);
19118 unsafe {
19119 b.launch(cfg)?;
19120 }
19121 }
19122 }
19123 let cfg2 = LaunchConfig {
19124 grid_dim: (n_head as u32, t_g as u32, 1),
19125 block_dim: (head_dim as u32, 1, 1),
19126 shared_mem_bytes: 0,
19127 };
19128 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
19129 if head_dim == 512 {
19130 let (bd, plus) = base_dev.unwrap();
19133 let plus_g = plus + r0 as i32;
19134 if let Some((oq, od)) = q8_out.as_mut() {
19135 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
19137 if Self::pdl_on() && Self::pdl_wb_on() {
19138 use cudarc::driver::{DevicePtr, DevicePtrMut};
19140 let s = &self.gpu.stream();
19141 let (po, _g0) = part_o.device_ptr(s);
19142 let (pm, _g1) = part_m.device_ptr(s);
19143 let (pl, _g2) = part_l.device_ptr(s);
19144 let (pq, _g3) = oq.device_ptr_mut(s);
19145 let (pd, _g4) = od.device_ptr_mut(s);
19146 let (pb, _g5) = bd.device_ptr(s);
19147 let mut ps = [
19148 &po as *const _ as *mut std::ffi::c_void,
19149 &pm as *const _ as *mut _,
19150 &pl as *const _ as *mut _,
19151 &pq as *const _ as *mut _,
19152 &pd as *const _ as *mut _,
19153 &hd as *const _ as *mut _,
19154 &nh as *const _ as *mut _,
19155 &pb as *const _ as *mut _,
19156 &plus_g as *const _ as *mut _,
19157 &nspm as *const _ as *mut _,
19158 &spk as *const _ as *mut _,
19159 ];
19160 unsafe {
19161 self.launch_pdl_flash(
19162 Self::gkv_on(),
19163 "fa_decode_combine_rows_dc_q8_1",
19164 cfg2.grid_dim,
19165 cfg2.block_dim,
19166 0,
19167 &mut ps,
19168 )?;
19169 }
19170 continue;
19171 }
19172 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
19173 let __s_b2 = self.gpu.stream();
19174 let mut b2 = __s_b2.launch_builder(&fc);
19175 b2.arg(&*part_o)
19176 .arg(&*part_m)
19177 .arg(&*part_l)
19178 .arg(&mut **oq)
19179 .arg(&mut **od)
19180 .arg(&hd)
19181 .arg(&nh)
19182 .arg(bd)
19183 .arg(&plus_g)
19184 .arg(&nspm)
19185 .arg(&spk);
19186 unsafe {
19187 b2.launch(cfg2)?;
19188 }
19189 continue;
19190 }
19191 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
19192 let __s_b2 = self.gpu.stream();
19193 let mut b2 = __s_b2.launch_builder(&fc);
19194 b2.arg(&*part_o)
19195 .arg(&*part_m)
19196 .arg(&*part_l)
19197 .arg(&mut o_g)
19198 .arg(&hd)
19199 .arg(&nh)
19200 .arg(bd)
19201 .arg(&plus_g)
19202 .arg(&nspm)
19203 .arg(&spk);
19204 unsafe {
19205 b2.launch(cfg2)?;
19206 }
19207 } else {
19208 assert!(
19211 q8_out.is_none(),
19212 "rows q8 emit requires the hd512 dc combine"
19213 );
19214 let fc = self.func("fa_decode_combine_rows");
19215 let __s_b2 = self.gpu.stream();
19216 let mut b2 = __s_b2.launch_builder(&fc);
19217 b2.arg(&*part_o)
19218 .arg(&*part_m)
19219 .arg(&*part_l)
19220 .arg(&mut o_g)
19221 .arg(&hd)
19222 .arg(&nh)
19223 .arg(&base_i)
19224 .arg(&nspm)
19225 .arg(&spk);
19226 unsafe {
19227 b2.launch(cfg2)?;
19228 }
19229 }
19230 }
19231 Ok(())
19232 }
19233
19234 #[allow(clippy::too_many_arguments)]
19238 pub fn fa_decode_rows_w(
19239 &self,
19240 q: &CudaSlice<f32>,
19241 k: &cudarc::driver::CudaView<u8>,
19242 v: &cudarc::driver::CudaView<u8>,
19243 o: &mut CudaSlice<f32>,
19244 head_dim: usize,
19245 n_head: usize,
19246 n_head_kv: usize,
19247 base_dev: &CudaSlice<i32>,
19248 base_plus: i32,
19249 t: usize,
19250 scale: f32,
19251 window: usize,
19252 k_tok_bytes: usize,
19253 v_tok_bytes: usize,
19254 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19255 ) -> Result<(), Box<dyn std::error::Error>> {
19256 debug_assert!(head_dim == 256);
19261 let sp = {
19269 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19270 let v = *SPW.get_or_init(|| {
19271 std::env::var("MEMRA_FA_SPW")
19272 .ok()
19273 .and_then(|x| x.parse().ok())
19274 .unwrap_or(0)
19275 });
19276 if v >= 8 {
19277 v
19278 } else {
19279 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
19280 }
19281 };
19282 let n_splits_max = (window + sp - 1) / sp;
19283 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19284 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
19285 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19286 let gqa = (n_head / n_head_kv).max(1) as u32;
19287 let o_len = t * n_head * n_splits_max * head_dim;
19288 let ml_len = t * n_head * n_splits_max;
19289 let mut part_guard = self.fa_part_pool.lock().unwrap();
19290 if part_guard
19291 .as_ref()
19292 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19293 .unwrap_or(true)
19294 {
19295 let old = part_guard.take();
19306 let (co, cm) = old
19307 .as_ref()
19308 .map(|pp| (pp.0.len(), pp.1.len()))
19309 .unwrap_or((0, 0));
19310 if let Some(old) = old {
19311 self.fa_part_retired.lock().unwrap().push(old);
19312 }
19313 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19314 eprintln!(
19315 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19316 co, o_len, cm, ml_len
19317 );
19318 }
19319 *part_guard = Some((
19320 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19321 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19322 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19323 ));
19324 }
19325 let pg = part_guard.as_mut().unwrap();
19326 self.gpu
19327 .stream()
19328 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19329 self.gpu
19330 .stream()
19331 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19332 self.gpu
19333 .stream()
19334 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19335 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19336 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19342 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
19343 std::env::var("MEMRA_FA_SMEM_TKV")
19344 .ok()
19345 .and_then(|v| v.parse().ok())
19346 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
19347 });
19348 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19354 let wg = Self::wkv_on();
19359 let sp2 =
19362 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
19363 if sp2 {
19364 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
19365 if Self::pdl_on() && Self::pdl_wb_on() {
19366 use cudarc::driver::{DevicePtr, DevicePtrMut};
19368 let s = &self.gpu.stream();
19369 let (pq, _b0) = q.device_ptr(s);
19370 let (pk, _b1) = k.device_ptr(s);
19371 let (pv, _b2) = v.device_ptr(s);
19372 let (po, _b3) = part_o.device_ptr_mut(s);
19373 let (pm, _b4) = part_m.device_ptr_mut(s);
19374 let (pl, _b5) = part_l.device_ptr_mut(s);
19375 let (pb, _b6) = base_dev.device_ptr(s);
19376 let mut ps = [
19377 &pq as *const _ as *mut std::ffi::c_void,
19378 &pk as *const _ as *mut _,
19379 &pv as *const _ as *mut _,
19380 &po as *const _ as *mut _,
19381 &pm as *const _ as *mut _,
19382 &pl as *const _ as *mut _,
19383 &hd as *const _ as *mut _,
19384 &nh as *const _ as *mut _,
19385 &nhkv as *const _ as *mut _,
19386 &pb as *const _ as *mut _,
19387 &base_plus as *const _ as *mut _,
19388 &scale as *const _ as *mut _,
19389 &nspm as *const _ as *mut _,
19390 &spk as *const _ as *mut _,
19391 &ktb as *const _ as *mut _,
19392 &vtb as *const _ as *mut _,
19393 &wini as *const _ as *mut _,
19394 ];
19395 unsafe {
19396 self.launch_pdl_flash(
19397 wg,
19398 "fa_decode_vec_q_rows_v4_w_sp",
19399 (n_head_kv as u32, n_splits_max as u32, t as u32),
19400 (32, gqa + 1, 1),
19401 sh,
19402 &mut ps,
19403 )?;
19404 }
19405 } else {
19406 let f = if wg {
19407 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
19408 } else {
19409 self.func("fa_decode_vec_q_rows_v4_w_sp")
19410 };
19411 f.set_attribute(
19412 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19413 sh as i32,
19414 )?;
19415 let cfg = LaunchConfig {
19416 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19417 block_dim: (32, gqa + 1, 1),
19418 shared_mem_bytes: sh,
19419 };
19420 let __s_b = self.gpu.stream();
19421 let mut b = __s_b.launch_builder(&f);
19422 b.arg(q)
19423 .arg(k)
19424 .arg(v)
19425 .arg(&mut *part_o)
19426 .arg(&mut *part_m)
19427 .arg(&mut *part_l)
19428 .arg(&hd)
19429 .arg(&nh)
19430 .arg(&nhkv)
19431 .arg(base_dev)
19432 .arg(&base_plus)
19433 .arg(&scale)
19434 .arg(&nspm)
19435 .arg(&spk)
19436 .arg(&ktb)
19437 .arg(&vtb)
19438 .arg(&wini);
19439 unsafe {
19440 b.launch(cfg)?;
19441 }
19442 }
19443 } else {
19444 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
19445 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
19447 use cudarc::driver::{DevicePtr, DevicePtrMut};
19448 let s = &self.gpu.stream();
19449 let (pq, _b0) = q.device_ptr(s);
19450 let (pk, _b1) = k.device_ptr(s);
19451 let (pv, _b2) = v.device_ptr(s);
19452 let (po, _b3) = part_o.device_ptr_mut(s);
19453 let (pm, _b4) = part_m.device_ptr_mut(s);
19454 let (pl, _b5) = part_l.device_ptr_mut(s);
19455 let (pb, _b6) = base_dev.device_ptr(s);
19456 let mut ps = [
19457 &pq as *const _ as *mut std::ffi::c_void,
19458 &pk as *const _ as *mut _,
19459 &pv as *const _ as *mut _,
19460 &po as *const _ as *mut _,
19461 &pm as *const _ as *mut _,
19462 &pl as *const _ as *mut _,
19463 &hd as *const _ as *mut _,
19464 &nh as *const _ as *mut _,
19465 &nhkv as *const _ as *mut _,
19466 &pb as *const _ as *mut _,
19467 &base_plus as *const _ as *mut _,
19468 &scale as *const _ as *mut _,
19469 &nspm as *const _ as *mut _,
19470 &spk as *const _ as *mut _,
19471 &ktb as *const _ as *mut _,
19472 &vtb as *const _ as *mut _,
19473 &wini as *const _ as *mut _,
19474 ];
19475 unsafe {
19476 self.launch_pdl_flash(
19477 wg,
19478 "fa_decode_vec_q_rows_v4_w",
19479 (n_head_kv as u32, n_splits_max as u32, t as u32),
19480 (32, gqa, 1),
19481 sh,
19482 &mut ps,
19483 )?;
19484 }
19485 } else {
19486 let pick = |name: &str| {
19487 if wg {
19488 self.func_g(name)
19489 } else {
19490 self.func(name)
19491 }
19492 };
19493 let (f, sh) = if fa_v4_at(window) {
19494 let f = pick("fa_decode_vec_q_rows_v4_w");
19495 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
19496 } else if smem_tkv > 0 && window >= smem_tkv {
19497 (
19500 pick("fa_decode_vec_q_rows_smem_w"),
19501 (2 * 32 * head_dim * 2) as u32,
19502 )
19503 } else {
19504 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
19505 };
19506 f.set_attribute(
19507 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19508 sh as i32,
19509 )?;
19510 let cfg = LaunchConfig {
19511 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19512 block_dim: (32, gqa, 1),
19513 shared_mem_bytes: sh,
19514 };
19515 let __s_b = self.gpu.stream();
19516 let mut b = __s_b.launch_builder(&f);
19517 b.arg(q)
19518 .arg(k)
19519 .arg(v)
19520 .arg(&mut *part_o)
19521 .arg(&mut *part_m)
19522 .arg(&mut *part_l)
19523 .arg(&hd)
19524 .arg(&nh)
19525 .arg(&nhkv)
19526 .arg(base_dev)
19527 .arg(&base_plus)
19528 .arg(&scale)
19529 .arg(&nspm)
19530 .arg(&spk)
19531 .arg(&ktb)
19532 .arg(&vtb)
19533 .arg(&wini);
19534 unsafe {
19535 b.launch(cfg)?;
19536 }
19537 }
19538 }
19539 let cfg2 = LaunchConfig {
19540 grid_dim: (n_head as u32, t as u32, 1),
19541 block_dim: (head_dim as u32, 1, 1),
19542 shared_mem_bytes: 0,
19543 };
19544 if let Some((oq, od)) = q8_out {
19545 if Self::pdl_on() && Self::pdl_wb_on() {
19548 use cudarc::driver::{DevicePtr, DevicePtrMut};
19550 let s = &self.gpu.stream();
19551 let (po, _g0) = part_o.device_ptr(s);
19552 let (pm, _g1) = part_m.device_ptr(s);
19553 let (pl, _g2) = part_l.device_ptr(s);
19554 let (pq, _g3) = oq.device_ptr_mut(s);
19555 let (pd, _g4) = od.device_ptr_mut(s);
19556 let mut ps = [
19557 &po as *const _ as *mut std::ffi::c_void,
19558 &pm as *const _ as *mut _,
19559 &pl as *const _ as *mut _,
19560 &pq as *const _ as *mut _,
19561 &pd as *const _ as *mut _,
19562 &hd as *const _ as *mut _,
19563 &nh as *const _ as *mut _,
19564 &nspm as *const _ as *mut _,
19565 &spk as *const _ as *mut _,
19566 &wini as *const _ as *mut _,
19567 ];
19568 unsafe {
19569 self.launch_pdl_flash(
19570 wg,
19571 "fa_decode_combine_rows_w_q8_1",
19572 cfg2.grid_dim,
19573 cfg2.block_dim,
19574 0,
19575 &mut ps,
19576 )?;
19577 }
19578 return Ok(());
19579 }
19580 let fc = if wg {
19581 self.func_g("fa_decode_combine_rows_w_q8_1")
19582 } else {
19583 self.func("fa_decode_combine_rows_w_q8_1")
19584 };
19585 let __s_b2 = self.gpu.stream();
19586 let mut b2 = __s_b2.launch_builder(&fc);
19587 b2.arg(&*part_o)
19588 .arg(&*part_m)
19589 .arg(&*part_l)
19590 .arg(oq)
19591 .arg(od)
19592 .arg(&hd)
19593 .arg(&nh)
19594 .arg(&nspm)
19595 .arg(&spk)
19596 .arg(&wini);
19597 unsafe {
19598 b2.launch(cfg2)?;
19599 }
19600 return Ok(());
19601 }
19602 let fc = if wg {
19603 self.func_g("fa_decode_combine_rows_w")
19604 } else {
19605 self.func("fa_decode_combine_rows_w")
19606 };
19607 let __s_b2 = self.gpu.stream();
19608 let mut b2 = __s_b2.launch_builder(&fc);
19609 b2.arg(&*part_o)
19610 .arg(&*part_m)
19611 .arg(&*part_l)
19612 .arg(o)
19613 .arg(&hd)
19614 .arg(&nh)
19615 .arg(&nspm)
19616 .arg(&spk)
19617 .arg(&wini);
19618 unsafe {
19619 b2.launch(cfg2)?;
19620 }
19621 Ok(())
19622 }
19623
19624 #[allow(clippy::too_many_arguments)]
19630 pub fn fa_decode_rows_dc(
19631 &self,
19632 q: &CudaSlice<f32>,
19633 k: &cudarc::driver::CudaView<u8>,
19634 v: &cudarc::driver::CudaView<u8>,
19635 o: &mut CudaSlice<f32>,
19636 head_dim: usize,
19637 n_head: usize,
19638 n_head_kv: usize,
19639 base_dev: &CudaSlice<i32>,
19640 t_kv_upper: usize,
19641 t: usize,
19642 scale: f32,
19643 k_tok_bytes: usize,
19644 v_tok_bytes: usize,
19645 base_plus: i32,
19646 g: bool,
19647 ) -> Result<(), Box<dyn std::error::Error>> {
19648 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
19649 assert!(
19650 v4 || fa_v3_active(head_dim),
19651 "stream fa rows requires the v3 or v4 lane"
19652 );
19653 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
19654 if v4 {
19655 let sp = fa_split_keys(t_kv_upper, n_head_kv);
19656 let n_splits_max = (t_kv_upper + sp - 1) / sp;
19657 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19658 let (nspm, spk) = (n_splits_max as i32, sp as i32);
19659 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19660 let gqa = (n_head / n_head_kv).max(1) as u32;
19661 let o_len = t * n_head * n_splits_max * head_dim;
19662 let ml_len = t * n_head * n_splits_max;
19663 let mut part_guard = self.fa_part_pool.lock().unwrap();
19664 if part_guard
19665 .as_ref()
19666 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19667 .unwrap_or(true)
19668 {
19669 let old = part_guard.take();
19680 let (co, cm) = old
19681 .as_ref()
19682 .map(|pp| (pp.0.len(), pp.1.len()))
19683 .unwrap_or((0, 0));
19684 if let Some(old) = old {
19685 self.fa_part_retired.lock().unwrap().push(old);
19686 }
19687 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19688 eprintln!(
19689 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19690 co, o_len, cm, ml_len
19691 );
19692 }
19693 *part_guard = Some((
19694 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19695 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19696 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19697 ));
19698 }
19699 let pg = part_guard.as_mut().unwrap();
19700 self.gpu
19701 .stream()
19702 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19703 self.gpu
19704 .stream()
19705 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19706 self.gpu
19707 .stream()
19708 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19709 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19710 let f = if g {
19711 self.func_g("fa_decode_vec_q_rows_v4_dc")
19712 } else {
19713 self.func("fa_decode_vec_q_rows_v4_dc")
19714 };
19715 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19716 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19717 f.set_attribute(
19718 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19719 sh as i32,
19720 )?;
19721 let cfg = LaunchConfig {
19722 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19723 block_dim: (32, gqa, 1),
19724 shared_mem_bytes: sh,
19725 };
19726 let __s_b = self.gpu.stream();
19727 let mut b = __s_b.launch_builder(&f);
19728 b.arg(q)
19729 .arg(k)
19730 .arg(v)
19731 .arg(&mut *part_o)
19732 .arg(&mut *part_m)
19733 .arg(&mut *part_l)
19734 .arg(&hd)
19735 .arg(&nh)
19736 .arg(&nhkv)
19737 .arg(base_dev)
19738 .arg(&base_plus)
19739 .arg(&scale)
19740 .arg(&nspm)
19741 .arg(&spk)
19742 .arg(&ktb)
19743 .arg(&vtb);
19744 unsafe {
19745 b.launch(cfg)?;
19746 }
19747 let fc = self.func("fa_decode_combine_rows_dc");
19748 let cfg2 = LaunchConfig {
19749 grid_dim: (n_head as u32, t as u32, 1),
19750 block_dim: (head_dim as u32, 1, 1),
19751 shared_mem_bytes: 0,
19752 };
19753 let __s_b2 = self.gpu.stream();
19754 let mut b2 = __s_b2.launch_builder(&fc);
19755 b2.arg(&*part_o)
19756 .arg(&*part_m)
19757 .arg(&*part_l)
19758 .arg(o)
19759 .arg(&hd)
19760 .arg(&nh)
19761 .arg(base_dev)
19762 .arg(&base_plus)
19763 .arg(&nspm)
19764 .arg(&spk);
19765 unsafe {
19766 b2.launch(cfg2)?;
19767 }
19768 return Ok(());
19769 }
19770 let sp = fa_split_keys(t_kv_upper, n_head_kv);
19771 let n_splits_max = (t_kv_upper + sp - 1) / sp;
19772 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19773 let (nspm, spk) = (n_splits_max as i32, sp as i32);
19774 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19775 let gqa = (n_head / n_head_kv).max(1) as u32;
19776 let o_len = t * n_head * n_splits_max * head_dim;
19777 let ml_len = t * n_head * n_splits_max;
19778 let mut part_guard = self.fa_part_pool.lock().unwrap();
19779 if part_guard
19780 .as_ref()
19781 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19782 .unwrap_or(true)
19783 {
19784 let old = part_guard.take();
19795 let (co, cm) = old
19796 .as_ref()
19797 .map(|pp| (pp.0.len(), pp.1.len()))
19798 .unwrap_or((0, 0));
19799 if let Some(old) = old {
19800 self.fa_part_retired.lock().unwrap().push(old);
19801 }
19802 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19803 eprintln!(
19804 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19805 co, o_len, cm, ml_len
19806 );
19807 }
19808 *part_guard = Some((
19809 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19810 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19811 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19812 ));
19813 }
19814 let pg = part_guard.as_mut().unwrap();
19815 self.gpu
19816 .stream()
19817 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19818 self.gpu
19819 .stream()
19820 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19821 self.gpu
19822 .stream()
19823 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19824 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19825 let f = self.func("fa_decode_vec_q_rows_v3_dc");
19826 let sh = (32 * head_dim * 2) as u32;
19827 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19828 f.set_attribute(
19829 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19830 sh as i32,
19831 )?;
19832 let cfg = LaunchConfig {
19833 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19834 block_dim: (32, gqa, 1),
19835 shared_mem_bytes: sh,
19836 };
19837 let __s_b = self.gpu.stream();
19838 let mut b = __s_b.launch_builder(&f);
19839 b.arg(q)
19840 .arg(k)
19841 .arg(v)
19842 .arg(&mut *part_o)
19843 .arg(&mut *part_m)
19844 .arg(&mut *part_l)
19845 .arg(&hd)
19846 .arg(&nh)
19847 .arg(&nhkv)
19848 .arg(base_dev)
19849 .arg(&scale)
19850 .arg(&nspm)
19851 .arg(&spk)
19852 .arg(&ktb)
19853 .arg(&vtb);
19854 unsafe {
19855 b.launch(cfg)?;
19856 }
19857 let fc = self.func("fa_decode_combine_rows_dc");
19858 let cfg2 = LaunchConfig {
19859 grid_dim: (n_head as u32, t as u32, 1),
19860 block_dim: (head_dim as u32, 1, 1),
19861 shared_mem_bytes: 0,
19862 };
19863 let plus0 = 0i32;
19864 let __s_b2 = self.gpu.stream();
19865 let mut b2 = __s_b2.launch_builder(&fc);
19866 b2.arg(&*part_o)
19867 .arg(&*part_m)
19868 .arg(&*part_l)
19869 .arg(o)
19870 .arg(&hd)
19871 .arg(&nh)
19872 .arg(base_dev)
19873 .arg(&plus0)
19874 .arg(&nspm)
19875 .arg(&spk);
19876 unsafe {
19877 b2.launch(cfg2)?;
19878 }
19879 Ok(())
19880 }
19881
19882 pub fn fa_decode_dc(
19893 &self,
19894 q: &CudaSlice<f32>,
19895 k: &cudarc::driver::CudaView<u8>,
19896 v: &cudarc::driver::CudaView<u8>,
19897 o: &mut CudaSlice<f32>,
19898 head_dim: usize,
19899 n_head: usize,
19900 n_head_kv: usize,
19901 t_kv_dev: &CudaSlice<i32>,
19902 bucket_max: usize,
19903 scale: f32,
19904 k_tok_bytes: usize,
19905 v_tok_bytes: usize,
19906 g: bool,
19907 ) -> Result<(), Box<dyn std::error::Error>> {
19908 self.fa_decode_dc_q8(
19909 q,
19910 k,
19911 v,
19912 o,
19913 head_dim,
19914 n_head,
19915 n_head_kv,
19916 t_kv_dev,
19917 bucket_max,
19918 scale,
19919 k_tok_bytes,
19920 v_tok_bytes,
19921 g,
19922 None,
19923 )
19924 }
19925
19926 #[allow(clippy::too_many_arguments)]
19929 pub fn fa_decode_dc_q8(
19930 &self,
19931 q: &CudaSlice<f32>,
19932 k: &cudarc::driver::CudaView<u8>,
19933 v: &cudarc::driver::CudaView<u8>,
19934 o: &mut CudaSlice<f32>,
19935 head_dim: usize,
19936 n_head: usize,
19937 n_head_kv: usize,
19938 t_kv_dev: &CudaSlice<i32>,
19939 bucket_max: usize,
19940 scale: f32,
19941 k_tok_bytes: usize,
19942 v_tok_bytes: usize,
19943 g: bool,
19944 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19945 ) -> Result<(), Box<dyn std::error::Error>> {
19946 let mut fa_vec =
19954 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19955 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19956 fa_vec = false;
19957 } let sp = fa_split_keys(bucket_max, n_head_kv);
19959 let n_splits = if fa_vec {
19960 ((bucket_max + sp - 1) / sp).max(1)
19961 } else {
19962 ((bucket_max + 255) / 256).max(1)
19963 };
19964 let o_len = n_head * n_splits * head_dim;
19965 let ml_len = n_head * n_splits;
19966 let mut part_guard = self.fa_part_pool.lock().unwrap();
19967 if part_guard
19968 .as_ref()
19969 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19970 .unwrap_or(true)
19971 {
19972 let old = part_guard.take();
19983 let (co, cm) = old
19984 .as_ref()
19985 .map(|pp| (pp.0.len(), pp.1.len()))
19986 .unwrap_or((0, 0));
19987 if let Some(old) = old {
19988 self.fa_part_retired.lock().unwrap().push(old);
19989 }
19990 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19991 eprintln!(
19992 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19993 co, o_len, cm, ml_len
19994 );
19995 }
19996 *part_guard = Some((
19997 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19998 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19999 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20000 ));
20001 }
20002 let pg = part_guard.as_mut().unwrap();
20003 self.gpu
20004 .stream()
20005 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20006 self.gpu
20007 .stream()
20008 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20009 self.gpu
20010 .stream()
20011 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20012 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20013 let (hd, nh, nhkv, nsp) = (
20014 head_dim as i32,
20015 n_head as i32,
20016 n_head_kv as i32,
20017 n_splits as i32,
20018 );
20019 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20020 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
20021 let deep = fa_vec
20024 && head_dim == 256
20025 && fa_v4_at(bucket_max)
20026 && !g
20027 && fa_deep_at(bucket_max)
20028 && !matches!(fa_v4_mode(), "noB3" | "stage");
20029 let (f, cfg) = if fa_vec
20030 && head_dim == 512
20031 && bucket_max >= {
20032 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20033 *FA512_MIN_DC.get_or_init(|| {
20034 std::env::var("MEMRA_FA512_MIN")
20035 .ok()
20036 .and_then(|v| v.parse().ok())
20037 .unwrap_or(512)
20038 })
20039 } {
20040 let gqa = (n_head / n_head_kv).max(1) as u32;
20042 (
20043 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
20044 LaunchConfig {
20045 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20046 block_dim: (32, gqa, 1),
20047 shared_mem_bytes: 0,
20048 },
20049 )
20050 } else if fa_vec && head_dim == 512 {
20051 let q_view = q.as_view();
20054 let mut o_view = o.as_view_mut();
20055 return self.fa_decode_scalar_unified(
20056 &q_view,
20057 k,
20058 v,
20059 &mut o_view,
20060 head_dim,
20061 n_head,
20062 n_head_kv,
20063 0,
20064 Some(t_kv_dev),
20065 scale,
20066 n_splits,
20067 sp,
20068 k_tok_bytes,
20069 v_tok_bytes,
20070 g,
20071 &mut *part_o,
20072 &mut *part_m,
20073 &mut *part_l,
20074 q8_out,
20075 );
20076 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
20077 let gqa = (n_head / n_head_kv).max(1) as u32;
20080 let fv = if g {
20081 self.func_g("fa_decode_vec_q_v4_dc")
20082 } else if deep {
20083 self.func("fa_decode_vec_q_v4_deep_dc")
20084 } else {
20085 self.func("fa_decode_vec_q_v4_dc")
20086 };
20087 let shmem =
20088 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
20089 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20090 fv.set_attribute(
20091 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20092 shmem as i32,
20093 )?;
20094 (
20095 fv,
20096 LaunchConfig {
20097 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20098 block_dim: (32, gqa, 1),
20099 shared_mem_bytes: shmem,
20100 },
20101 )
20102 } else if fa_vec && fa_v3_active(head_dim) {
20103 let gqa = (n_head / n_head_kv).max(1) as u32;
20106 let fv = if g {
20107 self.func_g("fa_decode_vec_q_v3_dc")
20108 } else {
20109 self.func("fa_decode_vec_q_v3_dc")
20110 };
20111 let shmem = (32 * head_dim * 2) as u32; (
20113 fv,
20114 LaunchConfig {
20115 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20116 block_dim: (32, gqa, 1),
20117 shared_mem_bytes: shmem,
20118 },
20119 )
20120 } else if fa_vec && fa_v2_on() {
20121 let gqa = (n_head / n_head_kv).max(1) as u32;
20125 let fv = if g {
20126 self.func_g("fa_decode_vec_q_v2_dc")
20127 } else {
20128 self.func("fa_decode_vec_q_v2_dc")
20129 };
20130 let shmem = (2 * 32 * head_dim * 2) as u32; (
20132 fv,
20133 LaunchConfig {
20134 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20135 block_dim: (32, gqa, 1),
20136 shared_mem_bytes: shmem,
20137 },
20138 )
20139 } else if fa_vec {
20140 let gqa = (n_head / n_head_kv).max(1) as u32;
20141 let fv = if g {
20143 self.func_g("fa_decode_vec_q_dc")
20144 } else {
20145 self.func("fa_decode_vec_q_dc")
20146 };
20147 (
20148 fv,
20149 LaunchConfig {
20150 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20151 block_dim: (32, gqa, 1),
20152 shared_mem_bytes: 0,
20153 },
20154 )
20155 } else {
20156 let q_view = q.as_view();
20157 let mut o_view = o.as_view_mut();
20158 return self.fa_decode_scalar_unified(
20159 &q_view,
20160 k,
20161 v,
20162 &mut o_view,
20163 head_dim,
20164 n_head,
20165 n_head_kv,
20166 0,
20167 Some(t_kv_dev),
20168 scale,
20169 n_splits,
20170 if fa_vec { sp } else { 256 },
20171 k_tok_bytes,
20172 v_tok_bytes,
20173 g,
20174 &mut *part_o,
20175 &mut *part_m,
20176 &mut *part_l,
20177 q8_out,
20178 );
20179 };
20180 let ski = sp as i32; let __s_b = self.gpu.stream();
20182 let mut b = __s_b.launch_builder(&f);
20183 b.arg(q)
20184 .arg(k)
20185 .arg(v)
20186 .arg(&mut *part_o)
20187 .arg(&mut *part_m)
20188 .arg(&mut *part_l)
20189 .arg(&hd)
20190 .arg(&nh)
20191 .arg(&nhkv)
20192 .arg(t_kv_dev)
20193 .arg(&scale)
20194 .arg(&nsp)
20195 .arg(&ski)
20196 .arg(&ktb)
20197 .arg(&vtb);
20198 unsafe {
20199 b.launch(cfg)?;
20200 }
20201 let cfg2 = LaunchConfig {
20202 grid_dim: (n_head as u32, 1, 1),
20203 block_dim: (head_dim as u32, 1, 1),
20204 shared_mem_bytes: 0,
20205 };
20206 if let Some((oq, od)) = q8_out {
20207 let fc = if g {
20208 self.func_g("fa_decode_combine_q8_1")
20209 } else {
20210 self.fa_func("fa_decode_combine_q8_1", head_dim)
20211 };
20212 let __s_b2 = self.gpu.stream();
20213 let mut b2 = __s_b2.launch_builder(&fc);
20214 b2.arg(&*part_o)
20215 .arg(&*part_m)
20216 .arg(&*part_l)
20217 .arg(oq)
20218 .arg(od)
20219 .arg(&hd)
20220 .arg(&nh)
20221 .arg(&nsp);
20222 unsafe {
20223 b2.launch(cfg2)?;
20224 }
20225 return Ok(());
20226 }
20227 let fc = if g {
20228 self.func_g("fa_decode_combine_f32")
20229 } else {
20230 self.fa_func("fa_decode_combine_f32", head_dim)
20231 };
20232 let __s_b2 = self.gpu.stream();
20233 let mut b2 = __s_b2.launch_builder(&fc);
20234 b2.arg(&*part_o)
20235 .arg(&*part_m)
20236 .arg(&*part_l)
20237 .arg(o)
20238 .arg(&hd)
20239 .arg(&nh)
20240 .arg(&nsp);
20241 unsafe {
20242 b2.launch(cfg2)?;
20243 }
20244 Ok(())
20245 }
20246
20247 pub fn fa_geom_eager(
20253 &self,
20254 t_kv: usize,
20255 head_dim: usize,
20256 n_head_kv: usize,
20257 g: bool,
20258 ) -> (bool, usize) {
20259 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20263 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
20269 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
20270 if g && head_dim == 256 && !fa_v4_at(t_kv) {
20276 fa_vec = false;
20277 }
20278 let sp = fa_split_keys(t_kv, n_head_kv);
20279 let n_splits = if fa_vec {
20280 ((t_kv + sp - 1) / sp).max(1)
20281 } else {
20282 ((t_kv + 255) / 256).max(1)
20283 };
20284 (fa_vec, n_splits)
20285 }
20286
20287 pub fn fa_bucket_key(
20293 &self,
20294 t_kv: usize,
20295 head_dim: usize,
20296 n_head_kv: usize,
20297 g: bool,
20298 ) -> (bool, usize) {
20299 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
20300 }
20301
20302 pub fn capture_graph_retained<F>(
20314 &self,
20315 step: F,
20316 ) -> Result<
20317 (
20318 cudarc::driver::CudaGraph,
20319 Vec<Box<dyn std::any::Any + Send>>,
20320 ),
20321 Box<dyn std::error::Error>,
20322 >
20323 where
20324 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20325 {
20326 use cudarc::driver::sys::CUgraphInstantiate_flags;
20327 self.capture_graph_retained_flags(
20328 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
20329 step,
20330 )
20331 }
20332
20333 pub fn capture_graph_retained_flags<F>(
20338 &self,
20339 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
20340 mut step: F,
20341 ) -> Result<
20342 (
20343 cudarc::driver::CudaGraph,
20344 Vec<Box<dyn std::any::Any + Send>>,
20345 ),
20346 Box<dyn std::error::Error>,
20347 >
20348 where
20349 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20350 {
20351 use cudarc::driver::sys::CUstreamCaptureMode;
20352 self.capture_keep.lock().unwrap().clear();
20360 let was_tracking = self.gpu.ctx.is_event_tracking();
20361 if was_tracking {
20362 unsafe {
20363 self.gpu.ctx.disable_event_tracking();
20364 }
20365 }
20366 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20367 self.capture_keep_on
20368 .store(true, std::sync::atomic::Ordering::Relaxed);
20369 let w = (|| {
20370 step(self)?;
20371 step(self)
20372 })();
20373 self.capture_keep_on
20374 .store(false, std::sync::atomic::Ordering::Relaxed);
20375 w?;
20376 self.gpu.stream().synchronize()?;
20377 self.gpu
20378 .stream()
20379 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20380 let r = step(self);
20381 let g = self.gpu.stream().end_capture(flags);
20382 r?;
20383 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20384 graph.upload()?;
20385 Ok(graph)
20386 };
20387 let result = run();
20388 self.capture_keep_on
20389 .store(false, std::sync::atomic::Ordering::Relaxed);
20390 if was_tracking {
20391 unsafe {
20392 self.gpu.ctx.enable_event_tracking();
20393 }
20394 }
20395 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
20396 Ok((result?, keeper))
20397 }
20398
20399 pub fn capture_graph<F>(
20400 &self,
20401 mut step: F,
20402 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
20403 where
20404 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20405 {
20406 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
20407 let was_tracking = self.gpu.ctx.is_event_tracking();
20415 if was_tracking {
20416 unsafe {
20417 self.gpu.ctx.disable_event_tracking();
20418 }
20419 }
20420 let iflag = {
20427 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
20428 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
20429 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
20432 Ok("priority") => {
20433 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
20434 }
20435 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
20436 })
20437 };
20438 let ct = {
20445 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20446 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
20447 };
20448 let warmups = {
20471 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20472 *W.get_or_init(|| {
20473 std::env::var("MEMRA_GRAPH_WARMUPS")
20474 .ok()
20475 .and_then(|v| v.parse().ok())
20476 .filter(|n| *n >= 1)
20477 .unwrap_or(1)
20478 })
20479 };
20480 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20481 let t_w = std::time::Instant::now();
20482 for _ in 0..warmups {
20484 step(self)?;
20485 }
20486 self.gpu.stream().synchronize()?;
20487 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
20488 let t_c = std::time::Instant::now();
20490 self.gpu
20491 .stream()
20492 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20493 let r = step(self);
20496 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
20497 let t_i = std::time::Instant::now();
20498 let g = self.gpu.stream().end_capture(iflag);
20499 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
20500 r?;
20501 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20502 let t_u = std::time::Instant::now();
20503 graph.upload()?;
20504 if ct {
20505 println!(
20506 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
20507 instantiate {ms_inst:.2} ms upload {:.2} ms",
20508 t_u.elapsed().as_secs_f64() * 1e3
20509 );
20510 }
20511 Ok(graph)
20512 };
20513 let result = run();
20514 if was_tracking {
20515 unsafe {
20516 self.gpu.ctx.enable_event_tracking();
20517 }
20518 }
20519 result
20520 }
20521
20522 pub fn gdn_scan_s128_view(
20524 &self,
20525 q: &CudaSlice<f32>,
20526 k: &CudaSlice<f32>,
20527 v: &CudaSlice<f32>,
20528 g: &CudaSlice<f32>,
20529 beta: &CudaSlice<f32>,
20530 state_in: &cudarc::driver::CudaView<f32>,
20531 state_out: &mut cudarc::driver::CudaViewMut<f32>,
20532 o: &mut CudaSlice<f32>,
20533 n_head: usize,
20534 t: usize,
20535 scale: f32,
20536 ) -> Result<(), Box<dyn std::error::Error>> {
20537 let f = self.func("gdn_scan_s128");
20538 const S_V: u32 = 128;
20539 const WARP: u32 = 32;
20540 const COLS: u32 = 4;
20541 let cfg = LaunchConfig {
20542 grid_dim: (n_head as u32, 1, S_V / COLS),
20543 block_dim: (WARP, COLS, 1),
20544 shared_mem_bytes: 0,
20545 };
20546 let (h, ti) = (n_head as i32, t as i32);
20547 let __s_b = self.gpu.stream();
20548 let mut b = __s_b.launch_builder(&f);
20549 b.arg(q)
20550 .arg(k)
20551 .arg(v)
20552 .arg(g)
20553 .arg(beta)
20554 .arg(state_in)
20555 .arg(state_out)
20556 .arg(o)
20557 .arg(&h)
20558 .arg(&ti)
20559 .arg(&scale);
20560 unsafe {
20561 b.launch(cfg)?;
20562 }
20563 Ok(())
20564 }
20565
20566 pub fn ssm_conv1d_view(
20568 &self,
20569 x: &cudarc::driver::CudaView<f32>,
20570 w: &CudaSlice<f32>,
20571 y: &mut CudaSlice<f32>,
20572 conv_dim: usize,
20573 t: usize,
20574 d_conv: usize,
20575 silu: bool,
20576 ) -> Result<(), Box<dyn std::error::Error>> {
20577 let f = self.func("ssm_conv1d_silu_f32");
20578 let cfg = LaunchConfig {
20580 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20581 block_dim: (256, 1, 1),
20582 shared_mem_bytes: 0,
20583 };
20584 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20585 let __s_b = self.gpu.stream();
20586 let mut b = __s_b.launch_builder(&f);
20587 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20588 unsafe {
20589 b.launch(cfg)?;
20590 }
20591 Ok(())
20592 }
20593
20594 pub fn ssm_conv1d_tm(
20601 &self,
20602 qkv_tm: &CudaSlice<f32>,
20603 w: &CudaSlice<f32>,
20604 y: &mut CudaSlice<f32>,
20605 conv_dim: usize,
20606 t: usize,
20607 d_conv: usize,
20608 ) -> Result<(), Box<dyn std::error::Error>> {
20609 let f = self.func("ssm_conv1d_tm_f32");
20610 let cfg = LaunchConfig {
20611 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20612 block_dim: (256, 1, 1),
20613 shared_mem_bytes: 0,
20614 };
20615 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20616 let __s_b = self.gpu.stream();
20617 let mut b = __s_b.launch_builder(&f);
20618 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
20619 unsafe {
20620 b.launch(cfg)?;
20621 }
20622 Ok(())
20623 }
20624
20625 pub fn ssm_conv1d_tm_state(
20633 &self,
20634 qkv_tm: &CudaSlice<f32>,
20635 conv_state: &mut CudaSlice<f32>,
20636 w: &CudaSlice<f32>,
20637 y: &mut CudaSlice<f32>,
20638 conv_dim: usize,
20639 t: usize,
20640 d_conv: usize,
20641 ) -> Result<(), Box<dyn std::error::Error>> {
20642 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
20643 }
20644
20645 #[allow(clippy::too_many_arguments)]
20648 pub fn ssm_conv1d_tm_state_pad(
20649 &self,
20650 qkv_tm: &CudaSlice<f32>,
20651 conv_state: &mut CudaSlice<f32>,
20652 w: &CudaSlice<f32>,
20653 y: &mut CudaSlice<f32>,
20654 conv_dim: usize,
20655 t: usize,
20656 d_conv: usize,
20657 pad_len: Option<&CudaSlice<i32>>,
20658 ) -> Result<(), Box<dyn std::error::Error>> {
20659 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20660 let ring_old = if t < d_conv - 1 {
20664 Some(self.clone_dtod(conv_state)?)
20665 } else {
20666 None
20667 };
20668 {
20669 let f = self.func("ssm_conv1d_tm_state_f32");
20670 let cfg = LaunchConfig {
20671 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20672 block_dim: (256, 1, 1),
20673 shared_mem_bytes: 0,
20674 };
20675 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20676 let __s_b = self.gpu.stream();
20677 let mut b = __s_b.launch_builder(&f);
20678 b.arg(qkv_tm)
20679 .arg(&*conv_state)
20680 .arg(w)
20681 .arg(y)
20682 .arg(&cd)
20683 .arg(&ti)
20684 .arg(&dc);
20685 unsafe {
20686 b.launch(cfg)?;
20687 }
20688 }
20689 match (ring_old, pad_len) {
20690 (None, Some(len_d)) => {
20691 let f = self.func("ssm_conv_ring_update_dev_f32");
20692 let n = conv_dim * (d_conv - 1);
20693 let cfg = LaunchConfig::for_num_elems(n as u32);
20694 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20695 let __s_b = self.gpu.stream();
20696 let mut b = __s_b.launch_builder(&f);
20697 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20698 unsafe {
20699 b.launch(cfg)?;
20700 }
20701 }
20702 (None, None) => {
20703 let f = self.func("ssm_conv_ring_update_f32");
20704 let n = conv_dim * (d_conv - 1);
20705 let cfg = LaunchConfig::for_num_elems(n as u32);
20706 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20707 let __s_b = self.gpu.stream();
20708 let mut b = __s_b.launch_builder(&f);
20709 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20710 unsafe {
20711 b.launch(cfg)?;
20712 }
20713 }
20714 (Some(old), _) => {
20715 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
20716 }
20717 }
20718 Ok(())
20719 }
20720
20721 pub fn ssm_conv1d_tm_state_pad_v(
20723 &self,
20724 qkv_tm: &cudarc::driver::CudaView<f32>,
20725 conv_state: &mut CudaSlice<f32>,
20726 w: &CudaSlice<f32>,
20727 y: &mut CudaSlice<f32>,
20728 conv_dim: usize,
20729 t: usize,
20730 d_conv: usize,
20731 pad_len: Option<&CudaSlice<i32>>,
20732 ) -> Result<(), Box<dyn std::error::Error>> {
20733 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20734 let ring_old = if t < d_conv - 1 {
20738 Some(self.clone_dtod(conv_state)?)
20739 } else {
20740 None
20741 };
20742 {
20743 let f = self.func("ssm_conv1d_tm_state_f32");
20744 let cfg = LaunchConfig {
20745 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20746 block_dim: (256, 1, 1),
20747 shared_mem_bytes: 0,
20748 };
20749 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20750 let __s_b = self.gpu.stream();
20751 let mut b = __s_b.launch_builder(&f);
20752 b.arg(qkv_tm)
20753 .arg(&*conv_state)
20754 .arg(w)
20755 .arg(y)
20756 .arg(&cd)
20757 .arg(&ti)
20758 .arg(&dc);
20759 unsafe {
20760 b.launch(cfg)?;
20761 }
20762 }
20763 match (ring_old, pad_len) {
20764 (None, Some(len_d)) => {
20765 let f = self.func("ssm_conv_ring_update_dev_f32");
20766 let n = conv_dim * (d_conv - 1);
20767 let cfg = LaunchConfig::for_num_elems(n as u32);
20768 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20769 let __s_b = self.gpu.stream();
20770 let mut b = __s_b.launch_builder(&f);
20771 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20772 unsafe {
20773 b.launch(cfg)?;
20774 }
20775 }
20776 (None, None) => {
20777 let f = self.func("ssm_conv_ring_update_f32");
20778 let n = conv_dim * (d_conv - 1);
20779 let cfg = LaunchConfig::for_num_elems(n as u32);
20780 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20781 let __s_b = self.gpu.stream();
20782 let mut b = __s_b.launch_builder(&f);
20783 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20784 unsafe {
20785 b.launch(cfg)?;
20786 }
20787 }
20788 (Some(_), _) => unreachable!(
20789 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
20790 ),
20791 }
20792 Ok(())
20793 }
20794
20795 pub fn ssm_conv_ring_rebuild(
20800 &self,
20801 qkv_tm: &CudaSlice<f32>,
20802 ring_old: &CudaSlice<f32>,
20803 conv_state: &mut CudaSlice<f32>,
20804 conv_dim: usize,
20805 tc: usize,
20806 d_conv: usize,
20807 ) -> Result<(), Box<dyn std::error::Error>> {
20808 let f = self.func("ssm_conv_ring_rebuild_f32");
20809 let n = conv_dim * (d_conv - 1);
20810 let cfg = LaunchConfig::for_num_elems(n as u32);
20811 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
20812 let __s_b = self.gpu.stream();
20813 let mut b = __s_b.launch_builder(&f);
20814 b.arg(qkv_tm)
20815 .arg(ring_old)
20816 .arg(conv_state)
20817 .arg(&cd)
20818 .arg(&ti)
20819 .arg(&dc);
20820 unsafe {
20821 b.launch(cfg)?;
20822 }
20823 Ok(())
20824 }
20825
20826 #[allow(clippy::too_many_arguments)]
20831 pub fn gdn_prep_decode(
20832 &self,
20833 conv_out: &CudaSlice<f32>,
20834 beta_raw: &CudaSlice<f32>,
20835 alpha: &CudaSlice<f32>,
20836 dt_bias: &CudaSlice<f32>,
20837 a: &CudaSlice<f32>,
20838 q_l2: &mut CudaSlice<f32>,
20839 k_l2: &mut CudaSlice<f32>,
20840 v_g: &mut CudaSlice<f32>,
20841 beta: &mut CudaSlice<f32>,
20842 g_log: &mut CudaSlice<f32>,
20843 d_state: usize,
20844 num_v: usize,
20845 num_k: usize,
20846 key_dim: usize,
20847 eps: f32,
20848 ) -> Result<(), Box<dyn std::error::Error>> {
20849 let f = self.func("gdn_prep_decode_f32");
20850 let cfg = LaunchConfig {
20851 grid_dim: (num_v as u32, 1, 1),
20852 block_dim: (32, 4, 1),
20853 shared_mem_bytes: 0,
20854 };
20855 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20856 let __s_b = self.gpu.stream();
20857 let mut b = __s_b.launch_builder(&f);
20858 b.arg(conv_out)
20859 .arg(beta_raw)
20860 .arg(alpha)
20861 .arg(dt_bias)
20862 .arg(a)
20863 .arg(q_l2)
20864 .arg(k_l2)
20865 .arg(v_g)
20866 .arg(beta)
20867 .arg(g_log)
20868 .arg(&ds)
20869 .arg(&nv)
20870 .arg(&nk)
20871 .arg(&kd)
20872 .arg(&eps);
20873 unsafe {
20874 b.launch(cfg)?;
20875 }
20876 Ok(())
20877 }
20878
20879 #[allow(clippy::too_many_arguments)]
20883 pub fn ssm_conv1d_gdn(
20884 &self,
20885 qkv_tm: &CudaSlice<f32>,
20886 w: &CudaSlice<f32>,
20887 q_g: &mut CudaSlice<f32>,
20888 k_g: &mut CudaSlice<f32>,
20889 v_g: &mut CudaSlice<f32>,
20890 conv_dim: usize,
20891 t: usize,
20892 d_conv: usize,
20893 d_state: usize,
20894 num_v: usize,
20895 num_k: usize,
20896 key_dim: usize,
20897 ) -> Result<(), Box<dyn std::error::Error>> {
20898 let f = self.func("ssm_conv1d_gdn_f32");
20899 let cfg = LaunchConfig {
20900 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20901 block_dim: (256, 1, 1),
20902 shared_mem_bytes: 0,
20903 };
20904 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20905 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20906 let __s_b = self.gpu.stream();
20907 let mut b = __s_b.launch_builder(&f);
20908 b.arg(qkv_tm)
20909 .arg(w)
20910 .arg(q_g)
20911 .arg(k_g)
20912 .arg(v_g)
20913 .arg(&cd)
20914 .arg(&ti)
20915 .arg(&dc)
20916 .arg(&ds)
20917 .arg(&nv)
20918 .arg(&nk)
20919 .arg(&kd);
20920 unsafe {
20921 b.launch(cfg)?;
20922 }
20923 Ok(())
20924 }
20925
20926 pub fn ssm_conv1d(
20927 &self,
20928 x: &CudaSlice<f32>,
20929 w: &CudaSlice<f32>,
20930 y: &mut CudaSlice<f32>,
20931 conv_dim: usize,
20932 t: usize,
20933 d_conv: usize,
20934 silu: bool,
20935 ) -> Result<(), Box<dyn std::error::Error>> {
20936 let f = self.func("ssm_conv1d_silu_f32");
20937 let cfg = LaunchConfig {
20938 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20939 block_dim: (256, 1, 1),
20940 shared_mem_bytes: 0,
20941 };
20942 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20943 let __s_b = self.gpu.stream();
20944 let mut b = __s_b.launch_builder(&f);
20945 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20946 unsafe {
20947 b.launch(cfg)?;
20948 }
20949 Ok(())
20950 }
20951
20952 pub fn gdn_scan_s128(
20955 &self,
20956 q: &CudaSlice<f32>,
20957 k: &CudaSlice<f32>,
20958 v: &CudaSlice<f32>,
20959 g: &CudaSlice<f32>,
20960 beta: &CudaSlice<f32>,
20961 state_in: &CudaSlice<f32>,
20962 state_out: &mut CudaSlice<f32>,
20963 o: &mut CudaSlice<f32>,
20964 n_head: usize,
20965 t: usize,
20966 scale: f32,
20967 ) -> Result<(), Box<dyn std::error::Error>> {
20968 let f = self.func("gdn_scan_s128");
20969 const S_V: u32 = 128;
20970 const WARP: u32 = 32;
20971 const COLS_PER_BLOCK: u32 = 4;
20972 let cfg = LaunchConfig {
20973 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20974 block_dim: (WARP, COLS_PER_BLOCK, 1),
20975 shared_mem_bytes: 0,
20976 };
20977 let (h, ti) = (n_head as i32, t as i32);
20978 let __s_b = self.gpu.stream();
20979 let mut b = __s_b.launch_builder(&f);
20980 b.arg(q)
20981 .arg(k)
20982 .arg(v)
20983 .arg(g)
20984 .arg(beta)
20985 .arg(state_in)
20986 .arg(state_out)
20987 .arg(o)
20988 .arg(&h)
20989 .arg(&ti)
20990 .arg(&scale);
20991 unsafe {
20992 b.launch(cfg)?;
20993 }
20994 Ok(())
20995 }
20996
20997 #[allow(clippy::too_many_arguments)]
21002 pub fn ssm_conv1d_fused_decode_b(
21003 &self,
21004 qkv_cols: &CudaSlice<f32>,
21005 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
21006 w: &CudaSlice<f32>,
21007 conv_outs: &mut CudaSlice<f32>,
21008 conv_dim: usize,
21009 d_conv: usize,
21010 b_n: usize,
21011 ) -> Result<(), Box<dyn std::error::Error>> {
21012 let f = self.func("ssm_conv1d_fused_decode_b_f32");
21013 let cfg = LaunchConfig {
21014 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
21015 block_dim: (256, 1, 1),
21016 shared_mem_bytes: 0,
21017 };
21018 let (cd, dc) = (conv_dim as i32, d_conv as i32);
21019 let __s_b = self.gpu.stream();
21020 let mut b = __s_b.launch_builder(&f);
21021 b.arg(qkv_cols)
21022 .arg(conv_state_ptrs)
21023 .arg(w)
21024 .arg(conv_outs)
21025 .arg(&cd)
21026 .arg(&dc);
21027 unsafe {
21028 b.launch(cfg)?;
21029 }
21030 Ok(())
21031 }
21032
21033 #[allow(clippy::too_many_arguments)]
21034 pub fn gdn_prep_decode_b(
21035 &self,
21036 conv_outs: &CudaSlice<f32>,
21037 beta_raws: &CudaSlice<f32>,
21038 alphas: &CudaSlice<f32>,
21039 dt_bias: &CudaSlice<f32>,
21040 a: &CudaSlice<f32>,
21041 q_l2: &mut CudaSlice<f32>,
21042 k_l2: &mut CudaSlice<f32>,
21043 v_g: &mut CudaSlice<f32>,
21044 beta: &mut CudaSlice<f32>,
21045 g_log: &mut CudaSlice<f32>,
21046 d_state: usize,
21047 num_v: usize,
21048 num_k: usize,
21049 key_dim: usize,
21050 eps: f32,
21051 conv_dim: usize,
21052 b_n: usize,
21053 ) -> Result<(), Box<dyn std::error::Error>> {
21054 let f = self.func("gdn_prep_decode_b_f32");
21055 let cfg = LaunchConfig {
21056 grid_dim: (num_v as u32, 1, b_n as u32),
21057 block_dim: (32, 4, 1),
21058 shared_mem_bytes: 0,
21059 };
21060 let (ds, nv, nk, kd, cd) = (
21061 d_state as i32,
21062 num_v as i32,
21063 num_k as i32,
21064 key_dim as i32,
21065 conv_dim as i32,
21066 );
21067 let __s_b = self.gpu.stream();
21068 let mut b = __s_b.launch_builder(&f);
21069 b.arg(conv_outs)
21070 .arg(beta_raws)
21071 .arg(alphas)
21072 .arg(dt_bias)
21073 .arg(a)
21074 .arg(q_l2)
21075 .arg(k_l2)
21076 .arg(v_g)
21077 .arg(beta)
21078 .arg(g_log)
21079 .arg(&ds)
21080 .arg(&nv)
21081 .arg(&nk)
21082 .arg(&kd)
21083 .arg(&eps)
21084 .arg(&cd);
21085 unsafe {
21086 b.launch(cfg)?;
21087 }
21088 Ok(())
21089 }
21090
21091 #[allow(clippy::too_many_arguments)]
21092 pub fn gdn_scan_s128_batched(
21093 &self,
21094 q: &CudaSlice<f32>,
21095 k: &CudaSlice<f32>,
21096 v: &CudaSlice<f32>,
21097 g: &CudaSlice<f32>,
21098 beta: &CudaSlice<f32>,
21099 state_in_ptrs: &cudarc::driver::CudaView<u64>,
21100 state_out_ptrs: &cudarc::driver::CudaView<u64>,
21101 o: &mut CudaSlice<f32>,
21102 n_head: usize,
21103 b_n: usize,
21104 scale: f32,
21105 ) -> Result<(), Box<dyn std::error::Error>> {
21106 let f = self.func("gdn_scan_s128_b");
21107 const S_V: u32 = 128;
21108 const WARP: u32 = 32;
21109 const COLS_PER_BLOCK: u32 = 4;
21110 let cfg = LaunchConfig {
21111 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
21112 block_dim: (WARP, COLS_PER_BLOCK, 1),
21113 shared_mem_bytes: 0,
21114 };
21115 let h = n_head as i32;
21116 let __s_b = self.gpu.stream();
21117 let mut b = __s_b.launch_builder(&f);
21118 b.arg(q)
21119 .arg(k)
21120 .arg(v)
21121 .arg(g)
21122 .arg(beta)
21123 .arg(state_in_ptrs)
21124 .arg(state_out_ptrs)
21125 .arg(o)
21126 .arg(&h)
21127 .arg(&scale);
21128 unsafe {
21129 b.launch(cfg)?;
21130 }
21131 Ok(())
21132 }
21133
21134 #[allow(clippy::too_many_arguments)]
21140 pub fn ssm_conv1d_fused_decode_b_view(
21141 &self,
21142 qkv_cols: &cudarc::driver::CudaView<f32>,
21143 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
21144 w: &CudaSlice<f32>,
21145 conv_outs: &mut CudaSlice<f32>,
21146 conv_dim: usize,
21147 d_conv: usize,
21148 b_n: usize,
21149 ) -> Result<(), Box<dyn std::error::Error>> {
21150 let f = self.func("ssm_conv1d_fused_decode_b_f32");
21151 let cfg = LaunchConfig {
21152 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
21153 block_dim: (256, 1, 1),
21154 shared_mem_bytes: 0,
21155 };
21156 let (cd, dc) = (conv_dim as i32, d_conv as i32);
21157 let __s_b = self.gpu.stream();
21158 let mut b = __s_b.launch_builder(&f);
21159 b.arg(qkv_cols)
21160 .arg(conv_state_ptrs)
21161 .arg(w)
21162 .arg(conv_outs)
21163 .arg(&cd)
21164 .arg(&dc);
21165 unsafe {
21166 b.launch(cfg)?;
21167 }
21168 Ok(())
21169 }
21170
21171 #[allow(clippy::too_many_arguments)]
21172 pub fn gdn_prep_decode_b_view(
21173 &self,
21174 conv_outs: &CudaSlice<f32>,
21175 beta_raws: &cudarc::driver::CudaView<f32>,
21176 alphas: &cudarc::driver::CudaView<f32>,
21177 dt_bias: &CudaSlice<f32>,
21178 a: &CudaSlice<f32>,
21179 q_l2: &mut CudaSlice<f32>,
21180 k_l2: &mut CudaSlice<f32>,
21181 v_g: &mut CudaSlice<f32>,
21182 beta: &mut CudaSlice<f32>,
21183 g_log: &mut CudaSlice<f32>,
21184 d_state: usize,
21185 num_v: usize,
21186 num_k: usize,
21187 key_dim: usize,
21188 eps: f32,
21189 conv_dim: usize,
21190 b_n: usize,
21191 ) -> Result<(), Box<dyn std::error::Error>> {
21192 let f = self.func("gdn_prep_decode_b_f32");
21193 let cfg = LaunchConfig {
21194 grid_dim: (num_v as u32, 1, b_n as u32),
21195 block_dim: (32, 4, 1),
21196 shared_mem_bytes: 0,
21197 };
21198 let (ds, nv, nk, kd, cd) = (
21199 d_state as i32,
21200 num_v as i32,
21201 num_k as i32,
21202 key_dim as i32,
21203 conv_dim as i32,
21204 );
21205 let __s_b = self.gpu.stream();
21206 let mut b = __s_b.launch_builder(&f);
21207 b.arg(conv_outs)
21208 .arg(beta_raws)
21209 .arg(alphas)
21210 .arg(dt_bias)
21211 .arg(a)
21212 .arg(q_l2)
21213 .arg(k_l2)
21214 .arg(v_g)
21215 .arg(beta)
21216 .arg(g_log)
21217 .arg(&ds)
21218 .arg(&nv)
21219 .arg(&nk)
21220 .arg(&kd)
21221 .arg(&eps)
21222 .arg(&cd);
21223 unsafe {
21224 b.launch(cfg)?;
21225 }
21226 Ok(())
21227 }
21228
21229 #[allow(clippy::too_many_arguments)]
21230 pub fn gdn_scan_s128_batched_view(
21231 &self,
21232 q: &CudaSlice<f32>,
21233 k: &CudaSlice<f32>,
21234 v: &CudaSlice<f32>,
21235 g: &CudaSlice<f32>,
21236 beta: &CudaSlice<f32>,
21237 state_in_ptrs: &cudarc::driver::CudaView<u64>,
21238 state_out_ptrs: &cudarc::driver::CudaView<u64>,
21239 o: &mut cudarc::driver::CudaViewMut<f32>,
21240 n_head: usize,
21241 b_n: usize,
21242 scale: f32,
21243 ) -> Result<(), Box<dyn std::error::Error>> {
21244 let f = self.func("gdn_scan_s128_b");
21245 const S_V: u32 = 128;
21246 const WARP: u32 = 32;
21247 const COLS_PER_BLOCK: u32 = 4;
21248 let cfg = LaunchConfig {
21249 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
21250 block_dim: (WARP, COLS_PER_BLOCK, 1),
21251 shared_mem_bytes: 0,
21252 };
21253 let h = n_head as i32;
21254 let __s_b = self.gpu.stream();
21255 let mut b = __s_b.launch_builder(&f);
21256 b.arg(q)
21257 .arg(k)
21258 .arg(v)
21259 .arg(g)
21260 .arg(beta)
21261 .arg(state_in_ptrs)
21262 .arg(state_out_ptrs)
21263 .arg(o)
21264 .arg(&h)
21265 .arg(&scale);
21266 unsafe {
21267 b.launch(cfg)?;
21268 }
21269 Ok(())
21270 }
21271
21272 pub fn gdn_chunked_enabled() -> bool {
21281 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21282 *E.get_or_init(|| {
21283 std::env::var("MEMRA_GDN_CHUNKED")
21284 .map(|v| v != "0")
21285 .unwrap_or(true)
21286 })
21287 }
21288
21289 pub fn gdn_chunk_size() -> usize {
21294 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21295 *C.get_or_init(|| {
21296 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
21297 .ok()
21298 .and_then(|v| v.parse().ok())
21299 .unwrap_or(32);
21300 c.clamp(32, 128) / 32 * 32
21301 })
21302 }
21303
21304 #[allow(clippy::too_many_arguments)]
21309 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
21312 #[allow(clippy::too_many_arguments)]
21313 pub fn gdn_chunk_k123(
21314 &self,
21315 q: &CudaSlice<f32>,
21316 k: &CudaSlice<f32>,
21317 v: &CudaSlice<f32>,
21318 g: &CudaSlice<f32>,
21319 beta: &CudaSlice<f32>,
21320 wb16: Option<&mut CudaSlice<u8>>,
21321 n_head: usize,
21322 t: usize,
21323 c: usize,
21324 hk: usize,
21325 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
21326 ) -> Result<
21327 (
21328 CudaSlice<f32>,
21329 CudaSlice<f32>,
21330 CudaSlice<f32>,
21331 CudaSlice<f32>,
21332 ),
21333 Box<dyn std::error::Error>,
21334 > {
21335 const D: usize = 128;
21336 let h = n_head;
21337 let nc = (t + c - 1) / c;
21338 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21339 let mut gcum = self.uninit(t * h)?;
21340 let mut a = self.uninit(nc * h * c * c)?;
21341 let mut p = self.uninit(nc * h * c * c)?;
21342 let mut u = self.uninit(nc * h * c * D)?;
21343 let mut w = self.uninit(nc * h * c * D)?;
21344 {
21345 let f = self.func("gdn_chunk_cumgate_f32");
21347 let cfg = LaunchConfig {
21348 grid_dim: (nc as u32, h as u32, 1),
21349 block_dim: (32, 1, 1),
21350 shared_mem_bytes: 0,
21351 };
21352 let __s_b = self.gpu.stream();
21353 let mut b = __s_b.launch_builder(&f);
21354 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
21355 unsafe {
21356 b.launch(cfg)?;
21357 }
21358 }
21359 if let Some((qb, kb, pb)) = k2w {
21360 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
21363 let f = self.func("gdn_k2_wgmma");
21364 let cfg = LaunchConfig {
21365 grid_dim: (nc as u32, h as u32, 1),
21366 block_dim: (128, 1, 1),
21367 shared_mem_bytes: 0,
21368 };
21369 let hki = hk as i32;
21370 let __s_b = self.gpu.stream();
21371 let mut b = __s_b.launch_builder(&f);
21372 b.arg(qb)
21373 .arg(kb)
21374 .arg(&gcum)
21375 .arg(beta)
21376 .arg(&mut a)
21377 .arg(&mut *pb)
21378 .arg(&hi)
21379 .arg(&ti)
21380 .arg(&ci)
21381 .arg(&hki);
21382 unsafe {
21383 b.launch(cfg)?;
21384 }
21385 } else if c <= 64 && !portable_mma_gated() {
21386 let f = self.func("gdn_chunk_attn_f32");
21388 let jt = ((c + 31) / 32) as u32;
21389 let cfg = LaunchConfig {
21390 grid_dim: (nc as u32, h as u32, jt),
21391 block_dim: (256, 1, 1),
21392 shared_mem_bytes: 0,
21393 };
21394 let hki = hk as i32;
21395 let __s_b = self.gpu.stream();
21396 let mut b = __s_b.launch_builder(&f);
21397 b.arg(q)
21398 .arg(k)
21399 .arg(&gcum)
21400 .arg(beta)
21401 .arg(&mut a)
21402 .arg(&mut p)
21403 .arg(&hi)
21404 .arg(&ti)
21405 .arg(&ci)
21406 .arg(&hki);
21407 unsafe {
21408 b.launch(cfg)?;
21409 }
21410 } else {
21411 assert!(
21413 hk == h,
21414 "generic K2 is broadcast-only (de-broadcast rides C==32)"
21415 );
21416 let f = self.func("gdn_chunk_attn_g_f32");
21417 let cfg = LaunchConfig {
21418 grid_dim: (nc as u32, h as u32, 1),
21419 block_dim: (32, 8, 1),
21420 shared_mem_bytes: 0,
21421 };
21422 let __s_b = self.gpu.stream();
21423 let mut b = __s_b.launch_builder(&f);
21424 b.arg(q)
21425 .arg(k)
21426 .arg(&gcum)
21427 .arg(beta)
21428 .arg(&mut a)
21429 .arg(&mut p)
21430 .arg(&hi)
21431 .arg(&ti)
21432 .arg(&ci);
21433 unsafe {
21434 b.launch(cfg)?;
21435 }
21436 }
21437 {
21438 let cfg = LaunchConfig {
21440 grid_dim: (nc as u32, h as u32, 1),
21441 block_dim: (256, 1, 1),
21442 shared_mem_bytes: 0,
21443 };
21444 match c {
21445 32 | 64 => {
21446 let f = self.func(if c == 32 {
21447 "gdn_chunk_solve32_f32"
21448 } else {
21449 "gdn_chunk_solve64_f32"
21450 });
21451 let wb: u64 = match wb16 {
21453 Some(d) => self.addr_u8(d),
21454 None => 0,
21455 };
21456 let hki = hk as i32;
21457 let __s_b = self.gpu.stream();
21458 let mut b = __s_b.launch_builder(&f);
21459 b.arg(v)
21460 .arg(k)
21461 .arg(&a)
21462 .arg(&gcum)
21463 .arg(&mut u)
21464 .arg(&mut w)
21465 .arg(&wb)
21466 .arg(&hi)
21467 .arg(&ti)
21468 .arg(&hki);
21469 unsafe {
21470 b.launch(cfg)?;
21471 }
21472 }
21473 _ => {
21474 assert!(hk == h, "generic K3 is broadcast-only");
21475 let f = self.func("gdn_chunk_solve_f32");
21476 let __s_b = self.gpu.stream();
21477 let mut b = __s_b.launch_builder(&f);
21478 b.arg(v)
21479 .arg(k)
21480 .arg(&a)
21481 .arg(&gcum)
21482 .arg(&mut u)
21483 .arg(&mut w)
21484 .arg(&hi)
21485 .arg(&ti)
21486 .arg(&ci);
21487 unsafe {
21488 b.launch(cfg)?;
21489 }
21490 }
21491 }
21492 }
21493 Ok((gcum, p, u, w))
21494 }
21495
21496 pub fn gdn_db_on() -> bool {
21500 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
21501 }
21502
21503 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
21512 !portable_mma_gated()
21513 && c == 32
21514 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21515 Ok("1") => true,
21516 Ok("0") => false,
21517 _ => gdn_mma_default_on(),
21518 }
21519 }
21520
21521 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
21528 cfg!(memra_hopper_mma)
21529 && self.gdn_mma_enabled(c)
21530 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
21531 }
21532
21533 #[allow(clippy::too_many_arguments)]
21538 pub fn ssm_conv1d_gdn_state_pad(
21539 &self,
21540 qkv_tm: &cudarc::driver::CudaView<f32>,
21541 conv_state: &mut CudaSlice<f32>,
21542 w: &CudaSlice<f32>,
21543 q_g: &mut CudaSlice<f32>,
21544 k_g: &mut CudaSlice<f32>,
21545 v_g: &mut CudaSlice<f32>,
21546 conv_dim: usize,
21547 t: usize,
21548 d_conv: usize,
21549 d_state: usize,
21550 num_v: usize,
21551 num_k: usize,
21552 key_dim: usize,
21553 hk: usize,
21554 pad_len: Option<&CudaSlice<i32>>,
21555 ) -> Result<(), Box<dyn std::error::Error>> {
21556 assert!(
21557 t >= d_conv - 1,
21558 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
21559 );
21560 {
21561 let f = self.func("ssm_conv1d_gdn_state_f32");
21562 let cfg = LaunchConfig {
21563 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
21564 block_dim: (256, 1, 1),
21565 shared_mem_bytes: 0,
21566 };
21567 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21568 let (ds, nv, nk, kd, hki) = (
21569 d_state as i32,
21570 num_v as i32,
21571 num_k as i32,
21572 key_dim as i32,
21573 hk as i32,
21574 );
21575 let __s_b = self.gpu.stream();
21576 let mut b = __s_b.launch_builder(&f);
21577 b.arg(qkv_tm)
21578 .arg(&*conv_state)
21579 .arg(w)
21580 .arg(q_g)
21581 .arg(k_g)
21582 .arg(v_g)
21583 .arg(&cd)
21584 .arg(&ti)
21585 .arg(&dc)
21586 .arg(&ds)
21587 .arg(&nv)
21588 .arg(&nk)
21589 .arg(&kd)
21590 .arg(&hki);
21591 unsafe {
21592 b.launch(cfg)?;
21593 }
21594 }
21595 match pad_len {
21596 Some(len_d) => {
21597 let f = self.func("ssm_conv_ring_update_dev_f32");
21598 let n = conv_dim * (d_conv - 1);
21599 let cfg = LaunchConfig::for_num_elems(n as u32);
21600 let (cd, dc) = (conv_dim as i32, d_conv as i32);
21601 let __s_b = self.gpu.stream();
21602 let mut b = __s_b.launch_builder(&f);
21603 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
21604 unsafe {
21605 b.launch(cfg)?;
21606 }
21607 }
21608 None => {
21609 let f = self.func("ssm_conv_ring_update_f32");
21610 let n = conv_dim * (d_conv - 1);
21611 let cfg = LaunchConfig::for_num_elems(n as u32);
21612 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21613 let __s_b = self.gpu.stream();
21614 let mut b = __s_b.launch_builder(&f);
21615 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
21616 unsafe {
21617 b.launch(cfg)?;
21618 }
21619 }
21620 }
21621 Ok(())
21622 }
21623
21624 pub fn gdn_chunk_alloc(
21628 &self,
21629 n_head: usize,
21630 t: usize,
21631 c: usize,
21632 hk: usize,
21633 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
21634 const D: usize = 128;
21635 assert!(
21636 c == 32,
21637 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
21638 );
21639 let h = n_head;
21640 let nc = (t + c - 1) / c;
21641 Ok(GdnChunkBufs {
21642 gcum: self.uninit(t * h)?,
21643 a: self.uninit(nc * h * c * c)?,
21644 p: self.uninit(nc * h * c * c)?,
21645 u: self.uninit(nc * h * c * D)?,
21646 w: self.uninit(nc * h * c * D)?,
21647 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21648 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21649 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21650 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
21651 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21652 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
21653 o: self.uninit(D * h * t)?,
21654 t,
21655 nc,
21656 })
21657 }
21658
21659 pub fn f32_to_bf16_v(
21661 &self,
21662 x: &cudarc::driver::CudaView<f32>,
21663 dst: &mut CudaSlice<u8>,
21664 n: usize,
21665 ) -> Result<(), Box<dyn std::error::Error>> {
21666 let f = self.func("f32_to_bf16_bulk");
21667 let ni = n as i64;
21668 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21669 let __s_b = self.gpu.stream();
21670 let mut b = __s_b.launch_builder(&f);
21671 b.arg(x).arg(dst).arg(&ni);
21672 unsafe {
21673 b.launch(cfg)?;
21674 }
21675 Ok(())
21676 }
21677
21678 pub fn f32_to_bf16_into(
21680 &self,
21681 x: &CudaSlice<f32>,
21682 dst: &mut CudaSlice<u8>,
21683 n: usize,
21684 ) -> Result<(), Box<dyn std::error::Error>> {
21685 let f = self.func("f32_to_bf16_bulk");
21686 let ni = n as i64;
21687 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21688 let __s_b = self.gpu.stream();
21689 let mut b = __s_b.launch_builder(&f);
21690 b.arg(x).arg(dst).arg(&ni);
21691 unsafe {
21692 b.launch(cfg)?;
21693 }
21694 Ok(())
21695 }
21696
21697 pub fn gdn_chunk_k123_vl8(
21700 &self,
21701 seqs: &[GdnSeqVl],
21702 n_head: usize,
21703 hk: usize,
21704 wq: Option<&GdnWVl8>,
21705 ) -> Result<(), Box<dyn std::error::Error>> {
21706 let b = seqs.len();
21707 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
21708 let mut packed = [GdnSeqVl::default(); 8];
21709 packed[..b].copy_from_slice(seqs);
21710 let v = GdnVl8(packed);
21711 let (hi, ci) = (n_head as i32, 32i32);
21712 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21713 {
21714 let f = self.func("gdn_chunk_cumgate_vl");
21715 let cfg = LaunchConfig {
21716 grid_dim: (max_nc, n_head as u32, b as u32),
21717 block_dim: (32, 1, 1),
21718 shared_mem_bytes: 0,
21719 };
21720 let __s_lb = self.gpu.stream();
21721 let mut lb = __s_lb.launch_builder(&f);
21722 lb.arg(&v).arg(&hi).arg(&ci);
21723 unsafe {
21724 lb.launch(cfg)?;
21725 }
21726 }
21727 let hki = hk as i32;
21728 if let Some(w) = wq {
21729 let f = self.func("gdn_k2_wgmma_vl");
21731 let cfg = LaunchConfig {
21732 grid_dim: (max_nc, n_head as u32, b as u32),
21733 block_dim: (128, 1, 1),
21734 shared_mem_bytes: 0,
21735 };
21736 let __s_lb = self.gpu.stream();
21737 let mut lb = __s_lb.launch_builder(&f);
21738 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
21739 unsafe {
21740 lb.launch(cfg)?;
21741 }
21742 } else {
21743 let f = self.func("gdn_chunk_attn_vl");
21744 let cfg = LaunchConfig {
21745 grid_dim: (max_nc, n_head as u32, b as u32),
21746 block_dim: (256, 1, 1),
21747 shared_mem_bytes: 0,
21748 };
21749 let __s_lb = self.gpu.stream();
21750 let mut lb = __s_lb.launch_builder(&f);
21751 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21752 unsafe {
21753 lb.launch(cfg)?;
21754 }
21755 }
21756 {
21757 let f = self.func("gdn_chunk_solve32_vl");
21758 let cfg = LaunchConfig {
21759 grid_dim: (max_nc, n_head as u32, b as u32),
21760 block_dim: (256, 1, 1),
21761 shared_mem_bytes: 0,
21762 };
21763 let __s_lb = self.gpu.stream();
21764 let mut lb = __s_lb.launch_builder(&f);
21765 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21766 unsafe {
21767 lb.launch(cfg)?;
21768 }
21769 }
21770 Ok(())
21771 }
21772
21773 #[allow(clippy::too_many_arguments)]
21777 pub fn gdn_prep_vl8(
21778 &self,
21779 seqs: &[GdnPrepVl],
21780 conv_w: &CudaSlice<f32>,
21781 dt_bias: &CudaSlice<f32>,
21782 a: &CudaSlice<f32>,
21783 conv_dim: usize,
21784 d_conv: usize,
21785 d_state: usize,
21786 num_v: usize,
21787 num_k: usize,
21788 key_dim: usize,
21789 hk: usize,
21790 eps: f32,
21791 ) -> Result<(), Box<dyn std::error::Error>> {
21792 let b = seqs.len();
21793 assert!(b >= 1 && b <= 8);
21794 let mut packed = [GdnPrepVl::default(); 8];
21795 packed[..b].copy_from_slice(seqs);
21796 let v = GdnPrepVl8(packed);
21797 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21798 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
21799 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
21800 assert!(
21801 conv_fuse || hk == num_v,
21802 "de-broadcast requires the fused conv"
21803 );
21804 if conv_fuse {
21805 let f = self.func("ssm_conv1d_gdn_state_vl");
21806 let cfg = LaunchConfig {
21807 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21808 block_dim: (256, 1, 1),
21809 shared_mem_bytes: 0,
21810 };
21811 let (dsi, nvi, nki, kdi, hki) = (
21812 d_state as i32,
21813 num_v as i32,
21814 num_k as i32,
21815 key_dim as i32,
21816 hk as i32,
21817 );
21818 let __s_lb = self.gpu.stream();
21819 let mut lb = __s_lb.launch_builder(&f);
21820 lb.arg(&v)
21821 .arg(conv_w)
21822 .arg(&cdi)
21823 .arg(&dci)
21824 .arg(&dsi)
21825 .arg(&nvi)
21826 .arg(&nki)
21827 .arg(&kdi)
21828 .arg(&hki);
21829 unsafe {
21830 lb.launch(cfg)?;
21831 }
21832 } else {
21833 let f = self.func("ssm_conv1d_tm_state_vl");
21834 let cfg = LaunchConfig {
21835 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21836 block_dim: (256, 1, 1),
21837 shared_mem_bytes: 0,
21838 };
21839 let __s_lb = self.gpu.stream();
21840 let mut lb = __s_lb.launch_builder(&f);
21841 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21842 unsafe {
21843 lb.launch(cfg)?;
21844 }
21845 }
21846 {
21847 let f = self.func("ssm_conv_ring_update_vl");
21848 let n = (conv_dim * (d_conv - 1)) as u32;
21849 let cfg = LaunchConfig {
21850 grid_dim: (n.div_ceil(256), 1, b as u32),
21851 block_dim: (256, 1, 1),
21852 shared_mem_bytes: 0,
21853 };
21854 let __s_lb = self.gpu.stream();
21855 let mut lb = __s_lb.launch_builder(&f);
21856 lb.arg(&v).arg(&cdi).arg(&dci);
21857 unsafe {
21858 lb.launch(cfg)?;
21859 }
21860 }
21861 if !conv_fuse {
21862 let f = self.func("qkv_to_gdn_repack_vl");
21863 let n = max_t * (num_v * d_state) as u32;
21864 let cfg = LaunchConfig {
21865 grid_dim: (n.div_ceil(256), 1, b as u32),
21866 block_dim: (256, 1, 1),
21867 shared_mem_bytes: 0,
21868 };
21869 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21870 let __s_lb = self.gpu.stream();
21871 let mut lb = __s_lb.launch_builder(&f);
21872 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21873 unsafe {
21874 lb.launch(cfg)?;
21875 }
21876 }
21877 if Self::l2_v2_on(d_state) {
21878 let f = self.func("gdn_l2_v2_vl");
21879 let cfg = LaunchConfig {
21880 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21881 block_dim: (256, 1, 1),
21882 shared_mem_bytes: 0,
21883 };
21884 let (dsi, nvi) = (d_state as i32, hk as i32);
21885 let __s_lb = self.gpu.stream();
21886 let mut lb = __s_lb.launch_builder(&f);
21887 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21888 unsafe {
21889 lb.launch(cfg)?;
21890 }
21891 } else {
21892 let f = self.func("gdn_l2_vl");
21893 let cfg = LaunchConfig {
21894 grid_dim: (max_t * hk as u32, 2, b as u32),
21895 block_dim: (256, 1, 1),
21896 shared_mem_bytes: 0,
21897 };
21898 let (dsi, nvi) = (d_state as i32, hk as i32);
21899 let __s_lb = self.gpu.stream();
21900 let mut lb = __s_lb.launch_builder(&f);
21901 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21902 unsafe {
21903 lb.launch(cfg)?;
21904 }
21905 }
21906 {
21907 let f = self.func("gdn_gate_prep_vl");
21908 let n = max_t * num_v as u32;
21909 let cfg = LaunchConfig {
21910 grid_dim: (n.div_ceil(256), 1, b as u32),
21911 block_dim: (256, 1, 1),
21912 shared_mem_bytes: 0,
21913 };
21914 let nvi = num_v as i32;
21915 let __s_lb = self.gpu.stream();
21916 let mut lb = __s_lb.launch_builder(&f);
21917 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21918 unsafe {
21919 lb.launch(cfg)?;
21920 }
21921 }
21922 Ok(())
21923 }
21924
21925 pub fn gdn_mirror_vl8(
21927 &self,
21928 seqs: &[GdnSeqVl],
21929 n_head: usize,
21930 which: i32,
21931 hk: usize,
21932 ) -> Result<(), Box<dyn std::error::Error>> {
21933 let b = seqs.len();
21934 assert!(b >= 1 && b <= 8);
21935 let mut packed = [GdnSeqVl::default(); 8];
21936 packed[..b].copy_from_slice(seqs);
21937 let v = GdnVl8(packed);
21938 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21939 let max_n = seqs
21940 .iter()
21941 .map(|s| {
21942 if which == 0 {
21943 s.t as i64 * ept as i64
21944 } else {
21945 s.nc as i64 * ept as i64 * 32
21946 }
21947 })
21948 .max()
21949 .unwrap();
21950 let f = self.func("gdn_mirror_vl");
21951 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21952 let cfg = LaunchConfig {
21953 grid_dim: (blocks, 1, b as u32),
21954 block_dim: (256, 1, 1),
21955 shared_mem_bytes: 0,
21956 };
21957 let __s_lb = self.gpu.stream();
21958 let mut lb = __s_lb.launch_builder(&f);
21959 lb.arg(&v).arg(&ept).arg(&which);
21960 unsafe {
21961 lb.launch(cfg)?;
21962 }
21963 Ok(())
21964 }
21965
21966 pub fn gdn_tail_vl8(
21968 &self,
21969 seqs: &[GdnPrepVl],
21970 norm_w: &CudaSlice<f32>,
21971 d_state: usize,
21972 num_v: usize,
21973 eps: f32,
21974 ) -> Result<(), Box<dyn std::error::Error>> {
21975 let b = seqs.len();
21976 assert!(b >= 1 && b <= 8);
21977 let mut packed = [GdnPrepVl::default(); 8];
21978 packed[..b].copy_from_slice(seqs);
21979 let v = GdnPrepVl8(packed);
21980 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21981 let f = self.func("gated_rmsnorm_f16out_vl");
21982 let cfg = LaunchConfig {
21984 grid_dim: (max_t * num_v as u32, 1, b as u32),
21985 block_dim: (128, 1, 1),
21986 shared_mem_bytes: 0,
21987 };
21988 let (dsi, nvi) = (d_state as i32, num_v as i32);
21989 let __s_lb = self.gpu.stream();
21990 let mut lb = __s_lb.launch_builder(&f);
21991 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21992 unsafe {
21993 lb.launch(cfg)?;
21994 }
21995 Ok(())
21996 }
21997
21998 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
22001 use cudarc::driver::DevicePtr;
22002 let s = self.gpu.stream();
22003 let (p, _g) = x.device_ptr(&s);
22004 p as u64
22005 }
22006 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
22007 use cudarc::driver::DevicePtrMut;
22008 let s = self.gpu.stream();
22009 let (p, _g) = x.device_ptr_mut(&s);
22010 p as u64
22011 }
22012 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
22013 use cudarc::driver::DevicePtr;
22014 let s = self.gpu.stream();
22015 let (p, _g) = x.device_ptr(&s);
22016 p as u64
22017 }
22018 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
22019 use cudarc::driver::DevicePtr;
22020 let s = self.gpu.stream();
22021 let (p, _g) = x.device_ptr(&s);
22022 p as u64
22023 }
22024
22025 pub fn gdn_chunk_vl8(
22029 &self,
22030 seqs: &[GdnSeqVl],
22031 n_head: usize,
22032 scale: f32,
22033 hk: usize,
22034 wq: Option<&GdnWVl8>,
22035 ) -> Result<(), Box<dyn std::error::Error>> {
22036 const NSPLIT: u32 = 4;
22037 let b = seqs.len();
22038 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
22039 let mut packed = [GdnSeqVl::default(); 8];
22040 packed[..b].copy_from_slice(seqs);
22041 let v = GdnVl8(packed);
22042 let (hi, ci) = (n_head as i32, 32i32);
22043 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
22044 let hki = hk as i32;
22045 if let Some(w) = wq {
22046 let f = self.func("gdn_k45_wgmma_vl");
22048 let cfg = LaunchConfig {
22049 grid_dim: (n_head as u32, NSPLIT, b as u32),
22050 block_dim: (256, 1, 1),
22051 shared_mem_bytes: 0,
22052 };
22053 let __s_lb = self.gpu.stream();
22054 let mut lb = __s_lb.launch_builder(&f);
22055 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
22056 unsafe {
22057 lb.launch(cfg)?;
22058 }
22059 let _ = max_nc;
22060 return Ok(());
22061 }
22062 {
22063 let f = self.func("gdn_chunk_state_mma_vl");
22064 let cfg = LaunchConfig {
22065 grid_dim: (n_head as u32, NSPLIT, b as u32),
22066 block_dim: (256, 1, 1),
22067 shared_mem_bytes: 0,
22068 };
22069 let __s_lb = self.gpu.stream();
22070 let mut lb = __s_lb.launch_builder(&f);
22071 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
22072 unsafe {
22073 lb.launch(cfg)?;
22074 }
22075 }
22076 {
22077 let f = self.func("gdn_chunk_output_mma_vl");
22078 let cfg = LaunchConfig {
22079 grid_dim: (max_nc, n_head as u32, b as u32),
22080 block_dim: (256, 1, 1),
22081 shared_mem_bytes: 0,
22082 };
22083 let __s_lb = self.gpu.stream();
22084 let mut lb = __s_lb.launch_builder(&f);
22085 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
22086 unsafe {
22087 lb.launch(cfg)?;
22088 }
22089 }
22090 Ok(())
22091 }
22092 pub fn gdn_scan_chunked(
22093 &self,
22094 q: &CudaSlice<f32>,
22095 k: &CudaSlice<f32>,
22096 v: &CudaSlice<f32>,
22097 g: &CudaSlice<f32>,
22098 beta: &CudaSlice<f32>,
22099 kb16_pre: Option<&CudaSlice<u8>>,
22100 qb16_pre: Option<&CudaSlice<u8>>,
22101 state_in: &CudaSlice<f32>,
22102 state_out: &mut CudaSlice<f32>,
22103 o: &mut CudaSlice<f32>,
22104 n_head: usize,
22105 t: usize,
22106 scale: f32,
22107 c: usize,
22108 hk: usize,
22109 ) -> Result<(), Box<dyn std::error::Error>> {
22110 const D: usize = 128;
22111 const NSPLIT: u32 = 4;
22112 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
22113 let h = n_head;
22114 let nc = (t + c - 1) / c;
22115 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
22116 let gdn_mma_pre = !portable_mma_gated()
22121 && c == 32
22122 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
22123 Ok("1") => true,
22124 Ok("0") => false,
22125 _ => gdn_mma_default_on(),
22126 };
22127 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
22128 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
22129 } else {
22130 None
22131 };
22132 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
22137 && gdn_mma_pre
22138 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
22139 let nk = t * hk * D;
22140 let mut kb16_local: Option<CudaSlice<u8>> = None;
22141 if gdn_mma_pre && kb16_pre.is_none() {
22142 let mut kb = self.alloc_u8_uninit(nk * 2)?;
22143 let f = self.func("f32_to_bf16_bulk");
22144 let n2 = nk as i64;
22145 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
22146 let __s_b = self.gpu.stream();
22147 let mut b = __s_b.launch_builder(&f);
22148 b.arg(k).arg(&mut kb).arg(&n2);
22149 unsafe {
22150 b.launch(cfg2)?;
22151 }
22152 kb16_local = Some(kb);
22153 }
22154 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
22155 if let Some(kb) = kb16_pre {
22156 assert!(kb.len() >= nk * 2, "kb16_pre too small");
22157 }
22158 let mut qb16: Option<CudaSlice<u8>> = None;
22159 let mut pb16: Option<CudaSlice<u8>> = None;
22160 if gdn_wgmma_pre {
22161 if qb16_pre.is_none() {
22164 let mut qb = self.alloc_u8_uninit(nk * 2)?;
22165 let f = self.func("f32_to_bf16_bulk");
22166 let n2 = nk as i64;
22167 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
22168 let __s_b = self.gpu.stream();
22169 let mut b = __s_b.launch_builder(&f);
22170 b.arg(q).arg(&mut qb).arg(&n2);
22171 unsafe {
22172 b.launch(cfg2)?;
22173 }
22174 qb16 = Some(qb);
22175 } else if let Some(qb) = qb16_pre {
22176 assert!(qb.len() >= nk * 2, "qb16_pre too small");
22177 }
22178 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
22179 }
22180 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
22181 let k2w = if gdn_wgmma_pre {
22182 Some((
22183 *qb16_ref0.as_ref().unwrap(),
22184 *kb16_ref0.as_ref().unwrap(),
22185 pb16.as_mut().unwrap(),
22186 ))
22187 } else {
22188 None
22189 };
22190 let (gcum, p, u, w) =
22191 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
22192 let _ = &w;
22193 let mut y = self.uninit(nc * h * c * D)?;
22194 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
22208 && c == 32
22209 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
22210 Ok("1") => true,
22211 Ok("0") => false,
22212 _ => gdn_mma_default_on(),
22213 };
22214 if gdn_mma {
22215 let wb16 = wb16_pre
22216 .take()
22217 .expect("mma path pre-allocates wb16 (K3 store fold)");
22218 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
22219 if gdn_wgmma_pre {
22231 let qb16 = qb16_ref0.unwrap();
22233 let pb16 = pb16.as_ref().unwrap();
22234 {
22235 let f = self.func("gdn_k45_wgmma");
22236 let cfg = LaunchConfig {
22237 grid_dim: (h as u32, 4, 1),
22238 block_dim: (256, 1, 1),
22239 shared_mem_bytes: 0,
22240 };
22241 let hki = hk as i32;
22242 let __s_b = self.gpu.stream();
22243 let mut b = __s_b.launch_builder(&f);
22244 b.arg(kb16_ref)
22245 .arg(&gcum)
22246 .arg(beta)
22247 .arg(&u)
22248 .arg(&wb16)
22249 .arg(qb16)
22250 .arg(pb16)
22251 .arg(o)
22252 .arg(&scale)
22253 .arg(state_in)
22254 .arg(&mut *state_out)
22255 .arg(&hi)
22256 .arg(&ti)
22257 .arg(&ci)
22258 .arg(&hki);
22259 unsafe {
22260 b.launch(cfg)?;
22261 }
22262 }
22263 return Ok(());
22264 }
22265 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
22269 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
22270 {
22271 let f = self.func("gdn_chunk_state_mma");
22272 let cfg = LaunchConfig {
22273 grid_dim: (h as u32, NSPLIT, 1),
22274 block_dim: (256, 1, 1),
22275 shared_mem_bytes: 0,
22276 };
22277 let hki = hk as i32;
22278 let __s_b = self.gpu.stream();
22279 let mut b = __s_b.launch_builder(&f);
22280 b.arg(kb16_ref)
22281 .arg(&gcum)
22282 .arg(beta)
22283 .arg(&u)
22284 .arg(&wb16)
22285 .arg(&mut y16)
22286 .arg(&mut ssnap16)
22287 .arg(state_in)
22288 .arg(&mut *state_out)
22289 .arg(&hi)
22290 .arg(&ti)
22291 .arg(&ci)
22292 .arg(&hki);
22293 unsafe {
22294 b.launch(cfg)?;
22295 }
22296 }
22297 {
22298 let f = self.func("gdn_chunk_output_mma");
22300 let jt = ((c + 31) / 32) as u32;
22301 let cfg = LaunchConfig {
22302 grid_dim: (nc as u32, h as u32, jt),
22303 block_dim: (256, 1, 1),
22304 shared_mem_bytes: 0,
22305 };
22306 let hki = hk as i32;
22307 let __s_b = self.gpu.stream();
22308 let mut b = __s_b.launch_builder(&f);
22309 b.arg(q)
22310 .arg(&gcum)
22311 .arg(&p)
22312 .arg(&y16)
22313 .arg(&ssnap16)
22314 .arg(o)
22315 .arg(&hi)
22316 .arg(&ti)
22317 .arg(&ci)
22318 .arg(&scale)
22319 .arg(&hki);
22320 unsafe {
22321 b.launch(cfg)?;
22322 }
22323 }
22324 return Ok(());
22325 }
22326 {
22327 let f = self.func("gdn_chunk_state_f32");
22329 let cfg = LaunchConfig {
22330 grid_dim: (h as u32, NSPLIT, 1),
22331 block_dim: (256, 1, 1),
22332 shared_mem_bytes: 0,
22333 };
22334 let __s_b = self.gpu.stream();
22335 let mut b = __s_b.launch_builder(&f);
22336 b.arg(k)
22337 .arg(&gcum)
22338 .arg(beta)
22339 .arg(&u)
22340 .arg(&w)
22341 .arg(&mut y)
22342 .arg(&mut ssnap)
22343 .arg(state_in)
22344 .arg(&mut *state_out)
22345 .arg(&hi)
22346 .arg(&ti)
22347 .arg(&ci);
22348 unsafe {
22349 b.launch(cfg)?;
22350 }
22351 }
22352 {
22353 let f = self.func("gdn_chunk_output_f32");
22355 let jt = ((c + 31) / 32) as u32;
22356 let cfg = LaunchConfig {
22357 grid_dim: (nc as u32, h as u32, jt),
22358 block_dim: (256, 1, 1),
22359 shared_mem_bytes: 0,
22360 };
22361 let __s_b = self.gpu.stream();
22362 let mut b = __s_b.launch_builder(&f);
22363 b.arg(q)
22364 .arg(&gcum)
22365 .arg(&p)
22366 .arg(&y)
22367 .arg(&ssnap)
22368 .arg(o)
22369 .arg(&hi)
22370 .arg(&ti)
22371 .arg(&ci)
22372 .arg(&scale);
22373 unsafe {
22374 b.launch(cfg)?;
22375 }
22376 }
22377 Ok(())
22378 }
22379
22380 #[allow(clippy::too_many_arguments)]
22389 #[allow(clippy::too_many_arguments)]
22390 pub fn gdn_scan_prefill(
22391 &self,
22392 q: &CudaSlice<f32>,
22393 k: &CudaSlice<f32>,
22394 v: &CudaSlice<f32>,
22395 g: &CudaSlice<f32>,
22396 beta: &CudaSlice<f32>,
22397 kb16_pre: Option<&CudaSlice<u8>>,
22398 qb16_pre: Option<&CudaSlice<u8>>,
22399 state_in: &CudaSlice<f32>,
22400 state_out: &mut CudaSlice<f32>,
22401 o: &mut CudaSlice<f32>,
22402 n_head: usize,
22403 t: usize,
22404 scale: f32,
22405 hk: usize,
22406 ) -> Result<(), Box<dyn std::error::Error>> {
22407 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
22408 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
22409 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
22410 }
22411 if Self::gdn_chunked_enabled() && t >= 16 {
22412 self.gdn_scan_chunked(
22413 q,
22414 k,
22415 v,
22416 g,
22417 beta,
22418 kb16_pre,
22419 qb16_pre,
22420 state_in,
22421 state_out,
22422 o,
22423 n_head,
22424 t,
22425 scale,
22426 Self::gdn_chunk_size(),
22427 hk,
22428 )
22429 } else {
22430 assert!(
22431 hk == n_head,
22432 "s128 scan is broadcast-only (prep guarantees by predicate)"
22433 );
22434 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
22435 }
22436 }
22437
22438 #[allow(clippy::too_many_arguments)]
22440 fn gdn_scan_diff(
22441 &self,
22442 q: &CudaSlice<f32>,
22443 k: &CudaSlice<f32>,
22444 v: &CudaSlice<f32>,
22445 g: &CudaSlice<f32>,
22446 beta: &CudaSlice<f32>,
22447 state_in: &CudaSlice<f32>,
22448 state_out: &mut CudaSlice<f32>,
22449 o: &mut CudaSlice<f32>,
22450 n_head: usize,
22451 t: usize,
22452 scale: f32,
22453 ) -> Result<(), Box<dyn std::error::Error>> {
22454 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
22455 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
22456 let mut o_c = self.uninit(o.len())?;
22457 let mut st_c = self.uninit(state_out.len())?;
22458 self.gdn_scan_chunked(
22459 q,
22460 k,
22461 v,
22462 g,
22463 beta,
22464 None,
22465 None,
22466 state_in,
22467 &mut st_c,
22468 &mut o_c,
22469 n_head,
22470 t,
22471 scale,
22472 Self::gdn_chunk_size(),
22473 n_head,
22474 )?;
22475 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
22476 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
22477 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
22478 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
22479 let mut max_abs = 0f32;
22480 let mut max_rel = 0f32;
22481 let mut sum_rel = 0f64;
22482 for (x, y) in a.iter().zip(b) {
22483 let ad = (x - y).abs();
22484 let rel = ad / x.abs().max(y.abs()).max(1e-3);
22485 if ad > max_abs {
22486 max_abs = ad;
22487 }
22488 if rel > max_rel {
22489 max_rel = rel;
22490 }
22491 sum_rel += rel as f64;
22492 }
22493 (max_abs, max_rel, sum_rel / a.len() as f64)
22494 };
22495 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
22496 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
22497 println!(
22498 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
22499 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
22500 Self::gdn_chunk_size()
22501 );
22502 Ok(())
22503 }
22504
22505 pub fn gdn_glog(
22507 &self,
22508 alpha: &CudaSlice<f32>,
22509 dt_bias: &CudaSlice<f32>,
22510 a: &CudaSlice<f32>,
22511 g_log: &mut CudaSlice<f32>,
22512 n_head: usize,
22513 t: usize,
22514 ) -> Result<(), Box<dyn std::error::Error>> {
22515 let f = self.func("gdn_glog_f32");
22516 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22517 let (h, ti) = (n_head as i32, t as i32);
22518 let __s_b = self.gpu.stream();
22519 let mut b = __s_b.launch_builder(&f);
22520 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22521 unsafe {
22522 b.launch(cfg)?;
22523 }
22524 Ok(())
22525 }
22526
22527 pub fn sigmoid_v(
22530 &self,
22531 x: &cudarc::driver::CudaView<f32>,
22532 y: &mut CudaSlice<f32>,
22533 n: usize,
22534 ) -> Result<(), Box<dyn std::error::Error>> {
22535 let f = self.func("sigmoid_f32");
22536 let cfg = LaunchConfig::for_num_elems(n as u32);
22537 let ni = n as i32;
22538 let __s_b = self.gpu.stream();
22539 let mut b = __s_b.launch_builder(&f);
22540 b.arg(x).arg(y).arg(&ni);
22541 unsafe {
22542 b.launch(cfg)?;
22543 }
22544 Ok(())
22545 }
22546
22547 pub fn gdn_glog_v(
22548 &self,
22549 alpha: &cudarc::driver::CudaView<f32>,
22550 dt_bias: &CudaSlice<f32>,
22551 a: &CudaSlice<f32>,
22552 g_log: &mut CudaSlice<f32>,
22553 n_head: usize,
22554 t: usize,
22555 ) -> Result<(), Box<dyn std::error::Error>> {
22556 let f = self.func("gdn_glog_f32");
22557 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22558 let (h, ti) = (n_head as i32, t as i32);
22559 let __s_b = self.gpu.stream();
22560 let mut b = __s_b.launch_builder(&f);
22561 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22562 unsafe {
22563 b.launch(cfg)?;
22564 }
22565 Ok(())
22566 }
22567
22568 pub fn sigmoid(
22569 &self,
22570 x: &CudaSlice<f32>,
22571 y: &mut CudaSlice<f32>,
22572 n: usize,
22573 ) -> Result<(), Box<dyn std::error::Error>> {
22574 let f = self.func("sigmoid_f32");
22575 let cfg = LaunchConfig::for_num_elems(n as u32);
22576 let ni = n as i32;
22577 let __s_b = self.gpu.stream();
22578 let mut b = __s_b.launch_builder(&f);
22579 b.arg(x).arg(y).arg(&ni);
22580 unsafe {
22581 b.launch(cfg)?;
22582 }
22583 Ok(())
22584 }
22585
22586 pub fn sig_mul_f16out(
22589 &self,
22590 a: &CudaSlice<f32>,
22591 g: &CudaSlice<f32>,
22592 dst: &mut CudaSlice<f32>,
22593 dst16: &mut CudaSlice<u8>,
22594 n: usize,
22595 ) -> Result<(), Box<dyn std::error::Error>> {
22596 let f = self.func("sig_mul_f16out_f32");
22597 let cfg = LaunchConfig::for_num_elems(n as u32);
22598 let ni = n as i32;
22599 let __s_b = self.gpu.stream();
22600 let mut b = __s_b.launch_builder(&f);
22601 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
22602 unsafe {
22603 b.launch(cfg)?;
22604 }
22605 Ok(())
22606 }
22607
22608 #[allow(clippy::too_many_arguments)]
22617 pub fn attn_head_gate(
22618 &self,
22619 a: &CudaSlice<f32>,
22620 g: &CudaSlice<f32>,
22621 dst: &mut CudaSlice<f32>,
22622 dst16: Option<&mut CudaSlice<u8>>,
22623 head_dim: usize,
22624 n_head: usize,
22625 t: usize,
22626 ) -> Result<(), Box<dyn std::error::Error>> {
22627 let f = self.func("attn_head_gate_f32");
22628 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22629 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22630 let d16: u64 = match dst16 {
22632 Some(d) => self.addr_u8(d),
22633 None => 0,
22634 };
22635 let __s_b = self.gpu.stream();
22636 let mut b = __s_b.launch_builder(&f);
22637 b.arg(a)
22638 .arg(g)
22639 .arg(dst)
22640 .arg(&d16)
22641 .arg(&hd)
22642 .arg(&nh)
22643 .arg(&ti);
22644 unsafe {
22645 b.launch(cfg)?;
22646 }
22647 Ok(())
22648 }
22649
22650 #[allow(clippy::too_many_arguments)]
22659 pub fn swiglu_clamped_mul_scaled(
22660 &self,
22661 gate: &CudaSlice<f32>,
22662 up: &CudaSlice<f32>,
22663 gs: f32,
22664 us: f32,
22665 limit: f32,
22666 dst: &mut CudaSlice<f32>,
22667 n: usize,
22668 ) -> Result<(), Box<dyn std::error::Error>> {
22669 debug_assert!(
22670 limit > 1e-6,
22671 "swiglu_clamped needs a live limit; use silu_mul_scaled"
22672 );
22673 let f = self.func("swiglu_clamped_mul_scaled_f32");
22674 let cfg = LaunchConfig::for_num_elems(n as u32);
22675 let ni = n as i32;
22676 let __s_b = self.gpu.stream();
22677 let mut b = __s_b.launch_builder(&f);
22678 b.arg(gate)
22679 .arg(up)
22680 .arg(&gs)
22681 .arg(&us)
22682 .arg(&limit)
22683 .arg(dst)
22684 .arg(&ni);
22685 unsafe {
22686 b.launch(cfg)?;
22687 }
22688 Ok(())
22689 }
22690
22691 pub fn gated_rmsnorm(
22693 &self,
22694 o: &CudaSlice<f32>,
22695 w: &CudaSlice<f32>,
22696 z: &CudaSlice<f32>,
22697 dst: &mut CudaSlice<f32>,
22698 ncols: usize,
22699 nrows: usize,
22700 eps: f32,
22701 ) -> Result<(), Box<dyn std::error::Error>> {
22702 let f = self.func("gated_rmsnorm_f32");
22703 let cfg = LaunchConfig {
22704 grid_dim: (nrows as u32, 1, 1),
22705 block_dim: (128, 1, 1),
22706 shared_mem_bytes: 0,
22707 };
22708 let (nc, e) = (ncols as i32, eps);
22709 let __s_b = self.gpu.stream();
22710 let mut b = __s_b.launch_builder(&f);
22711 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22712 unsafe {
22713 b.launch(cfg)?;
22714 }
22715 Ok(())
22716 }
22717
22718 pub fn gated_rmsnorm_f16out(
22721 &self,
22722 o: &CudaSlice<f32>,
22723 w: &CudaSlice<f32>,
22724 z: &CudaSlice<f32>,
22725 dst: &mut CudaSlice<f32>,
22726 dst16: &mut CudaSlice<u8>,
22727 ncols: usize,
22728 nrows: usize,
22729 eps: f32,
22730 ) -> Result<(), Box<dyn std::error::Error>> {
22731 let f = self.func("gated_rmsnorm_f16out_f32");
22732 let cfg = LaunchConfig {
22734 grid_dim: (nrows as u32, 1, 1),
22735 block_dim: (128, 1, 1),
22736 shared_mem_bytes: 0,
22737 };
22738 let (nc, e) = (ncols as i32, eps);
22739 let __s_b = self.gpu.stream();
22740 let mut b = __s_b.launch_builder(&f);
22741 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22742 unsafe {
22743 b.launch(cfg)?;
22744 }
22745 Ok(())
22746 }
22747
22748 #[allow(clippy::too_many_arguments)]
22752 pub fn add_rms_norm_zq8(
22753 &self,
22754 a: &CudaSlice<f32>,
22755 b_in: &CudaSlice<f32>,
22756 w: &CudaSlice<f32>,
22757 res: &mut CudaSlice<f32>,
22758 z: &mut CudaSlice<f32>,
22759 ncols: usize,
22760 nrows: usize,
22761 eps: f32,
22762 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22763 assert!(ncols % 32 == 0);
22764 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
22765 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22766 let f = self.func("add_rms_norm_zq8");
22767 let cfg = LaunchConfig {
22768 grid_dim: (nrows as u32, 1, 1),
22769 block_dim: (1024, 1, 1),
22770 shared_mem_bytes: 0,
22771 };
22772 let (nc, ep) = (ncols as i32, eps);
22773 let __s_b = self.gpu.stream();
22774 let mut b = __s_b.launch_builder(&f);
22775 b.arg(a)
22776 .arg(b_in)
22777 .arg(w)
22778 .arg(res)
22779 .arg(z)
22780 .arg(&mut q)
22781 .arg(&mut d)
22782 .arg(&nc)
22783 .arg(&ep);
22784 unsafe {
22785 b.launch(cfg)?;
22786 }
22787 Ok((q, d))
22788 }
22789
22790 pub fn gated_rmsnorm_zv(
22795 &self,
22796 o: &CudaSlice<f32>,
22797 w: &CudaSlice<f32>,
22798 z: &cudarc::driver::CudaView<f32>,
22799 dst: &mut CudaSlice<f32>,
22800 ncols: usize,
22801 nrows: usize,
22802 eps: f32,
22803 ) -> Result<(), Box<dyn std::error::Error>> {
22804 let f = self.func("gated_rmsnorm_f32");
22805 let cfg = LaunchConfig {
22806 grid_dim: (nrows as u32, 1, 1),
22807 block_dim: (128, 1, 1),
22808 shared_mem_bytes: 0,
22809 };
22810 let (nc, e) = (ncols as i32, eps);
22811 let __s_b = self.gpu.stream();
22812 let mut b = __s_b.launch_builder(&f);
22813 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22814 unsafe {
22815 b.launch(cfg)?;
22816 }
22817 Ok(())
22818 }
22819
22820 pub fn gated_rmsnorm_f16out_zv(
22821 &self,
22822 o: &CudaSlice<f32>,
22823 w: &CudaSlice<f32>,
22824 z: &cudarc::driver::CudaView<f32>,
22825 dst: &mut CudaSlice<f32>,
22826 dst16: &mut CudaSlice<u8>,
22827 ncols: usize,
22828 nrows: usize,
22829 eps: f32,
22830 ) -> Result<(), Box<dyn std::error::Error>> {
22831 let f = self.func("gated_rmsnorm_f16out_f32");
22832 let cfg = LaunchConfig {
22834 grid_dim: (nrows as u32, 1, 1),
22835 block_dim: (128, 1, 1),
22836 shared_mem_bytes: 0,
22837 };
22838 let (nc, e) = (ncols as i32, eps);
22839 let __s_b = self.gpu.stream();
22840 let mut b = __s_b.launch_builder(&f);
22841 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22842 unsafe {
22843 b.launch(cfg)?;
22844 }
22845 Ok(())
22846 }
22847
22848 pub fn gated_rmsnorm_q8_1(
22849 &self,
22850 o: &CudaSlice<f32>,
22851 w: &CudaSlice<f32>,
22852 z: &CudaSlice<f32>,
22853 ncols: usize,
22854 nrows: usize,
22855 eps: f32,
22856 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22857 assert!(ncols % 32 == 0);
22858 let f = self.func("gated_rmsnorm_q8_1");
22859 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22860 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22861 let cfg = LaunchConfig {
22862 grid_dim: (nrows as u32, 1, 1),
22863 block_dim: (128, 1, 1),
22864 shared_mem_bytes: 0,
22865 };
22866 let (nc, ep) = (ncols as i32, eps);
22867 let __s_b = self.gpu.stream();
22868 let mut b = __s_b.launch_builder(&f);
22869 b.arg(o)
22870 .arg(w)
22871 .arg(z)
22872 .arg(&mut out_q)
22873 .arg(&mut out_d)
22874 .arg(&nc)
22875 .arg(&ep);
22876 unsafe {
22877 b.launch(cfg)?;
22878 }
22879 Ok((out_q, out_d))
22880 }
22881
22882 pub fn transpose(
22884 &self,
22885 inp: &CudaSlice<f32>,
22886 rows: usize,
22887 cols: usize,
22888 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22889 let f = self.func("transpose_f32");
22890 let mut out = self.zeros(rows * cols)?;
22891 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22892 let (r, c) = (rows as i32, cols as i32);
22893 let __s_b = self.gpu.stream();
22894 let mut b = __s_b.launch_builder(&f);
22895 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22896 unsafe {
22897 b.launch(cfg)?;
22898 }
22899 Ok(out)
22900 }
22901
22902 pub fn repeat_heads(
22904 &self,
22905 inp: &CudaSlice<f32>,
22906 out: &mut CudaSlice<f32>,
22907 head_dim: usize,
22908 n_in: usize,
22909 n_out: usize,
22910 t: usize,
22911 ) -> Result<(), Box<dyn std::error::Error>> {
22912 let f = self.func("repeat_heads_f32");
22913 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22914 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22915 let __s_b = self.gpu.stream();
22916 let mut b = __s_b.launch_builder(&f);
22917 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22918 unsafe {
22919 b.launch(cfg)?;
22920 }
22921 Ok(())
22922 }
22923
22924 pub fn q_gate_split(
22931 &self,
22932 qf: &CudaSlice<f32>,
22933 q_out: &mut CudaSlice<f32>,
22934 gate_out: &mut CudaSlice<f32>,
22935 head_dim: usize,
22936 n_head: usize,
22937 t: usize,
22938 ) -> Result<(), Box<dyn std::error::Error>> {
22939 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
22940 let out_need = head_dim * n_head * t;
22941 if q_out.len() < out_need || gate_out.len() < out_need {
22942 return Err(format!(
22943 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
22944 q_out.len(),
22945 gate_out.len()
22946 )
22947 .into());
22948 }
22949 let f = self.func("q_gate_split_f32");
22950 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22951 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22952 let __s_b = self.gpu.stream();
22953 let mut b = __s_b.launch_builder(&f);
22954 b.arg(qf)
22955 .arg(q_out)
22956 .arg(gate_out)
22957 .arg(&hd)
22958 .arg(&nh)
22959 .arg(&ti);
22960 unsafe {
22961 b.launch(cfg)?;
22962 }
22963 Ok(())
22964 }
22965
22966 pub fn qkv_to_gdn_repack(
22970 &self,
22971 conv_out: &CudaSlice<f32>,
22972 q_g: &mut CudaSlice<f32>,
22973 k_g: &mut CudaSlice<f32>,
22974 v_g: &mut CudaSlice<f32>,
22975 d_state: usize,
22976 num_v: usize,
22977 num_k: usize,
22978 key_dim: usize,
22979 t: usize,
22980 ) -> Result<(), Box<dyn std::error::Error>> {
22981 let f = self.func("qkv_to_gdn_repack_f32");
22982 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22983 let (ds, nv, nk, kd, ti) = (
22984 d_state as i32,
22985 num_v as i32,
22986 num_k as i32,
22987 key_dim as i32,
22988 t as i32,
22989 );
22990 let __s_b = self.gpu.stream();
22991 let mut b = __s_b.launch_builder(&f);
22992 b.arg(conv_out)
22993 .arg(q_g)
22994 .arg(k_g)
22995 .arg(v_g)
22996 .arg(&ds)
22997 .arg(&nv)
22998 .arg(&nk)
22999 .arg(&kd)
23000 .arg(&ti);
23001 unsafe {
23002 b.launch(cfg)?;
23003 }
23004 Ok(())
23005 }
23006
23007 pub fn conv_left_pad(
23010 &self,
23011 src: &CudaSlice<f32>,
23012 dst: &mut CudaSlice<f32>,
23013 conv_dim: usize,
23014 t: usize,
23015 pad: usize,
23016 ) -> Result<(), Box<dyn std::error::Error>> {
23017 let f = self.func("conv_left_pad_f32");
23018 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
23019 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
23020 let __s_b = self.gpu.stream();
23021 let mut b = __s_b.launch_builder(&f);
23022 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
23023 unsafe {
23024 b.launch(cfg)?;
23025 }
23026 Ok(())
23027 }
23028
23029 pub fn conv_assemble_and_roll(
23033 &self,
23034 qkv_col: &CudaSlice<f32>,
23035 conv_state: &mut CudaSlice<f32>,
23036 conv_in: &mut CudaSlice<f32>,
23037 conv_dim: usize,
23038 pad: usize,
23039 ) -> Result<(), Box<dyn std::error::Error>> {
23040 let f = self.func("conv_assemble_and_roll_f32");
23041 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
23042 let (cd, p) = (conv_dim as i32, pad as i32);
23043 let __s_b = self.gpu.stream();
23044 let mut b = __s_b.launch_builder(&f);
23045 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
23046 unsafe {
23047 b.launch(cfg)?;
23048 }
23049 Ok(())
23050 }
23051
23052 pub fn ssm_conv1d_fused_decode(
23058 &self,
23059 qkv_col: &CudaSlice<f32>,
23060 conv_state: &mut CudaSlice<f32>,
23061 w: &CudaSlice<f32>,
23062 conv_out: &mut CudaSlice<f32>,
23063 conv_dim: usize,
23064 d_conv: usize,
23065 ) -> Result<(), Box<dyn std::error::Error>> {
23066 let f = self.func("ssm_conv1d_fused_decode_f32");
23067 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
23068 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23069 let __s_b = self.gpu.stream();
23070 let mut b = __s_b.launch_builder(&f);
23071 b.arg(qkv_col)
23072 .arg(conv_state)
23073 .arg(w)
23074 .arg(conv_out)
23075 .arg(&cd)
23076 .arg(&dc);
23077 unsafe {
23078 b.launch(cfg)?;
23079 }
23080 Ok(())
23081 }
23082
23083 pub fn slice_range(
23086 &self,
23087 src: &CudaSlice<f32>,
23088 start: usize,
23089 len: usize,
23090 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23091 let host = self.gpu.stream().clone_dtoh(src)?;
23092 self.gpu.stream().synchronize()?;
23093 Ok(self.htod(&host[start..start + len])?)
23094 }
23095}
23096
23097#[cfg(test)]
23098mod target_dispatch_tests {
23099 use super::legacy_quant_gemm_allowed;
23100
23101 #[test]
23102 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
23103 assert!(legacy_quant_gemm_allowed(false, false, false));
23105 assert!(!legacy_quant_gemm_allowed(false, false, true));
23106 assert!(!legacy_quant_gemm_allowed(true, false, false));
23108 assert!(!legacy_quant_gemm_allowed(true, false, true));
23109 assert!(legacy_quant_gemm_allowed(true, true, false));
23111 assert!(!legacy_quant_gemm_allowed(true, true, true));
23112 }
23113
23114 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
23115 #[test]
23116 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
23117 assert!(!legacy_quant_gemm_allowed(
23118 cfg!(memra_portable_cuda),
23119 cfg!(memra_hopper_mma),
23120 false
23121 ));
23122 }
23123
23124 #[cfg(memra_hopper_mma)]
23125 #[test]
23126 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
23127 assert!(legacy_quant_gemm_allowed(
23128 cfg!(memra_portable_cuda),
23129 cfg!(memra_hopper_mma),
23130 false
23131 ));
23132 assert!(super::portable_mma_gated() == false);
23133 }
23134}
23135
23136impl memra_kv::KvDev for Engine {
23139 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23140 Engine::zeros(self, n)
23141 }
23142 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23143 Engine::uninit(self, n)
23144 }
23145 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
23146 Engine::alloc_u8(self, n)
23147 }
23148 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
23149 Engine::htod_i32(self, v)
23150 }
23151 fn clone_dtod(
23152 &self,
23153 src: &CudaSlice<f32>,
23154 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23155 Engine::clone_dtod(self, src)
23156 }
23157 fn copy_into(
23158 &self,
23159 dst: &mut CudaSlice<f32>,
23160 off: usize,
23161 src: &CudaSlice<f32>,
23162 len: usize,
23163 ) -> Result<(), Box<dyn std::error::Error>> {
23164 Engine::copy_into(self, dst, off, src, len)
23165 }
23166 fn set_i32_one(
23167 &self,
23168 d: &mut CudaSlice<i32>,
23169 v: i32,
23170 ) -> Result<(), Box<dyn std::error::Error>> {
23171 Engine::set_i32_one(self, d, v)
23172 }
23173}
23174
23175#[cfg(test)]
23176mod fused_gate_bounds_tests {
23177 use super::*;
23178
23179 #[test]
23192 #[ignore = "requires a CUDA GPU"]
23193 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
23194 let e = Engine::new(0).unwrap();
23195 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
23196 let fused = 2 * head_dim * n_head * t;
23197 let out_n = head_dim * n_head * t;
23198
23199 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
23201 let mut q = e.uninit(out_n).unwrap();
23202 let mut gate = e.uninit(out_n).unwrap();
23203 let err = e
23204 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
23205 .expect_err("half-width wq must be refused, not read past")
23206 .to_string();
23207 assert!(err.contains("NO fused gate"), "{err}");
23208 assert!(err.contains(&format!("{fused}")), "{err}");
23209
23210 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
23213 let wide = e.htod(&host).unwrap();
23214 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
23215 .expect("full-width wq splits");
23216 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
23217 for tok in 0..t {
23218 for hh in 0..n_head {
23219 for d in 0..head_dim {
23220 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
23221 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
23222 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
23223 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
23224 }
23225 }
23226 }
23227
23228 let mut small = e.uninit(out_n - 1).unwrap();
23230 assert!(
23231 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
23232 .is_err()
23233 );
23234 }
23235}
23236
23237#[cfg(test)]
23241mod fused_rope_width_tests {
23242 use super::Engine;
23243
23244 #[test]
23247 fn full_width_is_accepted() {
23248 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
23249 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
23250 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
23251 }
23252
23253 #[test]
23266 fn gemma4_official_artifact_widths_pass() {
23267 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
23268 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
23269 }
23270
23271 #[test]
23274 fn partial_rotary_is_refused_with_the_geometry_named() {
23275 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
23277 .expect_err("partial rotary must refuse");
23278 let msg = err.to_string();
23279 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
23280 assert!(msg.contains("n_rot 64"), "{msg}");
23281 assert!(msg.contains("head_dim 256"), "{msg}");
23282 assert!(
23283 msg.contains("64..256"),
23284 "names the band it would corrupt: {msg}"
23285 );
23286 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
23288 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
23290 }
23291}