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 pub fn wpf_level() -> u32 {
1784 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1785 *ON.get_or_init(|| {
1786 std::env::var("MEMRA_WPF")
1787 .ok()
1788 .and_then(|v| v.parse().ok())
1789 .unwrap_or(1)
1790 })
1791 }
1792
1793 pub fn set_verify_exact(&self, on: bool) {
1805 self.verify_exact
1806 .store(on, std::sync::atomic::Ordering::Relaxed);
1807 }
1808 pub(crate) fn verify_exact_on(&self) -> bool {
1809 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1810 }
1811
1812 pub fn qkv_append_on() -> bool {
1815 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1816 *ON.get_or_init(|| {
1817 std::env::var("MEMRA_QKV_APPEND")
1818 .map(|v| v != "0")
1819 .unwrap_or(true)
1820 })
1821 }
1822
1823 pub fn pdl_wb_on() -> bool {
1826 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1827 *ON.get_or_init(|| {
1828 std::env::var("MEMRA_PDL_WB")
1829 .map(|v| v != "0")
1830 .unwrap_or(true)
1831 })
1832 }
1833
1834 pub fn norm_ilp_on() -> bool {
1842 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1843 *ON.get_or_init(|| {
1844 std::env::var("MEMRA_NORM_ILP")
1845 .map(|v| v != "0")
1846 .unwrap_or(true)
1847 })
1848 }
1849
1850 pub fn tk_ffn_dual_on() -> bool {
1858 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1859 *ON.get_or_init(|| {
1860 std::env::var("MEMRA_TK_FFN_DUAL")
1861 .map(|v| v != "0")
1862 .unwrap_or(true)
1863 })
1864 }
1865
1866 pub fn pdl_mmvq_on() -> bool {
1870 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1871 *ON.get_or_init(|| {
1872 std::env::var("MEMRA_PDL_MMVQ")
1873 .map(|v| v != "0")
1874 .unwrap_or(true)
1875 })
1876 }
1877
1878 pub fn pdl_on() -> bool {
1879 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1880 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1881 }
1882
1883 pub fn pdl_nvfp4q8_on() -> bool {
1889 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1890 *ON.get_or_init(|| {
1891 std::env::var("MEMRA_PDL_NVFP4")
1892 .map(|v| v != "0")
1893 .unwrap_or(true)
1894 })
1895 }
1896
1897 fn q40_mr1_on() -> bool {
1903 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1904 match *Q40MR.get_or_init(|| {
1905 std::env::var("MEMRA_Q40_MR")
1906 .ok()
1907 .and_then(|v| v.parse().ok())
1908 }) {
1909 Some(v) => v == 1,
1910 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1911 }
1912 }
1913
1914 fn pdl_func_flash(
1919 &self,
1920 g: bool,
1921 name: &'static str,
1922 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1923 use cudarc::driver::sys as cu;
1924 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1931 std::sync::Mutex::new(None);
1932 static FNS: std::sync::Mutex<
1933 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1934 > = std::sync::Mutex::new(None);
1935 let ctx_key = self.ctx().cu_ctx() as usize;
1936 if let Some(&f) = FNS
1937 .lock()
1938 .unwrap()
1939 .get_or_insert_with(Default::default)
1940 .get(&(ctx_key, g, name))
1941 {
1942 return Ok(f as cu::CUfunction);
1943 }
1944 let module = {
1945 let mut mods = MODS.lock().unwrap();
1946 let map = mods.get_or_insert_with(Default::default);
1947 match map.get(&(ctx_key, g)) {
1948 Some(&m) => m,
1949 None => {
1950 let m = self.pdl_load_module_in_ctx(if g {
1951 FLASH_FATBIN_KF8VF8
1952 } else {
1953 FLASH_FATBIN
1954 })?;
1955 map.insert((ctx_key, g), m);
1956 m
1957 }
1958 }
1959 };
1960 let cname = std::ffi::CString::new(name)?;
1961 let mut f: cu::CUfunction = std::ptr::null_mut();
1962 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1963 if r != cu::CUresult::CUDA_SUCCESS {
1964 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
1965 }
1966 FNS.lock()
1967 .unwrap()
1968 .get_or_insert_with(Default::default)
1969 .insert((ctx_key, g, name), f as usize);
1970 Ok(f)
1971 }
1972
1973 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1978 use cudarc::driver::sys as cu;
1979 let mut prev: cu::CUcontext = std::ptr::null_mut();
1980 unsafe {
1981 cu::cuCtxGetCurrent(&mut prev).result()?;
1982 }
1983 self.ctx().bind_to_thread()?;
1984 let mut m: cu::CUmodule = std::ptr::null_mut();
1985 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1986 let restore = if prev.is_null() {
1987 cu::CUresult::CUDA_SUCCESS
1988 } else {
1989 unsafe { cu::cuCtxSetCurrent(prev) }
1990 };
1991 if r != cu::CUresult::CUDA_SUCCESS {
1992 return Err(format!("pdl module load: {r:?}").into());
1993 }
1994 if restore != cu::CUresult::CUDA_SUCCESS {
1995 return Err(format!("pdl module load: ctx restore {restore:?}").into());
1996 }
1997 Ok(m as usize)
1998 }
1999
2000 fn pdl_func(
2001 &self,
2002 name: &'static str,
2003 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2004 use cudarc::driver::sys as cu;
2005 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2008 std::sync::Mutex::new(None);
2009 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2012 std::sync::Mutex::new(None);
2013 static FNS: std::sync::Mutex<
2014 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2015 > = std::sync::Mutex::new(None);
2016 let ctx_key = self.ctx().cu_ctx() as usize;
2017 if let Some(&f) = FNS
2018 .lock()
2019 .unwrap()
2020 .get_or_insert_with(Default::default)
2021 .get(&(ctx_key, name))
2022 {
2023 return Ok(f as cu::CUfunction);
2024 }
2025 let module = {
2026 let mut mods = MODULES.lock().unwrap();
2027 let map = mods.get_or_insert_with(Default::default);
2028 match map.get(&ctx_key) {
2029 Some(&m) => m,
2030 None => {
2031 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2032 map.insert(ctx_key, m);
2033 m
2034 }
2035 }
2036 };
2037 let cname = std::ffi::CString::new(name)?;
2038 let mut f: cu::CUfunction = std::ptr::null_mut();
2039 let mut r =
2040 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2041 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2042 let qmodule = {
2043 let mut mods = QMODULES.lock().unwrap();
2044 let map = mods.get_or_insert_with(Default::default);
2045 match map.get(&ctx_key) {
2046 Some(&m) => m,
2047 None => {
2048 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2049 map.insert(ctx_key, m);
2050 m
2051 }
2052 }
2053 };
2054 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2055 }
2056 if r != cu::CUresult::CUDA_SUCCESS {
2057 return Err(format!("pdl_func {name}: {r:?}").into());
2058 }
2059 FNS.lock()
2060 .unwrap()
2061 .get_or_insert_with(Default::default)
2062 .insert((ctx_key, name), f as usize);
2063 Ok(f)
2064 }
2065
2066 unsafe fn launch_pdl_flash(
2078 &self,
2079 g: bool,
2080 name: &'static str,
2081 grid: (u32, u32, u32),
2082 block: (u32, u32, u32),
2083 smem: u32,
2084 params: &mut [*mut std::ffi::c_void],
2085 ) -> Result<(), Box<dyn std::error::Error>> {
2086 use cudarc::driver::sys as cu;
2087 let f = self.pdl_func_flash(g, name)?;
2088 if smem > 0 {
2089 let r =
2091 unsafe {
2092 cu::cuFuncSetAttribute(f,
2093 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2094 smem as i32)
2095 };
2096 if r != cu::CUresult::CUDA_SUCCESS {
2097 return Err(format!("pdl smem attr {name}: {r:?}").into());
2098 }
2099 }
2100 let mut attr = cu::CUlaunchAttribute {
2101 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2102 pad: [0; 4],
2103 value: cu::CUlaunchAttributeValue {
2104 programmaticStreamSerializationAllowed: 1,
2105 },
2106 };
2107 let cfg = cu::CUlaunchConfig {
2108 gridDimX: grid.0,
2109 gridDimY: grid.1,
2110 gridDimZ: grid.2,
2111 blockDimX: block.0,
2112 blockDimY: block.1,
2113 blockDimZ: block.2,
2114 sharedMemBytes: smem,
2115 hStream: self.gpu.stream().cu_stream(),
2116 attrs: &mut attr,
2117 numAttrs: 1,
2118 };
2119 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2120 if r != cu::CUresult::CUDA_SUCCESS {
2121 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2122 }
2123 Ok(())
2124 }
2125
2126 unsafe fn launch_pdl(
2127 &self,
2128 name: &'static str,
2129 grid: (u32, u32, u32),
2130 block: (u32, u32, u32),
2131 params: &mut [*mut std::ffi::c_void],
2132 ) -> Result<(), Box<dyn std::error::Error>> {
2133 use cudarc::driver::sys as cu;
2134 let f = self.pdl_func(name)?;
2135 let mut attr = cu::CUlaunchAttribute {
2136 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2137 pad: [0; 4],
2138 value: cu::CUlaunchAttributeValue {
2139 programmaticStreamSerializationAllowed: 1,
2140 },
2141 };
2142 let cfg = cu::CUlaunchConfig {
2143 gridDimX: grid.0,
2144 gridDimY: grid.1,
2145 gridDimZ: grid.2,
2146 blockDimX: block.0,
2147 blockDimY: block.1,
2148 blockDimZ: block.2,
2149 sharedMemBytes: 0,
2150 hStream: self.gpu.stream().cu_stream(),
2151 attrs: &mut attr,
2152 numAttrs: 1,
2153 };
2154 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2155 if r != cu::CUresult::CUDA_SUCCESS {
2156 return Err(format!("launch_pdl {name}: {r:?}").into());
2157 }
2158 Ok(())
2159 }
2160
2161 pub fn prefetch_weight_l2(
2164 &self,
2165 w: &crate::model::GpuTensor,
2166 ) -> Result<(), Box<dyn std::error::Error>> {
2167 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2168 let p = rp4.as_ref().unwrap_or(bytes);
2169 self.prefetch_l2(p, p.len())?;
2170 }
2171 Ok(())
2172 }
2173
2174 pub fn gather_row_bf16(
2177 &self,
2178 table: &CudaSlice<u8>,
2179 tok: &CudaSlice<u32>,
2180 idx: usize,
2181 dst: &mut CudaSlice<f32>,
2182 ncols: usize,
2183 ) -> Result<(), Box<dyn std::error::Error>> {
2184 let f = self.func("gather_row_bf16_f32");
2185 let cfg = LaunchConfig {
2186 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2187 block_dim: (256, 1, 1),
2188 shared_mem_bytes: 0,
2189 };
2190 let (nc, ix) = (ncols as i32, idx as i32);
2191 let __s_b = self.gpu.stream();
2192 let mut b = __s_b.launch_builder(&f);
2193 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2194 unsafe {
2195 b.launch(cfg)?;
2196 }
2197 Ok(())
2198 }
2199
2200 #[allow(clippy::too_many_arguments)]
2206 pub fn dflash2_dynconv(
2207 &self,
2208 x: &CudaSlice<f32>,
2209 dyn_: &CudaSlice<f32>,
2210 base: &CudaSlice<f32>,
2211 out: &mut CudaSlice<f32>,
2212 rows: usize,
2213 hidden: usize,
2214 group_size: usize,
2215 ksize: usize,
2216 half: usize,
2217 ) -> Result<(), Box<dyn std::error::Error>> {
2218 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2219 let f = self.func("dflash2_dynconv_f32");
2220 let n = rows * hidden;
2221 let cfg = LaunchConfig {
2222 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2223 block_dim: (256, 1, 1),
2224 shared_mem_bytes: 0,
2225 };
2226 let (ri, hi, gi, ki, hf) = (
2227 rows as i32,
2228 hidden as i32,
2229 group_size as i32,
2230 ksize as i32,
2231 half as i32,
2232 );
2233 let __s_b = self.gpu.stream();
2234 let mut b = __s_b.launch_builder(&f);
2235 b.arg(x)
2236 .arg(dyn_)
2237 .arg(base)
2238 .arg(out)
2239 .arg(&ri)
2240 .arg(&hi)
2241 .arg(&gi)
2242 .arg(&ki)
2243 .arg(&hf);
2244 unsafe {
2245 b.launch(cfg)?;
2246 }
2247 Ok(())
2248 }
2249
2250 pub fn topk_rows(
2254 &self,
2255 logits: &CudaSlice<f32>,
2256 n_rows: usize,
2257 n_cols: usize,
2258 k: usize,
2259 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2260 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2261 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2262 let f = self.func("topk_rows_f32");
2263 let nth = 256usize;
2264 let mut vals = self.uninit(n_rows * k)?;
2265 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2266 let cfg = LaunchConfig {
2267 grid_dim: (n_rows as u32, 1, 1),
2268 block_dim: (nth as u32, 1, 1),
2269 shared_mem_bytes: (nth * k * 8) as u32,
2270 };
2271 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2272 let __s_b = self.gpu.stream();
2273 let mut b = __s_b.launch_builder(&f);
2274 b.arg(logits)
2275 .arg(&nr)
2276 .arg(&nc)
2277 .arg(&ki)
2278 .arg(&mut vals)
2279 .arg(&mut idxs);
2280 unsafe {
2281 b.launch(cfg)?;
2282 }
2283 Ok((vals, idxs))
2284 }
2285
2286 pub fn add_row_inplace(
2288 &self,
2289 logits: &mut CudaSlice<f32>,
2290 bias: &CudaSlice<f32>,
2291 n: usize,
2292 row_off: usize,
2293 ) -> Result<(), Box<dyn std::error::Error>> {
2294 let f = self.func("add_row_inplace_f32");
2295 let cfg = LaunchConfig {
2296 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2297 block_dim: (256, 1, 1),
2298 shared_mem_bytes: 0,
2299 };
2300 let (ni, off) = (n as i32, row_off as i64);
2301 let __s_b = self.gpu.stream();
2302 let mut b = __s_b.launch_builder(&f);
2303 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2304 unsafe {
2305 b.launch(cfg)?;
2306 }
2307 Ok(())
2308 }
2309
2310 pub fn prefetch_l2(
2312 &self,
2313 p: &CudaSlice<u8>,
2314 n: usize,
2315 ) -> Result<(), Box<dyn std::error::Error>> {
2316 let f = self.func("prefetch_l2_bytes");
2317 let lines = n.div_ceil(128);
2318 let ni = n as i64;
2319 let cfg = LaunchConfig {
2320 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2321 block_dim: (256, 1, 1),
2322 shared_mem_bytes: 0,
2323 };
2324 let __s_b = self.gpu.stream();
2325 let mut b = __s_b.launch_builder(&f);
2326 b.arg(p).arg(&ni);
2327 unsafe {
2328 b.launch(cfg)?;
2329 }
2330 Ok(())
2331 }
2332
2333 pub fn router_gemv(
2336 &self,
2337 w: &CudaSlice<f32>,
2338 x: &CudaSlice<f32>,
2339 n_embd: usize,
2340 n_experts: usize,
2341 t: usize,
2342 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2343 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2349 Ok("0") => false,
2350 Ok(_) => true,
2351 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2352 };
2353 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2363 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2364 }
2365
2366 pub fn router_gemv_form(
2369 &self,
2370 w: &CudaSlice<f32>,
2371 x: &CudaSlice<f32>,
2372 n_embd: usize,
2373 n_experts: usize,
2374 t: usize,
2375 w8: bool,
2376 batch: bool,
2377 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2378 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2379 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2380 let f = if batch {
2381 self.func("router_gemv_f32_w8_batch")
2382 } else if w8 {
2383 self.func("router_gemv_f32_w8")
2384 } else {
2385 self.func("router_gemv_f32")
2386 };
2387 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2388 let cfg = if batch {
2389 LaunchConfig {
2390 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2391 block_dim: (32, 8, 1),
2392 shared_mem_bytes: 0,
2393 }
2394 } else {
2395 LaunchConfig {
2396 grid_dim: (n_experts as u32, t as u32, 1),
2397 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2398 shared_mem_bytes: 0,
2399 }
2400 };
2401 let __s_b = self.gpu.stream();
2402 let mut b = __s_b.launch_builder(&f);
2403 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2404 unsafe {
2405 b.launch(cfg)?;
2406 }
2407 Ok(y)
2408 }
2409
2410 pub fn rows_permute(
2412 &self,
2413 src: &CudaSlice<f32>,
2414 idx: &CudaSlice<i32>,
2415 nrows: usize,
2416 ncols: usize,
2417 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2418 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2419 let f = self.func("rows_permute_f32");
2420 let (nc, nr) = (ncols as i32, nrows as i32);
2421 let cfg = LaunchConfig {
2422 grid_dim: (nrows as u32, 1, 1),
2423 block_dim: (256, 1, 1),
2424 shared_mem_bytes: 0,
2425 };
2426 let __s_b = self.gpu.stream();
2427 let mut b = __s_b.launch_builder(&f);
2428 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2429 unsafe {
2430 b.launch(cfg)?;
2431 }
2432 Ok(dst)
2433 }
2434
2435 pub fn sigmoid_dot_rows(
2440 &self,
2441 x: &CudaSlice<f32>,
2442 w: &CudaSlice<f32>,
2443 n_embd: usize,
2444 t: usize,
2445 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2446 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2449 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2450 let gs = self.linear(x, w, t, n_embd, 1)?;
2451 let mut g = self.uninit(t)?;
2452 self.sigmoid(&gs, &mut g, t)?;
2453 return Ok(g);
2454 }
2455 let mut g = self.alloc_uninit::<f32>(t)?;
2461 let f = self.func("sigmoid_dot_rows_f32");
2462 let (ne, ti) = (n_embd as i32, t as i32);
2463 let cfg = LaunchConfig {
2464 grid_dim: (t as u32, 1, 1),
2465 block_dim: (32, 8, 1),
2466 shared_mem_bytes: 0,
2467 };
2468 let __s_b = self.gpu.stream();
2469 let mut b = __s_b.launch_builder(&f);
2470 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2471 unsafe {
2472 b.launch(cfg)?;
2473 }
2474 Ok(g)
2475 }
2476
2477 pub fn spec_rollback_stream(
2479 &self,
2480 len_ptrs: &CudaSlice<u64>,
2481 pos_start: &CudaSlice<i32>,
2482 acc: &CudaSlice<u32>,
2483 base: usize,
2484 n_rows: usize,
2485 ) -> Result<(), Box<dyn std::error::Error>> {
2486 let f = self.func("spec_rollback_stream");
2487 let (b, nr) = (base as i32, n_rows as i32);
2488 let cfg = LaunchConfig {
2489 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2490 block_dim: (64, 1, 1),
2491 shared_mem_bytes: 0,
2492 };
2493 let __s_bl = self.gpu.stream();
2494 let mut bl = __s_bl.launch_builder(&f);
2495 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2496 unsafe {
2497 bl.launch(cfg)?;
2498 }
2499 Ok(())
2500 }
2501
2502 pub fn plain_tok_ring(
2504 &self,
2505 vam: &CudaSlice<u32>,
2506 pos_start: &CudaSlice<i32>,
2507 base: usize,
2508 ring: &mut CudaSlice<u32>,
2509 ) -> Result<(), Box<dyn std::error::Error>> {
2510 let f = self.func("plain_tok_ring");
2511 let (b, cap) = (base as i32, ring.len() as i32);
2512 let cfg = LaunchConfig {
2513 grid_dim: (1, 1, 1),
2514 block_dim: (32, 1, 1),
2515 shared_mem_bytes: 0,
2516 };
2517 let __s_bl = self.gpu.stream();
2518 let mut bl = __s_bl.launch_builder(&f);
2519 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2520 unsafe {
2521 bl.launch(cfg)?;
2522 }
2523 Ok(())
2524 }
2525
2526 pub fn spec_ring_commit(
2528 &self,
2529 vtok: &CudaSlice<u32>,
2530 acc: &CudaSlice<u32>,
2531 brk: &CudaSlice<u32>,
2532 ring: &mut CudaSlice<u32>,
2533 pend: &mut CudaSlice<u32>,
2534 ) -> Result<(), Box<dyn std::error::Error>> {
2535 let f = self.func("spec_ring_commit");
2536 let cfg = LaunchConfig {
2537 grid_dim: (1, 1, 1),
2538 block_dim: (32, 1, 1),
2539 shared_mem_bytes: 0,
2540 };
2541 let __s_b = self.gpu.stream();
2542 let mut b = __s_b.launch_builder(&f);
2543 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2544 unsafe {
2545 b.launch(cfg)?;
2546 }
2547 Ok(())
2548 }
2549 pub fn i32_copy_add(
2550 &self,
2551 src: &CudaSlice<i32>,
2552 dst: &mut CudaSlice<i32>,
2553 delta: i32,
2554 ) -> Result<(), Box<dyn std::error::Error>> {
2555 let f = self.func("i32_copy_add");
2556 let cfg = LaunchConfig {
2557 grid_dim: (1, 1, 1),
2558 block_dim: (32, 1, 1),
2559 shared_mem_bytes: 0,
2560 };
2561 let __s_b = self.gpu.stream();
2562 let mut b = __s_b.launch_builder(&f);
2563 b.arg(src).arg(dst).arg(&delta);
2564 unsafe {
2565 b.launch(cfg)?;
2566 }
2567 Ok(())
2568 }
2569 pub fn u32_copy(
2570 &self,
2571 src: &CudaSlice<u32>,
2572 dst: &mut CudaSlice<u32>,
2573 ) -> Result<(), Box<dyn std::error::Error>> {
2574 let f = self.func("u32_copy");
2575 let cfg = LaunchConfig {
2576 grid_dim: (1, 1, 1),
2577 block_dim: (32, 1, 1),
2578 shared_mem_bytes: 0,
2579 };
2580 let __s_b = self.gpu.stream();
2581 let mut b = __s_b.launch_builder(&f);
2582 b.arg(src).arg(dst);
2583 unsafe {
2584 b.launch(cfg)?;
2585 }
2586 Ok(())
2587 }
2588
2589 pub fn spec_adapt_k(
2593 &self,
2594 acc: &CudaSlice<u32>,
2595 brk: &mut CudaSlice<u32>,
2596 floor: usize,
2597 cap: usize,
2598 ) -> Result<(), Box<dyn std::error::Error>> {
2599 let f = self.func("spec_adapt_k");
2600 let (fl, cp) = (floor as i32, cap as i32);
2601 let cfg = LaunchConfig {
2602 grid_dim: (1, 1, 1),
2603 block_dim: (32, 1, 1),
2604 shared_mem_bytes: 0,
2605 };
2606 let __s_b = self.gpu.stream();
2607 let mut b = __s_b.launch_builder(&f);
2608 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2609 unsafe {
2610 b.launch(cfg)?;
2611 }
2612 Ok(())
2613 }
2614
2615 pub fn spec_accept_greedy_dc(
2617 &self,
2618 preds: &CudaSlice<u32>,
2619 vtok: &CudaSlice<u32>,
2620 last_pred: &CudaSlice<u32>,
2621 brk: &CudaSlice<u32>,
2622 out: &mut CudaSlice<u32>,
2623 ) -> Result<(), Box<dyn std::error::Error>> {
2624 let f = self.func("spec_accept_greedy_dc");
2625 let cfg = LaunchConfig {
2626 grid_dim: (1, 1, 1),
2627 block_dim: (32, 1, 1),
2628 shared_mem_bytes: 0,
2629 };
2630 let __s_b = self.gpu.stream();
2631 let mut b = __s_b.launch_builder(&f);
2632 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2633 unsafe {
2634 b.launch(cfg)?;
2635 }
2636 Ok(())
2637 }
2638
2639 pub fn pos_iota(
2641 &self,
2642 pos0: &CudaSlice<i32>,
2643 out: &mut CudaSlice<i32>,
2644 t: usize,
2645 ) -> Result<(), Box<dyn std::error::Error>> {
2646 let f = self.func("pos_iota_i32");
2647 let ti = t as i32;
2648 let cfg = LaunchConfig {
2649 grid_dim: (1, 1, 1),
2650 block_dim: (t.max(1) as u32, 1, 1),
2651 shared_mem_bytes: 0,
2652 };
2653 let __s_b = self.gpu.stream();
2654 let mut b = __s_b.launch_builder(&f);
2655 b.arg(pos0).arg(out).arg(&ti);
2656 unsafe {
2657 b.launch(cfg)?;
2658 }
2659 Ok(())
2660 }
2661 #[allow(clippy::too_many_arguments)]
2662 pub fn append_kv_quantized_rows_dc(
2663 &self,
2664 k_rows: &CudaSlice<f32>,
2665 v_rows: &CudaSlice<f32>,
2666 kc: &mut CudaSlice<u8>,
2667 vc: &mut CudaSlice<u8>,
2668 t0_dev: &CudaSlice<i32>,
2669 t: usize,
2670 kv_dim_k: usize,
2671 kv_dim_v: usize,
2672 k_tok_bytes: usize,
2673 v_tok_bytes: usize,
2674 g: bool,
2675 ) -> Result<(), Box<dyn std::error::Error>> {
2676 let f = if g {
2677 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2678 } else {
2679 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2680 };
2681 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2682 let cfg = LaunchConfig {
2683 grid_dim: (nblk, t as u32, 1),
2684 block_dim: (32, 1, 1),
2685 shared_mem_bytes: 0,
2686 };
2687 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2688 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2689 let __s_b = self.gpu.stream();
2690 let mut b = __s_b.launch_builder(&f);
2691 b.arg(k_rows)
2692 .arg(v_rows)
2693 .arg(kc)
2694 .arg(vc)
2695 .arg(t0_dev)
2696 .arg(&kdk)
2697 .arg(&kdv)
2698 .arg(&ktb)
2699 .arg(&vtb);
2700 unsafe {
2701 b.launch(cfg)?;
2702 }
2703 Ok(())
2704 }
2705
2706 #[allow(clippy::too_many_arguments)]
2709 pub fn append_kv_quantized_row_dc_inc(
2710 &self,
2711 k_row: &CudaSlice<f32>,
2712 v_row: &CudaSlice<f32>,
2713 kc: &mut CudaSlice<u8>,
2714 vc: &mut CudaSlice<u8>,
2715 t0_dev: &mut CudaSlice<i32>,
2716 kv_dim_k: usize,
2717 kv_dim_v: usize,
2718 k_tok_bytes: usize,
2719 v_tok_bytes: usize,
2720 g: bool,
2721 ) -> Result<(), Box<dyn std::error::Error>> {
2722 let f = if g {
2723 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2724 } else {
2725 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2726 };
2727 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2728 let cfg = LaunchConfig {
2729 grid_dim: (1, 1, 1),
2730 block_dim: (nthreads, 1, 1),
2731 shared_mem_bytes: 0,
2732 };
2733 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2734 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2735 let __s_b = self.gpu.stream();
2736 let mut b = __s_b.launch_builder(&f);
2737 b.arg(k_row)
2738 .arg(v_row)
2739 .arg(kc)
2740 .arg(vc)
2741 .arg(t0_dev)
2742 .arg(&kdk)
2743 .arg(&kdv)
2744 .arg(&ktb)
2745 .arg(&vtb);
2746 unsafe {
2747 b.launch(cfg)?;
2748 }
2749 Ok(())
2750 }
2751
2752 pub fn pack_tok_p(
2754 &self,
2755 tok: &CudaSlice<u32>,
2756 p: &CudaSlice<f32>,
2757 out: &mut CudaSlice<u32>,
2758 slot: usize,
2759 ) -> Result<(), Box<dyn std::error::Error>> {
2760 let f = self.func("pack_tok_p");
2761 let sl = slot as i32;
2762 let cfg = LaunchConfig {
2763 grid_dim: (1, 1, 1),
2764 block_dim: (32, 1, 1),
2765 shared_mem_bytes: 0,
2766 };
2767 let __s_b = self.gpu.stream();
2768 let mut b = __s_b.launch_builder(&f);
2769 b.arg(tok).arg(p).arg(out).arg(&sl);
2770 unsafe {
2771 b.launch(cfg)?;
2772 }
2773 Ok(())
2774 }
2775 pub fn tok_map_u32(
2776 &self,
2777 tok: &mut CudaSlice<u32>,
2778 map: &CudaSlice<u32>,
2779 ) -> Result<(), Box<dyn std::error::Error>> {
2780 let f = self.func("tok_map_u32");
2781 let cfg = LaunchConfig {
2782 grid_dim: (1, 1, 1),
2783 block_dim: (32, 1, 1),
2784 shared_mem_bytes: 0,
2785 };
2786 let __s_b = self.gpu.stream();
2787 let mut b = __s_b.launch_builder(&f);
2788 b.arg(tok).arg(map);
2789 unsafe {
2790 b.launch(cfg)?;
2791 }
2792 Ok(())
2793 }
2794
2795 #[allow(clippy::too_many_arguments)]
2797 pub fn spec_assemble_verify(
2798 &self,
2799 tokp: &CudaSlice<u32>,
2800 pend: &CudaSlice<u32>,
2801 d2t: Option<&CudaSlice<u32>>,
2802 vtok: &mut CudaSlice<u32>,
2803 brk: &mut CudaSlice<u32>,
2804 p_min: f32,
2805 k: usize,
2806 pmin0: bool,
2807 ) -> Result<(), Box<dyn std::error::Error>> {
2808 let f = self.func("spec_assemble_verify");
2809 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2810 let cfg = LaunchConfig {
2811 grid_dim: (1, 1, 1),
2812 block_dim: (32, 1, 1),
2813 shared_mem_bytes: 0,
2814 };
2815 let __s_b = self.gpu.stream();
2816 let mut b = __s_b.launch_builder(&f);
2817 match d2t {
2818 Some(m) => {
2819 b.arg(tokp)
2820 .arg(pend)
2821 .arg(m)
2822 .arg(vtok)
2823 .arg(brk)
2824 .arg(&p_min)
2825 .arg(&ki)
2826 .arg(&pm);
2827 unsafe {
2828 b.launch(cfg)?;
2829 }
2830 }
2831 None => {
2832 let null: u64 = 0;
2833 b.arg(tokp)
2834 .arg(pend)
2835 .arg(&null)
2836 .arg(vtok)
2837 .arg(brk)
2838 .arg(&p_min)
2839 .arg(&ki)
2840 .arg(&pm);
2841 unsafe {
2842 b.launch(cfg)?;
2843 }
2844 }
2845 }
2846 Ok(())
2847 }
2848
2849 #[allow(clippy::too_many_arguments)]
2851 pub fn ssm_conv_ring_rebuild_dc(
2852 &self,
2853 qkv_tm: &CudaSlice<f32>,
2854 ring_old: &CudaSlice<f32>,
2855 conv_state: &mut CudaSlice<f32>,
2856 conv_dim: usize,
2857 acc: &CudaSlice<u32>,
2858 base: usize,
2859 t_v: usize,
2860 d_conv: usize,
2861 ) -> Result<(), Box<dyn std::error::Error>> {
2862 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2863 let n = conv_dim * (d_conv - 1);
2864 let cfg = LaunchConfig::for_num_elems(n as u32);
2865 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2866 let __s_b = self.gpu.stream();
2867 let mut b = __s_b.launch_builder(&f);
2868 b.arg(qkv_tm)
2869 .arg(ring_old)
2870 .arg(conv_state)
2871 .arg(&cd)
2872 .arg(acc)
2873 .arg(&b0)
2874 .arg(&tv)
2875 .arg(&dc);
2876 unsafe {
2877 b.launch(cfg)?;
2878 }
2879 Ok(())
2880 }
2881 #[allow(clippy::too_many_arguments)]
2882 pub fn gdn_scan_s128_dc(
2883 &self,
2884 q: &CudaSlice<f32>,
2885 k: &CudaSlice<f32>,
2886 v: &CudaSlice<f32>,
2887 g: &CudaSlice<f32>,
2888 beta: &CudaSlice<f32>,
2889 state_in: &CudaSlice<f32>,
2890 state_out: &mut CudaSlice<f32>,
2891 o: &mut CudaSlice<f32>,
2892 n_head: usize,
2893 acc: &CudaSlice<u32>,
2894 base: usize,
2895 t_v: usize,
2896 scale: f32,
2897 ) -> Result<(), Box<dyn std::error::Error>> {
2898 let f = self.func("gdn_scan_s128_dc");
2899 const S_V: u32 = 128;
2900 const WARP: u32 = 32;
2901 const COLS_PER_BLOCK: u32 = 4;
2902 let cfg = LaunchConfig {
2903 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2904 block_dim: (WARP, COLS_PER_BLOCK, 1),
2905 shared_mem_bytes: 0,
2906 };
2907 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2908 let __s_b = self.gpu.stream();
2909 let mut b = __s_b.launch_builder(&f);
2910 b.arg(q)
2911 .arg(k)
2912 .arg(v)
2913 .arg(g)
2914 .arg(beta)
2915 .arg(state_in)
2916 .arg(state_out)
2917 .arg(o)
2918 .arg(&h)
2919 .arg(acc)
2920 .arg(&b0)
2921 .arg(&tv)
2922 .arg(&scale);
2923 unsafe {
2924 b.launch(cfg)?;
2925 }
2926 Ok(())
2927 }
2928
2929 pub fn spec_rollback_kv(
2931 &self,
2932 len_ptrs: &CudaSlice<u64>,
2933 saved: &CudaSlice<i32>,
2934 acc: &CudaSlice<u32>,
2935 base: usize,
2936 n_layer: usize,
2937 ) -> Result<(), Box<dyn std::error::Error>> {
2938 let f = self.func("spec_rollback_kv");
2939 let (b, nl) = (base as i32, n_layer as i32);
2940 let cfg = LaunchConfig {
2941 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2942 block_dim: (64, 1, 1),
2943 shared_mem_bytes: 0,
2944 };
2945 let __s_bl = self.gpu.stream();
2946 let mut bl = __s_bl.launch_builder(&f);
2947 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
2948 unsafe {
2949 bl.launch(cfg)?;
2950 }
2951 Ok(())
2952 }
2953
2954 pub fn spec_fork_valid(
2956 &self,
2957 acc: &CudaSlice<u32>,
2958 optimistic_pending: u32,
2959 valid: &mut CudaSlice<u32>,
2960 ) -> Result<(), Box<dyn std::error::Error>> {
2961 let f = self.func("spec_fork_valid");
2962 let cfg = LaunchConfig {
2963 grid_dim: (1, 1, 1),
2964 block_dim: (1, 1, 1),
2965 shared_mem_bytes: 0,
2966 };
2967 let __s_bl = self.gpu.stream();
2968 let mut bl = __s_bl.launch_builder(&f);
2969 bl.arg(acc).arg(&optimistic_pending).arg(valid);
2970 unsafe {
2971 bl.launch(cfg)?;
2972 }
2973 Ok(())
2974 }
2975
2976 pub fn spec_fork_reconcile_kv(
2978 &self,
2979 len_ptrs: &CudaSlice<u64>,
2980 saved: &CudaSlice<i32>,
2981 acc: &CudaSlice<u32>,
2982 valid: &CudaSlice<u32>,
2983 base: usize,
2984 n_layer: usize,
2985 ) -> Result<(), Box<dyn std::error::Error>> {
2986 let f = self.func("spec_fork_reconcile_kv");
2987 let (b, nl) = (base as i32, n_layer as i32);
2988 let cfg = LaunchConfig {
2989 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2990 block_dim: (64, 1, 1),
2991 shared_mem_bytes: 0,
2992 };
2993 let __s_bl = self.gpu.stream();
2994 let mut bl = __s_bl.launch_builder(&f);
2995 bl.arg(len_ptrs)
2996 .arg(saved)
2997 .arg(acc)
2998 .arg(valid)
2999 .arg(&b)
3000 .arg(&nl);
3001 unsafe {
3002 bl.launch(cfg)?;
3003 }
3004 Ok(())
3005 }
3006
3007 pub fn spec_fork_restore_f32(
3009 &self,
3010 snapshot: &CudaSlice<f32>,
3011 state: &mut CudaSlice<f32>,
3012 valid: &CudaSlice<u32>,
3013 ) -> Result<(), Box<dyn std::error::Error>> {
3014 assert_eq!(
3015 snapshot.len(),
3016 state.len(),
3017 "fork recurrent snapshot shape mismatch"
3018 );
3019 let f = self.func("spec_fork_restore_f32");
3020 let n = state.len() as i32;
3021 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3022 let cfg = LaunchConfig {
3023 grid_dim: (blocks, 1, 1),
3024 block_dim: (256, 1, 1),
3025 shared_mem_bytes: 0,
3026 };
3027 let __s_bl = self.gpu.stream();
3028 let mut bl = __s_bl.launch_builder(&f);
3029 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3030 unsafe {
3031 bl.launch(cfg)?;
3032 }
3033 Ok(())
3034 }
3035
3036 pub fn spec_seed_gather(
3039 &self,
3040 vx: &CudaSlice<f32>,
3041 fill_prev: &CudaSlice<f32>,
3042 acc: &CudaSlice<u32>,
3043 h_seed: &mut CudaSlice<f32>,
3044 base: usize,
3045 n_embd: usize,
3046 ) -> Result<(), Box<dyn std::error::Error>> {
3047 let f = self.func("spec_seed_gather");
3048 let (b, ne) = (base as i32, n_embd as i32);
3049 let cfg = LaunchConfig {
3050 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3051 block_dim: (256, 1, 1),
3052 shared_mem_bytes: 0,
3053 };
3054 let __s_bl = self.gpu.stream();
3055 let mut bl = __s_bl.launch_builder(&f);
3056 bl.arg(vx)
3057 .arg(fill_prev)
3058 .arg(acc)
3059 .arg(h_seed)
3060 .arg(&b)
3061 .arg(&ne);
3062 unsafe {
3063 bl.launch(cfg)?;
3064 }
3065 Ok(())
3066 }
3067
3068 pub fn spec_accept_greedy(
3070 &self,
3071 preds: &CudaSlice<u32>,
3072 draft: &CudaSlice<u32>,
3073 last_pred: u32,
3074 base: usize,
3075 k_round: usize,
3076 out: &mut CudaSlice<u32>,
3077 ) -> Result<(), Box<dyn std::error::Error>> {
3078 let f = self.func("spec_accept_greedy");
3079 let (b, k) = (base as i32, k_round as i32);
3080 let cfg = LaunchConfig {
3081 grid_dim: (1, 1, 1),
3082 block_dim: (32, 1, 1),
3083 shared_mem_bytes: 0,
3084 };
3085 let __s_bl = self.gpu.stream();
3086 let mut bl = __s_bl.launch_builder(&f);
3087 bl.arg(preds)
3088 .arg(draft)
3089 .arg(&last_pred)
3090 .arg(&b)
3091 .arg(&k)
3092 .arg(out);
3093 unsafe {
3094 bl.launch(cfg)?;
3095 }
3096 Ok(())
3097 }
3098
3099 pub fn gumbel_perturb(
3106 &self,
3107 x: &CudaSlice<f32>,
3108 y: &mut CudaSlice<f32>,
3109 n: usize,
3110 seed: u64,
3111 stream_pos: u32,
3112 temp: f32,
3113 ) -> Result<(), Box<dyn std::error::Error>> {
3114 let f = self.func("gumbel_perturb_f32");
3115 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3116 let cfg = LaunchConfig {
3117 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3118 block_dim: (256, 1, 1),
3119 shared_mem_bytes: 0,
3120 };
3121 let __s_b = self.gpu.stream();
3122 let mut b = __s_b.launch_builder(&f);
3123 b.arg(x)
3124 .arg(&mut *y)
3125 .arg(&ni)
3126 .arg(&slo)
3127 .arg(&shi)
3128 .arg(&stream_pos)
3129 .arg(&temp);
3130 unsafe {
3131 b.launch(cfg)?;
3132 }
3133 Ok(())
3134 }
3135
3136 pub fn mask_logits_col(
3144 &self,
3145 logits: &mut CudaSlice<f32>,
3146 mask: &CudaSlice<u32>,
3147 col: usize,
3148 n: usize,
3149 mask_words: usize,
3150 ) -> Result<(), Box<dyn std::error::Error>> {
3151 let f = self.func("mask_logits_f32");
3152 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3153 let cfg = LaunchConfig {
3154 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3155 block_dim: (256, 1, 1),
3156 shared_mem_bytes: 0,
3157 };
3158 let __s_b = self.gpu.stream();
3159 let mut b = __s_b.launch_builder(&f);
3160 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3161 unsafe {
3162 b.launch(cfg)?;
3163 }
3164 Ok(())
3165 }
3166
3167 pub fn gumbel_perturb_col(
3174 &self,
3175 x: &CudaSlice<f32>,
3176 col: usize,
3177 y: &mut CudaSlice<f32>,
3178 n: usize,
3179 seed: u64,
3180 stream_pos: u32,
3181 temp: f32,
3182 ) -> Result<(), Box<dyn std::error::Error>> {
3183 let f = self.func("gumbel_perturb_f32");
3184 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3185 let col_view = x.slice(col * n..(col + 1) * n);
3186 let cfg = LaunchConfig {
3187 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3188 block_dim: (256, 1, 1),
3189 shared_mem_bytes: 0,
3190 };
3191 let __s_b = self.gpu.stream();
3192 let mut b = __s_b.launch_builder(&f);
3193 b.arg(&col_view)
3194 .arg(&mut *y)
3195 .arg(&ni)
3196 .arg(&slo)
3197 .arg(&shi)
3198 .arg(&stream_pos)
3199 .arg(&temp);
3200 unsafe {
3201 b.launch(cfg)?;
3202 }
3203 Ok(())
3204 }
3205
3206 #[allow(clippy::too_many_arguments)]
3212 pub fn gumbel_perturb_filtered_col(
3213 &self,
3214 x: &CudaSlice<f32>,
3215 col: usize,
3216 y: &mut CudaSlice<f32>,
3217 n: usize,
3218 seed: u64,
3219 stream_pos: u32,
3220 temp: f32,
3221 stat_max: &CudaSlice<f32>,
3222 stat_th: &CudaSlice<f32>,
3223 stat_idx: usize,
3224 ) -> Result<(), Box<dyn std::error::Error>> {
3225 let f = self.func("gumbel_perturb_filtered_col_f32");
3226 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3227 let (ci, si) = (col as i32, stat_idx as i32);
3228 let cfg = LaunchConfig {
3229 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3230 block_dim: (256, 1, 1),
3231 shared_mem_bytes: 0,
3232 };
3233 let __s_b = self.gpu.stream();
3234 let mut b = __s_b.launch_builder(&f);
3235 b.arg(x)
3236 .arg(&ci)
3237 .arg(&mut *y)
3238 .arg(&ni)
3239 .arg(&slo)
3240 .arg(&shi)
3241 .arg(&stream_pos)
3242 .arg(&temp)
3243 .arg(stat_max)
3244 .arg(stat_th)
3245 .arg(&si);
3246 unsafe {
3247 b.launch(cfg)?;
3248 }
3249 Ok(())
3250 }
3251
3252 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3257 let f = self.func("memra_sctr_inc");
3258 let cfg = LaunchConfig {
3259 grid_dim: (1, 1, 1),
3260 block_dim: (1, 1, 1),
3261 shared_mem_bytes: 0,
3262 };
3263 let __s_b = self.gpu.stream();
3264 let mut b = __s_b.launch_builder(&f);
3265 b.arg(&mut *ctr);
3266 unsafe {
3267 b.launch(cfg)?;
3268 }
3269 Ok(())
3270 }
3271
3272 pub fn gumbel_perturb_ctr(
3277 &self,
3278 x: &CudaSlice<f32>,
3279 y: &mut CudaSlice<f32>,
3280 n: usize,
3281 seed: u64,
3282 ctr: &CudaSlice<u32>,
3283 temp: f32,
3284 ) -> Result<(), Box<dyn std::error::Error>> {
3285 let f = self.func("gumbel_perturb_ctr_f32");
3286 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3287 let cfg = LaunchConfig {
3288 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3289 block_dim: (256, 1, 1),
3290 shared_mem_bytes: 0,
3291 };
3292 let __s_b = self.gpu.stream();
3293 let mut b = __s_b.launch_builder(&f);
3294 b.arg(x)
3295 .arg(&mut *y)
3296 .arg(&ni)
3297 .arg(&slo)
3298 .arg(&shi)
3299 .arg(ctr)
3300 .arg(&temp);
3301 unsafe {
3302 b.launch(cfg)?;
3303 }
3304 Ok(())
3305 }
3306
3307 pub fn softmax_gather(
3311 &self,
3312 x: &CudaSlice<f32>,
3313 row_stride: usize,
3314 ids: &CudaSlice<u32>,
3315 rows: &CudaSlice<i32>,
3316 out: &mut CudaSlice<f32>,
3317 n: usize,
3318 npair: usize,
3319 temp: f32,
3320 ) -> Result<(), Box<dyn std::error::Error>> {
3321 let f = self.func("softmax_gather_f32");
3322 let (ni, rs) = (n as i32, row_stride as i64);
3323 let np = npair as i32;
3324 let cfg = LaunchConfig {
3325 grid_dim: (npair as u32, 1, 1),
3326 block_dim: (256, 1, 1),
3327 shared_mem_bytes: 0,
3328 };
3329 let __s_b = self.gpu.stream();
3330 let mut b = __s_b.launch_builder(&f);
3331 b.arg(x)
3332 .arg(&rs)
3333 .arg(ids)
3334 .arg(rows)
3335 .arg(&mut *out)
3336 .arg(&ni)
3337 .arg(&np)
3338 .arg(&temp);
3339 unsafe {
3340 b.launch(cfg)?;
3341 }
3342 Ok(())
3343 }
3344
3345 pub fn residual_sample(
3349 &self,
3350 p: &CudaSlice<f32>,
3351 q: Option<&CudaSlice<f32>>,
3352 n: usize,
3353 temp: f32,
3354 seed: u64,
3355 stream_pos: u32,
3356 out_tok: &mut CudaSlice<u32>,
3357 ) -> Result<(), Box<dyn std::error::Error>> {
3358 let f = self.func("residual_sample_f32");
3359 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3360 let nth = 1024u32;
3361 let cfg = LaunchConfig {
3362 grid_dim: (1, 1, 1),
3363 block_dim: (nth, 1, 1),
3364 shared_mem_bytes: 0,
3365 };
3366 let has_q: i32 = q.is_some() as i32;
3367 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3369 let mut b = __s_b.launch_builder(&f);
3370 b.arg(p)
3371 .arg(qbuf)
3372 .arg(&has_q)
3373 .arg(&ni)
3374 .arg(&temp)
3375 .arg(&slo)
3376 .arg(&shi)
3377 .arg(&stream_pos)
3378 .arg(&mut *out_tok);
3379 unsafe {
3380 b.launch(cfg)?;
3381 }
3382 Ok(())
3383 }
3384
3385 pub fn with_moe_cache<R>(
3390 &self,
3391 max_block_bytes: usize,
3392 f: impl FnOnce(
3393 &mut crate::moe_cache::MoeSlotCache,
3394 &Engine,
3395 ) -> Result<R, Box<dyn std::error::Error>>,
3396 ) -> Result<R, Box<dyn std::error::Error>> {
3397 let mut guard = self.moe_cache.lock().unwrap();
3398 if guard.is_none() {
3399 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3400 }
3401 let cache = guard.as_mut().unwrap();
3402 f(cache, self)
3403 }
3404
3405 pub fn freeze_moe_cache(&self) {
3408 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3409 cache.freeze();
3410 }
3411 }
3412
3413 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3416 self.moe_cache
3417 .lock()
3418 .unwrap()
3419 .as_ref()
3420 .map(crate::moe_cache::MoeSlotCache::export_residency)
3421 }
3422
3423 pub(crate) fn moe_cache_frozen(&self) -> bool {
3424 self.moe_cache
3425 .lock()
3426 .unwrap()
3427 .as_ref()
3428 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3429 }
3430
3431 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3438 crate::cpu_experts::configured()
3439 && self.moe_cache_frozen()
3440 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3441 }
3442
3443 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3445 assert!(
3446 self.moe_cache.lock().unwrap().is_none(),
3447 "MoE cache layout configured after cache construction"
3448 );
3449 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3450 }
3451
3452 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3453 self.moe_cache_layout.lock().unwrap().clone()
3454 }
3455
3456 pub fn moe_cache_enabled() -> bool {
3458 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3459 }
3460
3461 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3464 let guard = self.moe_cache.lock().unwrap();
3465 guard
3466 .as_ref()
3467 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3468 }
3469
3470 pub fn cpu_expert_stats(
3474 &self,
3475 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3476 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3477 }
3478
3479 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3482 crate::cpu_experts::predictor_stats()
3483 }
3484
3485 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3486 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3487 }
3488
3489 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3492 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3493 }
3494
3495 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3498 let guard = self.moe_cache.lock().unwrap();
3499 guard
3500 .as_ref()
3501 .and_then(|cache| cache.pread_stats())
3502 .map(|stats| {
3503 (
3504 stats.reads,
3505 stats.bytes,
3506 stats.read_errors,
3507 stats.short_reads,
3508 stats.fallbacks,
3509 stats.buffer_waits,
3510 stats.ring_full,
3511 )
3512 })
3513 }
3514
3515 pub fn spill_config_fallbacks(&self) -> u64 {
3517 crate::spill_pread::config_fallbacks()
3518 }
3519
3520 pub fn moe_cache_reset_counters(&self) {
3522 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3523 c.reset_counters();
3524 }
3525 }
3526
3527 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3528 Ok(self.gpu.stream().clone_htod(v)?)
3529 }
3530
3531 pub fn htod_bytes_padded(
3535 &self,
3536 v: &[u8],
3537 pad: usize,
3538 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3539 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3540 {
3541 let mut view = d.slice_mut(0..v.len());
3542 self.gpu.stream().memcpy_htod(v, &mut view)?;
3543 }
3544 Ok(d)
3545 }
3546
3547 pub fn copy_into(
3549 &self,
3550 dst: &mut CudaSlice<f32>,
3551 off: usize,
3552 src: &CudaSlice<f32>,
3553 len: usize,
3554 ) -> Result<(), Box<dyn std::error::Error>> {
3555 let mut view = dst.slice_mut(off..off + len);
3556 self.gpu
3557 .stream()
3558 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3559 Ok(())
3560 }
3561
3562 pub fn copy_u8_into(
3565 &self,
3566 dst: &mut CudaSlice<u8>,
3567 off: usize,
3568 src: &CudaSlice<u8>,
3569 len: usize,
3570 ) -> Result<(), Box<dyn std::error::Error>> {
3571 let mut view = dst.slice_mut(off..off + len);
3572 self.gpu
3573 .stream()
3574 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3575 Ok(())
3576 }
3577
3578 pub fn copy_u8_range_into(
3580 &self,
3581 dst: &mut CudaSlice<u8>,
3582 dst_off: usize,
3583 src: &CudaSlice<u8>,
3584 src_off: usize,
3585 len: usize,
3586 ) -> Result<(), Box<dyn std::error::Error>> {
3587 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3588 self.gpu
3589 .stream()
3590 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3591 Ok(())
3592 }
3593
3594 pub fn prepare_kv_append(
3598 &self,
3599 kv: &mut crate::cache::KvLayer,
3600 retain_from: usize,
3601 append_rows: usize,
3602 ) -> Result<usize, Box<dyn std::error::Error>> {
3603 let Some(plan) = kv
3604 .ring
3605 .as_ref()
3606 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3607 .transpose()?
3608 else {
3609 return Ok(kv.len);
3610 };
3611 match plan {
3612 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3613 crate::cache::KvRingAppend::Rebase {
3614 src_row,
3615 keep_rows,
3616 new_base,
3617 write_row,
3618 } => {
3619 if keep_rows > 0 {
3620 let k_len = keep_rows * kv.k_tok_bytes;
3621 let v_len = keep_rows * kv.v_tok_bytes;
3622 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3623 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3624 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3625 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3626 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3627 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3628 }
3629 kv.ring.as_mut().unwrap().apply_rebase(new_base);
3630 Ok(write_row)
3631 }
3632 }
3633 }
3634
3635 pub fn htod_u8_into(
3638 &self,
3639 dst: &mut CudaSlice<u8>,
3640 off: usize,
3641 src: &[u8],
3642 ) -> Result<(), Box<dyn std::error::Error>> {
3643 let mut view = dst.slice_mut(off..off + src.len());
3644 self.gpu.stream().memcpy_htod(src, &mut view)?;
3645 Ok(())
3646 }
3647
3648 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3649 b.slice(0..len)
3650 }
3651
3652 pub fn view_u8_range<'a>(
3655 &self,
3656 b: &'a CudaSlice<u8>,
3657 start: usize,
3658 end: usize,
3659 ) -> cudarc::driver::CudaView<'a, u8> {
3660 b.slice(start..end)
3661 }
3662 pub fn view_u8<'a>(
3663 &self,
3664 b: &'a CudaSlice<u8>,
3665 len: usize,
3666 ) -> cudarc::driver::CudaView<'a, u8> {
3667 b.slice(0..len)
3668 }
3669
3670 pub fn append_kv_quantized(
3674 &self,
3675 k_row: &CudaSlice<f32>,
3676 v_row: &CudaSlice<f32>,
3677 kc: &mut CudaSlice<u8>,
3678 vc: &mut CudaSlice<u8>,
3679 t: usize,
3680 kv_dim_k: usize,
3681 kv_dim_v: usize,
3682 k_tok_bytes: usize,
3683 v_tok_bytes: usize,
3684 g: bool,
3685 ) -> Result<(), Box<dyn std::error::Error>> {
3686 let f = if g {
3687 self.func_g("append_quantize_kv_q8_0_q5_1")
3688 } else {
3689 self.func("append_quantize_kv_q8_0_q5_1")
3690 };
3691 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3692 let cfg = LaunchConfig {
3693 grid_dim: (nblk, 1, 1),
3694 block_dim: (32, 1, 1),
3695 shared_mem_bytes: 0,
3696 };
3697 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3698 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3699 let __s_b = self.gpu.stream();
3700 let mut b = __s_b.launch_builder(&f);
3701 b.arg(k_row)
3702 .arg(v_row)
3703 .arg(kc)
3704 .arg(vc)
3705 .arg(&ti)
3706 .arg(&kdk)
3707 .arg(&kdv)
3708 .arg(&ktb)
3709 .arg(&vtb);
3710 unsafe {
3711 b.launch(cfg)?;
3712 }
3713 Ok(())
3714 }
3715
3716 pub fn append_kv_quantized_dc(
3720 &self,
3721 k_row: &CudaSlice<f32>,
3722 v_row: &CudaSlice<f32>,
3723 kc: &mut CudaSlice<u8>,
3724 vc: &mut CudaSlice<u8>,
3725 t_dev: &CudaSlice<i32>,
3726 kv_dim_k: usize,
3727 kv_dim_v: usize,
3728 k_tok_bytes: usize,
3729 v_tok_bytes: usize,
3730 g: bool,
3731 ) -> Result<(), Box<dyn std::error::Error>> {
3732 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3733 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3734 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3735 if Self::pdl_on() && Self::pdl_wb_on() {
3737 use cudarc::driver::{DevicePtr, DevicePtrMut};
3738 let s = &self.gpu.stream();
3739 let (pk, _g0) = k_row.device_ptr(s);
3740 let (pv, _g1) = v_row.device_ptr(s);
3741 let (pkc, _g2) = kc.device_ptr_mut(s);
3742 let (pvc, _g3) = vc.device_ptr_mut(s);
3743 let (pt, _g4) = t_dev.device_ptr(s);
3744 let mut ps = [
3745 &pk as *const _ as *mut std::ffi::c_void,
3746 &pv as *const _ as *mut _,
3747 &pkc as *const _ as *mut _,
3748 &pvc as *const _ as *mut _,
3749 &pt as *const _ as *mut _,
3750 &kdk as *const _ as *mut _,
3751 &kdv as *const _ as *mut _,
3752 &ktb as *const _ as *mut _,
3753 &vtb as *const _ as *mut _,
3754 ];
3755 unsafe {
3756 self.launch_pdl_flash(
3757 g,
3758 "append_quantize_kv_q8_0_q5_1_dc",
3759 (nblk, 1, 1),
3760 (32, 1, 1),
3761 0,
3762 &mut ps,
3763 )?;
3764 }
3765 return Ok(());
3766 }
3767 let f = if g {
3768 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3769 } else {
3770 self.func("append_quantize_kv_q8_0_q5_1_dc")
3771 };
3772 let cfg = LaunchConfig {
3773 grid_dim: (nblk, 1, 1),
3774 block_dim: (32, 1, 1),
3775 shared_mem_bytes: 0,
3776 };
3777 let __s_b = self.gpu.stream();
3778 let mut b = __s_b.launch_builder(&f);
3779 b.arg(k_row)
3780 .arg(v_row)
3781 .arg(kc)
3782 .arg(vc)
3783 .arg(t_dev)
3784 .arg(&kdk)
3785 .arg(&kdv)
3786 .arg(&ktb)
3787 .arg(&vtb);
3788 unsafe {
3789 b.launch(cfg)?;
3790 }
3791 Ok(())
3792 }
3793
3794 #[allow(clippy::too_many_arguments)]
3801 pub fn append_kv_quantized_rows(
3802 &self,
3803 k_rows: &CudaSlice<f32>,
3804 v_rows: &CudaSlice<f32>,
3805 kc: &mut CudaSlice<u8>,
3806 vc: &mut CudaSlice<u8>,
3807 t0: usize,
3808 t: usize,
3809 kv_dim_k: usize,
3810 kv_dim_v: usize,
3811 k_tok_bytes: usize,
3812 v_tok_bytes: usize,
3813 g: bool,
3814 ) -> Result<(), Box<dyn std::error::Error>> {
3815 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3816 for i in 0..t {
3817 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3818 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3819 self.append_kv_quantized_view(
3820 &k_row,
3821 &v_row,
3822 kc,
3823 vc,
3824 t0 + i,
3825 kv_dim_k,
3826 kv_dim_v,
3827 k_tok_bytes,
3828 v_tok_bytes,
3829 g,
3830 )?;
3831 }
3832 return Ok(());
3833 }
3834 let f = if g {
3835 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3836 } else {
3837 self.func("append_quantize_kv_q8_0_q5_1_rows")
3838 };
3839 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3840 let cfg = LaunchConfig {
3841 grid_dim: (nblk, t as u32, 1),
3842 block_dim: (32, 1, 1),
3843 shared_mem_bytes: 0,
3844 };
3845 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3846 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3847 let __s_b = self.gpu.stream();
3848 let mut b = __s_b.launch_builder(&f);
3849 b.arg(k_rows)
3850 .arg(v_rows)
3851 .arg(kc)
3852 .arg(vc)
3853 .arg(&t0i)
3854 .arg(&kdk)
3855 .arg(&kdv)
3856 .arg(&ktb)
3857 .arg(&vtb);
3858 unsafe {
3859 b.launch(cfg)?;
3860 }
3861 Ok(())
3862 }
3863
3864 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3868 let f = self.func("inc_i32");
3869 let cfg = LaunchConfig {
3870 grid_dim: (1, 1, 1),
3871 block_dim: (1, 1, 1),
3872 shared_mem_bytes: 0,
3873 };
3874 let __s_b = self.gpu.stream();
3875 let mut b = __s_b.launch_builder(&f);
3876 b.arg(p);
3877 unsafe {
3878 b.launch(cfg)?;
3879 }
3880 Ok(())
3881 }
3882
3883 pub fn append_kv_quantized_view(
3886 &self,
3887 k_row: &cudarc::driver::CudaView<f32>,
3888 v_row: &cudarc::driver::CudaView<f32>,
3889 kc: &mut CudaSlice<u8>,
3890 vc: &mut CudaSlice<u8>,
3891 t: usize,
3892 kv_dim_k: usize,
3893 kv_dim_v: usize,
3894 k_tok_bytes: usize,
3895 v_tok_bytes: usize,
3896 g: bool,
3897 ) -> Result<(), Box<dyn std::error::Error>> {
3898 let f = if g {
3899 self.func_g("append_quantize_kv_q8_0_q5_1")
3900 } else {
3901 self.func("append_quantize_kv_q8_0_q5_1")
3902 };
3903 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3904 let cfg = LaunchConfig {
3905 grid_dim: (nblk, 1, 1),
3906 block_dim: (32, 1, 1),
3907 shared_mem_bytes: 0,
3908 };
3909 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3910 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3911 let __s_b = self.gpu.stream();
3912 let mut b = __s_b.launch_builder(&f);
3913 b.arg(k_row)
3914 .arg(v_row)
3915 .arg(kc)
3916 .arg(vc)
3917 .arg(&ti)
3918 .arg(&kdk)
3919 .arg(&kdv)
3920 .arg(&ktb)
3921 .arg(&vtb);
3922 unsafe {
3923 b.launch(cfg)?;
3924 }
3925 Ok(())
3926 }
3927
3928 pub fn copy_view_into(
3931 &self,
3932 dst: &mut CudaSlice<f32>,
3933 off: usize,
3934 src: &cudarc::driver::CudaView<f32>,
3935 len: usize,
3936 ) -> Result<(), Box<dyn std::error::Error>> {
3937 let mut view = dst.slice_mut(off..off + len);
3938 self.gpu
3939 .stream()
3940 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3941 Ok(())
3942 }
3943
3944 pub fn clone_dtod(
3948 &self,
3949 src: &CudaSlice<f32>,
3950 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3951 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
3952 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
3953 Ok(dst)
3954 }
3955
3956 pub fn dtod_copy_view(
3959 &self,
3960 src: &cudarc::driver::CudaView<f32>,
3961 dst: &mut CudaSlice<f32>,
3962 ) -> Result<(), Box<dyn std::error::Error>> {
3963 self.gpu.stream().memcpy_dtod(src, dst)?;
3964 Ok(())
3965 }
3966
3967 pub fn dtod_copy_view_i8(
3969 &self,
3970 src: &cudarc::driver::CudaView<i8>,
3971 dst: &mut CudaSlice<i8>,
3972 ) -> Result<(), Box<dyn std::error::Error>> {
3973 self.gpu.stream().memcpy_dtod(src, dst)?;
3974 Ok(())
3975 }
3976
3977 pub fn dtod_copy_into(
3979 &self,
3980 src: &CudaSlice<f32>,
3981 dst: &mut CudaSlice<f32>,
3982 offset: usize,
3983 ) -> Result<(), Box<dyn std::error::Error>> {
3984 let n = src.len();
3985 let mut dv = dst.slice_mut(offset..offset + n);
3986 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
3987 Ok(())
3988 }
3989
3990 pub fn copy_batch_uniform_f32(
3996 &self,
3997 table: &CudaSlice<u64>,
3998 n: usize,
3999 words: usize,
4000 ) -> Result<(), Box<dyn std::error::Error>> {
4001 if n == 0 || words == 0 {
4002 return Ok(());
4003 }
4004 debug_assert!(
4005 table.len() >= 2 * n,
4006 "pointer table must hold n srcs + n dsts"
4007 );
4008 let f = self.func("copy_batch_uniform_f32");
4009 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4012 let (ni, wi) = (n as i32, words as i32);
4013 let cfg = LaunchConfig {
4014 grid_dim: (chunks, n as u32, 1),
4015 block_dim: (256, 1, 1),
4016 shared_mem_bytes: 0,
4017 };
4018 let __s = self.gpu.stream();
4019 let mut b = __s.launch_builder(&f);
4020 b.arg(table).arg(&ni).arg(&wi);
4021 unsafe {
4022 b.launch(cfg)?;
4023 }
4024 Ok(())
4025 }
4026
4027 pub fn htod_u64_into(
4030 &self,
4031 v: &[u64],
4032 dst: &mut CudaSlice<u64>,
4033 ) -> Result<(), Box<dyn std::error::Error>> {
4034 let mut view = dst.slice_mut(0..v.len());
4035 self.gpu.stream().memcpy_htod(v, &mut view)?;
4036 Ok(())
4037 }
4038
4039 pub fn copy_indirect_src_f32(
4044 &self,
4045 src_entry: &cudarc::driver::CudaView<u64>,
4046 dst: &mut CudaSlice<f32>,
4047 dst_off: usize,
4048 words: usize,
4049 ) -> Result<(), Box<dyn std::error::Error>> {
4050 let f = self.func("copy_indirect_src_f32");
4051 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4052 let wi = words as i32;
4053 let cfg = LaunchConfig {
4054 grid_dim: (chunks, 1, 1),
4055 block_dim: (256, 1, 1),
4056 shared_mem_bytes: 0,
4057 };
4058 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4059 let __s = self.gpu.stream();
4060 let mut b = __s.launch_builder(&f);
4061 b.arg(src_entry).arg(&mut dv).arg(&wi);
4062 unsafe {
4063 b.launch(cfg)?;
4064 }
4065 Ok(())
4066 }
4067
4068 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4070 self.alloc_uninit::<i8>(n)
4071 }
4072
4073 pub fn qmatvec(
4075 &self,
4076 w: &CudaSlice<u8>,
4077 x: &CudaSlice<f32>,
4078 m: usize,
4079 in_f: usize,
4080 out_f: usize,
4081 qtype: i32,
4082 row_bytes: usize,
4083 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4084 let f = self.func("qmatvec_f32");
4085 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4087 grid_dim: (out_f as u32, m as u32, 1),
4088 block_dim: (256, 1, 1),
4089 shared_mem_bytes: 0,
4090 };
4091 let (inf, outf, mi, qt, rb) =
4092 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4093 let __s_b = self.gpu.stream();
4094 let mut b = __s_b.launch_builder(&f);
4095 b.arg(w)
4096 .arg(x)
4097 .arg(&mut y)
4098 .arg(&inf)
4099 .arg(&outf)
4100 .arg(&mi)
4101 .arg(&qt)
4102 .arg(&rb);
4103 unsafe {
4104 b.launch(cfg)?;
4105 }
4106 Ok(y)
4107 }
4108
4109 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4111 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4112 self.keep_if_capturing(&s);
4113 Ok(s)
4114 }
4115
4116 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4120 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4121 self.keep_if_capturing(&s);
4122 Ok(s)
4123 }
4124
4125 pub fn memset_zeros_view(
4128 &self,
4129 dst: &mut cudarc::driver::CudaViewMut<f32>,
4130 ) -> Result<(), Box<dyn std::error::Error>> {
4131 self.gpu.stream().memset_zeros(dst)?;
4132 Ok(())
4133 }
4134
4135 pub fn stage_expert(
4141 &self,
4142 host_bytes: &[u8],
4143 scratch: &mut CudaSlice<u8>,
4144 off: usize,
4145 ) -> Result<(), Box<dyn std::error::Error>> {
4146 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4149 }
4150
4151 pub fn moe_router_topk(
4157 &self,
4158 logits: &CudaSlice<f32>,
4159 t: usize,
4160 n_expert: usize,
4161 n_used: usize,
4162 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4163 let f = self.func("moe_router_topk_f32");
4164 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 {
4167 grid_dim: (t as u32, 1, 1),
4168 block_dim: (n_expert as u32, 1, 1),
4169 shared_mem_bytes: 0,
4170 };
4171 let (ne, nu) = (n_expert as i32, n_used as i32);
4172 let __s_b = self.gpu.stream();
4173 let mut b = __s_b.launch_builder(&f);
4174 b.arg(logits)
4175 .arg(&mut sel_idx)
4176 .arg(&mut sel_w)
4177 .arg(&ne)
4178 .arg(&nu);
4179 unsafe {
4180 b.launch(cfg)?;
4181 }
4182 Ok((sel_idx, sel_w))
4183 }
4184
4185 pub fn moe_router_topk_scaled(
4188 &self,
4189 logits: &CudaSlice<f32>,
4190 t: usize,
4191 n_expert: usize,
4192 n_used: usize,
4193 ex_scale: &CudaSlice<f32>,
4194 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4195 let f = self.func("moe_router_topk_scaled_f32");
4200 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4201 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4202 let cfg = LaunchConfig {
4203 grid_dim: (t as u32, 1, 1),
4204 block_dim: (n_expert as u32, 1, 1),
4205 shared_mem_bytes: 0,
4206 };
4207 let (ne, nu) = (n_expert as i32, n_used as i32);
4208 let __s_b = self.gpu.stream();
4209 let mut b = __s_b.launch_builder(&f);
4210 b.arg(logits)
4211 .arg(&mut sel_idx)
4212 .arg(&mut sel_w)
4213 .arg(&ne)
4214 .arg(&nu)
4215 .arg(ex_scale);
4216 unsafe {
4217 b.launch(cfg)?;
4218 }
4219 Ok((sel_idx, sel_w))
4220 }
4221
4222 pub fn moe_router_topk_host(
4230 &self,
4231 logits: &CudaSlice<f32>,
4232 t: usize,
4233 n_expert: usize,
4234 n_used: usize,
4235 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4236 let f = self.func("moe_router_topk_f32");
4237 let n = t * n_used;
4238 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4239 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4240 let cfg = LaunchConfig {
4241 grid_dim: (t as u32, 1, 1),
4242 block_dim: (n_expert as u32, 1, 1),
4243 shared_mem_bytes: 0,
4244 };
4245 let (ne, nu) = (n_expert as i32, n_used as i32);
4246 let __s_b = self.gpu.stream();
4247 let mut b = __s_b.launch_builder(&f);
4248 b.arg(logits)
4249 .arg(&mut sel_idx)
4250 .arg(&mut sel_w)
4251 .arg(&ne)
4252 .arg(&nu);
4253 unsafe {
4254 b.launch(cfg)?;
4255 }
4256 let bytes = n * 8;
4258 let mut guard = self.router_stage.lock().unwrap();
4259 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4260 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4261 }
4262 let stage = guard.as_mut().unwrap();
4263 let (si, sw) = unsafe {
4264 (
4265 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4266 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4267 )
4268 };
4269 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()))
4273 }
4274
4275 #[allow(clippy::too_many_arguments)]
4279 pub fn moe_router_sigmoid_topk(
4280 &self,
4281 logits: &CudaSlice<f32>,
4282 t: usize,
4283 n_expert: usize,
4284 n_used: usize,
4285 active_count: usize,
4286 correction_bias: &CudaSlice<f32>,
4287 active: &CudaSlice<u8>,
4288 scaling_factor: f32,
4289 route_norm: bool,
4290 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4291 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4292 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4293 return Err(format!(
4294 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4295 )
4296 .into());
4297 }
4298 if logits.len() < t * n_expert
4299 || correction_bias.len() != n_expert
4300 || active.len() != n_expert
4301 {
4302 return Err(format!(
4303 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4304 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4305 ).into());
4306 }
4307 let f = self.func("moe_router_sigmoid_topk_f32");
4308 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4309 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4310 let threads = n_expert.div_ceil(32) * 32;
4311 let cfg = LaunchConfig {
4312 grid_dim: (t as u32, 1, 1),
4313 block_dim: (threads as u32, 1, 1),
4314 shared_mem_bytes: 0,
4315 };
4316 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4317 let __s_b = self.gpu.stream();
4318 let mut b = __s_b.launch_builder(&f);
4319 b.arg(logits)
4320 .arg(correction_bias)
4321 .arg(active)
4322 .arg(&mut sel_idx)
4323 .arg(&mut sel_w)
4324 .arg(&ne)
4325 .arg(&nu)
4326 .arg(&scaling_factor)
4327 .arg(&rn);
4328 unsafe {
4329 b.launch(cfg)?;
4330 }
4331 Ok((sel_idx, sel_w))
4332 }
4333
4334 #[allow(clippy::too_many_arguments)]
4337 pub fn moe_router_sigmoid_topk_host(
4338 &self,
4339 logits: &CudaSlice<f32>,
4340 t: usize,
4341 n_expert: usize,
4342 n_used: usize,
4343 active_count: usize,
4344 correction_bias: &CudaSlice<f32>,
4345 active: &CudaSlice<u8>,
4346 scaling_factor: f32,
4347 route_norm: bool,
4348 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4349 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4350 logits,
4351 t,
4352 n_expert,
4353 n_used,
4354 active_count,
4355 correction_bias,
4356 active,
4357 scaling_factor,
4358 route_norm,
4359 )?;
4360 let n = t * n_used;
4361 let bytes = n * 8;
4362 let mut guard = self.router_stage.lock().unwrap();
4363 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4364 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4365 }
4366 let stage = guard.as_mut().unwrap();
4367 let (si, sw) = unsafe {
4368 (
4369 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4370 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4371 )
4372 };
4373 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4374 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4375 self.gpu.stream().synchronize()?;
4376 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4377 }
4378
4379 pub fn stage_expert_async(
4383 &self,
4384 host_bytes: &[u8],
4385 scratch: &mut CudaSlice<u8>,
4386 off: usize,
4387 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4388 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4389 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4390 Ok(self.copy_stream.record_event(None)?)
4391 }
4392
4393 pub fn compute_wait(
4395 &self,
4396 ev: &cudarc::driver::CudaEvent,
4397 ) -> Result<(), Box<dyn std::error::Error>> {
4398 self.gpu.stream().wait(ev)?;
4399 Ok(())
4400 }
4401
4402 pub fn qmatvec_view(
4407 &self,
4408 w: &CudaSlice<u8>,
4409 range: std::ops::Range<usize>,
4410 x: &cudarc::driver::CudaView<f32>,
4411 m: usize,
4412 in_f: usize,
4413 out_f: usize,
4414 qtype: i32,
4415 row_bytes: usize,
4416 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4417 let f = self.func("qmatvec_f32");
4418 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4421 grid_dim: (out_f as u32, m as u32, 1),
4422 block_dim: (256, 1, 1),
4423 shared_mem_bytes: 0,
4424 };
4425 let (inf, outf, mi, qt, rb) =
4426 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4427 let __s_b = self.gpu.stream();
4428 let mut b = __s_b.launch_builder(&f);
4429 b.arg(&wv)
4430 .arg(x)
4431 .arg(&mut y)
4432 .arg(&inf)
4433 .arg(&outf)
4434 .arg(&mi)
4435 .arg(&qt)
4436 .arg(&rb);
4437 unsafe {
4438 b.launch(cfg)?;
4439 }
4440 Ok(y)
4441 }
4442
4443 #[allow(clippy::too_many_arguments)]
4450 pub fn moe_gate_up_silu8_q8(
4454 &self,
4455 gp: WPtr8,
4456 up: WPtr8,
4457 aq: &CudaSlice<i8>,
4458 ad: &CudaSlice<f32>,
4459 in_f: usize,
4460 n_ff: usize,
4461 n_used: usize,
4462 qt_g: i32,
4463 qt_u: i32,
4464 rb_g: usize,
4465 rb_u: usize,
4466 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4467 let f = self.func("moe_gate_up_silu8_q8");
4468 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4469 let cfg = LaunchConfig {
4470 grid_dim: (n_ff as u32, n_used as u32, 1),
4471 block_dim: (32, 1, 1),
4472 shared_mem_bytes: 0,
4473 };
4474 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4475 let __s_b = self.gpu.stream();
4476 let mut b = __s_b.launch_builder(&f);
4477 b.arg(&gp)
4478 .arg(&up)
4479 .arg(aq)
4480 .arg(ad)
4481 .arg(&mut act)
4482 .arg(&inf)
4483 .arg(&nff)
4484 .arg(&qt_g)
4485 .arg(&qt_u)
4486 .arg(&rbg)
4487 .arg(&rbu);
4488 unsafe {
4489 b.launch(cfg)?;
4490 }
4491 Ok(act)
4492 }
4493
4494 #[allow(clippy::too_many_arguments)]
4495 pub fn moe_down8_fma_q8(
4496 &self,
4497 dp: WPtr8,
4498 w: F32x8,
4499 aq2: &CudaSlice<i8>,
4500 ad2: &CudaSlice<f32>,
4501 dst: &mut cudarc::driver::CudaViewMut<f32>,
4502 in_f: usize,
4503 out_f: usize,
4504 n_used: usize,
4505 qt: i32,
4506 rb: usize,
4507 ) -> Result<(), Box<dyn std::error::Error>> {
4508 let f = self.func("moe_down8_fma_q8");
4509 let cfg = LaunchConfig {
4510 grid_dim: (out_f as u32, 1, 1),
4511 block_dim: (32, 1, 1),
4512 shared_mem_bytes: 0,
4513 };
4514 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4515 let __s_b = self.gpu.stream();
4516 let mut b = __s_b.launch_builder(&f);
4517 b.arg(&dp)
4518 .arg(&w)
4519 .arg(aq2)
4520 .arg(ad2)
4521 .arg(dst)
4522 .arg(&inf)
4523 .arg(&outf)
4524 .arg(&nu)
4525 .arg(&qt)
4526 .arg(&rbi);
4527 unsafe {
4528 b.launch(cfg)?;
4529 }
4530 Ok(())
4531 }
4532
4533 pub fn qmatvec_expert_q8(
4535 &self,
4536 w: &CudaSlice<u8>,
4537 range: std::ops::Range<usize>,
4538 aq: &CudaSlice<i8>,
4539 ad: &CudaSlice<f32>,
4540 m: usize,
4541 in_f: usize,
4542 out_f: usize,
4543 qtype: i32,
4544 row_bytes: usize,
4545 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4546 let f = self.func("qmatvec_expert_q8");
4547 let wv = w.slice(range);
4548 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4549 const ROWS: u32 = 4; let cfg = LaunchConfig {
4551 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4552 block_dim: (32, ROWS, 1),
4553 shared_mem_bytes: 0,
4554 };
4555 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4556 let __s_b = self.gpu.stream();
4557 let mut b = __s_b.launch_builder(&f);
4558 b.arg(&wv)
4559 .arg(aq)
4560 .arg(ad)
4561 .arg(&mut y)
4562 .arg(&inf)
4563 .arg(&outf)
4564 .arg(&mi)
4565 .arg(&qtype)
4566 .arg(&rbi);
4567 unsafe {
4568 b.launch(cfg)?;
4569 }
4570 Ok(y)
4571 }
4572
4573 pub fn moe_gate_up_silu8(
4574 &self,
4575 gp: WPtr8,
4576 up: WPtr8,
4577 x: &cudarc::driver::CudaView<f32>,
4578 in_f: usize,
4579 n_ff: usize,
4580 n_used: usize,
4581 qt_g: i32,
4582 qt_u: i32,
4583 rb_g: usize,
4584 rb_u: usize,
4585 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4586 let f = self.func("moe_gate_up_silu8_f32");
4587 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
4589 grid_dim: (n_ff as u32, n_used as u32, 1),
4590 block_dim: (256, 1, 1),
4591 shared_mem_bytes: 0,
4592 };
4593 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4594 let __s_b = self.gpu.stream();
4595 let mut b = __s_b.launch_builder(&f);
4596 b.arg(&gp)
4597 .arg(&up)
4598 .arg(x)
4599 .arg(&mut act)
4600 .arg(&inf)
4601 .arg(&nff)
4602 .arg(&qt_g)
4603 .arg(&qt_u)
4604 .arg(&rbg)
4605 .arg(&rbu);
4606 unsafe {
4607 b.launch(cfg)?;
4608 }
4609 Ok(act)
4610 }
4611
4612 #[allow(clippy::too_many_arguments)]
4618 pub fn moe_down8_fma_into(
4619 &self,
4620 dp: WPtr8,
4621 w: F32x8,
4622 act: &CudaSlice<f32>,
4623 dst: &mut cudarc::driver::CudaViewMut<f32>,
4624 in_f: usize,
4625 out_f: usize,
4626 n_used: usize,
4627 qt: i32,
4628 rb: usize,
4629 ) -> Result<(), Box<dyn std::error::Error>> {
4630 let f = self.func("moe_down8_fma_f32");
4631 let cfg = LaunchConfig {
4632 grid_dim: (out_f as u32, 1, 1),
4633 block_dim: (256, 1, 1),
4634 shared_mem_bytes: 0,
4635 };
4636 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4637 let __s_b = self.gpu.stream();
4638 let mut b = __s_b.launch_builder(&f);
4639 b.arg(&dp)
4640 .arg(&w)
4641 .arg(act)
4642 .arg(dst)
4643 .arg(&inf)
4644 .arg(&outf)
4645 .arg(&nu)
4646 .arg(&qt)
4647 .arg(&rbv);
4648 unsafe {
4649 b.launch(cfg)?;
4650 }
4651 Ok(())
4652 }
4653
4654 #[allow(clippy::too_many_arguments)]
4659 #[allow(clippy::too_many_arguments)]
4674 #[allow(clippy::too_many_arguments)]
4676 pub fn moe_pairs_matvec_q8(
4677 &self,
4678 table: &CudaSlice<u64>,
4679 proj: i32,
4680 pair_tok: &CudaSlice<i32>,
4681 pair_ex: &CudaSlice<i32>,
4682 aq: &CudaSlice<i8>,
4683 ad: &CudaSlice<f32>,
4684 in_f: usize,
4685 out_f: usize,
4686 n_expert: usize,
4687 n_pairs: usize,
4688 qtype: i32,
4689 row_bytes: usize,
4690 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4691 let f = self.func("moe_pairs_matvec_q8");
4692 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4693 const ROWS: u32 = 4;
4694 let cfg = LaunchConfig {
4695 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4696 block_dim: (32, ROWS, 1),
4697 shared_mem_bytes: 0,
4698 };
4699 let (inf, outf, ne, np, rbi) = (
4700 in_f as i32,
4701 out_f as i32,
4702 n_expert as i32,
4703 n_pairs as i32,
4704 row_bytes as i64,
4705 );
4706 let __s_b = self.gpu.stream();
4707 let mut b = __s_b.launch_builder(&f);
4708 b.arg(table)
4709 .arg(&proj)
4710 .arg(pair_tok)
4711 .arg(pair_ex)
4712 .arg(aq)
4713 .arg(ad)
4714 .arg(&mut y)
4715 .arg(&inf)
4716 .arg(&outf)
4717 .arg(&ne)
4718 .arg(&np)
4719 .arg(&qtype)
4720 .arg(&rbi);
4721 unsafe {
4722 b.launch(cfg)?;
4723 }
4724 Ok(y)
4725 }
4726
4727 #[allow(clippy::too_many_arguments)]
4729 pub fn moe_pairs_matvec_q8_em(
4730 &self,
4731 table: &CudaSlice<u64>,
4732 proj: i32,
4733 ex_ids: &CudaSlice<i32>,
4734 ex_off: &CudaSlice<i32>,
4735 ex_pairs: &CudaSlice<i32>,
4736 pair_tok: &CudaSlice<i32>,
4737 aq: &CudaSlice<i8>,
4738 ad: &CudaSlice<f32>,
4739 in_f: usize,
4740 out_f: usize,
4741 n_expert: usize,
4742 n_active: usize,
4743 n_pairs: usize,
4744 qtype: i32,
4745 row_bytes: usize,
4746 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4747 let f = self.func("moe_pairs_matvec_q8_em");
4748 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4749 const ROWS: u32 = 4;
4750 let cfg = LaunchConfig {
4751 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4752 block_dim: (32, ROWS, 1),
4753 shared_mem_bytes: 0,
4754 };
4755 let (inf, outf, ne, na, rbi) = (
4756 in_f as i32,
4757 out_f as i32,
4758 n_expert as i32,
4759 n_active as i32,
4760 row_bytes as i64,
4761 );
4762 let __s_b = self.gpu.stream();
4763 let mut b = __s_b.launch_builder(&f);
4764 b.arg(table)
4765 .arg(&proj)
4766 .arg(ex_ids)
4767 .arg(ex_off)
4768 .arg(ex_pairs)
4769 .arg(pair_tok)
4770 .arg(aq)
4771 .arg(ad)
4772 .arg(&mut y)
4773 .arg(&inf)
4774 .arg(&outf)
4775 .arg(&ne)
4776 .arg(&na)
4777 .arg(&qtype)
4778 .arg(&rbi);
4779 unsafe {
4780 b.launch(cfg)?;
4781 }
4782 Ok(y)
4783 }
4784
4785 #[allow(clippy::too_many_arguments)]
4788 pub fn moe_pairs_matvec_q8_dec(
4789 &self,
4790 table: &CudaSlice<u64>,
4791 proj: i32,
4792 ex_ids: &CudaSlice<i32>,
4793 ex_off: &CudaSlice<i32>,
4794 ex_pairs: &CudaSlice<i32>,
4795 pair_tok: &CudaSlice<i32>,
4796 aq: &CudaSlice<i8>,
4797 ad: &CudaSlice<f32>,
4798 in_f: usize,
4799 out_f: usize,
4800 n_expert: usize,
4801 n_active: usize,
4802 n_pairs: usize,
4803 qtype: i32,
4804 row_bytes: usize,
4805 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4806 let f = self.func("moe_pairs_matvec_q8_dec");
4807 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4808 const ROWS: u32 = 4;
4809 let cfg = LaunchConfig {
4810 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4811 block_dim: (32, ROWS, 1),
4812 shared_mem_bytes: 0,
4813 };
4814 let (inf, outf, ne, na, rbi) = (
4815 in_f as i32,
4816 out_f as i32,
4817 n_expert as i32,
4818 n_active as i32,
4819 row_bytes as i64,
4820 );
4821 let __s_b = self.gpu.stream();
4822 let mut b = __s_b.launch_builder(&f);
4823 b.arg(table)
4824 .arg(&proj)
4825 .arg(ex_ids)
4826 .arg(ex_off)
4827 .arg(ex_pairs)
4828 .arg(pair_tok)
4829 .arg(aq)
4830 .arg(ad)
4831 .arg(&mut y)
4832 .arg(&inf)
4833 .arg(&outf)
4834 .arg(&ne)
4835 .arg(&na)
4836 .arg(&qtype)
4837 .arg(&rbi);
4838 unsafe {
4839 b.launch(cfg)?;
4840 }
4841 Ok(y)
4842 }
4843
4844 pub fn moe_pairs_gelu_mul(
4845 &self,
4846 gate: &CudaSlice<f32>,
4847 up: &CudaSlice<f32>,
4848 n: usize,
4849 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4850 let f = self.func("moe_pairs_gelu_mul");
4851 let mut act = self.alloc_uninit::<f32>(n)?;
4852 let cfg = LaunchConfig::for_num_elems(n as u32);
4853 let nl = n as i64;
4854 let __s_b = self.gpu.stream();
4855 let mut b = __s_b.launch_builder(&f);
4856 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4857 unsafe {
4858 b.launch(cfg)?;
4859 }
4860 Ok(act)
4861 }
4862
4863 pub fn moe_pairs_silu_mul(
4864 &self,
4865 gate: &CudaSlice<f32>,
4866 up: &CudaSlice<f32>,
4867 n: usize,
4868 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4869 let f = self.func("moe_pairs_silu_mul");
4870 let mut act = self.alloc_uninit::<f32>(n)?;
4871 let cfg = LaunchConfig::for_num_elems(n as u32);
4872 let nl = n as i64;
4873 let __s_b = self.gpu.stream();
4874 let mut b = __s_b.launch_builder(&f);
4875 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4876 unsafe {
4877 b.launch(cfg)?;
4878 }
4879 Ok(act)
4880 }
4881
4882 #[allow(clippy::too_many_arguments)]
4883 pub fn moe_pairs_scatter(
4884 &self,
4885 y_down: &CudaSlice<f32>,
4886 pair_w: &CudaSlice<f32>,
4887 tok_pair_off: &CudaSlice<i32>,
4888 tok_pair_ids: &CudaSlice<i32>,
4889 moe_out: &mut CudaSlice<f32>,
4890 t: usize,
4891 n_embd: usize,
4892 ) -> Result<(), Box<dyn std::error::Error>> {
4893 let f = self.func("moe_pairs_scatter");
4894 let cfg = LaunchConfig {
4895 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4896 block_dim: (256, 1, 1),
4897 shared_mem_bytes: 0,
4898 };
4899 let ne = n_embd as i32;
4900 let __s_b = self.gpu.stream();
4901 let mut b = __s_b.launch_builder(&f);
4902 b.arg(y_down)
4903 .arg(pair_w)
4904 .arg(tok_pair_off)
4905 .arg(tok_pair_ids)
4906 .arg(moe_out)
4907 .arg(&ne);
4908 unsafe {
4909 b.launch(cfg)?;
4910 }
4911 Ok(())
4912 }
4913
4914 #[allow(clippy::too_many_arguments)]
4918 pub fn moe_gate_up_gelu8_dev_q8(
4919 &self,
4920 table: &CudaSlice<u64>,
4921 sel: &cudarc::driver::CudaView<i32>,
4922 aq: &CudaSlice<i8>,
4923 ad: &CudaSlice<f32>,
4924 in_f: usize,
4925 n_ff: usize,
4926 n_used: usize,
4927 n_expert: usize,
4928 qt_g: i32,
4929 qt_u: i32,
4930 rb_g: usize,
4931 rb_u: usize,
4932 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4933 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4934 let (inf, nff, ne, rbg, rbu) = (
4935 in_f as i32,
4936 n_ff as i32,
4937 n_expert as i32,
4938 rb_g as i64,
4939 rb_u as i64,
4940 );
4941 let f = self.func("moe_gate_up_gelu8_dev_q8");
4942 let cfg = LaunchConfig {
4943 grid_dim: (n_ff as u32, n_used as u32, 1),
4944 block_dim: (32, 1, 1),
4945 shared_mem_bytes: 0,
4946 };
4947 let __s_b = self.gpu.stream();
4948 let mut b = __s_b.launch_builder(&f);
4949 b.arg(table)
4950 .arg(sel)
4951 .arg(aq)
4952 .arg(ad)
4953 .arg(&mut act)
4954 .arg(&inf)
4955 .arg(&nff)
4956 .arg(&ne)
4957 .arg(&qt_g)
4958 .arg(&qt_u)
4959 .arg(&rbg)
4960 .arg(&rbu);
4961 unsafe {
4962 b.launch(cfg)?;
4963 }
4964 Ok(act)
4965 }
4966
4967 #[allow(clippy::too_many_arguments)]
4969 pub fn moe_gate_up_gelu8_dev_q8_rows(
4970 &self,
4971 table: &CudaSlice<u64>,
4972 sel: &CudaSlice<i32>,
4973 aq: &CudaSlice<i8>,
4974 ad: &CudaSlice<f32>,
4975 t: usize,
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>(t * n_used * n_ff)?;
4986 let (inf, nff, ne, rbg, rbu, nu) = (
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 n_used as i32,
4993 );
4994 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
4995 let cfg = LaunchConfig {
4996 grid_dim: (n_ff as u32, n_used as u32, t as u32),
4997 block_dim: (32, 1, 1),
4998 shared_mem_bytes: 0,
4999 };
5000 let __s_b = self.gpu.stream();
5001 let mut b = __s_b.launch_builder(&f);
5002 b.arg(table)
5003 .arg(sel)
5004 .arg(aq)
5005 .arg(ad)
5006 .arg(&mut act)
5007 .arg(&inf)
5008 .arg(&nff)
5009 .arg(&ne)
5010 .arg(&qt_g)
5011 .arg(&qt_u)
5012 .arg(&rbg)
5013 .arg(&rbu)
5014 .arg(&nu);
5015 unsafe {
5016 b.launch(cfg)?;
5017 }
5018 Ok(act)
5019 }
5020
5021 #[allow(clippy::too_many_arguments)]
5023 pub fn moe_gate_up_gelu8_dev_q8_csr(
5024 &self,
5025 table: &CudaSlice<u64>,
5026 sel: &CudaSlice<i32>,
5027 aq: &CudaSlice<i8>,
5028 ad: &CudaSlice<f32>,
5029 n_pairs: usize,
5030 in_f: usize,
5031 n_ff: usize,
5032 n_used: usize,
5033 n_expert: usize,
5034 qt_g: i32,
5035 qt_u: i32,
5036 rb_g: usize,
5037 rb_u: usize,
5038 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5039 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5040 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5041 in_f as i32,
5042 n_ff as i32,
5043 n_expert as i32,
5044 rb_g as i64,
5045 rb_u as i64,
5046 n_used as i32,
5047 n_pairs as i32,
5048 );
5049 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5050 let cfg = LaunchConfig {
5051 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5052 block_dim: (32, 1, 1),
5053 shared_mem_bytes: 0,
5054 };
5055 let __s_b = self.gpu.stream();
5056 let mut b = __s_b.launch_builder(&f);
5057 b.arg(table)
5058 .arg(sel)
5059 .arg(aq)
5060 .arg(ad)
5061 .arg(&mut act)
5062 .arg(&inf)
5063 .arg(&nff)
5064 .arg(&ne)
5065 .arg(&qt_g)
5066 .arg(&qt_u)
5067 .arg(&rbg)
5068 .arg(&rbu)
5069 .arg(&nu)
5070 .arg(&npi);
5071 unsafe {
5072 b.launch(cfg)?;
5073 }
5074 Ok(act)
5075 }
5076
5077 #[allow(clippy::too_many_arguments)]
5079 pub fn moe_down8_fma_dev_q8_rows_g(
5080 &self,
5081 table: &CudaSlice<u64>,
5082 sel: &CudaSlice<i32>,
5083 w: &CudaSlice<f32>,
5084 aq2: &CudaSlice<i8>,
5085 ad2: &CudaSlice<f32>,
5086 dst: &mut CudaSlice<f32>,
5087 t: usize,
5088 in_f: usize,
5089 out_f: usize,
5090 n_used: usize,
5091 n_expert: usize,
5092 qt: i32,
5093 rb: usize,
5094 ) -> Result<(), Box<dyn std::error::Error>> {
5095 let (inf, outf, nu, ne, rbi) = (
5096 in_f as i32,
5097 out_f as i32,
5098 n_used as i32,
5099 n_expert as i32,
5100 rb as i64,
5101 );
5102 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5106 let f = self.func(if step_b1_w8 {
5107 "moe_down8_fma_dev_q8_rows_w8"
5108 } else {
5109 "moe_down8_fma_dev_q8_rows_g"
5110 });
5111 let cfg = LaunchConfig {
5112 grid_dim: (out_f as u32, 1, t as u32),
5113 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5114 shared_mem_bytes: 0,
5115 };
5116 let __s_b = self.gpu.stream();
5117 let mut b = __s_b.launch_builder(&f);
5118 b.arg(table)
5119 .arg(sel)
5120 .arg(w)
5121 .arg(aq2)
5122 .arg(ad2)
5123 .arg(dst)
5124 .arg(&inf)
5125 .arg(&outf)
5126 .arg(&nu)
5127 .arg(&ne)
5128 .arg(&qt)
5129 .arg(&rbi);
5130 unsafe {
5131 b.launch(cfg)?;
5132 }
5133 Ok(())
5134 }
5135
5136 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5140 let (out_f, in_f) = (2048usize, 2816usize);
5141 let nblk = in_f / 32;
5142 let mut seed = 0x9E3779B97F4A7C15u64;
5143 let mut rng = move || {
5144 seed = seed
5145 .wrapping_mul(6364136223846793005)
5146 .wrapping_add(1442695040888963407);
5147 (seed >> 33) as u8
5148 };
5149 let mut w = vec![0u8; out_f * nblk * 18];
5150 for b in w.iter_mut() {
5151 *b = rng();
5152 }
5153 for r in 0..out_f {
5154 for g in 0..nblk {
5155 let off = (r * nblk + g) * 18;
5156 w[off] = 0x00;
5157 w[off + 1] = 0x2C; }
5159 }
5160 let qplane = out_f * nblk * 16;
5161 let mut wrp = vec![0u8; w.len()];
5162 for r in 0..out_f {
5163 for g in 0..nblk {
5164 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5165 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5166 .copy_from_slice(&src[0..2]);
5167 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5168 }
5169 }
5170 let w_d = self.htod_bytes(&w)?;
5171 let wrp_d = self.htod_bytes(&wrp)?;
5172 let mut aq = vec![0i8; m * in_f];
5173 for v in aq.iter_mut() {
5174 *v = rng() as i8;
5175 }
5176 let aq_d = self.htod_i8(&aq)?;
5177 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5178 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5179 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5180 const RPB: u32 = 4;
5181 let cfg = LaunchConfig {
5182 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5183 block_dim: (32, RPB, 1),
5184 shared_mem_bytes: 0,
5185 };
5186 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5187 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5188 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5189 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5190 {
5191 let __s_b = self.gpu.stream();
5192 let mut b = __s_b.launch_builder(&fb);
5193 b.arg(&w_d)
5194 .arg(&aq_d)
5195 .arg(&ad_d)
5196 .arg(&mut y0)
5197 .arg(&inf)
5198 .arg(&outf)
5199 .arg(&mi)
5200 .arg(&rb);
5201 unsafe {
5202 b.launch(cfg)?;
5203 }
5204 let __s_b = self.gpu.stream();
5205 let mut b = __s_b.launch_builder(&fr);
5206 b.arg(&wrp_d)
5207 .arg(&aq_d)
5208 .arg(&ad_d)
5209 .arg(&mut y1)
5210 .arg(&inf)
5211 .arg(&outf)
5212 .arg(&mi)
5213 .arg(&qp);
5214 unsafe {
5215 b.launch(cfg)?;
5216 }
5217 }
5218 self.gpu.stream().synchronize()?;
5219 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5220 let nd = h0
5221 .iter()
5222 .zip(&h1)
5223 .filter(|(a, b)| a.to_bits() != b.to_bits())
5224 .count();
5225 if nd != 0 {
5226 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5227 }
5228 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5229 self.gpu.stream().synchronize()?;
5230 let t0 = std::time::Instant::now();
5231 for _ in 0..500 {
5232 if rp {
5233 let __s_b = self.gpu.stream();
5234 let mut b = __s_b.launch_builder(&fr);
5235 b.arg(&wrp_d)
5236 .arg(&aq_d)
5237 .arg(&ad_d)
5238 .arg(&mut y1)
5239 .arg(&inf)
5240 .arg(&outf)
5241 .arg(&mi)
5242 .arg(&qp);
5243 unsafe {
5244 b.launch(cfg)?;
5245 }
5246 } else {
5247 let __s_b = self.gpu.stream();
5248 let mut b = __s_b.launch_builder(&fb);
5249 b.arg(&w_d)
5250 .arg(&aq_d)
5251 .arg(&ad_d)
5252 .arg(&mut y0)
5253 .arg(&inf)
5254 .arg(&outf)
5255 .arg(&mi)
5256 .arg(&rb);
5257 unsafe {
5258 b.launch(cfg)?;
5259 }
5260 }
5261 }
5262 self.gpu.stream().synchronize()?;
5263 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5264 };
5265 let _ = time(false)?;
5266 let _ = time(true)?; Ok((time(false)?, time(true)?))
5268 }
5269
5270 pub fn build_q4_rp4(
5275 &self,
5276 t: &mut crate::model::GpuTensor,
5277 ) -> Result<(), Box<dyn std::error::Error>> {
5278 use crate::model::GpuTensor;
5279 let GpuTensor::Quant {
5280 bytes,
5281 qtype,
5282 row_bytes,
5283 ne,
5284 rp4,
5285 ..
5286 } = t
5287 else {
5288 return Ok(());
5289 };
5290 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5291 return Ok(());
5292 }
5293 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5294 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5295 return Ok(());
5296 }
5297 let nblk = in_f / 32;
5298 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5299 let f = self.func("q4_0_split_rp_build");
5300 let n = (out_f * nblk) as i32;
5301 let cfg = LaunchConfig {
5302 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5303 block_dim: (256, 1, 1),
5304 shared_mem_bytes: 0,
5305 };
5306 let (of, nb) = (out_f as i32, nblk as i32);
5307 let _ = n;
5308 let __s_b = self.gpu.stream();
5309 let mut b = __s_b.launch_builder(&f);
5310 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5311 unsafe {
5312 b.launch(cfg)?;
5313 }
5314 *rp4 = Some(dst);
5315 Ok(())
5316 }
5317
5318 pub fn build_q8_rp4(
5323 &self,
5324 t: &mut crate::model::GpuTensor,
5325 ) -> Result<(), Box<dyn std::error::Error>> {
5326 use crate::model::GpuTensor;
5327 let GpuTensor::Quant {
5328 bytes,
5329 qtype,
5330 row_bytes,
5331 ne,
5332 rp4,
5333 ..
5334 } = t
5335 else {
5336 return Ok(());
5337 };
5338 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5339 return Ok(());
5340 }
5341 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5342 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5343 return Ok(());
5344 }
5345 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5346 Ok(())
5347 }
5348
5349 pub fn build_q8_rp4_raw(
5352 &self,
5353 bytes: &CudaSlice<u8>,
5354 in_f: usize,
5355 out_f: usize,
5356 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5357 assert!(in_f % 32 == 0);
5358 let nblk = in_f / 32;
5359 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5360 let f = self.func("q8_0_split_rp_build");
5361 let cfg = LaunchConfig {
5362 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5363 block_dim: (256, 1, 1),
5364 shared_mem_bytes: 0,
5365 };
5366 let (of, nb) = (out_f as i32, nblk as i32);
5367 let __s_b = self.gpu.stream();
5368 let mut b = __s_b.launch_builder(&f);
5369 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5370 unsafe {
5371 b.launch(cfg)?;
5372 }
5373 Ok(dst)
5374 }
5375
5376 pub fn build_q4k_rp4(
5384 &self,
5385 t: &mut crate::model::GpuTensor,
5386 ) -> Result<(), Box<dyn std::error::Error>> {
5387 use crate::model::GpuTensor;
5388 let GpuTensor::Quant {
5389 bytes,
5390 qtype,
5391 row_bytes,
5392 ne,
5393 rp4,
5394 ..
5395 } = t
5396 else {
5397 return Ok(());
5398 };
5399 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5400 return Ok(());
5401 }
5402 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5403 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5404 return Ok(());
5405 }
5406 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5407 Ok(())
5408 }
5409
5410 pub fn build_q6k_rp4(
5411 &self,
5412 t: &mut crate::model::GpuTensor,
5413 ) -> Result<(), Box<dyn std::error::Error>> {
5414 use crate::model::GpuTensor;
5415 let GpuTensor::Quant {
5416 bytes,
5417 qtype,
5418 row_bytes,
5419 ne,
5420 rp4,
5421 ..
5422 } = t
5423 else {
5424 return Ok(());
5425 };
5426 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5427 return Ok(());
5428 }
5429 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5430 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5431 return Ok(());
5432 }
5433 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5434 Ok(())
5435 }
5436
5437 pub fn build_kq_rp4_raw(
5439 &self,
5440 bytes: &CudaSlice<u8>,
5441 in_f: usize,
5442 out_f: usize,
5443 qtype: i32,
5444 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5445 assert!(in_f % 256 == 0);
5446 let nsbk = in_f / 256;
5447 let (sb_bytes, kname) = match qtype {
5448 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5449 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5450 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5451 };
5452 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5453 let f = self.func(kname);
5454 let cfg = LaunchConfig {
5455 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5456 block_dim: (256, 1, 1),
5457 shared_mem_bytes: 0,
5458 };
5459 let (of, nb) = (out_f as i32, nsbk as i32);
5460 let __s_b = self.gpu.stream();
5461 let mut b = __s_b.launch_builder(&f);
5462 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5463 unsafe {
5464 b.launch(cfg)?;
5465 }
5466 Ok(dst)
5467 }
5468
5469 pub fn kqrp_enabled() -> bool {
5473 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5474 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5475 Ok("0") => false,
5476 Ok(_) => true,
5477 Err(_) => cfg!(memra_hopper_mma),
5478 })
5479 }
5480
5481 pub fn build_q4_rp_swap(
5487 &self,
5488 t: &mut crate::model::GpuTensor,
5489 ) -> Result<bool, Box<dyn std::error::Error>> {
5490 use crate::model::GpuTensor;
5491 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5501 return Ok(false);
5502 }
5503 self.build_q4_rp4(t)?;
5504 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5506 return Ok(false);
5507 };
5508 match rp4.take() {
5509 Some(split) => {
5510 *bytes = split; *rp = true;
5512 Ok(true)
5513 }
5514 None => Ok(false),
5515 }
5516 }
5517
5518 pub fn q4rp_enabled() -> bool {
5520 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5521 *ON.get_or_init(|| {
5522 std::env::var("MEMRA_Q4RP")
5523 .map(|v| v != "0")
5524 .unwrap_or(true)
5525 })
5526 }
5527
5528 pub fn copy_rows_strided(
5531 &self,
5532 src: &CudaSlice<f32>,
5533 dst: &mut CudaSlice<f32>,
5534 row_elems: usize,
5535 n_rows: usize,
5536 src_stride: usize,
5537 src_off: usize,
5538 ) -> Result<(), Box<dyn std::error::Error>> {
5539 let f = self.func("copy_rows_strided_f32");
5540 let cfg = LaunchConfig {
5541 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5542 block_dim: (256, 1, 1),
5543 shared_mem_bytes: 0,
5544 };
5545 let (re, nr) = (row_elems as i32, n_rows as i32);
5546 let (st, off) = (src_stride as i64, src_off as i64);
5547 let __s_b = self.gpu.stream();
5548 let mut b = __s_b.launch_builder(&f);
5549 b.arg(src)
5550 .arg(&mut *dst)
5551 .arg(&re)
5552 .arg(&nr)
5553 .arg(&st)
5554 .arg(&off);
5555 unsafe {
5556 b.launch(cfg)?;
5557 }
5558 Ok(())
5559 }
5560
5561 pub fn u32_set_k(
5563 &self,
5564 dst: &mut CudaSlice<u32>,
5565 v: u32,
5566 idx: usize,
5567 ) -> Result<(), Box<dyn std::error::Error>> {
5568 let f = self.func("u32_set_k");
5569 let cfg = LaunchConfig {
5570 grid_dim: (1, 1, 1),
5571 block_dim: (1, 1, 1),
5572 shared_mem_bytes: 0,
5573 };
5574 let ii = idx as i32;
5575 let __s_b = self.gpu.stream();
5576 let mut b = __s_b.launch_builder(&f);
5577 b.arg(dst).arg(&v).arg(&ii);
5578 unsafe {
5579 b.launch(cfg)?;
5580 }
5581 Ok(())
5582 }
5583
5584 pub fn i32_add_k(
5586 &self,
5587 d: &mut CudaSlice<i32>,
5588 v: i32,
5589 ) -> Result<(), Box<dyn std::error::Error>> {
5590 let f = self.func("i32_add_k");
5591 let cfg = LaunchConfig {
5592 grid_dim: (1, 1, 1),
5593 block_dim: (32, 1, 1),
5594 shared_mem_bytes: 0,
5595 };
5596 let __s_b = self.gpu.stream();
5597 let mut b = __s_b.launch_builder(&f);
5598 b.arg(d).arg(&v);
5599 unsafe {
5600 b.launch(cfg)?;
5601 }
5602 Ok(())
5603 }
5604
5605 pub fn i32_iota_from(
5607 &self,
5608 ctr: &CudaSlice<i32>,
5609 dst: &mut CudaSlice<i32>,
5610 n: usize,
5611 ) -> Result<(), Box<dyn std::error::Error>> {
5612 let f = self.func("i32_iota_from");
5613 let cfg = LaunchConfig::for_num_elems(n as u32);
5614 let ni = n as i32;
5615 let __s_b = self.gpu.stream();
5616 let mut b = __s_b.launch_builder(&f);
5617 b.arg(ctr).arg(dst).arg(&ni);
5618 unsafe {
5619 b.launch(cfg)?;
5620 }
5621 Ok(())
5622 }
5623
5624 pub fn u32_map_k(
5626 &self,
5627 buf: &mut CudaSlice<u32>,
5628 map: &CudaSlice<u32>,
5629 idx: usize,
5630 ) -> Result<(), Box<dyn std::error::Error>> {
5631 let f = self.func("u32_map_k");
5632 let cfg = LaunchConfig {
5633 grid_dim: (1, 1, 1),
5634 block_dim: (1, 1, 1),
5635 shared_mem_bytes: 0,
5636 };
5637 let ii = idx as i32;
5638 let __s_b = self.gpu.stream();
5639 let mut b = __s_b.launch_builder(&f);
5640 b.arg(buf).arg(map).arg(&ii);
5641 unsafe {
5642 b.launch(cfg)?;
5643 }
5644 Ok(())
5645 }
5646
5647 #[allow(clippy::too_many_arguments)]
5649 pub fn u32_pack2(
5650 &self,
5651 a: &CudaSlice<u32>,
5652 off_a: usize,
5653 n1: usize,
5654 b_in: &CudaSlice<u32>,
5655 n2: usize,
5656 out: &mut CudaSlice<u32>,
5657 ) -> Result<(), Box<dyn std::error::Error>> {
5658 let f = self.func("u32_pack2");
5659 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5660 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5661 let __s_b = self.gpu.stream();
5662 let mut b = __s_b.launch_builder(&f);
5663 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5664 unsafe {
5665 b.launch(cfg)?;
5666 }
5667 Ok(())
5668 }
5669
5670 pub fn moe_w_exscale(
5672 &self,
5673 w: &mut CudaSlice<f32>,
5674 sel: &CudaSlice<i32>,
5675 s: &CudaSlice<f32>,
5676 n: usize,
5677 ) -> Result<(), Box<dyn std::error::Error>> {
5678 let f = self.func("moe_w_exscale");
5679 let cfg = LaunchConfig::for_num_elems(n as u32);
5680 let ni = n as i32;
5681 let __s_b = self.gpu.stream();
5682 let mut b = __s_b.launch_builder(&f);
5683 b.arg(w).arg(sel).arg(s).arg(&ni);
5684 unsafe {
5685 b.launch(cfg)?;
5686 }
5687 Ok(())
5688 }
5689
5690 pub fn moe_w_scale_by_expert(
5693 &self,
5694 w: &mut CudaSlice<f32>,
5695 sel: &CudaSlice<i32>,
5696 macros: &CudaSlice<f32>,
5697 n_expert: usize,
5698 n: usize,
5699 ) -> Result<(), Box<dyn std::error::Error>> {
5700 let f = self.func("moe_w_scale_by_expert");
5701 let cfg = LaunchConfig {
5702 grid_dim: (n.div_ceil(64) as u32, 1, 1),
5703 block_dim: (64, 1, 1),
5704 shared_mem_bytes: 0,
5705 };
5706 let (ne, nn) = (n_expert as i32, n as i32);
5707 let __s_b = self.gpu.stream();
5708 let mut b = __s_b.launch_builder(&f);
5709 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5710 unsafe {
5711 b.launch(cfg)?;
5712 }
5713 Ok(())
5714 }
5715
5716 pub fn moe_gate_up_silu8_dev_q8(
5717 &self,
5718 table: &CudaSlice<u64>,
5719 sel: &cudarc::driver::CudaView<i32>,
5720 aq: &CudaSlice<i8>,
5721 ad: &CudaSlice<f32>,
5722 in_f: usize,
5723 n_ff: usize,
5724 n_used: usize,
5725 n_expert: usize,
5726 qt_g: i32,
5727 qt_u: i32,
5728 rb_g: usize,
5729 rb_u: usize,
5730 macros: &CudaSlice<f32>,
5731 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5732 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5733 let (mode, wpb) = GU.get_or_init(|| {
5734 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5735 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5736 .ok()
5737 .and_then(|v| v.parse().ok())
5738 .unwrap_or(4u32)
5739 .clamp(1, 16);
5740 (mode, wpb)
5741 });
5742 let (mode, wpb) = (mode.as_str(), *wpb);
5743 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5744 let (inf, nff, ne, rbg, rbu) = (
5745 in_f as i32,
5746 n_ff as i32,
5747 n_expert as i32,
5748 rb_g as i64,
5749 rb_u as i64,
5750 );
5751 let (f, cfg) = match mode {
5752 "1" | "2" | "4" => {
5753 let rpw: u32 = mode.parse().unwrap();
5754 let f = self.func(match rpw {
5755 1 => "moe_gate_up_silu8_dev_q8_r1",
5756 2 => "moe_gate_up_silu8_dev_q8_r2",
5757 _ => "moe_gate_up_silu8_dev_q8_r4",
5758 });
5759 let rows_per_block = (rpw * wpb) as usize;
5760 let gx = n_ff.div_ceil(rows_per_block) as u32;
5761 (
5762 f,
5763 LaunchConfig {
5764 grid_dim: (gx, n_used as u32, 1),
5765 block_dim: (32, wpb, 1),
5766 shared_mem_bytes: 0,
5767 },
5768 )
5769 }
5770 "j8" if n_used <= 32 => (
5771 self.func("moe_gate_up_silu8_dev_q8_j8"),
5772 LaunchConfig {
5773 grid_dim: (n_ff as u32, 1, 1),
5774 block_dim: (32, n_used as u32, 1),
5775 shared_mem_bytes: 0,
5776 },
5777 ),
5778 "vsm2" => {
5780 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5781 let sh = (rb_g + rb_u) as u32;
5782 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5783 f.set_attribute(
5784 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5785 sh as i32,
5786 )?;
5787 (
5788 f,
5789 LaunchConfig {
5790 grid_dim: (n_ff as u32, n_used as u32, 1),
5791 block_dim: (32, 1, 1),
5792 shared_mem_bytes: sh,
5793 },
5794 )
5795 }
5796 "vsm" => {
5797 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5798 let sh = (rb_g + rb_u) as u32;
5799 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5800 f.set_attribute(
5801 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5802 sh as i32,
5803 )?;
5804 (
5805 f,
5806 LaunchConfig {
5807 grid_dim: (n_ff as u32, n_used as u32, 1),
5808 block_dim: (32, 1, 1),
5809 shared_mem_bytes: sh,
5810 },
5811 )
5812 }
5813 "sg" => (
5814 self.func("moe_gate_up_silu8_dev_q8_sg"),
5815 LaunchConfig {
5816 grid_dim: (n_ff as u32, n_used as u32, 1),
5817 block_dim: (32, 1, 1),
5818 shared_mem_bytes: 0,
5819 },
5820 ),
5821 "j8sg" if n_used <= 32 => (
5822 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5823 LaunchConfig {
5824 grid_dim: (n_ff as u32, 1, 1),
5825 block_dim: (32, n_used as u32, 1),
5826 shared_mem_bytes: 0,
5827 },
5828 ),
5829 "u64" if in_f == 2048 => (
5830 self.func("moe_gate_up_silu8_dev_q8_u64"),
5831 LaunchConfig {
5832 grid_dim: (n_ff as u32, n_used as u32, 1),
5833 block_dim: (32, 1, 1),
5834 shared_mem_bytes: 0,
5835 },
5836 ),
5837 "gs4" if in_f == 2048 => (
5838 self.func("moe_gate_up_silu8_dev_q8_gs4"),
5839 LaunchConfig {
5840 grid_dim: (n_ff as u32, n_used as u32, 1),
5841 block_dim: (32, 4, 1),
5842 shared_mem_bytes: 0,
5843 },
5844 ),
5845 "v" | "" => (
5847 self.func("moe_gate_up_silu8_dev_q8_v"),
5848 LaunchConfig {
5849 grid_dim: (n_ff as u32, n_used as u32, 1),
5850 block_dim: (32, 1, 1),
5851 shared_mem_bytes: 0,
5852 },
5853 ),
5854 "s2" => (
5855 self.func("moe_gate_up_silu8_dev_q8_s2"),
5856 LaunchConfig {
5857 grid_dim: (n_ff as u32, n_used as u32, 1),
5858 block_dim: (32, 2, 1),
5859 shared_mem_bytes: 0,
5860 },
5861 ),
5862 "s2z" => {
5863 let rz = wpb.min(16); (
5865 self.func("moe_gate_up_silu8_dev_q8_s2z"),
5866 LaunchConfig {
5867 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5868 block_dim: (32, 2, rz),
5869 shared_mem_bytes: 0,
5870 },
5871 )
5872 }
5873 _ => (
5874 self.func("moe_gate_up_silu8_dev_q8"),
5875 LaunchConfig {
5876 grid_dim: (n_ff as u32, n_used as u32, 1),
5877 block_dim: (32, 1, 1),
5878 shared_mem_bytes: 0,
5879 },
5880 ),
5881 };
5882 let __s_b = self.gpu.stream();
5883 let mut b = __s_b.launch_builder(&f);
5884 b.arg(table)
5885 .arg(sel)
5886 .arg(aq)
5887 .arg(ad)
5888 .arg(&mut act)
5889 .arg(&inf)
5890 .arg(&nff)
5891 .arg(&ne)
5892 .arg(&qt_g)
5893 .arg(&qt_u)
5894 .arg(&rbg)
5895 .arg(&rbu)
5896 .arg(macros);
5897 unsafe {
5898 b.launch(cfg)?;
5899 }
5900 Ok(act)
5901 }
5902
5903 #[allow(clippy::too_many_arguments)]
5904 pub fn moe_down8_fma_dev_q8(
5905 &self,
5906 table: &CudaSlice<u64>,
5907 sel: &cudarc::driver::CudaView<i32>,
5908 w: &cudarc::driver::CudaView<f32>,
5909 aq2: &CudaSlice<i8>,
5910 ad2: &CudaSlice<f32>,
5911 dst: &mut cudarc::driver::CudaViewMut<f32>,
5912 in_f: usize,
5913 out_f: usize,
5914 n_used: usize,
5915 n_expert: usize,
5916 qt: i32,
5917 rb: usize,
5918 ) -> Result<(), Box<dyn std::error::Error>> {
5919 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5920 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5921 let (inf, outf, nu, ne, rbi) = (
5922 in_f as i32,
5923 out_f as i32,
5924 n_used as i32,
5925 n_expert as i32,
5926 rb as i64,
5927 );
5928 let (f, cfg) = match mode.as_str() {
5931 m @ ("1" | "2" | "4") if n_used <= 8 => {
5932 let rpw: usize = m.parse().unwrap();
5933 let f = self.func(match rpw {
5934 1 => "moe_down8_fma_dev_q8_w8r1",
5935 2 => "moe_down8_fma_dev_q8_w8r2",
5936 _ => "moe_down8_fma_dev_q8_w8r4",
5937 });
5938 (
5939 f,
5940 LaunchConfig {
5941 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5942 block_dim: (32, n_used as u32, 1),
5943 shared_mem_bytes: 0,
5944 },
5945 )
5946 }
5947 "h2" if in_f == 512 => (
5948 self.func("moe_down8_fma_dev_q8_h2"),
5949 LaunchConfig {
5950 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5951 block_dim: (32, 1, 1),
5952 shared_mem_bytes: 0,
5953 },
5954 ),
5955 "" if in_f == 704 && n_used <= 8 => (
5958 self.func("moe_down8_fma_dev_q8_w8r2"),
5959 LaunchConfig {
5960 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5961 block_dim: (32, n_used as u32, 1),
5962 shared_mem_bytes: 0,
5963 },
5964 ),
5965 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
5969 self.func("moe_down8_fma_dev_q8_w8h2v"),
5970 LaunchConfig {
5971 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5972 block_dim: (32, n_used as u32, 1),
5973 shared_mem_bytes: 0,
5974 },
5975 ),
5976 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
5977 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
5978 LaunchConfig {
5979 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5980 block_dim: (32, n_used as u32, 1),
5981 shared_mem_bytes: 0,
5982 },
5983 ),
5984 "w8h2r2" if in_f == 512 && n_used <= 8 => (
5985 self.func("moe_down8_fma_dev_q8_w8h2r2"),
5986 LaunchConfig {
5987 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5988 block_dim: (32, n_used as u32, 1),
5989 shared_mem_bytes: 0,
5990 },
5991 ),
5992 "w8h2" if in_f == 512 && n_used <= 8 => (
5993 self.func("moe_down8_fma_dev_q8_w8h2"),
5994 LaunchConfig {
5995 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5996 block_dim: (32, n_used as u32, 1),
5997 shared_mem_bytes: 0,
5998 },
5999 ),
6000 _ => (
6001 self.func("moe_down8_fma_dev_q8"),
6002 LaunchConfig {
6003 grid_dim: (out_f as u32, 1, 1),
6004 block_dim: (32, 1, 1),
6005 shared_mem_bytes: 0,
6006 },
6007 ),
6008 };
6009 let __s_b = self.gpu.stream();
6010 let mut b = __s_b.launch_builder(&f);
6011 b.arg(table)
6012 .arg(sel)
6013 .arg(w)
6014 .arg(aq2)
6015 .arg(ad2)
6016 .arg(dst)
6017 .arg(&inf)
6018 .arg(&outf)
6019 .arg(&nu)
6020 .arg(&ne)
6021 .arg(&qt)
6022 .arg(&rbi);
6023 unsafe {
6024 b.launch(cfg)?;
6025 }
6026 Ok(())
6027 }
6028
6029 #[allow(clippy::too_many_arguments)]
6036 pub fn moe_gate_up_silu8_dev_q8_rows(
6037 &self,
6038 table: &CudaSlice<u64>,
6039 sel: &CudaSlice<i32>,
6040 aq: &CudaSlice<i8>,
6041 ad: &CudaSlice<f32>,
6042 t: usize,
6043 in_f: usize,
6044 n_ff: usize,
6045 n_used: usize,
6046 n_expert: usize,
6047 qt_g: i32,
6048 qt_u: i32,
6049 rb_g: usize,
6050 rb_u: usize,
6051 macros: &CudaSlice<f32>,
6052 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6053 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6054 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6055 let cfg = LaunchConfig {
6056 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6057 block_dim: (32, 1, 1),
6058 shared_mem_bytes: 0,
6059 };
6060 let (inf, nff, ne, nu, rbg, rbu) = (
6061 in_f as i32,
6062 n_ff as i32,
6063 n_expert as i32,
6064 n_used as i32,
6065 rb_g as i64,
6066 rb_u as i64,
6067 );
6068 let __s_b = self.gpu.stream();
6069 let mut b = __s_b.launch_builder(&f);
6070 b.arg(table)
6071 .arg(sel)
6072 .arg(aq)
6073 .arg(ad)
6074 .arg(&mut act)
6075 .arg(&inf)
6076 .arg(&nff)
6077 .arg(&ne)
6078 .arg(&qt_g)
6079 .arg(&qt_u)
6080 .arg(&rbg)
6081 .arg(&rbu)
6082 .arg(&nu)
6083 .arg(macros);
6084 unsafe {
6085 b.launch(cfg)?;
6086 }
6087 Ok(act)
6088 }
6089
6090 #[allow(clippy::too_many_arguments)]
6095 pub fn moe_down8_fma_dev_q8_rows(
6096 &self,
6097 table: &CudaSlice<u64>,
6098 sel: &CudaSlice<i32>,
6099 w: &CudaSlice<f32>,
6100 aq2: &CudaSlice<i8>,
6101 ad2: &CudaSlice<f32>,
6102 dst: &mut CudaSlice<f32>,
6103 t: usize,
6104 in_f: usize,
6105 out_f: usize,
6106 n_used: usize,
6107 n_expert: usize,
6108 qt: i32,
6109 rb: usize,
6110 ) -> Result<(), Box<dyn std::error::Error>> {
6111 assert!(
6112 in_f == 512 && n_used <= 8,
6113 "down rows twin is w8h2v shape-gated"
6114 );
6115 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6116 let cfg = LaunchConfig {
6117 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6118 block_dim: (32, n_used as u32, 1),
6119 shared_mem_bytes: 0,
6120 };
6121 let (inf, outf, nu, ne, rbi) = (
6122 in_f as i32,
6123 out_f as i32,
6124 n_used as i32,
6125 n_expert as i32,
6126 rb as i64,
6127 );
6128 let __s_b = self.gpu.stream();
6129 let mut b = __s_b.launch_builder(&f);
6130 b.arg(table)
6131 .arg(sel)
6132 .arg(w)
6133 .arg(aq2)
6134 .arg(ad2)
6135 .arg(dst)
6136 .arg(&inf)
6137 .arg(&outf)
6138 .arg(&nu)
6139 .arg(&ne)
6140 .arg(&qt)
6141 .arg(&rbi);
6142 unsafe {
6143 b.launch(cfg)?;
6144 }
6145 Ok(())
6146 }
6147
6148 #[allow(clippy::too_many_arguments)]
6152 pub fn moe_gate_up_silu8_dev_q8_csr(
6153 &self,
6154 table: &CudaSlice<u64>,
6155 sel: &CudaSlice<i32>,
6156 aq: &CudaSlice<i8>,
6157 ad: &CudaSlice<f32>,
6158 n_pairs: usize,
6159 in_f: usize,
6160 n_ff: usize,
6161 n_used: usize,
6162 n_expert: usize,
6163 qt_g: i32,
6164 qt_u: i32,
6165 rb_g: usize,
6166 rb_u: usize,
6167 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6168 let f = if qt_g == crate::QT_NVFP4 {
6171 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6172 } else {
6173 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6174 };
6175 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6176 let cfg = LaunchConfig {
6177 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6178 block_dim: (32, 1, 1),
6179 shared_mem_bytes: 0,
6180 };
6181 let (inf, nff, ne, nu, npi, rbg, rbu) = (
6182 in_f as i32,
6183 n_ff as i32,
6184 n_expert as i32,
6185 n_used as i32,
6186 n_pairs as i32,
6187 rb_g as i64,
6188 rb_u as i64,
6189 );
6190 let __s_b = self.gpu.stream();
6191 let mut b = __s_b.launch_builder(&f);
6192 b.arg(table)
6193 .arg(sel)
6194 .arg(aq)
6195 .arg(ad)
6196 .arg(&mut act)
6197 .arg(&inf)
6198 .arg(&nff)
6199 .arg(&ne)
6200 .arg(&qt_g)
6201 .arg(&qt_u)
6202 .arg(&rbg)
6203 .arg(&rbu)
6204 .arg(&nu)
6205 .arg(&npi);
6206 unsafe {
6207 b.launch(cfg)?;
6208 }
6209 Ok(act)
6210 }
6211
6212 #[allow(clippy::too_many_arguments)]
6216 pub fn moe_down8_fma_dev_q8_variant(
6217 &self,
6218 variant: &str,
6219 table: &CudaSlice<u64>,
6220 sel: &cudarc::driver::CudaView<i32>,
6221 w: &cudarc::driver::CudaView<f32>,
6222 aq2: &CudaSlice<i8>,
6223 ad2: &CudaSlice<f32>,
6224 dst: &mut cudarc::driver::CudaViewMut<f32>,
6225 in_f: usize,
6226 out_f: usize,
6227 n_used: usize,
6228 n_expert: usize,
6229 qt: i32,
6230 rb: usize,
6231 ) -> Result<(), Box<dyn std::error::Error>> {
6232 let (inf, outf, nu, ne, rbi) = (
6233 in_f as i32,
6234 out_f as i32,
6235 n_used as i32,
6236 n_expert as i32,
6237 rb as i64,
6238 );
6239 let (f, cfg) = match variant {
6240 "w8h2" | "w8h2v" => (
6241 self.func(if variant == "w8h2" {
6242 "moe_down8_fma_dev_q8_w8h2"
6243 } else {
6244 "moe_down8_fma_dev_q8_w8h2v"
6245 }),
6246 LaunchConfig {
6247 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6248 block_dim: (32, n_used as u32, 1),
6249 shared_mem_bytes: 0,
6250 },
6251 ),
6252 "w8h2r2" | "w8h2r2v" => (
6253 self.func(if variant == "w8h2r2" {
6254 "moe_down8_fma_dev_q8_w8h2r2"
6255 } else {
6256 "moe_down8_fma_dev_q8_w8h2r2v"
6257 }),
6258 LaunchConfig {
6259 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6260 block_dim: (32, n_used as u32, 1),
6261 shared_mem_bytes: 0,
6262 },
6263 ),
6264 _ => (
6265 self.func("moe_down8_fma_dev_q8"),
6266 LaunchConfig {
6267 grid_dim: (out_f as u32, 1, 1),
6268 block_dim: (32, 1, 1),
6269 shared_mem_bytes: 0,
6270 },
6271 ),
6272 };
6273 let __s_b = self.gpu.stream();
6274 let mut b = __s_b.launch_builder(&f);
6275 b.arg(table)
6276 .arg(sel)
6277 .arg(w)
6278 .arg(aq2)
6279 .arg(ad2)
6280 .arg(dst)
6281 .arg(&inf)
6282 .arg(&outf)
6283 .arg(&nu)
6284 .arg(&ne)
6285 .arg(&qt)
6286 .arg(&rbi);
6287 unsafe {
6288 b.launch(cfg)?;
6289 }
6290 Ok(())
6291 }
6292
6293 #[allow(clippy::too_many_arguments)]
6295 pub fn moe_gate_up_silu8_dev_q8_variant(
6296 &self,
6297 variant: &str,
6298 table: &CudaSlice<u64>,
6299 sel: &cudarc::driver::CudaView<i32>,
6300 aq: &CudaSlice<i8>,
6301 ad: &CudaSlice<f32>,
6302 in_f: usize,
6303 n_ff: usize,
6304 n_used: usize,
6305 n_expert: usize,
6306 qt_g: i32,
6307 qt_u: i32,
6308 rb_g: usize,
6309 rb_u: usize,
6310 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6311 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6312 let (inf, nff, ne, rbg, rbu) = (
6313 in_f as i32,
6314 n_ff as i32,
6315 n_expert as i32,
6316 rb_g as i64,
6317 rb_u as i64,
6318 );
6319 let f = self.func(if variant == "v" {
6320 "moe_gate_up_silu8_dev_q8_v"
6321 } else {
6322 "moe_gate_up_silu8_dev_q8"
6323 });
6324 let cfg = LaunchConfig {
6325 grid_dim: (n_ff as u32, n_used as u32, 1),
6326 block_dim: (32, 1, 1),
6327 shared_mem_bytes: 0,
6328 };
6329 let __s_b = self.gpu.stream();
6330 let mut b = __s_b.launch_builder(&f);
6331 b.arg(table)
6332 .arg(sel)
6333 .arg(aq)
6334 .arg(ad)
6335 .arg(&mut act)
6336 .arg(&inf)
6337 .arg(&nff)
6338 .arg(&ne)
6339 .arg(&qt_g)
6340 .arg(&qt_u)
6341 .arg(&rbg)
6342 .arg(&rbu);
6343 unsafe {
6344 b.launch(cfg)?;
6345 }
6346 Ok(act)
6347 }
6348
6349 pub fn moe_gate_up_silu8_dev(
6350 &self,
6351 table: &CudaSlice<u64>,
6352 sel: &cudarc::driver::CudaView<i32>,
6353 x: &cudarc::driver::CudaView<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 macros: &CudaSlice<f32>,
6363 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6364 let f = self.func("moe_gate_up_silu8_dev");
6365 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6367 grid_dim: (n_ff as u32, n_used as u32, 1),
6368 block_dim: (256, 1, 1),
6369 shared_mem_bytes: 0,
6370 };
6371 let (inf, nff, ne, rbg, rbu) = (
6372 in_f as i32,
6373 n_ff as i32,
6374 n_expert as i32,
6375 rb_g as i64,
6376 rb_u as i64,
6377 );
6378 let __s_b = self.gpu.stream();
6379 let mut b = __s_b.launch_builder(&f);
6380 b.arg(table)
6381 .arg(sel)
6382 .arg(x)
6383 .arg(&mut act)
6384 .arg(&inf)
6385 .arg(&nff)
6386 .arg(&ne)
6387 .arg(&qt_g)
6388 .arg(&qt_u)
6389 .arg(&rbg)
6390 .arg(&rbu)
6391 .arg(macros);
6392 unsafe {
6393 b.launch(cfg)?;
6394 }
6395 Ok(act)
6396 }
6397
6398 #[allow(clippy::too_many_arguments)]
6401 pub fn moe_down8_fma_dev(
6402 &self,
6403 table: &CudaSlice<u64>,
6404 sel: &cudarc::driver::CudaView<i32>,
6405 w: &cudarc::driver::CudaView<f32>,
6406 act: &CudaSlice<f32>,
6407 dst: &mut cudarc::driver::CudaViewMut<f32>,
6408 in_f: usize,
6409 out_f: usize,
6410 n_used: usize,
6411 n_expert: usize,
6412 qt: i32,
6413 rb: usize,
6414 ) -> Result<(), Box<dyn std::error::Error>> {
6415 let f = self.func("moe_down8_fma_dev");
6416 let cfg = LaunchConfig {
6417 grid_dim: (out_f as u32, 1, 1),
6418 block_dim: (256, 1, 1),
6419 shared_mem_bytes: 0,
6420 };
6421 let (inf, outf, nu, ne, rbv) = (
6422 in_f as i32,
6423 out_f as i32,
6424 n_used as i32,
6425 n_expert as i32,
6426 rb as i64,
6427 );
6428 let __s_b = self.gpu.stream();
6429 let mut b = __s_b.launch_builder(&f);
6430 b.arg(table)
6431 .arg(sel)
6432 .arg(w)
6433 .arg(act)
6434 .arg(dst)
6435 .arg(&inf)
6436 .arg(&outf)
6437 .arg(&nu)
6438 .arg(&ne)
6439 .arg(&qt)
6440 .arg(&rbv);
6441 unsafe {
6442 b.launch(cfg)?;
6443 }
6444 Ok(())
6445 }
6446
6447 pub fn axpy_into(
6449 &self,
6450 src: &CudaSlice<f32>,
6451 alpha: f32,
6452 dst: &mut cudarc::driver::CudaViewMut<f32>,
6453 n: usize,
6454 ) -> Result<(), Box<dyn std::error::Error>> {
6455 let f = self.func("axpy_f32");
6456 let cfg = LaunchConfig::for_num_elems(n as u32);
6457 let (a, ni) = (alpha, n as i32);
6458 let __s_b = self.gpu.stream();
6459 let mut b = __s_b.launch_builder(&f);
6460 b.arg(src).arg(dst).arg(&a).arg(&ni);
6461 unsafe {
6462 b.launch(cfg)?;
6463 }
6464 Ok(())
6465 }
6466
6467 pub fn add_scaled_rows(
6469 &self,
6470 src: &CudaSlice<f32>,
6471 scale: &CudaSlice<f32>,
6472 dst: &mut CudaSlice<f32>,
6473 ncols: usize,
6474 nrows: usize,
6475 ) -> Result<(), Box<dyn std::error::Error>> {
6476 let f = self.func("add_scaled_rows_f32");
6477 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6478 let (nc, nr) = (ncols as i32, nrows as i32);
6479 let __s_b = self.gpu.stream();
6480 let mut b = __s_b.launch_builder(&f);
6481 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6482 unsafe {
6483 b.launch(cfg)?;
6484 }
6485 Ok(())
6486 }
6487
6488 pub fn gather_rows(
6492 &self,
6493 src: &CudaSlice<f32>,
6494 idx: &CudaSlice<i32>,
6495 dst: &mut CudaSlice<f32>,
6496 ncols: usize,
6497 m_e: usize,
6498 ) -> Result<(), Box<dyn std::error::Error>> {
6499 let f = self.func("gather_rows_f32");
6500 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6501 let (nc, me) = (ncols as i32, m_e as i32);
6502 let __s_b = self.gpu.stream();
6503 let mut b = __s_b.launch_builder(&f);
6504 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6505 unsafe {
6506 b.launch(cfg)?;
6507 }
6508 Ok(())
6509 }
6510
6511 pub fn scatter_slot(
6516 &self,
6517 src: &CudaSlice<f32>,
6518 tok_idx: &CudaSlice<i32>,
6519 slot_idx: &CudaSlice<i32>,
6520 weight: &CudaSlice<f32>,
6521 dst: &mut CudaSlice<f32>,
6522 wbuf: &mut CudaSlice<f32>,
6523 ncols: usize,
6524 n_used: usize,
6525 m_e: usize,
6526 ) -> Result<(), Box<dyn std::error::Error>> {
6527 let f = self.func("scatter_add_slot_f32");
6528 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6529 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6530 let __s_b = self.gpu.stream();
6531 let mut b = __s_b.launch_builder(&f);
6532 b.arg(src)
6533 .arg(tok_idx)
6534 .arg(slot_idx)
6535 .arg(weight)
6536 .arg(dst)
6537 .arg(wbuf)
6538 .arg(&nc)
6539 .arg(&nu)
6540 .arg(&me);
6541 unsafe {
6542 b.launch(cfg)?;
6543 }
6544 Ok(())
6545 }
6546
6547 pub fn reduce_slots(
6551 &self,
6552 slots: &CudaSlice<f32>,
6553 wbuf: &CudaSlice<f32>,
6554 dst: &mut CudaSlice<f32>,
6555 ncols: usize,
6556 n_used: usize,
6557 t: usize,
6558 ) -> Result<(), Box<dyn std::error::Error>> {
6559 let f = self.func("reduce_slots_f32");
6560 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6561 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6562 let __s_b = self.gpu.stream();
6563 let mut b = __s_b.launch_builder(&f);
6564 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6565 unsafe {
6566 b.launch(cfg)?;
6567 }
6568 Ok(())
6569 }
6570
6571 pub fn quantize_q8_1_view(
6578 &self,
6579 x: &cudarc::driver::CudaView<f32>,
6580 m: usize,
6581 in_f: usize,
6582 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6583 let f = self.func("quantize_q8_1");
6584 let nblk = in_f / 32;
6585 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6586 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6587 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6588 let (inf, mi) = (in_f as i32, m as i32);
6589 let __s_b = self.gpu.stream();
6590 let mut b = __s_b.launch_builder(&f);
6591 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6592 unsafe {
6593 b.launch(cfg)?;
6594 }
6595 Ok((q, d))
6596 }
6597
6598 pub fn quantize_q8_1(
6599 &self,
6600 x: &CudaSlice<f32>,
6601 m: usize,
6602 in_f: usize,
6603 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6604 let nblk = in_f / 32;
6605 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);
6609 let (inf, mi) = (in_f as i32, m as i32);
6610 if Self::pdl_on() && Self::pdl_wb_on() {
6611 {
6612 use cudarc::driver::{DevicePtr, DevicePtrMut};
6613 let s = &self.gpu.stream();
6614 let (px, _g0) = x.device_ptr(s);
6615 let (pq, _g1) = q.device_ptr_mut(s);
6616 let (pd, _g2) = d.device_ptr_mut(s);
6617 let mut ps = [
6618 &px as *const _ as *mut std::ffi::c_void,
6619 &pq as *const _ as *mut _,
6620 &pd as *const _ as *mut _,
6621 &inf as *const _ as *mut _,
6622 &mi as *const _ as *mut _,
6623 ];
6624 unsafe {
6625 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6626 }
6627 }
6628 return Ok((q, d));
6629 }
6630 let f = self.func("quantize_q8_1");
6631 let __s_b = self.gpu.stream();
6632 let mut b = __s_b.launch_builder(&f);
6633 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6634 unsafe {
6635 b.launch(cfg)?;
6636 }
6637 Ok((q, d))
6638 }
6639
6640 pub fn quantize_fp4_act(
6644 &self,
6645 x: &CudaSlice<f32>,
6646 m: usize,
6647 in_f: usize,
6648 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6649 let f = self.func("quantize_fp4_act");
6650 let nb16 = in_f / 16;
6651 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);
6654 let (inf, mi) = (in_f as i32, m as i32);
6655 let __s_b = self.gpu.stream();
6656 let mut b = __s_b.launch_builder(&f);
6657 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6658 unsafe {
6659 b.launch(cfg)?;
6660 }
6661 Ok((aq4, ad4))
6662 }
6663
6664 pub fn qmatvec_gemm_nvfp4_fp4(
6669 &self,
6670 bytes: &CudaSlice<u8>,
6671 x: &CudaSlice<f32>,
6672 m: usize,
6673 in_f: usize,
6674 out_f: usize,
6675 row_bytes: usize,
6676 scale: f32,
6677 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6678 assert!(
6679 in_f % 64 == 0,
6680 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6681 );
6682 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6683 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6684 if scale != 1.0 {
6685 self.scale_inplace(&mut y, scale, m * out_f)?;
6686 }
6687 Ok(y)
6688 }
6689
6690 fn fp4_gemm_launch(
6693 &self,
6694 bytes: &CudaSlice<u8>,
6695 aq4: &CudaSlice<u32>,
6696 ad4: &CudaSlice<u8>,
6697 m: usize,
6698 in_f: usize,
6699 out_f: usize,
6700 row_bytes: usize,
6701 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6702 let f = self.func("qmatvec_gemm_nvfp4_fp4");
6703 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
6705 const BN: u32 = 256;
6706 let cfg = LaunchConfig {
6707 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6708 block_dim: (32, 4, 1),
6709 shared_mem_bytes: 0,
6710 };
6711 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6712 let __s_b = self.gpu.stream();
6713 let mut b = __s_b.launch_builder(&f);
6714 b.arg(bytes)
6715 .arg(aq4)
6716 .arg(ad4)
6717 .arg(&mut y)
6718 .arg(&inf)
6719 .arg(&outf)
6720 .arg(&mi)
6721 .arg(&rb);
6722 unsafe {
6723 b.launch(cfg)?;
6724 }
6725 Ok(y)
6726 }
6727
6728 pub fn qmatvec_gemm_nvfp4_fp4_raw(
6730 &self,
6731 bytes: &CudaSlice<u8>,
6732 x: &CudaSlice<f32>,
6733 m: usize,
6734 in_f: usize,
6735 out_f: usize,
6736 row_bytes: usize,
6737 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6738 assert!(
6739 in_f % 64 == 0,
6740 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6741 );
6742 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6743 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6744 }
6745
6746 pub fn qmatvec_q8_0_fast(
6748 &self,
6749 w: &CudaSlice<u8>,
6750 x: &CudaSlice<f32>,
6751 m: usize,
6752 in_f: usize,
6753 out_f: usize,
6754 row_bytes: usize,
6755 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6756 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6757 let f = self.func("qmatvec_q8_0_dp4a");
6758 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6760 grid_dim: (out_f as u32, m as u32, 1),
6761 block_dim: (128, 1, 1),
6762 shared_mem_bytes: 0,
6763 };
6764 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6765 let __s_b = self.gpu.stream();
6766 let mut b = __s_b.launch_builder(&f);
6767 b.arg(w)
6768 .arg(&aq)
6769 .arg(&ad)
6770 .arg(&mut y)
6771 .arg(&inf)
6772 .arg(&outf)
6773 .arg(&mi)
6774 .arg(&rb);
6775 unsafe {
6776 b.launch(cfg)?;
6777 }
6778 Ok(y)
6779 }
6780
6781 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
6784 &self,
6785 w: &CudaSlice<u8>,
6786 x: &CudaSlice<f32>,
6787 m: usize,
6788 in_f: usize,
6789 out_f: usize,
6790 row_bytes: usize,
6791 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6792 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6793 let f = self.func("qmatvec_q4_K_dp4a");
6794 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6796 grid_dim: (out_f as u32, m as u32, 1),
6797 block_dim: (128, 1, 1),
6798 shared_mem_bytes: 0,
6799 };
6800 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6801 let __s_b = self.gpu.stream();
6802 let mut b = __s_b.launch_builder(&f);
6803 b.arg(w)
6804 .arg(&aq)
6805 .arg(&ad)
6806 .arg(&mut y)
6807 .arg(&inf)
6808 .arg(&outf)
6809 .arg(&mi)
6810 .arg(&rb);
6811 unsafe {
6812 b.launch(cfg)?;
6813 }
6814 Ok(y)
6815 }
6816
6817 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
6820 &self,
6821 w: &CudaSlice<u8>,
6822 x: &CudaSlice<f32>,
6823 m: usize,
6824 in_f: usize,
6825 out_f: usize,
6826 row_bytes: usize,
6827 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6828 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6829 let f = self.func("qmatvec_q6_K_dp4a");
6830 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6832 grid_dim: (out_f as u32, m as u32, 1),
6833 block_dim: (128, 1, 1),
6834 shared_mem_bytes: 0,
6835 };
6836 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6837 let __s_b = self.gpu.stream();
6838 let mut b = __s_b.launch_builder(&f);
6839 b.arg(w)
6840 .arg(&aq)
6841 .arg(&ad)
6842 .arg(&mut y)
6843 .arg(&inf)
6844 .arg(&outf)
6845 .arg(&mi)
6846 .arg(&rb);
6847 unsafe {
6848 b.launch(cfg)?;
6849 }
6850 Ok(y)
6851 }
6852
6853 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
6856 &self,
6857 w: &CudaSlice<u8>,
6858 x: &CudaSlice<f32>,
6859 m: usize,
6860 in_f: usize,
6861 out_f: usize,
6862 row_bytes: usize,
6863 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6864 self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6865 }
6866 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
6869 &self,
6870 w: &CudaSlice<u8>,
6871 x: &CudaSlice<f32>,
6872 m: usize,
6873 in_f: usize,
6874 out_f: usize,
6875 row_bytes: usize,
6876 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6877 self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6878 }
6879 pub fn qmatvec_nvfp4_fast_rp(
6881 &self,
6882 w: &CudaSlice<u8>,
6883 x: &CudaSlice<f32>,
6884 m: usize,
6885 in_f: usize,
6886 out_f: usize,
6887 row_bytes: usize,
6888 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6889 assert!(
6890 in_f % 64 == 0,
6891 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6892 );
6893 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6894 }
6895 pub fn qmatvec_nvfp4_fast(
6897 &self,
6898 w: &CudaSlice<u8>,
6899 x: &CudaSlice<f32>,
6900 m: usize,
6901 in_f: usize,
6902 out_f: usize,
6903 row_bytes: usize,
6904 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6905 assert!(
6908 in_f % 64 == 0,
6909 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6910 );
6911 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6912 }
6913 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
6916 &self,
6917 w: &CudaSlice<u8>,
6918 x: &CudaSlice<f32>,
6919 m: usize,
6920 in_f: usize,
6921 out_f: usize,
6922 row_bytes: usize,
6923 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6924 self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6925 }
6926
6927 fn qmatvec_dp4a_named(
6929 &self,
6930 name: &str,
6931 w: &CudaSlice<u8>,
6932 x: &CudaSlice<f32>,
6933 m: usize,
6934 in_f: usize,
6935 out_f: usize,
6936 row_bytes: usize,
6937 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6938 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6939 let f = self.func(name);
6940 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6942 grid_dim: (out_f as u32, m as u32, 1),
6943 block_dim: (128, 1, 1),
6944 shared_mem_bytes: 0,
6945 };
6946 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6947 let __s_b = self.gpu.stream();
6948 let mut b = __s_b.launch_builder(&f);
6949 b.arg(w)
6950 .arg(&aq)
6951 .arg(&ad)
6952 .arg(&mut y)
6953 .arg(&inf)
6954 .arg(&outf)
6955 .arg(&mi)
6956 .arg(&rb);
6957 unsafe {
6958 b.launch(cfg)?;
6959 }
6960 Ok(y)
6961 }
6962
6963 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6964 Ok(self.gpu.stream().clone_htod(v)?)
6965 }
6966 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6967 Ok(self.gpu.stream().clone_htod(v)?)
6968 }
6969 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6971 Ok(self.gpu.stream().clone_htod(v)?)
6972 }
6973 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
6974 Ok(self.gpu.stream().clone_htod(v)?)
6975 }
6976 pub fn dtoh_view(
6978 &self,
6979 d: &cudarc::driver::CudaView<f32>,
6980 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6981 let v = self.gpu.stream().clone_dtoh(d)?;
6982 self.gpu.stream().synchronize()?;
6983 Ok(v)
6984 }
6985 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6986 let v = self.gpu.stream().clone_dtoh(d)?;
6987 self.gpu.stream().synchronize()?;
6988 Ok(v)
6989 }
6990 pub fn dtoh_pair(
6994 &self,
6995 a: &CudaSlice<f32>,
6996 b: &CudaSlice<f32>,
6997 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
6998 let av = self.gpu.stream().clone_dtoh(a)?;
6999 let bv = self.gpu.stream().clone_dtoh(b)?;
7000 self.gpu.stream().synchronize()?;
7001 Ok((av, bv))
7002 }
7003 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
7005 let v = self.gpu.stream().clone_dtoh(d)?;
7006 self.gpu.stream().synchronize()?;
7007 Ok(v)
7008 }
7009 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7011 let v = self.gpu.stream().clone_dtoh(d)?;
7012 self.gpu.stream().synchronize()?;
7013 Ok(v)
7014 }
7015 pub fn dtoh_u8_view(
7016 &self,
7017 d: &cudarc::driver::CudaView<u8>,
7018 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7019 let v = self.gpu.stream().clone_dtoh(d)?;
7020 self.gpu.stream().synchronize()?;
7021 Ok(v)
7022 }
7023 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7024 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
7025 self.keep_if_capturing(&s);
7026 Ok(s)
7027 }
7028
7029 pub fn prob_of_token_device(
7038 &self,
7039 logits: &CudaSlice<f32>,
7040 tok: &CudaSlice<u32>,
7041 n_vocab: usize,
7042 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7043 let nb = ARGMAX_NB;
7044 let mut part = self.alloc_uninit::<f32>(nb)?;
7045 let mut p = self.alloc_uninit::<f32>(1)?;
7046 let f1 = self.func("prob_of_token_partial_f32");
7047 let cfg1 = LaunchConfig {
7048 grid_dim: (nb as u32, 1, 1),
7049 block_dim: (256, 1, 1),
7050 shared_mem_bytes: 0,
7051 };
7052 let nv = n_vocab as i32;
7053 let __s_b1 = self.gpu.stream();
7054 let mut b1 = __s_b1.launch_builder(&f1);
7055 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7056 unsafe {
7057 b1.launch(cfg1)?;
7058 }
7059 let f2 = self.func("prob_of_token_final_f32");
7060 let cfg2 = LaunchConfig {
7061 grid_dim: (1, 1, 1),
7062 block_dim: (256, 1, 1),
7063 shared_mem_bytes: 0,
7064 };
7065 let nbi = nb as i32;
7066 let __s_b2 = self.gpu.stream();
7067 let mut b2 = __s_b2.launch_builder(&f2);
7068 b2.arg(&part).arg(&mut p).arg(&nbi);
7069 unsafe {
7070 b2.launch(cfg2)?;
7071 }
7072 Ok(p)
7073 }
7074
7075 pub fn prob_of_token_device_col(
7082 &self,
7083 logits: &CudaSlice<f32>,
7084 tok_all: &CudaSlice<u32>,
7085 tok_idx: usize,
7086 p_out: &mut CudaSlice<f32>,
7087 p_idx: usize,
7088 n_vocab: usize,
7089 ) -> Result<(), Box<dyn std::error::Error>> {
7090 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
7091 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
7092 let nb = ARGMAX_NB;
7093 let mut part = self.alloc_uninit::<f32>(nb)?;
7094 let f1 = self.func("prob_of_token_partial_f32");
7095 let cfg1 = LaunchConfig {
7096 grid_dim: (nb as u32, 1, 1),
7097 block_dim: (256, 1, 1),
7098 shared_mem_bytes: 0,
7099 };
7100 let nv = n_vocab as i32;
7101 let __s_b1 = self.gpu.stream();
7102 let mut b1 = __s_b1.launch_builder(&f1);
7103 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
7104 unsafe {
7105 b1.launch(cfg1)?;
7106 }
7107 let f2 = self.func("prob_of_token_final_f32");
7108 let cfg2 = LaunchConfig {
7109 grid_dim: (1, 1, 1),
7110 block_dim: (256, 1, 1),
7111 shared_mem_bytes: 0,
7112 };
7113 let nbi = nb as i32;
7114 let __s_b2 = self.gpu.stream();
7115 let mut b2 = __s_b2.launch_builder(&f2);
7116 b2.arg(&part).arg(&mut p_v).arg(&nbi);
7117 unsafe {
7118 b2.launch(cfg2)?;
7119 }
7120 Ok(())
7121 }
7122
7123 pub fn prob_of_token_device_into(
7124 &self,
7125 logits: &CudaSlice<f32>,
7126 tok: &CudaSlice<u32>,
7127 p_out: &mut CudaSlice<f32>,
7128 n_vocab: usize,
7129 ) -> Result<(), Box<dyn std::error::Error>> {
7130 let nb = ARGMAX_NB;
7131 let mut part = self.alloc_uninit::<f32>(nb)?;
7132 let f1 = self.func("prob_of_token_partial_f32");
7133 let cfg1 = LaunchConfig {
7134 grid_dim: (nb as u32, 1, 1),
7135 block_dim: (256, 1, 1),
7136 shared_mem_bytes: 0,
7137 };
7138 let nv = n_vocab as i32;
7139 let __s_b1 = self.gpu.stream();
7140 let mut b1 = __s_b1.launch_builder(&f1);
7141 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7142 unsafe {
7143 b1.launch(cfg1)?;
7144 }
7145 let f2 = self.func("prob_of_token_final_f32");
7146 let cfg2 = LaunchConfig {
7147 grid_dim: (1, 1, 1),
7148 block_dim: (256, 1, 1),
7149 shared_mem_bytes: 0,
7150 };
7151 let nbi = nb as i32;
7152 let __s_b2 = self.gpu.stream();
7153 let mut b2 = __s_b2.launch_builder(&f2);
7154 b2.arg(&part).arg(p_out).arg(&nbi);
7155 unsafe {
7156 b2.launch(cfg2)?;
7157 }
7158 Ok(())
7159 }
7160
7161 pub fn argmax_token_device(
7162 &self,
7163 logits: &CudaSlice<f32>,
7164 n_vocab: usize,
7165 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7166 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
7167 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
7168 Ok(tok)
7169 }
7170 pub fn argmax_token_device_into(
7177 &self,
7178 logits: &CudaSlice<f32>,
7179 tok: &mut CudaSlice<u32>,
7180 n_vocab: usize,
7181 ) -> Result<(), Box<dyn std::error::Error>> {
7182 let nb = ARGMAX_NB;
7183 let f1 = self.func("argmax_partial_f32");
7184 let f2 = self.func("argmax_final_f32");
7185 let mut guard = self.argmax_partials.lock().unwrap();
7186 if guard.is_none() {
7187 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7190 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7191 *guard = Some((pv, pi));
7192 }
7193 let (part_v, part_i) = guard.as_mut().unwrap();
7194 let nv = n_vocab as i32;
7195 let nbi = nb as i32;
7196 let cfg1 = LaunchConfig {
7198 grid_dim: (nb as u32, 1, 1),
7199 block_dim: (256, 1, 1),
7200 shared_mem_bytes: 0,
7201 };
7202 let __s_b1 = self.gpu.stream();
7203 let mut b1 = __s_b1.launch_builder(&f1);
7204 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
7205 unsafe {
7206 b1.launch(cfg1)?;
7207 }
7208 let cfg2 = LaunchConfig {
7210 grid_dim: (1, 1, 1),
7211 block_dim: (256, 1, 1),
7212 shared_mem_bytes: 0,
7213 };
7214 let __s_b2 = self.gpu.stream();
7215 let mut b2 = __s_b2.launch_builder(&f2);
7216 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
7217 unsafe {
7218 b2.launch(cfg2)?;
7219 }
7220 Ok(())
7221 }
7222 pub fn argmax_token_device_col(
7228 &self,
7229 logits: &CudaSlice<f32>,
7230 col: usize,
7231 n_vocab: usize,
7232 toks: &mut CudaSlice<u32>,
7233 out_idx: usize,
7234 ) -> Result<(), Box<dyn std::error::Error>> {
7235 let nb = ARGMAX_NB;
7236 let f1 = self.func("argmax_partial_f32");
7237 let f2 = self.func("argmax_final_f32");
7238 let mut guard = self.argmax_partials.lock().unwrap();
7239 if guard.is_none() {
7240 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7241 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7242 *guard = Some((pv, pi));
7243 }
7244 let (part_v, part_i) = guard.as_mut().unwrap();
7245 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
7246 let nv = n_vocab as i32;
7247 let nbi = nb as i32;
7248 let cfg1 = LaunchConfig {
7249 grid_dim: (nb as u32, 1, 1),
7250 block_dim: (256, 1, 1),
7251 shared_mem_bytes: 0,
7252 };
7253 let __s_b1 = self.gpu.stream();
7254 let mut b1 = __s_b1.launch_builder(&f1);
7255 b1.arg(&col_view)
7256 .arg(&mut *part_v)
7257 .arg(&mut *part_i)
7258 .arg(&nv);
7259 unsafe {
7260 b1.launch(cfg1)?;
7261 }
7262 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
7263 let cfg2 = LaunchConfig {
7264 grid_dim: (1, 1, 1),
7265 block_dim: (256, 1, 1),
7266 shared_mem_bytes: 0,
7267 };
7268 let __s_b2 = self.gpu.stream();
7269 let mut b2 = __s_b2.launch_builder(&f2);
7270 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
7271 unsafe {
7272 b2.launch(cfg2)?;
7273 }
7274 Ok(())
7275 }
7276 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7278 Ok(self.gpu.stream().clone_htod(v)?)
7279 }
7280 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7281 let v = self.gpu.stream().clone_dtoh(d)?;
7282 self.gpu.stream().synchronize()?;
7283 Ok(v)
7284 }
7285 pub fn htod_u32_into(
7289 &self,
7290 dst: &mut CudaSlice<u32>,
7291 src: &[u32],
7292 ) -> Result<(), Box<dyn std::error::Error>> {
7293 let mut view = dst.slice_mut(0..src.len());
7294 self.gpu.stream().memcpy_htod(src, &mut view)?;
7295 Ok(())
7296 }
7297
7298 pub fn htod_i32_into(
7301 &self,
7302 dst: &mut CudaSlice<i32>,
7303 src: &[i32],
7304 ) -> Result<(), Box<dyn std::error::Error>> {
7305 let mut view = dst.slice_mut(0..src.len());
7306 self.gpu.stream().memcpy_htod(src, &mut view)?;
7307 Ok(())
7308 }
7309
7310 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7311 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
7312 self.keep_if_capturing(&s);
7313 Ok(s)
7314 }
7315 pub fn embed_gather_device_into(
7318 &self,
7319 embd: &CudaSlice<u8>,
7320 token_d: &CudaSlice<u32>,
7321 x_out: &mut CudaSlice<f32>,
7322 n_embd: usize,
7323 qtype: i32,
7324 row_bytes: usize,
7325 ) -> Result<(), Box<dyn std::error::Error>> {
7326 let f = self.func("embed_gather_u32");
7327 let cfg = LaunchConfig {
7328 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7329 block_dim: (256, 1, 1),
7330 shared_mem_bytes: 0,
7331 };
7332 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7333 let __s_b = self.gpu.stream();
7334 let mut b = __s_b.launch_builder(&f);
7335 b.arg(embd)
7336 .arg(token_d)
7337 .arg(x_out)
7338 .arg(&ne)
7339 .arg(&qt)
7340 .arg(&rb);
7341 unsafe {
7342 b.launch(cfg)?;
7343 }
7344 Ok(())
7345 }
7346 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7348 let v = self.gpu.stream().clone_dtoh(d)?;
7349 self.gpu.stream().synchronize()?;
7350 Ok(v[0])
7351 }
7352 pub fn i32_set_k(
7359 &self,
7360 dst: &mut CudaSlice<i32>,
7361 v: i32,
7362 ) -> Result<(), Box<dyn std::error::Error>> {
7363 let f = self.func("i32_set_k");
7364 let cfg = LaunchConfig {
7365 grid_dim: (1, 1, 1),
7366 block_dim: (1, 1, 1),
7367 shared_mem_bytes: 0,
7368 };
7369 let idx = 0i32;
7370 let __s_b = self.gpu.stream();
7371 let mut b = __s_b.launch_builder(&f);
7372 b.arg(dst).arg(&v).arg(&idx);
7373 unsafe {
7374 b.launch(cfg)?;
7375 }
7376 Ok(())
7377 }
7378
7379 pub fn set_i32_one(
7380 &self,
7381 d: &mut CudaSlice<i32>,
7382 v: i32,
7383 ) -> Result<(), Box<dyn std::error::Error>> {
7384 self.gpu.stream().memcpy_htod(&[v], d)?;
7385 Ok(())
7386 }
7387 pub fn set_u32_one(
7390 &self,
7391 d: &mut CudaSlice<u32>,
7392 v: u32,
7393 ) -> Result<(), Box<dyn std::error::Error>> {
7394 self.gpu.stream().memcpy_htod(&[v], d)?;
7395 Ok(())
7396 }
7397 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7399 let v = self.gpu.stream().clone_dtoh(d)?;
7400 self.gpu.stream().synchronize()?;
7401 Ok(v[0])
7402 }
7403 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7405 Ok(self.gpu.stream().clone_htod(bytes)?)
7406 }
7407 pub fn embed_gather_device(
7411 &self,
7412 embd: &CudaSlice<u8>,
7413 token_d: &CudaSlice<u32>,
7414 n_embd: usize,
7415 qtype: i32,
7416 row_bytes: usize,
7417 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7418 let f = self.func("embed_gather_u32");
7419 let mut x = self.alloc_uninit::<f32>(n_embd)?;
7420 let cfg = LaunchConfig {
7421 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7422 block_dim: (256, 1, 1),
7423 shared_mem_bytes: 0,
7424 };
7425 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7426 let __s_b = self.gpu.stream();
7427 let mut b = __s_b.launch_builder(&f);
7428 b.arg(embd)
7429 .arg(token_d)
7430 .arg(&mut x)
7431 .arg(&ne)
7432 .arg(&qt)
7433 .arg(&rb);
7434 unsafe {
7435 b.launch(cfg)?;
7436 }
7437 Ok(x)
7438 }
7439
7440 pub fn embed_gather_device_t(
7444 &self,
7445 embd: &CudaSlice<u8>,
7446 tokens: &[u32],
7447 n_embd: usize,
7448 qtype: i32,
7449 row_bytes: usize,
7450 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7451 let t = tokens.len();
7452 let tok_d = self.gpu.stream().clone_htod(tokens)?;
7453 let f = self.func("embed_gather_u32_t");
7454 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7455 let cfg = LaunchConfig {
7456 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7457 block_dim: (256, 1, 1),
7458 shared_mem_bytes: 0,
7459 };
7460 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7461 let __s_b = self.gpu.stream();
7462 let mut b = __s_b.launch_builder(&f);
7463 b.arg(embd)
7464 .arg(&tok_d)
7465 .arg(&mut x)
7466 .arg(&ne)
7467 .arg(&qt)
7468 .arg(&rb)
7469 .arg(&ti);
7470 unsafe {
7471 b.launch(cfg)?;
7472 }
7473 Ok(x)
7474 }
7475
7476 pub fn embed_gather_device_tv(
7481 &self,
7482 embd: &CudaSlice<u8>,
7483 tok_v: &cudarc::driver::CudaView<u32>,
7484 t: usize,
7485 n_embd: usize,
7486 qtype: i32,
7487 row_bytes: usize,
7488 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7489 let f = self.func("embed_gather_u32_t");
7490 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7491 let cfg = LaunchConfig {
7492 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7493 block_dim: (256, 1, 1),
7494 shared_mem_bytes: 0,
7495 };
7496 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7497 let __s_b = self.gpu.stream();
7498 let mut b = __s_b.launch_builder(&f);
7499 b.arg(embd)
7500 .arg(tok_v)
7501 .arg(&mut x)
7502 .arg(&ne)
7503 .arg(&qt)
7504 .arg(&rb)
7505 .arg(&ti);
7506 unsafe {
7507 b.launch(cfg)?;
7508 }
7509 Ok(x)
7510 }
7511
7512 pub fn embed_gather_device_td(
7513 &self,
7514 embd: &CudaSlice<u8>,
7515 tok_d: &CudaSlice<u32>,
7516 t: usize,
7517 n_embd: usize,
7518 qtype: i32,
7519 row_bytes: usize,
7520 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7521 let f = self.func("embed_gather_u32_t");
7522 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7523 let cfg = LaunchConfig {
7524 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7525 block_dim: (256, 1, 1),
7526 shared_mem_bytes: 0,
7527 };
7528 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7529 let __s_b = self.gpu.stream();
7530 let mut b = __s_b.launch_builder(&f);
7531 b.arg(embd)
7532 .arg(tok_d)
7533 .arg(&mut x)
7534 .arg(&ne)
7535 .arg(&qt)
7536 .arg(&rb)
7537 .arg(&ti);
7538 unsafe {
7539 b.launch(cfg)?;
7540 }
7541 Ok(x)
7542 }
7543
7544 #[inline]
7550 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7552 if self
7553 .capture_keep_on
7554 .load(std::sync::atomic::Ordering::Relaxed)
7555 {
7556 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7557 }
7558 }
7559
7560 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7561 &self,
7562 n: usize,
7563 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7564 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7565 {
7569 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7570 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7571 use cudarc::driver::DevicePtrMut;
7573 let n_bytes = s.len() * std::mem::size_of::<T>();
7574 let stream = self.gpu.stream();
7575 let (p_, _g) = s.device_ptr_mut(&stream);
7576 unsafe {
7577 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7578 .result()?;
7579 }
7580 }
7581 }
7582 self.keep_if_capturing(&s);
7583 Ok(s)
7584 }
7585
7586 pub fn uninit_q8_pair(
7591 &self,
7592 n: usize,
7593 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7594 Ok((
7595 self.alloc_uninit::<i8>(n)?,
7596 self.alloc_uninit::<f32>(n / 32)?,
7597 ))
7598 }
7599
7600 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7601 self.alloc_uninit::<f32>(n)
7602 }
7603
7604 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7606 self.alloc_uninit::<i8>(n)
7607 }
7608
7609 #[allow(clippy::too_many_arguments)]
7613 pub fn rms_norm3(
7614 &self,
7615 x: &CudaSlice<f32>,
7616 w0: &CudaSlice<f32>,
7617 w1: &CudaSlice<f32>,
7618 w2: &CudaSlice<f32>,
7619 d0: &mut CudaSlice<f32>,
7620 d1: &mut CudaSlice<f32>,
7621 d2: &mut CudaSlice<f32>,
7622 ncols: usize,
7623 nrows: usize,
7624 eps: f32,
7625 ) -> Result<(), Box<dyn std::error::Error>> {
7626 let f = self.func("rms_norm3_f32");
7627 let cfg = LaunchConfig {
7628 grid_dim: (nrows as u32, 1, 1),
7629 block_dim: (rms_block(), 1, 1),
7630 shared_mem_bytes: 0,
7631 };
7632 let (nc, e) = (ncols as i32, eps);
7633 let __s_b = self.gpu.stream();
7634 let mut b = __s_b.launch_builder(&f);
7635 b.arg(x)
7636 .arg(w0)
7637 .arg(w1)
7638 .arg(w2)
7639 .arg(d0)
7640 .arg(d1)
7641 .arg(d2)
7642 .arg(&nc)
7643 .arg(&e);
7644 unsafe {
7645 b.launch(cfg)?;
7646 }
7647 Ok(())
7648 }
7649
7650 #[allow(clippy::too_many_arguments)]
7652 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7655 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7656 *WARP_ON.get_or_init(|| {
7657 std::env::var("MEMRA_QKVNORM_W")
7658 .map(|v| v != "0")
7659 .unwrap_or(true)
7660 }) && ncols % 4 == 0
7661 && rows >= 64
7662 }
7663
7664 #[allow(clippy::too_many_arguments)]
7667 pub fn rms_norm_qkv_w4b(
7668 &self,
7669 q: &CudaSlice<f32>,
7670 k: &CudaSlice<f32>,
7671 v: &CudaSlice<f32>,
7672 wq: &CudaSlice<f32>,
7673 wk: &CudaSlice<f32>,
7674 wv: &CudaSlice<f32>,
7675 dq: &mut CudaSlice<f32>,
7676 dk: &mut CudaSlice<f32>,
7677 dv: &mut CudaSlice<f32>,
7678 dvb: &mut CudaSlice<u8>,
7679 ncols: usize,
7680 rq: usize,
7681 rk: usize,
7682 eps: f32,
7683 vf16: bool,
7684 ) -> Result<(), Box<dyn std::error::Error>> {
7685 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7686 let f = self.func("rms_norm_qkv_w4b_f32");
7687 let rows = (rq + 2 * rk) as u32;
7688 let cfg = LaunchConfig {
7689 grid_dim: (rows.div_ceil(8), 1, 1),
7690 block_dim: (256, 1, 1),
7691 shared_mem_bytes: 0,
7692 };
7693 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7694 let vf = vf16 as i32;
7695 let __s_b = self.gpu.stream();
7696 let mut b = __s_b.launch_builder(&f);
7697 b.arg(q)
7698 .arg(k)
7699 .arg(v)
7700 .arg(wq)
7701 .arg(wk)
7702 .arg(wv)
7703 .arg(dq)
7704 .arg(dk)
7705 .arg(dv)
7706 .arg(&mut *dvb)
7707 .arg(&nc)
7708 .arg(&rqi)
7709 .arg(&rki)
7710 .arg(&rvi)
7711 .arg(&e)
7712 .arg(&vf);
7713 unsafe {
7714 b.launch(cfg)?;
7715 }
7716 Ok(())
7717 }
7718
7719 pub fn rms_norm_qkv(
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 ncols: usize,
7731 rq: usize,
7732 rk: usize,
7733 eps: f32,
7734 ) -> Result<(), Box<dyn std::error::Error>> {
7735 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7739 let warp_on = *WARP_ON.get_or_init(|| {
7740 std::env::var("MEMRA_QKVNORM_W")
7741 .map(|v| v != "0")
7742 .unwrap_or(true)
7743 });
7744 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7747 let f = self.func("rms_norm_qkv_w4_f32");
7748 let rows = (rq + 2 * rk) as u32;
7749 let cfg = LaunchConfig {
7750 grid_dim: (rows.div_ceil(8), 1, 1),
7751 block_dim: (256, 1, 1),
7752 shared_mem_bytes: 0,
7753 };
7754 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7755 let __s_b = self.gpu.stream();
7756 let mut b = __s_b.launch_builder(&f);
7757 b.arg(q)
7758 .arg(k)
7759 .arg(v)
7760 .arg(wq)
7761 .arg(wk)
7762 .arg(wv)
7763 .arg(dq)
7764 .arg(dk)
7765 .arg(dv)
7766 .arg(&nc)
7767 .arg(&rqi)
7768 .arg(&rki)
7769 .arg(&rvi)
7770 .arg(&e);
7771 unsafe {
7772 b.launch(cfg)?;
7773 }
7774 return Ok(());
7775 }
7776 let f = self.func("rms_norm_qkv_f32");
7777 let grid = (rq + 2 * rk) as u32;
7778 let cfg = LaunchConfig {
7779 grid_dim: (grid, 1, 1),
7780 block_dim: (rms_block(), 1, 1),
7781 shared_mem_bytes: 0,
7782 };
7783 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7784 let __s_b = self.gpu.stream();
7785 let mut b = __s_b.launch_builder(&f);
7786 b.arg(q)
7787 .arg(k)
7788 .arg(v)
7789 .arg(wq)
7790 .arg(wk)
7791 .arg(wv)
7792 .arg(dq)
7793 .arg(dk)
7794 .arg(dv)
7795 .arg(&nc)
7796 .arg(&rqi)
7797 .arg(&rki)
7798 .arg(&e);
7799 unsafe {
7800 b.launch(cfg)?;
7801 }
7802 Ok(())
7803 }
7804
7805 #[allow(clippy::too_many_arguments)]
7807 pub fn rms_norm2x(
7808 &self,
7809 a: &CudaSlice<f32>,
7810 bb: &CudaSlice<f32>,
7811 wa: &CudaSlice<f32>,
7812 wb: &CudaSlice<f32>,
7813 da: &mut CudaSlice<f32>,
7814 db: &mut CudaSlice<f32>,
7815 ncols: usize,
7816 nrows: usize,
7817 eps: f32,
7818 ) -> Result<(), Box<dyn std::error::Error>> {
7819 let f = self.func("rms_norm2x_f32");
7820 let cfg = LaunchConfig {
7821 grid_dim: (2 * nrows as u32, 1, 1),
7822 block_dim: (rms_block(), 1, 1),
7823 shared_mem_bytes: 0,
7824 };
7825 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7826 let __s_b = self.gpu.stream();
7827 let mut b = __s_b.launch_builder(&f);
7828 b.arg(a)
7829 .arg(bb)
7830 .arg(wa)
7831 .arg(wb)
7832 .arg(da)
7833 .arg(db)
7834 .arg(&nc)
7835 .arg(&nr)
7836 .arg(&e);
7837 unsafe {
7838 b.launch(cfg)?;
7839 }
7840 Ok(())
7841 }
7842
7843 pub fn softcap(
7845 &self,
7846 y: &mut CudaSlice<f32>,
7847 cap: f32,
7848 n: usize,
7849 ) -> Result<(), Box<dyn std::error::Error>> {
7850 let f = self.func("softcap_f32");
7851 let cfg = LaunchConfig::for_num_elems(n as u32);
7852 let ni = n as i32;
7853 let __s_b = self.gpu.stream();
7854 let mut b = __s_b.launch_builder(&f);
7855 b.arg(y).arg(&cap).arg(&ni);
7856 unsafe {
7857 b.launch(cfg)?;
7858 }
7859 Ok(())
7860 }
7861
7862 pub fn mask_ids_rows(
7865 &self,
7866 y: &mut CudaSlice<f32>,
7867 ids: &CudaSlice<i32>,
7868 n_ids: usize,
7869 n_vocab: usize,
7870 t: usize,
7871 ) -> Result<(), Box<dyn std::error::Error>> {
7872 let f = self.func("mask_ids_rows_f32");
7873 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7874 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7875 let __s_b = self.gpu.stream();
7876 let mut b = __s_b.launch_builder(&f);
7877 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7878 unsafe {
7879 b.launch(cfg)?;
7880 }
7881 Ok(())
7882 }
7883
7884 #[allow(clippy::too_many_arguments)]
7886 pub fn add_scale_rms_norm(
7887 &self,
7888 a: &CudaSlice<f32>,
7889 b_in: &CudaSlice<f32>,
7890 c: f32,
7891 w: &CudaSlice<f32>,
7892 res: &mut CudaSlice<f32>,
7893 dst: &mut CudaSlice<f32>,
7894 ncols: usize,
7895 nrows: usize,
7896 eps: f32,
7897 ) -> Result<(), Box<dyn std::error::Error>> {
7898 let f = self.func("add_scale_rms_norm_f32");
7899 let cfg = LaunchConfig {
7900 grid_dim: (nrows as u32, 1, 1),
7901 block_dim: (rms_block(), 1, 1),
7902 shared_mem_bytes: 0,
7903 };
7904 let (nc, e2) = (ncols as i32, eps);
7905 let __s_b = self.gpu.stream();
7906 let mut b = __s_b.launch_builder(&f);
7907 b.arg(a)
7908 .arg(b_in)
7909 .arg(&c)
7910 .arg(w)
7911 .arg(res)
7912 .arg(dst)
7913 .arg(&nc)
7914 .arg(&e2);
7915 unsafe {
7916 b.launch(cfg)?;
7917 }
7918 Ok(())
7919 }
7920
7921 #[allow(clippy::too_many_arguments)]
7924 pub fn add_scale_rms_norm_q8_1(
7925 &self,
7926 a: &CudaSlice<f32>,
7927 b_in: &CudaSlice<f32>,
7928 c: f32,
7929 w: &CudaSlice<f32>,
7930 res: &mut CudaSlice<f32>,
7931 ncols: usize,
7932 nrows: usize,
7933 eps: f32,
7934 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7935 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7936 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7937 let (nc, e2) = (ncols as i32, eps);
7938 if Self::pdl_on() && Self::pdl_wb_on() {
7939 {
7940 use cudarc::driver::{DevicePtr, DevicePtrMut};
7941 let s = &self.gpu.stream();
7942 let (pa, _g0) = a.device_ptr(s);
7943 let (pb, _g1) = b_in.device_ptr(s);
7944 let (pw, _g2) = w.device_ptr(s);
7945 let (pr, _g3) = res.device_ptr_mut(s);
7946 let (pq, _g4) = out_q.device_ptr_mut(s);
7947 let (pd, _g5) = out_d.device_ptr_mut(s);
7948 let mut ps = [
7949 &pa as *const _ as *mut std::ffi::c_void,
7950 &pb as *const _ as *mut _,
7951 &c as *const _ as *mut _,
7952 &pw as *const _ as *mut _,
7953 &pr as *const _ as *mut _,
7954 &pq as *const _ as *mut _,
7955 &pd as *const _ as *mut _,
7956 &nc as *const _ as *mut _,
7957 &e2 as *const _ as *mut _,
7958 ];
7959 unsafe {
7960 self.launch_pdl(
7961 "add_scale_rms_norm_q8_1",
7962 (nrows as u32, 1, 1),
7963 (rms_block(), 1, 1),
7964 &mut ps,
7965 )?;
7966 }
7967 }
7968 return Ok((out_q, out_d));
7969 }
7970 let f = self.func("add_scale_rms_norm_q8_1");
7971 let cfg = LaunchConfig {
7972 grid_dim: (nrows as u32, 1, 1),
7973 block_dim: (rms_block(), 1, 1),
7974 shared_mem_bytes: 0,
7975 };
7976 let __s_b = self.gpu.stream();
7977 let mut b = __s_b.launch_builder(&f);
7978 b.arg(a)
7979 .arg(b_in)
7980 .arg(&c)
7981 .arg(w)
7982 .arg(res)
7983 .arg(&mut out_q)
7984 .arg(&mut out_d)
7985 .arg(&nc)
7986 .arg(&e2);
7987 unsafe {
7988 b.launch(cfg)?;
7989 }
7990 Ok((out_q, out_d))
7991 }
7992
7993 #[allow(clippy::too_many_arguments)]
7995 pub fn add_scale_rms_norm_q8_1_into(
7996 &self,
7997 a: &CudaSlice<f32>,
7998 b_in: &CudaSlice<f32>,
7999 c: f32,
8000 w: &CudaSlice<f32>,
8001 res: &mut CudaSlice<f32>,
8002 ncols: usize,
8003 nrows: usize,
8004 eps: f32,
8005 out_q: &mut CudaSlice<i8>,
8006 out_d: &mut CudaSlice<f32>,
8007 ) -> Result<(), Box<dyn std::error::Error>> {
8008 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8009 let (nc, e2) = (ncols as i32, eps);
8010 if Self::pdl_on() && Self::pdl_wb_on() {
8011 use cudarc::driver::{DevicePtr, DevicePtrMut};
8012 let s = &self.gpu.stream();
8013 let (pa, _g0) = a.device_ptr(s);
8014 let (pb, _g1) = b_in.device_ptr(s);
8015 let (pw, _g2) = w.device_ptr(s);
8016 let (pr, _g3) = res.device_ptr_mut(s);
8017 let (pq, _g4) = out_q.device_ptr_mut(s);
8018 let (pd, _g5) = out_d.device_ptr_mut(s);
8019 let mut ps = [
8020 &pa as *const _ as *mut std::ffi::c_void,
8021 &pb as *const _ as *mut _,
8022 &c as *const _ as *mut _,
8023 &pw as *const _ as *mut _,
8024 &pr as *const _ as *mut _,
8025 &pq as *const _ as *mut _,
8026 &pd as *const _ as *mut _,
8027 &nc as *const _ as *mut _,
8028 &e2 as *const _ as *mut _,
8029 ];
8030 unsafe {
8031 self.launch_pdl(
8032 "add_scale_rms_norm_q8_1",
8033 (nrows as u32, 1, 1),
8034 (rms_block(), 1, 1),
8035 &mut ps,
8036 )?;
8037 }
8038 return Ok(());
8039 }
8040 let f = self.func("add_scale_rms_norm_q8_1");
8041 let cfg = LaunchConfig {
8042 grid_dim: (nrows as u32, 1, 1),
8043 block_dim: (rms_block(), 1, 1),
8044 shared_mem_bytes: 0,
8045 };
8046 let __s_b = self.gpu.stream();
8047 let mut b = __s_b.launch_builder(&f);
8048 b.arg(a)
8049 .arg(b_in)
8050 .arg(&c)
8051 .arg(w)
8052 .arg(res)
8053 .arg(&mut *out_q)
8054 .arg(&mut *out_d)
8055 .arg(&nc)
8056 .arg(&e2);
8057 unsafe {
8058 b.launch(cfg)?;
8059 }
8060 Ok(())
8061 }
8062
8063 #[allow(clippy::too_many_arguments)]
8066 pub fn rms_pre_add_scale_rms_norm_q8_1(
8067 &self,
8068 a: &CudaSlice<f32>,
8069 wa: &CudaSlice<f32>,
8070 b_in: &CudaSlice<f32>,
8071 c: f32,
8072 w: &CudaSlice<f32>,
8073 res: &mut CudaSlice<f32>,
8074 ncols: usize,
8075 nrows: usize,
8076 eps: f32,
8077 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8078 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8079 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8080 let (nc, e2) = (ncols as i32, eps);
8081 if Self::pdl_on() {
8082 {
8083 use cudarc::driver::{DevicePtr, DevicePtrMut};
8084 let s = &self.gpu.stream();
8085 let (pa, _g0) = a.device_ptr(s);
8086 let (pwa, _g1) = wa.device_ptr(s);
8087 let (pb, _g2) = b_in.device_ptr(s);
8088 let (pw, _g3) = w.device_ptr(s);
8089 let (pr, _g4) = res.device_ptr_mut(s);
8090 let (pq, _g5) = out_q.device_ptr_mut(s);
8091 let (pd, _g6) = out_d.device_ptr_mut(s);
8092 let mut ps = [
8093 &pa as *const _ as *mut std::ffi::c_void,
8094 &pwa as *const _ as *mut _,
8095 &pb as *const _ as *mut _,
8096 &c as *const _ as *mut _,
8097 &pw as *const _ as *mut _,
8098 &pr as *const _ as *mut _,
8099 &pq as *const _ as *mut _,
8100 &pd as *const _ as *mut _,
8101 &nc as *const _ as *mut _,
8102 &e2 as *const _ as *mut _,
8103 ];
8104 unsafe {
8105 self.launch_pdl(
8106 "rms_pre_add_scale_rms_norm_q8_1",
8107 (nrows as u32, 1, 1),
8108 (rms_block(), 1, 1),
8109 &mut ps,
8110 )?;
8111 }
8112 }
8113 return Ok((out_q, out_d));
8114 }
8115 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8116 let cfg = LaunchConfig {
8117 grid_dim: (nrows as u32, 1, 1),
8118 block_dim: (rms_block(), 1, 1),
8119 shared_mem_bytes: 0,
8120 };
8121 let __s_b = self.gpu.stream();
8122 let mut b = __s_b.launch_builder(&f);
8123 b.arg(a)
8124 .arg(wa)
8125 .arg(b_in)
8126 .arg(&c)
8127 .arg(w)
8128 .arg(res)
8129 .arg(&mut out_q)
8130 .arg(&mut out_d)
8131 .arg(&nc)
8132 .arg(&e2);
8133 unsafe {
8134 b.launch(cfg)?;
8135 }
8136 Ok((out_q, out_d))
8137 }
8138
8139 pub fn gelu_tanh_mul_q8_1(
8142 &self,
8143 gate: &CudaSlice<f32>,
8144 up: &cudarc::driver::CudaView<f32>,
8145 act: &mut CudaSlice<f32>,
8146 ncols: usize,
8147 nrows: usize,
8148 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8149 debug_assert!(ncols % 128 == 0);
8150 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8151 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8152 let nc = ncols as i32;
8153 if Self::pdl_on() {
8154 {
8155 use cudarc::driver::{DevicePtr, DevicePtrMut};
8156 let s = &self.gpu.stream();
8157 let (pg, _g0) = gate.device_ptr(s);
8158 let (pu, _g1) = up.device_ptr(s);
8159 let (pact, _g2) = act.device_ptr_mut(s);
8160 let (pq, _g3) = out_q.device_ptr_mut(s);
8161 let (pd, _g4) = out_d.device_ptr_mut(s);
8162 let mut ps = [
8163 &pg as *const _ as *mut std::ffi::c_void,
8164 &pu as *const _ as *mut _,
8165 &pact as *const _ as *mut _,
8166 &pq as *const _ as *mut _,
8167 &pd as *const _ as *mut _,
8168 &nc as *const _ as *mut _,
8169 ];
8170 unsafe {
8171 self.launch_pdl(
8172 "gelu_tanh_mul_q8_1",
8173 (nrows as u32, 1, 1),
8174 (rms_block(), 1, 1),
8175 &mut ps,
8176 )?;
8177 }
8178 }
8179 return Ok((out_q, out_d));
8180 }
8181 let f = self.func("gelu_tanh_mul_q8_1");
8182 let cfg = LaunchConfig {
8183 grid_dim: (nrows as u32, 1, 1),
8184 block_dim: (rms_block(), 1, 1),
8185 shared_mem_bytes: 0,
8186 };
8187 let __s_b = self.gpu.stream();
8188 let mut b = __s_b.launch_builder(&f);
8189 b.arg(gate)
8190 .arg(up)
8191 .arg(act)
8192 .arg(&mut out_q)
8193 .arg(&mut out_d)
8194 .arg(&nc);
8195 unsafe {
8196 b.launch(cfg)?;
8197 }
8198 Ok((out_q, out_d))
8199 }
8200
8201 #[allow(clippy::too_many_arguments)]
8203 pub fn gelu_tanh_mul_q8_1_into(
8204 &self,
8205 gate: &CudaSlice<f32>,
8206 up: &cudarc::driver::CudaView<f32>,
8207 act: &mut CudaSlice<f32>,
8208 ncols: usize,
8209 nrows: usize,
8210 out_q: &mut CudaSlice<i8>,
8211 out_d: &mut CudaSlice<f32>,
8212 ) -> Result<(), Box<dyn std::error::Error>> {
8213 debug_assert!(ncols % 128 == 0);
8214 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8215 let nc = ncols as i32;
8216 if Self::pdl_on() {
8217 use cudarc::driver::{DevicePtr, DevicePtrMut};
8218 let s = &self.gpu.stream();
8219 let (pg, _g0) = gate.device_ptr(s);
8220 let (pu, _g1) = up.device_ptr(s);
8221 let (pact, _g2) = act.device_ptr_mut(s);
8222 let (pq, _g3) = out_q.device_ptr_mut(s);
8223 let (pd, _g4) = out_d.device_ptr_mut(s);
8224 let mut ps = [
8225 &pg as *const _ as *mut std::ffi::c_void,
8226 &pu as *const _ as *mut _,
8227 &pact as *const _ as *mut _,
8228 &pq as *const _ as *mut _,
8229 &pd as *const _ as *mut _,
8230 &nc as *const _ as *mut _,
8231 ];
8232 unsafe {
8233 self.launch_pdl(
8234 "gelu_tanh_mul_q8_1",
8235 (nrows as u32, 1, 1),
8236 (rms_block(), 1, 1),
8237 &mut ps,
8238 )?;
8239 }
8240 return Ok(());
8241 }
8242 let f = self.func("gelu_tanh_mul_q8_1");
8243 let cfg = LaunchConfig {
8244 grid_dim: (nrows as u32, 1, 1),
8245 block_dim: (rms_block(), 1, 1),
8246 shared_mem_bytes: 0,
8247 };
8248 let __s_b = self.gpu.stream();
8249 let mut b = __s_b.launch_builder(&f);
8250 b.arg(gate)
8251 .arg(up)
8252 .arg(&mut *act)
8253 .arg(&mut *out_q)
8254 .arg(&mut *out_d)
8255 .arg(&nc);
8256 unsafe {
8257 b.launch(cfg)?;
8258 }
8259 Ok(())
8260 }
8261
8262 #[allow(clippy::too_many_arguments)]
8264 pub fn add_rms_norm3_q8z(
8265 &self,
8266 a: &CudaSlice<f32>,
8267 b_in: &CudaSlice<f32>,
8268 w0: &CudaSlice<f32>,
8269 w1: &CudaSlice<f32>,
8270 w2: &CudaSlice<f32>,
8271 res: &mut CudaSlice<f32>,
8272 out1: &mut CudaSlice<f32>,
8273 ncols: usize,
8274 nrows: usize,
8275 eps: f32,
8276 ) -> Result<
8277 (
8278 (CudaSlice<i8>, CudaSlice<f32>),
8279 (CudaSlice<i8>, CudaSlice<f32>),
8280 ),
8281 Box<dyn std::error::Error>,
8282 > {
8283 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
8284 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8285 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
8286 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8287 let f = self.func("add_rms_norm3_q8z_f32");
8288 let cfg = LaunchConfig {
8289 grid_dim: (nrows as u32, 1, 1),
8290 block_dim: (rms_block(), 1, 1),
8291 shared_mem_bytes: 0,
8292 };
8293 let (nc, e2) = (ncols as i32, eps);
8294 let __s_b = self.gpu.stream();
8295 let mut b = __s_b.launch_builder(&f);
8296 b.arg(a)
8297 .arg(b_in)
8298 .arg(w0)
8299 .arg(w1)
8300 .arg(w2)
8301 .arg(res)
8302 .arg(&mut q0)
8303 .arg(&mut d0)
8304 .arg(out1)
8305 .arg(&mut q2)
8306 .arg(&mut d2)
8307 .arg(&nc)
8308 .arg(&e2);
8309 unsafe {
8310 b.launch(cfg)?;
8311 }
8312 Ok(((q0, d0), (q2, d2)))
8313 }
8314
8315 #[allow(clippy::too_many_arguments)]
8317 pub fn add_rms_norm3(
8318 &self,
8319 a: &CudaSlice<f32>,
8320 b_in: &CudaSlice<f32>,
8321 w0: &CudaSlice<f32>,
8322 w1: &CudaSlice<f32>,
8323 w2: &CudaSlice<f32>,
8324 res: &mut CudaSlice<f32>,
8325 d0: &mut CudaSlice<f32>,
8326 d1: &mut CudaSlice<f32>,
8327 d2: &mut CudaSlice<f32>,
8328 ncols: usize,
8329 nrows: usize,
8330 eps: f32,
8331 ) -> Result<(), Box<dyn std::error::Error>> {
8332 let f = self.func("add_rms_norm3_f32");
8333 let cfg = LaunchConfig {
8334 grid_dim: (nrows as u32, 1, 1),
8335 block_dim: (rms_block(), 1, 1),
8336 shared_mem_bytes: 0,
8337 };
8338 let (nc, e2) = (ncols as i32, eps);
8339 let __s_b = self.gpu.stream();
8340 let mut b = __s_b.launch_builder(&f);
8341 b.arg(a)
8342 .arg(b_in)
8343 .arg(w0)
8344 .arg(w1)
8345 .arg(w2)
8346 .arg(res)
8347 .arg(d0)
8348 .arg(d1)
8349 .arg(d2)
8350 .arg(&nc)
8351 .arg(&e2);
8352 unsafe {
8353 b.launch(cfg)?;
8354 }
8355 Ok(())
8356 }
8357
8358 pub fn add_scale(
8360 &self,
8361 a: &CudaSlice<f32>,
8362 b_in: &CudaSlice<f32>,
8363 c: f32,
8364 dst: &mut CudaSlice<f32>,
8365 n: usize,
8366 ) -> Result<(), Box<dyn std::error::Error>> {
8367 let f = self.func("add_scale_f32");
8368 let cfg = LaunchConfig::for_num_elems(n as u32);
8369 let ni = n as i32;
8370 let __s_b = self.gpu.stream();
8371 let mut b = __s_b.launch_builder(&f);
8372 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8373 unsafe {
8374 b.launch(cfg)?;
8375 }
8376 Ok(())
8377 }
8378
8379 pub fn layer_norm_bias(
8381 &self,
8382 x: &CudaSlice<f32>,
8383 w: &CudaSlice<f32>,
8384 b: &CudaSlice<f32>,
8385 dst: &mut CudaSlice<f32>,
8386 ncols: usize,
8387 nrows: usize,
8388 eps: f32,
8389 ) -> Result<(), Box<dyn std::error::Error>> {
8390 let f = self.func("layer_norm_bias_f32");
8391 let (nc, e) = (ncols as i32, eps);
8392 let cfg = LaunchConfig {
8393 grid_dim: (nrows as u32, 1, 1),
8394 block_dim: (256, 1, 1),
8395 shared_mem_bytes: 0,
8396 };
8397 let __s_b = self.gpu.stream();
8398 let mut lb = __s_b.launch_builder(&f);
8399 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8400 unsafe {
8401 lb.launch(cfg)?;
8402 }
8403 Ok(())
8404 }
8405
8406 pub fn gelu_tanh(
8408 &self,
8409 x: &CudaSlice<f32>,
8410 dst: &mut CudaSlice<f32>,
8411 n: usize,
8412 ) -> Result<(), Box<dyn std::error::Error>> {
8413 let f = self.func("gelu_tanh_f32");
8414 let ni = n as i64;
8415 let cfg = LaunchConfig {
8416 grid_dim: (n.div_ceil(256) as u32, 1, 1),
8417 block_dim: (256, 1, 1),
8418 shared_mem_bytes: 0,
8419 };
8420 let __s_b = self.gpu.stream();
8421 let mut lb = __s_b.launch_builder(&f);
8422 lb.arg(x).arg(&mut *dst).arg(&ni);
8423 unsafe {
8424 lb.launch(cfg)?;
8425 }
8426 Ok(())
8427 }
8428
8429 pub fn row_softmax(
8431 &self,
8432 x: &mut CudaSlice<f32>,
8433 ncols: usize,
8434 nrows: usize,
8435 ) -> Result<(), Box<dyn std::error::Error>> {
8436 let f = self.func("row_softmax_f32");
8437 let nc = ncols as i32;
8438 let cfg = LaunchConfig {
8439 grid_dim: (nrows as u32, 1, 1),
8440 block_dim: (256, 1, 1),
8441 shared_mem_bytes: 0,
8442 };
8443 let __s_b = self.gpu.stream();
8444 let mut lb = __s_b.launch_builder(&f);
8445 lb.arg(&mut *x).arg(&nc);
8446 unsafe {
8447 lb.launch(cfg)?;
8448 }
8449 Ok(())
8450 }
8451
8452 pub fn rms_norm(
8453 &self,
8454 x: &CudaSlice<f32>,
8455 w: &CudaSlice<f32>,
8456 dst: &mut CudaSlice<f32>,
8457 ncols: usize,
8458 nrows: usize,
8459 eps: f32,
8460 ) -> Result<(), Box<dyn std::error::Error>> {
8461 let (nc, e) = (ncols as i32, eps);
8462 let kname = if Self::norm_ilp_on() {
8463 "rms_norm_f32_v2"
8464 } else {
8465 "rms_norm_f32"
8466 };
8467 if Self::pdl_on() && Self::pdl_wb_on() {
8468 use cudarc::driver::{DevicePtr, DevicePtrMut};
8469 let s = &self.gpu.stream();
8470 let (px, _g0) = x.device_ptr(s);
8471 let (pw, _g1) = w.device_ptr(s);
8472 let (pd, _g2) = dst.device_ptr_mut(s);
8473 let mut ps = [
8474 &px as *const _ as *mut std::ffi::c_void,
8475 &pw as *const _ as *mut _,
8476 &pd as *const _ as *mut _,
8477 &nc as *const _ as *mut _,
8478 &e as *const _ as *mut _,
8479 ];
8480 unsafe {
8481 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
8482 }
8483 return Ok(());
8484 }
8485 let f = self.func(kname);
8486 let cfg = LaunchConfig {
8487 grid_dim: (nrows as u32, 1, 1),
8488 block_dim: (rms_block(), 1, 1),
8489 shared_mem_bytes: 0,
8490 };
8491 let __s_b = self.gpu.stream();
8492 let mut b = __s_b.launch_builder(&f);
8493 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8494 unsafe {
8495 b.launch(cfg)?;
8496 }
8497 Ok(())
8498 }
8499
8500 pub fn rms_norm_decode(
8508 &self,
8509 x: &CudaSlice<f32>,
8510 w: &CudaSlice<f32>,
8511 dst: &mut CudaSlice<f32>,
8512 ncols: usize,
8513 nrows: usize,
8514 eps: f32,
8515 ) -> Result<(), Box<dyn std::error::Error>> {
8516 let f = self.func(if Self::norm_ilp_on() {
8517 "rms_norm_f32_v2"
8518 } else {
8519 "rms_norm_f32"
8520 });
8521 let cfg = LaunchConfig {
8522 grid_dim: (nrows as u32, 1, 1),
8523 block_dim: (1024, 1, 1),
8524 shared_mem_bytes: 0,
8525 };
8526 let (nc, e) = (ncols as i32, eps);
8527 let __s_b = self.gpu.stream();
8528 let mut b = __s_b.launch_builder(&f);
8529 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8530 unsafe {
8531 b.launch(cfg)?;
8532 }
8533 Ok(())
8534 }
8535
8536 pub fn rms_norm_q8_1(
8540 &self,
8541 x: &CudaSlice<f32>,
8542 w: &CudaSlice<f32>,
8543 ncols: usize,
8544 nrows: usize,
8545 eps: f32,
8546 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8547 let nblk = ncols / 32;
8548 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8549 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8550 let (nc, e) = (ncols as i32, eps);
8551 if Self::pdl_on() {
8552 {
8553 use cudarc::driver::{DevicePtr, DevicePtrMut};
8554 let s = &self.gpu.stream();
8555 let (px, _g0) = x.device_ptr(s);
8556 let (pw, _g1) = w.device_ptr(s);
8557 let (pq, _g2) = q.device_ptr_mut(s);
8558 let (pd, _g3) = d.device_ptr_mut(s);
8559 let mut ps = [
8560 &px as *const _ as *mut std::ffi::c_void,
8561 &pw as *const _ as *mut _,
8562 &pq as *const _ as *mut _,
8563 &pd as *const _ as *mut _,
8564 &nc as *const _ as *mut _,
8565 &e as *const _ as *mut _,
8566 ];
8567 unsafe {
8568 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8569 }
8570 }
8571 return Ok((q, d));
8572 }
8573 let f = self.func("rms_norm_q8_1");
8574 let cfg = LaunchConfig {
8577 grid_dim: (nrows as u32, 1, 1),
8578 block_dim: (1024, 1, 1),
8579 shared_mem_bytes: 0,
8580 };
8581 let __s_b = self.gpu.stream();
8582 let mut b = __s_b.launch_builder(&f);
8583 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8584 unsafe {
8585 b.launch(cfg)?;
8586 }
8587 Ok((q, d))
8588 }
8589
8590 pub fn rms_norm_q8_1_into(
8593 &self,
8594 x: &CudaSlice<f32>,
8595 w: &CudaSlice<f32>,
8596 ncols: usize,
8597 nrows: usize,
8598 eps: f32,
8599 q: &mut CudaSlice<i8>,
8600 d: &mut CudaSlice<f32>,
8601 ) -> Result<(), Box<dyn std::error::Error>> {
8602 let nblk = ncols / 32;
8603 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8604 let (nc, e) = (ncols as i32, eps);
8605 if Self::pdl_on() {
8606 use cudarc::driver::{DevicePtr, DevicePtrMut};
8607 let s = &self.gpu.stream();
8608 let (px, _g0) = x.device_ptr(s);
8609 let (pw, _g1) = w.device_ptr(s);
8610 let (pq, _g2) = q.device_ptr_mut(s);
8611 let (pd, _g3) = d.device_ptr_mut(s);
8612 let mut ps = [
8613 &px as *const _ as *mut std::ffi::c_void,
8614 &pw as *const _ as *mut _,
8615 &pq as *const _ as *mut _,
8616 &pd as *const _ as *mut _,
8617 &nc as *const _ as *mut _,
8618 &e as *const _ as *mut _,
8619 ];
8620 unsafe {
8621 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8622 }
8623 return Ok(());
8624 }
8625 let f = self.func("rms_norm_q8_1");
8626 let cfg = LaunchConfig {
8627 grid_dim: (nrows as u32, 1, 1),
8628 block_dim: (1024, 1, 1),
8629 shared_mem_bytes: 0,
8630 };
8631 let __s_b = self.gpu.stream();
8632 let mut b = __s_b.launch_builder(&f);
8633 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8634 unsafe {
8635 b.launch(cfg)?;
8636 }
8637 Ok(())
8638 }
8639
8640 pub fn quantize_q8_1_into(
8642 &self,
8643 x: &CudaSlice<f32>,
8644 m: usize,
8645 in_f: usize,
8646 q: &mut CudaSlice<i8>,
8647 d: &mut CudaSlice<f32>,
8648 ) -> Result<(), Box<dyn std::error::Error>> {
8649 let nblk = in_f / 32;
8650 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8651 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8652 let (inf, mi) = (in_f as i32, m as i32);
8653 if Self::pdl_on() && Self::pdl_wb_on() {
8654 use cudarc::driver::{DevicePtr, DevicePtrMut};
8655 let s = &self.gpu.stream();
8656 let (px, _g0) = x.device_ptr(s);
8657 let (pq, _g1) = q.device_ptr_mut(s);
8658 let (pd, _g2) = d.device_ptr_mut(s);
8659 let mut ps = [
8660 &px as *const _ as *mut std::ffi::c_void,
8661 &pq as *const _ as *mut _,
8662 &pd as *const _ as *mut _,
8663 &inf as *const _ as *mut _,
8664 &mi as *const _ as *mut _,
8665 ];
8666 unsafe {
8667 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8668 }
8669 return Ok(());
8670 }
8671 let f = self.func("quantize_q8_1");
8672 let __s_b = self.gpu.stream();
8673 let mut b = __s_b.launch_builder(&f);
8674 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8675 unsafe {
8676 b.launch(cfg)?;
8677 }
8678 Ok(())
8679 }
8680
8681 pub fn add_rms_norm_q8_1(
8685 &self,
8686 a: &CudaSlice<f32>,
8687 b_in: &CudaSlice<f32>,
8688 w: &CudaSlice<f32>,
8689 res: &mut CudaSlice<f32>,
8690 ncols: usize,
8691 nrows: usize,
8692 eps: f32,
8693 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8694 let nblk = ncols / 32;
8695 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8696 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8697 let f = self.func("add_rms_norm_q8_1");
8698 let cfg = LaunchConfig {
8700 grid_dim: (nrows as u32, 1, 1),
8701 block_dim: (1024, 1, 1),
8702 shared_mem_bytes: 0,
8703 };
8704 let (nc, e) = (ncols as i32, eps);
8705 let __s_bld = self.gpu.stream();
8706 let mut bld = __s_bld.launch_builder(&f);
8707 bld.arg(a)
8708 .arg(b_in)
8709 .arg(w)
8710 .arg(res)
8711 .arg(&mut q)
8712 .arg(&mut d)
8713 .arg(&nc)
8714 .arg(&e);
8715 unsafe {
8716 bld.launch(cfg)?;
8717 }
8718 Ok((q, d))
8719 }
8720
8721 pub fn add_rms_norm(
8725 &self,
8726 a: &CudaSlice<f32>,
8727 b: &CudaSlice<f32>,
8728 w: &CudaSlice<f32>,
8729 res: &mut CudaSlice<f32>,
8730 dst: &mut CudaSlice<f32>,
8731 ncols: usize,
8732 nrows: usize,
8733 eps: f32,
8734 ) -> Result<(), Box<dyn std::error::Error>> {
8735 let (nc, e) = (ncols as i32, eps);
8736 let kname = if Self::norm_ilp_on() {
8737 "add_rms_norm_f32_v2"
8738 } else {
8739 "add_rms_norm_f32"
8740 };
8741 if Self::pdl_on() && Self::pdl_wb_on() {
8742 use cudarc::driver::{DevicePtr, DevicePtrMut};
8743 let s = &self.gpu.stream();
8744 let (pa, _g0) = a.device_ptr(s);
8745 let (pb, _g1) = b.device_ptr(s);
8746 let (pw, _g2) = w.device_ptr(s);
8747 let (pr, _g3) = res.device_ptr_mut(s);
8748 let (pd, _g4) = dst.device_ptr_mut(s);
8749 let mut ps = [
8750 &pa as *const _ as *mut std::ffi::c_void,
8751 &pb as *const _ as *mut _,
8752 &pw as *const _ as *mut _,
8753 &pr as *const _ as *mut _,
8754 &pd as *const _ as *mut _,
8755 &nc as *const _ as *mut _,
8756 &e as *const _ as *mut _,
8757 ];
8758 unsafe {
8759 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
8760 }
8761 return Ok(());
8762 }
8763 let f = self.func(kname);
8764 let cfg = LaunchConfig {
8765 grid_dim: (nrows as u32, 1, 1),
8766 block_dim: (rms_block(), 1, 1),
8767 shared_mem_bytes: 0,
8768 };
8769 let __s_b2 = self.gpu.stream();
8770 let mut b2 = __s_b2.launch_builder(&f);
8771 b2.arg(a)
8772 .arg(b)
8773 .arg(w)
8774 .arg(&mut *res)
8775 .arg(&mut *dst)
8776 .arg(&nc)
8777 .arg(&e);
8778 unsafe {
8779 b2.launch(cfg)?;
8780 }
8781 Ok(())
8782 }
8783
8784 #[allow(clippy::too_many_arguments)]
8787 pub fn rms_pre_add_rms_norm(
8788 &self,
8789 a: &CudaSlice<f32>,
8790 wa: &CudaSlice<f32>,
8791 b: &CudaSlice<f32>,
8792 w: &CudaSlice<f32>,
8793 res: &mut CudaSlice<f32>,
8794 dst: &mut CudaSlice<f32>,
8795 ncols: usize,
8796 nrows: usize,
8797 eps: f32,
8798 ) -> Result<(), Box<dyn std::error::Error>> {
8799 let f = self.func("rms_pre_add_rms_norm_f32");
8800 let cfg = LaunchConfig {
8801 grid_dim: (nrows as u32, 1, 1),
8802 block_dim: (rms_block(), 1, 1),
8803 shared_mem_bytes: 0,
8804 };
8805 let (nc, e) = (ncols as i32, eps);
8806 let __s_b2 = self.gpu.stream();
8807 let mut b2 = __s_b2.launch_builder(&f);
8808 b2.arg(a)
8809 .arg(wa)
8810 .arg(b)
8811 .arg(w)
8812 .arg(&mut *res)
8813 .arg(&mut *dst)
8814 .arg(&nc)
8815 .arg(&e);
8816 unsafe {
8817 b2.launch(cfg)?;
8818 }
8819 Ok(())
8820 }
8821
8822 #[allow(clippy::too_many_arguments)]
8824 pub fn rms_pre_add_rms_norm_q8z(
8825 &self,
8826 a: &CudaSlice<f32>,
8827 wa: &CudaSlice<f32>,
8828 b: &CudaSlice<f32>,
8829 w: &CudaSlice<f32>,
8830 res: &mut CudaSlice<f32>,
8831 dst: &mut CudaSlice<f32>,
8832 ncols: usize,
8833 nrows: usize,
8834 eps: f32,
8835 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8836 debug_assert!(ncols % 128 == 0);
8837 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8838 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8839 let (nc, e) = (ncols as i32, eps);
8840 if Self::pdl_on() {
8841 {
8842 use cudarc::driver::{DevicePtr, DevicePtrMut};
8843 let s = &self.gpu.stream();
8844 let (pa, _g0) = a.device_ptr(s);
8845 let (pwa, _g1) = wa.device_ptr(s);
8846 let (pb, _g2) = b.device_ptr(s);
8847 let (pw, _g3) = w.device_ptr(s);
8848 let (pr, _g4) = res.device_ptr_mut(s);
8849 let (pdst, _g5) = dst.device_ptr_mut(s);
8850 let (pq, _g6) = out_q.device_ptr_mut(s);
8851 let (pd, _g7) = out_d.device_ptr_mut(s);
8852 let mut ps = [
8853 &pa as *const _ as *mut std::ffi::c_void,
8854 &pwa as *const _ as *mut _,
8855 &pb as *const _ as *mut _,
8856 &pw as *const _ as *mut _,
8857 &pr as *const _ as *mut _,
8858 &pdst as *const _ as *mut _,
8859 &pq as *const _ as *mut _,
8860 &pd as *const _ as *mut _,
8861 &nc as *const _ as *mut _,
8862 &e as *const _ as *mut _,
8863 ];
8864 unsafe {
8865 self.launch_pdl(
8866 "rms_pre_add_rms_norm_q8z_f32",
8867 (nrows as u32, 1, 1),
8868 (rms_block(), 1, 1),
8869 &mut ps,
8870 )?;
8871 }
8872 }
8873 return Ok((out_q, out_d));
8874 }
8875 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8876 let cfg = LaunchConfig {
8877 grid_dim: (nrows as u32, 1, 1),
8878 block_dim: (rms_block(), 1, 1),
8879 shared_mem_bytes: 0,
8880 };
8881 let __s_b2 = self.gpu.stream();
8882 let mut b2 = __s_b2.launch_builder(&f);
8883 b2.arg(a)
8884 .arg(wa)
8885 .arg(b)
8886 .arg(w)
8887 .arg(&mut *res)
8888 .arg(&mut *dst)
8889 .arg(&mut out_q)
8890 .arg(&mut out_d)
8891 .arg(&nc)
8892 .arg(&e);
8893 unsafe {
8894 b2.launch(cfg)?;
8895 }
8896 Ok((out_q, out_d))
8897 }
8898
8899 #[allow(clippy::too_many_arguments)]
8903 pub fn rms_pre_add_rms_norm_q8z_into(
8904 &self,
8905 a: &CudaSlice<f32>,
8906 wa: &CudaSlice<f32>,
8907 b: &CudaSlice<f32>,
8908 w: &CudaSlice<f32>,
8909 res: &mut CudaSlice<f32>,
8910 dst: &mut CudaSlice<f32>,
8911 ncols: usize,
8912 nrows: usize,
8913 eps: f32,
8914 out_q: &mut CudaSlice<i8>,
8915 out_d: &mut CudaSlice<f32>,
8916 ) -> Result<(), Box<dyn std::error::Error>> {
8917 debug_assert!(ncols % 128 == 0);
8918 let (nc, e) = (ncols as i32, eps);
8919 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8920 let cfg = LaunchConfig {
8921 grid_dim: (nrows as u32, 1, 1),
8922 block_dim: (rms_block(), 1, 1),
8923 shared_mem_bytes: 0,
8924 };
8925 let __s_b = self.gpu.stream();
8926 let mut b2 = __s_b.launch_builder(&f);
8927 b2.arg(a)
8928 .arg(wa)
8929 .arg(b)
8930 .arg(w)
8931 .arg(&mut *res)
8932 .arg(&mut *dst)
8933 .arg(&mut *out_q)
8934 .arg(&mut *out_d)
8935 .arg(&nc)
8936 .arg(&e);
8937 unsafe {
8938 b2.launch(cfg)?;
8939 }
8940 Ok(())
8941 }
8942
8943 #[allow(clippy::too_many_arguments)]
8946 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8947 &self,
8948 a: &CudaSlice<f32>,
8949 wa: &CudaSlice<f32>,
8950 b_in: &CudaSlice<f32>,
8951 c: f32,
8952 w: &CudaSlice<f32>,
8953 res: &mut CudaSlice<f32>,
8954 ncols: usize,
8955 nrows: usize,
8956 eps: f32,
8957 out_q: &mut CudaSlice<i8>,
8958 out_d: &mut CudaSlice<f32>,
8959 ) -> Result<(), Box<dyn std::error::Error>> {
8960 debug_assert!(ncols % 128 == 0);
8961 let (nc, e2) = (ncols as i32, eps);
8962 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8963 let cfg = LaunchConfig {
8964 grid_dim: (nrows as u32, 1, 1),
8965 block_dim: (rms_block(), 1, 1),
8966 shared_mem_bytes: 0,
8967 };
8968 let __s_b = self.gpu.stream();
8969 let mut b2 = __s_b.launch_builder(&f);
8970 b2.arg(a)
8971 .arg(wa)
8972 .arg(b_in)
8973 .arg(&c)
8974 .arg(w)
8975 .arg(&mut *res)
8976 .arg(&mut *out_q)
8977 .arg(&mut *out_d)
8978 .arg(&nc)
8979 .arg(&e2);
8980 unsafe {
8981 b2.launch(cfg)?;
8982 }
8983 Ok(())
8984 }
8985
8986 pub fn g4_pnfold_on() -> bool {
8994 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8995 *ON.get_or_init(|| {
8996 std::env::var("MEMRA_G4_PNFOLD")
8997 .map(|v| v != "0")
8998 .unwrap_or(true)
8999 })
9000 }
9001
9002 pub fn build_q4_out_concat3(
9006 &self,
9007 w0: &crate::model::GpuTensor,
9008 w1: &crate::model::GpuTensor,
9009 w2: &crate::model::GpuTensor,
9010 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
9011 use crate::model::GpuTensor;
9012 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
9013 match w {
9014 GpuTensor::Quant {
9015 qtype,
9016 row_bytes,
9017 rp,
9018 ..
9019 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
9020 _ => None,
9021 }
9022 };
9023 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
9024 else {
9025 return Ok(None);
9026 };
9027 if rb0 != rb1
9028 || rb0 != rb2
9029 || w0.in_features() != w1.in_features()
9030 || w0.in_features() != w2.in_features()
9031 {
9032 return Ok(None);
9033 }
9034 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
9035 match w {
9036 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
9037 _ => unreachable!(),
9038 }
9039 }
9040 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
9041 let total = rb0 * (o0 + o1 + o2);
9042 let mut cat = self.alloc_u8(total)?;
9043 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
9044 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
9045 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
9046 Ok(Some(GpuTensor::Quant {
9047 bytes: cat,
9048 qtype: QT_Q4_0,
9049 row_bytes: rb0,
9050 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
9051 scale: 1.0,
9052 rp: false,
9053 #[cfg(memra_cutlass)]
9054 cutlass: None,
9055 fp8: None,
9056 blk: None,
9057 rp4: None,
9058 f16: None,
9059 }))
9060 }
9061
9062 fn full_width_rope_only(
9080 kernel: &str,
9081 n_rot: usize,
9082 head_dim: usize,
9083 ) -> Result<(), Box<dyn std::error::Error>> {
9084 if n_rot == head_dim {
9085 return Ok(());
9086 }
9087 Err(format!(
9088 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
9089 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
9090 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
9091 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
9092 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
9093 )
9094 .into())
9095 }
9096
9097 #[allow(clippy::too_many_arguments)]
9101 pub fn rms_norm_qkv_rope_cat(
9102 &self,
9103 qkv: &CudaSlice<f32>,
9104 wq: &CudaSlice<f32>,
9105 wk: &CudaSlice<f32>,
9106 wv: &CudaSlice<f32>,
9107 q: &mut CudaSlice<f32>,
9108 k: &mut CudaSlice<f32>,
9109 v: &mut CudaSlice<f32>,
9110 head_dim: usize,
9111 n_rot: usize,
9112 rq: usize,
9113 rk: usize,
9114 pos: &CudaSlice<i32>,
9115 nh_q: usize,
9116 nh_k: usize,
9117 base: f32,
9118 freq_scale: f32,
9119 ff: Option<&CudaSlice<f32>>,
9120 eps: f32,
9121 ) -> Result<(), Box<dyn std::error::Error>> {
9122 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
9123 let rows = rq + rk + rk;
9124 let theta_scale = base.powf(-2.0 / head_dim as f32);
9125 let (nc, rqi, rki, nhq, nhk) = (
9126 head_dim as i32,
9127 rq as i32,
9128 rk as i32,
9129 nh_q as i32,
9130 nh_k as i32,
9131 );
9132 if Self::pdl_on() {
9133 use cudarc::driver::{DevicePtr, DevicePtrMut};
9134 let s = &self.gpu.stream();
9135 let (pqkv, _g0) = qkv.device_ptr(s);
9136 let (pwq, _g1) = wq.device_ptr(s);
9137 let (pwk, _g2) = wk.device_ptr(s);
9138 let (pwv, _g3) = wv.device_ptr(s);
9139 let (pq, _g4) = q.device_ptr_mut(s);
9140 let (pk, _g5) = k.device_ptr_mut(s);
9141 let (pv, _g6) = v.device_ptr_mut(s);
9142 let (ppos, _g7) = pos.device_ptr(s);
9143 let (pff, _g8) = match ff {
9144 Some(t) => {
9145 let (p, g) = t.device_ptr(s);
9146 (p, Some(g))
9147 }
9148 None => (0, None),
9149 };
9150 let mut ps = [
9151 &pqkv as *const _ as *mut std::ffi::c_void,
9152 &pwq as *const _ as *mut _,
9153 &pwk as *const _ as *mut _,
9154 &pwv as *const _ as *mut _,
9155 &pq as *const _ as *mut _,
9156 &pk as *const _ as *mut _,
9157 &pv as *const _ as *mut _,
9158 &nc as *const _ as *mut _,
9159 &rqi as *const _ as *mut _,
9160 &rki as *const _ as *mut _,
9161 &ppos as *const _ as *mut _,
9162 &nhq as *const _ as *mut _,
9163 &nhk as *const _ as *mut _,
9164 &theta_scale as *const _ as *mut _,
9165 &freq_scale as *const _ as *mut _,
9166 &pff as *const _ as *mut _,
9167 &eps as *const _ as *mut _,
9168 ];
9169 unsafe {
9170 self.launch_pdl(
9171 "rms_norm_qkv_rope_cat_f32",
9172 (rows as u32, 1, 1),
9173 (rms_block(), 1, 1),
9174 &mut ps,
9175 )?;
9176 }
9177 return Ok(());
9178 }
9179 let f = self.func("rms_norm_qkv_rope_cat_f32");
9180 let cfg = LaunchConfig {
9181 grid_dim: (rows as u32, 1, 1),
9182 block_dim: (rms_block(), 1, 1),
9183 shared_mem_bytes: 0,
9184 };
9185 let __s_b = self.gpu.stream();
9186 let mut b = __s_b.launch_builder(&f);
9187 match ff {
9188 Some(t) => {
9189 b.arg(qkv)
9190 .arg(wq)
9191 .arg(wk)
9192 .arg(wv)
9193 .arg(&mut *q)
9194 .arg(&mut *k)
9195 .arg(&mut *v)
9196 .arg(&nc)
9197 .arg(&rqi)
9198 .arg(&rki)
9199 .arg(pos)
9200 .arg(&nhq)
9201 .arg(&nhk)
9202 .arg(&theta_scale)
9203 .arg(&freq_scale)
9204 .arg(t)
9205 .arg(&eps);
9206 unsafe {
9207 b.launch(cfg)?;
9208 }
9209 }
9210 None => {
9211 let null: u64 = 0;
9212 b.arg(qkv)
9213 .arg(wq)
9214 .arg(wk)
9215 .arg(wv)
9216 .arg(&mut *q)
9217 .arg(&mut *k)
9218 .arg(&mut *v)
9219 .arg(&nc)
9220 .arg(&rqi)
9221 .arg(&rki)
9222 .arg(pos)
9223 .arg(&nhq)
9224 .arg(&nhk)
9225 .arg(&theta_scale)
9226 .arg(&freq_scale)
9227 .arg(&null)
9228 .arg(&eps);
9229 unsafe {
9230 b.launch(cfg)?;
9231 }
9232 }
9233 }
9234 Ok(())
9235 }
9236
9237 #[allow(clippy::too_many_arguments)]
9241 pub fn rms_norm_qkv_rope(
9242 &self,
9243 q0: &CudaSlice<f32>,
9244 k0: &CudaSlice<f32>,
9245 v0: &CudaSlice<f32>,
9246 wq: &CudaSlice<f32>,
9247 wk: &CudaSlice<f32>,
9248 wv: &CudaSlice<f32>,
9249 q: &mut CudaSlice<f32>,
9250 k: &mut CudaSlice<f32>,
9251 v: &mut CudaSlice<f32>,
9252 head_dim: usize,
9253 n_rot: usize,
9254 rq: usize,
9255 rk: usize,
9256 pos: &CudaSlice<i32>,
9257 nh_q: usize,
9258 nh_k: usize,
9259 base: f32,
9260 freq_scale: f32,
9261 ff: Option<&CudaSlice<f32>>,
9262 eps: f32,
9263 ) -> Result<(), Box<dyn std::error::Error>> {
9264 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
9265 let f = self.func("rms_norm_qkv_rope_f32");
9266 let rows = rq + rk + rk; let cfg = LaunchConfig {
9268 grid_dim: (rows as u32, 1, 1),
9269 block_dim: (rms_block(), 1, 1),
9270 shared_mem_bytes: 0,
9271 };
9272 let theta_scale = base.powf(-2.0 / head_dim as f32);
9273 let (nc, rqi, rki, nhq, nhk) = (
9274 head_dim as i32,
9275 rq as i32,
9276 rk as i32,
9277 nh_q as i32,
9278 nh_k as i32,
9279 );
9280 let __s_b = self.gpu.stream();
9281 let mut b = __s_b.launch_builder(&f);
9282 match ff {
9283 Some(t) => {
9284 b.arg(q0)
9285 .arg(k0)
9286 .arg(v0)
9287 .arg(wq)
9288 .arg(wk)
9289 .arg(wv)
9290 .arg(&mut *q)
9291 .arg(&mut *k)
9292 .arg(&mut *v)
9293 .arg(&nc)
9294 .arg(&rqi)
9295 .arg(&rki)
9296 .arg(pos)
9297 .arg(&nhq)
9298 .arg(&nhk)
9299 .arg(&theta_scale)
9300 .arg(&freq_scale)
9301 .arg(t)
9302 .arg(&eps);
9303 unsafe {
9304 b.launch(cfg)?;
9305 }
9306 }
9307 None => {
9308 let null: u64 = 0;
9309 b.arg(q0)
9310 .arg(k0)
9311 .arg(v0)
9312 .arg(wq)
9313 .arg(wk)
9314 .arg(wv)
9315 .arg(&mut *q)
9316 .arg(&mut *k)
9317 .arg(&mut *v)
9318 .arg(&nc)
9319 .arg(&rqi)
9320 .arg(&rki)
9321 .arg(pos)
9322 .arg(&nhq)
9323 .arg(&nhk)
9324 .arg(&theta_scale)
9325 .arg(&freq_scale)
9326 .arg(&null)
9327 .arg(&eps);
9328 unsafe {
9329 b.launch(cfg)?;
9330 }
9331 }
9332 }
9333 Ok(())
9334 }
9335
9336 #[allow(clippy::too_many_arguments)]
9342 pub fn rms_norm_qkv_rope_append_dc(
9343 &self,
9344 q0: &CudaSlice<f32>,
9345 k0: &CudaSlice<f32>,
9346 v0: &CudaSlice<f32>,
9347 wq: &CudaSlice<f32>,
9348 wk: &CudaSlice<f32>,
9349 wv: &CudaSlice<f32>,
9350 q: &mut CudaSlice<f32>,
9351 k: &mut CudaSlice<f32>,
9352 v: &mut CudaSlice<f32>,
9353 head_dim: usize,
9354 n_rot: usize,
9355 rq: usize,
9356 rk: usize,
9357 pos: &CudaSlice<i32>,
9358 nh_q: usize,
9359 nh_k: usize,
9360 base: f32,
9361 freq_scale: f32,
9362 ff: Option<&CudaSlice<f32>>,
9363 eps: f32,
9364 kc: &mut CudaSlice<u8>,
9365 vc: &mut CudaSlice<u8>,
9366 t_dev: &CudaSlice<i32>,
9367 k_tok_bytes: usize,
9368 v_tok_bytes: usize,
9369 g: bool,
9370 ) -> Result<(), Box<dyn std::error::Error>> {
9371 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
9372 let rows = rq + rk + rk;
9373 let theta_scale = base.powf(-2.0 / head_dim as f32);
9374 let (nc, rqi, rki, nhq, nhk) = (
9375 head_dim as i32,
9376 rq as i32,
9377 rk as i32,
9378 nh_q as i32,
9379 nh_k as i32,
9380 );
9381 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9382 if Self::pdl_on() && Self::pdl_wb_on() {
9383 use cudarc::driver::{DevicePtr, DevicePtrMut};
9384 let s = &self.gpu.stream();
9385 let (p0, _a0) = q0.device_ptr(s);
9386 let (p1, _a1) = k0.device_ptr(s);
9387 let (p2, _a2) = v0.device_ptr(s);
9388 let (pwq, _a3) = wq.device_ptr(s);
9389 let (pwk, _a4) = wk.device_ptr(s);
9390 let (pwv, _a5) = wv.device_ptr(s);
9391 let (pq, _a6) = q.device_ptr_mut(s);
9392 let (pk, _a7) = k.device_ptr_mut(s);
9393 let (pv, _a8) = v.device_ptr_mut(s);
9394 let (pp, _a9) = pos.device_ptr(s);
9395 let pff: u64 = match ff {
9396 Some(t) => {
9397 let (p, _gg) = t.device_ptr(s);
9398 p as u64
9399 }
9400 None => 0,
9401 };
9402 let (pkc, _a10) = kc.device_ptr_mut(s);
9403 let (pvc, _a11) = vc.device_ptr_mut(s);
9404 let (pt, _a12) = t_dev.device_ptr(s);
9405 let mut ps = [
9406 &p0 as *const _ as *mut std::ffi::c_void,
9407 &p1 as *const _ as *mut _,
9408 &p2 as *const _ as *mut _,
9409 &pwq as *const _ as *mut _,
9410 &pwk as *const _ as *mut _,
9411 &pwv as *const _ as *mut _,
9412 &pq as *const _ as *mut _,
9413 &pk as *const _ as *mut _,
9414 &pv as *const _ as *mut _,
9415 &nc as *const _ as *mut _,
9416 &rqi as *const _ as *mut _,
9417 &rki as *const _ as *mut _,
9418 &pp as *const _ as *mut _,
9419 &nhq as *const _ as *mut _,
9420 &nhk as *const _ as *mut _,
9421 &theta_scale as *const _ as *mut _,
9422 &freq_scale as *const _ as *mut _,
9423 &pff as *const _ as *mut _,
9424 &eps as *const _ as *mut _,
9425 &pkc as *const _ as *mut _,
9426 &pvc as *const _ as *mut _,
9427 &pt as *const _ as *mut _,
9428 &ktb as *const _ as *mut _,
9429 &vtb as *const _ as *mut _,
9430 ];
9431 unsafe {
9432 self.launch_pdl_flash(
9433 g,
9434 "rms_norm_qkv_rope_append_dc_f32",
9435 (rows as u32, 1, 1),
9436 (rms_block(), 1, 1),
9437 0,
9438 &mut ps,
9439 )?;
9440 }
9441 return Ok(());
9442 }
9443 let f = if g {
9444 self.func_g("rms_norm_qkv_rope_append_dc_f32")
9445 } else {
9446 self.func("rms_norm_qkv_rope_append_dc_f32")
9447 };
9448 let cfg = LaunchConfig {
9449 grid_dim: (rows as u32, 1, 1),
9450 block_dim: (rms_block(), 1, 1),
9451 shared_mem_bytes: 0,
9452 };
9453 let __s_b = self.gpu.stream();
9454 let mut b = __s_b.launch_builder(&f);
9455 match ff {
9456 Some(t) => {
9457 b.arg(q0)
9458 .arg(k0)
9459 .arg(v0)
9460 .arg(wq)
9461 .arg(wk)
9462 .arg(wv)
9463 .arg(&mut *q)
9464 .arg(&mut *k)
9465 .arg(&mut *v)
9466 .arg(&nc)
9467 .arg(&rqi)
9468 .arg(&rki)
9469 .arg(pos)
9470 .arg(&nhq)
9471 .arg(&nhk)
9472 .arg(&theta_scale)
9473 .arg(&freq_scale)
9474 .arg(t)
9475 .arg(&eps)
9476 .arg(&mut *kc)
9477 .arg(&mut *vc)
9478 .arg(t_dev)
9479 .arg(&ktb)
9480 .arg(&vtb);
9481 unsafe {
9482 b.launch(cfg)?;
9483 }
9484 }
9485 None => {
9486 let null: u64 = 0;
9487 b.arg(q0)
9488 .arg(k0)
9489 .arg(v0)
9490 .arg(wq)
9491 .arg(wk)
9492 .arg(wv)
9493 .arg(&mut *q)
9494 .arg(&mut *k)
9495 .arg(&mut *v)
9496 .arg(&nc)
9497 .arg(&rqi)
9498 .arg(&rki)
9499 .arg(pos)
9500 .arg(&nhq)
9501 .arg(&nhk)
9502 .arg(&theta_scale)
9503 .arg(&freq_scale)
9504 .arg(&null)
9505 .arg(&eps)
9506 .arg(&mut *kc)
9507 .arg(&mut *vc)
9508 .arg(t_dev)
9509 .arg(&ktb)
9510 .arg(&vtb);
9511 unsafe {
9512 b.launch(cfg)?;
9513 }
9514 }
9515 }
9516 Ok(())
9517 }
9518
9519 #[allow(clippy::too_many_arguments)]
9527 pub fn rms_norm_qkv_rope_append(
9528 &self,
9529 q0: &CudaSlice<f32>,
9530 k0: &CudaSlice<f32>,
9531 v0: &CudaSlice<f32>,
9532 wq: &CudaSlice<f32>,
9533 wk: &CudaSlice<f32>,
9534 wv: &CudaSlice<f32>,
9535 q: &mut CudaSlice<f32>,
9536 k: &mut CudaSlice<f32>,
9537 v: &mut CudaSlice<f32>,
9538 head_dim: usize,
9539 n_rot: usize,
9540 rq: usize,
9541 rk: usize,
9542 pos: &CudaSlice<i32>,
9543 nh_q: usize,
9544 nh_k: usize,
9545 base: f32,
9546 freq_scale: f32,
9547 ff: Option<&CudaSlice<f32>>,
9548 eps: f32,
9549 kc: &mut CudaSlice<u8>,
9550 vc: &mut CudaSlice<u8>,
9551 t: usize,
9552 k_tok_bytes: usize,
9553 v_tok_bytes: usize,
9554 g: bool,
9555 ) -> Result<(), Box<dyn std::error::Error>> {
9556 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
9557 let rows = rq + rk + rk;
9558 let theta_scale = base.powf(-2.0 / head_dim as f32);
9559 let (nc, rqi, rki, nhq, nhk) = (
9560 head_dim as i32,
9561 rq as i32,
9562 rk as i32,
9563 nh_q as i32,
9564 nh_k as i32,
9565 );
9566 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9567 let ti = t as i32;
9568 if Self::pdl_on() && Self::pdl_wb_on() {
9569 use cudarc::driver::{DevicePtr, DevicePtrMut};
9570 let s = &self.gpu.stream();
9571 let (p0, _a0) = q0.device_ptr(s);
9572 let (p1, _a1) = k0.device_ptr(s);
9573 let (p2, _a2) = v0.device_ptr(s);
9574 let (pwq, _a3) = wq.device_ptr(s);
9575 let (pwk, _a4) = wk.device_ptr(s);
9576 let (pwv, _a5) = wv.device_ptr(s);
9577 let (pq, _a6) = q.device_ptr_mut(s);
9578 let (pk, _a7) = k.device_ptr_mut(s);
9579 let (pv, _a8) = v.device_ptr_mut(s);
9580 let (pp, _a9) = pos.device_ptr(s);
9581 let pff: u64 = match ff {
9582 Some(t) => {
9583 let (p, _gg) = t.device_ptr(s);
9584 p as u64
9585 }
9586 None => 0,
9587 };
9588 let (pkc, _a10) = kc.device_ptr_mut(s);
9589 let (pvc, _a11) = vc.device_ptr_mut(s);
9590 let mut ps = [
9591 &p0 as *const _ as *mut std::ffi::c_void,
9592 &p1 as *const _ as *mut _,
9593 &p2 as *const _ as *mut _,
9594 &pwq as *const _ as *mut _,
9595 &pwk as *const _ as *mut _,
9596 &pwv as *const _ as *mut _,
9597 &pq as *const _ as *mut _,
9598 &pk as *const _ as *mut _,
9599 &pv as *const _ as *mut _,
9600 &nc as *const _ as *mut _,
9601 &rqi as *const _ as *mut _,
9602 &rki as *const _ as *mut _,
9603 &pp as *const _ as *mut _,
9604 &nhq as *const _ as *mut _,
9605 &nhk as *const _ as *mut _,
9606 &theta_scale as *const _ as *mut _,
9607 &freq_scale as *const _ as *mut _,
9608 &pff as *const _ as *mut _,
9609 &eps as *const _ as *mut _,
9610 &pkc as *const _ as *mut _,
9611 &pvc as *const _ as *mut _,
9612 &ti as *const _ as *mut _,
9613 &ktb as *const _ as *mut _,
9614 &vtb as *const _ as *mut _,
9615 ];
9616 unsafe {
9617 self.launch_pdl_flash(
9618 g,
9619 "rms_norm_qkv_rope_append_f32",
9620 (rows as u32, 1, 1),
9621 (rms_block(), 1, 1),
9622 0,
9623 &mut ps,
9624 )?;
9625 }
9626 return Ok(());
9627 }
9628 let f = if g {
9629 self.func_g("rms_norm_qkv_rope_append_f32")
9630 } else {
9631 self.func("rms_norm_qkv_rope_append_f32")
9632 };
9633 let cfg = LaunchConfig {
9634 grid_dim: (rows as u32, 1, 1),
9635 block_dim: (rms_block(), 1, 1),
9636 shared_mem_bytes: 0,
9637 };
9638 let __s_b = self.gpu.stream();
9639 let mut b = __s_b.launch_builder(&f);
9640 let null: u64 = 0;
9641 b.arg(q0)
9642 .arg(k0)
9643 .arg(v0)
9644 .arg(wq)
9645 .arg(wk)
9646 .arg(wv)
9647 .arg(&mut *q)
9648 .arg(&mut *k)
9649 .arg(&mut *v)
9650 .arg(&nc)
9651 .arg(&rqi)
9652 .arg(&rki)
9653 .arg(pos)
9654 .arg(&nhq)
9655 .arg(&nhk)
9656 .arg(&theta_scale)
9657 .arg(&freq_scale);
9658 match ff {
9659 Some(t) => {
9660 b.arg(t);
9661 }
9662 None => {
9663 b.arg(&null);
9664 }
9665 }
9666 b.arg(&eps)
9667 .arg(&mut *kc)
9668 .arg(&mut *vc)
9669 .arg(&ti)
9670 .arg(&ktb)
9671 .arg(&vtb);
9672 unsafe {
9673 b.launch(cfg)?;
9674 }
9675 Ok(())
9676 }
9677
9678 pub fn add_q8_1(
9679 &self,
9680 a: &CudaSlice<f32>,
9681 b: &CudaSlice<f32>,
9682 res: &mut CudaSlice<f32>,
9683 ncols: usize,
9684 nrows: usize,
9685 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9686 debug_assert!(ncols % 128 == 0);
9687 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9688 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9689 let f = self.func("add_q8_1_f32");
9690 let cfg = LaunchConfig {
9691 grid_dim: (nrows as u32, 1, 1),
9692 block_dim: (rms_block(), 1, 1),
9693 shared_mem_bytes: 0,
9694 };
9695 let nc = ncols as i32;
9696 let __s_b2 = self.gpu.stream();
9697 let mut b2 = __s_b2.launch_builder(&f);
9698 b2.arg(a)
9699 .arg(b)
9700 .arg(&mut *res)
9701 .arg(&mut out_q)
9702 .arg(&mut out_d)
9703 .arg(&nc);
9704 unsafe {
9705 b2.launch(cfg)?;
9706 }
9707 Ok((out_q, out_d))
9708 }
9709
9710 pub fn rms_pre_add_q8_1(
9714 &self,
9715 a: &CudaSlice<f32>,
9716 wa: &CudaSlice<f32>,
9717 b: &CudaSlice<f32>,
9718 res: &mut CudaSlice<f32>,
9719 ncols: usize,
9720 nrows: usize,
9721 eps: f32,
9722 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9723 debug_assert!(ncols % 128 == 0);
9724 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9725 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9726 let f = self.func("rms_pre_add_q8_1_f32");
9727 let cfg = LaunchConfig {
9728 grid_dim: (nrows as u32, 1, 1),
9729 block_dim: (rms_block(), 1, 1),
9730 shared_mem_bytes: 0,
9731 };
9732 let (nc, ep) = (ncols as i32, eps);
9733 let __s_b2 = self.gpu.stream();
9734 let mut b2 = __s_b2.launch_builder(&f);
9735 b2.arg(a)
9736 .arg(wa)
9737 .arg(b)
9738 .arg(&mut *res)
9739 .arg(&mut out_q)
9740 .arg(&mut out_d)
9741 .arg(&nc)
9742 .arg(&ep);
9743 unsafe {
9744 b2.launch(cfg)?;
9745 }
9746 Ok((out_q, out_d))
9747 }
9748
9749 pub fn l2_v2_on(ncols: usize) -> bool {
9753 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9754 }
9755
9756 pub fn l2_norm_pp(
9757 &self,
9758 x: &CudaSlice<f32>,
9759 dst: &mut CudaSlice<f32>,
9760 dst16: Option<&mut CudaSlice<u8>>,
9761 ncols: usize,
9762 nrows: usize,
9763 eps: f32,
9764 ) -> Result<(), Box<dyn std::error::Error>> {
9765 if Self::l2_v2_on(ncols) {
9766 let f = self.func("l2_norm_pp_v2_f32");
9767 let rows_per_block = 8u32; let cfg = LaunchConfig {
9769 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9770 block_dim: (256, 1, 1),
9771 shared_mem_bytes: 0,
9772 };
9773 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9774 let d16: u64 = match dst16 {
9776 Some(d) => self.addr_u8(d),
9777 None => 0,
9778 };
9779 let __s_b = self.gpu.stream();
9780 let mut b = __s_b.launch_builder(&f);
9781 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9782 unsafe {
9783 b.launch(cfg)?;
9784 }
9785 return Ok(());
9786 }
9787 self.l2_norm(x, dst, ncols, nrows, eps)
9788 }
9789
9790 pub fn l2_norm(
9791 &self,
9792 x: &CudaSlice<f32>,
9793 dst: &mut CudaSlice<f32>,
9794 ncols: usize,
9795 nrows: usize,
9796 eps: f32,
9797 ) -> Result<(), Box<dyn std::error::Error>> {
9798 let f = self.func("l2_norm_f32");
9799 let cfg = LaunchConfig {
9800 grid_dim: (nrows as u32, 1, 1),
9801 block_dim: (256, 1, 1),
9802 shared_mem_bytes: 0,
9803 };
9804 let (nc, e) = (ncols as i32, eps);
9805 let __s_b = self.gpu.stream();
9806 let mut b = __s_b.launch_builder(&f);
9807 b.arg(x).arg(dst).arg(&nc).arg(&e);
9808 unsafe {
9809 b.launch(cfg)?;
9810 }
9811 Ok(())
9812 }
9813
9814 pub fn l2_norm_decode(
9820 &self,
9821 x: &CudaSlice<f32>,
9822 dst: &mut CudaSlice<f32>,
9823 ncols: usize,
9824 nrows: usize,
9825 eps: f32,
9826 ) -> Result<(), Box<dyn std::error::Error>> {
9827 let f = self.func("l2_norm_f32");
9828 let cfg = LaunchConfig {
9829 grid_dim: (nrows as u32, 1, 1),
9830 block_dim: (32, 1, 1),
9831 shared_mem_bytes: 0,
9832 };
9833 let (nc, e) = (ncols as i32, eps);
9834 let __s_b = self.gpu.stream();
9835 let mut b = __s_b.launch_builder(&f);
9836 b.arg(x).arg(dst).arg(&nc).arg(&e);
9837 unsafe {
9838 b.launch(cfg)?;
9839 }
9840 Ok(())
9841 }
9842
9843 pub fn rope_neox(
9845 &self,
9846 x: &mut CudaSlice<f32>,
9847 pos: &CudaSlice<i32>,
9848 head_dim: usize,
9849 n_dims: usize,
9850 n_heads: usize,
9851 n_tokens: usize,
9852 freq_base: f32,
9853 freq_scale: f32,
9854 ) -> Result<(), Box<dyn std::error::Error>> {
9855 let f = self.func("rope_neox_f32");
9856 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9857 let grid = (n_heads * n_tokens) as u32;
9858 let cfg = LaunchConfig {
9859 grid_dim: (grid, 1, 1),
9860 block_dim: ((head_dim / 2) as u32, 1, 1),
9861 shared_mem_bytes: 0,
9862 };
9863 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9864 let __s_b = self.gpu.stream();
9865 let mut b = __s_b.launch_builder(&f);
9866 b.arg(x)
9867 .arg(pos)
9868 .arg(&hd)
9869 .arg(&nd)
9870 .arg(&nh)
9871 .arg(&theta_scale)
9872 .arg(&freq_scale);
9873 unsafe {
9874 b.launch(cfg)?;
9875 }
9876 Ok(())
9877 }
9878
9879 pub fn rope_neox_ff(
9881 &self,
9882 x: &mut CudaSlice<f32>,
9883 pos: &CudaSlice<i32>,
9884 head_dim: usize,
9885 n_dims: usize,
9886 n_heads: usize,
9887 n_tokens: usize,
9888 freq_base: f32,
9889 freq_scale: f32,
9890 ff: &CudaSlice<f32>,
9891 ) -> Result<(), Box<dyn std::error::Error>> {
9892 let f = self.func("rope_neox_ff_f32");
9893 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9894 let grid = (n_heads * n_tokens) as u32;
9895 let cfg = LaunchConfig {
9896 grid_dim: (grid, 1, 1),
9897 block_dim: ((head_dim / 2) as u32, 1, 1),
9898 shared_mem_bytes: 0,
9899 };
9900 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9901 let __s_b = self.gpu.stream();
9902 let mut b = __s_b.launch_builder(&f);
9903 b.arg(x)
9904 .arg(pos)
9905 .arg(&hd)
9906 .arg(&nd)
9907 .arg(&nh)
9908 .arg(&theta_scale)
9909 .arg(&freq_scale)
9910 .arg(ff);
9911 unsafe {
9912 b.launch(cfg)?;
9913 }
9914 Ok(())
9915 }
9916
9917 #[allow(clippy::too_many_arguments)]
9919 pub fn rope_neox2(
9920 &self,
9921 q: &mut CudaSlice<f32>,
9922 k: &mut CudaSlice<f32>,
9923 pos: &CudaSlice<i32>,
9924 head_dim: usize,
9925 n_dims: usize,
9926 nh_q: usize,
9927 nh_k: usize,
9928 n_tokens: usize,
9929 freq_base: f32,
9930 freq_scale: f32,
9931 ff: Option<&CudaSlice<f32>>,
9932 ) -> Result<(), Box<dyn std::error::Error>> {
9933 let f = self.func("rope_neox2_f32");
9934 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9935 let grid = ((nh_q + nh_k) * n_tokens) as u32;
9936 let cfg = LaunchConfig {
9937 grid_dim: (grid, 1, 1),
9938 block_dim: ((head_dim / 2) as u32, 1, 1),
9939 shared_mem_bytes: 0,
9940 };
9941 let (hd, nd, nq, nk, nt) = (
9942 head_dim as i32,
9943 n_dims as i32,
9944 nh_q as i32,
9945 nh_k as i32,
9946 n_tokens as i32,
9947 );
9948 let __s_b = self.gpu.stream();
9949 let mut b = __s_b.launch_builder(&f);
9950 b.arg(q)
9951 .arg(k)
9952 .arg(pos)
9953 .arg(&hd)
9954 .arg(&nd)
9955 .arg(&nq)
9956 .arg(&nk)
9957 .arg(&nt)
9958 .arg(&theta_scale)
9959 .arg(&freq_scale);
9960 match ff {
9961 Some(ffv) => {
9962 b.arg(ffv);
9963 unsafe {
9964 b.launch(cfg)?;
9965 }
9966 }
9967 None => {
9968 let null: u64 = 0;
9969 b.arg(&null);
9970 unsafe {
9971 b.launch(cfg)?;
9972 }
9973 }
9974 }
9975 Ok(())
9976 }
9977
9978 pub fn gelu_tanh_mul(
9980 &self,
9981 gate: &CudaSlice<f32>,
9982 up: &CudaSlice<f32>,
9983 dst: &mut CudaSlice<f32>,
9984 n: usize,
9985 ) -> Result<(), Box<dyn std::error::Error>> {
9986 let f = self.func("gelu_tanh_mul_f32");
9987 let cfg = LaunchConfig::for_num_elems(n as u32);
9988 let ni = n as i32;
9989 let __s_b = self.gpu.stream();
9990 let mut b = __s_b.launch_builder(&f);
9991 b.arg(gate).arg(up).arg(dst).arg(&ni);
9992 unsafe {
9993 b.launch(cfg)?;
9994 }
9995 Ok(())
9996 }
9997
9998 pub fn silu_mul(
9999 &self,
10000 gate: &CudaSlice<f32>,
10001 up: &CudaSlice<f32>,
10002 dst: &mut CudaSlice<f32>,
10003 n: usize,
10004 ) -> Result<(), Box<dyn std::error::Error>> {
10005 let f = self.func("silu_mul_f32");
10006 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10008 let ni = n as i32;
10009 let __s_b = self.gpu.stream();
10010 let mut b = __s_b.launch_builder(&f);
10011 b.arg(gate).arg(up).arg(dst).arg(&ni);
10012 unsafe {
10013 b.launch(cfg)?;
10014 }
10015 Ok(())
10016 }
10017
10018 pub fn silu_mul_f16out(
10021 &self,
10022 gate: &CudaSlice<f32>,
10023 up: &CudaSlice<f32>,
10024 dst: &mut CudaSlice<f32>,
10025 dst16: &mut CudaSlice<u8>,
10026 n: usize,
10027 ) -> Result<(), Box<dyn std::error::Error>> {
10028 let f = self.func("silu_mul_f16out_f32");
10029 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10030 let ni = n as i32;
10031 let __s_b = self.gpu.stream();
10032 let mut b = __s_b.launch_builder(&f);
10033 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
10034 unsafe {
10035 b.launch(cfg)?;
10036 }
10037 Ok(())
10038 }
10039
10040 pub fn silu_mul_scaled(
10047 &self,
10048 gate: &CudaSlice<f32>,
10049 up: &CudaSlice<f32>,
10050 gs: f32,
10051 us: f32,
10052 dst: &mut CudaSlice<f32>,
10053 n: usize,
10054 ) -> Result<(), Box<dyn std::error::Error>> {
10055 let f = self.func("silu_mul_scaled_f32");
10056 let cfg = LaunchConfig::for_num_elems(n as u32);
10057 let ni = n as i32;
10058 let (gsf, usf) = (gs, us);
10059 let __s_b = self.gpu.stream();
10060 let mut b = __s_b.launch_builder(&f);
10061 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
10062 unsafe {
10063 b.launch(cfg)?;
10064 }
10065 Ok(())
10066 }
10067
10068 #[allow(clippy::too_many_arguments)]
10072 pub fn swigluoai_mul_scaled(
10073 &self,
10074 gate: &CudaSlice<f32>,
10075 up: &CudaSlice<f32>,
10076 gs: f32,
10077 us: f32,
10078 alpha: f32,
10079 limit: f32,
10080 dst: &mut CudaSlice<f32>,
10081 n: usize,
10082 ) -> Result<(), Box<dyn std::error::Error>> {
10083 let f = self.func("swigluoai_mul_scaled_f32");
10084 let cfg = LaunchConfig::for_num_elems(n as u32);
10085 let ni = n as i32;
10086 let __s_b = self.gpu.stream();
10087 let mut b = __s_b.launch_builder(&f);
10088 b.arg(gate)
10089 .arg(up)
10090 .arg(&gs)
10091 .arg(&us)
10092 .arg(&alpha)
10093 .arg(&limit)
10094 .arg(dst)
10095 .arg(&ni);
10096 unsafe {
10097 b.launch(cfg)?;
10098 }
10099 Ok(())
10100 }
10101
10102 pub fn silu_mul_scaled_q8_1(
10110 &self,
10111 gate: &CudaSlice<f32>,
10112 up: &CudaSlice<f32>,
10113 gs: f32,
10114 us: f32,
10115 n: usize,
10116 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10117 let f = self.func("silu_mul_scaled_q8_1");
10118 let nblk = n / 32;
10119 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);
10123 let (gsf, usf, ni) = (gs, us, n as i32);
10124 let __s_b = self.gpu.stream();
10125 let mut b = __s_b.launch_builder(&f);
10126 b.arg(gate)
10127 .arg(up)
10128 .arg(&gsf)
10129 .arg(&usf)
10130 .arg(&mut aq)
10131 .arg(&mut ad)
10132 .arg(&ni);
10133 unsafe {
10134 b.launch(cfg)?;
10135 }
10136 Ok((aq, ad))
10137 }
10138
10139 pub fn add(
10140 &self,
10141 a: &CudaSlice<f32>,
10142 b_in: &CudaSlice<f32>,
10143 dst: &mut CudaSlice<f32>,
10144 n: usize,
10145 ) -> Result<(), Box<dyn std::error::Error>> {
10146 let f = self.func("add_f32");
10147 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10149 let ni = n as i32;
10150 let __s_bld = self.gpu.stream();
10151 let mut bld = __s_bld.launch_builder(&f);
10152 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
10153 unsafe {
10154 bld.launch(cfg)?;
10155 }
10156 Ok(())
10157 }
10158
10159 pub fn mul(
10160 &self,
10161 a: &CudaSlice<f32>,
10162 b_in: &CudaSlice<f32>,
10163 dst: &mut CudaSlice<f32>,
10164 n: usize,
10165 ) -> Result<(), Box<dyn std::error::Error>> {
10166 let f = self.func("mul_f32");
10167 let cfg = LaunchConfig::for_num_elems(n as u32);
10168 let ni = n as i32;
10169 let __s_bld = self.gpu.stream();
10170 let mut bld = __s_bld.launch_builder(&f);
10171 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
10172 unsafe {
10173 bld.launch(cfg)?;
10174 }
10175 Ok(())
10176 }
10177
10178 pub fn matmul(
10181 &self,
10182 w: &crate::model::GpuTensor,
10183 x: &CudaSlice<f32>,
10184 m: usize,
10185 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10186 use crate::model::GpuTensor;
10187 let in_f = w.in_features();
10188 let out_f = w.out_features();
10189 #[allow(non_snake_case)]
10197 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
10200 usize::MAX
10201 } else {
10202 16usize
10203 };
10204
10205 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
10230 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
10231 return Ok(y);
10232 }
10233 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
10240 return Ok(y);
10241 }
10242 if let Some(y) = self.try_f16_gemm(w, x, m)? {
10245 return Ok(y);
10246 }
10247 }
10248 if let GpuTensor::Quant { qtype, .. } = w {
10263 if *qtype == QT_F8_E4M3_BLK {
10264 if m >= GEMM_M_THRESHOLD {
10265 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
10266 return Ok(y);
10267 }
10268 }
10269 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10270 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10271 return Ok(y);
10272 }
10273 }
10274 }
10275 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
10276 return self.qmatvec_mmq(w, x, m);
10277 }
10278 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
10279 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10280 return self.qmatvec_gemm(w, &aq, &ad, m);
10281 }
10282 if m >= GEMM_M_THRESHOLD {
10285 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
10286 return Ok(y);
10287 }
10288 }
10289 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
10293 if m == 1 && fast {
10298 if let GpuTensor::Quant {
10299 bytes,
10300 qtype,
10301 row_bytes,
10302 rp,
10303 rp4,
10304 scale,
10305 ..
10306 } = w
10307 {
10308 if self.mmvq_supports(*qtype) {
10309 let (bytes, rp) = match rp4 {
10313 Some(m4) => (m4, true),
10314 None => (bytes, *rp),
10315 };
10316 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10317 return self.qmatvec_mmvq(
10318 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
10319 );
10320 }
10321 }
10322 }
10323 if (2..=16).contains(&m)
10339 && fast
10340 && std::env::var("MEMRA_NO_BATCHED").is_err()
10341 && (m <= 4 || Self::b8_enabled())
10342 {
10343 let m_ok = m <= 8
10353 || matches!(w, GpuTensor::Quant { qtype, .. }
10354 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
10355 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
10356 if m_ok {
10357 if let GpuTensor::Quant {
10358 bytes,
10359 qtype,
10360 row_bytes,
10361 rp,
10362 rp4,
10363 ..
10364 } = w
10365 {
10366 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
10367 let (bytes, rp) = match rp4 {
10368 Some(m4) => (m4, true),
10369 None => (bytes, *rp),
10370 };
10371 let mcols = Self::batched_mcols(m);
10372 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10373 let mut y = self.qmatvec_mmvq_batched(
10374 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
10375 )?;
10376 if let GpuTensor::Quant { scale, .. } = w {
10377 if *scale != 1.0 {
10378 self.scale_inplace(&mut y, *scale, m * out_f)?;
10379 }
10380 }
10381 return Ok(y);
10382 }
10383 }
10384 }
10385 }
10386 if fast {
10392 if let GpuTensor::Quant {
10393 bytes,
10394 qtype,
10395 row_bytes,
10396 scale,
10397 ..
10398 } = w
10399 {
10400 if *qtype == QT_F8_E4M3 {
10401 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10402 return self.qmatvec_mmvq(
10403 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10404 );
10405 }
10406 }
10407 }
10408 let mut y = match w {
10409 GpuTensor::Quant {
10410 bytes,
10411 qtype,
10412 row_bytes,
10413 ..
10414 } if fast && *qtype == QT_Q8_0 => {
10415 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10416 }
10417 GpuTensor::Quant {
10418 bytes,
10419 qtype,
10420 row_bytes,
10421 ..
10422 } if fast && *qtype == QT_Q4_K => {
10423 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10424 }
10425 GpuTensor::Quant {
10426 bytes,
10427 qtype,
10428 row_bytes,
10429 ..
10430 } if fast && *qtype == QT_Q6_K => {
10431 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10432 }
10433 GpuTensor::Quant {
10434 bytes,
10435 qtype,
10436 row_bytes,
10437 ..
10438 } if fast && *qtype == QT_Q5_K => {
10439 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10440 }
10441 GpuTensor::Quant {
10442 bytes,
10443 qtype,
10444 row_bytes,
10445 ..
10446 } if fast && *qtype == QT_Q3_K => {
10447 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10448 }
10449 GpuTensor::Quant {
10450 bytes,
10451 qtype,
10452 row_bytes,
10453 rp,
10454 ..
10455 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10456 if *rp {
10457 "qmatvec_nvfp4_dp4a_rp"
10458 } else {
10459 "qmatvec_nvfp4_dp4a"
10460 },
10461 bytes,
10462 x,
10463 m,
10464 in_f,
10465 out_f,
10466 *row_bytes,
10467 )?,
10468 GpuTensor::Quant {
10472 bytes,
10473 qtype,
10474 row_bytes,
10475 ..
10476 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10477 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10478 }
10479 GpuTensor::Quant {
10484 bytes,
10485 qtype,
10486 row_bytes,
10487 rp,
10488 ..
10489 } =>
10490 {
10493 self.qmatvec(
10494 bytes,
10495 x,
10496 m,
10497 in_f,
10498 out_f,
10499 if *rp && *qtype == QT_NVFP4 {
10500 QT_NVFP4_RP
10501 } else {
10502 *qtype
10503 },
10504 *row_bytes,
10505 )?
10506 }
10507 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10508 GpuTensor::FloatBf16 { data, .. } => {
10511 self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10512 }
10513 };
10514 if let GpuTensor::Quant { scale, .. } = w {
10516 if *scale != 1.0 {
10517 self.scale_inplace(&mut y, *scale, m * out_f)?;
10518 }
10519 }
10520 Ok(y)
10521 }
10522
10523 pub fn stage_a_raw_needed() -> bool {
10533 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10534 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
10535 }
10536
10537 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10540 use crate::model::GpuTensor;
10541 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10542 return false;
10543 }
10544 match w {
10545 GpuTensor::Quant { qtype, .. } => {
10552 matches!(
10553 *qtype,
10554 QT_Q8_0
10555 | QT_Q4_K
10556 | QT_Q6_K
10557 | QT_Q5_K
10558 | QT_Q3_K
10559 | QT_NVFP4
10560 | QT_F8_E4M3
10561 | QT_F8_E4M3_BLK
10562 | QT_Q4_0
10563 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10564 }
10565 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10566 }
10567 }
10568
10569 pub fn matmul_pre(
10574 &self,
10575 w: &crate::model::GpuTensor,
10576 aq: &CudaSlice<i8>,
10577 ad: &CudaSlice<f32>,
10578 x_fallback: &CudaSlice<f32>,
10579 m: usize,
10580 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10581 use crate::model::GpuTensor;
10582 let x_raw_ok = x_fallback.len() >= m * w.in_features();
10588 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10591 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10592 return Ok(y);
10593 }
10594 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10597 return Ok(y);
10598 }
10599 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10601 return Ok(y);
10602 }
10603 }
10604 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10610 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10611 return Ok(y);
10612 }
10613 }
10614 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10615 return Ok(y);
10616 }
10617 if m >= 16
10622 && w.out_features() >= 128
10623 && self.mmq_supports(w)
10624 && !self.verify_exact_on()
10625 && x_raw_ok
10626 {
10627 return self.qmatvec_mmq(w, x_fallback, m);
10628 }
10629 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10632 if let Some(y) =
10633 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10634 {
10635 return Ok(y);
10636 }
10637 }
10638 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10641 return self.qmatvec_gemm(w, aq, ad, m);
10642 }
10643 if !self.uses_q8_1_fast(w) {
10662 if !x_raw_ok {
10663 return Err(format!(
10664 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
10665 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
10666 dropped the f32, so there is nothing to fall back to — pass the real f32 \
10667 activation (see Engine::rms_norm_decode, which is bit-identical to \
10668 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
10669 x_fallback.len(),
10670 m,
10671 w.in_features(),
10672 m * w.in_features()
10673 )
10674 .into());
10675 }
10676 return self.matmul(w, x_fallback, m);
10677 }
10678 let in_f = w.in_features();
10679 let out_f = w.out_features();
10680 let (bytes, qtype, row_bytes, scale, rp) = match w {
10681 GpuTensor::Quant {
10682 bytes,
10683 qtype,
10684 row_bytes,
10685 scale,
10686 rp,
10687 ..
10688 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10689 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10690 };
10691 let (mbytes, mrp) = match w {
10694 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10695 _ => (bytes, rp),
10696 };
10697 if m == 1 && self.mmvq_supports(qtype) {
10701 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10702 }
10703 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10716 && std::env::var("MEMRA_NO_BATCHED").is_err()
10717 && (m <= 4 || Self::b8_enabled())
10718 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10722 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10723 {
10724 let mcols = Self::batched_mcols(m);
10725 return self.qmatvec_mmvq_batched(
10726 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10727 );
10728 }
10729 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10735 let (b2, r2) = if qtype == QT_Q4_0 {
10736 (mbytes, mrp)
10737 } else {
10738 (bytes, rp)
10739 };
10740 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10741 }
10742 let name = match qtype {
10743 QT_Q8_0 => "qmatvec_q8_0_dp4a",
10744 QT_Q4_K => "qmatvec_q4_K_dp4a",
10745 QT_Q6_K => "qmatvec_q6_K_dp4a",
10746 QT_Q5_K => "qmatvec_q5_K_dp4a",
10747 QT_Q3_K => "qmatvec_q3_K_dp4a",
10748 QT_NVFP4 => {
10749 if rp {
10750 "qmatvec_nvfp4_dp4a_rp"
10751 } else {
10752 "qmatvec_nvfp4_dp4a"
10753 }
10754 }
10755 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10756 _ => unreachable!(),
10757 };
10758 let f = self.func(name);
10759 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10761 grid_dim: (out_f as u32, m as u32, 1),
10762 block_dim: (128, 1, 1),
10763 shared_mem_bytes: 0,
10764 };
10765 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10766 let __s_b = self.gpu.stream();
10767 let mut b = __s_b.launch_builder(&f);
10768 b.arg(bytes)
10769 .arg(aq)
10770 .arg(ad)
10771 .arg(&mut y)
10772 .arg(&inf)
10773 .arg(&outf)
10774 .arg(&mi)
10775 .arg(&rb);
10776 unsafe {
10777 b.launch(cfg)?;
10778 }
10779 if scale != 1.0 {
10780 self.scale_inplace(&mut y, scale, m * out_f)?;
10781 }
10782 Ok(y)
10783 }
10784
10785 pub fn matmul_decode_exact(
10793 &self,
10794 w: &crate::model::GpuTensor,
10795 x: &CudaSlice<f32>,
10796 m: usize,
10797 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10798 use crate::model::GpuTensor;
10799 if let GpuTensor::Float { data, .. } = w {
10807 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10808 }
10809 if let GpuTensor::FloatBf16 { data, .. } = w {
10812 let (in_f, out_f) = (w.in_features(), w.out_features());
10813 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10814 }
10815 if !self.uses_q8_1_fast(w) {
10816 return self.matmul(w, x, m);
10817 }
10818 let in_f = w.in_features();
10819 let out_f = w.out_features();
10820 let (bytes, qtype, row_bytes, scale, rp) = match w {
10821 GpuTensor::Quant {
10822 bytes,
10823 qtype,
10824 row_bytes,
10825 scale,
10826 rp,
10827 ..
10828 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10829 _ => return self.matmul(w, x, m),
10830 };
10831 let (bytes, rp) = match w {
10834 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10835 _ => (bytes, rp),
10836 };
10837 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10838 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10842 return Ok(y);
10843 }
10844 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10853 && std::env::var("MEMRA_NO_BATCHED").is_err()
10854 && (m <= 4 || Self::b8_enabled())
10855 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10858 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10859 {
10860 let mcols = Self::batched_mcols(m);
10861 return self.qmatvec_mmvq_batched(
10862 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10863 );
10864 }
10865 if self.mmvq_supports(qtype) {
10866 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10869 }
10870 self.matmul_pre(w, &aq, &ad, x, m)
10873 }
10874
10875 pub fn matmul_decode_exact_pre(
10885 &self,
10886 w: &crate::model::GpuTensor,
10887 aq: &CudaSlice<i8>,
10888 ad: &CudaSlice<f32>,
10889 m: usize,
10890 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10891 use crate::model::GpuTensor;
10892 debug_assert!(
10893 self.uses_q8_1_fast(w),
10894 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10895 );
10896 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10898 return Ok(y);
10899 }
10900 let in_f = w.in_features();
10901 let out_f = w.out_features();
10902 let (bytes, qtype, row_bytes, scale, rp) = match w {
10903 GpuTensor::Quant {
10904 bytes,
10905 qtype,
10906 row_bytes,
10907 scale,
10908 rp,
10909 ..
10910 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10911 _ => {
10912 return Err(
10913 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10914 );
10915 }
10916 };
10917 let (bytes, rp) = match w {
10919 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10920 _ => (bytes, rp),
10921 };
10922 if (2..=16).contains(&m)
10924 && self.batched_supports(qtype)
10925 && self.mmvq_supports(qtype)
10926 && std::env::var("MEMRA_NO_BATCHED").is_err()
10927 && (m <= 4 || Self::b8_enabled())
10928 && (m <= 8
10929 || qtype == QT_Q4_0
10930 || qtype == QT_Q6_K
10931 || qtype == QT_F8_E4M3
10932 || qtype == QT_NVFP4
10933 || qtype == QT_Q4_K
10934 || qtype == QT_Q5_K
10935 || qtype == QT_Q8_0)
10936 {
10937 let mcols = Self::batched_mcols(m);
10938 return self.qmatvec_mmvq_batched(
10939 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10940 );
10941 }
10942 if self.mmvq_supports(qtype) {
10943 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10944 }
10945 let x0 = self.zeros(0)?;
10948 self.matmul_pre(w, aq, ad, &x0, m)
10949 }
10950
10951 pub fn matmul_decode_exact_dual_pre(
10960 &self,
10961 w0: &crate::model::GpuTensor,
10962 w1: &crate::model::GpuTensor,
10963 aq: &CudaSlice<i8>,
10964 ad: &CudaSlice<f32>,
10965 m: usize,
10966 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10967 {
10968 use crate::model::GpuTensor;
10969 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10970 let on = *ON.get_or_init(|| {
10971 std::env::var("MEMRA_SPEC_DUAL_T")
10972 .map(|v| v != "0")
10973 .unwrap_or(true)
10974 });
10975 if !on
10976 || !(2..=7).contains(&m)
10977 || std::env::var("MEMRA_NO_BATCHED").is_ok()
10978 || !self.uses_q8_1_fast(w0)
10979 || !self.uses_q8_1_fast(w1)
10980 {
10981 return Ok(None);
10982 }
10983 if !self.mmvq_supports(QT_NVFP4) {
10988 return Ok(None);
10989 }
10990 let (in_f, out_f) = (w0.in_features(), w0.out_features());
10991 if w1.in_features() != in_f || w1.out_features() != out_f {
10992 return Ok(None);
10993 }
10994 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10995 (
10996 GpuTensor::Quant {
10997 bytes: b0,
10998 qtype: q0,
10999 row_bytes: rb0,
11000 scale: s0,
11001 rp: rp0,
11002 rp4: None,
11003 ..
11004 },
11005 GpuTensor::Quant {
11006 bytes: b1,
11007 qtype: q1,
11008 row_bytes: rb1,
11009 scale: s1,
11010 rp: rp1,
11011 rp4: None,
11012 ..
11013 },
11014 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11015 (b0, b1, *rb0, *s0, *s1, *rp0)
11016 }
11017 _ => return Ok(None),
11018 };
11019 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
11022 {
11023 return Ok(None);
11024 }
11025 let (y0, y1) =
11026 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
11027 Ok(Some(((y0, s0), (y1, s1))))
11028 }
11029
11030 pub fn matmul_decode_exact_group4_pre(
11040 &self,
11041 ws: [&crate::model::GpuTensor; 4],
11042 aq: &CudaSlice<i8>,
11043 ad: &CudaSlice<f32>,
11044 m: usize,
11045 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11046 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11047 let on = *ON.get_or_init(|| {
11048 std::env::var("MEMRA_TK_GDN_GROUP")
11049 .map(|v| v != "0")
11050 .unwrap_or(true)
11051 });
11052 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
11053 }
11054
11055 pub fn matmul_decode_exact_group3_pre(
11060 &self,
11061 ws: [&crate::model::GpuTensor; 3],
11062 aq: &CudaSlice<i8>,
11063 ad: &CudaSlice<f32>,
11064 m: usize,
11065 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11066 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11067 let on = *ON.get_or_init(|| {
11068 std::env::var("MEMRA_TK_FA_GROUP")
11069 .map(|v| v != "0")
11070 .unwrap_or(true)
11071 });
11072 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
11073 }
11074
11075 fn matmul_decode_exact_group_pre(
11079 &self,
11080 ws: &[&crate::model::GpuTensor],
11081 aq: &CudaSlice<i8>,
11082 ad: &CudaSlice<f32>,
11083 m: usize,
11084 on: bool,
11085 tag: &'static str,
11086 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11087 use crate::model::GpuTensor;
11088 if !on
11089 || !(2..=8).contains(&m)
11090 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11091 || (m > 4 && !Self::b8_enabled())
11092 || !self.mmvq_supports(QT_NVFP4)
11093 || !self.batched_supports(QT_NVFP4)
11094 {
11095 return Ok(None);
11096 }
11097 let in_f = ws[0].in_features();
11098 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
11099 for w in ws {
11100 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
11101 return Ok(None);
11102 }
11103 match w {
11104 GpuTensor::Quant {
11105 bytes,
11106 qtype,
11107 scale,
11108 rp: true,
11109 rp4: None,
11110 ..
11111 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
11112 parts.push((bytes, w.out_features(), *scale));
11113 }
11114 _ => return Ok(None),
11115 }
11116 }
11117 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11119 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
11120 let mcols = if (5..=7).contains(&m) && b567 {
11121 m
11122 } else {
11123 Self::batched_mcols(m)
11124 };
11125 let kname: &'static str = match mcols {
11126 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
11127 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
11128 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
11129 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
11130 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
11131 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
11132 _ => return Ok(None),
11133 };
11134 if std::env::var("MEMRA_DEBUG").is_ok() {
11137 use std::sync::Mutex;
11138 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
11139 let mut seen = SEEN.lock().unwrap();
11140 if !seen.contains(&tag) {
11141 seen.push(tag);
11142 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
11143 }
11144 }
11145 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
11148 let three = parts.len() == 3;
11149 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
11150 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
11151 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
11152 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
11155 let cfg = LaunchConfig {
11156 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
11157 block_dim: (32, ROWS_PER_BLOCK, 1),
11158 shared_mem_bytes: 0,
11159 };
11160 let (inf, mi) = (in_f as i32, m as i32);
11161 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
11162 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
11163 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
11164 let s3 = if three { 1.0f32 } else { parts[3].2 };
11165 let w3 = if three { parts[0].0 } else { parts[3].0 };
11166 let f = self.func(kname);
11167 let __s_b = self.gpu.stream();
11168 let mut b = __s_b.launch_builder(&f);
11169 b.arg(parts[0].0)
11170 .arg(parts[1].0)
11171 .arg(parts[2].0)
11172 .arg(w3)
11173 .arg(aq)
11174 .arg(ad)
11175 .arg(&mut y0)
11176 .arg(&mut y1)
11177 .arg(&mut y2)
11178 .arg(&mut y3)
11179 .arg(&inf)
11180 .arg(&n0)
11181 .arg(&n1)
11182 .arg(&n2)
11183 .arg(&n3)
11184 .arg(&mi)
11185 .arg(&s0)
11186 .arg(&s1)
11187 .arg(&s2)
11188 .arg(&s3);
11189 unsafe {
11190 b.launch(cfg)?;
11191 }
11192 Ok(Some(if three {
11193 vec![y0, y1, y2]
11194 } else {
11195 vec![y0, y1, y2, y3]
11196 }))
11197 }
11198
11199 pub fn matmul_decode_exact_dual(
11215 &self,
11216 w0: &crate::model::GpuTensor,
11217 w1: &crate::model::GpuTensor,
11218 x: &CudaSlice<f32>,
11219 m: usize,
11220 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11221 use crate::model::GpuTensor;
11222 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11223 let on = *ON.get_or_init(|| {
11224 std::env::var("MEMRA_SPEC_DUAL_T")
11225 .map(|v| v != "0")
11226 .unwrap_or(true)
11227 });
11228 if !on
11229 || !(2..=4).contains(&m)
11230 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11231 || !self.uses_q8_1_fast(w0)
11232 || !self.uses_q8_1_fast(w1)
11233 {
11234 return Ok(None);
11235 }
11236 if !self.mmvq_supports(QT_NVFP4) {
11241 return Ok(None);
11242 }
11243 let (in_f, out_f) = (w0.in_features(), w0.out_features());
11244 if w1.in_features() != in_f || w1.out_features() != out_f {
11245 return Ok(None);
11246 }
11247 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
11248 (
11249 GpuTensor::Quant {
11250 bytes: b0,
11251 qtype: q0,
11252 row_bytes: rb0,
11253 scale: s0,
11254 rp: rp0,
11255 rp4: None,
11256 ..
11257 },
11258 GpuTensor::Quant {
11259 bytes: b1,
11260 qtype: q1,
11261 row_bytes: rb1,
11262 scale: s1,
11263 rp: rp1,
11264 rp4: None,
11265 ..
11266 },
11267 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11268 (b0, b1, *rb0, *s0, *s1, *rp0)
11269 }
11270 _ => return Ok(None),
11271 };
11272 if std::env::var("MEMRA_DEBUG").is_ok() {
11275 static ONCE: std::sync::Once = std::sync::Once::new();
11276 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
11277 }
11278 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11279 let (y0, y1) =
11280 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
11281 let mut y0 = y0;
11282 let mut y1 = y1;
11283 if s0 != 1.0 {
11284 self.scale_inplace(&mut y0, s0, m * out_f)?;
11285 }
11286 if s1 != 1.0 {
11287 self.scale_inplace(&mut y1, s1, m * out_f)?;
11288 }
11289 Ok(Some((y0, y1)))
11290 }
11291
11292 #[allow(clippy::too_many_arguments)]
11297 pub fn qmatvec_batched_dual_raw(
11298 &self,
11299 b0: &CudaSlice<u8>,
11300 b1: &CudaSlice<u8>,
11301 aq: &CudaSlice<i8>,
11302 ad: &CudaSlice<f32>,
11303 m: usize,
11304 in_f: usize,
11305 out_f: usize,
11306 row_bytes: usize,
11307 rp: bool,
11308 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11309 const ROWS_PER_BLOCK: u32 = 4;
11310 let mcols = Self::batched_mcols(m);
11311 let tiny_rp1 = rp
11314 && mcols == 4
11315 && out_f <= 128
11316 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
11317 let (name, rows_per_block) = if tiny_rp1 {
11318 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
11319 } else {
11320 match (mcols, rp, m) {
11321 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
11322 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
11323 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
11324 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
11325 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
11326 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
11327 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
11328 _ => {
11329 return Err(
11330 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
11331 );
11332 }
11333 }
11334 };
11335 let f = self.func(name);
11336 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
11337 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
11338 let cfg = LaunchConfig {
11339 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11340 block_dim: (32, ROWS_PER_BLOCK, 1),
11341 shared_mem_bytes: 0,
11342 };
11343 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11344 let __s_b = self.gpu.stream();
11345 let mut b = __s_b.launch_builder(&f);
11346 b.arg(b0)
11347 .arg(b1)
11348 .arg(aq)
11349 .arg(ad)
11350 .arg(&mut y0)
11351 .arg(&mut y1)
11352 .arg(&inf)
11353 .arg(&outf)
11354 .arg(&mi)
11355 .arg(&rb);
11356 unsafe {
11357 b.launch(cfg)?;
11358 }
11359 Ok((y0, y1))
11360 }
11361
11362 pub fn matmul_pre_dual_noscale(
11374 &self,
11375 w0: &crate::model::GpuTensor,
11376 w1: &crate::model::GpuTensor,
11377 aq: &CudaSlice<i8>,
11378 ad: &CudaSlice<f32>,
11379 m: usize,
11380 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11381 {
11382 use crate::model::GpuTensor;
11383 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11384 return Ok(None);
11385 }
11386 if !self.mmvq_supports(QT_NVFP4) {
11396 return Ok(None);
11397 }
11398 let (in_f, out_f) = (w0.in_features(), w0.out_features());
11399 if w1.in_features() != in_f || w1.out_features() != out_f {
11400 return Ok(None);
11401 }
11402 let no_mirror =
11415 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
11416 if self.q8_ffn_fuse2_on()
11417 && no_mirror(w0)
11418 && no_mirror(w1)
11419 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
11420 {
11421 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
11422 return Ok(Some(((y0, 1.0), (y1, 1.0))));
11423 }
11424 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11434 let (y0, y1) =
11435 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
11436 return Ok(Some(((y0, p0.3), (y1, p1.3))));
11437 }
11438 let (b0, q0, rb0, s0, rp0) = match w0 {
11439 GpuTensor::Quant {
11440 bytes,
11441 qtype,
11442 row_bytes,
11443 scale,
11444 rp,
11445 ..
11446 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11447 _ => return Ok(None),
11448 };
11449 let (b1, q1, rb1, s1, rp1) = match w1 {
11450 GpuTensor::Quant {
11451 bytes,
11452 qtype,
11453 row_bytes,
11454 scale,
11455 rp,
11456 ..
11457 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11458 _ => return Ok(None),
11459 };
11460 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
11461 return Ok(None);
11462 }
11463 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11465 let rows_per_block = ROWS_PER_BLOCK * RPW;
11466 let f = self.func(if rp0 {
11467 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
11468 } else {
11469 "qmatvec_nvfp4_mmvq_dual_mr2"
11470 });
11471 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
11472 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
11473 let cfg = LaunchConfig {
11474 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11475 block_dim: (32, ROWS_PER_BLOCK, 1),
11476 shared_mem_bytes: 0,
11477 };
11478 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
11479 let one = 1.0f32;
11482 let __s_b = self.gpu.stream();
11483 let mut b = __s_b.launch_builder(&f);
11484 b.arg(b0)
11485 .arg(b1)
11486 .arg(aq)
11487 .arg(ad)
11488 .arg(&mut y0)
11489 .arg(&mut y1)
11490 .arg(&inf)
11491 .arg(&outf)
11492 .arg(&mi)
11493 .arg(&rb)
11494 .arg(&one)
11495 .arg(&one);
11496 unsafe {
11497 b.launch(cfg)?;
11498 }
11499 Ok(Some(((y0, s0), (y1, s1))))
11500 }
11501
11502 #[allow(clippy::too_many_arguments)]
11510 pub fn matmul_nvfp4_fused3(
11511 &self,
11512 w0: &crate::model::GpuTensor,
11513 w1: &crate::model::GpuTensor,
11514 w2: &crate::model::GpuTensor,
11515 aq: &CudaSlice<i8>,
11516 ad: &CudaSlice<f32>,
11517 m: usize,
11518 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11519 {
11520 use crate::model::GpuTensor;
11521 if !(1..=8).contains(&m)
11528 || !self.mmvq_supports(QT_NVFP4)
11529 || !self.uses_q8_1_fast(w0)
11530 || !self.uses_q8_1_fast(w1)
11531 || !self.uses_q8_1_fast(w2)
11532 {
11533 return Ok(None);
11534 }
11535 if m > 1 {
11536 let in_f = w0.in_features();
11537 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
11538 || !self.batched_supports(QT_NVFP4)
11539 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11540 || (m > 4 && !Self::b8_enabled())
11541 || in_f % 512 != 0
11542 || in_f / 64 > 272
11543 {
11544 return Ok(None);
11545 }
11546 }
11547 let unpack = |w: &crate::model::GpuTensor| match w {
11548 GpuTensor::Quant {
11549 bytes,
11550 qtype,
11551 scale,
11552 rp,
11553 ..
11554 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11555 _ => None,
11556 };
11557 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
11558 return Ok(None);
11559 };
11560 let in_f = w0.in_features();
11561 if w1.in_features() != in_f || w2.in_features() != in_f {
11562 return Ok(None);
11563 }
11564 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
11565 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11567 let rows_pb = ROWS_PER_BLOCK * RPW;
11568 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11569 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11570 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11571 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11572 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
11573 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
11576 if m > 1 {
11577 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
11579 return Ok(None);
11580 }
11581 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
11582 let cfg = LaunchConfig {
11583 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
11584 block_dim: (32, ROWS_PER_BLOCK, 1),
11585 shared_mem_bytes: 0,
11586 };
11587 let __s_b = self.gpu.stream();
11588 let mut b = __s_b.launch_builder(&f);
11589 b.arg(b0)
11590 .arg(b1)
11591 .arg(b2)
11592 .arg(aq)
11593 .arg(ad)
11594 .arg(&mut y0)
11595 .arg(&mut y1)
11596 .arg(&mut y2)
11597 .arg(&inf)
11598 .arg(&oi0)
11599 .arg(&oi1)
11600 .arg(&oi2)
11601 .arg(&mi);
11602 unsafe {
11603 b.launch(cfg)?;
11604 }
11605 return Ok(Some((y0, y1, y2)));
11606 }
11607 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
11608 let cfg = LaunchConfig {
11609 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
11610 block_dim: (32, ROWS_PER_BLOCK, 1),
11611 shared_mem_bytes: 0,
11612 };
11613 let __s_b = self.gpu.stream();
11614 let mut b = __s_b.launch_builder(&f);
11615 b.arg(b0)
11616 .arg(b1)
11617 .arg(b2)
11618 .arg(aq)
11619 .arg(ad)
11620 .arg(&mut y0)
11621 .arg(&mut y1)
11622 .arg(&mut y2)
11623 .arg(&inf)
11624 .arg(&oi0)
11625 .arg(&oi1)
11626 .arg(&oi2)
11627 .arg(&mi)
11628 .arg(&p0.1)
11629 .arg(&p1.1)
11630 .arg(&p2.1);
11631 unsafe {
11632 b.launch(cfg)?;
11633 }
11634 Ok(Some((y0, y1, y2)))
11635 }
11636
11637 pub fn matmul_nvfp4_fused2(
11646 &self,
11647 w0: &crate::model::GpuTensor,
11648 w1: &crate::model::GpuTensor,
11649 aq: &CudaSlice<i8>,
11650 ad: &CudaSlice<f32>,
11651 m: usize,
11652 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11653 use crate::model::GpuTensor;
11654 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11655 let off =
11656 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11657 if off
11660 || m != 1
11661 || !self.mmvq_supports(QT_NVFP4)
11662 || !self.uses_q8_1_fast(w0)
11663 || !self.uses_q8_1_fast(w1)
11664 {
11665 return Ok(None);
11666 }
11667 let unpack = |w: &crate::model::GpuTensor| match w {
11668 GpuTensor::Quant {
11669 bytes,
11670 qtype,
11671 scale,
11672 rp,
11673 ..
11674 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11675 _ => None,
11676 };
11677 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11678 return Ok(None);
11679 };
11680 let in_f = w0.in_features();
11681 if w1.in_features() != in_f {
11682 return Ok(None);
11683 }
11684 let (o0, o1) = (w0.out_features(), w1.out_features());
11685 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11687 let rows_pb = ROWS_PER_BLOCK * RPW;
11688 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11689 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11690 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11691 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11692 let cfg = LaunchConfig {
11693 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11694 block_dim: (32, ROWS_PER_BLOCK, 1),
11695 shared_mem_bytes: 0,
11696 };
11697 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11698 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11701 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11704 {
11705 use cudarc::driver::{DevicePtr, DevicePtrMut};
11706 let s = &self.gpu.stream();
11707 let (pw0, _g0) = b0.device_ptr(s);
11708 let (pw1, _g1) = b1.device_ptr(s);
11709 let (paq, _g2) = aq.device_ptr(s);
11710 let (pad, _g3) = ad.device_ptr(s);
11711 let (py0, _g4) = y0.device_ptr_mut(s);
11712 let (py1, _g5) = y1.device_ptr_mut(s);
11713 let (s0, s1) = (p0.1, p1.1);
11714 let mut ps = [
11715 &pw0 as *const _ as *mut std::ffi::c_void,
11716 &pw1 as *const _ as *mut _,
11717 &paq as *const _ as *mut _,
11718 &pad as *const _ as *mut _,
11719 &py0 as *const _ as *mut _,
11720 &py1 as *const _ as *mut _,
11721 &inf as *const _ as *mut _,
11722 &oi0 as *const _ as *mut _,
11723 &oi1 as *const _ as *mut _,
11724 &mi as *const _ as *mut _,
11725 &s0 as *const _ as *mut _,
11726 &s1 as *const _ as *mut _,
11727 ];
11728 unsafe {
11729 self.launch_pdl(
11730 "qmatvec_nvfp4_mmvq_fused2_rp",
11731 cfg.grid_dim,
11732 cfg.block_dim,
11733 &mut ps,
11734 )?;
11735 }
11736 }
11737 return Ok(Some((y0, y1)));
11738 }
11739 let __s_b = self.gpu.stream();
11740 let mut b = __s_b.launch_builder(&f);
11741 b.arg(b0)
11742 .arg(b1)
11743 .arg(aq)
11744 .arg(ad)
11745 .arg(&mut y0)
11746 .arg(&mut y1)
11747 .arg(&inf)
11748 .arg(&oi0)
11749 .arg(&oi1)
11750 .arg(&mi)
11751 .arg(&p0.1)
11752 .arg(&p1.1);
11753 unsafe {
11754 b.launch(cfg)?;
11755 }
11756 Ok(Some((y0, y1)))
11757 }
11758
11759 pub fn matmul_nvfp4_fused2_into(
11764 &self,
11765 w0: &crate::model::GpuTensor,
11766 w1: &crate::model::GpuTensor,
11767 aq: &CudaSlice<i8>,
11768 ad: &CudaSlice<f32>,
11769 y0: &mut CudaSlice<f32>,
11770 y1: &mut CudaSlice<f32>,
11771 ) -> Result<bool, Box<dyn std::error::Error>> {
11772 use crate::model::GpuTensor;
11773 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11774 let off =
11775 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11776 if off
11777 || !self.mmvq_supports(QT_NVFP4)
11778 || !self.uses_q8_1_fast(w0)
11779 || !self.uses_q8_1_fast(w1)
11780 {
11781 return Ok(false);
11782 }
11783 let unpack = |w: &crate::model::GpuTensor| match w {
11784 GpuTensor::Quant {
11785 bytes,
11786 qtype,
11787 scale,
11788 rp,
11789 ..
11790 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11791 _ => None,
11792 };
11793 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11794 return Ok(false);
11795 };
11796 let in_f = w0.in_features();
11797 if w1.in_features() != in_f {
11798 return Ok(false);
11799 }
11800 let (o0, o1) = (w0.out_features(), w1.out_features());
11801 if y0.len() < o0 || y1.len() < o1 {
11802 return Ok(false);
11803 }
11804 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11806 let rows_pb = ROWS_PER_BLOCK * RPW;
11807 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11808 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11809 let cfg = LaunchConfig {
11810 grid_dim: (nb(o0) + nb(o1), 1, 1),
11811 block_dim: (32, ROWS_PER_BLOCK, 1),
11812 shared_mem_bytes: 0,
11813 };
11814 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11815 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11818 let __s_b = self.gpu.stream();
11819 let mut b = __s_b.launch_builder(&f);
11820 b.arg(b0)
11821 .arg(b1)
11822 .arg(aq)
11823 .arg(ad)
11824 .arg(&mut *y0)
11825 .arg(&mut *y1)
11826 .arg(&inf)
11827 .arg(&oi0)
11828 .arg(&oi1)
11829 .arg(&mi)
11830 .arg(&p0.1)
11831 .arg(&p1.1);
11832 unsafe {
11833 b.launch(cfg)?;
11834 }
11835 Ok(true)
11836 }
11837
11838 #[allow(clippy::type_complexity)]
11843 pub fn matmul_nvfp4_fused4(
11844 &self,
11845 w0: &crate::model::GpuTensor,
11846 w1: &crate::model::GpuTensor,
11847 w2: &crate::model::GpuTensor,
11848 w3: &crate::model::GpuTensor,
11849 aq: &CudaSlice<i8>,
11850 ad: &CudaSlice<f32>,
11851 m: usize,
11852 ) -> Result<
11853 Option<(
11854 CudaSlice<f32>,
11855 CudaSlice<f32>,
11856 CudaSlice<f32>,
11857 CudaSlice<f32>,
11858 )>,
11859 Box<dyn std::error::Error>,
11860 > {
11861 use crate::model::GpuTensor;
11862 if !(1..=8).contains(&m)
11869 || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11870 || !self.mmvq_supports(QT_NVFP4)
11871 || !self.uses_q8_1_fast(w0)
11872 || !self.uses_q8_1_fast(w1)
11873 || !self.uses_q8_1_fast(w2)
11874 || !self.uses_q8_1_fast(w3)
11875 {
11876 return Ok(None);
11877 }
11878 if m > 1 {
11879 let in_f = w0.in_features();
11882 if !self.batched_supports(QT_NVFP4)
11883 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11884 || (m > 4 && !Self::b8_enabled())
11885 || in_f % 512 != 0
11886 || in_f / 64 > 272
11887 {
11888 return Ok(None);
11889 }
11890 }
11891 let unpack = |w: &crate::model::GpuTensor| match w {
11892 GpuTensor::Quant {
11893 bytes,
11894 qtype,
11895 scale,
11896 rp,
11897 ..
11898 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11899 _ => None,
11900 };
11901 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
11902 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11903 else {
11904 return Ok(None);
11905 };
11906 let in_f = w0.in_features();
11907 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11908 return Ok(None);
11909 }
11910 let (o0, o1, o2, o3) = (
11911 w0.out_features(),
11912 w1.out_features(),
11913 w2.out_features(),
11914 w3.out_features(),
11915 );
11916 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11918 let rows_pb = ROWS_PER_BLOCK * RPW;
11919 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11920 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11921 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11922 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11923 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11924 let (inf, oi0, oi1, oi2, oi3, mi) = (
11925 in_f as i32,
11926 o0 as i32,
11927 o1 as i32,
11928 o2 as i32,
11929 o3 as i32,
11930 m as i32,
11931 );
11932 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11935 if m > 1 {
11936 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
11939 return Ok(None);
11940 }
11941 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
11942 let cfg = LaunchConfig {
11943 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
11944 block_dim: (32, ROWS_PER_BLOCK, 1),
11945 shared_mem_bytes: 0,
11946 };
11947 let __s_b = self.gpu.stream();
11948 let mut b = __s_b.launch_builder(&f);
11949 b.arg(b0)
11950 .arg(b1)
11951 .arg(b2)
11952 .arg(b3)
11953 .arg(aq)
11954 .arg(ad)
11955 .arg(&mut y0)
11956 .arg(&mut y1)
11957 .arg(&mut y2)
11958 .arg(&mut y3)
11959 .arg(&inf)
11960 .arg(&oi0)
11961 .arg(&oi1)
11962 .arg(&oi2)
11963 .arg(&oi3)
11964 .arg(&mi);
11965 unsafe {
11966 b.launch(cfg)?;
11967 }
11968 return Ok(Some((y0, y1, y2, y3)));
11969 }
11970 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
11971 let cfg = LaunchConfig {
11972 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
11973 block_dim: (32, ROWS_PER_BLOCK, 1),
11974 shared_mem_bytes: 0,
11975 };
11976 let __s_b = self.gpu.stream();
11977 let mut b = __s_b.launch_builder(&f);
11978 b.arg(b0)
11979 .arg(b1)
11980 .arg(b2)
11981 .arg(b3)
11982 .arg(aq)
11983 .arg(ad)
11984 .arg(&mut y0)
11985 .arg(&mut y1)
11986 .arg(&mut y2)
11987 .arg(&mut y3)
11988 .arg(&inf)
11989 .arg(&oi0)
11990 .arg(&oi1)
11991 .arg(&oi2)
11992 .arg(&oi3)
11993 .arg(&mi)
11994 .arg(&p0.1)
11995 .arg(&p1.1)
11996 .arg(&p2.1)
11997 .arg(&p3.1);
11998 unsafe {
11999 b.launch(cfg)?;
12000 }
12001 Ok(Some((y0, y1, y2, y3)))
12002 }
12003
12004 pub fn matmul_q8_fused2(
12012 &self,
12013 w0: &crate::model::GpuTensor,
12014 w1: &crate::model::GpuTensor,
12015 aq: &CudaSlice<i8>,
12016 ad: &CudaSlice<f32>,
12017 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12018 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12024 return Ok(Some(self.e4m3_fused2_core(
12025 p0.0,
12026 p1.0,
12027 aq,
12028 ad,
12029 w0.in_features(),
12030 p0.1,
12031 p1.1,
12032 p0.2,
12033 p0.3,
12034 p1.3,
12035 )?));
12036 }
12037 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12038 return Ok(None);
12039 };
12040 Ok(Some(self.q8_fused2_core(
12041 p0.0,
12042 p1.0,
12043 aq,
12044 ad,
12045 w0.in_features(),
12046 p0.1,
12047 p1.1,
12048 p0.2,
12049 )?))
12050 }
12051
12052 #[allow(clippy::too_many_arguments)]
12053 fn q8_fused2_core(
12054 &self,
12055 b0: &CudaSlice<u8>,
12056 b1: &CudaSlice<u8>,
12057 aq: &CudaSlice<i8>,
12058 ad: &CudaSlice<f32>,
12059 in_f: usize,
12060 out0: usize,
12061 out1: usize,
12062 row_bytes: usize,
12063 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12064 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12066 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12067 let f = self.func("qmatvec_q8_0_mmvq_fused2");
12068 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12069 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12070 let cfg = LaunchConfig {
12071 grid_dim: (nb0 + nb1, 1, 1),
12072 block_dim: (32, ROWS_PER_BLOCK, 1),
12073 shared_mem_bytes: 0,
12074 };
12075 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12076 let __s_b = self.gpu.stream();
12077 let mut b = __s_b.launch_builder(&f);
12078 b.arg(b0)
12079 .arg(b1)
12080 .arg(aq)
12081 .arg(ad)
12082 .arg(&mut y0)
12083 .arg(&mut y1)
12084 .arg(&inf)
12085 .arg(&o0)
12086 .arg(&o1)
12087 .arg(&rbl);
12088 unsafe {
12089 b.launch(cfg)?;
12090 }
12091 Ok((y0, y1))
12092 }
12093
12094 pub fn matmul_q8_fused2_x(
12100 &self,
12101 w0: &crate::model::GpuTensor,
12102 w1: &crate::model::GpuTensor,
12103 x: &CudaSlice<f32>,
12104 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12105 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12106 return Ok(None);
12107 }
12108 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12109 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
12110 return Ok(Some(self.e4m3_fused2_core(
12111 p0.0,
12112 p1.0,
12113 &aq,
12114 &ad,
12115 w0.in_features(),
12116 p0.1,
12117 p1.1,
12118 p0.2,
12119 p0.3,
12120 p1.3,
12121 )?));
12122 }
12123 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12124 return Ok(None);
12125 };
12126 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
12127 Ok(Some(self.q8_fused2_core(
12128 p0.0,
12129 p1.0,
12130 &aq,
12131 &ad,
12132 w0.in_features(),
12133 p0.1,
12134 p1.1,
12135 p0.2,
12136 )?))
12137 }
12138
12139 #[allow(clippy::too_many_arguments)]
12142 pub fn qmatvec_q8_fused2_raw(
12143 &self,
12144 b0: &CudaSlice<u8>,
12145 b1: &CudaSlice<u8>,
12146 x: &CudaSlice<f32>,
12147 in_f: usize,
12148 out0: usize,
12149 out1: usize,
12150 row_bytes: usize,
12151 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12152 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12153 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
12154 }
12155
12156 pub fn matmul_q4_fused3(
12162 &self,
12163 w0: &crate::model::GpuTensor,
12164 w1: &crate::model::GpuTensor,
12165 w2: &crate::model::GpuTensor,
12166 aq: &CudaSlice<i8>,
12167 ad: &CudaSlice<f32>,
12168 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12169 {
12170 use crate::model::GpuTensor;
12171 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12172 match w {
12173 GpuTensor::Quant {
12174 qtype, row_bytes, ..
12175 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12176 _ => None,
12177 }
12178 };
12179 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
12180 return Ok(None);
12181 };
12182 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12183 return Ok(None);
12184 }
12185 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12189 match w {
12190 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12191 Some(m) => (m, true),
12192 None => (bytes, *rp),
12193 },
12194 _ => unreachable!(),
12195 }
12196 }
12197 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12198 if rp0 != rp1 || rp1 != rp2 {
12199 return Ok(None);
12200 }
12201 let rp = rp0;
12202 let rpb: u32 = 4;
12203 let mr1 = rp && Self::q40_mr1_on();
12207 let nb = |o: usize| {
12208 if mr1 {
12209 (o as u32).div_ceil(rpb)
12210 } else {
12211 (o as u32).div_ceil(2).div_ceil(rpb)
12212 }
12213 };
12214 let grid = nb(o0) + nb(o1) + nb(o2);
12215 let mut y0 = self.alloc_uninit::<f32>(o0)?;
12216 let mut y1 = self.alloc_uninit::<f32>(o1)?;
12217 let mut y2 = self.alloc_uninit::<f32>(o2)?;
12218 let f = self.func(if mr1 {
12219 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12220 } else if rp {
12221 "qmatvec_q4_0_mmvq_fused3_rp"
12222 } else {
12223 "qmatvec_q4_0_mmvq_fused3"
12224 });
12225 let cfg = LaunchConfig {
12226 grid_dim: (grid, 1, 1),
12227 block_dim: (32, rpb, 1),
12228 shared_mem_bytes: 0,
12229 };
12230 let inf = w0.in_features() as i32;
12231 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12232 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12233 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12236 {
12237 use cudarc::driver::{DevicePtr, DevicePtrMut};
12238 let s = &self.gpu.stream();
12239 let (p0, _g0) = b0.device_ptr(s);
12240 let (p1, _g1) = b1.device_ptr(s);
12241 let (p2, _g2) = b2.device_ptr(s);
12242 let (paq, _g3) = aq.device_ptr(s);
12243 let (pad, _g4) = ad.device_ptr(s);
12244 let (py0, _g5) = y0.device_ptr_mut(s);
12245 let (py1, _g6) = y1.device_ptr_mut(s);
12246 let (py2, _g7) = y2.device_ptr_mut(s);
12247 let mut ps = [
12248 &p0 as *const _ as *mut std::ffi::c_void,
12249 &p1 as *const _ as *mut _,
12250 &p2 as *const _ as *mut _,
12251 &paq as *const _ as *mut _,
12252 &pad as *const _ as *mut _,
12253 &py0 as *const _ as *mut _,
12254 &py1 as *const _ as *mut _,
12255 &py2 as *const _ as *mut _,
12256 &inf as *const _ as *mut _,
12257 &oo0 as *const _ as *mut _,
12258 &oo1 as *const _ as *mut _,
12259 &oo2 as *const _ as *mut _,
12260 &r0 as *const _ as *mut _,
12261 &r1 as *const _ as *mut _,
12262 &r2 as *const _ as *mut _,
12263 ];
12264 unsafe {
12265 self.launch_pdl(
12266 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12267 (grid, 1, 1),
12268 (32, rpb, 1),
12269 &mut ps,
12270 )?;
12271 }
12272 }
12273 return Ok(Some((y0, y1, y2)));
12274 }
12275 let __s_b = self.gpu.stream();
12276 let mut b = __s_b.launch_builder(&f);
12277 b.arg(b0)
12278 .arg(b1)
12279 .arg(b2)
12280 .arg(aq)
12281 .arg(ad)
12282 .arg(&mut y0)
12283 .arg(&mut y1)
12284 .arg(&mut y2)
12285 .arg(&inf)
12286 .arg(&oo0)
12287 .arg(&oo1)
12288 .arg(&oo2)
12289 .arg(&r0)
12290 .arg(&r1)
12291 .arg(&r2);
12292 unsafe {
12293 b.launch(cfg)?;
12294 }
12295 Ok(Some((y0, y1, y2)))
12296 }
12297
12298 #[allow(clippy::too_many_arguments)]
12301 pub fn matmul_q4_fused3_into(
12302 &self,
12303 w0: &crate::model::GpuTensor,
12304 w1: &crate::model::GpuTensor,
12305 w2: &crate::model::GpuTensor,
12306 aq: &CudaSlice<i8>,
12307 ad: &CudaSlice<f32>,
12308 y0: &mut CudaSlice<f32>,
12309 y1: &mut CudaSlice<f32>,
12310 y2: &mut CudaSlice<f32>,
12311 ) -> Result<bool, Box<dyn std::error::Error>> {
12312 use crate::model::GpuTensor;
12313 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12314 match w {
12315 GpuTensor::Quant {
12316 qtype, row_bytes, ..
12317 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12318 _ => None,
12319 }
12320 };
12321 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
12322 return Ok(false);
12323 };
12324 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12325 return Ok(false);
12326 }
12327 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12328 match w {
12329 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12330 Some(m) => (m, true),
12331 None => (bytes, *rp),
12332 },
12333 _ => unreachable!(),
12334 }
12335 }
12336 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12337 if rp0 != rp1 || rp1 != rp2 {
12338 return Ok(false);
12339 }
12340 let rp = rp0;
12341 let rpb: u32 = 4;
12342 let mr1 = rp && Self::q40_mr1_on();
12343 let nb = |o: usize| {
12344 if mr1 {
12345 (o as u32).div_ceil(rpb)
12346 } else {
12347 (o as u32).div_ceil(2).div_ceil(rpb)
12348 }
12349 };
12350 let grid = nb(o0) + nb(o1) + nb(o2);
12351 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
12352 let f = self.func(if mr1 {
12353 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12354 } else if rp {
12355 "qmatvec_q4_0_mmvq_fused3_rp"
12356 } else {
12357 "qmatvec_q4_0_mmvq_fused3"
12358 });
12359 let cfg = LaunchConfig {
12360 grid_dim: (grid, 1, 1),
12361 block_dim: (32, rpb, 1),
12362 shared_mem_bytes: 0,
12363 };
12364 let inf = w0.in_features() as i32;
12365 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12366 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12367 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12369 use cudarc::driver::{DevicePtr, DevicePtrMut};
12370 let s = &self.gpu.stream();
12371 let (p0, _g0) = b0.device_ptr(s);
12372 let (p1, _g1) = b1.device_ptr(s);
12373 let (p2, _g2) = b2.device_ptr(s);
12374 let (paq, _g3) = aq.device_ptr(s);
12375 let (pad, _g4) = ad.device_ptr(s);
12376 let (py0, _g5) = y0.device_ptr_mut(s);
12377 let (py1, _g6) = y1.device_ptr_mut(s);
12378 let (py2, _g7) = y2.device_ptr_mut(s);
12379 let mut ps = [
12380 &p0 as *const _ as *mut std::ffi::c_void,
12381 &p1 as *const _ as *mut _,
12382 &p2 as *const _ as *mut _,
12383 &paq as *const _ as *mut _,
12384 &pad as *const _ as *mut _,
12385 &py0 as *const _ as *mut _,
12386 &py1 as *const _ as *mut _,
12387 &py2 as *const _ as *mut _,
12388 &inf as *const _ as *mut _,
12389 &oo0 as *const _ as *mut _,
12390 &oo1 as *const _ as *mut _,
12391 &oo2 as *const _ as *mut _,
12392 &r0 as *const _ as *mut _,
12393 &r1 as *const _ as *mut _,
12394 &r2 as *const _ as *mut _,
12395 ];
12396 unsafe {
12397 self.launch_pdl(
12398 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12399 (grid, 1, 1),
12400 (32, rpb, 1),
12401 &mut ps,
12402 )?;
12403 }
12404 return Ok(true);
12405 }
12406 let __s_b = self.gpu.stream();
12407 let mut b = __s_b.launch_builder(&f);
12408 b.arg(b0)
12409 .arg(b1)
12410 .arg(b2)
12411 .arg(aq)
12412 .arg(ad)
12413 .arg(&mut *y0)
12414 .arg(&mut *y1)
12415 .arg(&mut *y2)
12416 .arg(&inf)
12417 .arg(&oo0)
12418 .arg(&oo1)
12419 .arg(&oo2)
12420 .arg(&r0)
12421 .arg(&r1)
12422 .arg(&r2);
12423 unsafe {
12424 b.launch(cfg)?;
12425 }
12426 Ok(true)
12427 }
12428
12429 pub fn matmul_q4_fused2(
12431 &self,
12432 w0: &crate::model::GpuTensor,
12433 w1: &crate::model::GpuTensor,
12434 aq: &CudaSlice<i8>,
12435 ad: &CudaSlice<f32>,
12436 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12437 use crate::model::GpuTensor;
12438 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12439 match w {
12440 GpuTensor::Quant {
12441 qtype, row_bytes, ..
12442 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12443 _ => None,
12444 }
12445 };
12446 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12447 return Ok(None);
12448 };
12449 if w0.in_features() != w1.in_features() {
12450 return Ok(None);
12451 }
12452 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12454 match w {
12455 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12456 Some(m) => (m, true),
12457 None => (bytes, *rp),
12458 },
12459 _ => unreachable!(),
12460 }
12461 }
12462 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12463 if rp0 != rp1 {
12464 return Ok(None);
12465 }
12466 let rp = rp0;
12467 let rpb: u32 = 4;
12468 let mr1 = rp && Self::q40_mr1_on();
12470 let nb = |o: usize| {
12471 if mr1 {
12472 (o as u32).div_ceil(rpb)
12473 } else {
12474 (o as u32).div_ceil(2).div_ceil(rpb)
12475 }
12476 };
12477 let grid = nb(o0) + nb(o1);
12478 let mut y0 = self.alloc_uninit::<f32>(o0)?;
12479 let mut y1 = self.alloc_uninit::<f32>(o1)?;
12480 let f = self.func(if mr1 {
12481 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12482 } else if rp {
12483 "qmatvec_q4_0_mmvq_fused2_rp"
12484 } else {
12485 "qmatvec_q4_0_mmvq_fused2"
12486 });
12487 let cfg = LaunchConfig {
12488 grid_dim: (grid, 1, 1),
12489 block_dim: (32, rpb, 1),
12490 shared_mem_bytes: 0,
12491 };
12492 let inf = w0.in_features() as i32;
12493 let (oo0, oo1) = (o0 as i32, o1 as i32);
12494 let (r0, r1) = (rb0 as i64, rb1 as i64);
12495 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12497 {
12498 use cudarc::driver::{DevicePtr, DevicePtrMut};
12499 let s = &self.gpu.stream();
12500 let (p0, _g0) = b0.device_ptr(s);
12501 let (p1, _g1) = b1.device_ptr(s);
12502 let (paq, _g2) = aq.device_ptr(s);
12503 let (pad, _g3) = ad.device_ptr(s);
12504 let (py0, _g4) = y0.device_ptr_mut(s);
12505 let (py1, _g5) = y1.device_ptr_mut(s);
12506 let mut ps = [
12507 &p0 as *const _ as *mut std::ffi::c_void,
12508 &p1 as *const _ as *mut _,
12509 &paq as *const _ as *mut _,
12510 &pad as *const _ as *mut _,
12511 &py0 as *const _ as *mut _,
12512 &py1 as *const _ as *mut _,
12513 &inf as *const _ as *mut _,
12514 &oo0 as *const _ as *mut _,
12515 &oo1 as *const _ as *mut _,
12516 &r0 as *const _ as *mut _,
12517 &r1 as *const _ as *mut _,
12518 ];
12519 unsafe {
12520 self.launch_pdl(
12521 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12522 (grid, 1, 1),
12523 (32, rpb, 1),
12524 &mut ps,
12525 )?;
12526 }
12527 }
12528 return Ok(Some((y0, y1)));
12529 }
12530 let __s_b = self.gpu.stream();
12531 let mut b = __s_b.launch_builder(&f);
12532 b.arg(b0)
12533 .arg(b1)
12534 .arg(aq)
12535 .arg(ad)
12536 .arg(&mut y0)
12537 .arg(&mut y1)
12538 .arg(&inf)
12539 .arg(&oo0)
12540 .arg(&oo1)
12541 .arg(&r0)
12542 .arg(&r1);
12543 unsafe {
12544 b.launch(cfg)?;
12545 }
12546 Ok(Some((y0, y1)))
12547 }
12548
12549 pub fn matmul_q4_fused2_into(
12551 &self,
12552 w0: &crate::model::GpuTensor,
12553 w1: &crate::model::GpuTensor,
12554 aq: &CudaSlice<i8>,
12555 ad: &CudaSlice<f32>,
12556 y0: &mut CudaSlice<f32>,
12557 y1: &mut CudaSlice<f32>,
12558 ) -> Result<bool, Box<dyn std::error::Error>> {
12559 use crate::model::GpuTensor;
12560 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12561 match w {
12562 GpuTensor::Quant {
12563 qtype, row_bytes, ..
12564 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12565 _ => None,
12566 }
12567 };
12568 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12569 return Ok(false);
12570 };
12571 if w0.in_features() != w1.in_features() {
12572 return Ok(false);
12573 }
12574 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12575 match w {
12576 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12577 Some(m) => (m, true),
12578 None => (bytes, *rp),
12579 },
12580 _ => unreachable!(),
12581 }
12582 }
12583 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12584 if rp0 != rp1 {
12585 return Ok(false);
12586 }
12587 let rp = rp0;
12588 let rpb: u32 = 4;
12589 let mr1 = rp && Self::q40_mr1_on();
12590 let nb = |o: usize| {
12591 if mr1 {
12592 (o as u32).div_ceil(rpb)
12593 } else {
12594 (o as u32).div_ceil(2).div_ceil(rpb)
12595 }
12596 };
12597 let grid = nb(o0) + nb(o1);
12598 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
12599 let f = self.func(if mr1 {
12600 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12601 } else if rp {
12602 "qmatvec_q4_0_mmvq_fused2_rp"
12603 } else {
12604 "qmatvec_q4_0_mmvq_fused2"
12605 });
12606 let cfg = LaunchConfig {
12607 grid_dim: (grid, 1, 1),
12608 block_dim: (32, rpb, 1),
12609 shared_mem_bytes: 0,
12610 };
12611 let inf = w0.in_features() as i32;
12612 let (oo0, oo1) = (o0 as i32, o1 as i32);
12613 let (r0, r1) = (rb0 as i64, rb1 as i64);
12614 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12616 use cudarc::driver::{DevicePtr, DevicePtrMut};
12617 let s = &self.gpu.stream();
12618 let (p0, _g0) = b0.device_ptr(s);
12619 let (p1, _g1) = b1.device_ptr(s);
12620 let (paq, _g2) = aq.device_ptr(s);
12621 let (pad, _g3) = ad.device_ptr(s);
12622 let (py0, _g4) = y0.device_ptr_mut(s);
12623 let (py1, _g5) = y1.device_ptr_mut(s);
12624 let mut ps = [
12625 &p0 as *const _ as *mut std::ffi::c_void,
12626 &p1 as *const _ as *mut _,
12627 &paq as *const _ as *mut _,
12628 &pad as *const _ as *mut _,
12629 &py0 as *const _ as *mut _,
12630 &py1 as *const _ as *mut _,
12631 &inf as *const _ as *mut _,
12632 &oo0 as *const _ as *mut _,
12633 &oo1 as *const _ as *mut _,
12634 &r0 as *const _ as *mut _,
12635 &r1 as *const _ as *mut _,
12636 ];
12637 unsafe {
12638 self.launch_pdl(
12639 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12640 (grid, 1, 1),
12641 (32, rpb, 1),
12642 &mut ps,
12643 )?;
12644 }
12645 return Ok(true);
12646 }
12647 let __s_b = self.gpu.stream();
12648 let mut b = __s_b.launch_builder(&f);
12649 b.arg(b0)
12650 .arg(b1)
12651 .arg(aq)
12652 .arg(ad)
12653 .arg(&mut *y0)
12654 .arg(&mut *y1)
12655 .arg(&inf)
12656 .arg(&oo0)
12657 .arg(&oo1)
12658 .arg(&r0)
12659 .arg(&r1);
12660 unsafe {
12661 b.launch(cfg)?;
12662 }
12663 Ok(true)
12664 }
12665
12666 pub fn matmul_q4_fused2_batched(
12671 &self,
12672 w0: &crate::model::GpuTensor,
12673 w1: &crate::model::GpuTensor,
12674 aq: &CudaSlice<i8>,
12675 ad: &CudaSlice<f32>,
12676 m: usize,
12677 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12678 use crate::model::GpuTensor;
12679 if m < 2 || m > 8 {
12680 return Ok(None);
12681 }
12682 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12683 match w {
12684 GpuTensor::Quant {
12685 qtype, row_bytes, ..
12686 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12687 _ => None,
12688 }
12689 };
12690 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
12691 return Ok(None);
12692 };
12693 if w0.in_features() != w1.in_features() {
12694 return Ok(None);
12695 }
12696 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12697 match w {
12698 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12699 Some(mr) => (mr, true),
12700 None => (bytes, *rp),
12701 },
12702 _ => unreachable!(),
12703 }
12704 }
12705 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12706 if !rp0 || !rp1 {
12707 return Ok(None);
12708 }
12709 let mcols = Self::batched_mcols(m);
12710 let rpb: u32 = 4;
12711 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12712 let grid = nb(o0) + nb(o1);
12713 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12714 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12715 let f = self.func(match mcols {
12716 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12717 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12718 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12719 });
12720 let cfg = LaunchConfig {
12721 grid_dim: (grid, 1, 1),
12722 block_dim: (32, rpb, 1),
12723 shared_mem_bytes: 0,
12724 };
12725 let inf = w0.in_features() as i32;
12726 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12727 let rb = rb0 as i64;
12728 let __s_b = self.gpu.stream();
12729 let mut b = __s_b.launch_builder(&f);
12730 b.arg(b0)
12731 .arg(b1)
12732 .arg(aq)
12733 .arg(ad)
12734 .arg(&mut y0)
12735 .arg(&mut y1)
12736 .arg(&inf)
12737 .arg(&oo0)
12738 .arg(&oo1)
12739 .arg(&mi)
12740 .arg(&rb);
12741 unsafe {
12742 b.launch(cfg)?;
12743 }
12744 Ok(Some((y0, y1)))
12745 }
12746
12747 #[allow(clippy::too_many_arguments)]
12750 pub fn matmul_q4_fused3_batched(
12751 &self,
12752 w0: &crate::model::GpuTensor,
12753 w1: &crate::model::GpuTensor,
12754 w2: &crate::model::GpuTensor,
12755 aq: &CudaSlice<i8>,
12756 ad: &CudaSlice<f32>,
12757 m: usize,
12758 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12759 {
12760 use crate::model::GpuTensor;
12761 if m < 2 || m > 8 {
12762 return Ok(None);
12763 }
12764 let q4 = |w: &GpuTensor| -> Option<usize> {
12765 match w {
12766 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12767 _ => None,
12768 }
12769 };
12770 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12771 return Ok(None);
12772 };
12773 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12774 return Ok(None);
12775 }
12776 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12777 match w {
12778 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12779 Some(mr) => (mr, true),
12780 None => (bytes, *rp),
12781 },
12782 _ => unreachable!(),
12783 }
12784 }
12785 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12786 if !rp0 || !rp1 || !rp2 {
12787 return Ok(None);
12788 }
12789 let mcols = Self::batched_mcols(m);
12790 let rpb: u32 = 4;
12791 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12792 let grid = nb(o0) + nb(o1) + nb(o2);
12793 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12794 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12795 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12796 let f = self.func(match mcols {
12797 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12798 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12799 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12800 });
12801 let cfg = LaunchConfig {
12802 grid_dim: (grid, 1, 1),
12803 block_dim: (32, rpb, 1),
12804 shared_mem_bytes: 0,
12805 };
12806 let inf = w0.in_features() as i32;
12807 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12808 let rb = 0i64;
12809 let __s_b = self.gpu.stream();
12810 let mut b = __s_b.launch_builder(&f);
12811 b.arg(b0)
12812 .arg(b1)
12813 .arg(b2)
12814 .arg(aq)
12815 .arg(ad)
12816 .arg(&mut y0)
12817 .arg(&mut y1)
12818 .arg(&mut y2)
12819 .arg(&inf)
12820 .arg(&oo0)
12821 .arg(&oo1)
12822 .arg(&oo2)
12823 .arg(&mi)
12824 .arg(&rb);
12825 unsafe {
12826 b.launch(cfg)?;
12827 }
12828 Ok(Some((y0, y1, y2)))
12829 }
12830
12831 pub fn matmul_q8_fused3(
12832 &self,
12833 w0: &crate::model::GpuTensor,
12834 w1: &crate::model::GpuTensor,
12835 w2: &crate::model::GpuTensor,
12836 aq: &CudaSlice<i8>,
12837 ad: &CudaSlice<f32>,
12838 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12839 {
12840 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12843 return Ok(Some(self.e4m3_fused3_core(
12844 p0.0,
12845 p1.0,
12846 p2.0,
12847 aq,
12848 ad,
12849 w0.in_features(),
12850 p0.1,
12851 p1.1,
12852 p2.1,
12853 p0.2,
12854 p0.3,
12855 p1.3,
12856 p2.3,
12857 )?));
12858 }
12859 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12860 return Ok(None);
12861 };
12862 Ok(Some(self.q8_fused3_core(
12863 p0.0,
12864 p1.0,
12865 p2.0,
12866 aq,
12867 ad,
12868 w0.in_features(),
12869 p0.1,
12870 p1.1,
12871 p2.1,
12872 p0.2,
12873 )?))
12874 }
12875
12876 #[allow(clippy::too_many_arguments)]
12877 fn q8_fused3_core(
12878 &self,
12879 b0: &CudaSlice<u8>,
12880 b1: &CudaSlice<u8>,
12881 b2: &CudaSlice<u8>,
12882 aq: &CudaSlice<i8>,
12883 ad: &CudaSlice<f32>,
12884 in_f: usize,
12885 out0: usize,
12886 out1: usize,
12887 out2: usize,
12888 row_bytes: usize,
12889 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12890 const ROWS_PER_BLOCK: u32 = 4;
12891 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12892 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12893 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12894 let f = self.func("qmatvec_q8_0_mmvq_fused3");
12895 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12896 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12897 let mut y2 = self.alloc_uninit::<f32>(out2)?;
12898 let cfg = LaunchConfig {
12899 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12900 block_dim: (32, ROWS_PER_BLOCK, 1),
12901 shared_mem_bytes: 0,
12902 };
12903 let (inf, o0, o1, o2, rbl) = (
12904 in_f as i32,
12905 out0 as i32,
12906 out1 as i32,
12907 out2 as i32,
12908 row_bytes as i64,
12909 );
12910 let __s_b = self.gpu.stream();
12911 let mut b = __s_b.launch_builder(&f);
12912 b.arg(b0)
12913 .arg(b1)
12914 .arg(b2)
12915 .arg(aq)
12916 .arg(ad)
12917 .arg(&mut y0)
12918 .arg(&mut y1)
12919 .arg(&mut y2)
12920 .arg(&inf)
12921 .arg(&o0)
12922 .arg(&o1)
12923 .arg(&o2)
12924 .arg(&rbl);
12925 unsafe {
12926 b.launch(cfg)?;
12927 }
12928 Ok((y0, y1, y2))
12929 }
12930
12931 #[allow(clippy::too_many_arguments)]
12933 pub fn qmatvec_q8_fused3_raw(
12934 &self,
12935 b0: &CudaSlice<u8>,
12936 b1: &CudaSlice<u8>,
12937 b2: &CudaSlice<u8>,
12938 x: &CudaSlice<f32>,
12939 in_f: usize,
12940 out0: usize,
12941 out1: usize,
12942 out2: usize,
12943 row_bytes: usize,
12944 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12945 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12946 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12947 }
12948
12949 pub fn matmul_q8_fused2_t(
12960 &self,
12961 w0: &crate::model::GpuTensor,
12962 w1: &crate::model::GpuTensor,
12963 aq: &CudaSlice<i8>,
12964 ad: &CudaSlice<f32>,
12965 m: usize,
12966 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12967 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12971 return Ok(None);
12972 }
12973 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12976 if m > 4 && !Self::b8_enabled() {
12977 return Ok(None);
12978 }
12979 return Ok(Some(self.e4m3_fused2_t_core(
12980 p0.0,
12981 p1.0,
12982 aq,
12983 ad,
12984 m,
12985 w0.in_features(),
12986 p0.1,
12987 p1.1,
12988 p0.2,
12989 p0.3,
12990 p1.3,
12991 )?));
12992 }
12993 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12994 return Ok(None);
12995 };
12996 Ok(Some(self.q8_fused2_t_core(
12997 p0.0,
12998 p1.0,
12999 aq,
13000 ad,
13001 m,
13002 w0.in_features(),
13003 p0.1,
13004 p1.1,
13005 p0.2,
13006 )?))
13007 }
13008
13009 #[allow(clippy::too_many_arguments)]
13010 fn q8_fused2_t_core(
13011 &self,
13012 b0: &CudaSlice<u8>,
13013 b1: &CudaSlice<u8>,
13014 aq: &CudaSlice<i8>,
13015 ad: &CudaSlice<f32>,
13016 m: usize,
13017 in_f: usize,
13018 out0: usize,
13019 out1: usize,
13020 row_bytes: usize,
13021 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13022 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13024 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13025 let f = self.func(match Self::batched_mcols(m) {
13026 2 => "qmatvec_q8_0_mmvq_fused2_b2",
13027 4 => "qmatvec_q8_0_mmvq_fused2_b4",
13028 _ => "qmatvec_q8_0_mmvq_fused2_b8",
13030 });
13031 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13032 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13033 let cfg = LaunchConfig {
13034 grid_dim: (nb0 + nb1, 1, 1),
13035 block_dim: (32, ROWS_PER_BLOCK, 1),
13036 shared_mem_bytes: 0,
13037 };
13038 let (inf, o0, o1, mi, rbl) = (
13039 in_f as i32,
13040 out0 as i32,
13041 out1 as i32,
13042 m as i32,
13043 row_bytes as i64,
13044 );
13045 let __s_b = self.gpu.stream();
13046 let mut b = __s_b.launch_builder(&f);
13047 b.arg(b0)
13048 .arg(b1)
13049 .arg(aq)
13050 .arg(ad)
13051 .arg(&mut y0)
13052 .arg(&mut y1)
13053 .arg(&inf)
13054 .arg(&o0)
13055 .arg(&o1)
13056 .arg(&mi)
13057 .arg(&rbl);
13058 unsafe {
13059 b.launch(cfg)?;
13060 }
13061 Ok((y0, y1))
13062 }
13063
13064 #[allow(clippy::too_many_arguments)]
13067 pub fn qmatvec_q8_fused2_t_raw(
13068 &self,
13069 b0: &CudaSlice<u8>,
13070 b1: &CudaSlice<u8>,
13071 x: &CudaSlice<f32>,
13072 m: usize,
13073 in_f: usize,
13074 out0: usize,
13075 out1: usize,
13076 row_bytes: usize,
13077 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13078 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13079 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
13080 }
13081
13082 #[allow(clippy::too_many_arguments)]
13085 pub fn matmul_q8_fused3_t(
13086 &self,
13087 w0: &crate::model::GpuTensor,
13088 w1: &crate::model::GpuTensor,
13089 w2: &crate::model::GpuTensor,
13090 aq: &CudaSlice<i8>,
13091 ad: &CudaSlice<f32>,
13092 m: usize,
13093 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13094 {
13095 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
13096 return Ok(None);
13097 }
13098 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
13099 return Ok(Some(self.e4m3_fused3_t_core(
13100 p0.0,
13101 p1.0,
13102 p2.0,
13103 aq,
13104 ad,
13105 m,
13106 w0.in_features(),
13107 p0.1,
13108 p1.1,
13109 p2.1,
13110 p0.2,
13111 p0.3,
13112 p1.3,
13113 p2.3,
13114 )?));
13115 }
13116 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
13117 return Ok(None);
13118 };
13119 Ok(Some(self.q8_fused3_t_core(
13120 p0.0,
13121 p1.0,
13122 p2.0,
13123 aq,
13124 ad,
13125 m,
13126 w0.in_features(),
13127 p0.1,
13128 p1.1,
13129 p2.1,
13130 p0.2,
13131 )?))
13132 }
13133
13134 #[allow(clippy::too_many_arguments)]
13135 fn q8_fused3_t_core(
13136 &self,
13137 b0: &CudaSlice<u8>,
13138 b1: &CudaSlice<u8>,
13139 b2: &CudaSlice<u8>,
13140 aq: &CudaSlice<i8>,
13141 ad: &CudaSlice<f32>,
13142 m: usize,
13143 in_f: usize,
13144 out0: usize,
13145 out1: usize,
13146 out2: usize,
13147 row_bytes: usize,
13148 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13149 const ROWS_PER_BLOCK: u32 = 4;
13150 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13151 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13152 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13153 let f = self.func(if Self::batched_mcols(m) == 2 {
13154 "qmatvec_q8_0_mmvq_fused3_b2"
13155 } else {
13156 "qmatvec_q8_0_mmvq_fused3_b4"
13157 });
13158 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13159 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13160 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13161 let cfg = LaunchConfig {
13162 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13163 block_dim: (32, ROWS_PER_BLOCK, 1),
13164 shared_mem_bytes: 0,
13165 };
13166 let (inf, o0, o1, o2, mi, rbl) = (
13167 in_f as i32,
13168 out0 as i32,
13169 out1 as i32,
13170 out2 as i32,
13171 m as i32,
13172 row_bytes as i64,
13173 );
13174 let __s_b = self.gpu.stream();
13175 let mut b = __s_b.launch_builder(&f);
13176 b.arg(b0)
13177 .arg(b1)
13178 .arg(b2)
13179 .arg(aq)
13180 .arg(ad)
13181 .arg(&mut y0)
13182 .arg(&mut y1)
13183 .arg(&mut y2)
13184 .arg(&inf)
13185 .arg(&o0)
13186 .arg(&o1)
13187 .arg(&o2)
13188 .arg(&mi)
13189 .arg(&rbl);
13190 unsafe {
13191 b.launch(cfg)?;
13192 }
13193 Ok((y0, y1, y2))
13194 }
13195
13196 #[allow(clippy::too_many_arguments)]
13198 pub fn qmatvec_q8_fused3_t_raw(
13199 &self,
13200 b0: &CudaSlice<u8>,
13201 b1: &CudaSlice<u8>,
13202 b2: &CudaSlice<u8>,
13203 x: &CudaSlice<f32>,
13204 m: usize,
13205 in_f: usize,
13206 out0: usize,
13207 out1: usize,
13208 out2: usize,
13209 row_bytes: usize,
13210 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13211 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13212 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
13213 }
13214
13215 pub fn q8_ffn_fuse2_on(&self) -> bool {
13219 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13220 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
13221 }
13222
13223 #[allow(clippy::type_complexity)]
13229 fn q8_fused_params<'w, const N: usize>(
13230 &self,
13231 ws: &[&'w crate::model::GpuTensor; N],
13232 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
13233 use crate::model::GpuTensor;
13234 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13235 return None;
13236 }
13237 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
13238 return None;
13239 }
13240 let in_f = ws[0].in_features();
13241 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
13242 for (i, w) in ws.iter().enumerate() {
13243 match w {
13244 GpuTensor::Quant {
13245 bytes,
13246 qtype,
13247 row_bytes,
13248 scale,
13249 ..
13250 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
13251 out[i] = Some((bytes, w.out_features(), *row_bytes))
13252 }
13253 _ => return None,
13254 }
13255 }
13256 Some(out.map(|o| o.unwrap()))
13257 }
13258
13259 pub fn e4m3_dual_on(&self) -> bool {
13262 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13263 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
13264 }
13265
13266 #[allow(clippy::type_complexity)]
13278 fn e4m3_fused_params<'w, const N: usize>(
13279 &self,
13280 ws: &[&'w crate::model::GpuTensor; N],
13281 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
13282 use crate::model::GpuTensor;
13283 if !self.e4m3_dual_on() {
13284 return None;
13285 }
13286 let in_f = ws[0].in_features();
13287 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
13288 for (i, w) in ws.iter().enumerate() {
13289 match w {
13290 GpuTensor::Quant {
13291 bytes,
13292 qtype,
13293 row_bytes,
13294 scale,
13295 rp,
13296 rp4,
13297 ..
13298 } if *qtype == QT_F8_E4M3
13299 && w.in_features() == in_f
13300 && *row_bytes == in_f
13301 && !*rp
13302 && rp4.is_none() =>
13303 {
13304 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
13305 }
13306 _ => return None,
13307 }
13308 }
13309 Some(out.map(|o| o.unwrap()))
13310 }
13311
13312 #[allow(clippy::too_many_arguments)]
13316 fn e4m3_fused2_core(
13317 &self,
13318 b0: &CudaSlice<u8>,
13319 b1: &CudaSlice<u8>,
13320 aq: &CudaSlice<i8>,
13321 ad: &CudaSlice<f32>,
13322 in_f: usize,
13323 out0: usize,
13324 out1: usize,
13325 row_bytes: usize,
13326 ws0: f32,
13327 ws1: f32,
13328 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13329 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13331 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13332 let f = self.func("qmatvec_e4m3_mmvq_fused2");
13333 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13334 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13335 let cfg = LaunchConfig {
13336 grid_dim: (nb0 + nb1, 1, 1),
13337 block_dim: (32, ROWS_PER_BLOCK, 1),
13338 shared_mem_bytes: 0,
13339 };
13340 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13341 let __s_b = self.gpu.stream();
13342 let mut b = __s_b.launch_builder(&f);
13343 b.arg(b0)
13344 .arg(b1)
13345 .arg(aq)
13346 .arg(ad)
13347 .arg(&mut y0)
13348 .arg(&mut y1)
13349 .arg(&inf)
13350 .arg(&o0)
13351 .arg(&o1)
13352 .arg(&rbl)
13353 .arg(&ws0)
13354 .arg(&ws1);
13355 unsafe {
13356 b.launch(cfg)?;
13357 }
13358 Ok((y0, y1))
13359 }
13360
13361 #[allow(clippy::too_many_arguments)]
13363 fn e4m3_fused3_core(
13364 &self,
13365 b0: &CudaSlice<u8>,
13366 b1: &CudaSlice<u8>,
13367 b2: &CudaSlice<u8>,
13368 aq: &CudaSlice<i8>,
13369 ad: &CudaSlice<f32>,
13370 in_f: usize,
13371 out0: usize,
13372 out1: usize,
13373 out2: usize,
13374 row_bytes: usize,
13375 ws0: f32,
13376 ws1: f32,
13377 ws2: f32,
13378 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13379 const ROWS_PER_BLOCK: u32 = 4;
13380 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13381 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13382 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13383 let f = self.func("qmatvec_e4m3_mmvq_fused3");
13384 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13385 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13386 let mut y2 = self.alloc_uninit::<f32>(out2)?;
13387 let cfg = LaunchConfig {
13388 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13389 block_dim: (32, ROWS_PER_BLOCK, 1),
13390 shared_mem_bytes: 0,
13391 };
13392 let (inf, o0, o1, o2, rbl) = (
13393 in_f as i32,
13394 out0 as i32,
13395 out1 as i32,
13396 out2 as i32,
13397 row_bytes as i64,
13398 );
13399 let __s_b = self.gpu.stream();
13400 let mut b = __s_b.launch_builder(&f);
13401 b.arg(b0)
13402 .arg(b1)
13403 .arg(b2)
13404 .arg(aq)
13405 .arg(ad)
13406 .arg(&mut y0)
13407 .arg(&mut y1)
13408 .arg(&mut y2)
13409 .arg(&inf)
13410 .arg(&o0)
13411 .arg(&o1)
13412 .arg(&o2)
13413 .arg(&rbl)
13414 .arg(&ws0)
13415 .arg(&ws1)
13416 .arg(&ws2);
13417 unsafe {
13418 b.launch(cfg)?;
13419 }
13420 Ok((y0, y1, y2))
13421 }
13422
13423 #[allow(clippy::too_many_arguments)]
13427 fn e4m3_fused2_t_core(
13428 &self,
13429 b0: &CudaSlice<u8>,
13430 b1: &CudaSlice<u8>,
13431 aq: &CudaSlice<i8>,
13432 ad: &CudaSlice<f32>,
13433 m: usize,
13434 in_f: usize,
13435 out0: usize,
13436 out1: usize,
13437 row_bytes: usize,
13438 ws0: f32,
13439 ws1: f32,
13440 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13441 const ROWS_PER_BLOCK: u32 = 4;
13442 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13443 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13444 let f = self.func(match Self::batched_mcols(m) {
13445 2 => "qmatvec_e4m3_mmvq_fused2_b2",
13446 4 => "qmatvec_e4m3_mmvq_fused2_b4",
13447 _ => "qmatvec_e4m3_mmvq_fused2_b8",
13448 });
13449 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13450 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13451 let cfg = LaunchConfig {
13452 grid_dim: (nb0 + nb1, 1, 1),
13453 block_dim: (32, ROWS_PER_BLOCK, 1),
13454 shared_mem_bytes: 0,
13455 };
13456 let (inf, o0, o1, mi, rbl) = (
13457 in_f as i32,
13458 out0 as i32,
13459 out1 as i32,
13460 m as i32,
13461 row_bytes as i64,
13462 );
13463 let __s_b = self.gpu.stream();
13464 let mut b = __s_b.launch_builder(&f);
13465 b.arg(b0)
13466 .arg(b1)
13467 .arg(aq)
13468 .arg(ad)
13469 .arg(&mut y0)
13470 .arg(&mut y1)
13471 .arg(&inf)
13472 .arg(&o0)
13473 .arg(&o1)
13474 .arg(&mi)
13475 .arg(&rbl);
13476 unsafe {
13477 b.launch(cfg)?;
13478 }
13479 if ws0 != 1.0 {
13480 self.scale_inplace(&mut y0, ws0, m * out0)?;
13481 }
13482 if ws1 != 1.0 {
13483 self.scale_inplace(&mut y1, ws1, m * out1)?;
13484 }
13485 Ok((y0, y1))
13486 }
13487
13488 #[allow(clippy::too_many_arguments)]
13490 fn e4m3_fused3_t_core(
13491 &self,
13492 b0: &CudaSlice<u8>,
13493 b1: &CudaSlice<u8>,
13494 b2: &CudaSlice<u8>,
13495 aq: &CudaSlice<i8>,
13496 ad: &CudaSlice<f32>,
13497 m: usize,
13498 in_f: usize,
13499 out0: usize,
13500 out1: usize,
13501 out2: usize,
13502 row_bytes: usize,
13503 ws0: f32,
13504 ws1: f32,
13505 ws2: f32,
13506 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13507 const ROWS_PER_BLOCK: u32 = 4;
13508 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13509 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13510 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13511 let f = self.func(if Self::batched_mcols(m) == 2 {
13512 "qmatvec_e4m3_mmvq_fused3_b2"
13513 } else {
13514 "qmatvec_e4m3_mmvq_fused3_b4"
13515 });
13516 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13517 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13518 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13519 let cfg = LaunchConfig {
13520 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13521 block_dim: (32, ROWS_PER_BLOCK, 1),
13522 shared_mem_bytes: 0,
13523 };
13524 let (inf, o0, o1, o2, mi, rbl) = (
13525 in_f as i32,
13526 out0 as i32,
13527 out1 as i32,
13528 out2 as i32,
13529 m as i32,
13530 row_bytes as i64,
13531 );
13532 let __s_b = self.gpu.stream();
13533 let mut b = __s_b.launch_builder(&f);
13534 b.arg(b0)
13535 .arg(b1)
13536 .arg(b2)
13537 .arg(aq)
13538 .arg(ad)
13539 .arg(&mut y0)
13540 .arg(&mut y1)
13541 .arg(&mut y2)
13542 .arg(&inf)
13543 .arg(&o0)
13544 .arg(&o1)
13545 .arg(&o2)
13546 .arg(&mi)
13547 .arg(&rbl);
13548 unsafe {
13549 b.launch(cfg)?;
13550 }
13551 if ws0 != 1.0 {
13552 self.scale_inplace(&mut y0, ws0, m * out0)?;
13553 }
13554 if ws1 != 1.0 {
13555 self.scale_inplace(&mut y1, ws1, m * out1)?;
13556 }
13557 if ws2 != 1.0 {
13558 self.scale_inplace(&mut y2, ws2, m * out2)?;
13559 }
13560 Ok((y0, y1, y2))
13561 }
13562
13563 pub fn qmatvec_e4m3_blk_mmvq(
13573 &self,
13574 bytes: &CudaSlice<u8>,
13575 aq: &CudaSlice<i8>,
13576 ad: &CudaSlice<f32>,
13577 scales: &CudaSlice<f32>,
13578 m: usize,
13579 in_f: usize,
13580 out_f: usize,
13581 row_bytes: usize,
13582 scale_cols: usize,
13583 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13584 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
13586 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
13587 )?;
13588 Ok(y)
13589 }
13590
13591 #[allow(clippy::too_many_arguments)]
13593 pub fn qmatvec_e4m3_blk_mmvq_into(
13594 &self,
13595 bytes: &CudaSlice<u8>,
13596 aq: &CudaSlice<i8>,
13597 ad: &CudaSlice<f32>,
13598 scales: &CudaSlice<f32>,
13599 m: usize,
13600 in_f: usize,
13601 out_f: usize,
13602 row_bytes: usize,
13603 scale_cols: usize,
13604 y: &mut CudaSlice<f32>,
13605 ) -> Result<(), Box<dyn std::error::Error>> {
13606 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
13608 let cfg = LaunchConfig {
13609 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
13610 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
13613 let (inf, outf, mi, rb, sc) = (
13614 in_f as i32,
13615 out_f as i32,
13616 m as i32,
13617 row_bytes as i64,
13618 scale_cols as i32,
13619 );
13620 let __s_b = self.gpu.stream();
13621 let mut b = __s_b.launch_builder(&f);
13622 b.arg(bytes)
13623 .arg(aq)
13624 .arg(ad)
13625 .arg(scales)
13626 .arg(&mut *y)
13627 .arg(&inf)
13628 .arg(&outf)
13629 .arg(&mi)
13630 .arg(&rb)
13631 .arg(&sc);
13632 unsafe {
13633 b.launch(cfg)?;
13634 }
13635 Ok(())
13636 }
13637
13638 #[allow(clippy::too_many_arguments)]
13644 pub fn qmatvec_e4m3_blk_mmvq_batched(
13645 &self,
13646 bytes: &CudaSlice<u8>,
13647 aq: &CudaSlice<i8>,
13648 ad: &CudaSlice<f32>,
13649 scales: &CudaSlice<f32>,
13650 m: usize,
13651 in_f: usize,
13652 out_f: usize,
13653 row_bytes: usize,
13654 scale_cols: usize,
13655 mcols: usize,
13656 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13657 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
13659 let name = match mcols {
13660 2 => "qmatvec_e4m3_blk_mmvq_b2",
13661 4 => "qmatvec_e4m3_blk_mmvq_b4",
13662 8 => "qmatvec_e4m3_blk_mmvq_b8",
13663 16 => "qmatvec_e4m3_blk_mmvq_b16",
13664 _ => {
13665 return Err(
13666 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
13667 );
13668 }
13669 };
13670 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13671 let f = self.func(name);
13672 let cfg = LaunchConfig {
13673 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
13674 block_dim: (32, ROWS_PER_BLOCK, 1),
13675 shared_mem_bytes: 0,
13676 };
13677 let (inf, outf, mi, rb, sc) = (
13678 in_f as i32,
13679 out_f as i32,
13680 m as i32,
13681 row_bytes as i64,
13682 scale_cols as i32,
13683 );
13684 let __s_b = self.gpu.stream();
13685 let mut b = __s_b.launch_builder(&f);
13686 b.arg(bytes)
13687 .arg(aq)
13688 .arg(ad)
13689 .arg(scales)
13690 .arg(&mut y)
13691 .arg(&inf)
13692 .arg(&outf)
13693 .arg(&mi)
13694 .arg(&rb)
13695 .arg(&sc);
13696 unsafe {
13697 b.launch(cfg)?;
13698 }
13699 Ok(y)
13700 }
13701
13702 #[allow(clippy::too_many_arguments)]
13705 pub fn qmatvec_e4m3_blk_batched_raw(
13706 &self,
13707 bytes: &CudaSlice<u8>,
13708 x: &CudaSlice<f32>,
13709 scales: &CudaSlice<f32>,
13710 m: usize,
13711 in_f: usize,
13712 out_f: usize,
13713 row_bytes: usize,
13714 scale_cols: usize,
13715 mcols: usize,
13716 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13717 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13718 self.qmatvec_e4m3_blk_mmvq_batched(
13719 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13720 )
13721 }
13722
13723 #[allow(clippy::too_many_arguments)]
13726 pub fn qmatvec_e4m3_blk_mmvq_raw(
13727 &self,
13728 bytes: &CudaSlice<u8>,
13729 x: &CudaSlice<f32>,
13730 scales: &CudaSlice<f32>,
13731 m: usize,
13732 in_f: usize,
13733 out_f: usize,
13734 row_bytes: usize,
13735 scale_cols: usize,
13736 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13737 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13738 self.qmatvec_e4m3_blk_mmvq(
13739 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13740 )
13741 }
13742
13743 #[allow(clippy::too_many_arguments)]
13746 pub fn qmatvec_e4m3_fused2_raw(
13747 &self,
13748 b0: &CudaSlice<u8>,
13749 b1: &CudaSlice<u8>,
13750 x: &CudaSlice<f32>,
13751 in_f: usize,
13752 out0: usize,
13753 out1: usize,
13754 row_bytes: usize,
13755 ws0: f32,
13756 ws1: f32,
13757 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13758 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13759 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13760 }
13761
13762 #[allow(clippy::too_many_arguments)]
13763 pub fn qmatvec_e4m3_fused3_raw(
13764 &self,
13765 b0: &CudaSlice<u8>,
13766 b1: &CudaSlice<u8>,
13767 b2: &CudaSlice<u8>,
13768 x: &CudaSlice<f32>,
13769 in_f: usize,
13770 out0: usize,
13771 out1: usize,
13772 out2: usize,
13773 row_bytes: usize,
13774 ws0: f32,
13775 ws1: f32,
13776 ws2: f32,
13777 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13778 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13779 self.e4m3_fused3_core(
13780 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13781 )
13782 }
13783
13784 #[allow(clippy::too_many_arguments)]
13785 pub fn qmatvec_e4m3_fused2_t_raw(
13786 &self,
13787 b0: &CudaSlice<u8>,
13788 b1: &CudaSlice<u8>,
13789 x: &CudaSlice<f32>,
13790 m: usize,
13791 in_f: usize,
13792 out0: usize,
13793 out1: usize,
13794 row_bytes: usize,
13795 ws0: f32,
13796 ws1: f32,
13797 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13798 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13799 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13800 }
13801
13802 #[allow(clippy::too_many_arguments)]
13803 pub fn qmatvec_e4m3_fused3_t_raw(
13804 &self,
13805 b0: &CudaSlice<u8>,
13806 b1: &CudaSlice<u8>,
13807 b2: &CudaSlice<u8>,
13808 x: &CudaSlice<f32>,
13809 m: usize,
13810 in_f: usize,
13811 out0: usize,
13812 out1: usize,
13813 out2: usize,
13814 row_bytes: usize,
13815 ws0: f32,
13816 ws1: f32,
13817 ws2: f32,
13818 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13819 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13820 self.e4m3_fused3_t_core(
13821 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13822 )
13823 }
13824
13825 fn try_e4m3_blk_pre(
13836 &self,
13837 w: &crate::model::GpuTensor,
13838 aq: &CudaSlice<i8>,
13839 ad: &CudaSlice<f32>,
13840 m: usize,
13841 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13842 use crate::model::GpuTensor;
13843 if let GpuTensor::Quant {
13844 bytes,
13845 qtype,
13846 row_bytes,
13847 blk: Some(g),
13848 ..
13849 } = w
13850 {
13851 if *qtype == QT_F8_E4M3_BLK {
13852 if (2..=16).contains(&m)
13858 && std::env::var("MEMRA_NO_BATCHED").is_err()
13859 && (m <= 4 || Self::b8_enabled())
13860 {
13861 let mcols = Self::batched_mcols(m);
13862 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13863 bytes,
13864 aq,
13865 ad,
13866 &g.scales,
13867 m,
13868 w.in_features(),
13869 w.out_features(),
13870 *row_bytes,
13871 g.cols,
13872 mcols,
13873 )?));
13874 }
13875 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13876 bytes,
13877 aq,
13878 ad,
13879 &g.scales,
13880 m,
13881 w.in_features(),
13882 w.out_features(),
13883 *row_bytes,
13884 g.cols,
13885 )?));
13886 }
13887 }
13888 Ok(None)
13889 }
13890
13891 fn try_e4m3_blk_prefill(
13938 &self,
13939 w: &crate::model::GpuTensor,
13940 x: &CudaSlice<f32>,
13941 m: usize,
13942 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13943 use crate::model::GpuTensor;
13944 let GpuTensor::Quant {
13945 bytes,
13946 qtype,
13947 blk: Some(g),
13948 ..
13949 } = w
13950 else {
13951 return Ok(None);
13952 };
13953 if *qtype != QT_F8_E4M3_BLK {
13954 return Ok(None);
13955 }
13956 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
13961 return Ok(Some(y));
13962 }
13963 let (in_f, out_f) = (w.in_features(), w.out_features());
13964 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
13965 let tmp = GpuTensor::Quant {
13966 bytes: slab,
13967 qtype: QT_Q8_0,
13968 row_bytes: in_f / 32 * 34,
13969 ne: vec![in_f as u64, out_f as u64],
13970 scale: 1.0,
13971 rp: false,
13972 #[cfg(memra_cutlass)]
13973 cutlass: None,
13974 fp8: None,
13975 blk: None,
13976 f16: None,
13977 rp4: None,
13978 };
13979 Ok(Some(self.matmul(&tmp, x, m)?))
13981 }
13982
13983 pub fn matmul_pre_noscale(
13984 &self,
13985 w: &crate::model::GpuTensor,
13986 aq: &CudaSlice<i8>,
13987 ad: &CudaSlice<f32>,
13988 m: usize,
13989 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
13990 use crate::model::GpuTensor;
13991 if m == 1 {
13995 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
13996 return Ok(Some((y, 1.0)));
13997 }
13998 }
13999 if m != 1 || !self.uses_q8_1_fast(w) {
14001 return Ok(None);
14002 }
14003 let in_f = w.in_features();
14004 let out_f = w.out_features();
14005 let (bytes, qtype, row_bytes, scale, rp) = match w {
14006 GpuTensor::Quant {
14007 bytes,
14008 qtype,
14009 row_bytes,
14010 scale,
14011 rp,
14012 ..
14013 } => (bytes, *qtype, *row_bytes, *scale, *rp),
14014 _ => return Ok(None),
14015 };
14016 if self.mmvq_supports(qtype) {
14018 let (mbytes, mrp) = match w {
14020 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
14021 _ => (bytes, rp),
14022 };
14023 let y = self.qmatvec_mmvq(
14024 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
14025 )?;
14026 return Ok(Some((y, scale)));
14027 }
14028 let name = match qtype {
14030 QT_Q8_0 => "qmatvec_q8_0_dp4a",
14031 QT_Q4_K => "qmatvec_q4_K_dp4a",
14032 QT_Q6_K => "qmatvec_q6_K_dp4a",
14033 QT_Q5_K => "qmatvec_q5_K_dp4a",
14034 QT_Q3_K => "qmatvec_q3_K_dp4a",
14035 QT_NVFP4 => {
14036 if rp {
14037 "qmatvec_nvfp4_dp4a_rp"
14038 } else {
14039 "qmatvec_nvfp4_dp4a"
14040 }
14041 }
14042 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
14043 _ => return Ok(None),
14044 };
14045 let f = self.func(name);
14046 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14047 let cfg = LaunchConfig {
14048 grid_dim: (out_f as u32, m as u32, 1),
14049 block_dim: (128, 1, 1),
14050 shared_mem_bytes: 0,
14051 };
14052 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14053 let __s_b = self.gpu.stream();
14054 let mut b = __s_b.launch_builder(&f);
14055 b.arg(bytes)
14056 .arg(aq)
14057 .arg(ad)
14058 .arg(&mut y)
14059 .arg(&inf)
14060 .arg(&outf)
14061 .arg(&mi)
14062 .arg(&rb);
14063 unsafe {
14064 b.launch(cfg)?;
14065 }
14066 Ok(Some((y, scale)))
14067 }
14068
14069 pub fn mmvq_supports(&self, qtype: i32) -> bool {
14072 if qtype == QT_F8_E4M3 {
14077 return true;
14078 }
14079 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14080 return false;
14081 }
14082 matches!(
14083 qtype,
14084 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
14085 )
14086 }
14087
14088 pub fn qmatvec_mmvq(
14093 &self,
14094 bytes: &CudaSlice<u8>,
14095 aq: &CudaSlice<i8>,
14096 ad: &CudaSlice<f32>,
14097 m: usize,
14098 in_f: usize,
14099 out_f: usize,
14100 qtype: i32,
14101 row_bytes: usize,
14102 scale: f32,
14103 rp: bool,
14104 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14105 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
14107 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
14108 )?;
14109 Ok(y)
14110 }
14111
14112 #[allow(clippy::too_many_arguments)]
14114 pub fn qmatvec_mmvq_into(
14115 &self,
14116 bytes: &CudaSlice<u8>,
14117 aq: &CudaSlice<i8>,
14118 ad: &CudaSlice<f32>,
14119 m: usize,
14120 in_f: usize,
14121 out_f: usize,
14122 qtype: i32,
14123 row_bytes: usize,
14124 scale: f32,
14125 rp: bool,
14126 y: &mut CudaSlice<f32>,
14127 ) -> Result<(), Box<dyn std::error::Error>> {
14128 debug_assert!(y.len() >= m * out_f);
14129 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
14135 && rp
14136 && m == 1
14137 && out_f >= 64
14138 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
14139 && {
14140 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14141 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
14142 }
14143 {
14144 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
14145 let cfg = LaunchConfig {
14146 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
14147 block_dim: (32, 2, 1),
14148 shared_mem_bytes: 0,
14149 };
14150 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
14151 let __s_b = self.gpu.stream();
14152 let mut b = __s_b.launch_builder(&f);
14153 b.arg(bytes)
14154 .arg(aq)
14155 .arg(ad)
14156 .arg(&mut *y)
14157 .arg(&inf)
14158 .arg(&outf)
14159 .arg(&mi)
14160 .arg(&rb);
14161 unsafe {
14162 b.launch(cfg)?;
14163 }
14164 if scale != 1.0 {
14165 self.scale_inplace(y, scale, out_f)?;
14166 }
14167 return Ok(());
14168 }
14169 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
14178 2
14179 } else {
14180 1
14181 };
14182 if m == 1 && qtype == QT_Q4_0 {
14187 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
14188 mr = *Q40MR.get_or_init(|| {
14191 std::env::var("MEMRA_Q40_MR")
14192 .ok()
14193 .and_then(|v| v.parse().ok())
14194 .unwrap_or(1)
14195 });
14196 }
14197 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
14208 let q5_force = q5_mode.as_deref() == Some("2");
14209 let q5_il = qtype == QT_Q5_K
14212 && m == 1
14213 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
14214 if q5_il && !q5_force && out_f > 65536 {
14215 mr = 1;
14216 }
14217 if qtype == QT_Q4_0 && rp && mr != 1 {
14220 mr = 2;
14221 }
14222 if qtype == QT_Q8_0 && rp {
14226 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
14227 mr = *Q80MR.get_or_init(|| {
14228 std::env::var("MEMRA_Q80_MR")
14229 .ok()
14230 .and_then(|v| v.parse().ok())
14231 .unwrap_or(1)
14232 });
14233 }
14234 let name = match (qtype, mr, rp) {
14235 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
14236 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
14237 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
14238 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
14239 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
14240 (QT_Q5_K, 2, _) => {
14241 if q5_il {
14242 "qmatvec_q5_K_mmvq_mr2_il"
14243 } else {
14244 "qmatvec_q5_K_mmvq_mr2"
14245 }
14246 }
14247 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
14248 (QT_Q8_0, _, true)
14253 if in_f % 1024 == 0 && {
14254 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14255 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
14256 } =>
14257 {
14258 "qmatvec_q8_0_mmvq_rpca"
14259 }
14260 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
14261 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
14262 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
14266 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
14267 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
14268 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
14269 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
14270 (QT_Q5_K, _, _) => {
14271 if q5_il {
14272 "qmatvec_q5_K_mmvq_il"
14273 } else {
14274 "qmatvec_q5_K_mmvq"
14275 }
14276 }
14277 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
14278 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
14279 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
14280 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
14281 };
14282 let f = self.func(name);
14283 let rows_per_block = ROWS_PER_BLOCK * mr;
14285 let cfg = LaunchConfig {
14286 grid_dim: (
14287 (out_f as u32 + rows_per_block - 1) / rows_per_block,
14288 m as u32,
14289 1,
14290 ),
14291 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
14294 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14295 let __s_b = self.gpu.stream();
14296 let mut b = __s_b.launch_builder(&f);
14297 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
14302 if Self::pdl_on()
14305 && Self::pdl_mmvq_on()
14306 && Self::pdl_nvfp4q8_on()
14307 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
14308 {
14309 use cudarc::driver::{DevicePtr, DevicePtrMut};
14310 let s = &self.gpu.stream();
14311 let (pw, _g0) = bytes.device_ptr(s);
14312 let (paq, _g1) = aq.device_ptr(s);
14313 let (pad, _g2) = ad.device_ptr(s);
14314 let (py, _g3) = y.device_ptr_mut(s);
14315 let mut ps = [
14316 &pw as *const _ as *mut std::ffi::c_void,
14317 &paq as *const _ as *mut _,
14318 &pad as *const _ as *mut _,
14319 &py as *const _ as *mut _,
14320 &inf as *const _ as *mut _,
14321 &outf as *const _ as *mut _,
14322 &mi as *const _ as *mut _,
14323 &rb as *const _ as *mut _,
14324 &scale as *const _ as *mut _,
14325 ];
14326 unsafe {
14327 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14328 }
14329 return Ok(());
14330 }
14331 b.arg(bytes)
14332 .arg(aq)
14333 .arg(ad)
14334 .arg(&mut *y)
14335 .arg(&inf)
14336 .arg(&outf)
14337 .arg(&mi)
14338 .arg(&rb)
14339 .arg(&scale);
14340 unsafe {
14341 b.launch(cfg)?;
14342 }
14343 } else if Self::pdl_on()
14344 && Self::pdl_mmvq_on()
14345 && (matches!(
14346 name,
14347 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
14348 ) || (Self::pdl_nvfp4q8_on()
14349 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
14350 {
14351 {
14355 use cudarc::driver::{DevicePtr, DevicePtrMut};
14356 let s = &self.gpu.stream();
14357 let (pw, _g0) = bytes.device_ptr(s);
14358 let (paq, _g1) = aq.device_ptr(s);
14359 let (pad, _g2) = ad.device_ptr(s);
14360 let (py, _g3) = y.device_ptr_mut(s);
14361 let mut ps = [
14362 &pw as *const _ as *mut std::ffi::c_void,
14363 &paq as *const _ as *mut _,
14364 &pad as *const _ as *mut _,
14365 &py as *const _ as *mut _,
14366 &inf as *const _ as *mut _,
14367 &outf as *const _ as *mut _,
14368 &mi as *const _ as *mut _,
14369 &rb as *const _ as *mut _,
14370 ];
14371 unsafe {
14372 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14373 }
14374 }
14375 if scale != 1.0 {
14376 self.scale_inplace(y, scale, m * out_f)?;
14377 }
14378 } else {
14379 b.arg(bytes)
14380 .arg(aq)
14381 .arg(ad)
14382 .arg(&mut *y)
14383 .arg(&inf)
14384 .arg(&outf)
14385 .arg(&mi)
14386 .arg(&rb);
14387 unsafe {
14388 b.launch(cfg)?;
14389 }
14390 if scale != 1.0 {
14391 self.scale_inplace(y, scale, m * out_f)?;
14392 }
14393 }
14394 Ok(())
14395 }
14396
14397 pub fn qmatvec_mmvq_raw(
14401 &self,
14402 bytes: &CudaSlice<u8>,
14403 x: &CudaSlice<f32>,
14404 m: usize,
14405 in_f: usize,
14406 out_f: usize,
14407 qtype: i32,
14408 row_bytes: usize,
14409 rp: bool,
14410 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14411 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14412 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
14413 }
14414
14415 pub fn batched_supports(&self, qtype: i32) -> bool {
14419 matches!(
14420 qtype,
14421 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
14422 )
14423 }
14424
14425 pub fn iq_fast_enabled() -> bool {
14433 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14434 *ON.get_or_init(|| {
14435 std::env::var("MEMRA_IQ_FAST")
14436 .map(|v| v != "0")
14437 .unwrap_or(true)
14438 })
14439 }
14440
14441 pub fn b8_enabled() -> bool {
14444 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14445 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
14446 }
14447
14448 pub fn batched_mcols(m: usize) -> usize {
14450 if m == 2 {
14451 2
14452 } else if m <= 4 {
14453 4
14454 } else if m <= 8 {
14455 8
14456 } else {
14457 16
14458 }
14459 }
14460
14461 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
14466 Some(match (qtype, mcols) {
14467 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
14468 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
14469 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
14470 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
14476 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
14477 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
14478 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
14479 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
14482 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
14483 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
14484 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
14485 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
14488 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
14489 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
14490 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
14491 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
14492 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
14493 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
14494 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
14495 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
14499 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
14500 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
14501 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
14502 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
14506 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
14507 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
14508 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
14509 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
14510 _ => return None,
14511 })
14512 }
14513
14514 pub fn sm_count(&self) -> i32 {
14549 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14550 *SMS.get_or_init(|| {
14551 use cudarc::driver::sys::CUdevice_attribute_enum as A;
14552 self.gpu
14553 .ctx
14554 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14555 .unwrap_or(82)
14556 })
14557 }
14558
14559 pub fn batched_variant(
14560 &self,
14561 _m: usize,
14562 in_f: usize,
14563 out_f: usize,
14564 qtype: i32,
14565 row_bytes: usize,
14566 mcols: usize,
14567 rp: bool,
14568 ) -> &'static str {
14569 if qtype == QT_Q8_0 {
14574 return if rp { "rp" } else { "base" };
14575 }
14576 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14577 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
14578 Ok("base") => "base",
14579 Ok("pf") => "pf",
14580 Ok("r2") => "r2",
14581 Ok("r2w8") => "r2w8",
14582 Ok("pfr2") => "pfr2",
14583 Ok("ca") => "ca",
14584 Ok("car2") => "car2",
14585 Ok("rp") => "rp",
14588 Ok("rpr2") => "rpr2",
14589 Ok("rpr2w8") => "rpr2w8",
14590 Ok("rpca") => "rpca",
14593 Ok("rpcar2") => "rpcar2",
14594 Ok("rpsc") => "rpsc",
14601 Ok("rpms") => "rpms",
14602 Ok("rpmsc") => "rpmsc",
14603 Ok("rpks") => "rpks",
14604 Ok("rpksc") => "rpksc",
14605 _ => "auto",
14606 });
14607 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
14611 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14616 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
14617 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
14618 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
14619 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14620 let sms = *SMS.get_or_init(|| {
14621 use cudarc::driver::sys::CUdevice_attribute_enum as A;
14622 self.gpu
14623 .ctx
14624 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14625 .unwrap_or(82)
14626 });
14627 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
14647 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14650 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
14651 Ok("base") => "base",
14652 Ok("r2") => "r2",
14653 Ok("r2w8") => "r2w8",
14654 _ => "auto",
14655 });
14656 let variant: &'static str = if qtype == QT_Q4_0 {
14657 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14661 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
14662 Ok("base") => "base",
14668 Ok("r2") => "r2",
14669 Ok("ms") => "ms",
14670 Ok("sm") => "sm",
14671 Ok("la") => "la",
14672 _ => "auto",
14673 });
14674 let v = if q40 != "auto" {
14675 q40
14676 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
14677 "r2"
14678 } else {
14679 "base"
14680 };
14681 if rp {
14686 match v {
14687 "ms" => "r2ms_rp",
14688 "sm" => "r2sm_rp",
14689 "la" => "r2la_rp",
14690 "r2" => "r2_rp",
14691 _ => "rp",
14692 }
14693 } else if matches!(v, "ms" | "sm" | "la") {
14694 "r2"
14695 } else {
14696 v
14697 }
14698 } else if qtype != QT_NVFP4 && !kq_r2 {
14699 "base"
14700 } else if kq_r2 && rp {
14701 "rp"
14705 } else if kq_r2 {
14706 if kq_bv != "auto" {
14709 if kq_bv == "r2w8" && mcols != 4 {
14710 "r2"
14711 } else {
14712 kq_bv
14713 }
14714 } else if bv != "auto" {
14715 match bv {
14716 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14717 "r2w8" | "rpr2w8" => {
14718 if mcols != 4 {
14719 "r2"
14720 } else {
14721 "r2w8"
14722 }
14723 }
14724 _ => "base", }
14726 } else {
14727 let blocks = (out_f + 7) / 8;
14728 let waves = blocks as f64 / (7 * sms as usize) as f64;
14729 let filled = blocks >= 4 * sms as usize;
14730 let use_r2 = if qtype == QT_Q4_K {
14731 filled
14732 } else {
14733 waves >= 2.0
14734 };
14735 if use_r2 { "r2" } else { "base" }
14736 }
14737 } else if bv != "auto" {
14738 let v = if bv == "r2w8" && mcols == 2 {
14743 "r2"
14744 } else if bv == "ca" && (!ca_ok || mcols == 8) {
14745 "pf"
14746 } else if bv == "car2" && (!ca_ok || mcols == 8) {
14747 "r2"
14748 } else if bv == "pfr2" && mcols == 8 {
14749 "r2"
14750 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14751 "rpr2"
14752 }
14753 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14755 if mcols == 8 { "rpr2w8" } else { "rpr2" }
14756 } else if bv == "rpcar2" && mcols == 2 {
14757 "rpca"
14758 }
14759 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14762 "rpr2"
14763 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14764 "rpr2"
14765 } else {
14766 bv
14767 };
14768 if rp {
14769 match v {
14770 "base" | "pf" | "ca" | "rp" => "rp",
14771 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14772 "r2w8" | "rpr2w8" => {
14773 if mcols == 2 {
14774 "rpr2"
14775 } else {
14776 "rpr2w8"
14777 }
14778 }
14779 other => other, }
14781 } else {
14782 v
14783 }
14784 } else if mcols == 8 {
14785 if rp {
14796 if sc_ok { "rpsc" } else { "rpr2w8" }
14797 } else {
14798 "r2w8"
14799 }
14800 } else if mcols >= 4 {
14801 let blocks = (out_f + 7) / 8;
14805 let r7 = 7 * sms as usize;
14806 let r8 = 8 * sms as usize;
14807 let waves = blocks as f64 / r7 as f64;
14808 let filled = blocks >= 4 * sms as usize;
14809 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14813 if rp { "rpr2w8" } else { "r2w8" }
14817 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14818 if rp { "rpr2" } else { "r2" }
14821 } else {
14822 if rp { "rp" } else { "pf" }
14826 }
14827 } else if in_f >= 6144 {
14828 if rp { "rpr2" } else { "r2" }
14832 } else if rp {
14833 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14838 if sc_ok && waves >= 0.9 && waves <= 1.1 {
14839 "rpsc"
14840 } else {
14841 "rp"
14842 }
14843 } else {
14844 "base"
14845 };
14846 variant
14847 }
14848
14849 pub fn qmatvec_mmvq_batched(
14850 &self,
14851 bytes: &CudaSlice<u8>,
14852 aq: &CudaSlice<i8>,
14853 ad: &CudaSlice<f32>,
14854 m: usize,
14855 in_f: usize,
14856 out_f: usize,
14857 qtype: i32,
14858 row_bytes: usize,
14859 mcols: usize,
14860 scale: f32,
14861 rp: bool,
14862 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14863 const ROWS_PER_BLOCK: u32 = 4;
14864 let forced: Option<&'static str> = {
14869 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14870 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14871 .as_deref()
14872 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14873 };
14874 let variant = match forced {
14875 Some(v) if !rp || v.contains("rp") => v,
14876 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14877 };
14878 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14879 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14880 })?;
14881 let variant = if mcols == 16 {
14885 if rp { "rp" } else { "base" }
14886 } else {
14887 variant
14888 };
14889 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14896 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14897 if b567
14898 && qtype == QT_NVFP4
14899 && rp
14900 && mcols == 8
14901 && (5..=7).contains(&m)
14902 && matches!(variant, "rpsc" | "rpr2w8")
14903 {
14904 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14905 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14907 let cfg = LaunchConfig {
14908 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14909 block_dim: (32, ROWS_PER_BLOCK, 1),
14910 shared_mem_bytes: 0,
14911 };
14912 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14913 let __s_b = self.gpu.stream();
14914 let mut b = __s_b.launch_builder(&f);
14915 b.arg(bytes)
14916 .arg(aq)
14917 .arg(ad)
14918 .arg(&mut y)
14919 .arg(&inf)
14920 .arg(&outf)
14921 .arg(&mi)
14922 .arg(&rb);
14923 unsafe {
14924 b.launch(cfg)?;
14925 }
14926 if scale != 1.0 {
14927 self.scale_inplace(&mut y, scale, m * out_f)?;
14928 }
14929 return Ok(y);
14930 }
14931 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14932 "base" => (base_name.into(), ROWS_PER_BLOCK),
14933 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14934 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14935 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14936 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14940 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14941 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14942 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14943 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14944 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14945 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14946 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
14948 debug_assert!(
14949 !rp || name.contains("_rp"),
14950 "rp weight dispatched to a GGUF-layout kernel"
14951 );
14952 let f = self.func(&name);
14953 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14954 let smem = if name.contains("_r2sm_rp") {
14956 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
14957 } else {
14958 0
14959 };
14960 let cfg = LaunchConfig {
14961 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14962 block_dim: (32, ROWS_PER_BLOCK, 1),
14963 shared_mem_bytes: smem,
14964 };
14965 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14966 let __s_b = self.gpu.stream();
14967 let mut b = __s_b.launch_builder(&f);
14968 b.arg(bytes)
14969 .arg(aq)
14970 .arg(ad)
14971 .arg(&mut y)
14972 .arg(&inf)
14973 .arg(&outf)
14974 .arg(&mi)
14975 .arg(&rb);
14976 unsafe {
14977 b.launch(cfg)?;
14978 }
14979 if scale != 1.0 {
14980 self.scale_inplace(&mut y, scale, m * out_f)?;
14981 }
14982 Ok(y)
14983 }
14984
14985 pub fn qmatvec_batched_raw(
14989 &self,
14990 bytes: &CudaSlice<u8>,
14991 x: &CudaSlice<f32>,
14992 m: usize,
14993 in_f: usize,
14994 out_f: usize,
14995 qtype: i32,
14996 row_bytes: usize,
14997 mcols: usize,
14998 rp: bool,
14999 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15000 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15001 self.qmatvec_mmvq_batched(
15002 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
15003 )
15004 }
15005
15006 pub fn qmatvec_nvfp4_batched_raw(
15008 &self,
15009 bytes: &CudaSlice<u8>,
15010 x: &CudaSlice<f32>,
15011 m: usize,
15012 in_f: usize,
15013 out_f: usize,
15014 row_bytes: usize,
15015 mcols: usize,
15016 rp: bool,
15017 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15018 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
15019 }
15020
15021 fn try_fp4_gemm(
15025 &self,
15026 w: &crate::model::GpuTensor,
15027 x: &CudaSlice<f32>,
15028 m: usize,
15029 in_f: usize,
15030 out_f: usize,
15031 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15032 use crate::model::GpuTensor;
15033 if cfg!(memra_portable_cuda) {
15034 return Ok(None);
15035 }
15036 if std::env::var("MEMRA_FP4").is_err() {
15037 return Ok(None);
15038 }
15039 #[cfg(memra_cutlass)]
15048 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
15049 if let GpuTensor::Quant {
15050 bytes,
15051 qtype,
15052 scale,
15053 row_bytes,
15054 cutlass,
15055 ..
15056 } = w
15057 {
15058 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
15059 if let Some(cw) = cutlass {
15060 let y = self.cutlass_fp4_gemm(
15062 &cw.b_packed,
15063 &cw.sfb_swizzled,
15064 x,
15065 *scale,
15066 m,
15067 out_f,
15068 in_f,
15069 )?;
15070 return Ok(Some(y));
15071 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
15072 let (b_packed, sfb_sw) =
15077 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
15078 let y =
15079 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
15080 return Ok(Some(y));
15081 }
15082 }
15083 }
15084 }
15085 if let GpuTensor::Quant {
15086 bytes,
15087 qtype,
15088 row_bytes,
15089 scale,
15090 rp,
15091 ..
15092 } = w
15093 {
15094 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
15097 let y =
15098 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
15099 return Ok(Some(y));
15100 }
15101 }
15102 Ok(None)
15103 }
15104
15105 pub fn rms_norm_f16out(
15109 &self,
15110 x: &CudaSlice<f32>,
15111 w: &CudaSlice<f32>,
15112 dst: &mut CudaSlice<f32>,
15113 dst16: &mut CudaSlice<u8>,
15114 ncols: usize,
15115 nrows: usize,
15116 eps: f32,
15117 ) -> Result<(), Box<dyn std::error::Error>> {
15118 let f = self.func("rms_norm_f16out_f32");
15119 let cfg = LaunchConfig {
15120 grid_dim: (nrows as u32, 1, 1),
15121 block_dim: (rms_block(), 1, 1),
15122 shared_mem_bytes: 0,
15123 };
15124 let (nc, e) = (ncols as i32, eps);
15125 let __s_b = self.gpu.stream();
15126 let mut b = __s_b.launch_builder(&f);
15127 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
15128 unsafe {
15129 b.launch(cfg)?;
15130 }
15131 Ok(())
15132 }
15133
15134 #[allow(clippy::too_many_arguments)]
15137 pub fn add_rms_norm_f16out(
15138 &self,
15139 a: &CudaSlice<f32>,
15140 b: &CudaSlice<f32>,
15141 w: &CudaSlice<f32>,
15142 res: &mut CudaSlice<f32>,
15143 dst: &mut CudaSlice<f32>,
15144 dst16: &mut CudaSlice<u8>,
15145 ncols: usize,
15146 nrows: usize,
15147 eps: f32,
15148 ) -> Result<(), Box<dyn std::error::Error>> {
15149 let f = self.func("add_rms_norm_f16out_f32");
15150 let cfg = LaunchConfig {
15151 grid_dim: (nrows as u32, 1, 1),
15152 block_dim: (rms_block(), 1, 1),
15153 shared_mem_bytes: 0,
15154 };
15155 let (nc, e) = (ncols as i32, eps);
15156 let __s_lb = self.gpu.stream();
15157 let mut lb = __s_lb.launch_builder(&f);
15158 lb.arg(a)
15159 .arg(b)
15160 .arg(w)
15161 .arg(res)
15162 .arg(dst)
15163 .arg(dst16)
15164 .arg(&nc)
15165 .arg(&e);
15166 unsafe {
15167 lb.launch(cfg)?;
15168 }
15169 Ok(())
15170 }
15171
15172 pub fn matmul_group_xh(
15175 &self,
15176 ws: &[&crate::model::GpuTensor],
15177 x: &CudaSlice<f32>,
15178 xh: &CudaSlice<u8>,
15179 m: usize,
15180 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15181 let mut out = Vec::with_capacity(ws.len());
15182 let in_f = ws[0].in_features();
15183 for w in ws {
15184 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
15185 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
15186 out.push(y);
15187 continue;
15188 }
15189 }
15190 out.push(self.matmul(w, x, m)?);
15191 }
15192 Ok(out)
15193 }
15194
15195 pub fn gdn_pad_mask(
15198 &self,
15199 beta: &mut CudaSlice<f32>,
15200 g_log: &mut CudaSlice<f32>,
15201 len_d: &CudaSlice<i32>,
15202 h: usize,
15203 t: usize,
15204 ) -> Result<(), Box<dyn std::error::Error>> {
15205 let f = self.func("gdn_pad_mask_f32");
15206 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
15207 let (hi, ti) = (h as i32, t as i32);
15208 let __s_b = self.gpu.stream();
15209 let mut b = __s_b.launch_builder(&f);
15210 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
15211 unsafe {
15212 b.launch(cfg)?;
15213 }
15214 Ok(())
15215 }
15216
15217 pub fn row_gather_dev(
15220 &self,
15221 src: &CudaSlice<f32>,
15222 dst: &mut CudaSlice<f32>,
15223 len_d: &CudaSlice<i32>,
15224 ncols: usize,
15225 ) -> Result<(), Box<dyn std::error::Error>> {
15226 let f = self.func("row_gather_dev_f32");
15227 let cfg = LaunchConfig::for_num_elems(ncols as u32);
15228 let nc = ncols as i32;
15229 let __s_b = self.gpu.stream();
15230 let mut b = __s_b.launch_builder(&f);
15231 b.arg(src).arg(dst).arg(len_d).arg(&nc);
15232 unsafe {
15233 b.launch(cfg)?;
15234 }
15235 Ok(())
15236 }
15237
15238 pub fn matmul_group(
15245 &self,
15246 ws: &[&crate::model::GpuTensor],
15247 x: &CudaSlice<f32>,
15248 m: usize,
15249 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15250 use crate::model::GpuTensor;
15251 let mut out = Vec::with_capacity(ws.len());
15252 let any_mirror = ws
15253 .iter()
15254 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
15255 if m >= 16 && any_mirror && !self.verify_exact_on() {
15256 let in_f = ws[0].in_features();
15257 let xh = self.f16_act(x, m * in_f, in_f)?;
15258 for w in ws {
15259 if w.in_features() == in_f {
15260 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
15261 out.push(y);
15262 continue;
15263 }
15264 }
15265 out.push(self.matmul(w, x, m)?);
15266 }
15267 return Ok(out);
15268 }
15269 for w in ws {
15270 out.push(self.matmul(w, x, m)?);
15271 }
15272 Ok(out)
15273 }
15274
15275 pub fn matmul_group_multi(
15282 &self,
15283 ws: &[&crate::model::GpuTensor],
15284 xs: &[&CudaSlice<f32>],
15285 ms: &[usize],
15286 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
15287 assert_eq!(xs.len(), ms.len());
15288 let in_f = ws[0].in_features();
15289 let total: usize = ms.iter().sum();
15290 let mut xcat = self.uninit(total * in_f)?;
15291 let mut off = 0usize;
15292 for (x, &m) in xs.iter().zip(ms) {
15293 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
15294 off += m;
15295 }
15296 let ys = self.matmul_group(ws, &xcat, total)?;
15297 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
15298 for (w, y) in ws.iter().zip(ys) {
15299 let out_f = w.out_features();
15300 let mut off = 0usize;
15301 for (s, &m) in ms.iter().enumerate() {
15302 let mut ys_s = self.uninit(m * out_f)?;
15303 let src = y.slice(off * out_f..(off + m) * out_f);
15304 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
15305 out[s].push(ys_s);
15306 off += m;
15307 }
15308 }
15309 Ok(out)
15310 }
15311
15312 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
15322 use crate::model::GpuTensor;
15323 if !legacy_quant_gemm_allowed(
15324 cfg!(memra_portable_cuda),
15325 cfg!(memra_hopper_mma),
15326 std::env::var_os("MEMRA_NO_GEMM").is_some(),
15327 ) {
15328 return false;
15329 }
15330 match w {
15331 GpuTensor::Quant { qtype, .. } => {
15332 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
15333 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
15334 }
15335 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
15336 }
15337 }
15338
15339 pub fn qmatvec_gemm(
15346 &self,
15347 w: &crate::model::GpuTensor,
15348 aq: &CudaSlice<i8>,
15349 ad: &CudaSlice<f32>,
15350 m: usize,
15351 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15352 use crate::model::GpuTensor;
15353 let in_f = w.in_features();
15354 let out_f = w.out_features();
15355 let (bytes, qtype, row_bytes, scale, rp) = match w {
15356 GpuTensor::Quant {
15357 bytes,
15358 qtype,
15359 row_bytes,
15360 scale,
15361 rp,
15362 ..
15363 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15364 _ => unreachable!("gemm_supports guaranteed Quant"),
15365 };
15366 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
15372 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
15373 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
15374 if scale != 1.0 {
15375 self.scale_inplace(&mut y, scale, m * out_f)?;
15376 }
15377 return Ok(y);
15378 }
15379 }
15380 let name = match qtype {
15381 QT_Q8_0 => "qmatvec_gemm_q8_0",
15382 QT_Q4_K => "qmatvec_gemm_q4_K",
15383 QT_Q4_0 => {
15384 if rp {
15385 "qmatvec_gemm_q4_0_rp"
15386 } else {
15387 "qmatvec_gemm_q4_0"
15388 }
15389 }
15390 QT_Q5_K => "qmatvec_gemm_q5_K",
15391 QT_Q6_K => "qmatvec_gemm_q6_K",
15392 QT_NVFP4 => {
15393 if rp {
15394 "qmatvec_gemm_nvfp4_rp"
15395 } else {
15396 "qmatvec_gemm_nvfp4"
15397 }
15398 }
15399 _ => unreachable!(),
15400 };
15401 let f = self.func(name);
15402 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);
15407 let k1_tile = if is_k1 {
15409 k1_launch_override().unwrap_or((128, 128, 8))
15410 } else {
15411 (128, 128, 8)
15412 };
15413 let (bm, bn): (u32, u32) = if is_k1 {
15414 (k1_tile.0, k1_tile.1)
15415 } else {
15416 (64, 256)
15417 };
15418 let warps: u32 = if is_k1 {
15419 k1_tile.2
15420 } else {
15421 match qtype {
15422 QT_NVFP4 => 8,
15423 _ => 4,
15424 }
15425 };
15426 let cfg = LaunchConfig {
15427 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15428 block_dim: (32, warps, 1),
15429 shared_mem_bytes: 0,
15430 };
15431 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15432 let __s_b = self.gpu.stream();
15433 let mut b = __s_b.launch_builder(&f);
15434 b.arg(bytes)
15435 .arg(aq)
15436 .arg(ad)
15437 .arg(&mut y)
15438 .arg(&inf)
15439 .arg(&outf)
15440 .arg(&mi)
15441 .arg(&rb);
15442 unsafe {
15443 b.launch(cfg)?;
15444 }
15445 if scale != 1.0 {
15446 self.scale_inplace(&mut y, scale, m * out_f)?;
15447 }
15448 Ok(y)
15449 }
15450
15451 pub fn qmatvec_gemm_raw(
15456 &self,
15457 bytes: &CudaSlice<u8>,
15458 x: &CudaSlice<f32>,
15459 m: usize,
15460 in_f: usize,
15461 out_f: usize,
15462 qtype: i32,
15463 row_bytes: usize,
15464 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15465 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15466 let name = match qtype {
15467 QT_Q8_0 => "qmatvec_gemm_q8_0",
15468 QT_Q4_K => "qmatvec_gemm_q4_K",
15469 QT_Q4_0 => "qmatvec_gemm_q4_0",
15470 QT_Q5_K => "qmatvec_gemm_q5_K",
15471 QT_Q6_K => "qmatvec_gemm_q6_K",
15472 QT_NVFP4 => "qmatvec_gemm_nvfp4",
15473 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
15474 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
15475 };
15476 let f = self.func(name);
15477 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);
15481 let k1_tile = if is_k1 {
15483 k1_launch_override().unwrap_or((128, 128, 8))
15484 } else {
15485 (128, 128, 8)
15486 };
15487 let (bm, bn): (u32, u32) = if is_k1 {
15488 (k1_tile.0, k1_tile.1)
15489 } else {
15490 (64, 256)
15491 };
15492 let warps: u32 = if is_k1 {
15493 k1_tile.2
15494 } else {
15495 match qtype {
15496 QT_NVFP4 | QT_NVFP4_RP => 8,
15497 _ => 4,
15498 }
15499 };
15500 let cfg = LaunchConfig {
15501 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15502 block_dim: (32, warps, 1),
15503 shared_mem_bytes: 0,
15504 };
15505 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15506 let __s_b = self.gpu.stream();
15507 let mut b = __s_b.launch_builder(&f);
15508 b.arg(bytes)
15509 .arg(&aq)
15510 .arg(&ad)
15511 .arg(&mut y)
15512 .arg(&inf)
15513 .arg(&outf)
15514 .arg(&mi)
15515 .arg(&rb);
15516 unsafe {
15517 b.launch(cfg)?;
15518 }
15519 Ok(y)
15520 }
15521
15522 pub fn qmatvec_gemm_q8_0_wgmma_raw(
15529 &self,
15530 rp4: &CudaSlice<u8>,
15531 aq: &CudaSlice<i8>,
15532 ad: &CudaSlice<f32>,
15533 m: usize,
15534 in_f: usize,
15535 out_f: usize,
15536 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15537 assert!(
15538 out_f % 64 == 0 && in_f % 32 == 0,
15539 "wgmma GEMM needs out_f%64==0, in_f%32==0"
15540 );
15541 let f = self.func("qmatvec_gemm_q8_0_wgmma");
15542 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
15544 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
15545 block_dim: (128, 1, 1),
15546 shared_mem_bytes: 0,
15547 };
15548 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
15549 let __s_b = self.gpu.stream();
15550 let mut b = __s_b.launch_builder(&f);
15551 b.arg(rp4)
15552 .arg(aq)
15553 .arg(ad)
15554 .arg(&mut y)
15555 .arg(&inf)
15556 .arg(&outf)
15557 .arg(&mi);
15558 unsafe {
15559 b.launch(cfg)?;
15560 }
15561 Ok(y)
15562 }
15563
15564 pub fn scale_inplace(
15566 &self,
15567 y: &mut CudaSlice<f32>,
15568 s: f32,
15569 n: usize,
15570 ) -> Result<(), Box<dyn std::error::Error>> {
15571 let f = self.func("scale_f32");
15572 let cfg = LaunchConfig::for_num_elems(n as u32);
15573 let (sf, ni) = (s, n as i32);
15574 let __s_b = self.gpu.stream();
15575 let mut b = __s_b.launch_builder(&f);
15576 b.arg(y).arg(&sf).arg(&ni);
15577 unsafe {
15578 b.launch(cfg)?;
15579 }
15580 Ok(())
15581 }
15582
15583 pub fn bf16_to_f32(
15588 &self,
15589 data: &cudarc::driver::CudaView<'_, u8>,
15590 n: usize,
15591 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15592 let mut out = self.alloc_uninit::<f32>(n)?;
15593 let f = self.func("bf16_to_f32");
15594 let cfg = LaunchConfig::for_num_elems(n as u32);
15595 let ni = n as i32;
15596 let __s_b = self.gpu.stream();
15597 let mut b = __s_b.launch_builder(&f);
15598 b.arg(data).arg(&mut out).arg(&ni);
15599 unsafe {
15600 b.launch(cfg)?;
15601 }
15602 Ok(out)
15603 }
15604
15605 fn linear_bf16_chunked(
15612 &self,
15613 x: &CudaSlice<f32>,
15614 data: &CudaSlice<u8>,
15615 m: usize,
15616 in_f: usize,
15617 out_f: usize,
15618 exact: bool,
15619 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15620 const CHUNK_BYTES: usize = 256 << 20;
15621 let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
15622 if chunk_rows >= out_f {
15623 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
15624 return if exact {
15625 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
15626 } else {
15627 self.linear(x, &wf32, m, in_f, out_f)
15628 };
15629 }
15630 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15631 let mut r0 = 0usize;
15632 while r0 < out_f {
15633 let rows = chunk_rows.min(out_f - r0);
15634 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
15635 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
15636 let yc = if exact {
15637 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
15638 } else {
15639 self.linear(x, &wf32, m, in_f, rows)?
15640 };
15641 for mi in 0..m {
15643 let src = yc.slice(mi * rows..(mi + 1) * rows);
15644 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
15645 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
15646 }
15647 r0 += rows;
15648 }
15649 Ok(y)
15650 }
15651
15652 pub fn linear_decode_exact(
15659 &self,
15660 x: &CudaSlice<f32>,
15661 w: &CudaSlice<f32>,
15662 m_tokens: usize,
15663 in_f: usize,
15664 out_f: usize,
15665 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15666 if m_tokens == 1 {
15667 return self.linear(x, w, 1, in_f, out_f);
15668 }
15669 let xv = self.view(x, m_tokens * in_f);
15670 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
15671 for t in 0..m_tokens {
15672 let row = xv.slice(t * in_f..(t + 1) * in_f);
15673 let mut xr = self.alloc_uninit::<f32>(in_f)?;
15674 self.copy_view_into(&mut xr, 0, &row, in_f)?;
15675 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
15676 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
15677 }
15678 Ok(y)
15679 }
15680
15681 pub fn linear(
15682 &self,
15683 x: &CudaSlice<f32>,
15684 w: &CudaSlice<f32>,
15685 m_tokens: usize,
15686 in_f: usize,
15687 out_f: usize,
15688 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15689 use cudarc::cublaslt::{Matmul, MatmulConfig};
15690 let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; let cfg = MatmulConfig {
15692 transa: true,
15693 transb: false,
15694 transc: false,
15695 m: out_f as u64,
15696 n: m_tokens as u64,
15697 k: in_f as u64,
15698 alpha: 1.0,
15699 lda: in_f as i64,
15700 ldb: in_f as i64,
15701 beta: 0.0,
15702 ldc: out_f as i64,
15703 stride_a: None,
15704 stride_b: None,
15705 stride_c: None,
15706 stride_bias: None,
15707 batch_size: None,
15708 };
15709 unsafe {
15710 self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15711 }
15712 Ok(c)
15713 }
15714
15715 pub fn sdpa_naive(
15717 &self,
15718 q: &CudaSlice<f32>,
15719 k: &CudaSlice<f32>,
15720 v: &CudaSlice<f32>,
15721 o: &mut CudaSlice<f32>,
15722 head_dim: usize,
15723 n_head: usize,
15724 n_head_kv: usize,
15725 t: usize,
15726 t_kv: usize,
15727 scale: f32,
15728 causal: bool,
15729 ) -> Result<(), Box<dyn std::error::Error>> {
15730 let f = self.func("sdpa_naive_f32");
15731 let cfg = LaunchConfig {
15732 grid_dim: (n_head as u32, t as u32, 1),
15733 block_dim: (128, 1, 1),
15734 shared_mem_bytes: (t_kv * 4) as u32,
15735 };
15736 let (hd, nh, nhkv, ti, tkvi, cz) = (
15737 head_dim as i32,
15738 n_head as i32,
15739 n_head_kv as i32,
15740 t as i32,
15741 t_kv as i32,
15742 causal as i32,
15743 );
15744 let __s_b = self.gpu.stream();
15745 let mut b = __s_b.launch_builder(&f);
15746 b.arg(q)
15747 .arg(k)
15748 .arg(v)
15749 .arg(o)
15750 .arg(&hd)
15751 .arg(&nh)
15752 .arg(&nhkv)
15753 .arg(&ti)
15754 .arg(&tkvi)
15755 .arg(&scale)
15756 .arg(&cz);
15757 unsafe {
15758 b.launch(cfg)?;
15759 }
15760 Ok(())
15761 }
15762
15763 #[allow(clippy::too_many_arguments)]
15768 pub fn sdpa_naive_island(
15769 &self,
15770 q: &CudaSlice<f32>,
15771 k: &CudaSlice<f32>,
15772 v: &CudaSlice<f32>,
15773 o: &mut CudaSlice<f32>,
15774 span_id: &CudaSlice<i32>,
15775 head_dim: usize,
15776 n_head: usize,
15777 n_head_kv: usize,
15778 t: usize,
15779 t_kv: usize,
15780 scale: f32,
15781 window: usize,
15782 ) -> Result<(), Box<dyn std::error::Error>> {
15783 let f = self.func("sdpa_naive_island_f32");
15784 let cfg = LaunchConfig {
15785 grid_dim: (n_head as u32, t as u32, 1),
15786 block_dim: (128, 1, 1),
15787 shared_mem_bytes: (t_kv * 4) as u32,
15788 };
15789 let (hd, nh, nhkv, ti, tkvi, wi) = (
15790 head_dim as i32,
15791 n_head as i32,
15792 n_head_kv as i32,
15793 t as i32,
15794 t_kv as i32,
15795 window as i32,
15796 );
15797 let __s_b = self.gpu.stream();
15798 let mut b = __s_b.launch_builder(&f);
15799 b.arg(q)
15800 .arg(k)
15801 .arg(v)
15802 .arg(o)
15803 .arg(span_id)
15804 .arg(&hd)
15805 .arg(&nh)
15806 .arg(&nhkv)
15807 .arg(&ti)
15808 .arg(&tkvi)
15809 .arg(&scale)
15810 .arg(&wi);
15811 unsafe {
15812 b.launch(cfg)?;
15813 }
15814 Ok(())
15815 }
15816
15817 #[allow(clippy::too_many_arguments)]
15819 pub fn sdpa_naive_w(
15820 &self,
15821 q: &CudaSlice<f32>,
15822 k: &CudaSlice<f32>,
15823 v: &CudaSlice<f32>,
15824 o: &mut CudaSlice<f32>,
15825 head_dim: usize,
15826 n_head: usize,
15827 n_head_kv: usize,
15828 t: usize,
15829 t_kv: usize,
15830 scale: f32,
15831 causal: bool,
15832 window: usize,
15833 ) -> Result<(), Box<dyn std::error::Error>> {
15834 let f = self.func("sdpa_naive_w_f32");
15835 let cfg = LaunchConfig {
15836 grid_dim: (n_head as u32, t as u32, 1),
15837 block_dim: (128, 1, 1),
15838 shared_mem_bytes: (t_kv * 4) as u32,
15839 };
15840 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15841 head_dim as i32,
15842 n_head as i32,
15843 n_head_kv as i32,
15844 t as i32,
15845 t_kv as i32,
15846 causal 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(&hd)
15856 .arg(&nh)
15857 .arg(&nhkv)
15858 .arg(&ti)
15859 .arg(&tkvi)
15860 .arg(&scale)
15861 .arg(&cz)
15862 .arg(&wi);
15863 unsafe {
15864 b.launch(cfg)?;
15865 }
15866 Ok(())
15867 }
15868
15869 #[allow(clippy::too_many_arguments)]
15879 pub fn sdpa_naive_w_lo(
15880 &self,
15881 q: &CudaSlice<f32>,
15882 k: &CudaSlice<f32>,
15883 v: &CudaSlice<f32>,
15884 o: &mut CudaSlice<f32>,
15885 head_dim: usize,
15886 n_head: usize,
15887 n_head_kv: usize,
15888 t: usize,
15889 t_kv: usize,
15890 scale: f32,
15891 causal: bool,
15892 window: usize,
15893 ) -> Result<(), Box<dyn std::error::Error>> {
15894 let kv_lo = if window > 0 {
15895 (t_kv - t + 1).saturating_sub(window)
15896 } else {
15897 0
15898 };
15899 let smem = (t_kv - kv_lo) * 4;
15900 if smem > 48 * 1024 {
15901 return Err(format!(
15902 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
15903 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
15904 a window this wide needs the multi-pass long-ctx kernel"
15905 )
15906 .into());
15907 }
15908 let f = self.func("sdpa_naive_w_lo_f32");
15909 let cfg = LaunchConfig {
15910 grid_dim: (n_head as u32, t as u32, 1),
15911 block_dim: (128, 1, 1),
15912 shared_mem_bytes: smem as u32,
15913 };
15914 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
15915 head_dim as i32,
15916 n_head as i32,
15917 n_head_kv as i32,
15918 t as i32,
15919 t_kv as i32,
15920 causal as i32,
15921 window as i32,
15922 kv_lo as i32,
15923 );
15924 let __s_b = self.gpu.stream();
15925 let mut b = __s_b.launch_builder(&f);
15926 b.arg(q)
15927 .arg(k)
15928 .arg(v)
15929 .arg(o)
15930 .arg(&hd)
15931 .arg(&nh)
15932 .arg(&nhkv)
15933 .arg(&ti)
15934 .arg(&tkvi)
15935 .arg(&scale)
15936 .arg(&cz)
15937 .arg(&wi)
15938 .arg(&lo);
15939 unsafe {
15940 b.launch(cfg)?;
15941 }
15942 Ok(())
15943 }
15944
15945 pub fn sdpa_naive_view(
15947 &self,
15948 q: &CudaSlice<f32>,
15949 k: &cudarc::driver::CudaView<f32>,
15950 v: &cudarc::driver::CudaView<f32>,
15951 o: &mut CudaSlice<f32>,
15952 head_dim: usize,
15953 n_head: usize,
15954 n_head_kv: usize,
15955 t: usize,
15956 t_kv: usize,
15957 scale: f32,
15958 causal: bool,
15959 ) -> Result<(), Box<dyn std::error::Error>> {
15960 let f = self.func("sdpa_naive_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: (t_kv * 4) as u32,
15965 };
15966 let (hd, nh, nhkv, ti, tkvi, cz) = (
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 );
15974 let __s_b = self.gpu.stream();
15975 let mut b = __s_b.launch_builder(&f);
15976 b.arg(q)
15977 .arg(k)
15978 .arg(v)
15979 .arg(o)
15980 .arg(&hd)
15981 .arg(&nh)
15982 .arg(&nhkv)
15983 .arg(&ti)
15984 .arg(&tkvi)
15985 .arg(&scale)
15986 .arg(&cz);
15987 unsafe {
15988 b.launch(cfg)?;
15989 }
15990 Ok(())
15991 }
15992
15993 #[allow(clippy::too_many_arguments)]
16001 pub fn fa_dequant_kv_view_f32(
16002 &self,
16003 k: &cudarc::driver::CudaView<u8>,
16004 v: &cudarc::driver::CudaView<u8>,
16005 kf: &mut CudaSlice<f32>,
16006 vf: &mut CudaSlice<f32>,
16007 kv_dim_k: usize,
16008 kv_dim_v: usize,
16009 t_kv: usize,
16010 k_tok_bytes: usize,
16011 v_tok_bytes: usize,
16012 g: bool,
16013 ) -> Result<(), Box<dyn std::error::Error>> {
16014 let f = if g {
16015 self.func_g("fa_dequant_kv_ws_f32")
16016 } else {
16017 self.func("fa_dequant_kv_ws_f32")
16018 };
16019 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
16020 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16021 let cfg = LaunchConfig {
16022 grid_dim: (nblk.max(1), 1, 1),
16023 block_dim: (256, 1, 1),
16024 shared_mem_bytes: 0,
16025 };
16026 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
16027 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16028 let __s_b = self.gpu.stream();
16029 let mut b = __s_b.launch_builder(&f);
16030 b.arg(k)
16031 .arg(v)
16032 .arg(&mut *kf)
16033 .arg(&mut *vf)
16034 .arg(&kdk)
16035 .arg(&kdv)
16036 .arg(&tkvi)
16037 .arg(&ktb)
16038 .arg(&vtb);
16039 unsafe {
16040 b.launch(cfg)?;
16041 }
16042 Ok(())
16043 }
16044
16045 #[allow(clippy::too_many_arguments)]
16046 pub fn sdpa_naive_quantized_view(
16047 &self,
16048 q: &CudaSlice<f32>,
16049 k: &cudarc::driver::CudaView<u8>,
16050 v: &cudarc::driver::CudaView<u8>,
16051 o: &mut CudaSlice<f32>,
16052 head_dim: usize,
16053 n_head: usize,
16054 n_head_kv: usize,
16055 t: usize,
16056 t_kv: usize,
16057 scale: f32,
16058 causal: bool,
16059 k_tok_bytes: usize,
16060 v_tok_bytes: usize,
16061 ) -> Result<(), Box<dyn std::error::Error>> {
16062 let kv_dim = n_head_kv * head_dim;
16063 let mut kf = self.uninit(t_kv * kv_dim)?;
16064 let mut vf = self.uninit(t_kv * kv_dim)?;
16065 let f = self.func("fa_dequant_kv_ws_f32");
16066 let total = (2 * t_kv * kv_dim) as u64;
16067 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16068 let cfg = LaunchConfig {
16069 grid_dim: (nblk.max(1), 1, 1),
16070 block_dim: (256, 1, 1),
16071 shared_mem_bytes: 0,
16072 };
16073 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
16074 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
16075 let __s_b = self.gpu.stream();
16076 let mut b = __s_b.launch_builder(&f);
16077 b.arg(k)
16078 .arg(v)
16079 .arg(&mut kf)
16080 .arg(&mut vf)
16081 .arg(&kv_dim_i)
16082 .arg(&kv_dim_i)
16083 .arg(&t_kv_i)
16084 .arg(&k_tok_bytes_i)
16085 .arg(&v_tok_bytes_i);
16086 unsafe { b.launch(cfg)? };
16087 self.sdpa_naive(
16088 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16089 )
16090 }
16091
16092 #[allow(clippy::too_many_arguments)]
16104 pub fn sdpa_naive_w_quantized_view(
16105 &self,
16106 q: &CudaSlice<f32>,
16107 k: &cudarc::driver::CudaView<u8>,
16108 v: &cudarc::driver::CudaView<u8>,
16109 o: &mut CudaSlice<f32>,
16110 head_dim: usize,
16111 n_head: usize,
16112 n_head_kv: usize,
16113 t: usize,
16114 t_kv: usize,
16115 scale: f32,
16116 causal: bool,
16117 window: usize,
16118 k_tok_bytes: usize,
16119 v_tok_bytes: usize,
16120 ) -> Result<(), Box<dyn std::error::Error>> {
16121 let kv_dim = n_head_kv * head_dim;
16122 let mut kf = self.uninit(t_kv * kv_dim)?;
16123 let mut vf = self.uninit(t_kv * kv_dim)?;
16124 let f = self.func("fa_dequant_kv_ws_f32");
16125 let total = (2 * t_kv * kv_dim) as u64;
16126 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
16127 let cfg = LaunchConfig {
16128 grid_dim: (nblk.max(1), 1, 1),
16129 block_dim: (256, 1, 1),
16130 shared_mem_bytes: 0,
16131 };
16132 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
16133 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
16134 let __s_b = self.gpu.stream();
16135 let mut b = __s_b.launch_builder(&f);
16136 b.arg(k)
16137 .arg(v)
16138 .arg(&mut kf)
16139 .arg(&mut vf)
16140 .arg(&kv_dim_i)
16141 .arg(&kv_dim_i)
16142 .arg(&t_kv_i)
16143 .arg(&k_tok_bytes_i)
16144 .arg(&v_tok_bytes_i);
16145 unsafe { b.launch(cfg)? };
16146 self.sdpa_naive_w(
16147 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
16148 )
16149 }
16150
16151 pub fn fa_prefill(
16155 &self,
16156 q: &CudaSlice<f32>,
16157 k: &CudaSlice<f32>,
16158 v: &CudaSlice<f32>,
16159 o: &mut CudaSlice<f32>,
16160 head_dim: usize,
16161 n_head: usize,
16162 n_head_kv: usize,
16163 t: usize,
16164 t_kv: usize,
16165 scale: f32,
16166 causal: bool,
16167 ) -> Result<(), Box<dyn std::error::Error>> {
16168 if portable_mma_gated() {
16169 return self.sdpa_naive(
16170 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16171 );
16172 }
16173 let fa3_on = head_dim == 256
16181 && causal
16182 && t == t_kv
16183 && match std::env::var("MEMRA_FA3").as_deref() {
16184 Ok("0") => false,
16185 Ok("1") => true,
16186 _ => cfg!(memra_hopper_mma),
16187 };
16188 if fa3_on {
16189 let n = t * n_head * head_dim;
16190 let nkv = t * n_head_kv * head_dim;
16191 let mut q16 = self.alloc_u8_uninit(n * 2)?;
16192 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
16193 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
16194 self.f32_to_bf16_into(q, &mut q16, n)?;
16195 self.f32_to_bf16_into(k, &mut k16, nkv)?;
16196 self.f32_to_bf16_into(v, &mut v16, nkv)?;
16197 let rc = {
16198 use cudarc::driver::{DevicePtr, DevicePtrMut};
16199 let stream = self.gpu.stream();
16200 let (qp, _g1) = q16.device_ptr(&stream);
16201 let (kp, _g2) = k16.device_ptr(&stream);
16202 let (vp, _g3) = v16.device_ptr(&stream);
16203 let (op, _g4) = o.device_ptr_mut(&stream);
16204 unsafe {
16205 memra_fa3_prefill(
16206 qp as *const core::ffi::c_void,
16207 kp as *const core::ffi::c_void,
16208 vp as *const core::ffi::c_void,
16209 op as *mut f32,
16210 t as i32,
16211 n_head as i32,
16212 n_head_kv as i32,
16213 head_dim as i32,
16214 scale,
16215 stream.cu_stream() as *mut core::ffi::c_void,
16216 )
16217 }
16218 };
16219 if rc != 0 {
16220 return Err(format!("memra_fa3_prefill rc={rc}").into());
16221 }
16222 return Ok(());
16223 }
16224 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16229 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
16230 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
16231 const BLOCK_Q: usize = 64;
16232 const BKX: usize = 32;
16233 let f = self.func("fa_prefill_bf16_p1");
16234 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
16235 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
16236 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16237 f.set_attribute(
16238 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16239 shmem as i32,
16240 )?;
16241 let cfg = LaunchConfig {
16242 grid_dim: (
16243 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16244 n_head as u32,
16245 1,
16246 ),
16247 block_dim: (32, 4, 1),
16248 shared_mem_bytes: shmem,
16249 };
16250 let (hd, nh, nhkv, ti, tkvi, cz) = (
16251 head_dim as i32,
16252 n_head as i32,
16253 n_head_kv as i32,
16254 t as i32,
16255 t_kv as i32,
16256 causal as i32,
16257 );
16258 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16259 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16260 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16261 let __s_b = self.gpu.stream();
16262 let mut b = __s_b.launch_builder(&f);
16263 b.arg(&qb)
16264 .arg(&kb)
16265 .arg(&vb)
16266 .arg(o)
16267 .arg(&hd)
16268 .arg(&nh)
16269 .arg(&nhkv)
16270 .arg(&ti)
16271 .arg(&tkvi)
16272 .arg(&scale)
16273 .arg(&cz);
16274 unsafe {
16275 b.launch(cfg)?;
16276 }
16277 return Ok(());
16278 }
16279 const BK: usize = 32;
16285 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
16288 let (block_q, warps, w2_sfx): (usize, u32, &str) =
16289 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
16290 let hd_sfx = fa_hd_suffix(head_dim)?;
16294 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
16295 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
16300 let (kb16, vb16) = if bf16kv {
16301 let n = t_kv * n_head_kv * head_dim;
16302 let mut kb = self.alloc_u8_uninit(n * 2)?;
16303 let mut vb = self.alloc_u8_uninit(n * 2)?;
16304 let fcv = self.func("f32_to_bf16_bulk");
16305 let ni = n as i64;
16306 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16307 let __s_b = self.gpu.stream();
16308 let mut b = __s_b.launch_builder(&fcv);
16309 b.arg(k).arg(&mut kb).arg(&ni);
16310 unsafe {
16311 b.launch(cfgc)?;
16312 }
16313 let __s_b = self.gpu.stream();
16314 let mut b = __s_b.launch_builder(&fcv);
16315 b.arg(v).arg(&mut vb).arg(&ni);
16316 unsafe {
16317 b.launch(cfgc)?;
16318 }
16319 (Some(kb), Some(vb))
16320 } else {
16321 (None, None)
16322 };
16323 let f = self.func(&if bf16kv {
16324 format!("fa_prefill_bf16kv_pp{hd_sfx}")
16325 } else {
16326 format!(
16327 "fa_prefill_f32{}{}{hd_sfx}",
16328 if floor { "" } else { "_pp" },
16329 if floor { "" } else { w2_sfx }
16330 )
16331 });
16332 let kv_stages = if bf16kv { 2 } else { 1 };
16335 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16336 + 4 * (block_q * BK + 2 * block_q)) as u32;
16337 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16338 f.set_attribute(
16339 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16340 shmem as i32,
16341 )?;
16342 let cfg = LaunchConfig {
16343 grid_dim: (
16344 (t as u32 + block_q as u32 - 1) / block_q as u32,
16345 n_head as u32,
16346 1,
16347 ),
16348 block_dim: (32, warps, 1),
16349 shared_mem_bytes: shmem,
16350 };
16351 let (hd, nh, nhkv, ti, tkvi, cz) = (
16352 head_dim as i32,
16353 n_head as i32,
16354 n_head_kv as i32,
16355 t as i32,
16356 t_kv as i32,
16357 causal as i32,
16358 );
16359 let __s_b = self.gpu.stream();
16360 let mut b = __s_b.launch_builder(&f);
16361 b.arg(q);
16362 match (&kb16, &vb16) {
16363 (Some(kb), Some(vb)) => {
16364 b.arg(kb).arg(vb);
16365 }
16366 _ => {
16367 b.arg(k).arg(v);
16368 }
16369 }
16370 b.arg(o)
16371 .arg(&hd)
16372 .arg(&nh)
16373 .arg(&nhkv)
16374 .arg(&ti)
16375 .arg(&tkvi)
16376 .arg(&scale)
16377 .arg(&cz);
16378 unsafe {
16379 b.launch(cfg)?;
16380 }
16381 Ok(())
16382 }
16383
16384 #[allow(clippy::too_many_arguments)]
16388 pub fn fa_prefill_w(
16389 &self,
16390 q: &CudaSlice<f32>,
16391 k: &CudaSlice<f32>,
16392 v: &CudaSlice<f32>,
16393 o: &mut CudaSlice<f32>,
16394 head_dim: usize,
16395 n_head: usize,
16396 n_head_kv: usize,
16397 t: usize,
16398 t_kv: usize,
16399 scale: f32,
16400 causal: bool,
16401 window: usize,
16402 ) -> Result<(), Box<dyn std::error::Error>> {
16403 if portable_mma_gated() {
16406 return self.sdpa_naive_w(
16407 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
16408 );
16409 }
16410 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16414 let faw_f32 =
16415 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
16416 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
16417 self.fa_prefill_w_arm(
16418 q,
16419 k,
16420 v,
16421 o,
16422 head_dim,
16423 n_head,
16424 n_head_kv,
16425 t,
16426 t_kv,
16427 scale,
16428 causal,
16429 window,
16430 floor || faw_f32,
16431 floor,
16432 )
16433 }
16434
16435 #[allow(clippy::too_many_arguments)]
16438 pub fn fa_prefill_w_pre(
16439 &self,
16440 qb: &CudaSlice<u8>,
16441 kb: &CudaSlice<u8>,
16442 vb: &CudaSlice<u8>,
16443 o: &mut CudaSlice<f32>,
16444 head_dim: usize,
16445 n_head: usize,
16446 n_head_kv: usize,
16447 t: usize,
16448 t_kv: usize,
16449 scale: f32,
16450 causal: bool,
16451 window: usize,
16452 v_f16: bool,
16453 ) -> Result<(), Box<dyn std::error::Error>> {
16454 const BLOCK_Q: usize = 64;
16455 const BK: usize = 32;
16456 debug_assert_eq!(head_dim, 256);
16457 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16458 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
16459 if hp {
16460 const BLOCK_QH: usize = 32;
16461 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16464 let vh: &CudaSlice<u8> = if v_f16 {
16465 vb
16466 } else {
16467 let n = t_kv * n_head_kv * head_dim;
16468 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
16469 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
16470 }
16471 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
16472 vguard.as_ref().unwrap()
16473 };
16474 let f = self.func("fa_prefill_w_bf16_p1h2");
16475 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16476 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16477 f.set_attribute(
16478 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16479 shmem as i32,
16480 )?;
16481 let cfg = LaunchConfig {
16482 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16483 block_dim: (32, 4, 1),
16484 shared_mem_bytes: shmem,
16485 };
16486 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16487 head_dim as i32,
16488 n_head as i32,
16489 n_head_kv as i32,
16490 t as i32,
16491 t_kv as i32,
16492 causal as i32,
16493 window as i32,
16494 );
16495 let __s_b = self.gpu.stream();
16496 let mut b = __s_b.launch_builder(&f);
16497 b.arg(qb)
16498 .arg(kb)
16499 .arg(vh)
16500 .arg(o)
16501 .arg(&hd)
16502 .arg(&nh)
16503 .arg(&nhkv)
16504 .arg(&ti)
16505 .arg(&tkvi)
16506 .arg(&scale)
16507 .arg(&cz)
16508 .arg(&wi);
16509 unsafe {
16510 b.launch(cfg)?;
16511 }
16512 return Ok(());
16513 }
16514 let f = self.func("fa_prefill_w_bf16_p1");
16515 let shmem =
16516 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16517 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16518 f.set_attribute(
16519 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16520 shmem as i32,
16521 )?;
16522 let cfg = LaunchConfig {
16523 grid_dim: (
16524 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16525 n_head as u32,
16526 1,
16527 ),
16528 block_dim: (32, 4, 1),
16529 shared_mem_bytes: shmem,
16530 };
16531 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16532 head_dim as i32,
16533 n_head as i32,
16534 n_head_kv as i32,
16535 t as i32,
16536 t_kv as i32,
16537 causal as i32,
16538 window as i32,
16539 );
16540 let __s_b = self.gpu.stream();
16541 let mut b = __s_b.launch_builder(&f);
16542 b.arg(qb)
16543 .arg(kb)
16544 .arg(vb)
16545 .arg(o)
16546 .arg(&hd)
16547 .arg(&nh)
16548 .arg(&nhkv)
16549 .arg(&ti)
16550 .arg(&tkvi)
16551 .arg(&scale)
16552 .arg(&cz)
16553 .arg(&wi);
16554 unsafe {
16555 b.launch(cfg)?;
16556 }
16557 Ok(())
16558 }
16559
16560 #[allow(clippy::too_many_arguments)]
16562 pub fn fa_prefill_w_arm(
16563 &self,
16564 q: &CudaSlice<f32>,
16565 k: &CudaSlice<f32>,
16566 v: &CudaSlice<f32>,
16567 o: &mut CudaSlice<f32>,
16568 head_dim: usize,
16569 n_head: usize,
16570 n_head_kv: usize,
16571 t: usize,
16572 t_kv: usize,
16573 scale: f32,
16574 causal: bool,
16575 window: usize,
16576 f32_stage: bool,
16577 floor: bool,
16578 ) -> Result<(), Box<dyn std::error::Error>> {
16579 const BLOCK_Q: usize = 64;
16580 const BK: usize = 32;
16581 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
16582 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16586 let p1 = !floor
16587 && !f32_stage
16588 && *P1_ON.get_or_init(|| {
16589 std::env::var("MEMRA_FAW_P1")
16590 .map(|v| v != "0")
16591 .unwrap_or(true)
16592 });
16593 let hp =
16594 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16595 if hp {
16596 const BLOCK_QH: usize = 32;
16597 let f = self.func("fa_prefill_w_bf16_p1h2");
16598 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16599 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16600 f.set_attribute(
16601 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16602 shmem as i32,
16603 )?;
16604 let cfg = LaunchConfig {
16605 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16606 block_dim: (32, 4, 1),
16607 shared_mem_bytes: shmem,
16608 };
16609 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16610 head_dim as i32,
16611 n_head as i32,
16612 n_head_kv as i32,
16613 t as i32,
16614 t_kv as i32,
16615 causal as i32,
16616 window as i32,
16617 );
16618 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16619 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16620 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
16621 let __s_b = self.gpu.stream();
16622 let mut b = __s_b.launch_builder(&f);
16623 b.arg(&qb)
16624 .arg(&kb)
16625 .arg(&vh)
16626 .arg(o)
16627 .arg(&hd)
16628 .arg(&nh)
16629 .arg(&nhkv)
16630 .arg(&ti)
16631 .arg(&tkvi)
16632 .arg(&scale)
16633 .arg(&cz)
16634 .arg(&wi);
16635 unsafe {
16636 b.launch(cfg)?;
16637 }
16638 return Ok(());
16639 }
16640 if p1 {
16641 let f = self.func("fa_prefill_w_bf16_p1");
16642 let shmem =
16643 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16644 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16645 f.set_attribute(
16646 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16647 shmem as i32,
16648 )?;
16649 let cfg = LaunchConfig {
16650 grid_dim: (
16651 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16652 n_head as u32,
16653 1,
16654 ),
16655 block_dim: (32, 4, 1),
16656 shared_mem_bytes: shmem,
16657 };
16658 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16659 head_dim as i32,
16660 n_head as i32,
16661 n_head_kv as i32,
16662 t as i32,
16663 t_kv as i32,
16664 causal as i32,
16665 window as i32,
16666 );
16667 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16668 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16669 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16670 let __s_b = self.gpu.stream();
16671 let mut b = __s_b.launch_builder(&f);
16672 b.arg(&qb)
16673 .arg(&kb)
16674 .arg(&vb)
16675 .arg(o)
16676 .arg(&hd)
16677 .arg(&nh)
16678 .arg(&nhkv)
16679 .arg(&ti)
16680 .arg(&tkvi)
16681 .arg(&scale)
16682 .arg(&cz)
16683 .arg(&wi);
16684 unsafe {
16685 b.launch(cfg)?;
16686 }
16687 return Ok(());
16688 }
16689 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16692 let g4 = !floor
16693 && !f32_stage
16694 && n_head_kv == 1
16695 && n_head % 4 == 0
16696 && *G4_ON.get_or_init(|| {
16697 std::env::var("MEMRA_FAW_G4")
16698 .map(|v| v != "0")
16699 .unwrap_or(true)
16700 });
16701 if g4 {
16702 const SP_M: usize = 16;
16703 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16706 let o2 = *O2_ON.get_or_init(|| {
16707 std::env::var("MEMRA_FAW_O2")
16708 .map(|v| v != "0")
16709 .unwrap_or(true)
16710 });
16711 let f = self.func(if o2 {
16712 "fa_prefill_w_bf16_g4o2"
16713 } else {
16714 "fa_prefill_w_bf16_g4"
16715 });
16716 let shmem = if o2 {
16717 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
16718 } else {
16719 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
16720 as u32
16721 };
16722 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16723 f.set_attribute(
16724 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16725 shmem as i32,
16726 )?;
16727 let cfg = LaunchConfig {
16728 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
16729 block_dim: (32, 4, 1),
16730 shared_mem_bytes: shmem,
16731 };
16732 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16733 head_dim as i32,
16734 n_head as i32,
16735 n_head_kv as i32,
16736 t as i32,
16737 t_kv as i32,
16738 causal as i32,
16739 window as i32,
16740 );
16741 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16742 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16743 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16744 let __s_b = self.gpu.stream();
16745 let mut b = __s_b.launch_builder(&f);
16746 b.arg(&qb)
16747 .arg(&kb)
16748 .arg(&vb)
16749 .arg(o)
16750 .arg(&hd)
16751 .arg(&nh)
16752 .arg(&nhkv)
16753 .arg(&ti)
16754 .arg(&tkvi)
16755 .arg(&scale)
16756 .arg(&cz)
16757 .arg(&wi);
16758 unsafe {
16759 b.launch(cfg)?;
16760 }
16761 return Ok(());
16762 }
16763 let f = self.func(if floor {
16764 "fa_prefill_w_f32"
16765 } else if f32_stage {
16766 "fa_prefill_w_f32_pp"
16767 } else {
16768 "fa_prefill_w_bf16_pp"
16769 });
16770 let shmem =
16771 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16772 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16773 f.set_attribute(
16774 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16775 shmem as i32,
16776 )?;
16777 let cfg = LaunchConfig {
16778 grid_dim: (
16779 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16780 n_head as u32,
16781 1,
16782 ),
16783 block_dim: (32, 4, 1),
16784 shared_mem_bytes: shmem,
16785 };
16786 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16787 head_dim as i32,
16788 n_head as i32,
16789 n_head_kv as i32,
16790 t as i32,
16791 t_kv as i32,
16792 causal as i32,
16793 window as i32,
16794 );
16795 if f32_stage {
16796 let __s_b = self.gpu.stream();
16797 let mut b = __s_b.launch_builder(&f);
16798 b.arg(q)
16799 .arg(k)
16800 .arg(v)
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 } else {
16814 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16815 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16816 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16817 let __s_b = self.gpu.stream();
16818 let mut b = __s_b.launch_builder(&f);
16819 b.arg(&qb)
16820 .arg(&kb)
16821 .arg(&vb)
16822 .arg(o)
16823 .arg(&hd)
16824 .arg(&nh)
16825 .arg(&nhkv)
16826 .arg(&ti)
16827 .arg(&tkvi)
16828 .arg(&scale)
16829 .arg(&cz)
16830 .arg(&wi);
16831 unsafe {
16832 b.launch(cfg)?;
16833 }
16834 }
16835 Ok(())
16836 }
16837
16838 #[allow(clippy::too_many_arguments)]
16842 pub fn fa_prefill_hd512(
16843 &self,
16844 q: &CudaSlice<f32>,
16845 k: &CudaSlice<f32>,
16846 v: &CudaSlice<f32>,
16847 o: &mut CudaSlice<f32>,
16848 head_dim: usize,
16849 n_head: usize,
16850 n_head_kv: usize,
16851 t: usize,
16852 t_kv: usize,
16853 scale: f32,
16854 causal: bool,
16855 ) -> Result<(), Box<dyn std::error::Error>> {
16856 if portable_mma_gated() {
16858 return self.sdpa_naive(
16859 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16860 );
16861 }
16862 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16868 let f32_stage =
16869 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16870 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16874 let sp = !f32_stage
16875 && *SP_ON.get_or_init(|| {
16876 std::env::var("MEMRA_FA512_SP")
16877 .map(|v| v != "0")
16878 .unwrap_or(true)
16879 });
16880 self.fa_prefill_hd512_arm(
16881 q,
16882 k,
16883 v,
16884 o,
16885 head_dim,
16886 n_head,
16887 n_head_kv,
16888 t,
16889 t_kv,
16890 scale,
16891 causal,
16892 f32_stage,
16893 sp,
16894 sp && fa_f16pv_on(),
16895 )
16896 }
16897
16898 #[allow(clippy::too_many_arguments)]
16900 pub fn fa_prefill_hd512_pre(
16901 &self,
16902 qb: &CudaSlice<u8>,
16903 kb: &CudaSlice<u8>,
16904 vb: &CudaSlice<u8>,
16905 o: &mut CudaSlice<f32>,
16906 head_dim: usize,
16907 n_head: usize,
16908 n_head_kv: usize,
16909 t: usize,
16910 t_kv: usize,
16911 scale: f32,
16912 causal: bool,
16913 v_f16: bool,
16914 ) -> Result<(), Box<dyn std::error::Error>> {
16915 debug_assert_eq!(head_dim, 512);
16916 const SP_M: usize = 16;
16917 const BKS: usize = 32;
16918 let f16pv = fa_f16pv_on();
16922 let nw = if f16pv { fa512_wide_warps() } else { 2 };
16923 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16924 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16925 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16926 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16927 let n = t_kv * n_head_kv * head_dim;
16929 let need = n * 2;
16930 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16931 *vguard = Some(self.alloc_uninit::<u8>(need)?);
16932 }
16933 let dst = vguard.as_mut().unwrap();
16934 self.bf16_to_f16_into(vb, n, dst)?;
16935 vguard.as_ref().unwrap()
16936 } else {
16937 vb
16938 };
16939 let f = self.func(if hp {
16940 "fa_prefill_bf16_hd512_sp16h2"
16941 } else {
16942 match (f16pv, nw) {
16943 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16944 (true, _) => "fa_prefill_bf16_hd512_sp16",
16945 _ => "fa_prefill_bf16_hd512_sp",
16946 }
16947 });
16948 let (nwarp, npart) = if hp {
16949 (4usize, 4usize)
16950 } else if nw > 2 {
16951 (nw, nw)
16952 } else {
16953 (2, 1)
16954 };
16955 let shmem = if hp {
16957 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
16958 as u32
16959 } else {
16960 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16961 + 4 * (npart * SP_M * BKS + SP_M)) as u32
16962 };
16963 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16964 f.set_attribute(
16965 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16966 shmem as i32,
16967 )?;
16968 let grid_y = if hp {
16969 (n_head / 2) as u32
16970 } else {
16971 n_head as u32
16972 };
16973 let cfg = LaunchConfig {
16974 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16975 block_dim: (32, nwarp as u32, 1),
16976 shared_mem_bytes: shmem,
16977 };
16978 let (hd, nh, nhkv, ti, tkvi, cz) = (
16979 head_dim as i32,
16980 n_head as i32,
16981 n_head_kv as i32,
16982 t as i32,
16983 t_kv as i32,
16984 causal as i32,
16985 );
16986 let __s_b = self.gpu.stream();
16987 let mut b = __s_b.launch_builder(&f);
16988 b.arg(qb)
16989 .arg(kb)
16990 .arg(vref)
16991 .arg(o)
16992 .arg(&hd)
16993 .arg(&nh)
16994 .arg(&nhkv)
16995 .arg(&ti)
16996 .arg(&tkvi)
16997 .arg(&scale)
16998 .arg(&cz);
16999 unsafe {
17000 b.launch(cfg)?;
17001 }
17002 Ok(())
17003 }
17004
17005 #[allow(clippy::too_many_arguments)]
17008 pub fn fa_prefill_hd512_arm(
17009 &self,
17010 q: &CudaSlice<f32>,
17011 k: &CudaSlice<f32>,
17012 v: &CudaSlice<f32>,
17013 o: &mut CudaSlice<f32>,
17014 head_dim: usize,
17015 n_head: usize,
17016 n_head_kv: usize,
17017 t: usize,
17018 t_kv: usize,
17019 scale: f32,
17020 causal: bool,
17021 f32_stage: bool,
17022 sp: bool,
17023 f16pv: bool,
17024 ) -> Result<(), Box<dyn std::error::Error>> {
17025 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
17026 if sp && !f32_stage {
17027 const SP_M: usize = 16;
17031 const BKS: usize = 32;
17032 let nw = if f16pv { fa512_wide_warps() } else { 2 };
17033 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
17034 let f = self.func(if hp {
17035 "fa_prefill_bf16_hd512_sp16h2"
17036 } else {
17037 match (f16pv, nw) {
17038 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
17039 (true, _) => "fa_prefill_bf16_hd512_sp16",
17040 _ => "fa_prefill_bf16_hd512_sp",
17041 }
17042 });
17043 let (nwarp, npart) = if hp {
17044 (4usize, 4usize)
17045 } else if nw > 2 {
17046 (nw, nw)
17047 } else {
17048 (2, 1)
17049 };
17050 let shmem = if hp {
17051 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
17052 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
17053 } else {
17054 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
17055 + 4 * (npart * SP_M * BKS + SP_M)) as u32
17056 };
17057 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17058 f.set_attribute(
17059 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17060 shmem as i32,
17061 )?;
17062 let grid_y = if hp {
17063 (n_head / 2) as u32
17064 } else {
17065 n_head as u32
17066 };
17067 let cfg = LaunchConfig {
17068 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
17069 block_dim: (32, nwarp as u32, 1),
17070 shared_mem_bytes: shmem,
17071 };
17072 let (hd, nh, nhkv, ti, tkvi, cz) = (
17073 head_dim as i32,
17074 n_head as i32,
17075 n_head_kv as i32,
17076 t as i32,
17077 t_kv as i32,
17078 causal as i32,
17079 );
17080 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
17081 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
17082 let vb = if f16pv {
17083 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
17084 } else {
17085 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
17086 };
17087 let __s_b = self.gpu.stream();
17088 let mut b = __s_b.launch_builder(&f);
17089 b.arg(&qb)
17090 .arg(&kb)
17091 .arg(&vb)
17092 .arg(o)
17093 .arg(&hd)
17094 .arg(&nh)
17095 .arg(&nhkv)
17096 .arg(&ti)
17097 .arg(&tkvi)
17098 .arg(&scale)
17099 .arg(&cz);
17100 unsafe {
17101 b.launch(cfg)?;
17102 }
17103 return Ok(());
17104 }
17105 const BLOCK_Q: usize = 32;
17106 const BK: usize = 32;
17107 const HALF: usize = 256;
17108 let f = self.func(if f32_stage {
17109 "fa_prefill_f32_hd512"
17110 } else {
17111 "fa_prefill_bf16_hd512"
17112 });
17113 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
17115 + 4 * BLOCK_Q) as u32;
17116 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17117 f.set_attribute(
17118 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17119 shmem as i32,
17120 )?;
17121 let cfg = LaunchConfig {
17122 grid_dim: (
17123 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17124 n_head as u32,
17125 2,
17126 ),
17127 block_dim: (32, 2, 1),
17128 shared_mem_bytes: shmem,
17129 };
17130 let (hd, nh, nhkv, ti, tkvi, cz) = (
17131 head_dim as i32,
17132 n_head as i32,
17133 n_head_kv as i32,
17134 t as i32,
17135 t_kv as i32,
17136 causal as i32,
17137 );
17138 if f32_stage {
17139 let __s_b = self.gpu.stream();
17140 let mut b = __s_b.launch_builder(&f);
17141 b.arg(q)
17142 .arg(k)
17143 .arg(v)
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 } else {
17156 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
17157 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
17158 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
17159 let __s_b = self.gpu.stream();
17160 let mut b = __s_b.launch_builder(&f);
17161 b.arg(&qb)
17162 .arg(&kb)
17163 .arg(&vb)
17164 .arg(o)
17165 .arg(&hd)
17166 .arg(&nh)
17167 .arg(&nhkv)
17168 .arg(&ti)
17169 .arg(&tkvi)
17170 .arg(&scale)
17171 .arg(&cz);
17172 unsafe {
17173 b.launch(cfg)?;
17174 }
17175 }
17176 Ok(())
17177 }
17178
17179 #[allow(clippy::too_many_arguments)]
17183 pub fn rope_neox2_bf16e(
17184 &self,
17185 q: &mut CudaSlice<f32>,
17186 k: &mut CudaSlice<f32>,
17187 qb: &mut CudaSlice<u8>,
17188 kb: &mut CudaSlice<u8>,
17189 pos: &CudaSlice<i32>,
17190 head_dim: usize,
17191 n_dims: usize,
17192 nh_q: usize,
17193 nh_k: usize,
17194 n_tokens: usize,
17195 base: f32,
17196 freq_scale: f32,
17197 ff: Option<&CudaSlice<f32>>,
17198 ) -> Result<(), Box<dyn std::error::Error>> {
17199 let f = self.func("rope_neox2_bf16e_f32");
17200 let rows = ((nh_q + nh_k) * n_tokens) as u32;
17201 let cfg = LaunchConfig {
17202 grid_dim: (rows, 1, 1),
17203 block_dim: ((head_dim / 2) as u32, 1, 1),
17204 shared_mem_bytes: 0,
17205 };
17206 let theta_scale = base.powf(-2.0 / n_dims as f32);
17207 let (hd, nd, nhq, nhk, nt) = (
17208 head_dim as i32,
17209 n_dims as i32,
17210 nh_q as i32,
17211 nh_k as i32,
17212 n_tokens as i32,
17213 );
17214 let __s_b = self.gpu.stream();
17215 let mut b = __s_b.launch_builder(&f);
17216 match ff {
17217 Some(t) => {
17218 b.arg(&mut *q)
17219 .arg(&mut *k)
17220 .arg(&mut *qb)
17221 .arg(&mut *kb)
17222 .arg(pos)
17223 .arg(&hd)
17224 .arg(&nd)
17225 .arg(&nhq)
17226 .arg(&nhk)
17227 .arg(&nt)
17228 .arg(&theta_scale)
17229 .arg(&freq_scale)
17230 .arg(t);
17231 unsafe {
17232 b.launch(cfg)?;
17233 }
17234 }
17235 None => {
17236 let null: u64 = 0;
17237 b.arg(&mut *q)
17238 .arg(&mut *k)
17239 .arg(&mut *qb)
17240 .arg(&mut *kb)
17241 .arg(pos)
17242 .arg(&hd)
17243 .arg(&nd)
17244 .arg(&nhq)
17245 .arg(&nhk)
17246 .arg(&nt)
17247 .arg(&theta_scale)
17248 .arg(&freq_scale)
17249 .arg(&null);
17250 unsafe {
17251 b.launch(cfg)?;
17252 }
17253 }
17254 }
17255 Ok(())
17256 }
17257
17258 pub fn f32_to_bf16(
17261 &self,
17262 x: &CudaSlice<f32>,
17263 n: usize,
17264 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17265 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
17266 let mut y = self.alloc_uninit::<u8>(n * 2)?;
17267 let f = self.func("f32_to_bf16_flat");
17268 let n_i = n as i64;
17269 let cfg = LaunchConfig {
17270 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
17271 block_dim: (256, 1, 1),
17272 shared_mem_bytes: 0,
17273 };
17274 let __s_b = self.gpu.stream();
17275 let mut b = __s_b.launch_builder(&f);
17276 b.arg(x).arg(&mut y).arg(&n_i);
17277 unsafe {
17278 b.launch(cfg)?;
17279 }
17280 Ok(y)
17281 }
17282
17283 pub fn f32_to_f16(
17284 &self,
17285 x: &CudaSlice<f32>,
17286 n: usize,
17287 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17288 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
17289 let mut y = self.alloc_uninit::<u8>(n * 2)?;
17290 let f = self.func("f32_to_f16_flat");
17291 let n_i = n as i64;
17292 let cfg = LaunchConfig {
17293 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
17294 block_dim: (256, 1, 1),
17295 shared_mem_bytes: 0,
17296 };
17297 let __s_b = self.gpu.stream();
17298 let mut b = __s_b.launch_builder(&f);
17299 b.arg(x).arg(&mut y).arg(&n_i);
17300 unsafe {
17301 b.launch(cfg)?;
17302 }
17303 Ok(y)
17304 }
17305
17306 pub fn bf16_to_f16(
17308 &self,
17309 xb: &CudaSlice<u8>,
17310 n: usize,
17311 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17312 let mut y = self.alloc_uninit::<u8>(n * 2)?;
17313 self.bf16_to_f16_into(xb, n, &mut y)?;
17314 Ok(y)
17315 }
17316
17317 pub fn bf16_to_f16_into(
17319 &self,
17320 xb: &CudaSlice<u8>,
17321 n: usize,
17322 y: &mut CudaSlice<u8>,
17323 ) -> Result<(), Box<dyn std::error::Error>> {
17324 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
17325 assert!(y.len() >= n * 2);
17326 let f = self.func("bf16_to_f16_flat");
17327 let n2 = (n / 2) as i64;
17328 let cfg = LaunchConfig {
17329 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
17330 block_dim: (256, 1, 1),
17331 shared_mem_bytes: 0,
17332 };
17333 let __s_b = self.gpu.stream();
17334 let mut b = __s_b.launch_builder(&f);
17335 b.arg(xb).arg(y).arg(&n2);
17336 unsafe {
17337 b.launch(cfg)?;
17338 }
17339 Ok(())
17340 }
17341
17342 #[allow(clippy::too_many_arguments)]
17347 pub fn fa_prefill_vl8(
17348 &self,
17349 seqs: &[FaSeqVl],
17350 head_dim: usize,
17351 n_head: usize,
17352 n_head_kv: usize,
17353 scale: f32,
17354 ) -> Result<(), Box<dyn std::error::Error>> {
17355 const BK: usize = 32;
17356 let b = seqs.len();
17357 assert!(b >= 1 && b <= 8);
17358 let mut packed = [FaSeqVl::default(); 8];
17359 packed[..b].copy_from_slice(seqs);
17360 let v = FaVl8(packed);
17361 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17362 let ept = (n_head_kv * head_dim) as i32;
17363 {
17364 let f = self.func("fa_mirror_vl");
17365 let max_n = (max_t as i64) * ept as i64;
17366 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
17367 for which in 0..2i32 {
17368 let cfg = LaunchConfig {
17369 grid_dim: (blocks, 1, b as u32),
17370 block_dim: (256, 1, 1),
17371 shared_mem_bytes: 0,
17372 };
17373 let __s_lb = self.gpu.stream();
17374 let mut lb = __s_lb.launch_builder(&f);
17375 lb.arg(&v).arg(&ept).arg(&which);
17376 unsafe {
17377 lb.launch(cfg)?;
17378 }
17379 }
17380 }
17381 let hd_sfx = fa_hd_suffix(head_dim)?;
17382 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
17383 let block_q = 64usize;
17384 let kv_stages = 2usize;
17385 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
17386 + 4 * (block_q * BK + 2 * block_q)) as u32;
17387 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17388 f.set_attribute(
17389 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17390 shmem as i32,
17391 )?;
17392 let cfg = LaunchConfig {
17393 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
17394 block_dim: (32, 4, 1),
17395 shared_mem_bytes: shmem,
17396 };
17397 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17398 let __s_lb = self.gpu.stream();
17399 let mut lb = __s_lb.launch_builder(&f);
17400 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
17401 unsafe {
17402 lb.launch(cfg)?;
17403 }
17404 Ok(())
17405 }
17406
17407 #[allow(clippy::too_many_arguments)]
17411 pub fn attn_pre_vl8(
17412 &self,
17413 seqs: &[AttnPreVl],
17414 wq: &CudaSlice<f32>,
17415 wk: &CudaSlice<f32>,
17416 head_dim: usize,
17417 rope_dims: usize,
17418 n_head: usize,
17419 n_head_kv: usize,
17420 eps: f32,
17421 freq_base: f32,
17422 freq_scale: f32,
17423 kv_dim_k: usize,
17424 kv_dim_v: usize,
17425 k_tok_bytes: usize,
17426 v_tok_bytes: usize,
17427 ) -> Result<(), Box<dyn std::error::Error>> {
17428 let b = seqs.len();
17429 assert!(b >= 1 && b <= 8);
17430 let mut packed = [AttnPreVl::default(); 8];
17431 packed[..b].copy_from_slice(seqs);
17432 let v = AttnPreVl8(packed);
17433 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17434 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17435 {
17436 let f = self.func("q_gate_split_vl");
17437 let n = max_t * (n_head * head_dim) as u32;
17438 let cfg = LaunchConfig {
17439 grid_dim: (n.div_ceil(256), 1, b as u32),
17440 block_dim: (256, 1, 1),
17441 shared_mem_bytes: 0,
17442 };
17443 let __s_lb = self.gpu.stream();
17444 let mut lb = __s_lb.launch_builder(&f);
17445 lb.arg(&v).arg(&hd).arg(&nh);
17446 unsafe {
17447 lb.launch(cfg)?;
17448 }
17449 }
17450 {
17451 let f = self.func("attn_rms_vl");
17452 let cfg = LaunchConfig {
17453 grid_dim: (max_t * n_head as u32, 2, b as u32),
17454 block_dim: (rms_block(), 1, 1),
17455 shared_mem_bytes: 0,
17456 };
17457 let __s_lb = self.gpu.stream();
17458 let mut lb = __s_lb.launch_builder(&f);
17459 lb.arg(&v)
17460 .arg(wq)
17461 .arg(wk)
17462 .arg(&hd)
17463 .arg(&nh)
17464 .arg(&nhkv)
17465 .arg(&eps);
17466 unsafe {
17467 lb.launch(cfg)?;
17468 }
17469 }
17470 {
17471 let f = self.func("attn_rope_vl");
17472 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
17473 let nd = rope_dims as i32;
17474 let cfg = LaunchConfig {
17475 grid_dim: (max_t * n_head as u32, 2, b as u32),
17476 block_dim: ((head_dim / 2) as u32, 1, 1),
17477 shared_mem_bytes: 0,
17478 };
17479 let __s_lb = self.gpu.stream();
17480 let mut lb = __s_lb.launch_builder(&f);
17481 lb.arg(&v)
17482 .arg(&hd)
17483 .arg(&nd)
17484 .arg(&nh)
17485 .arg(&nhkv)
17486 .arg(&theta_scale)
17487 .arg(&freq_scale);
17488 unsafe {
17489 lb.launch(cfg)?;
17490 }
17491 }
17492 {
17493 let f = self.func("append_kv_vl");
17494 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
17495 let cfg = LaunchConfig {
17496 grid_dim: (nblk, max_t, b as u32),
17497 block_dim: (32, 1, 1),
17498 shared_mem_bytes: 0,
17499 };
17500 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17501 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17502 let __s_lb = self.gpu.stream();
17503 let mut lb = __s_lb.launch_builder(&f);
17504 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
17505 unsafe {
17506 lb.launch(cfg)?;
17507 }
17508 }
17509 Ok(())
17510 }
17511
17512 pub fn fa_prefill_view(
17517 &self,
17518 q: &CudaSlice<f32>,
17519 k: &cudarc::driver::CudaView<u8>,
17520 v: &cudarc::driver::CudaView<u8>,
17521 o: &mut CudaSlice<f32>,
17522 head_dim: usize,
17523 n_head: usize,
17524 n_head_kv: usize,
17525 t: usize,
17526 t_kv: usize,
17527 scale: f32,
17528 causal: bool,
17529 k_tok_bytes: usize,
17530 v_tok_bytes: usize,
17531 g: bool,
17532 ) -> Result<(), Box<dyn std::error::Error>> {
17533 if portable_mma_gated() {
17534 return self.sdpa_naive_quantized_view(
17535 q,
17536 k,
17537 v,
17538 o,
17539 head_dim,
17540 n_head,
17541 n_head_kv,
17542 t,
17543 t_kv,
17544 scale,
17545 causal,
17546 k_tok_bytes,
17547 v_tok_bytes,
17548 );
17549 }
17550 const BLOCK_Q: usize = 64;
17551 const BK: usize = 32;
17552 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
17555 let f = if g {
17556 self.func_g(&name)
17557 } else {
17558 self.func(&name)
17559 };
17560 let shmem =
17561 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
17562 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17563 f.set_attribute(
17564 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17565 shmem as i32,
17566 )?;
17567 let cfg = LaunchConfig {
17568 grid_dim: (
17569 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17570 n_head as u32,
17571 1,
17572 ),
17573 block_dim: (32, 4, 1),
17574 shared_mem_bytes: shmem,
17575 };
17576 let (hd, nh, nhkv, ti, tkvi, cz) = (
17577 head_dim as i32,
17578 n_head as i32,
17579 n_head_kv as i32,
17580 t as i32,
17581 t_kv as i32,
17582 causal as i32,
17583 );
17584 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17585 let __s_b = self.gpu.stream();
17586 let mut b = __s_b.launch_builder(&f);
17587 b.arg(q)
17588 .arg(k)
17589 .arg(v)
17590 .arg(o)
17591 .arg(&hd)
17592 .arg(&nh)
17593 .arg(&nhkv)
17594 .arg(&ti)
17595 .arg(&tkvi)
17596 .arg(&scale)
17597 .arg(&cz)
17598 .arg(&ktb)
17599 .arg(&vtb);
17600 unsafe {
17601 b.launch(cfg)?;
17602 }
17603 Ok(())
17604 }
17605
17606 #[allow(clippy::too_many_arguments)]
17616 pub fn fa_prefill_view_ws(
17617 &self,
17618 q: &CudaSlice<f32>,
17619 k: &cudarc::driver::CudaView<u8>,
17620 v: &cudarc::driver::CudaView<u8>,
17621 o: &mut CudaSlice<f32>,
17622 head_dim: usize,
17623 n_head: usize,
17624 n_head_kv: usize,
17625 t: usize,
17626 t_kv: usize,
17627 scale: f32,
17628 causal: bool,
17629 k_tok_bytes: usize,
17630 v_tok_bytes: usize,
17631 g: bool,
17632 ) -> Result<(), Box<dyn std::error::Error>> {
17633 if portable_mma_gated() {
17634 return self.sdpa_naive_quantized_view(
17635 q,
17636 k,
17637 v,
17638 o,
17639 head_dim,
17640 n_head,
17641 n_head_kv,
17642 t,
17643 t_kv,
17644 scale,
17645 causal,
17646 k_tok_bytes,
17647 v_tok_bytes,
17648 );
17649 }
17650 const BLOCK_Q: usize = 64;
17651 const BK: usize = 32;
17652 let kv_dim_k = n_head_kv * head_dim;
17653 let kv_dim_v = n_head_kv * head_dim;
17654 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17656 let mut guard = self.prime_deqw_ws.lock().unwrap();
17658 let need_grow = match guard.as_ref() {
17659 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17660 None => true,
17661 };
17662 if need_grow {
17663 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17664 let (ck, cv) = guard
17665 .as_ref()
17666 .map(|(a, b)| (a.len(), b.len()))
17667 .unwrap_or((0, 0));
17668 *guard = Some((
17669 self.alloc_u8(grow(ck, k_ws_bytes))?,
17670 self.alloc_u8(grow(cv, v_ws_bytes))?,
17671 ));
17672 }
17673 let (kw, vw) = guard.as_mut().unwrap();
17674 {
17676 let f = if g {
17678 self.func_g("fa_dequant_kv_ws_bf16")
17679 } else {
17680 self.func("fa_dequant_kv_ws_bf16")
17681 };
17682 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17683 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17684 let cfg = LaunchConfig {
17685 grid_dim: (nblk.max(1), 1, 1),
17686 block_dim: (256, 1, 1),
17687 shared_mem_bytes: 0,
17688 };
17689 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17690 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17691 let __s_b = self.gpu.stream();
17692 let mut b = __s_b.launch_builder(&f);
17693 b.arg(k)
17694 .arg(v)
17695 .arg(&mut *kw)
17696 .arg(&mut *vw)
17697 .arg(&kdk)
17698 .arg(&kdv)
17699 .arg(&tkvi)
17700 .arg(&ktb)
17701 .arg(&vtb);
17702 unsafe {
17703 b.launch(cfg)?;
17704 }
17705 }
17706 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17714 .map(|v| v != "0")
17715 .unwrap_or(true);
17716 {
17717 let hd_sfx = fa_hd_suffix(head_dim)?;
17718 let f = self.func(&format!(
17719 "fa_prefill_qw{}{hd_sfx}",
17720 if db { "_db" } else { "" }
17721 ));
17722 let shmem = if db {
17723 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17725 } else {
17726 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17727 };
17728 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17729 f.set_attribute(
17730 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17731 shmem as i32,
17732 )?;
17733 let cfg = LaunchConfig {
17734 grid_dim: (
17735 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17736 n_head as u32,
17737 1,
17738 ),
17739 block_dim: (32, 4, 1),
17740 shared_mem_bytes: shmem,
17741 };
17742 let (hd, nh, nhkv, ti, tkvi, cz) = (
17743 head_dim as i32,
17744 n_head as i32,
17745 n_head_kv as i32,
17746 t as i32,
17747 t_kv as i32,
17748 causal as i32,
17749 );
17750 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17751 let __s_b = self.gpu.stream();
17752 let mut b = __s_b.launch_builder(&f);
17753 b.arg(q)
17754 .arg(&*kw)
17755 .arg(&*vw)
17756 .arg(o)
17757 .arg(&hd)
17758 .arg(&nh)
17759 .arg(&nhkv)
17760 .arg(&ti)
17761 .arg(&tkvi)
17762 .arg(&scale)
17763 .arg(&cz)
17764 .arg(&kdk)
17765 .arg(&kdv);
17766 unsafe {
17767 b.launch(cfg)?;
17768 }
17769 }
17770 Ok(())
17771 }
17772
17773 #[allow(clippy::too_many_arguments)]
17789 pub fn fa_prefill_view_ws_w_hd128(
17790 &self,
17791 q: &CudaSlice<f32>,
17792 k: &cudarc::driver::CudaView<u8>,
17793 v: &cudarc::driver::CudaView<u8>,
17794 o: &mut CudaSlice<f32>,
17795 head_dim: usize,
17796 n_head: usize,
17797 n_head_kv: usize,
17798 t: usize,
17799 t_kv: usize,
17800 scale: f32,
17801 causal: bool,
17802 window: usize,
17803 k_tok_bytes: usize,
17804 v_tok_bytes: usize,
17805 ) -> Result<(), Box<dyn std::error::Error>> {
17806 assert_eq!(
17807 head_dim, 128,
17808 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17809 );
17810 if portable_mma_gated() {
17811 return self.sdpa_naive_w_quantized_view(
17812 q,
17813 k,
17814 v,
17815 o,
17816 head_dim,
17817 n_head,
17818 n_head_kv,
17819 t,
17820 t_kv,
17821 scale,
17822 causal,
17823 window,
17824 k_tok_bytes,
17825 v_tok_bytes,
17826 );
17827 }
17828 const BLOCK_Q: usize = 64;
17829 const BK: usize = 32;
17830 let kv_dim_k = n_head_kv * head_dim;
17831 let kv_dim_v = n_head_kv * head_dim;
17832 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17834 let mut guard = self.prime_deqw_ws.lock().unwrap();
17835 let need_grow = match guard.as_ref() {
17836 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17837 None => true,
17838 };
17839 if need_grow {
17840 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17841 let (ck, cv) = guard
17842 .as_ref()
17843 .map(|(a, b)| (a.len(), b.len()))
17844 .unwrap_or((0, 0));
17845 *guard = Some((
17846 self.alloc_u8(grow(ck, k_ws_bytes))?,
17847 self.alloc_u8(grow(cv, v_ws_bytes))?,
17848 ));
17849 }
17850 let (kw, vw) = guard.as_mut().unwrap();
17851 {
17854 let f = self.func("fa_dequant_kv_ws_bf16");
17855 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17856 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17857 let cfg = LaunchConfig {
17858 grid_dim: (nblk.max(1), 1, 1),
17859 block_dim: (256, 1, 1),
17860 shared_mem_bytes: 0,
17861 };
17862 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17863 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17864 let __s_b = self.gpu.stream();
17865 let mut b = __s_b.launch_builder(&f);
17866 b.arg(k)
17867 .arg(v)
17868 .arg(&mut *kw)
17869 .arg(&mut *vw)
17870 .arg(&kdk)
17871 .arg(&kdv)
17872 .arg(&tkvi)
17873 .arg(&ktb)
17874 .arg(&vtb);
17875 unsafe {
17876 b.launch(cfg)?;
17877 }
17878 }
17879 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17881 .map(|v| v != "0")
17882 .unwrap_or(true);
17883 {
17884 let f = self.func(if db {
17885 "fa_prefill_qw_db_w_hd128"
17886 } else {
17887 "fa_prefill_qw_w_hd128"
17888 });
17889 let shmem = if db {
17890 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17891 } else {
17892 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17893 };
17894 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17895 f.set_attribute(
17896 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17897 shmem as i32,
17898 )?;
17899 let cfg = LaunchConfig {
17900 grid_dim: (
17901 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17902 n_head as u32,
17903 1,
17904 ),
17905 block_dim: (32, 4, 1),
17906 shared_mem_bytes: shmem,
17907 };
17908 let (hd, nh, nhkv, ti, tkvi, cz) = (
17909 head_dim as i32,
17910 n_head as i32,
17911 n_head_kv as i32,
17912 t as i32,
17913 t_kv as i32,
17914 causal as i32,
17915 );
17916 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17917 let __s_b = self.gpu.stream();
17918 let mut b = __s_b.launch_builder(&f);
17919 b.arg(q)
17920 .arg(&*kw)
17921 .arg(&*vw)
17922 .arg(o)
17923 .arg(&hd)
17924 .arg(&nh)
17925 .arg(&nhkv)
17926 .arg(&ti)
17927 .arg(&tkvi)
17928 .arg(&scale)
17929 .arg(&cz)
17930 .arg(&kdk)
17931 .arg(&kdv)
17932 .arg(&wnd);
17933 unsafe {
17934 b.launch(cfg)?;
17935 }
17936 }
17937 Ok(())
17938 }
17939
17940 pub fn fa_decode(
17944 &self,
17945 q: &CudaSlice<f32>,
17946 k: &cudarc::driver::CudaView<u8>,
17947 v: &cudarc::driver::CudaView<u8>,
17948 o: &mut CudaSlice<f32>,
17949 head_dim: usize,
17950 n_head: usize,
17951 n_head_kv: usize,
17952 t_kv: usize,
17953 scale: f32,
17954 k_tok_bytes: usize,
17955 v_tok_bytes: usize,
17956 ) -> Result<(), Box<dyn std::error::Error>> {
17957 self.fa_decode_kvmod(
17958 q,
17959 k,
17960 v,
17961 o,
17962 head_dim,
17963 n_head,
17964 n_head_kv,
17965 t_kv,
17966 scale,
17967 k_tok_bytes,
17968 v_tok_bytes,
17969 false,
17970 )
17971 }
17972
17973 #[allow(clippy::too_many_arguments)]
17977 #[allow(clippy::too_many_arguments)]
17981 #[allow(clippy::too_many_arguments)]
17982 fn fa_decode_scalar_unified(
17983 &self,
17984 q: &cudarc::driver::CudaView<f32>,
17985 k: &cudarc::driver::CudaView<u8>,
17986 v: &cudarc::driver::CudaView<u8>,
17987 o: &mut cudarc::driver::CudaViewMut<f32>,
17988 head_dim: usize,
17989 n_head: usize,
17990 n_head_kv: usize,
17991 t_kv_host: usize,
17992 t_kv_dev: Option<&CudaSlice<i32>>,
17993 scale: f32,
17994 n_splits: usize,
17995 split_keys: usize,
17996 k_tok_bytes: usize,
17997 v_tok_bytes: usize,
17998 g: bool,
17999 part_o: &mut CudaSlice<f32>,
18000 part_m: &mut CudaSlice<f32>,
18001 part_l: &mut CudaSlice<f32>,
18002 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18003 ) -> Result<(), Box<dyn std::error::Error>> {
18004 let f = if g {
18005 self.func_g("fa_decode_f32")
18006 } else {
18007 self.fa_func("fa_decode_f32", head_dim)
18008 };
18009 let cfg = LaunchConfig {
18010 grid_dim: (n_head as u32, n_splits as u32, 1),
18011 block_dim: (head_dim as u32, 1, 1),
18012 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
18013 };
18014 let (hd, nh, nhkv, nsp) = (
18015 head_dim as i32,
18016 n_head as i32,
18017 n_head_kv as i32,
18018 n_splits as i32,
18019 );
18020 let (ktb, vtb, tkvi, ski) = (
18021 k_tok_bytes as i64,
18022 v_tok_bytes as i64,
18023 t_kv_host as i32,
18024 split_keys as i32,
18025 );
18026 let __s_b = self.gpu.stream();
18027 let mut b = __s_b.launch_builder(&f);
18028 match t_kv_dev {
18029 Some(d) => {
18030 b.arg(q)
18031 .arg(k)
18032 .arg(v)
18033 .arg(&mut *part_o)
18034 .arg(&mut *part_m)
18035 .arg(&mut *part_l)
18036 .arg(&hd)
18037 .arg(&nh)
18038 .arg(&nhkv)
18039 .arg(&tkvi)
18040 .arg(d)
18041 .arg(&scale)
18042 .arg(&nsp)
18043 .arg(&ski)
18044 .arg(&ktb)
18045 .arg(&vtb);
18046 unsafe {
18047 b.launch(cfg)?;
18048 }
18049 }
18050 None => {
18051 let null: u64 = 0;
18052 b.arg(q)
18053 .arg(k)
18054 .arg(v)
18055 .arg(&mut *part_o)
18056 .arg(&mut *part_m)
18057 .arg(&mut *part_l)
18058 .arg(&hd)
18059 .arg(&nh)
18060 .arg(&nhkv)
18061 .arg(&tkvi)
18062 .arg(&null)
18063 .arg(&scale)
18064 .arg(&nsp)
18065 .arg(&ski)
18066 .arg(&ktb)
18067 .arg(&vtb);
18068 unsafe {
18069 b.launch(cfg)?;
18070 }
18071 }
18072 }
18073 let cfg2 = LaunchConfig {
18074 grid_dim: (n_head as u32, 1, 1),
18075 block_dim: (head_dim as u32, 1, 1),
18076 shared_mem_bytes: 0,
18077 };
18078 if let Some((oq, od)) = q8_out {
18079 let fc = if g {
18081 self.func_g("fa_decode_combine_q8_1")
18082 } else {
18083 self.fa_func("fa_decode_combine_q8_1", head_dim)
18084 };
18085 let __s_b2 = self.gpu.stream();
18086 let mut b2 = __s_b2.launch_builder(&fc);
18087 b2.arg(&*part_o)
18088 .arg(&*part_m)
18089 .arg(&*part_l)
18090 .arg(oq)
18091 .arg(od)
18092 .arg(&hd)
18093 .arg(&nh)
18094 .arg(&nsp);
18095 unsafe {
18096 b2.launch(cfg2)?;
18097 }
18098 return Ok(());
18099 }
18100 let fc = if g {
18101 self.func_g("fa_decode_combine_f32")
18102 } else {
18103 self.fa_func("fa_decode_combine_f32", head_dim)
18104 };
18105 let __s_b2 = self.gpu.stream();
18106 let mut b2 = __s_b2.launch_builder(&fc);
18107 b2.arg(&*part_o)
18108 .arg(&*part_m)
18109 .arg(&*part_l)
18110 .arg(o)
18111 .arg(&hd)
18112 .arg(&nh)
18113 .arg(&nsp);
18114 unsafe {
18115 b2.launch(cfg2)?;
18116 }
18117 Ok(())
18118 }
18119
18120 pub fn fa_decode_kvmod(
18121 &self,
18122 q: &CudaSlice<f32>,
18123 k: &cudarc::driver::CudaView<u8>,
18124 v: &cudarc::driver::CudaView<u8>,
18125 o: &mut CudaSlice<f32>,
18126 head_dim: usize,
18127 n_head: usize,
18128 n_head_kv: usize,
18129 t_kv: usize,
18130 scale: f32,
18131 k_tok_bytes: usize,
18132 v_tok_bytes: usize,
18133 g: bool,
18134 ) -> Result<(), Box<dyn std::error::Error>> {
18135 let q_view = q.as_view();
18136 let mut o_view = o.as_view_mut();
18137 self.fa_decode_kvmod_view(
18138 &q_view,
18139 k,
18140 v,
18141 &mut o_view,
18142 head_dim,
18143 n_head,
18144 n_head_kv,
18145 t_kv,
18146 scale,
18147 k_tok_bytes,
18148 v_tok_bytes,
18149 g,
18150 )
18151 }
18152
18153 #[allow(clippy::too_many_arguments)]
18158 pub fn fa_decode_kvmod_view(
18159 &self,
18160 q: &cudarc::driver::CudaView<f32>,
18161 k: &cudarc::driver::CudaView<u8>,
18162 v: &cudarc::driver::CudaView<u8>,
18163 o: &mut cudarc::driver::CudaViewMut<f32>,
18164 head_dim: usize,
18165 n_head: usize,
18166 n_head_kv: usize,
18167 t_kv: usize,
18168 scale: f32,
18169 k_tok_bytes: usize,
18170 v_tok_bytes: usize,
18171 g: bool,
18172 ) -> Result<(), Box<dyn std::error::Error>> {
18173 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
18194 if g && head_dim == 256 && !fa_v4_at(t_kv) {
18198 fa_vec = false;
18199 }
18200 let sp = fa_split_keys(t_kv, n_head_kv);
18201 let n_splits = if fa_vec {
18202 ((t_kv + sp - 1) / sp).max(1)
18203 } else {
18204 ((t_kv + 255) / 256).max(1)
18205 };
18206 let o_len = n_head * n_splits * head_dim;
18207 let ml_len = n_head * n_splits;
18208 let mut part_guard = self.fa_part_pool.lock().unwrap();
18209 if part_guard
18210 .as_ref()
18211 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18212 .unwrap_or(true)
18213 {
18214 let old = part_guard.take();
18225 let (co, cm) = old
18226 .as_ref()
18227 .map(|pp| (pp.0.len(), pp.1.len()))
18228 .unwrap_or((0, 0));
18229 if let Some(old) = old {
18230 self.fa_part_retired.lock().unwrap().push(old);
18231 }
18232 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18233 eprintln!(
18234 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18235 co, o_len, cm, ml_len
18236 );
18237 }
18238 *part_guard = Some((
18239 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18240 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18241 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18242 ));
18243 }
18244 let pg = part_guard.as_mut().unwrap();
18245 self.gpu
18246 .stream()
18247 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18248 self.gpu
18249 .stream()
18250 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18251 self.gpu
18252 .stream()
18253 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18254 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18255 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18256 let (hd, nh, nhkv, tkvi, nsp) = (
18257 head_dim as i32,
18258 n_head as i32,
18259 n_head_kv as i32,
18260 t_kv as i32,
18261 n_splits as i32,
18262 );
18263 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18264 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
18268 let fa512_min = fa512_min_tkv();
18273 let deep = fa_vec
18276 && head_dim == 256
18277 && fa_v4_at(t_kv)
18278 && !g
18279 && fa_deep_at(t_kv)
18280 && !matches!(fa_v4_mode(), "noB3" | "stage");
18281 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
18282 let gqa = (n_head / n_head_kv).max(1) as u32;
18285 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
18286 (
18287 fv,
18288 LaunchConfig {
18289 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18290 block_dim: (32, gqa, 1),
18291 shared_mem_bytes: 0,
18292 },
18293 )
18294 } else if fa_vec && head_dim <= 256 {
18295 let gqa = (n_head / n_head_kv).max(1) as u32;
18296 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18307 let smem_tkv = *SMEM_TKV.get_or_init(|| {
18308 std::env::var("MEMRA_FA_SMEM_TKV")
18309 .ok()
18310 .and_then(|v| v.parse().ok())
18311 .unwrap_or_else(|| {
18312 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18313 })
18314 });
18315 if fa_v4_at(t_kv) && head_dim == 256 {
18316 let v4name = match fa_v4_mode() {
18320 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
18323 _ => "fa_decode_vec_q_v4",
18324 };
18325 let fv = if g {
18326 self.func_g(v4name)
18327 } else {
18328 self.func(v4name)
18329 };
18330 let shmem = (if deep { 12160 } else { 11520 }
18333 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
18334 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18335 fv.set_attribute(
18336 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18337 shmem as i32,
18338 )?;
18339 (
18340 fv,
18341 LaunchConfig {
18342 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18343 block_dim: (32, gqa, 1),
18344 shared_mem_bytes: shmem,
18345 },
18346 )
18347 } else if fa_v3_active(head_dim) {
18348 let fv = if g {
18351 self.func_g("fa_decode_vec_q_v3")
18352 } else {
18353 self.func("fa_decode_vec_q_v3")
18354 };
18355 let shmem = (32 * head_dim * 2) as u32; (
18357 fv,
18358 LaunchConfig {
18359 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18360 block_dim: (32, gqa, 1),
18361 shared_mem_bytes: shmem,
18362 },
18363 )
18364 } else if fa_v2_on() {
18365 let fv = if g {
18369 self.func_g("fa_decode_vec_q_v2")
18370 } else {
18371 self.func("fa_decode_vec_q_v2")
18372 };
18373 let shmem = (2 * 32 * head_dim * 2) as u32; (
18375 fv,
18376 LaunchConfig {
18377 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18378 block_dim: (32, gqa, 1),
18379 shared_mem_bytes: shmem,
18380 },
18381 )
18382 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
18383 {
18384 let fv = if g {
18388 self.func_g("fa_decode_vec_q_smem")
18389 } else {
18390 self.func("fa_decode_vec_q_smem")
18391 };
18392 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
18394 fv.set_attribute(
18395 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18396 shmem as i32,
18397 )?;
18398 (
18399 fv,
18400 LaunchConfig {
18401 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18402 block_dim: (32, gqa, 1),
18403 shared_mem_bytes: shmem,
18404 },
18405 )
18406 } else {
18407 let fv = if g {
18410 self.func_g("fa_decode_vec_q")
18411 } else {
18412 self.func("fa_decode_vec_q")
18413 };
18414 (
18415 fv,
18416 LaunchConfig {
18417 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
18418 block_dim: (32, gqa, 1),
18419 shared_mem_bytes: 0,
18420 },
18421 )
18422 }
18423 } else {
18424 return self.fa_decode_scalar_unified(
18427 q,
18428 k,
18429 v,
18430 o,
18431 head_dim,
18432 n_head,
18433 n_head_kv,
18434 t_kv,
18435 None,
18436 scale,
18437 n_splits,
18438 if fa_vec { sp } else { 256 },
18439 k_tok_bytes,
18440 v_tok_bytes,
18441 g,
18442 part_o,
18443 part_m,
18444 part_l,
18445 None,
18446 );
18447 };
18448 let __s_b = self.gpu.stream();
18449 let mut b = __s_b.launch_builder(&f);
18450 b.arg(q)
18451 .arg(k)
18452 .arg(v)
18453 .arg(&mut *part_o)
18454 .arg(&mut *part_m)
18455 .arg(&mut *part_l)
18456 .arg(&hd)
18457 .arg(&nh)
18458 .arg(&nhkv)
18459 .arg(&tkvi)
18460 .arg(&scale)
18461 .arg(&nsp)
18462 .arg(&ktb)
18463 .arg(&vtb);
18464 unsafe {
18465 b.launch(cfg)?;
18466 }
18467 let (fc, cfg2) = (
18470 if g {
18471 self.func_g("fa_decode_combine_f32")
18472 } else {
18473 self.fa_func("fa_decode_combine_f32", head_dim)
18474 },
18475 LaunchConfig {
18476 grid_dim: (n_head as u32, 1, 1),
18477 block_dim: (head_dim as u32, 1, 1),
18478 shared_mem_bytes: 0,
18479 },
18480 );
18481 let __s_b2 = self.gpu.stream();
18482 let mut b2 = __s_b2.launch_builder(&fc);
18483 b2.arg(&*part_o)
18484 .arg(&*part_m)
18485 .arg(&*part_l)
18486 .arg(o)
18487 .arg(&hd)
18488 .arg(&nh)
18489 .arg(&nsp);
18490 unsafe {
18491 b2.launch(cfg2)?;
18492 }
18493 Ok(())
18494 }
18495
18496 #[allow(clippy::too_many_arguments)]
18507 pub fn fa_decode_batch_seqs_v4(
18508 &self,
18509 q: &CudaSlice<f32>,
18510 kv_ptrs: &cudarc::driver::CudaView<u64>,
18511 pos_seq: &CudaSlice<i32>,
18512 o: &mut CudaSlice<f32>,
18513 head_dim: usize,
18514 n_head: usize,
18515 n_head_kv: usize,
18516 b_n: usize,
18517 t_kv_max: usize,
18518 scale: f32,
18519 split_keys: usize,
18520 k_tok_bytes: usize,
18521 v_tok_bytes: usize,
18522 ) -> Result<(), Box<dyn std::error::Error>> {
18523 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
18524 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
18525 let o_len = b_n * n_head * n_splits_max * head_dim;
18526 let ml_len = b_n * n_head * n_splits_max;
18527 let mut part_guard = self.fa_part_pool.lock().unwrap();
18528 if part_guard
18529 .as_ref()
18530 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18531 .unwrap_or(true)
18532 {
18533 let old = part_guard.take();
18544 let (co, cm) = old
18545 .as_ref()
18546 .map(|pp| (pp.0.len(), pp.1.len()))
18547 .unwrap_or((0, 0));
18548 if let Some(old) = old {
18549 self.fa_part_retired.lock().unwrap().push(old);
18550 }
18551 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18552 eprintln!(
18553 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18554 co, o_len, cm, ml_len
18555 );
18556 }
18557 *part_guard = Some((
18558 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18559 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18560 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18561 ));
18562 }
18563 let pg = part_guard.as_mut().unwrap();
18564 self.gpu
18565 .stream()
18566 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18567 self.gpu
18568 .stream()
18569 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18570 self.gpu
18571 .stream()
18572 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18573 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18574 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18575 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
18576 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18577 let gqa = (n_head / n_head_kv).max(1) as u32;
18578 let f = self.func("fa_decode_vec_q_seqs_v4");
18579 let shmem = (11520 + 32 * head_dim * 2) as u32;
18581 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18582 f.set_attribute(
18583 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18584 shmem as i32,
18585 )?;
18586 let cfg = LaunchConfig {
18587 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
18588 block_dim: (32, gqa, 1),
18589 shared_mem_bytes: shmem,
18590 };
18591 {
18592 let __s_b = self.gpu.stream();
18593 let mut b = __s_b.launch_builder(&f);
18594 b.arg(q)
18595 .arg(kv_ptrs)
18596 .arg(pos_seq)
18597 .arg(&mut *part_o)
18598 .arg(&mut *part_m)
18599 .arg(&mut *part_l)
18600 .arg(&hd)
18601 .arg(&nh)
18602 .arg(&nhkv)
18603 .arg(&scale)
18604 .arg(&nspm)
18605 .arg(&spk)
18606 .arg(&ktb)
18607 .arg(&vtb);
18608 unsafe {
18609 b.launch(cfg)?;
18610 }
18611 }
18612 let fc = self.func("fa_decode_combine_seqs");
18613 let cfg2 = LaunchConfig {
18614 grid_dim: (n_head as u32, b_n as u32, 1),
18615 block_dim: (head_dim as u32, 1, 1),
18616 shared_mem_bytes: 0,
18617 };
18618 let __s_b2 = self.gpu.stream();
18619 let mut b2 = __s_b2.launch_builder(&fc);
18620 b2.arg(&*part_o)
18621 .arg(&*part_m)
18622 .arg(&*part_l)
18623 .arg(o)
18624 .arg(&hd)
18625 .arg(&nh)
18626 .arg(pos_seq)
18627 .arg(&nspm)
18628 .arg(&spk);
18629 unsafe {
18630 b2.launch(cfg2)?;
18631 }
18632 Ok(())
18633 }
18634
18635 #[allow(clippy::too_many_arguments)]
18642 pub fn append_kv_quantized_seqs(
18643 &self,
18644 k_rows: &CudaSlice<f32>,
18645 v_rows: &CudaSlice<f32>,
18646 kv_ptrs: &cudarc::driver::CudaView<u64>,
18647 pos_seq: &CudaSlice<i32>,
18648 b_n: usize,
18649 kv_dim_k: usize,
18650 kv_dim_v: usize,
18651 k_tok_bytes: usize,
18652 v_tok_bytes: usize,
18653 ) -> Result<(), Box<dyn std::error::Error>> {
18654 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
18655 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
18656 let cfg = LaunchConfig {
18657 grid_dim: (nblk, b_n as u32, 1),
18658 block_dim: (32, 1, 1),
18659 shared_mem_bytes: 0,
18660 };
18661 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
18662 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18663 let __s_b = self.gpu.stream();
18664 let mut b = __s_b.launch_builder(&f);
18665 b.arg(k_rows)
18666 .arg(v_rows)
18667 .arg(kv_ptrs)
18668 .arg(pos_seq)
18669 .arg(&kdk)
18670 .arg(&kdv)
18671 .arg(&ktb)
18672 .arg(&vtb);
18673 unsafe {
18674 b.launch(cfg)?;
18675 }
18676 Ok(())
18677 }
18678
18679 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
18685 std::env::var("MEMRA_NO_FA_VEC").is_err()
18686 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
18687 && base_len + 1 >= fa_vec_min_tkv()
18688 && head_dim <= 256
18689 && head_dim % 32 == 0
18690 }
18691
18692 #[allow(clippy::too_many_arguments)]
18701 pub fn fa_decode_rows(
18702 &self,
18703 q: &CudaSlice<f32>,
18704 k: &cudarc::driver::CudaView<u8>,
18705 v: &cudarc::driver::CudaView<u8>,
18706 o: &mut CudaSlice<f32>,
18707 head_dim: usize,
18708 n_head: usize,
18709 n_head_kv: usize,
18710 base_len: usize,
18711 t: usize,
18712 scale: f32,
18713 k_tok_bytes: usize,
18714 v_tok_bytes: usize,
18715 base_dev: Option<(&CudaSlice<i32>, i32)>,
18719 kv_shared: bool,
18722 g: bool,
18726 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18729 ) -> Result<(), Box<dyn std::error::Error>> {
18730 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
18731 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
18738 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18739 let v = *SP512.get_or_init(|| {
18742 std::env::var("MEMRA_FA_SP512")
18743 .ok()
18744 .and_then(|x| x.parse().ok())
18745 .unwrap_or(0)
18746 });
18747 sp = if v >= 8 {
18748 v
18749 } else {
18750 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18751 };
18752 }
18753 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18754 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18755 let gqa = (n_head / n_head_kv).max(1) as u32;
18756 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
18767 groups.push((0, t, sp));
18768 } else {
18769 let mut r0 = 0usize;
18770 while r0 < t {
18771 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
18772 let mut r1 = r0 + 1;
18773 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18774 r1 += 1;
18775 }
18776 groups.push((r0, r1 - r0, sp_g));
18777 r0 = r1;
18778 }
18779 }
18780 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18784 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18785 std::env::var("MEMRA_FA_SMEM_TKV")
18786 .ok()
18787 .and_then(|v| v.parse().ok())
18788 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18789 });
18790 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18791 let v3 = fa_v3_active(head_dim);
18792 let smem_rows =
18793 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18794 let _ = kv_shared;
18799 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18802 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18816 let tb512 = head_dim == 512
18818 && sp <= 32
18819 && n_head / n_head_kv.max(1) <= 16
18820 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18821 let fname = if tb512 {
18822 "fa_decode_vec_q_rows_v4_512_tb"
18823 } else if i2 {
18824 "fa_decode_vec_q_rows_dpl16_i2"
18825 } else if head_dim == 512 {
18826 "fa_decode_vec_q_rows_dpl16"
18827 }
18828 else if v4 {
18830 "fa_decode_vec_q_rows_v4"
18831 } else if v3 {
18832 "fa_decode_vec_q_rows_v3"
18833 } else if fa_v2_on() {
18834 "fa_decode_vec_q_rows_v2"
18835 } else if smem_rows {
18836 "fa_decode_vec_q_rows_smem"
18837 } else {
18838 "fa_decode_vec_q_rows"
18839 };
18840 let f = if head_dim == 512 {
18841 self.fa_func(fname, head_dim)
18842 } else if g {
18843 self.func_g(if smem_rows {
18851 "fa_decode_vec_q_rows"
18852 } else {
18853 fname
18854 })
18855 } else {
18856 self.func(fname)
18857 };
18858 let shmem = if tb512 {
18859 let gk = Self::gkv_on();
18861 let sh =
18862 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18863 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18864 f.set_attribute(
18865 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18866 sh as i32,
18867 )?;
18868 sh
18869 } else if v4 || v3 || smem_rows || fa_v2_on() {
18870 let sh = (if v4 {
18872 11520 + 32 * head_dim * if g { 1 } else { 2 }
18873 } else if v3 {
18874 32 * head_dim * 2
18875 } else {
18876 2 * 32 * head_dim * 2
18877 }) as u32;
18878 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18879 f.set_attribute(
18880 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18881 sh as i32,
18882 )?;
18883 sh
18884 } else {
18885 0
18886 };
18887 for &(r0, t_g, sp_g) in &groups {
18891 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18892 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18893 let base_i = (base_len + r0) as i32;
18894 let o_len = t_g * n_head * n_splits_g * head_dim;
18895 let ml_len = t_g * n_head * n_splits_g;
18896 let mut part_guard = self.fa_part_pool.lock().unwrap();
18897 if part_guard
18898 .as_ref()
18899 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18900 .unwrap_or(true)
18901 {
18902 let old = part_guard.take();
18913 let (co, cm) = old
18914 .as_ref()
18915 .map(|pp| (pp.0.len(), pp.1.len()))
18916 .unwrap_or((0, 0));
18917 if let Some(old) = old {
18918 self.fa_part_retired.lock().unwrap().push(old);
18919 }
18920 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18921 eprintln!(
18922 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18923 co, o_len, cm, ml_len
18924 );
18925 }
18926 *part_guard = Some((
18927 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18928 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18929 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18930 ));
18931 }
18932 let pg = part_guard.as_mut().unwrap();
18933 self.gpu
18934 .stream()
18935 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18936 self.gpu
18937 .stream()
18938 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18939 self.gpu
18940 .stream()
18941 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18942 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18943 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18944 let qv = self.view(q, t * n_head * head_dim);
18945 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18946 let cfg = LaunchConfig {
18947 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
18948 block_dim: (32, gqa, 1),
18949 shared_mem_bytes: shmem,
18950 };
18951 {
18952 let __s_b = self.gpu.stream();
18953 let mut b = __s_b.launch_builder(&f);
18954 if tb512 {
18955 let (bd, plus) =
18957 base_dev.expect("hd512 rows twin requires a device base counter");
18958 let plus_g = plus + r0 as i32;
18959 let nr = t_g as i32;
18960 if Self::pdl_on() && Self::pdl_wb_on() {
18961 use cudarc::driver::{DevicePtr, DevicePtrMut};
18963 let s = &self.gpu.stream();
18964 let (pq, _b0) = q_g.device_ptr(s);
18965 let (pk, _b1) = k.device_ptr(s);
18966 let (pv, _b2) = v.device_ptr(s);
18967 let (po, _b3) = part_o.device_ptr_mut(s);
18968 let (pm, _b4) = part_m.device_ptr_mut(s);
18969 let (pl, _b5) = part_l.device_ptr_mut(s);
18970 let (pb, _b6) = bd.device_ptr(s);
18971 let mut ps = [
18972 &pq as *const _ as *mut std::ffi::c_void,
18973 &pk as *const _ as *mut _,
18974 &pv as *const _ as *mut _,
18975 &po as *const _ as *mut _,
18976 &pm as *const _ as *mut _,
18977 &pl as *const _ as *mut _,
18978 &hd as *const _ as *mut _,
18979 &nh as *const _ as *mut _,
18980 &nhkv as *const _ as *mut _,
18981 &pb as *const _ as *mut _,
18982 &plus_g as *const _ as *mut _,
18983 &scale as *const _ as *mut _,
18984 &nspm as *const _ as *mut _,
18985 &spk as *const _ as *mut _,
18986 &ktb as *const _ as *mut _,
18987 &vtb as *const _ as *mut _,
18988 &nr as *const _ as *mut _,
18989 ];
18990 unsafe {
18991 self.launch_pdl_flash(
18992 Self::gkv_on(),
18993 "fa_decode_vec_q_rows_v4_512_tb",
18994 (n_head_kv as u32, n_splits_g as u32, 1),
18995 (32, gqa, 1),
18996 shmem,
18997 &mut ps,
18998 )?;
18999 }
19000 } else {
19001 let cfg_tb = LaunchConfig {
19002 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
19003 block_dim: (32, gqa, 1),
19004 shared_mem_bytes: shmem,
19005 };
19006 b.arg(&q_g)
19007 .arg(k)
19008 .arg(v)
19009 .arg(&mut *part_o)
19010 .arg(&mut *part_m)
19011 .arg(&mut *part_l)
19012 .arg(&hd)
19013 .arg(&nh)
19014 .arg(&nhkv)
19015 .arg(bd)
19016 .arg(&plus_g)
19017 .arg(&scale)
19018 .arg(&nspm)
19019 .arg(&spk)
19020 .arg(&ktb)
19021 .arg(&vtb)
19022 .arg(&nr);
19023 unsafe {
19024 b.launch(cfg_tb)?;
19025 }
19026 }
19027 } else if head_dim == 512 {
19028 let (bd, plus) =
19029 base_dev.expect("hd512 rows twin requires a device base counter");
19030 let plus_g = plus + r0 as i32;
19031 b.arg(&q_g)
19032 .arg(k)
19033 .arg(v)
19034 .arg(&mut *part_o)
19035 .arg(&mut *part_m)
19036 .arg(&mut *part_l)
19037 .arg(&hd)
19038 .arg(&nh)
19039 .arg(&nhkv)
19040 .arg(bd)
19041 .arg(&plus_g)
19042 .arg(&scale)
19043 .arg(&nspm)
19044 .arg(&spk)
19045 .arg(&ktb)
19046 .arg(&vtb);
19047 unsafe {
19048 b.launch(cfg)?;
19049 }
19050 } else {
19051 b.arg(&q_g)
19052 .arg(k)
19053 .arg(v)
19054 .arg(&mut *part_o)
19055 .arg(&mut *part_m)
19056 .arg(&mut *part_l)
19057 .arg(&hd)
19058 .arg(&nh)
19059 .arg(&nhkv)
19060 .arg(&base_i)
19061 .arg(&scale)
19062 .arg(&nspm)
19063 .arg(&spk)
19064 .arg(&ktb)
19065 .arg(&vtb);
19066 unsafe {
19067 b.launch(cfg)?;
19068 }
19069 }
19070 }
19071 let cfg2 = LaunchConfig {
19072 grid_dim: (n_head as u32, t_g as u32, 1),
19073 block_dim: (head_dim as u32, 1, 1),
19074 shared_mem_bytes: 0,
19075 };
19076 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
19077 if head_dim == 512 {
19078 let (bd, plus) = base_dev.unwrap();
19081 let plus_g = plus + r0 as i32;
19082 if let Some((oq, od)) = q8_out.as_mut() {
19083 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
19085 if Self::pdl_on() && Self::pdl_wb_on() {
19086 use cudarc::driver::{DevicePtr, DevicePtrMut};
19088 let s = &self.gpu.stream();
19089 let (po, _g0) = part_o.device_ptr(s);
19090 let (pm, _g1) = part_m.device_ptr(s);
19091 let (pl, _g2) = part_l.device_ptr(s);
19092 let (pq, _g3) = oq.device_ptr_mut(s);
19093 let (pd, _g4) = od.device_ptr_mut(s);
19094 let (pb, _g5) = bd.device_ptr(s);
19095 let mut ps = [
19096 &po as *const _ as *mut std::ffi::c_void,
19097 &pm as *const _ as *mut _,
19098 &pl as *const _ as *mut _,
19099 &pq as *const _ as *mut _,
19100 &pd as *const _ as *mut _,
19101 &hd as *const _ as *mut _,
19102 &nh as *const _ as *mut _,
19103 &pb as *const _ as *mut _,
19104 &plus_g as *const _ as *mut _,
19105 &nspm as *const _ as *mut _,
19106 &spk as *const _ as *mut _,
19107 ];
19108 unsafe {
19109 self.launch_pdl_flash(
19110 Self::gkv_on(),
19111 "fa_decode_combine_rows_dc_q8_1",
19112 cfg2.grid_dim,
19113 cfg2.block_dim,
19114 0,
19115 &mut ps,
19116 )?;
19117 }
19118 continue;
19119 }
19120 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
19121 let __s_b2 = self.gpu.stream();
19122 let mut b2 = __s_b2.launch_builder(&fc);
19123 b2.arg(&*part_o)
19124 .arg(&*part_m)
19125 .arg(&*part_l)
19126 .arg(&mut **oq)
19127 .arg(&mut **od)
19128 .arg(&hd)
19129 .arg(&nh)
19130 .arg(bd)
19131 .arg(&plus_g)
19132 .arg(&nspm)
19133 .arg(&spk);
19134 unsafe {
19135 b2.launch(cfg2)?;
19136 }
19137 continue;
19138 }
19139 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
19140 let __s_b2 = self.gpu.stream();
19141 let mut b2 = __s_b2.launch_builder(&fc);
19142 b2.arg(&*part_o)
19143 .arg(&*part_m)
19144 .arg(&*part_l)
19145 .arg(&mut o_g)
19146 .arg(&hd)
19147 .arg(&nh)
19148 .arg(bd)
19149 .arg(&plus_g)
19150 .arg(&nspm)
19151 .arg(&spk);
19152 unsafe {
19153 b2.launch(cfg2)?;
19154 }
19155 } else {
19156 assert!(
19159 q8_out.is_none(),
19160 "rows q8 emit requires the hd512 dc combine"
19161 );
19162 let fc = self.func("fa_decode_combine_rows");
19163 let __s_b2 = self.gpu.stream();
19164 let mut b2 = __s_b2.launch_builder(&fc);
19165 b2.arg(&*part_o)
19166 .arg(&*part_m)
19167 .arg(&*part_l)
19168 .arg(&mut o_g)
19169 .arg(&hd)
19170 .arg(&nh)
19171 .arg(&base_i)
19172 .arg(&nspm)
19173 .arg(&spk);
19174 unsafe {
19175 b2.launch(cfg2)?;
19176 }
19177 }
19178 }
19179 Ok(())
19180 }
19181
19182 #[allow(clippy::too_many_arguments)]
19186 pub fn fa_decode_rows_w(
19187 &self,
19188 q: &CudaSlice<f32>,
19189 k: &cudarc::driver::CudaView<u8>,
19190 v: &cudarc::driver::CudaView<u8>,
19191 o: &mut CudaSlice<f32>,
19192 head_dim: usize,
19193 n_head: usize,
19194 n_head_kv: usize,
19195 base_dev: &CudaSlice<i32>,
19196 base_plus: i32,
19197 t: usize,
19198 scale: f32,
19199 window: usize,
19200 k_tok_bytes: usize,
19201 v_tok_bytes: usize,
19202 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19203 ) -> Result<(), Box<dyn std::error::Error>> {
19204 debug_assert!(head_dim == 256);
19209 let sp = {
19217 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19218 let v = *SPW.get_or_init(|| {
19219 std::env::var("MEMRA_FA_SPW")
19220 .ok()
19221 .and_then(|x| x.parse().ok())
19222 .unwrap_or(0)
19223 });
19224 if v >= 8 {
19225 v
19226 } else {
19227 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
19228 }
19229 };
19230 let n_splits_max = (window + sp - 1) / sp;
19231 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19232 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
19233 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19234 let gqa = (n_head / n_head_kv).max(1) as u32;
19235 let o_len = t * n_head * n_splits_max * head_dim;
19236 let ml_len = t * n_head * n_splits_max;
19237 let mut part_guard = self.fa_part_pool.lock().unwrap();
19238 if part_guard
19239 .as_ref()
19240 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19241 .unwrap_or(true)
19242 {
19243 let old = part_guard.take();
19254 let (co, cm) = old
19255 .as_ref()
19256 .map(|pp| (pp.0.len(), pp.1.len()))
19257 .unwrap_or((0, 0));
19258 if let Some(old) = old {
19259 self.fa_part_retired.lock().unwrap().push(old);
19260 }
19261 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19262 eprintln!(
19263 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19264 co, o_len, cm, ml_len
19265 );
19266 }
19267 *part_guard = Some((
19268 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19269 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19270 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19271 ));
19272 }
19273 let pg = part_guard.as_mut().unwrap();
19274 self.gpu
19275 .stream()
19276 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19277 self.gpu
19278 .stream()
19279 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19280 self.gpu
19281 .stream()
19282 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19283 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19284 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19290 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
19291 std::env::var("MEMRA_FA_SMEM_TKV")
19292 .ok()
19293 .and_then(|v| v.parse().ok())
19294 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
19295 });
19296 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19302 let wg = Self::wkv_on();
19307 let sp2 =
19310 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
19311 if sp2 {
19312 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
19313 if Self::pdl_on() && Self::pdl_wb_on() {
19314 use cudarc::driver::{DevicePtr, DevicePtrMut};
19316 let s = &self.gpu.stream();
19317 let (pq, _b0) = q.device_ptr(s);
19318 let (pk, _b1) = k.device_ptr(s);
19319 let (pv, _b2) = v.device_ptr(s);
19320 let (po, _b3) = part_o.device_ptr_mut(s);
19321 let (pm, _b4) = part_m.device_ptr_mut(s);
19322 let (pl, _b5) = part_l.device_ptr_mut(s);
19323 let (pb, _b6) = base_dev.device_ptr(s);
19324 let mut ps = [
19325 &pq as *const _ as *mut std::ffi::c_void,
19326 &pk as *const _ as *mut _,
19327 &pv as *const _ as *mut _,
19328 &po as *const _ as *mut _,
19329 &pm as *const _ as *mut _,
19330 &pl as *const _ as *mut _,
19331 &hd as *const _ as *mut _,
19332 &nh as *const _ as *mut _,
19333 &nhkv as *const _ as *mut _,
19334 &pb as *const _ as *mut _,
19335 &base_plus as *const _ as *mut _,
19336 &scale as *const _ as *mut _,
19337 &nspm as *const _ as *mut _,
19338 &spk as *const _ as *mut _,
19339 &ktb as *const _ as *mut _,
19340 &vtb as *const _ as *mut _,
19341 &wini as *const _ as *mut _,
19342 ];
19343 unsafe {
19344 self.launch_pdl_flash(
19345 wg,
19346 "fa_decode_vec_q_rows_v4_w_sp",
19347 (n_head_kv as u32, n_splits_max as u32, t as u32),
19348 (32, gqa + 1, 1),
19349 sh,
19350 &mut ps,
19351 )?;
19352 }
19353 } else {
19354 let f = if wg {
19355 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
19356 } else {
19357 self.func("fa_decode_vec_q_rows_v4_w_sp")
19358 };
19359 f.set_attribute(
19360 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19361 sh as i32,
19362 )?;
19363 let cfg = LaunchConfig {
19364 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19365 block_dim: (32, gqa + 1, 1),
19366 shared_mem_bytes: sh,
19367 };
19368 let __s_b = self.gpu.stream();
19369 let mut b = __s_b.launch_builder(&f);
19370 b.arg(q)
19371 .arg(k)
19372 .arg(v)
19373 .arg(&mut *part_o)
19374 .arg(&mut *part_m)
19375 .arg(&mut *part_l)
19376 .arg(&hd)
19377 .arg(&nh)
19378 .arg(&nhkv)
19379 .arg(base_dev)
19380 .arg(&base_plus)
19381 .arg(&scale)
19382 .arg(&nspm)
19383 .arg(&spk)
19384 .arg(&ktb)
19385 .arg(&vtb)
19386 .arg(&wini);
19387 unsafe {
19388 b.launch(cfg)?;
19389 }
19390 }
19391 } else {
19392 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
19393 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
19395 use cudarc::driver::{DevicePtr, DevicePtrMut};
19396 let s = &self.gpu.stream();
19397 let (pq, _b0) = q.device_ptr(s);
19398 let (pk, _b1) = k.device_ptr(s);
19399 let (pv, _b2) = v.device_ptr(s);
19400 let (po, _b3) = part_o.device_ptr_mut(s);
19401 let (pm, _b4) = part_m.device_ptr_mut(s);
19402 let (pl, _b5) = part_l.device_ptr_mut(s);
19403 let (pb, _b6) = base_dev.device_ptr(s);
19404 let mut ps = [
19405 &pq as *const _ as *mut std::ffi::c_void,
19406 &pk as *const _ as *mut _,
19407 &pv as *const _ as *mut _,
19408 &po as *const _ as *mut _,
19409 &pm as *const _ as *mut _,
19410 &pl as *const _ as *mut _,
19411 &hd as *const _ as *mut _,
19412 &nh as *const _ as *mut _,
19413 &nhkv as *const _ as *mut _,
19414 &pb as *const _ as *mut _,
19415 &base_plus as *const _ as *mut _,
19416 &scale as *const _ as *mut _,
19417 &nspm as *const _ as *mut _,
19418 &spk as *const _ as *mut _,
19419 &ktb as *const _ as *mut _,
19420 &vtb as *const _ as *mut _,
19421 &wini as *const _ as *mut _,
19422 ];
19423 unsafe {
19424 self.launch_pdl_flash(
19425 wg,
19426 "fa_decode_vec_q_rows_v4_w",
19427 (n_head_kv as u32, n_splits_max as u32, t as u32),
19428 (32, gqa, 1),
19429 sh,
19430 &mut ps,
19431 )?;
19432 }
19433 } else {
19434 let pick = |name: &str| {
19435 if wg {
19436 self.func_g(name)
19437 } else {
19438 self.func(name)
19439 }
19440 };
19441 let (f, sh) = if fa_v4_at(window) {
19442 let f = pick("fa_decode_vec_q_rows_v4_w");
19443 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
19444 } else if smem_tkv > 0 && window >= smem_tkv {
19445 (
19448 pick("fa_decode_vec_q_rows_smem_w"),
19449 (2 * 32 * head_dim * 2) as u32,
19450 )
19451 } else {
19452 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
19453 };
19454 f.set_attribute(
19455 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19456 sh as i32,
19457 )?;
19458 let cfg = LaunchConfig {
19459 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19460 block_dim: (32, gqa, 1),
19461 shared_mem_bytes: sh,
19462 };
19463 let __s_b = self.gpu.stream();
19464 let mut b = __s_b.launch_builder(&f);
19465 b.arg(q)
19466 .arg(k)
19467 .arg(v)
19468 .arg(&mut *part_o)
19469 .arg(&mut *part_m)
19470 .arg(&mut *part_l)
19471 .arg(&hd)
19472 .arg(&nh)
19473 .arg(&nhkv)
19474 .arg(base_dev)
19475 .arg(&base_plus)
19476 .arg(&scale)
19477 .arg(&nspm)
19478 .arg(&spk)
19479 .arg(&ktb)
19480 .arg(&vtb)
19481 .arg(&wini);
19482 unsafe {
19483 b.launch(cfg)?;
19484 }
19485 }
19486 }
19487 let cfg2 = LaunchConfig {
19488 grid_dim: (n_head as u32, t as u32, 1),
19489 block_dim: (head_dim as u32, 1, 1),
19490 shared_mem_bytes: 0,
19491 };
19492 if let Some((oq, od)) = q8_out {
19493 if Self::pdl_on() && Self::pdl_wb_on() {
19496 use cudarc::driver::{DevicePtr, DevicePtrMut};
19498 let s = &self.gpu.stream();
19499 let (po, _g0) = part_o.device_ptr(s);
19500 let (pm, _g1) = part_m.device_ptr(s);
19501 let (pl, _g2) = part_l.device_ptr(s);
19502 let (pq, _g3) = oq.device_ptr_mut(s);
19503 let (pd, _g4) = od.device_ptr_mut(s);
19504 let mut ps = [
19505 &po as *const _ as *mut std::ffi::c_void,
19506 &pm as *const _ as *mut _,
19507 &pl as *const _ as *mut _,
19508 &pq as *const _ as *mut _,
19509 &pd as *const _ as *mut _,
19510 &hd as *const _ as *mut _,
19511 &nh as *const _ as *mut _,
19512 &nspm as *const _ as *mut _,
19513 &spk as *const _ as *mut _,
19514 &wini as *const _ as *mut _,
19515 ];
19516 unsafe {
19517 self.launch_pdl_flash(
19518 wg,
19519 "fa_decode_combine_rows_w_q8_1",
19520 cfg2.grid_dim,
19521 cfg2.block_dim,
19522 0,
19523 &mut ps,
19524 )?;
19525 }
19526 return Ok(());
19527 }
19528 let fc = if wg {
19529 self.func_g("fa_decode_combine_rows_w_q8_1")
19530 } else {
19531 self.func("fa_decode_combine_rows_w_q8_1")
19532 };
19533 let __s_b2 = self.gpu.stream();
19534 let mut b2 = __s_b2.launch_builder(&fc);
19535 b2.arg(&*part_o)
19536 .arg(&*part_m)
19537 .arg(&*part_l)
19538 .arg(oq)
19539 .arg(od)
19540 .arg(&hd)
19541 .arg(&nh)
19542 .arg(&nspm)
19543 .arg(&spk)
19544 .arg(&wini);
19545 unsafe {
19546 b2.launch(cfg2)?;
19547 }
19548 return Ok(());
19549 }
19550 let fc = if wg {
19551 self.func_g("fa_decode_combine_rows_w")
19552 } else {
19553 self.func("fa_decode_combine_rows_w")
19554 };
19555 let __s_b2 = self.gpu.stream();
19556 let mut b2 = __s_b2.launch_builder(&fc);
19557 b2.arg(&*part_o)
19558 .arg(&*part_m)
19559 .arg(&*part_l)
19560 .arg(o)
19561 .arg(&hd)
19562 .arg(&nh)
19563 .arg(&nspm)
19564 .arg(&spk)
19565 .arg(&wini);
19566 unsafe {
19567 b2.launch(cfg2)?;
19568 }
19569 Ok(())
19570 }
19571
19572 #[allow(clippy::too_many_arguments)]
19578 pub fn fa_decode_rows_dc(
19579 &self,
19580 q: &CudaSlice<f32>,
19581 k: &cudarc::driver::CudaView<u8>,
19582 v: &cudarc::driver::CudaView<u8>,
19583 o: &mut CudaSlice<f32>,
19584 head_dim: usize,
19585 n_head: usize,
19586 n_head_kv: usize,
19587 base_dev: &CudaSlice<i32>,
19588 t_kv_upper: usize,
19589 t: usize,
19590 scale: f32,
19591 k_tok_bytes: usize,
19592 v_tok_bytes: usize,
19593 base_plus: i32,
19594 g: bool,
19595 ) -> Result<(), Box<dyn std::error::Error>> {
19596 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
19597 assert!(
19598 v4 || fa_v3_active(head_dim),
19599 "stream fa rows requires the v3 or v4 lane"
19600 );
19601 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
19602 if v4 {
19603 let sp = fa_split_keys(t_kv_upper, n_head_kv);
19604 let n_splits_max = (t_kv_upper + sp - 1) / sp;
19605 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19606 let (nspm, spk) = (n_splits_max as i32, sp as i32);
19607 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19608 let gqa = (n_head / n_head_kv).max(1) as u32;
19609 let o_len = t * n_head * n_splits_max * head_dim;
19610 let ml_len = t * n_head * n_splits_max;
19611 let mut part_guard = self.fa_part_pool.lock().unwrap();
19612 if part_guard
19613 .as_ref()
19614 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19615 .unwrap_or(true)
19616 {
19617 let old = part_guard.take();
19628 let (co, cm) = old
19629 .as_ref()
19630 .map(|pp| (pp.0.len(), pp.1.len()))
19631 .unwrap_or((0, 0));
19632 if let Some(old) = old {
19633 self.fa_part_retired.lock().unwrap().push(old);
19634 }
19635 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19636 eprintln!(
19637 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19638 co, o_len, cm, ml_len
19639 );
19640 }
19641 *part_guard = Some((
19642 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19643 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19644 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19645 ));
19646 }
19647 let pg = part_guard.as_mut().unwrap();
19648 self.gpu
19649 .stream()
19650 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19651 self.gpu
19652 .stream()
19653 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19654 self.gpu
19655 .stream()
19656 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19657 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19658 let f = if g {
19659 self.func_g("fa_decode_vec_q_rows_v4_dc")
19660 } else {
19661 self.func("fa_decode_vec_q_rows_v4_dc")
19662 };
19663 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19664 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19665 f.set_attribute(
19666 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19667 sh as i32,
19668 )?;
19669 let cfg = LaunchConfig {
19670 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19671 block_dim: (32, gqa, 1),
19672 shared_mem_bytes: sh,
19673 };
19674 let __s_b = self.gpu.stream();
19675 let mut b = __s_b.launch_builder(&f);
19676 b.arg(q)
19677 .arg(k)
19678 .arg(v)
19679 .arg(&mut *part_o)
19680 .arg(&mut *part_m)
19681 .arg(&mut *part_l)
19682 .arg(&hd)
19683 .arg(&nh)
19684 .arg(&nhkv)
19685 .arg(base_dev)
19686 .arg(&base_plus)
19687 .arg(&scale)
19688 .arg(&nspm)
19689 .arg(&spk)
19690 .arg(&ktb)
19691 .arg(&vtb);
19692 unsafe {
19693 b.launch(cfg)?;
19694 }
19695 let fc = self.func("fa_decode_combine_rows_dc");
19696 let cfg2 = LaunchConfig {
19697 grid_dim: (n_head as u32, t as u32, 1),
19698 block_dim: (head_dim as u32, 1, 1),
19699 shared_mem_bytes: 0,
19700 };
19701 let __s_b2 = self.gpu.stream();
19702 let mut b2 = __s_b2.launch_builder(&fc);
19703 b2.arg(&*part_o)
19704 .arg(&*part_m)
19705 .arg(&*part_l)
19706 .arg(o)
19707 .arg(&hd)
19708 .arg(&nh)
19709 .arg(base_dev)
19710 .arg(&base_plus)
19711 .arg(&nspm)
19712 .arg(&spk);
19713 unsafe {
19714 b2.launch(cfg2)?;
19715 }
19716 return Ok(());
19717 }
19718 let sp = fa_split_keys(t_kv_upper, n_head_kv);
19719 let n_splits_max = (t_kv_upper + sp - 1) / sp;
19720 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19721 let (nspm, spk) = (n_splits_max as i32, sp as i32);
19722 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19723 let gqa = (n_head / n_head_kv).max(1) as u32;
19724 let o_len = t * n_head * n_splits_max * head_dim;
19725 let ml_len = t * n_head * n_splits_max;
19726 let mut part_guard = self.fa_part_pool.lock().unwrap();
19727 if part_guard
19728 .as_ref()
19729 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19730 .unwrap_or(true)
19731 {
19732 let old = part_guard.take();
19743 let (co, cm) = old
19744 .as_ref()
19745 .map(|pp| (pp.0.len(), pp.1.len()))
19746 .unwrap_or((0, 0));
19747 if let Some(old) = old {
19748 self.fa_part_retired.lock().unwrap().push(old);
19749 }
19750 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19751 eprintln!(
19752 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19753 co, o_len, cm, ml_len
19754 );
19755 }
19756 *part_guard = Some((
19757 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19758 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19759 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19760 ));
19761 }
19762 let pg = part_guard.as_mut().unwrap();
19763 self.gpu
19764 .stream()
19765 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19766 self.gpu
19767 .stream()
19768 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19769 self.gpu
19770 .stream()
19771 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19772 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19773 let f = self.func("fa_decode_vec_q_rows_v3_dc");
19774 let sh = (32 * head_dim * 2) as u32;
19775 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19776 f.set_attribute(
19777 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19778 sh as i32,
19779 )?;
19780 let cfg = LaunchConfig {
19781 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19782 block_dim: (32, gqa, 1),
19783 shared_mem_bytes: sh,
19784 };
19785 let __s_b = self.gpu.stream();
19786 let mut b = __s_b.launch_builder(&f);
19787 b.arg(q)
19788 .arg(k)
19789 .arg(v)
19790 .arg(&mut *part_o)
19791 .arg(&mut *part_m)
19792 .arg(&mut *part_l)
19793 .arg(&hd)
19794 .arg(&nh)
19795 .arg(&nhkv)
19796 .arg(base_dev)
19797 .arg(&scale)
19798 .arg(&nspm)
19799 .arg(&spk)
19800 .arg(&ktb)
19801 .arg(&vtb);
19802 unsafe {
19803 b.launch(cfg)?;
19804 }
19805 let fc = self.func("fa_decode_combine_rows_dc");
19806 let cfg2 = LaunchConfig {
19807 grid_dim: (n_head as u32, t as u32, 1),
19808 block_dim: (head_dim as u32, 1, 1),
19809 shared_mem_bytes: 0,
19810 };
19811 let plus0 = 0i32;
19812 let __s_b2 = self.gpu.stream();
19813 let mut b2 = __s_b2.launch_builder(&fc);
19814 b2.arg(&*part_o)
19815 .arg(&*part_m)
19816 .arg(&*part_l)
19817 .arg(o)
19818 .arg(&hd)
19819 .arg(&nh)
19820 .arg(base_dev)
19821 .arg(&plus0)
19822 .arg(&nspm)
19823 .arg(&spk);
19824 unsafe {
19825 b2.launch(cfg2)?;
19826 }
19827 Ok(())
19828 }
19829
19830 pub fn fa_decode_dc(
19841 &self,
19842 q: &CudaSlice<f32>,
19843 k: &cudarc::driver::CudaView<u8>,
19844 v: &cudarc::driver::CudaView<u8>,
19845 o: &mut CudaSlice<f32>,
19846 head_dim: usize,
19847 n_head: usize,
19848 n_head_kv: usize,
19849 t_kv_dev: &CudaSlice<i32>,
19850 bucket_max: usize,
19851 scale: f32,
19852 k_tok_bytes: usize,
19853 v_tok_bytes: usize,
19854 g: bool,
19855 ) -> Result<(), Box<dyn std::error::Error>> {
19856 self.fa_decode_dc_q8(
19857 q,
19858 k,
19859 v,
19860 o,
19861 head_dim,
19862 n_head,
19863 n_head_kv,
19864 t_kv_dev,
19865 bucket_max,
19866 scale,
19867 k_tok_bytes,
19868 v_tok_bytes,
19869 g,
19870 None,
19871 )
19872 }
19873
19874 #[allow(clippy::too_many_arguments)]
19877 pub fn fa_decode_dc_q8(
19878 &self,
19879 q: &CudaSlice<f32>,
19880 k: &cudarc::driver::CudaView<u8>,
19881 v: &cudarc::driver::CudaView<u8>,
19882 o: &mut CudaSlice<f32>,
19883 head_dim: usize,
19884 n_head: usize,
19885 n_head_kv: usize,
19886 t_kv_dev: &CudaSlice<i32>,
19887 bucket_max: usize,
19888 scale: f32,
19889 k_tok_bytes: usize,
19890 v_tok_bytes: usize,
19891 g: bool,
19892 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19893 ) -> Result<(), Box<dyn std::error::Error>> {
19894 let mut fa_vec =
19902 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19903 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19904 fa_vec = false;
19905 } let sp = fa_split_keys(bucket_max, n_head_kv);
19907 let n_splits = if fa_vec {
19908 ((bucket_max + sp - 1) / sp).max(1)
19909 } else {
19910 ((bucket_max + 255) / 256).max(1)
19911 };
19912 let o_len = n_head * n_splits * head_dim;
19913 let ml_len = n_head * n_splits;
19914 let mut part_guard = self.fa_part_pool.lock().unwrap();
19915 if part_guard
19916 .as_ref()
19917 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19918 .unwrap_or(true)
19919 {
19920 let old = part_guard.take();
19931 let (co, cm) = old
19932 .as_ref()
19933 .map(|pp| (pp.0.len(), pp.1.len()))
19934 .unwrap_or((0, 0));
19935 if let Some(old) = old {
19936 self.fa_part_retired.lock().unwrap().push(old);
19937 }
19938 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19939 eprintln!(
19940 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19941 co, o_len, cm, ml_len
19942 );
19943 }
19944 *part_guard = Some((
19945 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19946 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19947 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19948 ));
19949 }
19950 let pg = part_guard.as_mut().unwrap();
19951 self.gpu
19952 .stream()
19953 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19954 self.gpu
19955 .stream()
19956 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19957 self.gpu
19958 .stream()
19959 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19960 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19961 let (hd, nh, nhkv, nsp) = (
19962 head_dim as i32,
19963 n_head as i32,
19964 n_head_kv as i32,
19965 n_splits as i32,
19966 );
19967 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19968 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
19969 let deep = fa_vec
19972 && head_dim == 256
19973 && fa_v4_at(bucket_max)
19974 && !g
19975 && fa_deep_at(bucket_max)
19976 && !matches!(fa_v4_mode(), "noB3" | "stage");
19977 let (f, cfg) = if fa_vec
19978 && head_dim == 512
19979 && bucket_max >= {
19980 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19981 *FA512_MIN_DC.get_or_init(|| {
19982 std::env::var("MEMRA_FA512_MIN")
19983 .ok()
19984 .and_then(|v| v.parse().ok())
19985 .unwrap_or(512)
19986 })
19987 } {
19988 let gqa = (n_head / n_head_kv).max(1) as u32;
19990 (
19991 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
19992 LaunchConfig {
19993 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19994 block_dim: (32, gqa, 1),
19995 shared_mem_bytes: 0,
19996 },
19997 )
19998 } else if fa_vec && head_dim == 512 {
19999 let q_view = q.as_view();
20002 let mut o_view = o.as_view_mut();
20003 return self.fa_decode_scalar_unified(
20004 &q_view,
20005 k,
20006 v,
20007 &mut o_view,
20008 head_dim,
20009 n_head,
20010 n_head_kv,
20011 0,
20012 Some(t_kv_dev),
20013 scale,
20014 n_splits,
20015 sp,
20016 k_tok_bytes,
20017 v_tok_bytes,
20018 g,
20019 &mut *part_o,
20020 &mut *part_m,
20021 &mut *part_l,
20022 q8_out,
20023 );
20024 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
20025 let gqa = (n_head / n_head_kv).max(1) as u32;
20028 let fv = if g {
20029 self.func_g("fa_decode_vec_q_v4_dc")
20030 } else if deep {
20031 self.func("fa_decode_vec_q_v4_deep_dc")
20032 } else {
20033 self.func("fa_decode_vec_q_v4_dc")
20034 };
20035 let shmem =
20036 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
20037 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20038 fv.set_attribute(
20039 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20040 shmem as i32,
20041 )?;
20042 (
20043 fv,
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: shmem,
20048 },
20049 )
20050 } else if fa_vec && fa_v3_active(head_dim) {
20051 let gqa = (n_head / n_head_kv).max(1) as u32;
20054 let fv = if g {
20055 self.func_g("fa_decode_vec_q_v3_dc")
20056 } else {
20057 self.func("fa_decode_vec_q_v3_dc")
20058 };
20059 let shmem = (32 * head_dim * 2) as u32; (
20061 fv,
20062 LaunchConfig {
20063 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20064 block_dim: (32, gqa, 1),
20065 shared_mem_bytes: shmem,
20066 },
20067 )
20068 } else if fa_vec && fa_v2_on() {
20069 let gqa = (n_head / n_head_kv).max(1) as u32;
20073 let fv = if g {
20074 self.func_g("fa_decode_vec_q_v2_dc")
20075 } else {
20076 self.func("fa_decode_vec_q_v2_dc")
20077 };
20078 let shmem = (2 * 32 * head_dim * 2) as u32; (
20080 fv,
20081 LaunchConfig {
20082 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20083 block_dim: (32, gqa, 1),
20084 shared_mem_bytes: shmem,
20085 },
20086 )
20087 } else if fa_vec {
20088 let gqa = (n_head / n_head_kv).max(1) as u32;
20089 let fv = if g {
20091 self.func_g("fa_decode_vec_q_dc")
20092 } else {
20093 self.func("fa_decode_vec_q_dc")
20094 };
20095 (
20096 fv,
20097 LaunchConfig {
20098 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20099 block_dim: (32, gqa, 1),
20100 shared_mem_bytes: 0,
20101 },
20102 )
20103 } else {
20104 let q_view = q.as_view();
20105 let mut o_view = o.as_view_mut();
20106 return self.fa_decode_scalar_unified(
20107 &q_view,
20108 k,
20109 v,
20110 &mut o_view,
20111 head_dim,
20112 n_head,
20113 n_head_kv,
20114 0,
20115 Some(t_kv_dev),
20116 scale,
20117 n_splits,
20118 if fa_vec { sp } else { 256 },
20119 k_tok_bytes,
20120 v_tok_bytes,
20121 g,
20122 &mut *part_o,
20123 &mut *part_m,
20124 &mut *part_l,
20125 q8_out,
20126 );
20127 };
20128 let ski = sp as i32; let __s_b = self.gpu.stream();
20130 let mut b = __s_b.launch_builder(&f);
20131 b.arg(q)
20132 .arg(k)
20133 .arg(v)
20134 .arg(&mut *part_o)
20135 .arg(&mut *part_m)
20136 .arg(&mut *part_l)
20137 .arg(&hd)
20138 .arg(&nh)
20139 .arg(&nhkv)
20140 .arg(t_kv_dev)
20141 .arg(&scale)
20142 .arg(&nsp)
20143 .arg(&ski)
20144 .arg(&ktb)
20145 .arg(&vtb);
20146 unsafe {
20147 b.launch(cfg)?;
20148 }
20149 let cfg2 = LaunchConfig {
20150 grid_dim: (n_head as u32, 1, 1),
20151 block_dim: (head_dim as u32, 1, 1),
20152 shared_mem_bytes: 0,
20153 };
20154 if let Some((oq, od)) = q8_out {
20155 let fc = if g {
20156 self.func_g("fa_decode_combine_q8_1")
20157 } else {
20158 self.fa_func("fa_decode_combine_q8_1", head_dim)
20159 };
20160 let __s_b2 = self.gpu.stream();
20161 let mut b2 = __s_b2.launch_builder(&fc);
20162 b2.arg(&*part_o)
20163 .arg(&*part_m)
20164 .arg(&*part_l)
20165 .arg(oq)
20166 .arg(od)
20167 .arg(&hd)
20168 .arg(&nh)
20169 .arg(&nsp);
20170 unsafe {
20171 b2.launch(cfg2)?;
20172 }
20173 return Ok(());
20174 }
20175 let fc = if g {
20176 self.func_g("fa_decode_combine_f32")
20177 } else {
20178 self.fa_func("fa_decode_combine_f32", head_dim)
20179 };
20180 let __s_b2 = self.gpu.stream();
20181 let mut b2 = __s_b2.launch_builder(&fc);
20182 b2.arg(&*part_o)
20183 .arg(&*part_m)
20184 .arg(&*part_l)
20185 .arg(o)
20186 .arg(&hd)
20187 .arg(&nh)
20188 .arg(&nsp);
20189 unsafe {
20190 b2.launch(cfg2)?;
20191 }
20192 Ok(())
20193 }
20194
20195 pub fn fa_geom_eager(
20201 &self,
20202 t_kv: usize,
20203 head_dim: usize,
20204 n_head_kv: usize,
20205 g: bool,
20206 ) -> (bool, usize) {
20207 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20211 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
20217 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
20218 if g && head_dim == 256 && !fa_v4_at(t_kv) {
20224 fa_vec = false;
20225 }
20226 let sp = fa_split_keys(t_kv, n_head_kv);
20227 let n_splits = if fa_vec {
20228 ((t_kv + sp - 1) / sp).max(1)
20229 } else {
20230 ((t_kv + 255) / 256).max(1)
20231 };
20232 (fa_vec, n_splits)
20233 }
20234
20235 pub fn fa_bucket_key(
20241 &self,
20242 t_kv: usize,
20243 head_dim: usize,
20244 n_head_kv: usize,
20245 g: bool,
20246 ) -> (bool, usize) {
20247 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
20248 }
20249
20250 pub fn capture_graph_retained<F>(
20262 &self,
20263 step: F,
20264 ) -> Result<
20265 (
20266 cudarc::driver::CudaGraph,
20267 Vec<Box<dyn std::any::Any + Send>>,
20268 ),
20269 Box<dyn std::error::Error>,
20270 >
20271 where
20272 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20273 {
20274 use cudarc::driver::sys::CUgraphInstantiate_flags;
20275 self.capture_graph_retained_flags(
20276 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
20277 step,
20278 )
20279 }
20280
20281 pub fn capture_graph_retained_flags<F>(
20286 &self,
20287 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
20288 mut step: F,
20289 ) -> Result<
20290 (
20291 cudarc::driver::CudaGraph,
20292 Vec<Box<dyn std::any::Any + Send>>,
20293 ),
20294 Box<dyn std::error::Error>,
20295 >
20296 where
20297 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20298 {
20299 use cudarc::driver::sys::CUstreamCaptureMode;
20300 self.capture_keep.lock().unwrap().clear();
20308 let was_tracking = self.gpu.ctx.is_event_tracking();
20309 if was_tracking {
20310 unsafe {
20311 self.gpu.ctx.disable_event_tracking();
20312 }
20313 }
20314 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20315 self.capture_keep_on
20316 .store(true, std::sync::atomic::Ordering::Relaxed);
20317 let w = (|| {
20318 step(self)?;
20319 step(self)
20320 })();
20321 self.capture_keep_on
20322 .store(false, std::sync::atomic::Ordering::Relaxed);
20323 w?;
20324 self.gpu.stream().synchronize()?;
20325 self.gpu
20326 .stream()
20327 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20328 let r = step(self);
20329 let g = self.gpu.stream().end_capture(flags);
20330 r?;
20331 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20332 graph.upload()?;
20333 Ok(graph)
20334 };
20335 let result = run();
20336 self.capture_keep_on
20337 .store(false, std::sync::atomic::Ordering::Relaxed);
20338 if was_tracking {
20339 unsafe {
20340 self.gpu.ctx.enable_event_tracking();
20341 }
20342 }
20343 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
20344 Ok((result?, keeper))
20345 }
20346
20347 pub fn capture_graph<F>(
20348 &self,
20349 mut step: F,
20350 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
20351 where
20352 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
20353 {
20354 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
20355 let was_tracking = self.gpu.ctx.is_event_tracking();
20363 if was_tracking {
20364 unsafe {
20365 self.gpu.ctx.disable_event_tracking();
20366 }
20367 }
20368 let iflag = {
20375 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
20376 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
20377 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
20380 Ok("priority") => {
20381 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
20382 }
20383 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
20384 })
20385 };
20386 let ct = {
20393 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20394 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
20395 };
20396 let warmups = {
20419 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20420 *W.get_or_init(|| {
20421 std::env::var("MEMRA_GRAPH_WARMUPS")
20422 .ok()
20423 .and_then(|v| v.parse().ok())
20424 .filter(|n| *n >= 1)
20425 .unwrap_or(1)
20426 })
20427 };
20428 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20429 let t_w = std::time::Instant::now();
20430 for _ in 0..warmups {
20432 step(self)?;
20433 }
20434 self.gpu.stream().synchronize()?;
20435 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
20436 let t_c = std::time::Instant::now();
20438 self.gpu
20439 .stream()
20440 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20441 let r = step(self);
20444 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
20445 let t_i = std::time::Instant::now();
20446 let g = self.gpu.stream().end_capture(iflag);
20447 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
20448 r?;
20449 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20450 let t_u = std::time::Instant::now();
20451 graph.upload()?;
20452 if ct {
20453 println!(
20454 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
20455 instantiate {ms_inst:.2} ms upload {:.2} ms",
20456 t_u.elapsed().as_secs_f64() * 1e3
20457 );
20458 }
20459 Ok(graph)
20460 };
20461 let result = run();
20462 if was_tracking {
20463 unsafe {
20464 self.gpu.ctx.enable_event_tracking();
20465 }
20466 }
20467 result
20468 }
20469
20470 pub fn gdn_scan_s128_view(
20472 &self,
20473 q: &CudaSlice<f32>,
20474 k: &CudaSlice<f32>,
20475 v: &CudaSlice<f32>,
20476 g: &CudaSlice<f32>,
20477 beta: &CudaSlice<f32>,
20478 state_in: &cudarc::driver::CudaView<f32>,
20479 state_out: &mut cudarc::driver::CudaViewMut<f32>,
20480 o: &mut CudaSlice<f32>,
20481 n_head: usize,
20482 t: usize,
20483 scale: f32,
20484 ) -> Result<(), Box<dyn std::error::Error>> {
20485 let f = self.func("gdn_scan_s128");
20486 const S_V: u32 = 128;
20487 const WARP: u32 = 32;
20488 const COLS: u32 = 4;
20489 let cfg = LaunchConfig {
20490 grid_dim: (n_head as u32, 1, S_V / COLS),
20491 block_dim: (WARP, COLS, 1),
20492 shared_mem_bytes: 0,
20493 };
20494 let (h, ti) = (n_head as i32, t as i32);
20495 let __s_b = self.gpu.stream();
20496 let mut b = __s_b.launch_builder(&f);
20497 b.arg(q)
20498 .arg(k)
20499 .arg(v)
20500 .arg(g)
20501 .arg(beta)
20502 .arg(state_in)
20503 .arg(state_out)
20504 .arg(o)
20505 .arg(&h)
20506 .arg(&ti)
20507 .arg(&scale);
20508 unsafe {
20509 b.launch(cfg)?;
20510 }
20511 Ok(())
20512 }
20513
20514 pub fn ssm_conv1d_view(
20516 &self,
20517 x: &cudarc::driver::CudaView<f32>,
20518 w: &CudaSlice<f32>,
20519 y: &mut CudaSlice<f32>,
20520 conv_dim: usize,
20521 t: usize,
20522 d_conv: usize,
20523 silu: bool,
20524 ) -> Result<(), Box<dyn std::error::Error>> {
20525 let f = self.func("ssm_conv1d_silu_f32");
20526 let cfg = LaunchConfig {
20528 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20529 block_dim: (256, 1, 1),
20530 shared_mem_bytes: 0,
20531 };
20532 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20533 let __s_b = self.gpu.stream();
20534 let mut b = __s_b.launch_builder(&f);
20535 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20536 unsafe {
20537 b.launch(cfg)?;
20538 }
20539 Ok(())
20540 }
20541
20542 pub fn ssm_conv1d_tm(
20549 &self,
20550 qkv_tm: &CudaSlice<f32>,
20551 w: &CudaSlice<f32>,
20552 y: &mut CudaSlice<f32>,
20553 conv_dim: usize,
20554 t: usize,
20555 d_conv: usize,
20556 ) -> Result<(), Box<dyn std::error::Error>> {
20557 let f = self.func("ssm_conv1d_tm_f32");
20558 let cfg = LaunchConfig {
20559 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20560 block_dim: (256, 1, 1),
20561 shared_mem_bytes: 0,
20562 };
20563 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20564 let __s_b = self.gpu.stream();
20565 let mut b = __s_b.launch_builder(&f);
20566 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
20567 unsafe {
20568 b.launch(cfg)?;
20569 }
20570 Ok(())
20571 }
20572
20573 pub fn ssm_conv1d_tm_state(
20581 &self,
20582 qkv_tm: &CudaSlice<f32>,
20583 conv_state: &mut CudaSlice<f32>,
20584 w: &CudaSlice<f32>,
20585 y: &mut CudaSlice<f32>,
20586 conv_dim: usize,
20587 t: usize,
20588 d_conv: usize,
20589 ) -> Result<(), Box<dyn std::error::Error>> {
20590 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
20591 }
20592
20593 #[allow(clippy::too_many_arguments)]
20596 pub fn ssm_conv1d_tm_state_pad(
20597 &self,
20598 qkv_tm: &CudaSlice<f32>,
20599 conv_state: &mut CudaSlice<f32>,
20600 w: &CudaSlice<f32>,
20601 y: &mut CudaSlice<f32>,
20602 conv_dim: usize,
20603 t: usize,
20604 d_conv: usize,
20605 pad_len: Option<&CudaSlice<i32>>,
20606 ) -> Result<(), Box<dyn std::error::Error>> {
20607 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20608 let ring_old = if t < d_conv - 1 {
20612 Some(self.clone_dtod(conv_state)?)
20613 } else {
20614 None
20615 };
20616 {
20617 let f = self.func("ssm_conv1d_tm_state_f32");
20618 let cfg = LaunchConfig {
20619 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20620 block_dim: (256, 1, 1),
20621 shared_mem_bytes: 0,
20622 };
20623 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20624 let __s_b = self.gpu.stream();
20625 let mut b = __s_b.launch_builder(&f);
20626 b.arg(qkv_tm)
20627 .arg(&*conv_state)
20628 .arg(w)
20629 .arg(y)
20630 .arg(&cd)
20631 .arg(&ti)
20632 .arg(&dc);
20633 unsafe {
20634 b.launch(cfg)?;
20635 }
20636 }
20637 match (ring_old, pad_len) {
20638 (None, Some(len_d)) => {
20639 let f = self.func("ssm_conv_ring_update_dev_f32");
20640 let n = conv_dim * (d_conv - 1);
20641 let cfg = LaunchConfig::for_num_elems(n as u32);
20642 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20643 let __s_b = self.gpu.stream();
20644 let mut b = __s_b.launch_builder(&f);
20645 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20646 unsafe {
20647 b.launch(cfg)?;
20648 }
20649 }
20650 (None, None) => {
20651 let f = self.func("ssm_conv_ring_update_f32");
20652 let n = conv_dim * (d_conv - 1);
20653 let cfg = LaunchConfig::for_num_elems(n as u32);
20654 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20655 let __s_b = self.gpu.stream();
20656 let mut b = __s_b.launch_builder(&f);
20657 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20658 unsafe {
20659 b.launch(cfg)?;
20660 }
20661 }
20662 (Some(old), _) => {
20663 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
20664 }
20665 }
20666 Ok(())
20667 }
20668
20669 pub fn ssm_conv1d_tm_state_pad_v(
20671 &self,
20672 qkv_tm: &cudarc::driver::CudaView<f32>,
20673 conv_state: &mut CudaSlice<f32>,
20674 w: &CudaSlice<f32>,
20675 y: &mut CudaSlice<f32>,
20676 conv_dim: usize,
20677 t: usize,
20678 d_conv: usize,
20679 pad_len: Option<&CudaSlice<i32>>,
20680 ) -> Result<(), Box<dyn std::error::Error>> {
20681 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20682 let ring_old = if t < d_conv - 1 {
20686 Some(self.clone_dtod(conv_state)?)
20687 } else {
20688 None
20689 };
20690 {
20691 let f = self.func("ssm_conv1d_tm_state_f32");
20692 let cfg = LaunchConfig {
20693 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20694 block_dim: (256, 1, 1),
20695 shared_mem_bytes: 0,
20696 };
20697 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20698 let __s_b = self.gpu.stream();
20699 let mut b = __s_b.launch_builder(&f);
20700 b.arg(qkv_tm)
20701 .arg(&*conv_state)
20702 .arg(w)
20703 .arg(y)
20704 .arg(&cd)
20705 .arg(&ti)
20706 .arg(&dc);
20707 unsafe {
20708 b.launch(cfg)?;
20709 }
20710 }
20711 match (ring_old, pad_len) {
20712 (None, Some(len_d)) => {
20713 let f = self.func("ssm_conv_ring_update_dev_f32");
20714 let n = conv_dim * (d_conv - 1);
20715 let cfg = LaunchConfig::for_num_elems(n as u32);
20716 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20717 let __s_b = self.gpu.stream();
20718 let mut b = __s_b.launch_builder(&f);
20719 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20720 unsafe {
20721 b.launch(cfg)?;
20722 }
20723 }
20724 (None, None) => {
20725 let f = self.func("ssm_conv_ring_update_f32");
20726 let n = conv_dim * (d_conv - 1);
20727 let cfg = LaunchConfig::for_num_elems(n as u32);
20728 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20729 let __s_b = self.gpu.stream();
20730 let mut b = __s_b.launch_builder(&f);
20731 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20732 unsafe {
20733 b.launch(cfg)?;
20734 }
20735 }
20736 (Some(_), _) => unreachable!(
20737 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
20738 ),
20739 }
20740 Ok(())
20741 }
20742
20743 pub fn ssm_conv_ring_rebuild(
20748 &self,
20749 qkv_tm: &CudaSlice<f32>,
20750 ring_old: &CudaSlice<f32>,
20751 conv_state: &mut CudaSlice<f32>,
20752 conv_dim: usize,
20753 tc: usize,
20754 d_conv: usize,
20755 ) -> Result<(), Box<dyn std::error::Error>> {
20756 let f = self.func("ssm_conv_ring_rebuild_f32");
20757 let n = conv_dim * (d_conv - 1);
20758 let cfg = LaunchConfig::for_num_elems(n as u32);
20759 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
20760 let __s_b = self.gpu.stream();
20761 let mut b = __s_b.launch_builder(&f);
20762 b.arg(qkv_tm)
20763 .arg(ring_old)
20764 .arg(conv_state)
20765 .arg(&cd)
20766 .arg(&ti)
20767 .arg(&dc);
20768 unsafe {
20769 b.launch(cfg)?;
20770 }
20771 Ok(())
20772 }
20773
20774 #[allow(clippy::too_many_arguments)]
20779 pub fn gdn_prep_decode(
20780 &self,
20781 conv_out: &CudaSlice<f32>,
20782 beta_raw: &CudaSlice<f32>,
20783 alpha: &CudaSlice<f32>,
20784 dt_bias: &CudaSlice<f32>,
20785 a: &CudaSlice<f32>,
20786 q_l2: &mut CudaSlice<f32>,
20787 k_l2: &mut CudaSlice<f32>,
20788 v_g: &mut CudaSlice<f32>,
20789 beta: &mut CudaSlice<f32>,
20790 g_log: &mut CudaSlice<f32>,
20791 d_state: usize,
20792 num_v: usize,
20793 num_k: usize,
20794 key_dim: usize,
20795 eps: f32,
20796 ) -> Result<(), Box<dyn std::error::Error>> {
20797 let f = self.func("gdn_prep_decode_f32");
20798 let cfg = LaunchConfig {
20799 grid_dim: (num_v as u32, 1, 1),
20800 block_dim: (32, 4, 1),
20801 shared_mem_bytes: 0,
20802 };
20803 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20804 let __s_b = self.gpu.stream();
20805 let mut b = __s_b.launch_builder(&f);
20806 b.arg(conv_out)
20807 .arg(beta_raw)
20808 .arg(alpha)
20809 .arg(dt_bias)
20810 .arg(a)
20811 .arg(q_l2)
20812 .arg(k_l2)
20813 .arg(v_g)
20814 .arg(beta)
20815 .arg(g_log)
20816 .arg(&ds)
20817 .arg(&nv)
20818 .arg(&nk)
20819 .arg(&kd)
20820 .arg(&eps);
20821 unsafe {
20822 b.launch(cfg)?;
20823 }
20824 Ok(())
20825 }
20826
20827 #[allow(clippy::too_many_arguments)]
20831 pub fn ssm_conv1d_gdn(
20832 &self,
20833 qkv_tm: &CudaSlice<f32>,
20834 w: &CudaSlice<f32>,
20835 q_g: &mut CudaSlice<f32>,
20836 k_g: &mut CudaSlice<f32>,
20837 v_g: &mut CudaSlice<f32>,
20838 conv_dim: usize,
20839 t: usize,
20840 d_conv: usize,
20841 d_state: usize,
20842 num_v: usize,
20843 num_k: usize,
20844 key_dim: usize,
20845 ) -> Result<(), Box<dyn std::error::Error>> {
20846 let f = self.func("ssm_conv1d_gdn_f32");
20847 let cfg = LaunchConfig {
20848 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20849 block_dim: (256, 1, 1),
20850 shared_mem_bytes: 0,
20851 };
20852 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20853 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20854 let __s_b = self.gpu.stream();
20855 let mut b = __s_b.launch_builder(&f);
20856 b.arg(qkv_tm)
20857 .arg(w)
20858 .arg(q_g)
20859 .arg(k_g)
20860 .arg(v_g)
20861 .arg(&cd)
20862 .arg(&ti)
20863 .arg(&dc)
20864 .arg(&ds)
20865 .arg(&nv)
20866 .arg(&nk)
20867 .arg(&kd);
20868 unsafe {
20869 b.launch(cfg)?;
20870 }
20871 Ok(())
20872 }
20873
20874 pub fn ssm_conv1d(
20875 &self,
20876 x: &CudaSlice<f32>,
20877 w: &CudaSlice<f32>,
20878 y: &mut CudaSlice<f32>,
20879 conv_dim: usize,
20880 t: usize,
20881 d_conv: usize,
20882 silu: bool,
20883 ) -> Result<(), Box<dyn std::error::Error>> {
20884 let f = self.func("ssm_conv1d_silu_f32");
20885 let cfg = LaunchConfig {
20886 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20887 block_dim: (256, 1, 1),
20888 shared_mem_bytes: 0,
20889 };
20890 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20891 let __s_b = self.gpu.stream();
20892 let mut b = __s_b.launch_builder(&f);
20893 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20894 unsafe {
20895 b.launch(cfg)?;
20896 }
20897 Ok(())
20898 }
20899
20900 pub fn gdn_scan_s128(
20903 &self,
20904 q: &CudaSlice<f32>,
20905 k: &CudaSlice<f32>,
20906 v: &CudaSlice<f32>,
20907 g: &CudaSlice<f32>,
20908 beta: &CudaSlice<f32>,
20909 state_in: &CudaSlice<f32>,
20910 state_out: &mut CudaSlice<f32>,
20911 o: &mut CudaSlice<f32>,
20912 n_head: usize,
20913 t: usize,
20914 scale: f32,
20915 ) -> Result<(), Box<dyn std::error::Error>> {
20916 let f = self.func("gdn_scan_s128");
20917 const S_V: u32 = 128;
20918 const WARP: u32 = 32;
20919 const COLS_PER_BLOCK: u32 = 4;
20920 let cfg = LaunchConfig {
20921 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20922 block_dim: (WARP, COLS_PER_BLOCK, 1),
20923 shared_mem_bytes: 0,
20924 };
20925 let (h, ti) = (n_head as i32, t as i32);
20926 let __s_b = self.gpu.stream();
20927 let mut b = __s_b.launch_builder(&f);
20928 b.arg(q)
20929 .arg(k)
20930 .arg(v)
20931 .arg(g)
20932 .arg(beta)
20933 .arg(state_in)
20934 .arg(state_out)
20935 .arg(o)
20936 .arg(&h)
20937 .arg(&ti)
20938 .arg(&scale);
20939 unsafe {
20940 b.launch(cfg)?;
20941 }
20942 Ok(())
20943 }
20944
20945 #[allow(clippy::too_many_arguments)]
20950 pub fn ssm_conv1d_fused_decode_b(
20951 &self,
20952 qkv_cols: &CudaSlice<f32>,
20953 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20954 w: &CudaSlice<f32>,
20955 conv_outs: &mut CudaSlice<f32>,
20956 conv_dim: usize,
20957 d_conv: usize,
20958 b_n: usize,
20959 ) -> Result<(), Box<dyn std::error::Error>> {
20960 let f = self.func("ssm_conv1d_fused_decode_b_f32");
20961 let cfg = LaunchConfig {
20962 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20963 block_dim: (256, 1, 1),
20964 shared_mem_bytes: 0,
20965 };
20966 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20967 let __s_b = self.gpu.stream();
20968 let mut b = __s_b.launch_builder(&f);
20969 b.arg(qkv_cols)
20970 .arg(conv_state_ptrs)
20971 .arg(w)
20972 .arg(conv_outs)
20973 .arg(&cd)
20974 .arg(&dc);
20975 unsafe {
20976 b.launch(cfg)?;
20977 }
20978 Ok(())
20979 }
20980
20981 #[allow(clippy::too_many_arguments)]
20982 pub fn gdn_prep_decode_b(
20983 &self,
20984 conv_outs: &CudaSlice<f32>,
20985 beta_raws: &CudaSlice<f32>,
20986 alphas: &CudaSlice<f32>,
20987 dt_bias: &CudaSlice<f32>,
20988 a: &CudaSlice<f32>,
20989 q_l2: &mut CudaSlice<f32>,
20990 k_l2: &mut CudaSlice<f32>,
20991 v_g: &mut CudaSlice<f32>,
20992 beta: &mut CudaSlice<f32>,
20993 g_log: &mut CudaSlice<f32>,
20994 d_state: usize,
20995 num_v: usize,
20996 num_k: usize,
20997 key_dim: usize,
20998 eps: f32,
20999 conv_dim: usize,
21000 b_n: usize,
21001 ) -> Result<(), Box<dyn std::error::Error>> {
21002 let f = self.func("gdn_prep_decode_b_f32");
21003 let cfg = LaunchConfig {
21004 grid_dim: (num_v as u32, 1, b_n as u32),
21005 block_dim: (32, 4, 1),
21006 shared_mem_bytes: 0,
21007 };
21008 let (ds, nv, nk, kd, cd) = (
21009 d_state as i32,
21010 num_v as i32,
21011 num_k as i32,
21012 key_dim as i32,
21013 conv_dim as i32,
21014 );
21015 let __s_b = self.gpu.stream();
21016 let mut b = __s_b.launch_builder(&f);
21017 b.arg(conv_outs)
21018 .arg(beta_raws)
21019 .arg(alphas)
21020 .arg(dt_bias)
21021 .arg(a)
21022 .arg(q_l2)
21023 .arg(k_l2)
21024 .arg(v_g)
21025 .arg(beta)
21026 .arg(g_log)
21027 .arg(&ds)
21028 .arg(&nv)
21029 .arg(&nk)
21030 .arg(&kd)
21031 .arg(&eps)
21032 .arg(&cd);
21033 unsafe {
21034 b.launch(cfg)?;
21035 }
21036 Ok(())
21037 }
21038
21039 #[allow(clippy::too_many_arguments)]
21040 pub fn gdn_scan_s128_batched(
21041 &self,
21042 q: &CudaSlice<f32>,
21043 k: &CudaSlice<f32>,
21044 v: &CudaSlice<f32>,
21045 g: &CudaSlice<f32>,
21046 beta: &CudaSlice<f32>,
21047 state_in_ptrs: &cudarc::driver::CudaView<u64>,
21048 state_out_ptrs: &cudarc::driver::CudaView<u64>,
21049 o: &mut CudaSlice<f32>,
21050 n_head: usize,
21051 b_n: usize,
21052 scale: f32,
21053 ) -> Result<(), Box<dyn std::error::Error>> {
21054 let f = self.func("gdn_scan_s128_b");
21055 const S_V: u32 = 128;
21056 const WARP: u32 = 32;
21057 const COLS_PER_BLOCK: u32 = 4;
21058 let cfg = LaunchConfig {
21059 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
21060 block_dim: (WARP, COLS_PER_BLOCK, 1),
21061 shared_mem_bytes: 0,
21062 };
21063 let h = n_head as i32;
21064 let __s_b = self.gpu.stream();
21065 let mut b = __s_b.launch_builder(&f);
21066 b.arg(q)
21067 .arg(k)
21068 .arg(v)
21069 .arg(g)
21070 .arg(beta)
21071 .arg(state_in_ptrs)
21072 .arg(state_out_ptrs)
21073 .arg(o)
21074 .arg(&h)
21075 .arg(&scale);
21076 unsafe {
21077 b.launch(cfg)?;
21078 }
21079 Ok(())
21080 }
21081
21082 #[allow(clippy::too_many_arguments)]
21088 pub fn ssm_conv1d_fused_decode_b_view(
21089 &self,
21090 qkv_cols: &cudarc::driver::CudaView<f32>,
21091 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
21092 w: &CudaSlice<f32>,
21093 conv_outs: &mut CudaSlice<f32>,
21094 conv_dim: usize,
21095 d_conv: usize,
21096 b_n: usize,
21097 ) -> Result<(), Box<dyn std::error::Error>> {
21098 let f = self.func("ssm_conv1d_fused_decode_b_f32");
21099 let cfg = LaunchConfig {
21100 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
21101 block_dim: (256, 1, 1),
21102 shared_mem_bytes: 0,
21103 };
21104 let (cd, dc) = (conv_dim as i32, d_conv as i32);
21105 let __s_b = self.gpu.stream();
21106 let mut b = __s_b.launch_builder(&f);
21107 b.arg(qkv_cols)
21108 .arg(conv_state_ptrs)
21109 .arg(w)
21110 .arg(conv_outs)
21111 .arg(&cd)
21112 .arg(&dc);
21113 unsafe {
21114 b.launch(cfg)?;
21115 }
21116 Ok(())
21117 }
21118
21119 #[allow(clippy::too_many_arguments)]
21120 pub fn gdn_prep_decode_b_view(
21121 &self,
21122 conv_outs: &CudaSlice<f32>,
21123 beta_raws: &cudarc::driver::CudaView<f32>,
21124 alphas: &cudarc::driver::CudaView<f32>,
21125 dt_bias: &CudaSlice<f32>,
21126 a: &CudaSlice<f32>,
21127 q_l2: &mut CudaSlice<f32>,
21128 k_l2: &mut CudaSlice<f32>,
21129 v_g: &mut CudaSlice<f32>,
21130 beta: &mut CudaSlice<f32>,
21131 g_log: &mut CudaSlice<f32>,
21132 d_state: usize,
21133 num_v: usize,
21134 num_k: usize,
21135 key_dim: usize,
21136 eps: f32,
21137 conv_dim: usize,
21138 b_n: usize,
21139 ) -> Result<(), Box<dyn std::error::Error>> {
21140 let f = self.func("gdn_prep_decode_b_f32");
21141 let cfg = LaunchConfig {
21142 grid_dim: (num_v as u32, 1, b_n as u32),
21143 block_dim: (32, 4, 1),
21144 shared_mem_bytes: 0,
21145 };
21146 let (ds, nv, nk, kd, cd) = (
21147 d_state as i32,
21148 num_v as i32,
21149 num_k as i32,
21150 key_dim as i32,
21151 conv_dim as i32,
21152 );
21153 let __s_b = self.gpu.stream();
21154 let mut b = __s_b.launch_builder(&f);
21155 b.arg(conv_outs)
21156 .arg(beta_raws)
21157 .arg(alphas)
21158 .arg(dt_bias)
21159 .arg(a)
21160 .arg(q_l2)
21161 .arg(k_l2)
21162 .arg(v_g)
21163 .arg(beta)
21164 .arg(g_log)
21165 .arg(&ds)
21166 .arg(&nv)
21167 .arg(&nk)
21168 .arg(&kd)
21169 .arg(&eps)
21170 .arg(&cd);
21171 unsafe {
21172 b.launch(cfg)?;
21173 }
21174 Ok(())
21175 }
21176
21177 #[allow(clippy::too_many_arguments)]
21178 pub fn gdn_scan_s128_batched_view(
21179 &self,
21180 q: &CudaSlice<f32>,
21181 k: &CudaSlice<f32>,
21182 v: &CudaSlice<f32>,
21183 g: &CudaSlice<f32>,
21184 beta: &CudaSlice<f32>,
21185 state_in_ptrs: &cudarc::driver::CudaView<u64>,
21186 state_out_ptrs: &cudarc::driver::CudaView<u64>,
21187 o: &mut cudarc::driver::CudaViewMut<f32>,
21188 n_head: usize,
21189 b_n: usize,
21190 scale: f32,
21191 ) -> Result<(), Box<dyn std::error::Error>> {
21192 let f = self.func("gdn_scan_s128_b");
21193 const S_V: u32 = 128;
21194 const WARP: u32 = 32;
21195 const COLS_PER_BLOCK: u32 = 4;
21196 let cfg = LaunchConfig {
21197 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
21198 block_dim: (WARP, COLS_PER_BLOCK, 1),
21199 shared_mem_bytes: 0,
21200 };
21201 let h = n_head as i32;
21202 let __s_b = self.gpu.stream();
21203 let mut b = __s_b.launch_builder(&f);
21204 b.arg(q)
21205 .arg(k)
21206 .arg(v)
21207 .arg(g)
21208 .arg(beta)
21209 .arg(state_in_ptrs)
21210 .arg(state_out_ptrs)
21211 .arg(o)
21212 .arg(&h)
21213 .arg(&scale);
21214 unsafe {
21215 b.launch(cfg)?;
21216 }
21217 Ok(())
21218 }
21219
21220 pub fn gdn_chunked_enabled() -> bool {
21229 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21230 *E.get_or_init(|| {
21231 std::env::var("MEMRA_GDN_CHUNKED")
21232 .map(|v| v != "0")
21233 .unwrap_or(true)
21234 })
21235 }
21236
21237 pub fn gdn_chunk_size() -> usize {
21242 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21243 *C.get_or_init(|| {
21244 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
21245 .ok()
21246 .and_then(|v| v.parse().ok())
21247 .unwrap_or(32);
21248 c.clamp(32, 128) / 32 * 32
21249 })
21250 }
21251
21252 #[allow(clippy::too_many_arguments)]
21257 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
21260 #[allow(clippy::too_many_arguments)]
21261 pub fn gdn_chunk_k123(
21262 &self,
21263 q: &CudaSlice<f32>,
21264 k: &CudaSlice<f32>,
21265 v: &CudaSlice<f32>,
21266 g: &CudaSlice<f32>,
21267 beta: &CudaSlice<f32>,
21268 wb16: Option<&mut CudaSlice<u8>>,
21269 n_head: usize,
21270 t: usize,
21271 c: usize,
21272 hk: usize,
21273 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
21274 ) -> Result<
21275 (
21276 CudaSlice<f32>,
21277 CudaSlice<f32>,
21278 CudaSlice<f32>,
21279 CudaSlice<f32>,
21280 ),
21281 Box<dyn std::error::Error>,
21282 > {
21283 const D: usize = 128;
21284 let h = n_head;
21285 let nc = (t + c - 1) / c;
21286 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21287 let mut gcum = self.uninit(t * h)?;
21288 let mut a = self.uninit(nc * h * c * c)?;
21289 let mut p = self.uninit(nc * h * c * c)?;
21290 let mut u = self.uninit(nc * h * c * D)?;
21291 let mut w = self.uninit(nc * h * c * D)?;
21292 {
21293 let f = self.func("gdn_chunk_cumgate_f32");
21295 let cfg = LaunchConfig {
21296 grid_dim: (nc as u32, h as u32, 1),
21297 block_dim: (32, 1, 1),
21298 shared_mem_bytes: 0,
21299 };
21300 let __s_b = self.gpu.stream();
21301 let mut b = __s_b.launch_builder(&f);
21302 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
21303 unsafe {
21304 b.launch(cfg)?;
21305 }
21306 }
21307 if let Some((qb, kb, pb)) = k2w {
21308 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
21311 let f = self.func("gdn_k2_wgmma");
21312 let cfg = LaunchConfig {
21313 grid_dim: (nc as u32, h as u32, 1),
21314 block_dim: (128, 1, 1),
21315 shared_mem_bytes: 0,
21316 };
21317 let hki = hk as i32;
21318 let __s_b = self.gpu.stream();
21319 let mut b = __s_b.launch_builder(&f);
21320 b.arg(qb)
21321 .arg(kb)
21322 .arg(&gcum)
21323 .arg(beta)
21324 .arg(&mut a)
21325 .arg(&mut *pb)
21326 .arg(&hi)
21327 .arg(&ti)
21328 .arg(&ci)
21329 .arg(&hki);
21330 unsafe {
21331 b.launch(cfg)?;
21332 }
21333 } else if c <= 64 && !portable_mma_gated() {
21334 let f = self.func("gdn_chunk_attn_f32");
21336 let jt = ((c + 31) / 32) as u32;
21337 let cfg = LaunchConfig {
21338 grid_dim: (nc as u32, h as u32, jt),
21339 block_dim: (256, 1, 1),
21340 shared_mem_bytes: 0,
21341 };
21342 let hki = hk as i32;
21343 let __s_b = self.gpu.stream();
21344 let mut b = __s_b.launch_builder(&f);
21345 b.arg(q)
21346 .arg(k)
21347 .arg(&gcum)
21348 .arg(beta)
21349 .arg(&mut a)
21350 .arg(&mut p)
21351 .arg(&hi)
21352 .arg(&ti)
21353 .arg(&ci)
21354 .arg(&hki);
21355 unsafe {
21356 b.launch(cfg)?;
21357 }
21358 } else {
21359 assert!(
21361 hk == h,
21362 "generic K2 is broadcast-only (de-broadcast rides C==32)"
21363 );
21364 let f = self.func("gdn_chunk_attn_g_f32");
21365 let cfg = LaunchConfig {
21366 grid_dim: (nc as u32, h as u32, 1),
21367 block_dim: (32, 8, 1),
21368 shared_mem_bytes: 0,
21369 };
21370 let __s_b = self.gpu.stream();
21371 let mut b = __s_b.launch_builder(&f);
21372 b.arg(q)
21373 .arg(k)
21374 .arg(&gcum)
21375 .arg(beta)
21376 .arg(&mut a)
21377 .arg(&mut p)
21378 .arg(&hi)
21379 .arg(&ti)
21380 .arg(&ci);
21381 unsafe {
21382 b.launch(cfg)?;
21383 }
21384 }
21385 {
21386 let cfg = LaunchConfig {
21388 grid_dim: (nc as u32, h as u32, 1),
21389 block_dim: (256, 1, 1),
21390 shared_mem_bytes: 0,
21391 };
21392 match c {
21393 32 | 64 => {
21394 let f = self.func(if c == 32 {
21395 "gdn_chunk_solve32_f32"
21396 } else {
21397 "gdn_chunk_solve64_f32"
21398 });
21399 let wb: u64 = match wb16 {
21401 Some(d) => self.addr_u8(d),
21402 None => 0,
21403 };
21404 let hki = hk as i32;
21405 let __s_b = self.gpu.stream();
21406 let mut b = __s_b.launch_builder(&f);
21407 b.arg(v)
21408 .arg(k)
21409 .arg(&a)
21410 .arg(&gcum)
21411 .arg(&mut u)
21412 .arg(&mut w)
21413 .arg(&wb)
21414 .arg(&hi)
21415 .arg(&ti)
21416 .arg(&hki);
21417 unsafe {
21418 b.launch(cfg)?;
21419 }
21420 }
21421 _ => {
21422 assert!(hk == h, "generic K3 is broadcast-only");
21423 let f = self.func("gdn_chunk_solve_f32");
21424 let __s_b = self.gpu.stream();
21425 let mut b = __s_b.launch_builder(&f);
21426 b.arg(v)
21427 .arg(k)
21428 .arg(&a)
21429 .arg(&gcum)
21430 .arg(&mut u)
21431 .arg(&mut w)
21432 .arg(&hi)
21433 .arg(&ti)
21434 .arg(&ci);
21435 unsafe {
21436 b.launch(cfg)?;
21437 }
21438 }
21439 }
21440 }
21441 Ok((gcum, p, u, w))
21442 }
21443
21444 pub fn gdn_db_on() -> bool {
21448 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
21449 }
21450
21451 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
21460 !portable_mma_gated()
21461 && c == 32
21462 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21463 Ok("1") => true,
21464 Ok("0") => false,
21465 _ => gdn_mma_default_on(),
21466 }
21467 }
21468
21469 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
21476 cfg!(memra_hopper_mma)
21477 && self.gdn_mma_enabled(c)
21478 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
21479 }
21480
21481 #[allow(clippy::too_many_arguments)]
21486 pub fn ssm_conv1d_gdn_state_pad(
21487 &self,
21488 qkv_tm: &cudarc::driver::CudaView<f32>,
21489 conv_state: &mut CudaSlice<f32>,
21490 w: &CudaSlice<f32>,
21491 q_g: &mut CudaSlice<f32>,
21492 k_g: &mut CudaSlice<f32>,
21493 v_g: &mut CudaSlice<f32>,
21494 conv_dim: usize,
21495 t: usize,
21496 d_conv: usize,
21497 d_state: usize,
21498 num_v: usize,
21499 num_k: usize,
21500 key_dim: usize,
21501 hk: usize,
21502 pad_len: Option<&CudaSlice<i32>>,
21503 ) -> Result<(), Box<dyn std::error::Error>> {
21504 assert!(
21505 t >= d_conv - 1,
21506 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
21507 );
21508 {
21509 let f = self.func("ssm_conv1d_gdn_state_f32");
21510 let cfg = LaunchConfig {
21511 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
21512 block_dim: (256, 1, 1),
21513 shared_mem_bytes: 0,
21514 };
21515 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21516 let (ds, nv, nk, kd, hki) = (
21517 d_state as i32,
21518 num_v as i32,
21519 num_k as i32,
21520 key_dim as i32,
21521 hk as i32,
21522 );
21523 let __s_b = self.gpu.stream();
21524 let mut b = __s_b.launch_builder(&f);
21525 b.arg(qkv_tm)
21526 .arg(&*conv_state)
21527 .arg(w)
21528 .arg(q_g)
21529 .arg(k_g)
21530 .arg(v_g)
21531 .arg(&cd)
21532 .arg(&ti)
21533 .arg(&dc)
21534 .arg(&ds)
21535 .arg(&nv)
21536 .arg(&nk)
21537 .arg(&kd)
21538 .arg(&hki);
21539 unsafe {
21540 b.launch(cfg)?;
21541 }
21542 }
21543 match pad_len {
21544 Some(len_d) => {
21545 let f = self.func("ssm_conv_ring_update_dev_f32");
21546 let n = conv_dim * (d_conv - 1);
21547 let cfg = LaunchConfig::for_num_elems(n as u32);
21548 let (cd, dc) = (conv_dim as i32, d_conv as i32);
21549 let __s_b = self.gpu.stream();
21550 let mut b = __s_b.launch_builder(&f);
21551 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
21552 unsafe {
21553 b.launch(cfg)?;
21554 }
21555 }
21556 None => {
21557 let f = self.func("ssm_conv_ring_update_f32");
21558 let n = conv_dim * (d_conv - 1);
21559 let cfg = LaunchConfig::for_num_elems(n as u32);
21560 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21561 let __s_b = self.gpu.stream();
21562 let mut b = __s_b.launch_builder(&f);
21563 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
21564 unsafe {
21565 b.launch(cfg)?;
21566 }
21567 }
21568 }
21569 Ok(())
21570 }
21571
21572 pub fn gdn_chunk_alloc(
21576 &self,
21577 n_head: usize,
21578 t: usize,
21579 c: usize,
21580 hk: usize,
21581 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
21582 const D: usize = 128;
21583 assert!(
21584 c == 32,
21585 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
21586 );
21587 let h = n_head;
21588 let nc = (t + c - 1) / c;
21589 Ok(GdnChunkBufs {
21590 gcum: self.uninit(t * h)?,
21591 a: self.uninit(nc * h * c * c)?,
21592 p: self.uninit(nc * h * c * c)?,
21593 u: self.uninit(nc * h * c * D)?,
21594 w: self.uninit(nc * h * c * D)?,
21595 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21596 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21597 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21598 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
21599 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21600 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
21601 o: self.uninit(D * h * t)?,
21602 t,
21603 nc,
21604 })
21605 }
21606
21607 pub fn f32_to_bf16_v(
21609 &self,
21610 x: &cudarc::driver::CudaView<f32>,
21611 dst: &mut CudaSlice<u8>,
21612 n: usize,
21613 ) -> Result<(), Box<dyn std::error::Error>> {
21614 let f = self.func("f32_to_bf16_bulk");
21615 let ni = n as i64;
21616 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21617 let __s_b = self.gpu.stream();
21618 let mut b = __s_b.launch_builder(&f);
21619 b.arg(x).arg(dst).arg(&ni);
21620 unsafe {
21621 b.launch(cfg)?;
21622 }
21623 Ok(())
21624 }
21625
21626 pub fn f32_to_bf16_into(
21628 &self,
21629 x: &CudaSlice<f32>,
21630 dst: &mut CudaSlice<u8>,
21631 n: usize,
21632 ) -> Result<(), Box<dyn std::error::Error>> {
21633 let f = self.func("f32_to_bf16_bulk");
21634 let ni = n as i64;
21635 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21636 let __s_b = self.gpu.stream();
21637 let mut b = __s_b.launch_builder(&f);
21638 b.arg(x).arg(dst).arg(&ni);
21639 unsafe {
21640 b.launch(cfg)?;
21641 }
21642 Ok(())
21643 }
21644
21645 pub fn gdn_chunk_k123_vl8(
21648 &self,
21649 seqs: &[GdnSeqVl],
21650 n_head: usize,
21651 hk: usize,
21652 wq: Option<&GdnWVl8>,
21653 ) -> Result<(), Box<dyn std::error::Error>> {
21654 let b = seqs.len();
21655 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
21656 let mut packed = [GdnSeqVl::default(); 8];
21657 packed[..b].copy_from_slice(seqs);
21658 let v = GdnVl8(packed);
21659 let (hi, ci) = (n_head as i32, 32i32);
21660 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21661 {
21662 let f = self.func("gdn_chunk_cumgate_vl");
21663 let cfg = LaunchConfig {
21664 grid_dim: (max_nc, n_head as u32, b as u32),
21665 block_dim: (32, 1, 1),
21666 shared_mem_bytes: 0,
21667 };
21668 let __s_lb = self.gpu.stream();
21669 let mut lb = __s_lb.launch_builder(&f);
21670 lb.arg(&v).arg(&hi).arg(&ci);
21671 unsafe {
21672 lb.launch(cfg)?;
21673 }
21674 }
21675 let hki = hk as i32;
21676 if let Some(w) = wq {
21677 let f = self.func("gdn_k2_wgmma_vl");
21679 let cfg = LaunchConfig {
21680 grid_dim: (max_nc, n_head as u32, b as u32),
21681 block_dim: (128, 1, 1),
21682 shared_mem_bytes: 0,
21683 };
21684 let __s_lb = self.gpu.stream();
21685 let mut lb = __s_lb.launch_builder(&f);
21686 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
21687 unsafe {
21688 lb.launch(cfg)?;
21689 }
21690 } else {
21691 let f = self.func("gdn_chunk_attn_vl");
21692 let cfg = LaunchConfig {
21693 grid_dim: (max_nc, n_head as u32, b as u32),
21694 block_dim: (256, 1, 1),
21695 shared_mem_bytes: 0,
21696 };
21697 let __s_lb = self.gpu.stream();
21698 let mut lb = __s_lb.launch_builder(&f);
21699 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21700 unsafe {
21701 lb.launch(cfg)?;
21702 }
21703 }
21704 {
21705 let f = self.func("gdn_chunk_solve32_vl");
21706 let cfg = LaunchConfig {
21707 grid_dim: (max_nc, n_head as u32, b as u32),
21708 block_dim: (256, 1, 1),
21709 shared_mem_bytes: 0,
21710 };
21711 let __s_lb = self.gpu.stream();
21712 let mut lb = __s_lb.launch_builder(&f);
21713 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21714 unsafe {
21715 lb.launch(cfg)?;
21716 }
21717 }
21718 Ok(())
21719 }
21720
21721 #[allow(clippy::too_many_arguments)]
21725 pub fn gdn_prep_vl8(
21726 &self,
21727 seqs: &[GdnPrepVl],
21728 conv_w: &CudaSlice<f32>,
21729 dt_bias: &CudaSlice<f32>,
21730 a: &CudaSlice<f32>,
21731 conv_dim: usize,
21732 d_conv: usize,
21733 d_state: usize,
21734 num_v: usize,
21735 num_k: usize,
21736 key_dim: usize,
21737 hk: usize,
21738 eps: f32,
21739 ) -> Result<(), Box<dyn std::error::Error>> {
21740 let b = seqs.len();
21741 assert!(b >= 1 && b <= 8);
21742 let mut packed = [GdnPrepVl::default(); 8];
21743 packed[..b].copy_from_slice(seqs);
21744 let v = GdnPrepVl8(packed);
21745 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21746 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
21747 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
21748 assert!(
21749 conv_fuse || hk == num_v,
21750 "de-broadcast requires the fused conv"
21751 );
21752 if conv_fuse {
21753 let f = self.func("ssm_conv1d_gdn_state_vl");
21754 let cfg = LaunchConfig {
21755 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21756 block_dim: (256, 1, 1),
21757 shared_mem_bytes: 0,
21758 };
21759 let (dsi, nvi, nki, kdi, hki) = (
21760 d_state as i32,
21761 num_v as i32,
21762 num_k as i32,
21763 key_dim as i32,
21764 hk as i32,
21765 );
21766 let __s_lb = self.gpu.stream();
21767 let mut lb = __s_lb.launch_builder(&f);
21768 lb.arg(&v)
21769 .arg(conv_w)
21770 .arg(&cdi)
21771 .arg(&dci)
21772 .arg(&dsi)
21773 .arg(&nvi)
21774 .arg(&nki)
21775 .arg(&kdi)
21776 .arg(&hki);
21777 unsafe {
21778 lb.launch(cfg)?;
21779 }
21780 } else {
21781 let f = self.func("ssm_conv1d_tm_state_vl");
21782 let cfg = LaunchConfig {
21783 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21784 block_dim: (256, 1, 1),
21785 shared_mem_bytes: 0,
21786 };
21787 let __s_lb = self.gpu.stream();
21788 let mut lb = __s_lb.launch_builder(&f);
21789 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21790 unsafe {
21791 lb.launch(cfg)?;
21792 }
21793 }
21794 {
21795 let f = self.func("ssm_conv_ring_update_vl");
21796 let n = (conv_dim * (d_conv - 1)) as u32;
21797 let cfg = LaunchConfig {
21798 grid_dim: (n.div_ceil(256), 1, b as u32),
21799 block_dim: (256, 1, 1),
21800 shared_mem_bytes: 0,
21801 };
21802 let __s_lb = self.gpu.stream();
21803 let mut lb = __s_lb.launch_builder(&f);
21804 lb.arg(&v).arg(&cdi).arg(&dci);
21805 unsafe {
21806 lb.launch(cfg)?;
21807 }
21808 }
21809 if !conv_fuse {
21810 let f = self.func("qkv_to_gdn_repack_vl");
21811 let n = max_t * (num_v * d_state) as u32;
21812 let cfg = LaunchConfig {
21813 grid_dim: (n.div_ceil(256), 1, b as u32),
21814 block_dim: (256, 1, 1),
21815 shared_mem_bytes: 0,
21816 };
21817 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21818 let __s_lb = self.gpu.stream();
21819 let mut lb = __s_lb.launch_builder(&f);
21820 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21821 unsafe {
21822 lb.launch(cfg)?;
21823 }
21824 }
21825 if Self::l2_v2_on(d_state) {
21826 let f = self.func("gdn_l2_v2_vl");
21827 let cfg = LaunchConfig {
21828 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21829 block_dim: (256, 1, 1),
21830 shared_mem_bytes: 0,
21831 };
21832 let (dsi, nvi) = (d_state as i32, hk as i32);
21833 let __s_lb = self.gpu.stream();
21834 let mut lb = __s_lb.launch_builder(&f);
21835 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21836 unsafe {
21837 lb.launch(cfg)?;
21838 }
21839 } else {
21840 let f = self.func("gdn_l2_vl");
21841 let cfg = LaunchConfig {
21842 grid_dim: (max_t * hk as u32, 2, b as u32),
21843 block_dim: (256, 1, 1),
21844 shared_mem_bytes: 0,
21845 };
21846 let (dsi, nvi) = (d_state as i32, hk as i32);
21847 let __s_lb = self.gpu.stream();
21848 let mut lb = __s_lb.launch_builder(&f);
21849 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21850 unsafe {
21851 lb.launch(cfg)?;
21852 }
21853 }
21854 {
21855 let f = self.func("gdn_gate_prep_vl");
21856 let n = max_t * num_v as u32;
21857 let cfg = LaunchConfig {
21858 grid_dim: (n.div_ceil(256), 1, b as u32),
21859 block_dim: (256, 1, 1),
21860 shared_mem_bytes: 0,
21861 };
21862 let nvi = num_v as i32;
21863 let __s_lb = self.gpu.stream();
21864 let mut lb = __s_lb.launch_builder(&f);
21865 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21866 unsafe {
21867 lb.launch(cfg)?;
21868 }
21869 }
21870 Ok(())
21871 }
21872
21873 pub fn gdn_mirror_vl8(
21875 &self,
21876 seqs: &[GdnSeqVl],
21877 n_head: usize,
21878 which: i32,
21879 hk: usize,
21880 ) -> Result<(), Box<dyn std::error::Error>> {
21881 let b = seqs.len();
21882 assert!(b >= 1 && b <= 8);
21883 let mut packed = [GdnSeqVl::default(); 8];
21884 packed[..b].copy_from_slice(seqs);
21885 let v = GdnVl8(packed);
21886 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21887 let max_n = seqs
21888 .iter()
21889 .map(|s| {
21890 if which == 0 {
21891 s.t as i64 * ept as i64
21892 } else {
21893 s.nc as i64 * ept as i64 * 32
21894 }
21895 })
21896 .max()
21897 .unwrap();
21898 let f = self.func("gdn_mirror_vl");
21899 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21900 let cfg = LaunchConfig {
21901 grid_dim: (blocks, 1, b as u32),
21902 block_dim: (256, 1, 1),
21903 shared_mem_bytes: 0,
21904 };
21905 let __s_lb = self.gpu.stream();
21906 let mut lb = __s_lb.launch_builder(&f);
21907 lb.arg(&v).arg(&ept).arg(&which);
21908 unsafe {
21909 lb.launch(cfg)?;
21910 }
21911 Ok(())
21912 }
21913
21914 pub fn gdn_tail_vl8(
21916 &self,
21917 seqs: &[GdnPrepVl],
21918 norm_w: &CudaSlice<f32>,
21919 d_state: usize,
21920 num_v: usize,
21921 eps: f32,
21922 ) -> Result<(), Box<dyn std::error::Error>> {
21923 let b = seqs.len();
21924 assert!(b >= 1 && b <= 8);
21925 let mut packed = [GdnPrepVl::default(); 8];
21926 packed[..b].copy_from_slice(seqs);
21927 let v = GdnPrepVl8(packed);
21928 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21929 let f = self.func("gated_rmsnorm_f16out_vl");
21930 let cfg = LaunchConfig {
21932 grid_dim: (max_t * num_v as u32, 1, b as u32),
21933 block_dim: (128, 1, 1),
21934 shared_mem_bytes: 0,
21935 };
21936 let (dsi, nvi) = (d_state as i32, num_v as i32);
21937 let __s_lb = self.gpu.stream();
21938 let mut lb = __s_lb.launch_builder(&f);
21939 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21940 unsafe {
21941 lb.launch(cfg)?;
21942 }
21943 Ok(())
21944 }
21945
21946 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
21949 use cudarc::driver::DevicePtr;
21950 let s = self.gpu.stream();
21951 let (p, _g) = x.device_ptr(&s);
21952 p as u64
21953 }
21954 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
21955 use cudarc::driver::DevicePtrMut;
21956 let s = self.gpu.stream();
21957 let (p, _g) = x.device_ptr_mut(&s);
21958 p as u64
21959 }
21960 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
21961 use cudarc::driver::DevicePtr;
21962 let s = self.gpu.stream();
21963 let (p, _g) = x.device_ptr(&s);
21964 p as u64
21965 }
21966 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
21967 use cudarc::driver::DevicePtr;
21968 let s = self.gpu.stream();
21969 let (p, _g) = x.device_ptr(&s);
21970 p as u64
21971 }
21972
21973 pub fn gdn_chunk_vl8(
21977 &self,
21978 seqs: &[GdnSeqVl],
21979 n_head: usize,
21980 scale: f32,
21981 hk: usize,
21982 wq: Option<&GdnWVl8>,
21983 ) -> Result<(), Box<dyn std::error::Error>> {
21984 const NSPLIT: u32 = 4;
21985 let b = seqs.len();
21986 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
21987 let mut packed = [GdnSeqVl::default(); 8];
21988 packed[..b].copy_from_slice(seqs);
21989 let v = GdnVl8(packed);
21990 let (hi, ci) = (n_head as i32, 32i32);
21991 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21992 let hki = hk as i32;
21993 if let Some(w) = wq {
21994 let f = self.func("gdn_k45_wgmma_vl");
21996 let cfg = LaunchConfig {
21997 grid_dim: (n_head as u32, NSPLIT, b as u32),
21998 block_dim: (256, 1, 1),
21999 shared_mem_bytes: 0,
22000 };
22001 let __s_lb = self.gpu.stream();
22002 let mut lb = __s_lb.launch_builder(&f);
22003 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
22004 unsafe {
22005 lb.launch(cfg)?;
22006 }
22007 let _ = max_nc;
22008 return Ok(());
22009 }
22010 {
22011 let f = self.func("gdn_chunk_state_mma_vl");
22012 let cfg = LaunchConfig {
22013 grid_dim: (n_head as u32, NSPLIT, b as u32),
22014 block_dim: (256, 1, 1),
22015 shared_mem_bytes: 0,
22016 };
22017 let __s_lb = self.gpu.stream();
22018 let mut lb = __s_lb.launch_builder(&f);
22019 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
22020 unsafe {
22021 lb.launch(cfg)?;
22022 }
22023 }
22024 {
22025 let f = self.func("gdn_chunk_output_mma_vl");
22026 let cfg = LaunchConfig {
22027 grid_dim: (max_nc, n_head as u32, b as u32),
22028 block_dim: (256, 1, 1),
22029 shared_mem_bytes: 0,
22030 };
22031 let __s_lb = self.gpu.stream();
22032 let mut lb = __s_lb.launch_builder(&f);
22033 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
22034 unsafe {
22035 lb.launch(cfg)?;
22036 }
22037 }
22038 Ok(())
22039 }
22040 pub fn gdn_scan_chunked(
22041 &self,
22042 q: &CudaSlice<f32>,
22043 k: &CudaSlice<f32>,
22044 v: &CudaSlice<f32>,
22045 g: &CudaSlice<f32>,
22046 beta: &CudaSlice<f32>,
22047 kb16_pre: Option<&CudaSlice<u8>>,
22048 qb16_pre: Option<&CudaSlice<u8>>,
22049 state_in: &CudaSlice<f32>,
22050 state_out: &mut CudaSlice<f32>,
22051 o: &mut CudaSlice<f32>,
22052 n_head: usize,
22053 t: usize,
22054 scale: f32,
22055 c: usize,
22056 hk: usize,
22057 ) -> Result<(), Box<dyn std::error::Error>> {
22058 const D: usize = 128;
22059 const NSPLIT: u32 = 4;
22060 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
22061 let h = n_head;
22062 let nc = (t + c - 1) / c;
22063 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
22064 let gdn_mma_pre = !portable_mma_gated()
22069 && c == 32
22070 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
22071 Ok("1") => true,
22072 Ok("0") => false,
22073 _ => gdn_mma_default_on(),
22074 };
22075 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
22076 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
22077 } else {
22078 None
22079 };
22080 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
22085 && gdn_mma_pre
22086 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
22087 let nk = t * hk * D;
22088 let mut kb16_local: Option<CudaSlice<u8>> = None;
22089 if gdn_mma_pre && kb16_pre.is_none() {
22090 let mut kb = self.alloc_u8_uninit(nk * 2)?;
22091 let f = self.func("f32_to_bf16_bulk");
22092 let n2 = nk as i64;
22093 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
22094 let __s_b = self.gpu.stream();
22095 let mut b = __s_b.launch_builder(&f);
22096 b.arg(k).arg(&mut kb).arg(&n2);
22097 unsafe {
22098 b.launch(cfg2)?;
22099 }
22100 kb16_local = Some(kb);
22101 }
22102 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
22103 if let Some(kb) = kb16_pre {
22104 assert!(kb.len() >= nk * 2, "kb16_pre too small");
22105 }
22106 let mut qb16: Option<CudaSlice<u8>> = None;
22107 let mut pb16: Option<CudaSlice<u8>> = None;
22108 if gdn_wgmma_pre {
22109 if qb16_pre.is_none() {
22112 let mut qb = self.alloc_u8_uninit(nk * 2)?;
22113 let f = self.func("f32_to_bf16_bulk");
22114 let n2 = nk as i64;
22115 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
22116 let __s_b = self.gpu.stream();
22117 let mut b = __s_b.launch_builder(&f);
22118 b.arg(q).arg(&mut qb).arg(&n2);
22119 unsafe {
22120 b.launch(cfg2)?;
22121 }
22122 qb16 = Some(qb);
22123 } else if let Some(qb) = qb16_pre {
22124 assert!(qb.len() >= nk * 2, "qb16_pre too small");
22125 }
22126 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
22127 }
22128 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
22129 let k2w = if gdn_wgmma_pre {
22130 Some((
22131 *qb16_ref0.as_ref().unwrap(),
22132 *kb16_ref0.as_ref().unwrap(),
22133 pb16.as_mut().unwrap(),
22134 ))
22135 } else {
22136 None
22137 };
22138 let (gcum, p, u, w) =
22139 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
22140 let _ = &w;
22141 let mut y = self.uninit(nc * h * c * D)?;
22142 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
22156 && c == 32
22157 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
22158 Ok("1") => true,
22159 Ok("0") => false,
22160 _ => gdn_mma_default_on(),
22161 };
22162 if gdn_mma {
22163 let wb16 = wb16_pre
22164 .take()
22165 .expect("mma path pre-allocates wb16 (K3 store fold)");
22166 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
22167 if gdn_wgmma_pre {
22179 let qb16 = qb16_ref0.unwrap();
22181 let pb16 = pb16.as_ref().unwrap();
22182 {
22183 let f = self.func("gdn_k45_wgmma");
22184 let cfg = LaunchConfig {
22185 grid_dim: (h as u32, 4, 1),
22186 block_dim: (256, 1, 1),
22187 shared_mem_bytes: 0,
22188 };
22189 let hki = hk as i32;
22190 let __s_b = self.gpu.stream();
22191 let mut b = __s_b.launch_builder(&f);
22192 b.arg(kb16_ref)
22193 .arg(&gcum)
22194 .arg(beta)
22195 .arg(&u)
22196 .arg(&wb16)
22197 .arg(qb16)
22198 .arg(pb16)
22199 .arg(o)
22200 .arg(&scale)
22201 .arg(state_in)
22202 .arg(&mut *state_out)
22203 .arg(&hi)
22204 .arg(&ti)
22205 .arg(&ci)
22206 .arg(&hki);
22207 unsafe {
22208 b.launch(cfg)?;
22209 }
22210 }
22211 return Ok(());
22212 }
22213 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
22217 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
22218 {
22219 let f = self.func("gdn_chunk_state_mma");
22220 let cfg = LaunchConfig {
22221 grid_dim: (h as u32, NSPLIT, 1),
22222 block_dim: (256, 1, 1),
22223 shared_mem_bytes: 0,
22224 };
22225 let hki = hk as i32;
22226 let __s_b = self.gpu.stream();
22227 let mut b = __s_b.launch_builder(&f);
22228 b.arg(kb16_ref)
22229 .arg(&gcum)
22230 .arg(beta)
22231 .arg(&u)
22232 .arg(&wb16)
22233 .arg(&mut y16)
22234 .arg(&mut ssnap16)
22235 .arg(state_in)
22236 .arg(&mut *state_out)
22237 .arg(&hi)
22238 .arg(&ti)
22239 .arg(&ci)
22240 .arg(&hki);
22241 unsafe {
22242 b.launch(cfg)?;
22243 }
22244 }
22245 {
22246 let f = self.func("gdn_chunk_output_mma");
22248 let jt = ((c + 31) / 32) as u32;
22249 let cfg = LaunchConfig {
22250 grid_dim: (nc as u32, h as u32, jt),
22251 block_dim: (256, 1, 1),
22252 shared_mem_bytes: 0,
22253 };
22254 let hki = hk as i32;
22255 let __s_b = self.gpu.stream();
22256 let mut b = __s_b.launch_builder(&f);
22257 b.arg(q)
22258 .arg(&gcum)
22259 .arg(&p)
22260 .arg(&y16)
22261 .arg(&ssnap16)
22262 .arg(o)
22263 .arg(&hi)
22264 .arg(&ti)
22265 .arg(&ci)
22266 .arg(&scale)
22267 .arg(&hki);
22268 unsafe {
22269 b.launch(cfg)?;
22270 }
22271 }
22272 return Ok(());
22273 }
22274 {
22275 let f = self.func("gdn_chunk_state_f32");
22277 let cfg = LaunchConfig {
22278 grid_dim: (h as u32, NSPLIT, 1),
22279 block_dim: (256, 1, 1),
22280 shared_mem_bytes: 0,
22281 };
22282 let __s_b = self.gpu.stream();
22283 let mut b = __s_b.launch_builder(&f);
22284 b.arg(k)
22285 .arg(&gcum)
22286 .arg(beta)
22287 .arg(&u)
22288 .arg(&w)
22289 .arg(&mut y)
22290 .arg(&mut ssnap)
22291 .arg(state_in)
22292 .arg(&mut *state_out)
22293 .arg(&hi)
22294 .arg(&ti)
22295 .arg(&ci);
22296 unsafe {
22297 b.launch(cfg)?;
22298 }
22299 }
22300 {
22301 let f = self.func("gdn_chunk_output_f32");
22303 let jt = ((c + 31) / 32) as u32;
22304 let cfg = LaunchConfig {
22305 grid_dim: (nc as u32, h as u32, jt),
22306 block_dim: (256, 1, 1),
22307 shared_mem_bytes: 0,
22308 };
22309 let __s_b = self.gpu.stream();
22310 let mut b = __s_b.launch_builder(&f);
22311 b.arg(q)
22312 .arg(&gcum)
22313 .arg(&p)
22314 .arg(&y)
22315 .arg(&ssnap)
22316 .arg(o)
22317 .arg(&hi)
22318 .arg(&ti)
22319 .arg(&ci)
22320 .arg(&scale);
22321 unsafe {
22322 b.launch(cfg)?;
22323 }
22324 }
22325 Ok(())
22326 }
22327
22328 #[allow(clippy::too_many_arguments)]
22337 #[allow(clippy::too_many_arguments)]
22338 pub fn gdn_scan_prefill(
22339 &self,
22340 q: &CudaSlice<f32>,
22341 k: &CudaSlice<f32>,
22342 v: &CudaSlice<f32>,
22343 g: &CudaSlice<f32>,
22344 beta: &CudaSlice<f32>,
22345 kb16_pre: Option<&CudaSlice<u8>>,
22346 qb16_pre: Option<&CudaSlice<u8>>,
22347 state_in: &CudaSlice<f32>,
22348 state_out: &mut CudaSlice<f32>,
22349 o: &mut CudaSlice<f32>,
22350 n_head: usize,
22351 t: usize,
22352 scale: f32,
22353 hk: usize,
22354 ) -> Result<(), Box<dyn std::error::Error>> {
22355 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
22356 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
22357 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
22358 }
22359 if Self::gdn_chunked_enabled() && t >= 16 {
22360 self.gdn_scan_chunked(
22361 q,
22362 k,
22363 v,
22364 g,
22365 beta,
22366 kb16_pre,
22367 qb16_pre,
22368 state_in,
22369 state_out,
22370 o,
22371 n_head,
22372 t,
22373 scale,
22374 Self::gdn_chunk_size(),
22375 hk,
22376 )
22377 } else {
22378 assert!(
22379 hk == n_head,
22380 "s128 scan is broadcast-only (prep guarantees by predicate)"
22381 );
22382 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
22383 }
22384 }
22385
22386 #[allow(clippy::too_many_arguments)]
22388 fn gdn_scan_diff(
22389 &self,
22390 q: &CudaSlice<f32>,
22391 k: &CudaSlice<f32>,
22392 v: &CudaSlice<f32>,
22393 g: &CudaSlice<f32>,
22394 beta: &CudaSlice<f32>,
22395 state_in: &CudaSlice<f32>,
22396 state_out: &mut CudaSlice<f32>,
22397 o: &mut CudaSlice<f32>,
22398 n_head: usize,
22399 t: usize,
22400 scale: f32,
22401 ) -> Result<(), Box<dyn std::error::Error>> {
22402 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
22403 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
22404 let mut o_c = self.uninit(o.len())?;
22405 let mut st_c = self.uninit(state_out.len())?;
22406 self.gdn_scan_chunked(
22407 q,
22408 k,
22409 v,
22410 g,
22411 beta,
22412 None,
22413 None,
22414 state_in,
22415 &mut st_c,
22416 &mut o_c,
22417 n_head,
22418 t,
22419 scale,
22420 Self::gdn_chunk_size(),
22421 n_head,
22422 )?;
22423 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
22424 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
22425 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
22426 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
22427 let mut max_abs = 0f32;
22428 let mut max_rel = 0f32;
22429 let mut sum_rel = 0f64;
22430 for (x, y) in a.iter().zip(b) {
22431 let ad = (x - y).abs();
22432 let rel = ad / x.abs().max(y.abs()).max(1e-3);
22433 if ad > max_abs {
22434 max_abs = ad;
22435 }
22436 if rel > max_rel {
22437 max_rel = rel;
22438 }
22439 sum_rel += rel as f64;
22440 }
22441 (max_abs, max_rel, sum_rel / a.len() as f64)
22442 };
22443 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
22444 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
22445 println!(
22446 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
22447 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
22448 Self::gdn_chunk_size()
22449 );
22450 Ok(())
22451 }
22452
22453 pub fn gdn_glog(
22455 &self,
22456 alpha: &CudaSlice<f32>,
22457 dt_bias: &CudaSlice<f32>,
22458 a: &CudaSlice<f32>,
22459 g_log: &mut CudaSlice<f32>,
22460 n_head: usize,
22461 t: usize,
22462 ) -> Result<(), Box<dyn std::error::Error>> {
22463 let f = self.func("gdn_glog_f32");
22464 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22465 let (h, ti) = (n_head as i32, t as i32);
22466 let __s_b = self.gpu.stream();
22467 let mut b = __s_b.launch_builder(&f);
22468 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22469 unsafe {
22470 b.launch(cfg)?;
22471 }
22472 Ok(())
22473 }
22474
22475 pub fn sigmoid_v(
22478 &self,
22479 x: &cudarc::driver::CudaView<f32>,
22480 y: &mut CudaSlice<f32>,
22481 n: usize,
22482 ) -> Result<(), Box<dyn std::error::Error>> {
22483 let f = self.func("sigmoid_f32");
22484 let cfg = LaunchConfig::for_num_elems(n as u32);
22485 let ni = n as i32;
22486 let __s_b = self.gpu.stream();
22487 let mut b = __s_b.launch_builder(&f);
22488 b.arg(x).arg(y).arg(&ni);
22489 unsafe {
22490 b.launch(cfg)?;
22491 }
22492 Ok(())
22493 }
22494
22495 pub fn gdn_glog_v(
22496 &self,
22497 alpha: &cudarc::driver::CudaView<f32>,
22498 dt_bias: &CudaSlice<f32>,
22499 a: &CudaSlice<f32>,
22500 g_log: &mut CudaSlice<f32>,
22501 n_head: usize,
22502 t: usize,
22503 ) -> Result<(), Box<dyn std::error::Error>> {
22504 let f = self.func("gdn_glog_f32");
22505 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22506 let (h, ti) = (n_head as i32, t as i32);
22507 let __s_b = self.gpu.stream();
22508 let mut b = __s_b.launch_builder(&f);
22509 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22510 unsafe {
22511 b.launch(cfg)?;
22512 }
22513 Ok(())
22514 }
22515
22516 pub fn sigmoid(
22517 &self,
22518 x: &CudaSlice<f32>,
22519 y: &mut CudaSlice<f32>,
22520 n: usize,
22521 ) -> Result<(), Box<dyn std::error::Error>> {
22522 let f = self.func("sigmoid_f32");
22523 let cfg = LaunchConfig::for_num_elems(n as u32);
22524 let ni = n as i32;
22525 let __s_b = self.gpu.stream();
22526 let mut b = __s_b.launch_builder(&f);
22527 b.arg(x).arg(y).arg(&ni);
22528 unsafe {
22529 b.launch(cfg)?;
22530 }
22531 Ok(())
22532 }
22533
22534 pub fn sig_mul_f16out(
22537 &self,
22538 a: &CudaSlice<f32>,
22539 g: &CudaSlice<f32>,
22540 dst: &mut CudaSlice<f32>,
22541 dst16: &mut CudaSlice<u8>,
22542 n: usize,
22543 ) -> Result<(), Box<dyn std::error::Error>> {
22544 let f = self.func("sig_mul_f16out_f32");
22545 let cfg = LaunchConfig::for_num_elems(n as u32);
22546 let ni = n as i32;
22547 let __s_b = self.gpu.stream();
22548 let mut b = __s_b.launch_builder(&f);
22549 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
22550 unsafe {
22551 b.launch(cfg)?;
22552 }
22553 Ok(())
22554 }
22555
22556 #[allow(clippy::too_many_arguments)]
22565 pub fn attn_head_gate(
22566 &self,
22567 a: &CudaSlice<f32>,
22568 g: &CudaSlice<f32>,
22569 dst: &mut CudaSlice<f32>,
22570 dst16: Option<&mut CudaSlice<u8>>,
22571 head_dim: usize,
22572 n_head: usize,
22573 t: usize,
22574 ) -> Result<(), Box<dyn std::error::Error>> {
22575 let f = self.func("attn_head_gate_f32");
22576 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22577 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22578 let d16: u64 = match dst16 {
22580 Some(d) => self.addr_u8(d),
22581 None => 0,
22582 };
22583 let __s_b = self.gpu.stream();
22584 let mut b = __s_b.launch_builder(&f);
22585 b.arg(a)
22586 .arg(g)
22587 .arg(dst)
22588 .arg(&d16)
22589 .arg(&hd)
22590 .arg(&nh)
22591 .arg(&ti);
22592 unsafe {
22593 b.launch(cfg)?;
22594 }
22595 Ok(())
22596 }
22597
22598 #[allow(clippy::too_many_arguments)]
22607 pub fn swiglu_clamped_mul_scaled(
22608 &self,
22609 gate: &CudaSlice<f32>,
22610 up: &CudaSlice<f32>,
22611 gs: f32,
22612 us: f32,
22613 limit: f32,
22614 dst: &mut CudaSlice<f32>,
22615 n: usize,
22616 ) -> Result<(), Box<dyn std::error::Error>> {
22617 debug_assert!(
22618 limit > 1e-6,
22619 "swiglu_clamped needs a live limit; use silu_mul_scaled"
22620 );
22621 let f = self.func("swiglu_clamped_mul_scaled_f32");
22622 let cfg = LaunchConfig::for_num_elems(n as u32);
22623 let ni = n as i32;
22624 let __s_b = self.gpu.stream();
22625 let mut b = __s_b.launch_builder(&f);
22626 b.arg(gate)
22627 .arg(up)
22628 .arg(&gs)
22629 .arg(&us)
22630 .arg(&limit)
22631 .arg(dst)
22632 .arg(&ni);
22633 unsafe {
22634 b.launch(cfg)?;
22635 }
22636 Ok(())
22637 }
22638
22639 pub fn gated_rmsnorm(
22641 &self,
22642 o: &CudaSlice<f32>,
22643 w: &CudaSlice<f32>,
22644 z: &CudaSlice<f32>,
22645 dst: &mut CudaSlice<f32>,
22646 ncols: usize,
22647 nrows: usize,
22648 eps: f32,
22649 ) -> Result<(), Box<dyn std::error::Error>> {
22650 let f = self.func("gated_rmsnorm_f32");
22651 let cfg = LaunchConfig {
22652 grid_dim: (nrows as u32, 1, 1),
22653 block_dim: (128, 1, 1),
22654 shared_mem_bytes: 0,
22655 };
22656 let (nc, e) = (ncols as i32, eps);
22657 let __s_b = self.gpu.stream();
22658 let mut b = __s_b.launch_builder(&f);
22659 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22660 unsafe {
22661 b.launch(cfg)?;
22662 }
22663 Ok(())
22664 }
22665
22666 pub fn gated_rmsnorm_f16out(
22669 &self,
22670 o: &CudaSlice<f32>,
22671 w: &CudaSlice<f32>,
22672 z: &CudaSlice<f32>,
22673 dst: &mut CudaSlice<f32>,
22674 dst16: &mut CudaSlice<u8>,
22675 ncols: usize,
22676 nrows: usize,
22677 eps: f32,
22678 ) -> Result<(), Box<dyn std::error::Error>> {
22679 let f = self.func("gated_rmsnorm_f16out_f32");
22680 let cfg = LaunchConfig {
22682 grid_dim: (nrows as u32, 1, 1),
22683 block_dim: (128, 1, 1),
22684 shared_mem_bytes: 0,
22685 };
22686 let (nc, e) = (ncols as i32, eps);
22687 let __s_b = self.gpu.stream();
22688 let mut b = __s_b.launch_builder(&f);
22689 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22690 unsafe {
22691 b.launch(cfg)?;
22692 }
22693 Ok(())
22694 }
22695
22696 #[allow(clippy::too_many_arguments)]
22700 pub fn add_rms_norm_zq8(
22701 &self,
22702 a: &CudaSlice<f32>,
22703 b_in: &CudaSlice<f32>,
22704 w: &CudaSlice<f32>,
22705 res: &mut CudaSlice<f32>,
22706 z: &mut CudaSlice<f32>,
22707 ncols: usize,
22708 nrows: usize,
22709 eps: f32,
22710 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22711 assert!(ncols % 32 == 0);
22712 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
22713 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22714 let f = self.func("add_rms_norm_zq8");
22715 let cfg = LaunchConfig {
22716 grid_dim: (nrows as u32, 1, 1),
22717 block_dim: (1024, 1, 1),
22718 shared_mem_bytes: 0,
22719 };
22720 let (nc, ep) = (ncols as i32, eps);
22721 let __s_b = self.gpu.stream();
22722 let mut b = __s_b.launch_builder(&f);
22723 b.arg(a)
22724 .arg(b_in)
22725 .arg(w)
22726 .arg(res)
22727 .arg(z)
22728 .arg(&mut q)
22729 .arg(&mut d)
22730 .arg(&nc)
22731 .arg(&ep);
22732 unsafe {
22733 b.launch(cfg)?;
22734 }
22735 Ok((q, d))
22736 }
22737
22738 pub fn gated_rmsnorm_zv(
22743 &self,
22744 o: &CudaSlice<f32>,
22745 w: &CudaSlice<f32>,
22746 z: &cudarc::driver::CudaView<f32>,
22747 dst: &mut CudaSlice<f32>,
22748 ncols: usize,
22749 nrows: usize,
22750 eps: f32,
22751 ) -> Result<(), Box<dyn std::error::Error>> {
22752 let f = self.func("gated_rmsnorm_f32");
22753 let cfg = LaunchConfig {
22754 grid_dim: (nrows as u32, 1, 1),
22755 block_dim: (128, 1, 1),
22756 shared_mem_bytes: 0,
22757 };
22758 let (nc, e) = (ncols as i32, eps);
22759 let __s_b = self.gpu.stream();
22760 let mut b = __s_b.launch_builder(&f);
22761 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22762 unsafe {
22763 b.launch(cfg)?;
22764 }
22765 Ok(())
22766 }
22767
22768 pub fn gated_rmsnorm_f16out_zv(
22769 &self,
22770 o: &CudaSlice<f32>,
22771 w: &CudaSlice<f32>,
22772 z: &cudarc::driver::CudaView<f32>,
22773 dst: &mut CudaSlice<f32>,
22774 dst16: &mut CudaSlice<u8>,
22775 ncols: usize,
22776 nrows: usize,
22777 eps: f32,
22778 ) -> Result<(), Box<dyn std::error::Error>> {
22779 let f = self.func("gated_rmsnorm_f16out_f32");
22780 let cfg = LaunchConfig {
22782 grid_dim: (nrows as u32, 1, 1),
22783 block_dim: (128, 1, 1),
22784 shared_mem_bytes: 0,
22785 };
22786 let (nc, e) = (ncols as i32, eps);
22787 let __s_b = self.gpu.stream();
22788 let mut b = __s_b.launch_builder(&f);
22789 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22790 unsafe {
22791 b.launch(cfg)?;
22792 }
22793 Ok(())
22794 }
22795
22796 pub fn gated_rmsnorm_q8_1(
22797 &self,
22798 o: &CudaSlice<f32>,
22799 w: &CudaSlice<f32>,
22800 z: &CudaSlice<f32>,
22801 ncols: usize,
22802 nrows: usize,
22803 eps: f32,
22804 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22805 assert!(ncols % 32 == 0);
22806 let f = self.func("gated_rmsnorm_q8_1");
22807 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22808 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22809 let cfg = LaunchConfig {
22810 grid_dim: (nrows as u32, 1, 1),
22811 block_dim: (128, 1, 1),
22812 shared_mem_bytes: 0,
22813 };
22814 let (nc, ep) = (ncols as i32, eps);
22815 let __s_b = self.gpu.stream();
22816 let mut b = __s_b.launch_builder(&f);
22817 b.arg(o)
22818 .arg(w)
22819 .arg(z)
22820 .arg(&mut out_q)
22821 .arg(&mut out_d)
22822 .arg(&nc)
22823 .arg(&ep);
22824 unsafe {
22825 b.launch(cfg)?;
22826 }
22827 Ok((out_q, out_d))
22828 }
22829
22830 pub fn transpose(
22832 &self,
22833 inp: &CudaSlice<f32>,
22834 rows: usize,
22835 cols: usize,
22836 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22837 let f = self.func("transpose_f32");
22838 let mut out = self.zeros(rows * cols)?;
22839 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22840 let (r, c) = (rows as i32, cols as i32);
22841 let __s_b = self.gpu.stream();
22842 let mut b = __s_b.launch_builder(&f);
22843 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22844 unsafe {
22845 b.launch(cfg)?;
22846 }
22847 Ok(out)
22848 }
22849
22850 pub fn repeat_heads(
22852 &self,
22853 inp: &CudaSlice<f32>,
22854 out: &mut CudaSlice<f32>,
22855 head_dim: usize,
22856 n_in: usize,
22857 n_out: usize,
22858 t: usize,
22859 ) -> Result<(), Box<dyn std::error::Error>> {
22860 let f = self.func("repeat_heads_f32");
22861 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22862 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22863 let __s_b = self.gpu.stream();
22864 let mut b = __s_b.launch_builder(&f);
22865 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22866 unsafe {
22867 b.launch(cfg)?;
22868 }
22869 Ok(())
22870 }
22871
22872 pub fn q_gate_split(
22879 &self,
22880 qf: &CudaSlice<f32>,
22881 q_out: &mut CudaSlice<f32>,
22882 gate_out: &mut CudaSlice<f32>,
22883 head_dim: usize,
22884 n_head: usize,
22885 t: usize,
22886 ) -> Result<(), Box<dyn std::error::Error>> {
22887 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
22888 let out_need = head_dim * n_head * t;
22889 if q_out.len() < out_need || gate_out.len() < out_need {
22890 return Err(format!(
22891 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
22892 q_out.len(),
22893 gate_out.len()
22894 )
22895 .into());
22896 }
22897 let f = self.func("q_gate_split_f32");
22898 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22899 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22900 let __s_b = self.gpu.stream();
22901 let mut b = __s_b.launch_builder(&f);
22902 b.arg(qf)
22903 .arg(q_out)
22904 .arg(gate_out)
22905 .arg(&hd)
22906 .arg(&nh)
22907 .arg(&ti);
22908 unsafe {
22909 b.launch(cfg)?;
22910 }
22911 Ok(())
22912 }
22913
22914 pub fn qkv_to_gdn_repack(
22918 &self,
22919 conv_out: &CudaSlice<f32>,
22920 q_g: &mut CudaSlice<f32>,
22921 k_g: &mut CudaSlice<f32>,
22922 v_g: &mut CudaSlice<f32>,
22923 d_state: usize,
22924 num_v: usize,
22925 num_k: usize,
22926 key_dim: usize,
22927 t: usize,
22928 ) -> Result<(), Box<dyn std::error::Error>> {
22929 let f = self.func("qkv_to_gdn_repack_f32");
22930 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22931 let (ds, nv, nk, kd, ti) = (
22932 d_state as i32,
22933 num_v as i32,
22934 num_k as i32,
22935 key_dim as i32,
22936 t as i32,
22937 );
22938 let __s_b = self.gpu.stream();
22939 let mut b = __s_b.launch_builder(&f);
22940 b.arg(conv_out)
22941 .arg(q_g)
22942 .arg(k_g)
22943 .arg(v_g)
22944 .arg(&ds)
22945 .arg(&nv)
22946 .arg(&nk)
22947 .arg(&kd)
22948 .arg(&ti);
22949 unsafe {
22950 b.launch(cfg)?;
22951 }
22952 Ok(())
22953 }
22954
22955 pub fn conv_left_pad(
22958 &self,
22959 src: &CudaSlice<f32>,
22960 dst: &mut CudaSlice<f32>,
22961 conv_dim: usize,
22962 t: usize,
22963 pad: usize,
22964 ) -> Result<(), Box<dyn std::error::Error>> {
22965 let f = self.func("conv_left_pad_f32");
22966 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
22967 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
22968 let __s_b = self.gpu.stream();
22969 let mut b = __s_b.launch_builder(&f);
22970 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
22971 unsafe {
22972 b.launch(cfg)?;
22973 }
22974 Ok(())
22975 }
22976
22977 pub fn conv_assemble_and_roll(
22981 &self,
22982 qkv_col: &CudaSlice<f32>,
22983 conv_state: &mut CudaSlice<f32>,
22984 conv_in: &mut CudaSlice<f32>,
22985 conv_dim: usize,
22986 pad: usize,
22987 ) -> Result<(), Box<dyn std::error::Error>> {
22988 let f = self.func("conv_assemble_and_roll_f32");
22989 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22990 let (cd, p) = (conv_dim as i32, pad as i32);
22991 let __s_b = self.gpu.stream();
22992 let mut b = __s_b.launch_builder(&f);
22993 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
22994 unsafe {
22995 b.launch(cfg)?;
22996 }
22997 Ok(())
22998 }
22999
23000 pub fn ssm_conv1d_fused_decode(
23006 &self,
23007 qkv_col: &CudaSlice<f32>,
23008 conv_state: &mut CudaSlice<f32>,
23009 w: &CudaSlice<f32>,
23010 conv_out: &mut CudaSlice<f32>,
23011 conv_dim: usize,
23012 d_conv: usize,
23013 ) -> Result<(), Box<dyn std::error::Error>> {
23014 let f = self.func("ssm_conv1d_fused_decode_f32");
23015 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
23016 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23017 let __s_b = self.gpu.stream();
23018 let mut b = __s_b.launch_builder(&f);
23019 b.arg(qkv_col)
23020 .arg(conv_state)
23021 .arg(w)
23022 .arg(conv_out)
23023 .arg(&cd)
23024 .arg(&dc);
23025 unsafe {
23026 b.launch(cfg)?;
23027 }
23028 Ok(())
23029 }
23030
23031 pub fn slice_range(
23034 &self,
23035 src: &CudaSlice<f32>,
23036 start: usize,
23037 len: usize,
23038 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23039 let host = self.gpu.stream().clone_dtoh(src)?;
23040 self.gpu.stream().synchronize()?;
23041 Ok(self.htod(&host[start..start + len])?)
23042 }
23043}
23044
23045#[cfg(test)]
23046mod target_dispatch_tests {
23047 use super::legacy_quant_gemm_allowed;
23048
23049 #[test]
23050 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
23051 assert!(legacy_quant_gemm_allowed(false, false, false));
23053 assert!(!legacy_quant_gemm_allowed(false, false, true));
23054 assert!(!legacy_quant_gemm_allowed(true, false, false));
23056 assert!(!legacy_quant_gemm_allowed(true, false, true));
23057 assert!(legacy_quant_gemm_allowed(true, true, false));
23059 assert!(!legacy_quant_gemm_allowed(true, true, true));
23060 }
23061
23062 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
23063 #[test]
23064 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
23065 assert!(!legacy_quant_gemm_allowed(
23066 cfg!(memra_portable_cuda),
23067 cfg!(memra_hopper_mma),
23068 false
23069 ));
23070 }
23071
23072 #[cfg(memra_hopper_mma)]
23073 #[test]
23074 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
23075 assert!(legacy_quant_gemm_allowed(
23076 cfg!(memra_portable_cuda),
23077 cfg!(memra_hopper_mma),
23078 false
23079 ));
23080 assert!(super::portable_mma_gated() == false);
23081 }
23082}
23083
23084impl memra_kv::KvDev for Engine {
23087 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23088 Engine::zeros(self, n)
23089 }
23090 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23091 Engine::uninit(self, n)
23092 }
23093 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
23094 Engine::alloc_u8(self, n)
23095 }
23096 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
23097 Engine::htod_i32(self, v)
23098 }
23099 fn clone_dtod(
23100 &self,
23101 src: &CudaSlice<f32>,
23102 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23103 Engine::clone_dtod(self, src)
23104 }
23105 fn copy_into(
23106 &self,
23107 dst: &mut CudaSlice<f32>,
23108 off: usize,
23109 src: &CudaSlice<f32>,
23110 len: usize,
23111 ) -> Result<(), Box<dyn std::error::Error>> {
23112 Engine::copy_into(self, dst, off, src, len)
23113 }
23114 fn set_i32_one(
23115 &self,
23116 d: &mut CudaSlice<i32>,
23117 v: i32,
23118 ) -> Result<(), Box<dyn std::error::Error>> {
23119 Engine::set_i32_one(self, d, v)
23120 }
23121}
23122
23123#[cfg(test)]
23124mod fused_gate_bounds_tests {
23125 use super::*;
23126
23127 #[test]
23140 #[ignore = "requires a CUDA GPU"]
23141 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
23142 let e = Engine::new(0).unwrap();
23143 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
23144 let fused = 2 * head_dim * n_head * t;
23145 let out_n = head_dim * n_head * t;
23146
23147 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
23149 let mut q = e.uninit(out_n).unwrap();
23150 let mut gate = e.uninit(out_n).unwrap();
23151 let err = e
23152 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
23153 .expect_err("half-width wq must be refused, not read past")
23154 .to_string();
23155 assert!(err.contains("NO fused gate"), "{err}");
23156 assert!(err.contains(&format!("{fused}")), "{err}");
23157
23158 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
23161 let wide = e.htod(&host).unwrap();
23162 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
23163 .expect("full-width wq splits");
23164 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
23165 for tok in 0..t {
23166 for hh in 0..n_head {
23167 for d in 0..head_dim {
23168 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
23169 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
23170 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
23171 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
23172 }
23173 }
23174 }
23175
23176 let mut small = e.uninit(out_n - 1).unwrap();
23178 assert!(
23179 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
23180 .is_err()
23181 );
23182 }
23183}
23184
23185#[cfg(test)]
23189mod fused_rope_width_tests {
23190 use super::Engine;
23191
23192 #[test]
23195 fn full_width_is_accepted() {
23196 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
23197 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
23198 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
23199 }
23200
23201 #[test]
23214 fn gemma4_official_artifact_widths_pass() {
23215 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
23216 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
23217 }
23218
23219 #[test]
23222 fn partial_rotary_is_refused_with_the_geometry_named() {
23223 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
23225 .expect_err("partial rotary must refuse");
23226 let msg = err.to_string();
23227 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
23228 assert!(msg.contains("n_rot 64"), "{msg}");
23229 assert!(msg.contains("head_dim 256"), "{msg}");
23230 assert!(
23231 msg.contains("64..256"),
23232 "names the band it would corrupt: {msg}"
23233 );
23234 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
23236 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
23238 }
23239}