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)]
1610 pub fn gumbel_perturb_filtered(
1611 &self,
1612 x: &CudaSlice<f32>,
1613 y: &mut CudaSlice<f32>,
1614 n: usize,
1615 seed: u64,
1616 stream_pos: u32,
1617 temp: f32,
1618 row_max: f32,
1619 th: f32,
1620 ) -> Result<(), Box<dyn std::error::Error>> {
1621 let f = self.func("gumbel_perturb_filtered_f32");
1622 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1623 let cfg = LaunchConfig {
1624 grid_dim: (n.div_ceil(256) as u32, 1, 1),
1625 block_dim: (256, 1, 1),
1626 shared_mem_bytes: 0,
1627 };
1628 let __s_b = self.gpu.stream();
1629 let mut b = __s_b.launch_builder(&f);
1630 b.arg(x)
1631 .arg(&mut *y)
1632 .arg(&ni)
1633 .arg(&slo)
1634 .arg(&shi)
1635 .arg(&stream_pos)
1636 .arg(&temp)
1637 .arg(&row_max)
1638 .arg(&th);
1639 unsafe {
1640 b.launch(cfg)?;
1641 }
1642 Ok(())
1643 }
1644
1645 #[allow(clippy::too_many_arguments)]
1649 pub fn penalize_logits(
1650 &self,
1651 x: &mut CudaSlice<f32>,
1652 hist: &CudaSlice<u32>,
1653 n_hist: usize,
1654 rep: f32,
1655 freq: f32,
1656 present: f32,
1657 n: usize,
1658 ) -> Result<(), Box<dyn std::error::Error>> {
1659 if n_hist == 0 {
1660 return Ok(());
1661 }
1662 let f = self.func("penalize_logits_f32");
1663 let (nh, ni) = (n_hist as i32, n as i32);
1664 let cfg = LaunchConfig {
1665 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1666 block_dim: (128, 1, 1),
1667 shared_mem_bytes: 0,
1668 };
1669 let __s_b = self.gpu.stream();
1670 let mut b = __s_b.launch_builder(&f);
1671 b.arg(&mut *x)
1672 .arg(hist)
1673 .arg(&nh)
1674 .arg(&rep)
1675 .arg(&freq)
1676 .arg(&present)
1677 .arg(&ni);
1678 unsafe {
1679 b.launch(cfg)?;
1680 }
1681 Ok(())
1682 }
1683
1684 #[allow(clippy::too_many_arguments)]
1686 pub fn penalize_logits_rows(
1687 &self,
1688 x: &mut CudaSlice<f32>,
1689 hist: &CudaSlice<u32>,
1690 n_hist: usize,
1691 rep: f32,
1692 freq: f32,
1693 present: f32,
1694 n: usize,
1695 nrow: usize,
1696 ) -> Result<(), Box<dyn std::error::Error>> {
1697 if n_hist == 0 || nrow == 0 {
1698 return Ok(());
1699 }
1700 let f = self.func("penalize_logits_rows_f32");
1701 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1702 let cfg = LaunchConfig {
1703 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1704 block_dim: (128, 1, 1),
1705 shared_mem_bytes: 0,
1706 };
1707 let __s_b = self.gpu.stream();
1708 let mut b = __s_b.launch_builder(&f);
1709 b.arg(&mut *x)
1710 .arg(hist)
1711 .arg(&nh)
1712 .arg(&rep)
1713 .arg(&freq)
1714 .arg(&present)
1715 .arg(&ni)
1716 .arg(&nr);
1717 unsafe {
1718 b.launch(cfg)?;
1719 }
1720 Ok(())
1721 }
1722
1723 pub fn wpf_level() -> u32 {
1731 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1732 *ON.get_or_init(|| {
1733 std::env::var("MEMRA_WPF")
1734 .ok()
1735 .and_then(|v| v.parse().ok())
1736 .unwrap_or(1)
1737 })
1738 }
1739
1740 pub fn set_verify_exact(&self, on: bool) {
1752 self.verify_exact
1753 .store(on, std::sync::atomic::Ordering::Relaxed);
1754 }
1755 pub(crate) fn verify_exact_on(&self) -> bool {
1756 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1757 }
1758
1759 pub fn qkv_append_on() -> bool {
1762 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1763 *ON.get_or_init(|| {
1764 std::env::var("MEMRA_QKV_APPEND")
1765 .map(|v| v != "0")
1766 .unwrap_or(true)
1767 })
1768 }
1769
1770 pub fn pdl_wb_on() -> bool {
1773 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1774 *ON.get_or_init(|| {
1775 std::env::var("MEMRA_PDL_WB")
1776 .map(|v| v != "0")
1777 .unwrap_or(true)
1778 })
1779 }
1780
1781 pub fn pdl_mmvq_on() -> bool {
1785 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1786 *ON.get_or_init(|| {
1787 std::env::var("MEMRA_PDL_MMVQ")
1788 .map(|v| v != "0")
1789 .unwrap_or(true)
1790 })
1791 }
1792
1793 pub fn pdl_on() -> bool {
1794 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1795 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1796 }
1797
1798 pub fn pdl_nvfp4q8_on() -> bool {
1804 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1805 *ON.get_or_init(|| {
1806 std::env::var("MEMRA_PDL_NVFP4")
1807 .map(|v| v != "0")
1808 .unwrap_or(true)
1809 })
1810 }
1811
1812 fn q40_mr1_on() -> bool {
1818 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1819 match *Q40MR.get_or_init(|| {
1820 std::env::var("MEMRA_Q40_MR")
1821 .ok()
1822 .and_then(|v| v.parse().ok())
1823 }) {
1824 Some(v) => v == 1,
1825 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1826 }
1827 }
1828
1829 fn pdl_func_flash(
1834 &self,
1835 g: bool,
1836 name: &'static str,
1837 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1838 use cudarc::driver::sys as cu;
1839 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1846 std::sync::Mutex::new(None);
1847 static FNS: std::sync::Mutex<
1848 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
1849 > = std::sync::Mutex::new(None);
1850 let ctx_key = self.ctx().cu_ctx() as usize;
1851 if let Some(&f) = FNS
1852 .lock()
1853 .unwrap()
1854 .get_or_insert_with(Default::default)
1855 .get(&(ctx_key, g, name))
1856 {
1857 return Ok(f as cu::CUfunction);
1858 }
1859 let module = {
1860 let mut mods = MODS.lock().unwrap();
1861 let map = mods.get_or_insert_with(Default::default);
1862 match map.get(&(ctx_key, g)) {
1863 Some(&m) => m,
1864 None => {
1865 let m = self.pdl_load_module_in_ctx(if g {
1866 FLASH_FATBIN_KF8VF8
1867 } else {
1868 FLASH_FATBIN
1869 })?;
1870 map.insert((ctx_key, g), m);
1871 m
1872 }
1873 }
1874 };
1875 let cname = std::ffi::CString::new(name)?;
1876 let mut f: cu::CUfunction = std::ptr::null_mut();
1877 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1878 if r != cu::CUresult::CUDA_SUCCESS {
1879 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
1880 }
1881 FNS.lock()
1882 .unwrap()
1883 .get_or_insert_with(Default::default)
1884 .insert((ctx_key, g, name), f as usize);
1885 Ok(f)
1886 }
1887
1888 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1893 use cudarc::driver::sys as cu;
1894 let mut prev: cu::CUcontext = std::ptr::null_mut();
1895 unsafe {
1896 cu::cuCtxGetCurrent(&mut prev).result()?;
1897 }
1898 self.ctx().bind_to_thread()?;
1899 let mut m: cu::CUmodule = std::ptr::null_mut();
1900 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1901 let restore = if prev.is_null() {
1902 cu::CUresult::CUDA_SUCCESS
1903 } else {
1904 unsafe { cu::cuCtxSetCurrent(prev) }
1905 };
1906 if r != cu::CUresult::CUDA_SUCCESS {
1907 return Err(format!("pdl module load: {r:?}").into());
1908 }
1909 if restore != cu::CUresult::CUDA_SUCCESS {
1910 return Err(format!("pdl module load: ctx restore {restore:?}").into());
1911 }
1912 Ok(m as usize)
1913 }
1914
1915 fn pdl_func(
1916 &self,
1917 name: &'static str,
1918 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1919 use cudarc::driver::sys as cu;
1920 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1923 std::sync::Mutex::new(None);
1924 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1927 std::sync::Mutex::new(None);
1928 static FNS: std::sync::Mutex<
1929 Option<std::collections::HashMap<(usize, &'static str), usize>>,
1930 > = std::sync::Mutex::new(None);
1931 let ctx_key = self.ctx().cu_ctx() as usize;
1932 if let Some(&f) = FNS
1933 .lock()
1934 .unwrap()
1935 .get_or_insert_with(Default::default)
1936 .get(&(ctx_key, name))
1937 {
1938 return Ok(f as cu::CUfunction);
1939 }
1940 let module = {
1941 let mut mods = MODULES.lock().unwrap();
1942 let map = mods.get_or_insert_with(Default::default);
1943 match map.get(&ctx_key) {
1944 Some(&m) => m,
1945 None => {
1946 let m = self.pdl_load_module_in_ctx(FATBIN)?;
1947 map.insert(ctx_key, m);
1948 m
1949 }
1950 }
1951 };
1952 let cname = std::ffi::CString::new(name)?;
1953 let mut f: cu::CUfunction = std::ptr::null_mut();
1954 let mut r =
1955 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1956 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
1957 let qmodule = {
1958 let mut mods = QMODULES.lock().unwrap();
1959 let map = mods.get_or_insert_with(Default::default);
1960 match map.get(&ctx_key) {
1961 Some(&m) => m,
1962 None => {
1963 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
1964 map.insert(ctx_key, m);
1965 m
1966 }
1967 }
1968 };
1969 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
1970 }
1971 if r != cu::CUresult::CUDA_SUCCESS {
1972 return Err(format!("pdl_func {name}: {r:?}").into());
1973 }
1974 FNS.lock()
1975 .unwrap()
1976 .get_or_insert_with(Default::default)
1977 .insert((ctx_key, name), f as usize);
1978 Ok(f)
1979 }
1980
1981 unsafe fn launch_pdl_flash(
1993 &self,
1994 g: bool,
1995 name: &'static str,
1996 grid: (u32, u32, u32),
1997 block: (u32, u32, u32),
1998 smem: u32,
1999 params: &mut [*mut std::ffi::c_void],
2000 ) -> Result<(), Box<dyn std::error::Error>> {
2001 use cudarc::driver::sys as cu;
2002 let f = self.pdl_func_flash(g, name)?;
2003 if smem > 0 {
2004 let r =
2006 unsafe {
2007 cu::cuFuncSetAttribute(f,
2008 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2009 smem as i32)
2010 };
2011 if r != cu::CUresult::CUDA_SUCCESS {
2012 return Err(format!("pdl smem attr {name}: {r:?}").into());
2013 }
2014 }
2015 let mut attr = cu::CUlaunchAttribute {
2016 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2017 pad: [0; 4],
2018 value: cu::CUlaunchAttributeValue {
2019 programmaticStreamSerializationAllowed: 1,
2020 },
2021 };
2022 let cfg = cu::CUlaunchConfig {
2023 gridDimX: grid.0,
2024 gridDimY: grid.1,
2025 gridDimZ: grid.2,
2026 blockDimX: block.0,
2027 blockDimY: block.1,
2028 blockDimZ: block.2,
2029 sharedMemBytes: smem,
2030 hStream: self.gpu.stream().cu_stream(),
2031 attrs: &mut attr,
2032 numAttrs: 1,
2033 };
2034 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2035 if r != cu::CUresult::CUDA_SUCCESS {
2036 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2037 }
2038 Ok(())
2039 }
2040
2041 unsafe fn launch_pdl(
2042 &self,
2043 name: &'static str,
2044 grid: (u32, u32, u32),
2045 block: (u32, u32, u32),
2046 params: &mut [*mut std::ffi::c_void],
2047 ) -> Result<(), Box<dyn std::error::Error>> {
2048 use cudarc::driver::sys as cu;
2049 let f = self.pdl_func(name)?;
2050 let mut attr = cu::CUlaunchAttribute {
2051 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2052 pad: [0; 4],
2053 value: cu::CUlaunchAttributeValue {
2054 programmaticStreamSerializationAllowed: 1,
2055 },
2056 };
2057 let cfg = cu::CUlaunchConfig {
2058 gridDimX: grid.0,
2059 gridDimY: grid.1,
2060 gridDimZ: grid.2,
2061 blockDimX: block.0,
2062 blockDimY: block.1,
2063 blockDimZ: block.2,
2064 sharedMemBytes: 0,
2065 hStream: self.gpu.stream().cu_stream(),
2066 attrs: &mut attr,
2067 numAttrs: 1,
2068 };
2069 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2070 if r != cu::CUresult::CUDA_SUCCESS {
2071 return Err(format!("launch_pdl {name}: {r:?}").into());
2072 }
2073 Ok(())
2074 }
2075
2076 pub fn prefetch_weight_l2(
2079 &self,
2080 w: &crate::model::GpuTensor,
2081 ) -> Result<(), Box<dyn std::error::Error>> {
2082 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2083 let p = rp4.as_ref().unwrap_or(bytes);
2084 self.prefetch_l2(p, p.len())?;
2085 }
2086 Ok(())
2087 }
2088
2089 pub fn gather_row_bf16(
2092 &self,
2093 table: &CudaSlice<u8>,
2094 tok: &CudaSlice<u32>,
2095 idx: usize,
2096 dst: &mut CudaSlice<f32>,
2097 ncols: usize,
2098 ) -> Result<(), Box<dyn std::error::Error>> {
2099 let f = self.func("gather_row_bf16_f32");
2100 let cfg = LaunchConfig {
2101 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2102 block_dim: (256, 1, 1),
2103 shared_mem_bytes: 0,
2104 };
2105 let (nc, ix) = (ncols as i32, idx as i32);
2106 let __s_b = self.gpu.stream();
2107 let mut b = __s_b.launch_builder(&f);
2108 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2109 unsafe {
2110 b.launch(cfg)?;
2111 }
2112 Ok(())
2113 }
2114
2115 pub fn add_row_inplace(
2117 &self,
2118 logits: &mut CudaSlice<f32>,
2119 bias: &CudaSlice<f32>,
2120 n: usize,
2121 row_off: usize,
2122 ) -> Result<(), Box<dyn std::error::Error>> {
2123 let f = self.func("add_row_inplace_f32");
2124 let cfg = LaunchConfig {
2125 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2126 block_dim: (256, 1, 1),
2127 shared_mem_bytes: 0,
2128 };
2129 let (ni, off) = (n as i32, row_off as i64);
2130 let __s_b = self.gpu.stream();
2131 let mut b = __s_b.launch_builder(&f);
2132 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2133 unsafe {
2134 b.launch(cfg)?;
2135 }
2136 Ok(())
2137 }
2138
2139 pub fn prefetch_l2(
2141 &self,
2142 p: &CudaSlice<u8>,
2143 n: usize,
2144 ) -> Result<(), Box<dyn std::error::Error>> {
2145 let f = self.func("prefetch_l2_bytes");
2146 let lines = n.div_ceil(128);
2147 let ni = n as i64;
2148 let cfg = LaunchConfig {
2149 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2150 block_dim: (256, 1, 1),
2151 shared_mem_bytes: 0,
2152 };
2153 let __s_b = self.gpu.stream();
2154 let mut b = __s_b.launch_builder(&f);
2155 b.arg(p).arg(&ni);
2156 unsafe {
2157 b.launch(cfg)?;
2158 }
2159 Ok(())
2160 }
2161
2162 pub fn router_gemv(
2165 &self,
2166 w: &CudaSlice<f32>,
2167 x: &CudaSlice<f32>,
2168 n_embd: usize,
2169 n_experts: usize,
2170 t: usize,
2171 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2172 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2178 Ok("0") => false,
2179 Ok(_) => true,
2180 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2181 };
2182 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2192 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2193 }
2194
2195 pub fn router_gemv_form(
2198 &self,
2199 w: &CudaSlice<f32>,
2200 x: &CudaSlice<f32>,
2201 n_embd: usize,
2202 n_experts: usize,
2203 t: usize,
2204 w8: bool,
2205 batch: bool,
2206 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2207 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2208 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2209 let f = if batch {
2210 self.func("router_gemv_f32_w8_batch")
2211 } else if w8 {
2212 self.func("router_gemv_f32_w8")
2213 } else {
2214 self.func("router_gemv_f32")
2215 };
2216 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2217 let cfg = if batch {
2218 LaunchConfig {
2219 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2220 block_dim: (32, 8, 1),
2221 shared_mem_bytes: 0,
2222 }
2223 } else {
2224 LaunchConfig {
2225 grid_dim: (n_experts as u32, t as u32, 1),
2226 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2227 shared_mem_bytes: 0,
2228 }
2229 };
2230 let __s_b = self.gpu.stream();
2231 let mut b = __s_b.launch_builder(&f);
2232 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2233 unsafe {
2234 b.launch(cfg)?;
2235 }
2236 Ok(y)
2237 }
2238
2239 pub fn rows_permute(
2241 &self,
2242 src: &CudaSlice<f32>,
2243 idx: &CudaSlice<i32>,
2244 nrows: usize,
2245 ncols: usize,
2246 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2247 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2248 let f = self.func("rows_permute_f32");
2249 let (nc, nr) = (ncols as i32, nrows as i32);
2250 let cfg = LaunchConfig {
2251 grid_dim: (nrows as u32, 1, 1),
2252 block_dim: (256, 1, 1),
2253 shared_mem_bytes: 0,
2254 };
2255 let __s_b = self.gpu.stream();
2256 let mut b = __s_b.launch_builder(&f);
2257 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2258 unsafe {
2259 b.launch(cfg)?;
2260 }
2261 Ok(dst)
2262 }
2263
2264 pub fn sigmoid_dot_rows(
2269 &self,
2270 x: &CudaSlice<f32>,
2271 w: &CudaSlice<f32>,
2272 n_embd: usize,
2273 t: usize,
2274 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2275 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2278 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2279 let gs = self.linear(x, w, t, n_embd, 1)?;
2280 let mut g = self.uninit(t)?;
2281 self.sigmoid(&gs, &mut g, t)?;
2282 return Ok(g);
2283 }
2284 let mut g = self.alloc_uninit::<f32>(t)?;
2290 let f = self.func("sigmoid_dot_rows_f32");
2291 let (ne, ti) = (n_embd as i32, t as i32);
2292 let cfg = LaunchConfig {
2293 grid_dim: (t as u32, 1, 1),
2294 block_dim: (32, 8, 1),
2295 shared_mem_bytes: 0,
2296 };
2297 let __s_b = self.gpu.stream();
2298 let mut b = __s_b.launch_builder(&f);
2299 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2300 unsafe {
2301 b.launch(cfg)?;
2302 }
2303 Ok(g)
2304 }
2305
2306 pub fn spec_rollback_stream(
2308 &self,
2309 len_ptrs: &CudaSlice<u64>,
2310 pos_start: &CudaSlice<i32>,
2311 acc: &CudaSlice<u32>,
2312 base: usize,
2313 n_rows: usize,
2314 ) -> Result<(), Box<dyn std::error::Error>> {
2315 let f = self.func("spec_rollback_stream");
2316 let (b, nr) = (base as i32, n_rows as i32);
2317 let cfg = LaunchConfig {
2318 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2319 block_dim: (64, 1, 1),
2320 shared_mem_bytes: 0,
2321 };
2322 let __s_bl = self.gpu.stream();
2323 let mut bl = __s_bl.launch_builder(&f);
2324 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2325 unsafe {
2326 bl.launch(cfg)?;
2327 }
2328 Ok(())
2329 }
2330
2331 pub fn plain_tok_ring(
2333 &self,
2334 vam: &CudaSlice<u32>,
2335 pos_start: &CudaSlice<i32>,
2336 base: usize,
2337 ring: &mut CudaSlice<u32>,
2338 ) -> Result<(), Box<dyn std::error::Error>> {
2339 let f = self.func("plain_tok_ring");
2340 let (b, cap) = (base as i32, ring.len() as i32);
2341 let cfg = LaunchConfig {
2342 grid_dim: (1, 1, 1),
2343 block_dim: (32, 1, 1),
2344 shared_mem_bytes: 0,
2345 };
2346 let __s_bl = self.gpu.stream();
2347 let mut bl = __s_bl.launch_builder(&f);
2348 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2349 unsafe {
2350 bl.launch(cfg)?;
2351 }
2352 Ok(())
2353 }
2354
2355 pub fn spec_ring_commit(
2357 &self,
2358 vtok: &CudaSlice<u32>,
2359 acc: &CudaSlice<u32>,
2360 brk: &CudaSlice<u32>,
2361 ring: &mut CudaSlice<u32>,
2362 pend: &mut CudaSlice<u32>,
2363 ) -> Result<(), Box<dyn std::error::Error>> {
2364 let f = self.func("spec_ring_commit");
2365 let cfg = LaunchConfig {
2366 grid_dim: (1, 1, 1),
2367 block_dim: (32, 1, 1),
2368 shared_mem_bytes: 0,
2369 };
2370 let __s_b = self.gpu.stream();
2371 let mut b = __s_b.launch_builder(&f);
2372 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2373 unsafe {
2374 b.launch(cfg)?;
2375 }
2376 Ok(())
2377 }
2378 pub fn i32_copy_add(
2379 &self,
2380 src: &CudaSlice<i32>,
2381 dst: &mut CudaSlice<i32>,
2382 delta: i32,
2383 ) -> Result<(), Box<dyn std::error::Error>> {
2384 let f = self.func("i32_copy_add");
2385 let cfg = LaunchConfig {
2386 grid_dim: (1, 1, 1),
2387 block_dim: (32, 1, 1),
2388 shared_mem_bytes: 0,
2389 };
2390 let __s_b = self.gpu.stream();
2391 let mut b = __s_b.launch_builder(&f);
2392 b.arg(src).arg(dst).arg(&delta);
2393 unsafe {
2394 b.launch(cfg)?;
2395 }
2396 Ok(())
2397 }
2398 pub fn u32_copy(
2399 &self,
2400 src: &CudaSlice<u32>,
2401 dst: &mut CudaSlice<u32>,
2402 ) -> Result<(), Box<dyn std::error::Error>> {
2403 let f = self.func("u32_copy");
2404 let cfg = LaunchConfig {
2405 grid_dim: (1, 1, 1),
2406 block_dim: (32, 1, 1),
2407 shared_mem_bytes: 0,
2408 };
2409 let __s_b = self.gpu.stream();
2410 let mut b = __s_b.launch_builder(&f);
2411 b.arg(src).arg(dst);
2412 unsafe {
2413 b.launch(cfg)?;
2414 }
2415 Ok(())
2416 }
2417
2418 pub fn spec_adapt_k(
2422 &self,
2423 acc: &CudaSlice<u32>,
2424 brk: &mut CudaSlice<u32>,
2425 floor: usize,
2426 cap: usize,
2427 ) -> Result<(), Box<dyn std::error::Error>> {
2428 let f = self.func("spec_adapt_k");
2429 let (fl, cp) = (floor as i32, cap as i32);
2430 let cfg = LaunchConfig {
2431 grid_dim: (1, 1, 1),
2432 block_dim: (32, 1, 1),
2433 shared_mem_bytes: 0,
2434 };
2435 let __s_b = self.gpu.stream();
2436 let mut b = __s_b.launch_builder(&f);
2437 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2438 unsafe {
2439 b.launch(cfg)?;
2440 }
2441 Ok(())
2442 }
2443
2444 pub fn spec_accept_greedy_dc(
2446 &self,
2447 preds: &CudaSlice<u32>,
2448 vtok: &CudaSlice<u32>,
2449 last_pred: &CudaSlice<u32>,
2450 brk: &CudaSlice<u32>,
2451 out: &mut CudaSlice<u32>,
2452 ) -> Result<(), Box<dyn std::error::Error>> {
2453 let f = self.func("spec_accept_greedy_dc");
2454 let cfg = LaunchConfig {
2455 grid_dim: (1, 1, 1),
2456 block_dim: (32, 1, 1),
2457 shared_mem_bytes: 0,
2458 };
2459 let __s_b = self.gpu.stream();
2460 let mut b = __s_b.launch_builder(&f);
2461 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2462 unsafe {
2463 b.launch(cfg)?;
2464 }
2465 Ok(())
2466 }
2467
2468 pub fn pos_iota(
2470 &self,
2471 pos0: &CudaSlice<i32>,
2472 out: &mut CudaSlice<i32>,
2473 t: usize,
2474 ) -> Result<(), Box<dyn std::error::Error>> {
2475 let f = self.func("pos_iota_i32");
2476 let ti = t as i32;
2477 let cfg = LaunchConfig {
2478 grid_dim: (1, 1, 1),
2479 block_dim: (t.max(1) as u32, 1, 1),
2480 shared_mem_bytes: 0,
2481 };
2482 let __s_b = self.gpu.stream();
2483 let mut b = __s_b.launch_builder(&f);
2484 b.arg(pos0).arg(out).arg(&ti);
2485 unsafe {
2486 b.launch(cfg)?;
2487 }
2488 Ok(())
2489 }
2490 #[allow(clippy::too_many_arguments)]
2491 pub fn append_kv_quantized_rows_dc(
2492 &self,
2493 k_rows: &CudaSlice<f32>,
2494 v_rows: &CudaSlice<f32>,
2495 kc: &mut CudaSlice<u8>,
2496 vc: &mut CudaSlice<u8>,
2497 t0_dev: &CudaSlice<i32>,
2498 t: usize,
2499 kv_dim_k: usize,
2500 kv_dim_v: usize,
2501 k_tok_bytes: usize,
2502 v_tok_bytes: usize,
2503 g: bool,
2504 ) -> Result<(), Box<dyn std::error::Error>> {
2505 let f = if g {
2506 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2507 } else {
2508 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2509 };
2510 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2511 let cfg = LaunchConfig {
2512 grid_dim: (nblk, t as u32, 1),
2513 block_dim: (32, 1, 1),
2514 shared_mem_bytes: 0,
2515 };
2516 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2517 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2518 let __s_b = self.gpu.stream();
2519 let mut b = __s_b.launch_builder(&f);
2520 b.arg(k_rows)
2521 .arg(v_rows)
2522 .arg(kc)
2523 .arg(vc)
2524 .arg(t0_dev)
2525 .arg(&kdk)
2526 .arg(&kdv)
2527 .arg(&ktb)
2528 .arg(&vtb);
2529 unsafe {
2530 b.launch(cfg)?;
2531 }
2532 Ok(())
2533 }
2534
2535 #[allow(clippy::too_many_arguments)]
2538 pub fn append_kv_quantized_row_dc_inc(
2539 &self,
2540 k_row: &CudaSlice<f32>,
2541 v_row: &CudaSlice<f32>,
2542 kc: &mut CudaSlice<u8>,
2543 vc: &mut CudaSlice<u8>,
2544 t0_dev: &mut CudaSlice<i32>,
2545 kv_dim_k: usize,
2546 kv_dim_v: usize,
2547 k_tok_bytes: usize,
2548 v_tok_bytes: usize,
2549 g: bool,
2550 ) -> Result<(), Box<dyn std::error::Error>> {
2551 let f = if g {
2552 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2553 } else {
2554 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2555 };
2556 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2557 let cfg = LaunchConfig {
2558 grid_dim: (1, 1, 1),
2559 block_dim: (nthreads, 1, 1),
2560 shared_mem_bytes: 0,
2561 };
2562 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2563 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2564 let __s_b = self.gpu.stream();
2565 let mut b = __s_b.launch_builder(&f);
2566 b.arg(k_row)
2567 .arg(v_row)
2568 .arg(kc)
2569 .arg(vc)
2570 .arg(t0_dev)
2571 .arg(&kdk)
2572 .arg(&kdv)
2573 .arg(&ktb)
2574 .arg(&vtb);
2575 unsafe {
2576 b.launch(cfg)?;
2577 }
2578 Ok(())
2579 }
2580
2581 pub fn pack_tok_p(
2583 &self,
2584 tok: &CudaSlice<u32>,
2585 p: &CudaSlice<f32>,
2586 out: &mut CudaSlice<u32>,
2587 slot: usize,
2588 ) -> Result<(), Box<dyn std::error::Error>> {
2589 let f = self.func("pack_tok_p");
2590 let sl = slot as i32;
2591 let cfg = LaunchConfig {
2592 grid_dim: (1, 1, 1),
2593 block_dim: (32, 1, 1),
2594 shared_mem_bytes: 0,
2595 };
2596 let __s_b = self.gpu.stream();
2597 let mut b = __s_b.launch_builder(&f);
2598 b.arg(tok).arg(p).arg(out).arg(&sl);
2599 unsafe {
2600 b.launch(cfg)?;
2601 }
2602 Ok(())
2603 }
2604 pub fn tok_map_u32(
2605 &self,
2606 tok: &mut CudaSlice<u32>,
2607 map: &CudaSlice<u32>,
2608 ) -> Result<(), Box<dyn std::error::Error>> {
2609 let f = self.func("tok_map_u32");
2610 let cfg = LaunchConfig {
2611 grid_dim: (1, 1, 1),
2612 block_dim: (32, 1, 1),
2613 shared_mem_bytes: 0,
2614 };
2615 let __s_b = self.gpu.stream();
2616 let mut b = __s_b.launch_builder(&f);
2617 b.arg(tok).arg(map);
2618 unsafe {
2619 b.launch(cfg)?;
2620 }
2621 Ok(())
2622 }
2623
2624 #[allow(clippy::too_many_arguments)]
2626 pub fn spec_assemble_verify(
2627 &self,
2628 tokp: &CudaSlice<u32>,
2629 pend: &CudaSlice<u32>,
2630 d2t: Option<&CudaSlice<u32>>,
2631 vtok: &mut CudaSlice<u32>,
2632 brk: &mut CudaSlice<u32>,
2633 p_min: f32,
2634 k: usize,
2635 pmin0: bool,
2636 ) -> Result<(), Box<dyn std::error::Error>> {
2637 let f = self.func("spec_assemble_verify");
2638 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2639 let cfg = LaunchConfig {
2640 grid_dim: (1, 1, 1),
2641 block_dim: (32, 1, 1),
2642 shared_mem_bytes: 0,
2643 };
2644 let __s_b = self.gpu.stream();
2645 let mut b = __s_b.launch_builder(&f);
2646 match d2t {
2647 Some(m) => {
2648 b.arg(tokp)
2649 .arg(pend)
2650 .arg(m)
2651 .arg(vtok)
2652 .arg(brk)
2653 .arg(&p_min)
2654 .arg(&ki)
2655 .arg(&pm);
2656 unsafe {
2657 b.launch(cfg)?;
2658 }
2659 }
2660 None => {
2661 let null: u64 = 0;
2662 b.arg(tokp)
2663 .arg(pend)
2664 .arg(&null)
2665 .arg(vtok)
2666 .arg(brk)
2667 .arg(&p_min)
2668 .arg(&ki)
2669 .arg(&pm);
2670 unsafe {
2671 b.launch(cfg)?;
2672 }
2673 }
2674 }
2675 Ok(())
2676 }
2677
2678 #[allow(clippy::too_many_arguments)]
2680 pub fn ssm_conv_ring_rebuild_dc(
2681 &self,
2682 qkv_tm: &CudaSlice<f32>,
2683 ring_old: &CudaSlice<f32>,
2684 conv_state: &mut CudaSlice<f32>,
2685 conv_dim: usize,
2686 acc: &CudaSlice<u32>,
2687 base: usize,
2688 t_v: usize,
2689 d_conv: usize,
2690 ) -> Result<(), Box<dyn std::error::Error>> {
2691 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
2692 let n = conv_dim * (d_conv - 1);
2693 let cfg = LaunchConfig::for_num_elems(n as u32);
2694 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
2695 let __s_b = self.gpu.stream();
2696 let mut b = __s_b.launch_builder(&f);
2697 b.arg(qkv_tm)
2698 .arg(ring_old)
2699 .arg(conv_state)
2700 .arg(&cd)
2701 .arg(acc)
2702 .arg(&b0)
2703 .arg(&tv)
2704 .arg(&dc);
2705 unsafe {
2706 b.launch(cfg)?;
2707 }
2708 Ok(())
2709 }
2710 #[allow(clippy::too_many_arguments)]
2711 pub fn gdn_scan_s128_dc(
2712 &self,
2713 q: &CudaSlice<f32>,
2714 k: &CudaSlice<f32>,
2715 v: &CudaSlice<f32>,
2716 g: &CudaSlice<f32>,
2717 beta: &CudaSlice<f32>,
2718 state_in: &CudaSlice<f32>,
2719 state_out: &mut CudaSlice<f32>,
2720 o: &mut CudaSlice<f32>,
2721 n_head: usize,
2722 acc: &CudaSlice<u32>,
2723 base: usize,
2724 t_v: usize,
2725 scale: f32,
2726 ) -> Result<(), Box<dyn std::error::Error>> {
2727 let f = self.func("gdn_scan_s128_dc");
2728 const S_V: u32 = 128;
2729 const WARP: u32 = 32;
2730 const COLS_PER_BLOCK: u32 = 4;
2731 let cfg = LaunchConfig {
2732 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
2733 block_dim: (WARP, COLS_PER_BLOCK, 1),
2734 shared_mem_bytes: 0,
2735 };
2736 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
2737 let __s_b = self.gpu.stream();
2738 let mut b = __s_b.launch_builder(&f);
2739 b.arg(q)
2740 .arg(k)
2741 .arg(v)
2742 .arg(g)
2743 .arg(beta)
2744 .arg(state_in)
2745 .arg(state_out)
2746 .arg(o)
2747 .arg(&h)
2748 .arg(acc)
2749 .arg(&b0)
2750 .arg(&tv)
2751 .arg(&scale);
2752 unsafe {
2753 b.launch(cfg)?;
2754 }
2755 Ok(())
2756 }
2757
2758 pub fn spec_rollback_kv(
2760 &self,
2761 len_ptrs: &CudaSlice<u64>,
2762 saved: &CudaSlice<i32>,
2763 acc: &CudaSlice<u32>,
2764 base: usize,
2765 n_layer: usize,
2766 ) -> Result<(), Box<dyn std::error::Error>> {
2767 let f = self.func("spec_rollback_kv");
2768 let (b, nl) = (base as i32, n_layer as i32);
2769 let cfg = LaunchConfig {
2770 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2771 block_dim: (64, 1, 1),
2772 shared_mem_bytes: 0,
2773 };
2774 let __s_bl = self.gpu.stream();
2775 let mut bl = __s_bl.launch_builder(&f);
2776 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
2777 unsafe {
2778 bl.launch(cfg)?;
2779 }
2780 Ok(())
2781 }
2782
2783 pub fn spec_fork_valid(
2785 &self,
2786 acc: &CudaSlice<u32>,
2787 optimistic_pending: u32,
2788 valid: &mut CudaSlice<u32>,
2789 ) -> Result<(), Box<dyn std::error::Error>> {
2790 let f = self.func("spec_fork_valid");
2791 let cfg = LaunchConfig {
2792 grid_dim: (1, 1, 1),
2793 block_dim: (1, 1, 1),
2794 shared_mem_bytes: 0,
2795 };
2796 let __s_bl = self.gpu.stream();
2797 let mut bl = __s_bl.launch_builder(&f);
2798 bl.arg(acc).arg(&optimistic_pending).arg(valid);
2799 unsafe {
2800 bl.launch(cfg)?;
2801 }
2802 Ok(())
2803 }
2804
2805 pub fn spec_fork_reconcile_kv(
2807 &self,
2808 len_ptrs: &CudaSlice<u64>,
2809 saved: &CudaSlice<i32>,
2810 acc: &CudaSlice<u32>,
2811 valid: &CudaSlice<u32>,
2812 base: usize,
2813 n_layer: usize,
2814 ) -> Result<(), Box<dyn std::error::Error>> {
2815 let f = self.func("spec_fork_reconcile_kv");
2816 let (b, nl) = (base as i32, n_layer as i32);
2817 let cfg = LaunchConfig {
2818 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
2819 block_dim: (64, 1, 1),
2820 shared_mem_bytes: 0,
2821 };
2822 let __s_bl = self.gpu.stream();
2823 let mut bl = __s_bl.launch_builder(&f);
2824 bl.arg(len_ptrs)
2825 .arg(saved)
2826 .arg(acc)
2827 .arg(valid)
2828 .arg(&b)
2829 .arg(&nl);
2830 unsafe {
2831 bl.launch(cfg)?;
2832 }
2833 Ok(())
2834 }
2835
2836 pub fn spec_fork_restore_f32(
2838 &self,
2839 snapshot: &CudaSlice<f32>,
2840 state: &mut CudaSlice<f32>,
2841 valid: &CudaSlice<u32>,
2842 ) -> Result<(), Box<dyn std::error::Error>> {
2843 assert_eq!(
2844 snapshot.len(),
2845 state.len(),
2846 "fork recurrent snapshot shape mismatch"
2847 );
2848 let f = self.func("spec_fork_restore_f32");
2849 let n = state.len() as i32;
2850 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
2851 let cfg = LaunchConfig {
2852 grid_dim: (blocks, 1, 1),
2853 block_dim: (256, 1, 1),
2854 shared_mem_bytes: 0,
2855 };
2856 let __s_bl = self.gpu.stream();
2857 let mut bl = __s_bl.launch_builder(&f);
2858 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
2859 unsafe {
2860 bl.launch(cfg)?;
2861 }
2862 Ok(())
2863 }
2864
2865 pub fn spec_seed_gather(
2868 &self,
2869 vx: &CudaSlice<f32>,
2870 fill_prev: &CudaSlice<f32>,
2871 acc: &CudaSlice<u32>,
2872 h_seed: &mut CudaSlice<f32>,
2873 base: usize,
2874 n_embd: usize,
2875 ) -> Result<(), Box<dyn std::error::Error>> {
2876 let f = self.func("spec_seed_gather");
2877 let (b, ne) = (base as i32, n_embd as i32);
2878 let cfg = LaunchConfig {
2879 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
2880 block_dim: (256, 1, 1),
2881 shared_mem_bytes: 0,
2882 };
2883 let __s_bl = self.gpu.stream();
2884 let mut bl = __s_bl.launch_builder(&f);
2885 bl.arg(vx)
2886 .arg(fill_prev)
2887 .arg(acc)
2888 .arg(h_seed)
2889 .arg(&b)
2890 .arg(&ne);
2891 unsafe {
2892 bl.launch(cfg)?;
2893 }
2894 Ok(())
2895 }
2896
2897 pub fn spec_accept_greedy(
2899 &self,
2900 preds: &CudaSlice<u32>,
2901 draft: &CudaSlice<u32>,
2902 last_pred: u32,
2903 base: usize,
2904 k_round: usize,
2905 out: &mut CudaSlice<u32>,
2906 ) -> Result<(), Box<dyn std::error::Error>> {
2907 let f = self.func("spec_accept_greedy");
2908 let (b, k) = (base as i32, k_round as i32);
2909 let cfg = LaunchConfig {
2910 grid_dim: (1, 1, 1),
2911 block_dim: (32, 1, 1),
2912 shared_mem_bytes: 0,
2913 };
2914 let __s_bl = self.gpu.stream();
2915 let mut bl = __s_bl.launch_builder(&f);
2916 bl.arg(preds)
2917 .arg(draft)
2918 .arg(&last_pred)
2919 .arg(&b)
2920 .arg(&k)
2921 .arg(out);
2922 unsafe {
2923 bl.launch(cfg)?;
2924 }
2925 Ok(())
2926 }
2927
2928 pub fn gumbel_perturb(
2935 &self,
2936 x: &CudaSlice<f32>,
2937 y: &mut CudaSlice<f32>,
2938 n: usize,
2939 seed: u64,
2940 stream_pos: u32,
2941 temp: f32,
2942 ) -> Result<(), Box<dyn std::error::Error>> {
2943 let f = self.func("gumbel_perturb_f32");
2944 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2945 let cfg = LaunchConfig {
2946 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2947 block_dim: (256, 1, 1),
2948 shared_mem_bytes: 0,
2949 };
2950 let __s_b = self.gpu.stream();
2951 let mut b = __s_b.launch_builder(&f);
2952 b.arg(x)
2953 .arg(&mut *y)
2954 .arg(&ni)
2955 .arg(&slo)
2956 .arg(&shi)
2957 .arg(&stream_pos)
2958 .arg(&temp);
2959 unsafe {
2960 b.launch(cfg)?;
2961 }
2962 Ok(())
2963 }
2964
2965 pub fn mask_logits_col(
2973 &self,
2974 logits: &mut CudaSlice<f32>,
2975 mask: &CudaSlice<u32>,
2976 col: usize,
2977 n: usize,
2978 mask_words: usize,
2979 ) -> Result<(), Box<dyn std::error::Error>> {
2980 let f = self.func("mask_logits_f32");
2981 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
2982 let cfg = LaunchConfig {
2983 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
2984 block_dim: (256, 1, 1),
2985 shared_mem_bytes: 0,
2986 };
2987 let __s_b = self.gpu.stream();
2988 let mut b = __s_b.launch_builder(&f);
2989 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
2990 unsafe {
2991 b.launch(cfg)?;
2992 }
2993 Ok(())
2994 }
2995
2996 pub fn gumbel_perturb_col(
3003 &self,
3004 x: &CudaSlice<f32>,
3005 col: usize,
3006 y: &mut CudaSlice<f32>,
3007 n: usize,
3008 seed: u64,
3009 stream_pos: u32,
3010 temp: f32,
3011 ) -> Result<(), Box<dyn std::error::Error>> {
3012 let f = self.func("gumbel_perturb_f32");
3013 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3014 let col_view = x.slice(col * n..(col + 1) * n);
3015 let cfg = LaunchConfig {
3016 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3017 block_dim: (256, 1, 1),
3018 shared_mem_bytes: 0,
3019 };
3020 let __s_b = self.gpu.stream();
3021 let mut b = __s_b.launch_builder(&f);
3022 b.arg(&col_view)
3023 .arg(&mut *y)
3024 .arg(&ni)
3025 .arg(&slo)
3026 .arg(&shi)
3027 .arg(&stream_pos)
3028 .arg(&temp);
3029 unsafe {
3030 b.launch(cfg)?;
3031 }
3032 Ok(())
3033 }
3034
3035 #[allow(clippy::too_many_arguments)]
3041 pub fn gumbel_perturb_filtered_col(
3042 &self,
3043 x: &CudaSlice<f32>,
3044 col: usize,
3045 y: &mut CudaSlice<f32>,
3046 n: usize,
3047 seed: u64,
3048 stream_pos: u32,
3049 temp: f32,
3050 stat_max: &CudaSlice<f32>,
3051 stat_th: &CudaSlice<f32>,
3052 stat_idx: usize,
3053 ) -> Result<(), Box<dyn std::error::Error>> {
3054 let f = self.func("gumbel_perturb_filtered_col_f32");
3055 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3056 let (ci, si) = (col as i32, stat_idx as i32);
3057 let cfg = LaunchConfig {
3058 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3059 block_dim: (256, 1, 1),
3060 shared_mem_bytes: 0,
3061 };
3062 let __s_b = self.gpu.stream();
3063 let mut b = __s_b.launch_builder(&f);
3064 b.arg(x)
3065 .arg(&ci)
3066 .arg(&mut *y)
3067 .arg(&ni)
3068 .arg(&slo)
3069 .arg(&shi)
3070 .arg(&stream_pos)
3071 .arg(&temp)
3072 .arg(stat_max)
3073 .arg(stat_th)
3074 .arg(&si);
3075 unsafe {
3076 b.launch(cfg)?;
3077 }
3078 Ok(())
3079 }
3080
3081 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3086 let f = self.func("memra_sctr_inc");
3087 let cfg = LaunchConfig {
3088 grid_dim: (1, 1, 1),
3089 block_dim: (1, 1, 1),
3090 shared_mem_bytes: 0,
3091 };
3092 let __s_b = self.gpu.stream();
3093 let mut b = __s_b.launch_builder(&f);
3094 b.arg(&mut *ctr);
3095 unsafe {
3096 b.launch(cfg)?;
3097 }
3098 Ok(())
3099 }
3100
3101 pub fn gumbel_perturb_ctr(
3106 &self,
3107 x: &CudaSlice<f32>,
3108 y: &mut CudaSlice<f32>,
3109 n: usize,
3110 seed: u64,
3111 ctr: &CudaSlice<u32>,
3112 temp: f32,
3113 ) -> Result<(), Box<dyn std::error::Error>> {
3114 let f = self.func("gumbel_perturb_ctr_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(ctr)
3129 .arg(&temp);
3130 unsafe {
3131 b.launch(cfg)?;
3132 }
3133 Ok(())
3134 }
3135
3136 pub fn softmax_gather(
3140 &self,
3141 x: &CudaSlice<f32>,
3142 row_stride: usize,
3143 ids: &CudaSlice<u32>,
3144 rows: &CudaSlice<i32>,
3145 out: &mut CudaSlice<f32>,
3146 n: usize,
3147 npair: usize,
3148 temp: f32,
3149 ) -> Result<(), Box<dyn std::error::Error>> {
3150 let f = self.func("softmax_gather_f32");
3151 let (ni, rs) = (n as i32, row_stride as i64);
3152 let np = npair as i32;
3153 let cfg = LaunchConfig {
3154 grid_dim: (npair 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(x)
3161 .arg(&rs)
3162 .arg(ids)
3163 .arg(rows)
3164 .arg(&mut *out)
3165 .arg(&ni)
3166 .arg(&np)
3167 .arg(&temp);
3168 unsafe {
3169 b.launch(cfg)?;
3170 }
3171 Ok(())
3172 }
3173
3174 pub fn residual_sample(
3178 &self,
3179 p: &CudaSlice<f32>,
3180 q: Option<&CudaSlice<f32>>,
3181 n: usize,
3182 temp: f32,
3183 seed: u64,
3184 stream_pos: u32,
3185 out_tok: &mut CudaSlice<u32>,
3186 ) -> Result<(), Box<dyn std::error::Error>> {
3187 let f = self.func("residual_sample_f32");
3188 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3189 let nth = 1024u32;
3190 let cfg = LaunchConfig {
3191 grid_dim: (1, 1, 1),
3192 block_dim: (nth, 1, 1),
3193 shared_mem_bytes: 0,
3194 };
3195 let has_q: i32 = q.is_some() as i32;
3196 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3198 let mut b = __s_b.launch_builder(&f);
3199 b.arg(p)
3200 .arg(qbuf)
3201 .arg(&has_q)
3202 .arg(&ni)
3203 .arg(&temp)
3204 .arg(&slo)
3205 .arg(&shi)
3206 .arg(&stream_pos)
3207 .arg(&mut *out_tok);
3208 unsafe {
3209 b.launch(cfg)?;
3210 }
3211 Ok(())
3212 }
3213
3214 pub fn with_moe_cache<R>(
3219 &self,
3220 max_block_bytes: usize,
3221 f: impl FnOnce(
3222 &mut crate::moe_cache::MoeSlotCache,
3223 &Engine,
3224 ) -> Result<R, Box<dyn std::error::Error>>,
3225 ) -> Result<R, Box<dyn std::error::Error>> {
3226 let mut guard = self.moe_cache.lock().unwrap();
3227 if guard.is_none() {
3228 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3229 }
3230 let cache = guard.as_mut().unwrap();
3231 f(cache, self)
3232 }
3233
3234 pub fn freeze_moe_cache(&self) {
3237 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3238 cache.freeze();
3239 }
3240 }
3241
3242 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3245 self.moe_cache
3246 .lock()
3247 .unwrap()
3248 .as_ref()
3249 .map(crate::moe_cache::MoeSlotCache::export_residency)
3250 }
3251
3252 pub(crate) fn moe_cache_frozen(&self) -> bool {
3253 self.moe_cache
3254 .lock()
3255 .unwrap()
3256 .as_ref()
3257 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3258 }
3259
3260 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3267 crate::cpu_experts::configured()
3268 && self.moe_cache_frozen()
3269 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3270 }
3271
3272 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3274 assert!(
3275 self.moe_cache.lock().unwrap().is_none(),
3276 "MoE cache layout configured after cache construction"
3277 );
3278 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3279 }
3280
3281 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3282 self.moe_cache_layout.lock().unwrap().clone()
3283 }
3284
3285 pub fn moe_cache_enabled() -> bool {
3287 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3288 }
3289
3290 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3293 let guard = self.moe_cache.lock().unwrap();
3294 guard
3295 .as_ref()
3296 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3297 }
3298
3299 pub fn cpu_expert_stats(
3303 &self,
3304 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3305 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3306 }
3307
3308 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3311 crate::cpu_experts::predictor_stats()
3312 }
3313
3314 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3315 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3316 }
3317
3318 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3321 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3322 }
3323
3324 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3327 let guard = self.moe_cache.lock().unwrap();
3328 guard
3329 .as_ref()
3330 .and_then(|cache| cache.pread_stats())
3331 .map(|stats| {
3332 (
3333 stats.reads,
3334 stats.bytes,
3335 stats.read_errors,
3336 stats.short_reads,
3337 stats.fallbacks,
3338 stats.buffer_waits,
3339 stats.ring_full,
3340 )
3341 })
3342 }
3343
3344 pub fn spill_config_fallbacks(&self) -> u64 {
3346 crate::spill_pread::config_fallbacks()
3347 }
3348
3349 pub fn moe_cache_reset_counters(&self) {
3351 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3352 c.reset_counters();
3353 }
3354 }
3355
3356 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3357 Ok(self.gpu.stream().clone_htod(v)?)
3358 }
3359
3360 pub fn htod_bytes_padded(
3364 &self,
3365 v: &[u8],
3366 pad: usize,
3367 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3368 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3369 {
3370 let mut view = d.slice_mut(0..v.len());
3371 self.gpu.stream().memcpy_htod(v, &mut view)?;
3372 }
3373 Ok(d)
3374 }
3375
3376 pub fn copy_into(
3378 &self,
3379 dst: &mut CudaSlice<f32>,
3380 off: usize,
3381 src: &CudaSlice<f32>,
3382 len: usize,
3383 ) -> Result<(), Box<dyn std::error::Error>> {
3384 let mut view = dst.slice_mut(off..off + len);
3385 self.gpu
3386 .stream()
3387 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3388 Ok(())
3389 }
3390
3391 pub fn copy_u8_into(
3394 &self,
3395 dst: &mut CudaSlice<u8>,
3396 off: usize,
3397 src: &CudaSlice<u8>,
3398 len: usize,
3399 ) -> Result<(), Box<dyn std::error::Error>> {
3400 let mut view = dst.slice_mut(off..off + len);
3401 self.gpu
3402 .stream()
3403 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3404 Ok(())
3405 }
3406
3407 pub fn copy_u8_range_into(
3409 &self,
3410 dst: &mut CudaSlice<u8>,
3411 dst_off: usize,
3412 src: &CudaSlice<u8>,
3413 src_off: usize,
3414 len: usize,
3415 ) -> Result<(), Box<dyn std::error::Error>> {
3416 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3417 self.gpu
3418 .stream()
3419 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3420 Ok(())
3421 }
3422
3423 pub fn prepare_kv_append(
3427 &self,
3428 kv: &mut crate::cache::KvLayer,
3429 retain_from: usize,
3430 append_rows: usize,
3431 ) -> Result<usize, Box<dyn std::error::Error>> {
3432 let Some(plan) = kv
3433 .ring
3434 .as_ref()
3435 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3436 .transpose()?
3437 else {
3438 return Ok(kv.len);
3439 };
3440 match plan {
3441 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3442 crate::cache::KvRingAppend::Rebase {
3443 src_row,
3444 keep_rows,
3445 new_base,
3446 write_row,
3447 } => {
3448 if keep_rows > 0 {
3449 let k_len = keep_rows * kv.k_tok_bytes;
3450 let v_len = keep_rows * kv.v_tok_bytes;
3451 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3452 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3453 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3454 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3455 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3456 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3457 }
3458 kv.ring.as_mut().unwrap().apply_rebase(new_base);
3459 Ok(write_row)
3460 }
3461 }
3462 }
3463
3464 pub fn htod_u8_into(
3467 &self,
3468 dst: &mut CudaSlice<u8>,
3469 off: usize,
3470 src: &[u8],
3471 ) -> Result<(), Box<dyn std::error::Error>> {
3472 let mut view = dst.slice_mut(off..off + src.len());
3473 self.gpu.stream().memcpy_htod(src, &mut view)?;
3474 Ok(())
3475 }
3476
3477 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3478 b.slice(0..len)
3479 }
3480
3481 pub fn view_u8_range<'a>(
3484 &self,
3485 b: &'a CudaSlice<u8>,
3486 start: usize,
3487 end: usize,
3488 ) -> cudarc::driver::CudaView<'a, u8> {
3489 b.slice(start..end)
3490 }
3491 pub fn view_u8<'a>(
3492 &self,
3493 b: &'a CudaSlice<u8>,
3494 len: usize,
3495 ) -> cudarc::driver::CudaView<'a, u8> {
3496 b.slice(0..len)
3497 }
3498
3499 pub fn append_kv_quantized(
3503 &self,
3504 k_row: &CudaSlice<f32>,
3505 v_row: &CudaSlice<f32>,
3506 kc: &mut CudaSlice<u8>,
3507 vc: &mut CudaSlice<u8>,
3508 t: usize,
3509 kv_dim_k: usize,
3510 kv_dim_v: usize,
3511 k_tok_bytes: usize,
3512 v_tok_bytes: usize,
3513 g: bool,
3514 ) -> Result<(), Box<dyn std::error::Error>> {
3515 let f = if g {
3516 self.func_g("append_quantize_kv_q8_0_q5_1")
3517 } else {
3518 self.func("append_quantize_kv_q8_0_q5_1")
3519 };
3520 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3521 let cfg = LaunchConfig {
3522 grid_dim: (nblk, 1, 1),
3523 block_dim: (32, 1, 1),
3524 shared_mem_bytes: 0,
3525 };
3526 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3527 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3528 let __s_b = self.gpu.stream();
3529 let mut b = __s_b.launch_builder(&f);
3530 b.arg(k_row)
3531 .arg(v_row)
3532 .arg(kc)
3533 .arg(vc)
3534 .arg(&ti)
3535 .arg(&kdk)
3536 .arg(&kdv)
3537 .arg(&ktb)
3538 .arg(&vtb);
3539 unsafe {
3540 b.launch(cfg)?;
3541 }
3542 Ok(())
3543 }
3544
3545 pub fn append_kv_quantized_dc(
3549 &self,
3550 k_row: &CudaSlice<f32>,
3551 v_row: &CudaSlice<f32>,
3552 kc: &mut CudaSlice<u8>,
3553 vc: &mut CudaSlice<u8>,
3554 t_dev: &CudaSlice<i32>,
3555 kv_dim_k: usize,
3556 kv_dim_v: usize,
3557 k_tok_bytes: usize,
3558 v_tok_bytes: usize,
3559 g: bool,
3560 ) -> Result<(), Box<dyn std::error::Error>> {
3561 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3562 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3563 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3564 if Self::pdl_on() && Self::pdl_wb_on() {
3566 use cudarc::driver::{DevicePtr, DevicePtrMut};
3567 let s = &self.gpu.stream();
3568 let (pk, _g0) = k_row.device_ptr(s);
3569 let (pv, _g1) = v_row.device_ptr(s);
3570 let (pkc, _g2) = kc.device_ptr_mut(s);
3571 let (pvc, _g3) = vc.device_ptr_mut(s);
3572 let (pt, _g4) = t_dev.device_ptr(s);
3573 let mut ps = [
3574 &pk as *const _ as *mut std::ffi::c_void,
3575 &pv as *const _ as *mut _,
3576 &pkc as *const _ as *mut _,
3577 &pvc as *const _ as *mut _,
3578 &pt as *const _ as *mut _,
3579 &kdk as *const _ as *mut _,
3580 &kdv as *const _ as *mut _,
3581 &ktb as *const _ as *mut _,
3582 &vtb as *const _ as *mut _,
3583 ];
3584 unsafe {
3585 self.launch_pdl_flash(
3586 g,
3587 "append_quantize_kv_q8_0_q5_1_dc",
3588 (nblk, 1, 1),
3589 (32, 1, 1),
3590 0,
3591 &mut ps,
3592 )?;
3593 }
3594 return Ok(());
3595 }
3596 let f = if g {
3597 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3598 } else {
3599 self.func("append_quantize_kv_q8_0_q5_1_dc")
3600 };
3601 let cfg = LaunchConfig {
3602 grid_dim: (nblk, 1, 1),
3603 block_dim: (32, 1, 1),
3604 shared_mem_bytes: 0,
3605 };
3606 let __s_b = self.gpu.stream();
3607 let mut b = __s_b.launch_builder(&f);
3608 b.arg(k_row)
3609 .arg(v_row)
3610 .arg(kc)
3611 .arg(vc)
3612 .arg(t_dev)
3613 .arg(&kdk)
3614 .arg(&kdv)
3615 .arg(&ktb)
3616 .arg(&vtb);
3617 unsafe {
3618 b.launch(cfg)?;
3619 }
3620 Ok(())
3621 }
3622
3623 #[allow(clippy::too_many_arguments)]
3630 pub fn append_kv_quantized_rows(
3631 &self,
3632 k_rows: &CudaSlice<f32>,
3633 v_rows: &CudaSlice<f32>,
3634 kc: &mut CudaSlice<u8>,
3635 vc: &mut CudaSlice<u8>,
3636 t0: usize,
3637 t: usize,
3638 kv_dim_k: usize,
3639 kv_dim_v: usize,
3640 k_tok_bytes: usize,
3641 v_tok_bytes: usize,
3642 g: bool,
3643 ) -> Result<(), Box<dyn std::error::Error>> {
3644 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3645 for i in 0..t {
3646 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3647 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3648 self.append_kv_quantized_view(
3649 &k_row,
3650 &v_row,
3651 kc,
3652 vc,
3653 t0 + i,
3654 kv_dim_k,
3655 kv_dim_v,
3656 k_tok_bytes,
3657 v_tok_bytes,
3658 g,
3659 )?;
3660 }
3661 return Ok(());
3662 }
3663 let f = if g {
3664 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
3665 } else {
3666 self.func("append_quantize_kv_q8_0_q5_1_rows")
3667 };
3668 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3669 let cfg = LaunchConfig {
3670 grid_dim: (nblk, t as u32, 1),
3671 block_dim: (32, 1, 1),
3672 shared_mem_bytes: 0,
3673 };
3674 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
3675 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3676 let __s_b = self.gpu.stream();
3677 let mut b = __s_b.launch_builder(&f);
3678 b.arg(k_rows)
3679 .arg(v_rows)
3680 .arg(kc)
3681 .arg(vc)
3682 .arg(&t0i)
3683 .arg(&kdk)
3684 .arg(&kdv)
3685 .arg(&ktb)
3686 .arg(&vtb);
3687 unsafe {
3688 b.launch(cfg)?;
3689 }
3690 Ok(())
3691 }
3692
3693 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
3697 let f = self.func("inc_i32");
3698 let cfg = LaunchConfig {
3699 grid_dim: (1, 1, 1),
3700 block_dim: (1, 1, 1),
3701 shared_mem_bytes: 0,
3702 };
3703 let __s_b = self.gpu.stream();
3704 let mut b = __s_b.launch_builder(&f);
3705 b.arg(p);
3706 unsafe {
3707 b.launch(cfg)?;
3708 }
3709 Ok(())
3710 }
3711
3712 pub fn append_kv_quantized_view(
3715 &self,
3716 k_row: &cudarc::driver::CudaView<f32>,
3717 v_row: &cudarc::driver::CudaView<f32>,
3718 kc: &mut CudaSlice<u8>,
3719 vc: &mut CudaSlice<u8>,
3720 t: usize,
3721 kv_dim_k: usize,
3722 kv_dim_v: usize,
3723 k_tok_bytes: usize,
3724 v_tok_bytes: usize,
3725 g: bool,
3726 ) -> Result<(), Box<dyn std::error::Error>> {
3727 let f = if g {
3728 self.func_g("append_quantize_kv_q8_0_q5_1")
3729 } else {
3730 self.func("append_quantize_kv_q8_0_q5_1")
3731 };
3732 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3733 let cfg = LaunchConfig {
3734 grid_dim: (nblk, 1, 1),
3735 block_dim: (32, 1, 1),
3736 shared_mem_bytes: 0,
3737 };
3738 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3739 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3740 let __s_b = self.gpu.stream();
3741 let mut b = __s_b.launch_builder(&f);
3742 b.arg(k_row)
3743 .arg(v_row)
3744 .arg(kc)
3745 .arg(vc)
3746 .arg(&ti)
3747 .arg(&kdk)
3748 .arg(&kdv)
3749 .arg(&ktb)
3750 .arg(&vtb);
3751 unsafe {
3752 b.launch(cfg)?;
3753 }
3754 Ok(())
3755 }
3756
3757 pub fn copy_view_into(
3760 &self,
3761 dst: &mut CudaSlice<f32>,
3762 off: usize,
3763 src: &cudarc::driver::CudaView<f32>,
3764 len: usize,
3765 ) -> Result<(), Box<dyn std::error::Error>> {
3766 let mut view = dst.slice_mut(off..off + len);
3767 self.gpu
3768 .stream()
3769 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3770 Ok(())
3771 }
3772
3773 pub fn clone_dtod(
3777 &self,
3778 src: &CudaSlice<f32>,
3779 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3780 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
3781 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
3782 Ok(dst)
3783 }
3784
3785 pub fn dtod_copy_view(
3788 &self,
3789 src: &cudarc::driver::CudaView<f32>,
3790 dst: &mut CudaSlice<f32>,
3791 ) -> Result<(), Box<dyn std::error::Error>> {
3792 self.gpu.stream().memcpy_dtod(src, dst)?;
3793 Ok(())
3794 }
3795
3796 pub fn dtod_copy_view_i8(
3798 &self,
3799 src: &cudarc::driver::CudaView<i8>,
3800 dst: &mut CudaSlice<i8>,
3801 ) -> Result<(), Box<dyn std::error::Error>> {
3802 self.gpu.stream().memcpy_dtod(src, dst)?;
3803 Ok(())
3804 }
3805
3806 pub fn dtod_copy_into(
3808 &self,
3809 src: &CudaSlice<f32>,
3810 dst: &mut CudaSlice<f32>,
3811 offset: usize,
3812 ) -> Result<(), Box<dyn std::error::Error>> {
3813 let n = src.len();
3814 let mut dv = dst.slice_mut(offset..offset + n);
3815 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
3816 Ok(())
3817 }
3818
3819 pub fn copy_batch_uniform_f32(
3825 &self,
3826 table: &CudaSlice<u64>,
3827 n: usize,
3828 words: usize,
3829 ) -> Result<(), Box<dyn std::error::Error>> {
3830 if n == 0 || words == 0 {
3831 return Ok(());
3832 }
3833 debug_assert!(
3834 table.len() >= 2 * n,
3835 "pointer table must hold n srcs + n dsts"
3836 );
3837 let f = self.func("copy_batch_uniform_f32");
3838 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
3841 let (ni, wi) = (n as i32, words as i32);
3842 let cfg = LaunchConfig {
3843 grid_dim: (chunks, n as u32, 1),
3844 block_dim: (256, 1, 1),
3845 shared_mem_bytes: 0,
3846 };
3847 let __s = self.gpu.stream();
3848 let mut b = __s.launch_builder(&f);
3849 b.arg(table).arg(&ni).arg(&wi);
3850 unsafe {
3851 b.launch(cfg)?;
3852 }
3853 Ok(())
3854 }
3855
3856 pub fn htod_u64_into(
3859 &self,
3860 v: &[u64],
3861 dst: &mut CudaSlice<u64>,
3862 ) -> Result<(), Box<dyn std::error::Error>> {
3863 let mut view = dst.slice_mut(0..v.len());
3864 self.gpu.stream().memcpy_htod(v, &mut view)?;
3865 Ok(())
3866 }
3867
3868 pub fn copy_indirect_src_f32(
3873 &self,
3874 src_entry: &cudarc::driver::CudaView<u64>,
3875 dst: &mut CudaSlice<f32>,
3876 dst_off: usize,
3877 words: usize,
3878 ) -> Result<(), Box<dyn std::error::Error>> {
3879 let f = self.func("copy_indirect_src_f32");
3880 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
3881 let wi = words as i32;
3882 let cfg = LaunchConfig {
3883 grid_dim: (chunks, 1, 1),
3884 block_dim: (256, 1, 1),
3885 shared_mem_bytes: 0,
3886 };
3887 let mut dv = dst.slice_mut(dst_off..dst_off + words);
3888 let __s = self.gpu.stream();
3889 let mut b = __s.launch_builder(&f);
3890 b.arg(src_entry).arg(&mut dv).arg(&wi);
3891 unsafe {
3892 b.launch(cfg)?;
3893 }
3894 Ok(())
3895 }
3896
3897 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
3899 self.alloc_uninit::<i8>(n)
3900 }
3901
3902 pub fn qmatvec(
3904 &self,
3905 w: &CudaSlice<u8>,
3906 x: &CudaSlice<f32>,
3907 m: usize,
3908 in_f: usize,
3909 out_f: usize,
3910 qtype: i32,
3911 row_bytes: usize,
3912 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3913 let f = self.func("qmatvec_f32");
3914 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
3916 grid_dim: (out_f as u32, m as u32, 1),
3917 block_dim: (256, 1, 1),
3918 shared_mem_bytes: 0,
3919 };
3920 let (inf, outf, mi, qt, rb) =
3921 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
3922 let __s_b = self.gpu.stream();
3923 let mut b = __s_b.launch_builder(&f);
3924 b.arg(w)
3925 .arg(x)
3926 .arg(&mut y)
3927 .arg(&inf)
3928 .arg(&outf)
3929 .arg(&mi)
3930 .arg(&qt)
3931 .arg(&rb);
3932 unsafe {
3933 b.launch(cfg)?;
3934 }
3935 Ok(y)
3936 }
3937
3938 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3940 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
3941 self.keep_if_capturing(&s);
3942 Ok(s)
3943 }
3944
3945 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3949 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
3950 self.keep_if_capturing(&s);
3951 Ok(s)
3952 }
3953
3954 pub fn memset_zeros_view(
3957 &self,
3958 dst: &mut cudarc::driver::CudaViewMut<f32>,
3959 ) -> Result<(), Box<dyn std::error::Error>> {
3960 self.gpu.stream().memset_zeros(dst)?;
3961 Ok(())
3962 }
3963
3964 pub fn stage_expert(
3970 &self,
3971 host_bytes: &[u8],
3972 scratch: &mut CudaSlice<u8>,
3973 off: usize,
3974 ) -> Result<(), Box<dyn std::error::Error>> {
3975 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
3978 }
3979
3980 pub fn moe_router_topk(
3986 &self,
3987 logits: &CudaSlice<f32>,
3988 t: usize,
3989 n_expert: usize,
3990 n_used: usize,
3991 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3992 let f = self.func("moe_router_topk_f32");
3993 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 {
3996 grid_dim: (t as u32, 1, 1),
3997 block_dim: (n_expert as u32, 1, 1),
3998 shared_mem_bytes: 0,
3999 };
4000 let (ne, nu) = (n_expert as i32, n_used as i32);
4001 let __s_b = self.gpu.stream();
4002 let mut b = __s_b.launch_builder(&f);
4003 b.arg(logits)
4004 .arg(&mut sel_idx)
4005 .arg(&mut sel_w)
4006 .arg(&ne)
4007 .arg(&nu);
4008 unsafe {
4009 b.launch(cfg)?;
4010 }
4011 Ok((sel_idx, sel_w))
4012 }
4013
4014 pub fn moe_router_topk_scaled(
4017 &self,
4018 logits: &CudaSlice<f32>,
4019 t: usize,
4020 n_expert: usize,
4021 n_used: usize,
4022 ex_scale: &CudaSlice<f32>,
4023 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4024 let f = self.func("moe_router_topk_scaled_f32");
4029 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4030 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4031 let cfg = LaunchConfig {
4032 grid_dim: (t as u32, 1, 1),
4033 block_dim: (n_expert as u32, 1, 1),
4034 shared_mem_bytes: 0,
4035 };
4036 let (ne, nu) = (n_expert as i32, n_used as i32);
4037 let __s_b = self.gpu.stream();
4038 let mut b = __s_b.launch_builder(&f);
4039 b.arg(logits)
4040 .arg(&mut sel_idx)
4041 .arg(&mut sel_w)
4042 .arg(&ne)
4043 .arg(&nu)
4044 .arg(ex_scale);
4045 unsafe {
4046 b.launch(cfg)?;
4047 }
4048 Ok((sel_idx, sel_w))
4049 }
4050
4051 pub fn moe_router_topk_host(
4059 &self,
4060 logits: &CudaSlice<f32>,
4061 t: usize,
4062 n_expert: usize,
4063 n_used: usize,
4064 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4065 let f = self.func("moe_router_topk_f32");
4066 let n = t * n_used;
4067 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4068 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4069 let cfg = LaunchConfig {
4070 grid_dim: (t as u32, 1, 1),
4071 block_dim: (n_expert as u32, 1, 1),
4072 shared_mem_bytes: 0,
4073 };
4074 let (ne, nu) = (n_expert as i32, n_used as i32);
4075 let __s_b = self.gpu.stream();
4076 let mut b = __s_b.launch_builder(&f);
4077 b.arg(logits)
4078 .arg(&mut sel_idx)
4079 .arg(&mut sel_w)
4080 .arg(&ne)
4081 .arg(&nu);
4082 unsafe {
4083 b.launch(cfg)?;
4084 }
4085 let bytes = n * 8;
4087 let mut guard = self.router_stage.lock().unwrap();
4088 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4089 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4090 }
4091 let stage = guard.as_mut().unwrap();
4092 let (si, sw) = unsafe {
4093 (
4094 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4095 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4096 )
4097 };
4098 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()))
4102 }
4103
4104 #[allow(clippy::too_many_arguments)]
4108 pub fn moe_router_sigmoid_topk(
4109 &self,
4110 logits: &CudaSlice<f32>,
4111 t: usize,
4112 n_expert: usize,
4113 n_used: usize,
4114 active_count: usize,
4115 correction_bias: &CudaSlice<f32>,
4116 active: &CudaSlice<u8>,
4117 scaling_factor: f32,
4118 route_norm: bool,
4119 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4120 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4121 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4122 return Err(format!(
4123 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4124 )
4125 .into());
4126 }
4127 if logits.len() < t * n_expert
4128 || correction_bias.len() != n_expert
4129 || active.len() != n_expert
4130 {
4131 return Err(format!(
4132 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4133 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4134 ).into());
4135 }
4136 let f = self.func("moe_router_sigmoid_topk_f32");
4137 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4138 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4139 let threads = n_expert.div_ceil(32) * 32;
4140 let cfg = LaunchConfig {
4141 grid_dim: (t as u32, 1, 1),
4142 block_dim: (threads as u32, 1, 1),
4143 shared_mem_bytes: 0,
4144 };
4145 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4146 let __s_b = self.gpu.stream();
4147 let mut b = __s_b.launch_builder(&f);
4148 b.arg(logits)
4149 .arg(correction_bias)
4150 .arg(active)
4151 .arg(&mut sel_idx)
4152 .arg(&mut sel_w)
4153 .arg(&ne)
4154 .arg(&nu)
4155 .arg(&scaling_factor)
4156 .arg(&rn);
4157 unsafe {
4158 b.launch(cfg)?;
4159 }
4160 Ok((sel_idx, sel_w))
4161 }
4162
4163 #[allow(clippy::too_many_arguments)]
4166 pub fn moe_router_sigmoid_topk_host(
4167 &self,
4168 logits: &CudaSlice<f32>,
4169 t: usize,
4170 n_expert: usize,
4171 n_used: usize,
4172 active_count: usize,
4173 correction_bias: &CudaSlice<f32>,
4174 active: &CudaSlice<u8>,
4175 scaling_factor: f32,
4176 route_norm: bool,
4177 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4178 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4179 logits,
4180 t,
4181 n_expert,
4182 n_used,
4183 active_count,
4184 correction_bias,
4185 active,
4186 scaling_factor,
4187 route_norm,
4188 )?;
4189 let n = t * n_used;
4190 let bytes = n * 8;
4191 let mut guard = self.router_stage.lock().unwrap();
4192 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4193 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4194 }
4195 let stage = guard.as_mut().unwrap();
4196 let (si, sw) = unsafe {
4197 (
4198 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4199 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4200 )
4201 };
4202 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4203 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4204 self.gpu.stream().synchronize()?;
4205 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4206 }
4207
4208 pub fn stage_expert_async(
4212 &self,
4213 host_bytes: &[u8],
4214 scratch: &mut CudaSlice<u8>,
4215 off: usize,
4216 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4217 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4218 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4219 Ok(self.copy_stream.record_event(None)?)
4220 }
4221
4222 pub fn compute_wait(
4224 &self,
4225 ev: &cudarc::driver::CudaEvent,
4226 ) -> Result<(), Box<dyn std::error::Error>> {
4227 self.gpu.stream().wait(ev)?;
4228 Ok(())
4229 }
4230
4231 pub fn qmatvec_view(
4236 &self,
4237 w: &CudaSlice<u8>,
4238 range: std::ops::Range<usize>,
4239 x: &cudarc::driver::CudaView<f32>,
4240 m: usize,
4241 in_f: usize,
4242 out_f: usize,
4243 qtype: i32,
4244 row_bytes: usize,
4245 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4246 let f = self.func("qmatvec_f32");
4247 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4250 grid_dim: (out_f as u32, m as u32, 1),
4251 block_dim: (256, 1, 1),
4252 shared_mem_bytes: 0,
4253 };
4254 let (inf, outf, mi, qt, rb) =
4255 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4256 let __s_b = self.gpu.stream();
4257 let mut b = __s_b.launch_builder(&f);
4258 b.arg(&wv)
4259 .arg(x)
4260 .arg(&mut y)
4261 .arg(&inf)
4262 .arg(&outf)
4263 .arg(&mi)
4264 .arg(&qt)
4265 .arg(&rb);
4266 unsafe {
4267 b.launch(cfg)?;
4268 }
4269 Ok(y)
4270 }
4271
4272 #[allow(clippy::too_many_arguments)]
4279 pub fn moe_gate_up_silu8_q8(
4283 &self,
4284 gp: WPtr8,
4285 up: WPtr8,
4286 aq: &CudaSlice<i8>,
4287 ad: &CudaSlice<f32>,
4288 in_f: usize,
4289 n_ff: usize,
4290 n_used: usize,
4291 qt_g: i32,
4292 qt_u: i32,
4293 rb_g: usize,
4294 rb_u: usize,
4295 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4296 let f = self.func("moe_gate_up_silu8_q8");
4297 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4298 let cfg = LaunchConfig {
4299 grid_dim: (n_ff as u32, n_used as u32, 1),
4300 block_dim: (32, 1, 1),
4301 shared_mem_bytes: 0,
4302 };
4303 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4304 let __s_b = self.gpu.stream();
4305 let mut b = __s_b.launch_builder(&f);
4306 b.arg(&gp)
4307 .arg(&up)
4308 .arg(aq)
4309 .arg(ad)
4310 .arg(&mut act)
4311 .arg(&inf)
4312 .arg(&nff)
4313 .arg(&qt_g)
4314 .arg(&qt_u)
4315 .arg(&rbg)
4316 .arg(&rbu);
4317 unsafe {
4318 b.launch(cfg)?;
4319 }
4320 Ok(act)
4321 }
4322
4323 #[allow(clippy::too_many_arguments)]
4324 pub fn moe_down8_fma_q8(
4325 &self,
4326 dp: WPtr8,
4327 w: F32x8,
4328 aq2: &CudaSlice<i8>,
4329 ad2: &CudaSlice<f32>,
4330 dst: &mut cudarc::driver::CudaViewMut<f32>,
4331 in_f: usize,
4332 out_f: usize,
4333 n_used: usize,
4334 qt: i32,
4335 rb: usize,
4336 ) -> Result<(), Box<dyn std::error::Error>> {
4337 let f = self.func("moe_down8_fma_q8");
4338 let cfg = LaunchConfig {
4339 grid_dim: (out_f as u32, 1, 1),
4340 block_dim: (32, 1, 1),
4341 shared_mem_bytes: 0,
4342 };
4343 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4344 let __s_b = self.gpu.stream();
4345 let mut b = __s_b.launch_builder(&f);
4346 b.arg(&dp)
4347 .arg(&w)
4348 .arg(aq2)
4349 .arg(ad2)
4350 .arg(dst)
4351 .arg(&inf)
4352 .arg(&outf)
4353 .arg(&nu)
4354 .arg(&qt)
4355 .arg(&rbi);
4356 unsafe {
4357 b.launch(cfg)?;
4358 }
4359 Ok(())
4360 }
4361
4362 pub fn qmatvec_expert_q8(
4364 &self,
4365 w: &CudaSlice<u8>,
4366 range: std::ops::Range<usize>,
4367 aq: &CudaSlice<i8>,
4368 ad: &CudaSlice<f32>,
4369 m: usize,
4370 in_f: usize,
4371 out_f: usize,
4372 qtype: i32,
4373 row_bytes: usize,
4374 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4375 let f = self.func("qmatvec_expert_q8");
4376 let wv = w.slice(range);
4377 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4378 const ROWS: u32 = 4; let cfg = LaunchConfig {
4380 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4381 block_dim: (32, ROWS, 1),
4382 shared_mem_bytes: 0,
4383 };
4384 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4385 let __s_b = self.gpu.stream();
4386 let mut b = __s_b.launch_builder(&f);
4387 b.arg(&wv)
4388 .arg(aq)
4389 .arg(ad)
4390 .arg(&mut y)
4391 .arg(&inf)
4392 .arg(&outf)
4393 .arg(&mi)
4394 .arg(&qtype)
4395 .arg(&rbi);
4396 unsafe {
4397 b.launch(cfg)?;
4398 }
4399 Ok(y)
4400 }
4401
4402 pub fn moe_gate_up_silu8(
4403 &self,
4404 gp: WPtr8,
4405 up: WPtr8,
4406 x: &cudarc::driver::CudaView<f32>,
4407 in_f: usize,
4408 n_ff: usize,
4409 n_used: usize,
4410 qt_g: i32,
4411 qt_u: i32,
4412 rb_g: usize,
4413 rb_u: usize,
4414 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4415 let f = self.func("moe_gate_up_silu8_f32");
4416 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
4418 grid_dim: (n_ff as u32, n_used as u32, 1),
4419 block_dim: (256, 1, 1),
4420 shared_mem_bytes: 0,
4421 };
4422 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4423 let __s_b = self.gpu.stream();
4424 let mut b = __s_b.launch_builder(&f);
4425 b.arg(&gp)
4426 .arg(&up)
4427 .arg(x)
4428 .arg(&mut act)
4429 .arg(&inf)
4430 .arg(&nff)
4431 .arg(&qt_g)
4432 .arg(&qt_u)
4433 .arg(&rbg)
4434 .arg(&rbu);
4435 unsafe {
4436 b.launch(cfg)?;
4437 }
4438 Ok(act)
4439 }
4440
4441 #[allow(clippy::too_many_arguments)]
4447 pub fn moe_down8_fma_into(
4448 &self,
4449 dp: WPtr8,
4450 w: F32x8,
4451 act: &CudaSlice<f32>,
4452 dst: &mut cudarc::driver::CudaViewMut<f32>,
4453 in_f: usize,
4454 out_f: usize,
4455 n_used: usize,
4456 qt: i32,
4457 rb: usize,
4458 ) -> Result<(), Box<dyn std::error::Error>> {
4459 let f = self.func("moe_down8_fma_f32");
4460 let cfg = LaunchConfig {
4461 grid_dim: (out_f as u32, 1, 1),
4462 block_dim: (256, 1, 1),
4463 shared_mem_bytes: 0,
4464 };
4465 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4466 let __s_b = self.gpu.stream();
4467 let mut b = __s_b.launch_builder(&f);
4468 b.arg(&dp)
4469 .arg(&w)
4470 .arg(act)
4471 .arg(dst)
4472 .arg(&inf)
4473 .arg(&outf)
4474 .arg(&nu)
4475 .arg(&qt)
4476 .arg(&rbv);
4477 unsafe {
4478 b.launch(cfg)?;
4479 }
4480 Ok(())
4481 }
4482
4483 #[allow(clippy::too_many_arguments)]
4488 #[allow(clippy::too_many_arguments)]
4503 #[allow(clippy::too_many_arguments)]
4505 pub fn moe_pairs_matvec_q8(
4506 &self,
4507 table: &CudaSlice<u64>,
4508 proj: i32,
4509 pair_tok: &CudaSlice<i32>,
4510 pair_ex: &CudaSlice<i32>,
4511 aq: &CudaSlice<i8>,
4512 ad: &CudaSlice<f32>,
4513 in_f: usize,
4514 out_f: usize,
4515 n_expert: usize,
4516 n_pairs: usize,
4517 qtype: i32,
4518 row_bytes: usize,
4519 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4520 let f = self.func("moe_pairs_matvec_q8");
4521 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4522 const ROWS: u32 = 4;
4523 let cfg = LaunchConfig {
4524 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4525 block_dim: (32, ROWS, 1),
4526 shared_mem_bytes: 0,
4527 };
4528 let (inf, outf, ne, np, rbi) = (
4529 in_f as i32,
4530 out_f as i32,
4531 n_expert as i32,
4532 n_pairs as i32,
4533 row_bytes as i64,
4534 );
4535 let __s_b = self.gpu.stream();
4536 let mut b = __s_b.launch_builder(&f);
4537 b.arg(table)
4538 .arg(&proj)
4539 .arg(pair_tok)
4540 .arg(pair_ex)
4541 .arg(aq)
4542 .arg(ad)
4543 .arg(&mut y)
4544 .arg(&inf)
4545 .arg(&outf)
4546 .arg(&ne)
4547 .arg(&np)
4548 .arg(&qtype)
4549 .arg(&rbi);
4550 unsafe {
4551 b.launch(cfg)?;
4552 }
4553 Ok(y)
4554 }
4555
4556 #[allow(clippy::too_many_arguments)]
4558 pub fn moe_pairs_matvec_q8_em(
4559 &self,
4560 table: &CudaSlice<u64>,
4561 proj: i32,
4562 ex_ids: &CudaSlice<i32>,
4563 ex_off: &CudaSlice<i32>,
4564 ex_pairs: &CudaSlice<i32>,
4565 pair_tok: &CudaSlice<i32>,
4566 aq: &CudaSlice<i8>,
4567 ad: &CudaSlice<f32>,
4568 in_f: usize,
4569 out_f: usize,
4570 n_expert: usize,
4571 n_active: usize,
4572 n_pairs: usize,
4573 qtype: i32,
4574 row_bytes: usize,
4575 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4576 let f = self.func("moe_pairs_matvec_q8_em");
4577 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4578 const ROWS: u32 = 4;
4579 let cfg = LaunchConfig {
4580 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4581 block_dim: (32, ROWS, 1),
4582 shared_mem_bytes: 0,
4583 };
4584 let (inf, outf, ne, na, rbi) = (
4585 in_f as i32,
4586 out_f as i32,
4587 n_expert as i32,
4588 n_active as i32,
4589 row_bytes as i64,
4590 );
4591 let __s_b = self.gpu.stream();
4592 let mut b = __s_b.launch_builder(&f);
4593 b.arg(table)
4594 .arg(&proj)
4595 .arg(ex_ids)
4596 .arg(ex_off)
4597 .arg(ex_pairs)
4598 .arg(pair_tok)
4599 .arg(aq)
4600 .arg(ad)
4601 .arg(&mut y)
4602 .arg(&inf)
4603 .arg(&outf)
4604 .arg(&ne)
4605 .arg(&na)
4606 .arg(&qtype)
4607 .arg(&rbi);
4608 unsafe {
4609 b.launch(cfg)?;
4610 }
4611 Ok(y)
4612 }
4613
4614 #[allow(clippy::too_many_arguments)]
4617 pub fn moe_pairs_matvec_q8_dec(
4618 &self,
4619 table: &CudaSlice<u64>,
4620 proj: i32,
4621 ex_ids: &CudaSlice<i32>,
4622 ex_off: &CudaSlice<i32>,
4623 ex_pairs: &CudaSlice<i32>,
4624 pair_tok: &CudaSlice<i32>,
4625 aq: &CudaSlice<i8>,
4626 ad: &CudaSlice<f32>,
4627 in_f: usize,
4628 out_f: usize,
4629 n_expert: usize,
4630 n_active: usize,
4631 n_pairs: usize,
4632 qtype: i32,
4633 row_bytes: usize,
4634 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4635 let f = self.func("moe_pairs_matvec_q8_dec");
4636 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4637 const ROWS: u32 = 4;
4638 let cfg = LaunchConfig {
4639 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
4640 block_dim: (32, ROWS, 1),
4641 shared_mem_bytes: 0,
4642 };
4643 let (inf, outf, ne, na, rbi) = (
4644 in_f as i32,
4645 out_f as i32,
4646 n_expert as i32,
4647 n_active as i32,
4648 row_bytes as i64,
4649 );
4650 let __s_b = self.gpu.stream();
4651 let mut b = __s_b.launch_builder(&f);
4652 b.arg(table)
4653 .arg(&proj)
4654 .arg(ex_ids)
4655 .arg(ex_off)
4656 .arg(ex_pairs)
4657 .arg(pair_tok)
4658 .arg(aq)
4659 .arg(ad)
4660 .arg(&mut y)
4661 .arg(&inf)
4662 .arg(&outf)
4663 .arg(&ne)
4664 .arg(&na)
4665 .arg(&qtype)
4666 .arg(&rbi);
4667 unsafe {
4668 b.launch(cfg)?;
4669 }
4670 Ok(y)
4671 }
4672
4673 pub fn moe_pairs_gelu_mul(
4674 &self,
4675 gate: &CudaSlice<f32>,
4676 up: &CudaSlice<f32>,
4677 n: usize,
4678 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4679 let f = self.func("moe_pairs_gelu_mul");
4680 let mut act = self.alloc_uninit::<f32>(n)?;
4681 let cfg = LaunchConfig::for_num_elems(n as u32);
4682 let nl = n as i64;
4683 let __s_b = self.gpu.stream();
4684 let mut b = __s_b.launch_builder(&f);
4685 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4686 unsafe {
4687 b.launch(cfg)?;
4688 }
4689 Ok(act)
4690 }
4691
4692 pub fn moe_pairs_silu_mul(
4693 &self,
4694 gate: &CudaSlice<f32>,
4695 up: &CudaSlice<f32>,
4696 n: usize,
4697 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4698 let f = self.func("moe_pairs_silu_mul");
4699 let mut act = self.alloc_uninit::<f32>(n)?;
4700 let cfg = LaunchConfig::for_num_elems(n as u32);
4701 let nl = n as i64;
4702 let __s_b = self.gpu.stream();
4703 let mut b = __s_b.launch_builder(&f);
4704 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
4705 unsafe {
4706 b.launch(cfg)?;
4707 }
4708 Ok(act)
4709 }
4710
4711 #[allow(clippy::too_many_arguments)]
4712 pub fn moe_pairs_scatter(
4713 &self,
4714 y_down: &CudaSlice<f32>,
4715 pair_w: &CudaSlice<f32>,
4716 tok_pair_off: &CudaSlice<i32>,
4717 tok_pair_ids: &CudaSlice<i32>,
4718 moe_out: &mut CudaSlice<f32>,
4719 t: usize,
4720 n_embd: usize,
4721 ) -> Result<(), Box<dyn std::error::Error>> {
4722 let f = self.func("moe_pairs_scatter");
4723 let cfg = LaunchConfig {
4724 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
4725 block_dim: (256, 1, 1),
4726 shared_mem_bytes: 0,
4727 };
4728 let ne = n_embd as i32;
4729 let __s_b = self.gpu.stream();
4730 let mut b = __s_b.launch_builder(&f);
4731 b.arg(y_down)
4732 .arg(pair_w)
4733 .arg(tok_pair_off)
4734 .arg(tok_pair_ids)
4735 .arg(moe_out)
4736 .arg(&ne);
4737 unsafe {
4738 b.launch(cfg)?;
4739 }
4740 Ok(())
4741 }
4742
4743 #[allow(clippy::too_many_arguments)]
4747 pub fn moe_gate_up_gelu8_dev_q8(
4748 &self,
4749 table: &CudaSlice<u64>,
4750 sel: &cudarc::driver::CudaView<i32>,
4751 aq: &CudaSlice<i8>,
4752 ad: &CudaSlice<f32>,
4753 in_f: usize,
4754 n_ff: usize,
4755 n_used: usize,
4756 n_expert: usize,
4757 qt_g: i32,
4758 qt_u: i32,
4759 rb_g: usize,
4760 rb_u: usize,
4761 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4762 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4763 let (inf, nff, ne, rbg, rbu) = (
4764 in_f as i32,
4765 n_ff as i32,
4766 n_expert as i32,
4767 rb_g as i64,
4768 rb_u as i64,
4769 );
4770 let f = self.func("moe_gate_up_gelu8_dev_q8");
4771 let cfg = LaunchConfig {
4772 grid_dim: (n_ff as u32, n_used as u32, 1),
4773 block_dim: (32, 1, 1),
4774 shared_mem_bytes: 0,
4775 };
4776 let __s_b = self.gpu.stream();
4777 let mut b = __s_b.launch_builder(&f);
4778 b.arg(table)
4779 .arg(sel)
4780 .arg(aq)
4781 .arg(ad)
4782 .arg(&mut act)
4783 .arg(&inf)
4784 .arg(&nff)
4785 .arg(&ne)
4786 .arg(&qt_g)
4787 .arg(&qt_u)
4788 .arg(&rbg)
4789 .arg(&rbu);
4790 unsafe {
4791 b.launch(cfg)?;
4792 }
4793 Ok(act)
4794 }
4795
4796 #[allow(clippy::too_many_arguments)]
4798 pub fn moe_gate_up_gelu8_dev_q8_rows(
4799 &self,
4800 table: &CudaSlice<u64>,
4801 sel: &CudaSlice<i32>,
4802 aq: &CudaSlice<i8>,
4803 ad: &CudaSlice<f32>,
4804 t: usize,
4805 in_f: usize,
4806 n_ff: usize,
4807 n_used: usize,
4808 n_expert: usize,
4809 qt_g: i32,
4810 qt_u: i32,
4811 rb_g: usize,
4812 rb_u: usize,
4813 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4814 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
4815 let (inf, nff, ne, rbg, rbu, nu) = (
4816 in_f as i32,
4817 n_ff as i32,
4818 n_expert as i32,
4819 rb_g as i64,
4820 rb_u as i64,
4821 n_used as i32,
4822 );
4823 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
4824 let cfg = LaunchConfig {
4825 grid_dim: (n_ff as u32, n_used as u32, t as u32),
4826 block_dim: (32, 1, 1),
4827 shared_mem_bytes: 0,
4828 };
4829 let __s_b = self.gpu.stream();
4830 let mut b = __s_b.launch_builder(&f);
4831 b.arg(table)
4832 .arg(sel)
4833 .arg(aq)
4834 .arg(ad)
4835 .arg(&mut act)
4836 .arg(&inf)
4837 .arg(&nff)
4838 .arg(&ne)
4839 .arg(&qt_g)
4840 .arg(&qt_u)
4841 .arg(&rbg)
4842 .arg(&rbu)
4843 .arg(&nu);
4844 unsafe {
4845 b.launch(cfg)?;
4846 }
4847 Ok(act)
4848 }
4849
4850 #[allow(clippy::too_many_arguments)]
4852 pub fn moe_gate_up_gelu8_dev_q8_csr(
4853 &self,
4854 table: &CudaSlice<u64>,
4855 sel: &CudaSlice<i32>,
4856 aq: &CudaSlice<i8>,
4857 ad: &CudaSlice<f32>,
4858 n_pairs: usize,
4859 in_f: usize,
4860 n_ff: usize,
4861 n_used: usize,
4862 n_expert: usize,
4863 qt_g: i32,
4864 qt_u: i32,
4865 rb_g: usize,
4866 rb_u: usize,
4867 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4868 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
4869 let (inf, nff, ne, rbg, rbu, nu, npi) = (
4870 in_f as i32,
4871 n_ff as i32,
4872 n_expert as i32,
4873 rb_g as i64,
4874 rb_u as i64,
4875 n_used as i32,
4876 n_pairs as i32,
4877 );
4878 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
4879 let cfg = LaunchConfig {
4880 grid_dim: (n_ff as u32, n_pairs as u32, 1),
4881 block_dim: (32, 1, 1),
4882 shared_mem_bytes: 0,
4883 };
4884 let __s_b = self.gpu.stream();
4885 let mut b = __s_b.launch_builder(&f);
4886 b.arg(table)
4887 .arg(sel)
4888 .arg(aq)
4889 .arg(ad)
4890 .arg(&mut act)
4891 .arg(&inf)
4892 .arg(&nff)
4893 .arg(&ne)
4894 .arg(&qt_g)
4895 .arg(&qt_u)
4896 .arg(&rbg)
4897 .arg(&rbu)
4898 .arg(&nu)
4899 .arg(&npi);
4900 unsafe {
4901 b.launch(cfg)?;
4902 }
4903 Ok(act)
4904 }
4905
4906 #[allow(clippy::too_many_arguments)]
4908 pub fn moe_down8_fma_dev_q8_rows_g(
4909 &self,
4910 table: &CudaSlice<u64>,
4911 sel: &CudaSlice<i32>,
4912 w: &CudaSlice<f32>,
4913 aq2: &CudaSlice<i8>,
4914 ad2: &CudaSlice<f32>,
4915 dst: &mut CudaSlice<f32>,
4916 t: usize,
4917 in_f: usize,
4918 out_f: usize,
4919 n_used: usize,
4920 n_expert: usize,
4921 qt: i32,
4922 rb: usize,
4923 ) -> Result<(), Box<dyn std::error::Error>> {
4924 let (inf, outf, nu, ne, rbi) = (
4925 in_f as i32,
4926 out_f as i32,
4927 n_used as i32,
4928 n_expert as i32,
4929 rb as i64,
4930 );
4931 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
4935 let f = self.func(if step_b1_w8 {
4936 "moe_down8_fma_dev_q8_rows_w8"
4937 } else {
4938 "moe_down8_fma_dev_q8_rows_g"
4939 });
4940 let cfg = LaunchConfig {
4941 grid_dim: (out_f as u32, 1, t as u32),
4942 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
4943 shared_mem_bytes: 0,
4944 };
4945 let __s_b = self.gpu.stream();
4946 let mut b = __s_b.launch_builder(&f);
4947 b.arg(table)
4948 .arg(sel)
4949 .arg(w)
4950 .arg(aq2)
4951 .arg(ad2)
4952 .arg(dst)
4953 .arg(&inf)
4954 .arg(&outf)
4955 .arg(&nu)
4956 .arg(&ne)
4957 .arg(&qt)
4958 .arg(&rbi);
4959 unsafe {
4960 b.launch(cfg)?;
4961 }
4962 Ok(())
4963 }
4964
4965 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
4969 let (out_f, in_f) = (2048usize, 2816usize);
4970 let nblk = in_f / 32;
4971 let mut seed = 0x9E3779B97F4A7C15u64;
4972 let mut rng = move || {
4973 seed = seed
4974 .wrapping_mul(6364136223846793005)
4975 .wrapping_add(1442695040888963407);
4976 (seed >> 33) as u8
4977 };
4978 let mut w = vec![0u8; out_f * nblk * 18];
4979 for b in w.iter_mut() {
4980 *b = rng();
4981 }
4982 for r in 0..out_f {
4983 for g in 0..nblk {
4984 let off = (r * nblk + g) * 18;
4985 w[off] = 0x00;
4986 w[off + 1] = 0x2C; }
4988 }
4989 let qplane = out_f * nblk * 16;
4990 let mut wrp = vec![0u8; w.len()];
4991 for r in 0..out_f {
4992 for g in 0..nblk {
4993 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
4994 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
4995 .copy_from_slice(&src[0..2]);
4996 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
4997 }
4998 }
4999 let w_d = self.htod_bytes(&w)?;
5000 let wrp_d = self.htod_bytes(&wrp)?;
5001 let mut aq = vec![0i8; m * in_f];
5002 for v in aq.iter_mut() {
5003 *v = rng() as i8;
5004 }
5005 let aq_d = self.htod_i8(&aq)?;
5006 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5007 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5008 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5009 const RPB: u32 = 4;
5010 let cfg = LaunchConfig {
5011 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5012 block_dim: (32, RPB, 1),
5013 shared_mem_bytes: 0,
5014 };
5015 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5016 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5017 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5018 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5019 {
5020 let __s_b = self.gpu.stream();
5021 let mut b = __s_b.launch_builder(&fb);
5022 b.arg(&w_d)
5023 .arg(&aq_d)
5024 .arg(&ad_d)
5025 .arg(&mut y0)
5026 .arg(&inf)
5027 .arg(&outf)
5028 .arg(&mi)
5029 .arg(&rb);
5030 unsafe {
5031 b.launch(cfg)?;
5032 }
5033 let __s_b = self.gpu.stream();
5034 let mut b = __s_b.launch_builder(&fr);
5035 b.arg(&wrp_d)
5036 .arg(&aq_d)
5037 .arg(&ad_d)
5038 .arg(&mut y1)
5039 .arg(&inf)
5040 .arg(&outf)
5041 .arg(&mi)
5042 .arg(&qp);
5043 unsafe {
5044 b.launch(cfg)?;
5045 }
5046 }
5047 self.gpu.stream().synchronize()?;
5048 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5049 let nd = h0
5050 .iter()
5051 .zip(&h1)
5052 .filter(|(a, b)| a.to_bits() != b.to_bits())
5053 .count();
5054 if nd != 0 {
5055 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5056 }
5057 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5058 self.gpu.stream().synchronize()?;
5059 let t0 = std::time::Instant::now();
5060 for _ in 0..500 {
5061 if rp {
5062 let __s_b = self.gpu.stream();
5063 let mut b = __s_b.launch_builder(&fr);
5064 b.arg(&wrp_d)
5065 .arg(&aq_d)
5066 .arg(&ad_d)
5067 .arg(&mut y1)
5068 .arg(&inf)
5069 .arg(&outf)
5070 .arg(&mi)
5071 .arg(&qp);
5072 unsafe {
5073 b.launch(cfg)?;
5074 }
5075 } else {
5076 let __s_b = self.gpu.stream();
5077 let mut b = __s_b.launch_builder(&fb);
5078 b.arg(&w_d)
5079 .arg(&aq_d)
5080 .arg(&ad_d)
5081 .arg(&mut y0)
5082 .arg(&inf)
5083 .arg(&outf)
5084 .arg(&mi)
5085 .arg(&rb);
5086 unsafe {
5087 b.launch(cfg)?;
5088 }
5089 }
5090 }
5091 self.gpu.stream().synchronize()?;
5092 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5093 };
5094 let _ = time(false)?;
5095 let _ = time(true)?; Ok((time(false)?, time(true)?))
5097 }
5098
5099 pub fn build_q4_rp4(
5104 &self,
5105 t: &mut crate::model::GpuTensor,
5106 ) -> Result<(), Box<dyn std::error::Error>> {
5107 use crate::model::GpuTensor;
5108 let GpuTensor::Quant {
5109 bytes,
5110 qtype,
5111 row_bytes,
5112 ne,
5113 rp4,
5114 ..
5115 } = t
5116 else {
5117 return Ok(());
5118 };
5119 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5120 return Ok(());
5121 }
5122 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5123 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5124 return Ok(());
5125 }
5126 let nblk = in_f / 32;
5127 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5128 let f = self.func("q4_0_split_rp_build");
5129 let n = (out_f * nblk) as i32;
5130 let cfg = LaunchConfig {
5131 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5132 block_dim: (256, 1, 1),
5133 shared_mem_bytes: 0,
5134 };
5135 let (of, nb) = (out_f as i32, nblk as i32);
5136 let _ = n;
5137 let __s_b = self.gpu.stream();
5138 let mut b = __s_b.launch_builder(&f);
5139 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5140 unsafe {
5141 b.launch(cfg)?;
5142 }
5143 *rp4 = Some(dst);
5144 Ok(())
5145 }
5146
5147 pub fn build_q8_rp4(
5152 &self,
5153 t: &mut crate::model::GpuTensor,
5154 ) -> Result<(), Box<dyn std::error::Error>> {
5155 use crate::model::GpuTensor;
5156 let GpuTensor::Quant {
5157 bytes,
5158 qtype,
5159 row_bytes,
5160 ne,
5161 rp4,
5162 ..
5163 } = t
5164 else {
5165 return Ok(());
5166 };
5167 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5168 return Ok(());
5169 }
5170 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5171 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5172 return Ok(());
5173 }
5174 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5175 Ok(())
5176 }
5177
5178 pub fn build_q8_rp4_raw(
5181 &self,
5182 bytes: &CudaSlice<u8>,
5183 in_f: usize,
5184 out_f: usize,
5185 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5186 assert!(in_f % 32 == 0);
5187 let nblk = in_f / 32;
5188 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5189 let f = self.func("q8_0_split_rp_build");
5190 let cfg = LaunchConfig {
5191 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5192 block_dim: (256, 1, 1),
5193 shared_mem_bytes: 0,
5194 };
5195 let (of, nb) = (out_f as i32, nblk as i32);
5196 let __s_b = self.gpu.stream();
5197 let mut b = __s_b.launch_builder(&f);
5198 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5199 unsafe {
5200 b.launch(cfg)?;
5201 }
5202 Ok(dst)
5203 }
5204
5205 pub fn build_q4k_rp4(
5213 &self,
5214 t: &mut crate::model::GpuTensor,
5215 ) -> Result<(), Box<dyn std::error::Error>> {
5216 use crate::model::GpuTensor;
5217 let GpuTensor::Quant {
5218 bytes,
5219 qtype,
5220 row_bytes,
5221 ne,
5222 rp4,
5223 ..
5224 } = t
5225 else {
5226 return Ok(());
5227 };
5228 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5229 return Ok(());
5230 }
5231 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5232 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5233 return Ok(());
5234 }
5235 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5236 Ok(())
5237 }
5238
5239 pub fn build_q6k_rp4(
5240 &self,
5241 t: &mut crate::model::GpuTensor,
5242 ) -> Result<(), Box<dyn std::error::Error>> {
5243 use crate::model::GpuTensor;
5244 let GpuTensor::Quant {
5245 bytes,
5246 qtype,
5247 row_bytes,
5248 ne,
5249 rp4,
5250 ..
5251 } = t
5252 else {
5253 return Ok(());
5254 };
5255 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5256 return Ok(());
5257 }
5258 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5259 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5260 return Ok(());
5261 }
5262 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5263 Ok(())
5264 }
5265
5266 pub fn build_kq_rp4_raw(
5268 &self,
5269 bytes: &CudaSlice<u8>,
5270 in_f: usize,
5271 out_f: usize,
5272 qtype: i32,
5273 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5274 assert!(in_f % 256 == 0);
5275 let nsbk = in_f / 256;
5276 let (sb_bytes, kname) = match qtype {
5277 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5278 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5279 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5280 };
5281 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5282 let f = self.func(kname);
5283 let cfg = LaunchConfig {
5284 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5285 block_dim: (256, 1, 1),
5286 shared_mem_bytes: 0,
5287 };
5288 let (of, nb) = (out_f as i32, nsbk as i32);
5289 let __s_b = self.gpu.stream();
5290 let mut b = __s_b.launch_builder(&f);
5291 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5292 unsafe {
5293 b.launch(cfg)?;
5294 }
5295 Ok(dst)
5296 }
5297
5298 pub fn kqrp_enabled() -> bool {
5302 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5303 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5304 Ok("0") => false,
5305 Ok(_) => true,
5306 Err(_) => cfg!(memra_hopper_mma),
5307 })
5308 }
5309
5310 pub fn build_q4_rp_swap(
5316 &self,
5317 t: &mut crate::model::GpuTensor,
5318 ) -> Result<bool, Box<dyn std::error::Error>> {
5319 use crate::model::GpuTensor;
5320 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5330 return Ok(false);
5331 }
5332 self.build_q4_rp4(t)?;
5333 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5335 return Ok(false);
5336 };
5337 match rp4.take() {
5338 Some(split) => {
5339 *bytes = split; *rp = true;
5341 Ok(true)
5342 }
5343 None => Ok(false),
5344 }
5345 }
5346
5347 pub fn q4rp_enabled() -> bool {
5349 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5350 *ON.get_or_init(|| {
5351 std::env::var("MEMRA_Q4RP")
5352 .map(|v| v != "0")
5353 .unwrap_or(true)
5354 })
5355 }
5356
5357 pub fn copy_rows_strided(
5360 &self,
5361 src: &CudaSlice<f32>,
5362 dst: &mut CudaSlice<f32>,
5363 row_elems: usize,
5364 n_rows: usize,
5365 src_stride: usize,
5366 src_off: usize,
5367 ) -> Result<(), Box<dyn std::error::Error>> {
5368 let f = self.func("copy_rows_strided_f32");
5369 let cfg = LaunchConfig {
5370 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5371 block_dim: (256, 1, 1),
5372 shared_mem_bytes: 0,
5373 };
5374 let (re, nr) = (row_elems as i32, n_rows as i32);
5375 let (st, off) = (src_stride as i64, src_off as i64);
5376 let __s_b = self.gpu.stream();
5377 let mut b = __s_b.launch_builder(&f);
5378 b.arg(src)
5379 .arg(&mut *dst)
5380 .arg(&re)
5381 .arg(&nr)
5382 .arg(&st)
5383 .arg(&off);
5384 unsafe {
5385 b.launch(cfg)?;
5386 }
5387 Ok(())
5388 }
5389
5390 pub fn u32_set_k(
5392 &self,
5393 dst: &mut CudaSlice<u32>,
5394 v: u32,
5395 idx: usize,
5396 ) -> Result<(), Box<dyn std::error::Error>> {
5397 let f = self.func("u32_set_k");
5398 let cfg = LaunchConfig {
5399 grid_dim: (1, 1, 1),
5400 block_dim: (1, 1, 1),
5401 shared_mem_bytes: 0,
5402 };
5403 let ii = idx as i32;
5404 let __s_b = self.gpu.stream();
5405 let mut b = __s_b.launch_builder(&f);
5406 b.arg(dst).arg(&v).arg(&ii);
5407 unsafe {
5408 b.launch(cfg)?;
5409 }
5410 Ok(())
5411 }
5412
5413 pub fn i32_add_k(
5415 &self,
5416 d: &mut CudaSlice<i32>,
5417 v: i32,
5418 ) -> Result<(), Box<dyn std::error::Error>> {
5419 let f = self.func("i32_add_k");
5420 let cfg = LaunchConfig {
5421 grid_dim: (1, 1, 1),
5422 block_dim: (32, 1, 1),
5423 shared_mem_bytes: 0,
5424 };
5425 let __s_b = self.gpu.stream();
5426 let mut b = __s_b.launch_builder(&f);
5427 b.arg(d).arg(&v);
5428 unsafe {
5429 b.launch(cfg)?;
5430 }
5431 Ok(())
5432 }
5433
5434 pub fn i32_iota_from(
5436 &self,
5437 ctr: &CudaSlice<i32>,
5438 dst: &mut CudaSlice<i32>,
5439 n: usize,
5440 ) -> Result<(), Box<dyn std::error::Error>> {
5441 let f = self.func("i32_iota_from");
5442 let cfg = LaunchConfig::for_num_elems(n as u32);
5443 let ni = n as i32;
5444 let __s_b = self.gpu.stream();
5445 let mut b = __s_b.launch_builder(&f);
5446 b.arg(ctr).arg(dst).arg(&ni);
5447 unsafe {
5448 b.launch(cfg)?;
5449 }
5450 Ok(())
5451 }
5452
5453 pub fn u32_map_k(
5455 &self,
5456 buf: &mut CudaSlice<u32>,
5457 map: &CudaSlice<u32>,
5458 idx: usize,
5459 ) -> Result<(), Box<dyn std::error::Error>> {
5460 let f = self.func("u32_map_k");
5461 let cfg = LaunchConfig {
5462 grid_dim: (1, 1, 1),
5463 block_dim: (1, 1, 1),
5464 shared_mem_bytes: 0,
5465 };
5466 let ii = idx as i32;
5467 let __s_b = self.gpu.stream();
5468 let mut b = __s_b.launch_builder(&f);
5469 b.arg(buf).arg(map).arg(&ii);
5470 unsafe {
5471 b.launch(cfg)?;
5472 }
5473 Ok(())
5474 }
5475
5476 #[allow(clippy::too_many_arguments)]
5478 pub fn u32_pack2(
5479 &self,
5480 a: &CudaSlice<u32>,
5481 off_a: usize,
5482 n1: usize,
5483 b_in: &CudaSlice<u32>,
5484 n2: usize,
5485 out: &mut CudaSlice<u32>,
5486 ) -> Result<(), Box<dyn std::error::Error>> {
5487 let f = self.func("u32_pack2");
5488 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5489 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5490 let __s_b = self.gpu.stream();
5491 let mut b = __s_b.launch_builder(&f);
5492 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5493 unsafe {
5494 b.launch(cfg)?;
5495 }
5496 Ok(())
5497 }
5498
5499 pub fn moe_w_exscale(
5501 &self,
5502 w: &mut CudaSlice<f32>,
5503 sel: &CudaSlice<i32>,
5504 s: &CudaSlice<f32>,
5505 n: usize,
5506 ) -> Result<(), Box<dyn std::error::Error>> {
5507 let f = self.func("moe_w_exscale");
5508 let cfg = LaunchConfig::for_num_elems(n as u32);
5509 let ni = n as i32;
5510 let __s_b = self.gpu.stream();
5511 let mut b = __s_b.launch_builder(&f);
5512 b.arg(w).arg(sel).arg(s).arg(&ni);
5513 unsafe {
5514 b.launch(cfg)?;
5515 }
5516 Ok(())
5517 }
5518
5519 pub fn moe_w_scale_by_expert(
5522 &self,
5523 w: &mut CudaSlice<f32>,
5524 sel: &CudaSlice<i32>,
5525 macros: &CudaSlice<f32>,
5526 n_expert: usize,
5527 n: usize,
5528 ) -> Result<(), Box<dyn std::error::Error>> {
5529 let f = self.func("moe_w_scale_by_expert");
5530 let cfg = LaunchConfig {
5531 grid_dim: (n.div_ceil(64) as u32, 1, 1),
5532 block_dim: (64, 1, 1),
5533 shared_mem_bytes: 0,
5534 };
5535 let (ne, nn) = (n_expert as i32, n as i32);
5536 let __s_b = self.gpu.stream();
5537 let mut b = __s_b.launch_builder(&f);
5538 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
5539 unsafe {
5540 b.launch(cfg)?;
5541 }
5542 Ok(())
5543 }
5544
5545 pub fn moe_gate_up_silu8_dev_q8(
5546 &self,
5547 table: &CudaSlice<u64>,
5548 sel: &cudarc::driver::CudaView<i32>,
5549 aq: &CudaSlice<i8>,
5550 ad: &CudaSlice<f32>,
5551 in_f: usize,
5552 n_ff: usize,
5553 n_used: usize,
5554 n_expert: usize,
5555 qt_g: i32,
5556 qt_u: i32,
5557 rb_g: usize,
5558 rb_u: usize,
5559 macros: &CudaSlice<f32>,
5560 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5561 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
5562 let (mode, wpb) = GU.get_or_init(|| {
5563 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
5564 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
5565 .ok()
5566 .and_then(|v| v.parse().ok())
5567 .unwrap_or(4u32)
5568 .clamp(1, 16);
5569 (mode, wpb)
5570 });
5571 let (mode, wpb) = (mode.as_str(), *wpb);
5572 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5573 let (inf, nff, ne, rbg, rbu) = (
5574 in_f as i32,
5575 n_ff as i32,
5576 n_expert as i32,
5577 rb_g as i64,
5578 rb_u as i64,
5579 );
5580 let (f, cfg) = match mode {
5581 "1" | "2" | "4" => {
5582 let rpw: u32 = mode.parse().unwrap();
5583 let f = self.func(match rpw {
5584 1 => "moe_gate_up_silu8_dev_q8_r1",
5585 2 => "moe_gate_up_silu8_dev_q8_r2",
5586 _ => "moe_gate_up_silu8_dev_q8_r4",
5587 });
5588 let rows_per_block = (rpw * wpb) as usize;
5589 let gx = n_ff.div_ceil(rows_per_block) as u32;
5590 (
5591 f,
5592 LaunchConfig {
5593 grid_dim: (gx, n_used as u32, 1),
5594 block_dim: (32, wpb, 1),
5595 shared_mem_bytes: 0,
5596 },
5597 )
5598 }
5599 "j8" if n_used <= 32 => (
5600 self.func("moe_gate_up_silu8_dev_q8_j8"),
5601 LaunchConfig {
5602 grid_dim: (n_ff as u32, 1, 1),
5603 block_dim: (32, n_used as u32, 1),
5604 shared_mem_bytes: 0,
5605 },
5606 ),
5607 "vsm2" => {
5609 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
5610 let sh = (rb_g + rb_u) as u32;
5611 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5612 f.set_attribute(
5613 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5614 sh as i32,
5615 )?;
5616 (
5617 f,
5618 LaunchConfig {
5619 grid_dim: (n_ff as u32, n_used as u32, 1),
5620 block_dim: (32, 1, 1),
5621 shared_mem_bytes: sh,
5622 },
5623 )
5624 }
5625 "vsm" => {
5626 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
5627 let sh = (rb_g + rb_u) as u32;
5628 use cudarc::driver::sys::CUfunction_attribute_enum as A;
5629 f.set_attribute(
5630 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
5631 sh as i32,
5632 )?;
5633 (
5634 f,
5635 LaunchConfig {
5636 grid_dim: (n_ff as u32, n_used as u32, 1),
5637 block_dim: (32, 1, 1),
5638 shared_mem_bytes: sh,
5639 },
5640 )
5641 }
5642 "sg" => (
5643 self.func("moe_gate_up_silu8_dev_q8_sg"),
5644 LaunchConfig {
5645 grid_dim: (n_ff as u32, n_used as u32, 1),
5646 block_dim: (32, 1, 1),
5647 shared_mem_bytes: 0,
5648 },
5649 ),
5650 "j8sg" if n_used <= 32 => (
5651 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
5652 LaunchConfig {
5653 grid_dim: (n_ff as u32, 1, 1),
5654 block_dim: (32, n_used as u32, 1),
5655 shared_mem_bytes: 0,
5656 },
5657 ),
5658 "u64" if in_f == 2048 => (
5659 self.func("moe_gate_up_silu8_dev_q8_u64"),
5660 LaunchConfig {
5661 grid_dim: (n_ff as u32, n_used as u32, 1),
5662 block_dim: (32, 1, 1),
5663 shared_mem_bytes: 0,
5664 },
5665 ),
5666 "gs4" if in_f == 2048 => (
5667 self.func("moe_gate_up_silu8_dev_q8_gs4"),
5668 LaunchConfig {
5669 grid_dim: (n_ff as u32, n_used as u32, 1),
5670 block_dim: (32, 4, 1),
5671 shared_mem_bytes: 0,
5672 },
5673 ),
5674 "v" | "" => (
5676 self.func("moe_gate_up_silu8_dev_q8_v"),
5677 LaunchConfig {
5678 grid_dim: (n_ff as u32, n_used as u32, 1),
5679 block_dim: (32, 1, 1),
5680 shared_mem_bytes: 0,
5681 },
5682 ),
5683 "s2" => (
5684 self.func("moe_gate_up_silu8_dev_q8_s2"),
5685 LaunchConfig {
5686 grid_dim: (n_ff as u32, n_used as u32, 1),
5687 block_dim: (32, 2, 1),
5688 shared_mem_bytes: 0,
5689 },
5690 ),
5691 "s2z" => {
5692 let rz = wpb.min(16); (
5694 self.func("moe_gate_up_silu8_dev_q8_s2z"),
5695 LaunchConfig {
5696 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
5697 block_dim: (32, 2, rz),
5698 shared_mem_bytes: 0,
5699 },
5700 )
5701 }
5702 _ => (
5703 self.func("moe_gate_up_silu8_dev_q8"),
5704 LaunchConfig {
5705 grid_dim: (n_ff as u32, n_used as u32, 1),
5706 block_dim: (32, 1, 1),
5707 shared_mem_bytes: 0,
5708 },
5709 ),
5710 };
5711 let __s_b = self.gpu.stream();
5712 let mut b = __s_b.launch_builder(&f);
5713 b.arg(table)
5714 .arg(sel)
5715 .arg(aq)
5716 .arg(ad)
5717 .arg(&mut act)
5718 .arg(&inf)
5719 .arg(&nff)
5720 .arg(&ne)
5721 .arg(&qt_g)
5722 .arg(&qt_u)
5723 .arg(&rbg)
5724 .arg(&rbu)
5725 .arg(macros);
5726 unsafe {
5727 b.launch(cfg)?;
5728 }
5729 Ok(act)
5730 }
5731
5732 #[allow(clippy::too_many_arguments)]
5733 pub fn moe_down8_fma_dev_q8(
5734 &self,
5735 table: &CudaSlice<u64>,
5736 sel: &cudarc::driver::CudaView<i32>,
5737 w: &cudarc::driver::CudaView<f32>,
5738 aq2: &CudaSlice<i8>,
5739 ad2: &CudaSlice<f32>,
5740 dst: &mut cudarc::driver::CudaViewMut<f32>,
5741 in_f: usize,
5742 out_f: usize,
5743 n_used: usize,
5744 n_expert: usize,
5745 qt: i32,
5746 rb: usize,
5747 ) -> Result<(), Box<dyn std::error::Error>> {
5748 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
5749 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
5750 let (inf, outf, nu, ne, rbi) = (
5751 in_f as i32,
5752 out_f as i32,
5753 n_used as i32,
5754 n_expert as i32,
5755 rb as i64,
5756 );
5757 let (f, cfg) = match mode.as_str() {
5760 m @ ("1" | "2" | "4") if n_used <= 8 => {
5761 let rpw: usize = m.parse().unwrap();
5762 let f = self.func(match rpw {
5763 1 => "moe_down8_fma_dev_q8_w8r1",
5764 2 => "moe_down8_fma_dev_q8_w8r2",
5765 _ => "moe_down8_fma_dev_q8_w8r4",
5766 });
5767 (
5768 f,
5769 LaunchConfig {
5770 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
5771 block_dim: (32, n_used as u32, 1),
5772 shared_mem_bytes: 0,
5773 },
5774 )
5775 }
5776 "h2" if in_f == 512 => (
5777 self.func("moe_down8_fma_dev_q8_h2"),
5778 LaunchConfig {
5779 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5780 block_dim: (32, 1, 1),
5781 shared_mem_bytes: 0,
5782 },
5783 ),
5784 "" if in_f == 704 && n_used <= 8 => (
5787 self.func("moe_down8_fma_dev_q8_w8r2"),
5788 LaunchConfig {
5789 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5790 block_dim: (32, n_used as u32, 1),
5791 shared_mem_bytes: 0,
5792 },
5793 ),
5794 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
5798 self.func("moe_down8_fma_dev_q8_w8h2v"),
5799 LaunchConfig {
5800 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5801 block_dim: (32, n_used as u32, 1),
5802 shared_mem_bytes: 0,
5803 },
5804 ),
5805 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
5806 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
5807 LaunchConfig {
5808 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5809 block_dim: (32, n_used as u32, 1),
5810 shared_mem_bytes: 0,
5811 },
5812 ),
5813 "w8h2r2" if in_f == 512 && n_used <= 8 => (
5814 self.func("moe_down8_fma_dev_q8_w8h2r2"),
5815 LaunchConfig {
5816 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
5817 block_dim: (32, n_used as u32, 1),
5818 shared_mem_bytes: 0,
5819 },
5820 ),
5821 "w8h2" if in_f == 512 && n_used <= 8 => (
5822 self.func("moe_down8_fma_dev_q8_w8h2"),
5823 LaunchConfig {
5824 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
5825 block_dim: (32, n_used as u32, 1),
5826 shared_mem_bytes: 0,
5827 },
5828 ),
5829 _ => (
5830 self.func("moe_down8_fma_dev_q8"),
5831 LaunchConfig {
5832 grid_dim: (out_f as u32, 1, 1),
5833 block_dim: (32, 1, 1),
5834 shared_mem_bytes: 0,
5835 },
5836 ),
5837 };
5838 let __s_b = self.gpu.stream();
5839 let mut b = __s_b.launch_builder(&f);
5840 b.arg(table)
5841 .arg(sel)
5842 .arg(w)
5843 .arg(aq2)
5844 .arg(ad2)
5845 .arg(dst)
5846 .arg(&inf)
5847 .arg(&outf)
5848 .arg(&nu)
5849 .arg(&ne)
5850 .arg(&qt)
5851 .arg(&rbi);
5852 unsafe {
5853 b.launch(cfg)?;
5854 }
5855 Ok(())
5856 }
5857
5858 #[allow(clippy::too_many_arguments)]
5865 pub fn moe_gate_up_silu8_dev_q8_rows(
5866 &self,
5867 table: &CudaSlice<u64>,
5868 sel: &CudaSlice<i32>,
5869 aq: &CudaSlice<i8>,
5870 ad: &CudaSlice<f32>,
5871 t: usize,
5872 in_f: usize,
5873 n_ff: usize,
5874 n_used: usize,
5875 n_expert: usize,
5876 qt_g: i32,
5877 qt_u: i32,
5878 rb_g: usize,
5879 rb_u: usize,
5880 macros: &CudaSlice<f32>,
5881 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5882 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
5883 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5884 let cfg = LaunchConfig {
5885 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5886 block_dim: (32, 1, 1),
5887 shared_mem_bytes: 0,
5888 };
5889 let (inf, nff, ne, nu, rbg, rbu) = (
5890 in_f as i32,
5891 n_ff as i32,
5892 n_expert as i32,
5893 n_used as i32,
5894 rb_g as i64,
5895 rb_u as i64,
5896 );
5897 let __s_b = self.gpu.stream();
5898 let mut b = __s_b.launch_builder(&f);
5899 b.arg(table)
5900 .arg(sel)
5901 .arg(aq)
5902 .arg(ad)
5903 .arg(&mut act)
5904 .arg(&inf)
5905 .arg(&nff)
5906 .arg(&ne)
5907 .arg(&qt_g)
5908 .arg(&qt_u)
5909 .arg(&rbg)
5910 .arg(&rbu)
5911 .arg(&nu)
5912 .arg(macros);
5913 unsafe {
5914 b.launch(cfg)?;
5915 }
5916 Ok(act)
5917 }
5918
5919 #[allow(clippy::too_many_arguments)]
5924 pub fn moe_down8_fma_dev_q8_rows(
5925 &self,
5926 table: &CudaSlice<u64>,
5927 sel: &CudaSlice<i32>,
5928 w: &CudaSlice<f32>,
5929 aq2: &CudaSlice<i8>,
5930 ad2: &CudaSlice<f32>,
5931 dst: &mut CudaSlice<f32>,
5932 t: usize,
5933 in_f: usize,
5934 out_f: usize,
5935 n_used: usize,
5936 n_expert: usize,
5937 qt: i32,
5938 rb: usize,
5939 ) -> Result<(), Box<dyn std::error::Error>> {
5940 assert!(
5941 in_f == 512 && n_used <= 8,
5942 "down rows twin is w8h2v shape-gated"
5943 );
5944 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
5945 let cfg = LaunchConfig {
5946 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
5947 block_dim: (32, n_used as u32, 1),
5948 shared_mem_bytes: 0,
5949 };
5950 let (inf, outf, nu, ne, rbi) = (
5951 in_f as i32,
5952 out_f as i32,
5953 n_used as i32,
5954 n_expert as i32,
5955 rb as i64,
5956 );
5957 let __s_b = self.gpu.stream();
5958 let mut b = __s_b.launch_builder(&f);
5959 b.arg(table)
5960 .arg(sel)
5961 .arg(w)
5962 .arg(aq2)
5963 .arg(ad2)
5964 .arg(dst)
5965 .arg(&inf)
5966 .arg(&outf)
5967 .arg(&nu)
5968 .arg(&ne)
5969 .arg(&qt)
5970 .arg(&rbi);
5971 unsafe {
5972 b.launch(cfg)?;
5973 }
5974 Ok(())
5975 }
5976
5977 #[allow(clippy::too_many_arguments)]
5981 pub fn moe_gate_up_silu8_dev_q8_csr(
5982 &self,
5983 table: &CudaSlice<u64>,
5984 sel: &CudaSlice<i32>,
5985 aq: &CudaSlice<i8>,
5986 ad: &CudaSlice<f32>,
5987 n_pairs: usize,
5988 in_f: usize,
5989 n_ff: usize,
5990 n_used: usize,
5991 n_expert: usize,
5992 qt_g: i32,
5993 qt_u: i32,
5994 rb_g: usize,
5995 rb_u: usize,
5996 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5997 let f = if qt_g == crate::QT_NVFP4 {
6000 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6001 } else {
6002 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6003 };
6004 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6005 let cfg = LaunchConfig {
6006 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6007 block_dim: (32, 1, 1),
6008 shared_mem_bytes: 0,
6009 };
6010 let (inf, nff, ne, nu, npi, rbg, rbu) = (
6011 in_f as i32,
6012 n_ff as i32,
6013 n_expert as i32,
6014 n_used as i32,
6015 n_pairs as i32,
6016 rb_g as i64,
6017 rb_u as i64,
6018 );
6019 let __s_b = self.gpu.stream();
6020 let mut b = __s_b.launch_builder(&f);
6021 b.arg(table)
6022 .arg(sel)
6023 .arg(aq)
6024 .arg(ad)
6025 .arg(&mut act)
6026 .arg(&inf)
6027 .arg(&nff)
6028 .arg(&ne)
6029 .arg(&qt_g)
6030 .arg(&qt_u)
6031 .arg(&rbg)
6032 .arg(&rbu)
6033 .arg(&nu)
6034 .arg(&npi);
6035 unsafe {
6036 b.launch(cfg)?;
6037 }
6038 Ok(act)
6039 }
6040
6041 #[allow(clippy::too_many_arguments)]
6045 pub fn moe_down8_fma_dev_q8_variant(
6046 &self,
6047 variant: &str,
6048 table: &CudaSlice<u64>,
6049 sel: &cudarc::driver::CudaView<i32>,
6050 w: &cudarc::driver::CudaView<f32>,
6051 aq2: &CudaSlice<i8>,
6052 ad2: &CudaSlice<f32>,
6053 dst: &mut cudarc::driver::CudaViewMut<f32>,
6054 in_f: usize,
6055 out_f: usize,
6056 n_used: usize,
6057 n_expert: usize,
6058 qt: i32,
6059 rb: usize,
6060 ) -> Result<(), Box<dyn std::error::Error>> {
6061 let (inf, outf, nu, ne, rbi) = (
6062 in_f as i32,
6063 out_f as i32,
6064 n_used as i32,
6065 n_expert as i32,
6066 rb as i64,
6067 );
6068 let (f, cfg) = match variant {
6069 "w8h2" | "w8h2v" => (
6070 self.func(if variant == "w8h2" {
6071 "moe_down8_fma_dev_q8_w8h2"
6072 } else {
6073 "moe_down8_fma_dev_q8_w8h2v"
6074 }),
6075 LaunchConfig {
6076 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6077 block_dim: (32, n_used as u32, 1),
6078 shared_mem_bytes: 0,
6079 },
6080 ),
6081 "w8h2r2" | "w8h2r2v" => (
6082 self.func(if variant == "w8h2r2" {
6083 "moe_down8_fma_dev_q8_w8h2r2"
6084 } else {
6085 "moe_down8_fma_dev_q8_w8h2r2v"
6086 }),
6087 LaunchConfig {
6088 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6089 block_dim: (32, n_used as u32, 1),
6090 shared_mem_bytes: 0,
6091 },
6092 ),
6093 _ => (
6094 self.func("moe_down8_fma_dev_q8"),
6095 LaunchConfig {
6096 grid_dim: (out_f as u32, 1, 1),
6097 block_dim: (32, 1, 1),
6098 shared_mem_bytes: 0,
6099 },
6100 ),
6101 };
6102 let __s_b = self.gpu.stream();
6103 let mut b = __s_b.launch_builder(&f);
6104 b.arg(table)
6105 .arg(sel)
6106 .arg(w)
6107 .arg(aq2)
6108 .arg(ad2)
6109 .arg(dst)
6110 .arg(&inf)
6111 .arg(&outf)
6112 .arg(&nu)
6113 .arg(&ne)
6114 .arg(&qt)
6115 .arg(&rbi);
6116 unsafe {
6117 b.launch(cfg)?;
6118 }
6119 Ok(())
6120 }
6121
6122 #[allow(clippy::too_many_arguments)]
6124 pub fn moe_gate_up_silu8_dev_q8_variant(
6125 &self,
6126 variant: &str,
6127 table: &CudaSlice<u64>,
6128 sel: &cudarc::driver::CudaView<i32>,
6129 aq: &CudaSlice<i8>,
6130 ad: &CudaSlice<f32>,
6131 in_f: usize,
6132 n_ff: usize,
6133 n_used: usize,
6134 n_expert: usize,
6135 qt_g: i32,
6136 qt_u: i32,
6137 rb_g: usize,
6138 rb_u: usize,
6139 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6140 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6141 let (inf, nff, ne, rbg, rbu) = (
6142 in_f as i32,
6143 n_ff as i32,
6144 n_expert as i32,
6145 rb_g as i64,
6146 rb_u as i64,
6147 );
6148 let f = self.func(if variant == "v" {
6149 "moe_gate_up_silu8_dev_q8_v"
6150 } else {
6151 "moe_gate_up_silu8_dev_q8"
6152 });
6153 let cfg = LaunchConfig {
6154 grid_dim: (n_ff as u32, n_used as u32, 1),
6155 block_dim: (32, 1, 1),
6156 shared_mem_bytes: 0,
6157 };
6158 let __s_b = self.gpu.stream();
6159 let mut b = __s_b.launch_builder(&f);
6160 b.arg(table)
6161 .arg(sel)
6162 .arg(aq)
6163 .arg(ad)
6164 .arg(&mut act)
6165 .arg(&inf)
6166 .arg(&nff)
6167 .arg(&ne)
6168 .arg(&qt_g)
6169 .arg(&qt_u)
6170 .arg(&rbg)
6171 .arg(&rbu);
6172 unsafe {
6173 b.launch(cfg)?;
6174 }
6175 Ok(act)
6176 }
6177
6178 pub fn moe_gate_up_silu8_dev(
6179 &self,
6180 table: &CudaSlice<u64>,
6181 sel: &cudarc::driver::CudaView<i32>,
6182 x: &cudarc::driver::CudaView<f32>,
6183 in_f: usize,
6184 n_ff: usize,
6185 n_used: usize,
6186 n_expert: usize,
6187 qt_g: i32,
6188 qt_u: i32,
6189 rb_g: usize,
6190 rb_u: usize,
6191 macros: &CudaSlice<f32>,
6192 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6193 let f = self.func("moe_gate_up_silu8_dev");
6194 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6196 grid_dim: (n_ff as u32, n_used as u32, 1),
6197 block_dim: (256, 1, 1),
6198 shared_mem_bytes: 0,
6199 };
6200 let (inf, nff, ne, rbg, rbu) = (
6201 in_f as i32,
6202 n_ff as i32,
6203 n_expert as i32,
6204 rb_g as i64,
6205 rb_u as i64,
6206 );
6207 let __s_b = self.gpu.stream();
6208 let mut b = __s_b.launch_builder(&f);
6209 b.arg(table)
6210 .arg(sel)
6211 .arg(x)
6212 .arg(&mut act)
6213 .arg(&inf)
6214 .arg(&nff)
6215 .arg(&ne)
6216 .arg(&qt_g)
6217 .arg(&qt_u)
6218 .arg(&rbg)
6219 .arg(&rbu)
6220 .arg(macros);
6221 unsafe {
6222 b.launch(cfg)?;
6223 }
6224 Ok(act)
6225 }
6226
6227 #[allow(clippy::too_many_arguments)]
6230 pub fn moe_down8_fma_dev(
6231 &self,
6232 table: &CudaSlice<u64>,
6233 sel: &cudarc::driver::CudaView<i32>,
6234 w: &cudarc::driver::CudaView<f32>,
6235 act: &CudaSlice<f32>,
6236 dst: &mut cudarc::driver::CudaViewMut<f32>,
6237 in_f: usize,
6238 out_f: usize,
6239 n_used: usize,
6240 n_expert: usize,
6241 qt: i32,
6242 rb: usize,
6243 ) -> Result<(), Box<dyn std::error::Error>> {
6244 let f = self.func("moe_down8_fma_dev");
6245 let cfg = LaunchConfig {
6246 grid_dim: (out_f as u32, 1, 1),
6247 block_dim: (256, 1, 1),
6248 shared_mem_bytes: 0,
6249 };
6250 let (inf, outf, nu, ne, rbv) = (
6251 in_f as i32,
6252 out_f as i32,
6253 n_used as i32,
6254 n_expert as i32,
6255 rb as i64,
6256 );
6257 let __s_b = self.gpu.stream();
6258 let mut b = __s_b.launch_builder(&f);
6259 b.arg(table)
6260 .arg(sel)
6261 .arg(w)
6262 .arg(act)
6263 .arg(dst)
6264 .arg(&inf)
6265 .arg(&outf)
6266 .arg(&nu)
6267 .arg(&ne)
6268 .arg(&qt)
6269 .arg(&rbv);
6270 unsafe {
6271 b.launch(cfg)?;
6272 }
6273 Ok(())
6274 }
6275
6276 pub fn axpy_into(
6278 &self,
6279 src: &CudaSlice<f32>,
6280 alpha: f32,
6281 dst: &mut cudarc::driver::CudaViewMut<f32>,
6282 n: usize,
6283 ) -> Result<(), Box<dyn std::error::Error>> {
6284 let f = self.func("axpy_f32");
6285 let cfg = LaunchConfig::for_num_elems(n as u32);
6286 let (a, ni) = (alpha, n as i32);
6287 let __s_b = self.gpu.stream();
6288 let mut b = __s_b.launch_builder(&f);
6289 b.arg(src).arg(dst).arg(&a).arg(&ni);
6290 unsafe {
6291 b.launch(cfg)?;
6292 }
6293 Ok(())
6294 }
6295
6296 pub fn add_scaled_rows(
6298 &self,
6299 src: &CudaSlice<f32>,
6300 scale: &CudaSlice<f32>,
6301 dst: &mut CudaSlice<f32>,
6302 ncols: usize,
6303 nrows: usize,
6304 ) -> Result<(), Box<dyn std::error::Error>> {
6305 let f = self.func("add_scaled_rows_f32");
6306 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6307 let (nc, nr) = (ncols as i32, nrows as i32);
6308 let __s_b = self.gpu.stream();
6309 let mut b = __s_b.launch_builder(&f);
6310 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6311 unsafe {
6312 b.launch(cfg)?;
6313 }
6314 Ok(())
6315 }
6316
6317 pub fn gather_rows(
6321 &self,
6322 src: &CudaSlice<f32>,
6323 idx: &CudaSlice<i32>,
6324 dst: &mut CudaSlice<f32>,
6325 ncols: usize,
6326 m_e: usize,
6327 ) -> Result<(), Box<dyn std::error::Error>> {
6328 let f = self.func("gather_rows_f32");
6329 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6330 let (nc, me) = (ncols as i32, m_e as i32);
6331 let __s_b = self.gpu.stream();
6332 let mut b = __s_b.launch_builder(&f);
6333 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6334 unsafe {
6335 b.launch(cfg)?;
6336 }
6337 Ok(())
6338 }
6339
6340 pub fn scatter_slot(
6345 &self,
6346 src: &CudaSlice<f32>,
6347 tok_idx: &CudaSlice<i32>,
6348 slot_idx: &CudaSlice<i32>,
6349 weight: &CudaSlice<f32>,
6350 dst: &mut CudaSlice<f32>,
6351 wbuf: &mut CudaSlice<f32>,
6352 ncols: usize,
6353 n_used: usize,
6354 m_e: usize,
6355 ) -> Result<(), Box<dyn std::error::Error>> {
6356 let f = self.func("scatter_add_slot_f32");
6357 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6358 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6359 let __s_b = self.gpu.stream();
6360 let mut b = __s_b.launch_builder(&f);
6361 b.arg(src)
6362 .arg(tok_idx)
6363 .arg(slot_idx)
6364 .arg(weight)
6365 .arg(dst)
6366 .arg(wbuf)
6367 .arg(&nc)
6368 .arg(&nu)
6369 .arg(&me);
6370 unsafe {
6371 b.launch(cfg)?;
6372 }
6373 Ok(())
6374 }
6375
6376 pub fn reduce_slots(
6380 &self,
6381 slots: &CudaSlice<f32>,
6382 wbuf: &CudaSlice<f32>,
6383 dst: &mut CudaSlice<f32>,
6384 ncols: usize,
6385 n_used: usize,
6386 t: usize,
6387 ) -> Result<(), Box<dyn std::error::Error>> {
6388 let f = self.func("reduce_slots_f32");
6389 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6390 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6391 let __s_b = self.gpu.stream();
6392 let mut b = __s_b.launch_builder(&f);
6393 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6394 unsafe {
6395 b.launch(cfg)?;
6396 }
6397 Ok(())
6398 }
6399
6400 pub fn quantize_q8_1_view(
6407 &self,
6408 x: &cudarc::driver::CudaView<f32>,
6409 m: usize,
6410 in_f: usize,
6411 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6412 let f = self.func("quantize_q8_1");
6413 let nblk = in_f / 32;
6414 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6415 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6416 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6417 let (inf, mi) = (in_f as i32, m as i32);
6418 let __s_b = self.gpu.stream();
6419 let mut b = __s_b.launch_builder(&f);
6420 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6421 unsafe {
6422 b.launch(cfg)?;
6423 }
6424 Ok((q, d))
6425 }
6426
6427 pub fn quantize_q8_1(
6428 &self,
6429 x: &CudaSlice<f32>,
6430 m: usize,
6431 in_f: usize,
6432 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6433 let nblk = in_f / 32;
6434 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);
6438 let (inf, mi) = (in_f as i32, m as i32);
6439 if Self::pdl_on() && Self::pdl_wb_on() {
6440 {
6441 use cudarc::driver::{DevicePtr, DevicePtrMut};
6442 let s = &self.gpu.stream();
6443 let (px, _g0) = x.device_ptr(s);
6444 let (pq, _g1) = q.device_ptr_mut(s);
6445 let (pd, _g2) = d.device_ptr_mut(s);
6446 let mut ps = [
6447 &px as *const _ as *mut std::ffi::c_void,
6448 &pq as *const _ as *mut _,
6449 &pd as *const _ as *mut _,
6450 &inf as *const _ as *mut _,
6451 &mi as *const _ as *mut _,
6452 ];
6453 unsafe {
6454 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6455 }
6456 }
6457 return Ok((q, d));
6458 }
6459 let f = self.func("quantize_q8_1");
6460 let __s_b = self.gpu.stream();
6461 let mut b = __s_b.launch_builder(&f);
6462 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6463 unsafe {
6464 b.launch(cfg)?;
6465 }
6466 Ok((q, d))
6467 }
6468
6469 pub fn quantize_fp4_act(
6473 &self,
6474 x: &CudaSlice<f32>,
6475 m: usize,
6476 in_f: usize,
6477 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
6478 let f = self.func("quantize_fp4_act");
6479 let nb16 = in_f / 16;
6480 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);
6483 let (inf, mi) = (in_f as i32, m as i32);
6484 let __s_b = self.gpu.stream();
6485 let mut b = __s_b.launch_builder(&f);
6486 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
6487 unsafe {
6488 b.launch(cfg)?;
6489 }
6490 Ok((aq4, ad4))
6491 }
6492
6493 pub fn qmatvec_gemm_nvfp4_fp4(
6498 &self,
6499 bytes: &CudaSlice<u8>,
6500 x: &CudaSlice<f32>,
6501 m: usize,
6502 in_f: usize,
6503 out_f: usize,
6504 row_bytes: usize,
6505 scale: f32,
6506 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6507 assert!(
6508 in_f % 64 == 0,
6509 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6510 );
6511 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6512 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
6513 if scale != 1.0 {
6514 self.scale_inplace(&mut y, scale, m * out_f)?;
6515 }
6516 Ok(y)
6517 }
6518
6519 fn fp4_gemm_launch(
6522 &self,
6523 bytes: &CudaSlice<u8>,
6524 aq4: &CudaSlice<u32>,
6525 ad4: &CudaSlice<u8>,
6526 m: usize,
6527 in_f: usize,
6528 out_f: usize,
6529 row_bytes: usize,
6530 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6531 let f = self.func("qmatvec_gemm_nvfp4_fp4");
6532 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
6534 const BN: u32 = 256;
6535 let cfg = LaunchConfig {
6536 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
6537 block_dim: (32, 4, 1),
6538 shared_mem_bytes: 0,
6539 };
6540 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6541 let __s_b = self.gpu.stream();
6542 let mut b = __s_b.launch_builder(&f);
6543 b.arg(bytes)
6544 .arg(aq4)
6545 .arg(ad4)
6546 .arg(&mut y)
6547 .arg(&inf)
6548 .arg(&outf)
6549 .arg(&mi)
6550 .arg(&rb);
6551 unsafe {
6552 b.launch(cfg)?;
6553 }
6554 Ok(y)
6555 }
6556
6557 pub fn qmatvec_gemm_nvfp4_fp4_raw(
6559 &self,
6560 bytes: &CudaSlice<u8>,
6561 x: &CudaSlice<f32>,
6562 m: usize,
6563 in_f: usize,
6564 out_f: usize,
6565 row_bytes: usize,
6566 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6567 assert!(
6568 in_f % 64 == 0,
6569 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
6570 );
6571 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
6572 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
6573 }
6574
6575 pub fn qmatvec_q8_0_fast(
6577 &self,
6578 w: &CudaSlice<u8>,
6579 x: &CudaSlice<f32>,
6580 m: usize,
6581 in_f: usize,
6582 out_f: usize,
6583 row_bytes: usize,
6584 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6585 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6586 let f = self.func("qmatvec_q8_0_dp4a");
6587 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6589 grid_dim: (out_f as u32, m as u32, 1),
6590 block_dim: (128, 1, 1),
6591 shared_mem_bytes: 0,
6592 };
6593 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6594 let __s_b = self.gpu.stream();
6595 let mut b = __s_b.launch_builder(&f);
6596 b.arg(w)
6597 .arg(&aq)
6598 .arg(&ad)
6599 .arg(&mut y)
6600 .arg(&inf)
6601 .arg(&outf)
6602 .arg(&mi)
6603 .arg(&rb);
6604 unsafe {
6605 b.launch(cfg)?;
6606 }
6607 Ok(y)
6608 }
6609
6610 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
6613 &self,
6614 w: &CudaSlice<u8>,
6615 x: &CudaSlice<f32>,
6616 m: usize,
6617 in_f: usize,
6618 out_f: usize,
6619 row_bytes: usize,
6620 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6621 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6622 let f = self.func("qmatvec_q4_K_dp4a");
6623 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6625 grid_dim: (out_f as u32, m as u32, 1),
6626 block_dim: (128, 1, 1),
6627 shared_mem_bytes: 0,
6628 };
6629 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6630 let __s_b = self.gpu.stream();
6631 let mut b = __s_b.launch_builder(&f);
6632 b.arg(w)
6633 .arg(&aq)
6634 .arg(&ad)
6635 .arg(&mut y)
6636 .arg(&inf)
6637 .arg(&outf)
6638 .arg(&mi)
6639 .arg(&rb);
6640 unsafe {
6641 b.launch(cfg)?;
6642 }
6643 Ok(y)
6644 }
6645
6646 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
6649 &self,
6650 w: &CudaSlice<u8>,
6651 x: &CudaSlice<f32>,
6652 m: usize,
6653 in_f: usize,
6654 out_f: usize,
6655 row_bytes: usize,
6656 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6657 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6658 let f = self.func("qmatvec_q6_K_dp4a");
6659 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6661 grid_dim: (out_f as u32, m as u32, 1),
6662 block_dim: (128, 1, 1),
6663 shared_mem_bytes: 0,
6664 };
6665 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6666 let __s_b = self.gpu.stream();
6667 let mut b = __s_b.launch_builder(&f);
6668 b.arg(w)
6669 .arg(&aq)
6670 .arg(&ad)
6671 .arg(&mut y)
6672 .arg(&inf)
6673 .arg(&outf)
6674 .arg(&mi)
6675 .arg(&rb);
6676 unsafe {
6677 b.launch(cfg)?;
6678 }
6679 Ok(y)
6680 }
6681
6682 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
6685 &self,
6686 w: &CudaSlice<u8>,
6687 x: &CudaSlice<f32>,
6688 m: usize,
6689 in_f: usize,
6690 out_f: usize,
6691 row_bytes: usize,
6692 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6693 self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6694 }
6695 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
6698 &self,
6699 w: &CudaSlice<u8>,
6700 x: &CudaSlice<f32>,
6701 m: usize,
6702 in_f: usize,
6703 out_f: usize,
6704 row_bytes: usize,
6705 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6706 self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
6707 }
6708 pub fn qmatvec_nvfp4_fast_rp(
6710 &self,
6711 w: &CudaSlice<u8>,
6712 x: &CudaSlice<f32>,
6713 m: usize,
6714 in_f: usize,
6715 out_f: usize,
6716 row_bytes: usize,
6717 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6718 assert!(
6719 in_f % 64 == 0,
6720 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6721 );
6722 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
6723 }
6724 pub fn qmatvec_nvfp4_fast(
6726 &self,
6727 w: &CudaSlice<u8>,
6728 x: &CudaSlice<f32>,
6729 m: usize,
6730 in_f: usize,
6731 out_f: usize,
6732 row_bytes: usize,
6733 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6734 assert!(
6737 in_f % 64 == 0,
6738 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
6739 );
6740 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
6741 }
6742 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
6745 &self,
6746 w: &CudaSlice<u8>,
6747 x: &CudaSlice<f32>,
6748 m: usize,
6749 in_f: usize,
6750 out_f: usize,
6751 row_bytes: usize,
6752 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6753 self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
6754 }
6755
6756 fn qmatvec_dp4a_named(
6758 &self,
6759 name: &str,
6760 w: &CudaSlice<u8>,
6761 x: &CudaSlice<f32>,
6762 m: usize,
6763 in_f: usize,
6764 out_f: usize,
6765 row_bytes: usize,
6766 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6767 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6768 let f = self.func(name);
6769 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6771 grid_dim: (out_f as u32, m as u32, 1),
6772 block_dim: (128, 1, 1),
6773 shared_mem_bytes: 0,
6774 };
6775 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6776 let __s_b = self.gpu.stream();
6777 let mut b = __s_b.launch_builder(&f);
6778 b.arg(w)
6779 .arg(&aq)
6780 .arg(&ad)
6781 .arg(&mut y)
6782 .arg(&inf)
6783 .arg(&outf)
6784 .arg(&mi)
6785 .arg(&rb);
6786 unsafe {
6787 b.launch(cfg)?;
6788 }
6789 Ok(y)
6790 }
6791
6792 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6793 Ok(self.gpu.stream().clone_htod(v)?)
6794 }
6795 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6796 Ok(self.gpu.stream().clone_htod(v)?)
6797 }
6798 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6800 Ok(self.gpu.stream().clone_htod(v)?)
6801 }
6802 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
6803 Ok(self.gpu.stream().clone_htod(v)?)
6804 }
6805 pub fn dtoh_view(
6807 &self,
6808 d: &cudarc::driver::CudaView<f32>,
6809 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6810 let v = self.gpu.stream().clone_dtoh(d)?;
6811 self.gpu.stream().synchronize()?;
6812 Ok(v)
6813 }
6814 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6815 let v = self.gpu.stream().clone_dtoh(d)?;
6816 self.gpu.stream().synchronize()?;
6817 Ok(v)
6818 }
6819 pub fn dtoh_pair(
6823 &self,
6824 a: &CudaSlice<f32>,
6825 b: &CudaSlice<f32>,
6826 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
6827 let av = self.gpu.stream().clone_dtoh(a)?;
6828 let bv = self.gpu.stream().clone_dtoh(b)?;
6829 self.gpu.stream().synchronize()?;
6830 Ok((av, bv))
6831 }
6832 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
6834 let v = self.gpu.stream().clone_dtoh(d)?;
6835 self.gpu.stream().synchronize()?;
6836 Ok(v)
6837 }
6838 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6840 let v = self.gpu.stream().clone_dtoh(d)?;
6841 self.gpu.stream().synchronize()?;
6842 Ok(v)
6843 }
6844 pub fn dtoh_u8_view(
6845 &self,
6846 d: &cudarc::driver::CudaView<u8>,
6847 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
6848 let v = self.gpu.stream().clone_dtoh(d)?;
6849 self.gpu.stream().synchronize()?;
6850 Ok(v)
6851 }
6852 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6853 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
6854 self.keep_if_capturing(&s);
6855 Ok(s)
6856 }
6857
6858 pub fn prob_of_token_device(
6867 &self,
6868 logits: &CudaSlice<f32>,
6869 tok: &CudaSlice<u32>,
6870 n_vocab: usize,
6871 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6872 let nb = ARGMAX_NB;
6873 let mut part = self.alloc_uninit::<f32>(nb)?;
6874 let mut p = self.alloc_uninit::<f32>(1)?;
6875 let f1 = self.func("prob_of_token_partial_f32");
6876 let cfg1 = LaunchConfig {
6877 grid_dim: (nb as u32, 1, 1),
6878 block_dim: (256, 1, 1),
6879 shared_mem_bytes: 0,
6880 };
6881 let nv = n_vocab as i32;
6882 let __s_b1 = self.gpu.stream();
6883 let mut b1 = __s_b1.launch_builder(&f1);
6884 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6885 unsafe {
6886 b1.launch(cfg1)?;
6887 }
6888 let f2 = self.func("prob_of_token_final_f32");
6889 let cfg2 = LaunchConfig {
6890 grid_dim: (1, 1, 1),
6891 block_dim: (256, 1, 1),
6892 shared_mem_bytes: 0,
6893 };
6894 let nbi = nb as i32;
6895 let __s_b2 = self.gpu.stream();
6896 let mut b2 = __s_b2.launch_builder(&f2);
6897 b2.arg(&part).arg(&mut p).arg(&nbi);
6898 unsafe {
6899 b2.launch(cfg2)?;
6900 }
6901 Ok(p)
6902 }
6903
6904 pub fn prob_of_token_device_col(
6911 &self,
6912 logits: &CudaSlice<f32>,
6913 tok_all: &CudaSlice<u32>,
6914 tok_idx: usize,
6915 p_out: &mut CudaSlice<f32>,
6916 p_idx: usize,
6917 n_vocab: usize,
6918 ) -> Result<(), Box<dyn std::error::Error>> {
6919 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
6920 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
6921 let nb = ARGMAX_NB;
6922 let mut part = self.alloc_uninit::<f32>(nb)?;
6923 let f1 = self.func("prob_of_token_partial_f32");
6924 let cfg1 = LaunchConfig {
6925 grid_dim: (nb as u32, 1, 1),
6926 block_dim: (256, 1, 1),
6927 shared_mem_bytes: 0,
6928 };
6929 let nv = n_vocab as i32;
6930 let __s_b1 = self.gpu.stream();
6931 let mut b1 = __s_b1.launch_builder(&f1);
6932 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
6933 unsafe {
6934 b1.launch(cfg1)?;
6935 }
6936 let f2 = self.func("prob_of_token_final_f32");
6937 let cfg2 = LaunchConfig {
6938 grid_dim: (1, 1, 1),
6939 block_dim: (256, 1, 1),
6940 shared_mem_bytes: 0,
6941 };
6942 let nbi = nb as i32;
6943 let __s_b2 = self.gpu.stream();
6944 let mut b2 = __s_b2.launch_builder(&f2);
6945 b2.arg(&part).arg(&mut p_v).arg(&nbi);
6946 unsafe {
6947 b2.launch(cfg2)?;
6948 }
6949 Ok(())
6950 }
6951
6952 pub fn prob_of_token_device_into(
6953 &self,
6954 logits: &CudaSlice<f32>,
6955 tok: &CudaSlice<u32>,
6956 p_out: &mut CudaSlice<f32>,
6957 n_vocab: usize,
6958 ) -> Result<(), Box<dyn std::error::Error>> {
6959 let nb = ARGMAX_NB;
6960 let mut part = self.alloc_uninit::<f32>(nb)?;
6961 let f1 = self.func("prob_of_token_partial_f32");
6962 let cfg1 = LaunchConfig {
6963 grid_dim: (nb as u32, 1, 1),
6964 block_dim: (256, 1, 1),
6965 shared_mem_bytes: 0,
6966 };
6967 let nv = n_vocab as i32;
6968 let __s_b1 = self.gpu.stream();
6969 let mut b1 = __s_b1.launch_builder(&f1);
6970 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
6971 unsafe {
6972 b1.launch(cfg1)?;
6973 }
6974 let f2 = self.func("prob_of_token_final_f32");
6975 let cfg2 = LaunchConfig {
6976 grid_dim: (1, 1, 1),
6977 block_dim: (256, 1, 1),
6978 shared_mem_bytes: 0,
6979 };
6980 let nbi = nb as i32;
6981 let __s_b2 = self.gpu.stream();
6982 let mut b2 = __s_b2.launch_builder(&f2);
6983 b2.arg(&part).arg(p_out).arg(&nbi);
6984 unsafe {
6985 b2.launch(cfg2)?;
6986 }
6987 Ok(())
6988 }
6989
6990 pub fn argmax_token_device(
6991 &self,
6992 logits: &CudaSlice<f32>,
6993 n_vocab: usize,
6994 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
6995 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
6996 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
6997 Ok(tok)
6998 }
6999 pub fn argmax_token_device_into(
7006 &self,
7007 logits: &CudaSlice<f32>,
7008 tok: &mut CudaSlice<u32>,
7009 n_vocab: usize,
7010 ) -> Result<(), Box<dyn std::error::Error>> {
7011 let nb = ARGMAX_NB;
7012 let f1 = self.func("argmax_partial_f32");
7013 let f2 = self.func("argmax_final_f32");
7014 let mut guard = self.argmax_partials.lock().unwrap();
7015 if guard.is_none() {
7016 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7019 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7020 *guard = Some((pv, pi));
7021 }
7022 let (part_v, part_i) = guard.as_mut().unwrap();
7023 let nv = n_vocab as i32;
7024 let nbi = nb as i32;
7025 let cfg1 = LaunchConfig {
7027 grid_dim: (nb as u32, 1, 1),
7028 block_dim: (256, 1, 1),
7029 shared_mem_bytes: 0,
7030 };
7031 let __s_b1 = self.gpu.stream();
7032 let mut b1 = __s_b1.launch_builder(&f1);
7033 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
7034 unsafe {
7035 b1.launch(cfg1)?;
7036 }
7037 let cfg2 = LaunchConfig {
7039 grid_dim: (1, 1, 1),
7040 block_dim: (256, 1, 1),
7041 shared_mem_bytes: 0,
7042 };
7043 let __s_b2 = self.gpu.stream();
7044 let mut b2 = __s_b2.launch_builder(&f2);
7045 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
7046 unsafe {
7047 b2.launch(cfg2)?;
7048 }
7049 Ok(())
7050 }
7051 pub fn argmax_token_device_col(
7057 &self,
7058 logits: &CudaSlice<f32>,
7059 col: usize,
7060 n_vocab: usize,
7061 toks: &mut CudaSlice<u32>,
7062 out_idx: usize,
7063 ) -> Result<(), Box<dyn std::error::Error>> {
7064 let nb = ARGMAX_NB;
7065 let f1 = self.func("argmax_partial_f32");
7066 let f2 = self.func("argmax_final_f32");
7067 let mut guard = self.argmax_partials.lock().unwrap();
7068 if guard.is_none() {
7069 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
7070 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
7071 *guard = Some((pv, pi));
7072 }
7073 let (part_v, part_i) = guard.as_mut().unwrap();
7074 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
7075 let nv = n_vocab as i32;
7076 let nbi = nb as i32;
7077 let cfg1 = LaunchConfig {
7078 grid_dim: (nb as u32, 1, 1),
7079 block_dim: (256, 1, 1),
7080 shared_mem_bytes: 0,
7081 };
7082 let __s_b1 = self.gpu.stream();
7083 let mut b1 = __s_b1.launch_builder(&f1);
7084 b1.arg(&col_view)
7085 .arg(&mut *part_v)
7086 .arg(&mut *part_i)
7087 .arg(&nv);
7088 unsafe {
7089 b1.launch(cfg1)?;
7090 }
7091 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
7092 let cfg2 = LaunchConfig {
7093 grid_dim: (1, 1, 1),
7094 block_dim: (256, 1, 1),
7095 shared_mem_bytes: 0,
7096 };
7097 let __s_b2 = self.gpu.stream();
7098 let mut b2 = __s_b2.launch_builder(&f2);
7099 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
7100 unsafe {
7101 b2.launch(cfg2)?;
7102 }
7103 Ok(())
7104 }
7105 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7107 Ok(self.gpu.stream().clone_htod(v)?)
7108 }
7109 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
7110 let v = self.gpu.stream().clone_dtoh(d)?;
7111 self.gpu.stream().synchronize()?;
7112 Ok(v)
7113 }
7114 pub fn htod_u32_into(
7118 &self,
7119 dst: &mut CudaSlice<u32>,
7120 src: &[u32],
7121 ) -> Result<(), Box<dyn std::error::Error>> {
7122 let mut view = dst.slice_mut(0..src.len());
7123 self.gpu.stream().memcpy_htod(src, &mut view)?;
7124 Ok(())
7125 }
7126
7127 pub fn htod_i32_into(
7130 &self,
7131 dst: &mut CudaSlice<i32>,
7132 src: &[i32],
7133 ) -> Result<(), Box<dyn std::error::Error>> {
7134 let mut view = dst.slice_mut(0..src.len());
7135 self.gpu.stream().memcpy_htod(src, &mut view)?;
7136 Ok(())
7137 }
7138
7139 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7140 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
7141 self.keep_if_capturing(&s);
7142 Ok(s)
7143 }
7144 pub fn embed_gather_device_into(
7147 &self,
7148 embd: &CudaSlice<u8>,
7149 token_d: &CudaSlice<u32>,
7150 x_out: &mut CudaSlice<f32>,
7151 n_embd: usize,
7152 qtype: i32,
7153 row_bytes: usize,
7154 ) -> Result<(), Box<dyn std::error::Error>> {
7155 let f = self.func("embed_gather_u32");
7156 let cfg = LaunchConfig {
7157 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7158 block_dim: (256, 1, 1),
7159 shared_mem_bytes: 0,
7160 };
7161 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7162 let __s_b = self.gpu.stream();
7163 let mut b = __s_b.launch_builder(&f);
7164 b.arg(embd)
7165 .arg(token_d)
7166 .arg(x_out)
7167 .arg(&ne)
7168 .arg(&qt)
7169 .arg(&rb);
7170 unsafe {
7171 b.launch(cfg)?;
7172 }
7173 Ok(())
7174 }
7175 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
7177 let v = self.gpu.stream().clone_dtoh(d)?;
7178 self.gpu.stream().synchronize()?;
7179 Ok(v[0])
7180 }
7181 pub fn i32_set_k(
7188 &self,
7189 dst: &mut CudaSlice<i32>,
7190 v: i32,
7191 ) -> Result<(), Box<dyn std::error::Error>> {
7192 let f = self.func("i32_set_k");
7193 let cfg = LaunchConfig {
7194 grid_dim: (1, 1, 1),
7195 block_dim: (1, 1, 1),
7196 shared_mem_bytes: 0,
7197 };
7198 let idx = 0i32;
7199 let __s_b = self.gpu.stream();
7200 let mut b = __s_b.launch_builder(&f);
7201 b.arg(dst).arg(&v).arg(&idx);
7202 unsafe {
7203 b.launch(cfg)?;
7204 }
7205 Ok(())
7206 }
7207
7208 pub fn set_i32_one(
7209 &self,
7210 d: &mut CudaSlice<i32>,
7211 v: i32,
7212 ) -> Result<(), Box<dyn std::error::Error>> {
7213 self.gpu.stream().memcpy_htod(&[v], d)?;
7214 Ok(())
7215 }
7216 pub fn set_u32_one(
7219 &self,
7220 d: &mut CudaSlice<u32>,
7221 v: u32,
7222 ) -> Result<(), Box<dyn std::error::Error>> {
7223 self.gpu.stream().memcpy_htod(&[v], d)?;
7224 Ok(())
7225 }
7226 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
7228 let v = self.gpu.stream().clone_dtoh(d)?;
7229 self.gpu.stream().synchronize()?;
7230 Ok(v[0])
7231 }
7232 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
7234 Ok(self.gpu.stream().clone_htod(bytes)?)
7235 }
7236 pub fn embed_gather_device(
7240 &self,
7241 embd: &CudaSlice<u8>,
7242 token_d: &CudaSlice<u32>,
7243 n_embd: usize,
7244 qtype: i32,
7245 row_bytes: usize,
7246 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7247 let f = self.func("embed_gather_u32");
7248 let mut x = self.alloc_uninit::<f32>(n_embd)?;
7249 let cfg = LaunchConfig {
7250 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
7251 block_dim: (256, 1, 1),
7252 shared_mem_bytes: 0,
7253 };
7254 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
7255 let __s_b = self.gpu.stream();
7256 let mut b = __s_b.launch_builder(&f);
7257 b.arg(embd)
7258 .arg(token_d)
7259 .arg(&mut x)
7260 .arg(&ne)
7261 .arg(&qt)
7262 .arg(&rb);
7263 unsafe {
7264 b.launch(cfg)?;
7265 }
7266 Ok(x)
7267 }
7268
7269 pub fn embed_gather_device_t(
7273 &self,
7274 embd: &CudaSlice<u8>,
7275 tokens: &[u32],
7276 n_embd: usize,
7277 qtype: i32,
7278 row_bytes: usize,
7279 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7280 let t = tokens.len();
7281 let tok_d = self.gpu.stream().clone_htod(tokens)?;
7282 let f = self.func("embed_gather_u32_t");
7283 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7284 let cfg = LaunchConfig {
7285 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7286 block_dim: (256, 1, 1),
7287 shared_mem_bytes: 0,
7288 };
7289 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7290 let __s_b = self.gpu.stream();
7291 let mut b = __s_b.launch_builder(&f);
7292 b.arg(embd)
7293 .arg(&tok_d)
7294 .arg(&mut x)
7295 .arg(&ne)
7296 .arg(&qt)
7297 .arg(&rb)
7298 .arg(&ti);
7299 unsafe {
7300 b.launch(cfg)?;
7301 }
7302 Ok(x)
7303 }
7304
7305 pub fn embed_gather_device_tv(
7310 &self,
7311 embd: &CudaSlice<u8>,
7312 tok_v: &cudarc::driver::CudaView<u32>,
7313 t: usize,
7314 n_embd: usize,
7315 qtype: i32,
7316 row_bytes: usize,
7317 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7318 let f = self.func("embed_gather_u32_t");
7319 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7320 let cfg = LaunchConfig {
7321 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7322 block_dim: (256, 1, 1),
7323 shared_mem_bytes: 0,
7324 };
7325 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7326 let __s_b = self.gpu.stream();
7327 let mut b = __s_b.launch_builder(&f);
7328 b.arg(embd)
7329 .arg(tok_v)
7330 .arg(&mut x)
7331 .arg(&ne)
7332 .arg(&qt)
7333 .arg(&rb)
7334 .arg(&ti);
7335 unsafe {
7336 b.launch(cfg)?;
7337 }
7338 Ok(x)
7339 }
7340
7341 pub fn embed_gather_device_td(
7342 &self,
7343 embd: &CudaSlice<u8>,
7344 tok_d: &CudaSlice<u32>,
7345 t: usize,
7346 n_embd: usize,
7347 qtype: i32,
7348 row_bytes: usize,
7349 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7350 let f = self.func("embed_gather_u32_t");
7351 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
7352 let cfg = LaunchConfig {
7353 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
7354 block_dim: (256, 1, 1),
7355 shared_mem_bytes: 0,
7356 };
7357 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
7358 let __s_b = self.gpu.stream();
7359 let mut b = __s_b.launch_builder(&f);
7360 b.arg(embd)
7361 .arg(tok_d)
7362 .arg(&mut x)
7363 .arg(&ne)
7364 .arg(&qt)
7365 .arg(&rb)
7366 .arg(&ti);
7367 unsafe {
7368 b.launch(cfg)?;
7369 }
7370 Ok(x)
7371 }
7372
7373 #[inline]
7379 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
7381 if self
7382 .capture_keep_on
7383 .load(std::sync::atomic::Ordering::Relaxed)
7384 {
7385 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
7386 }
7387 }
7388
7389 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
7390 &self,
7391 n: usize,
7392 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
7393 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
7394 {
7398 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7399 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
7400 use cudarc::driver::DevicePtrMut;
7402 let n_bytes = s.len() * std::mem::size_of::<T>();
7403 let stream = self.gpu.stream();
7404 let (p_, _g) = s.device_ptr_mut(&stream);
7405 unsafe {
7406 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
7407 .result()?;
7408 }
7409 }
7410 }
7411 self.keep_if_capturing(&s);
7412 Ok(s)
7413 }
7414
7415 pub fn uninit_q8_pair(
7420 &self,
7421 n: usize,
7422 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7423 Ok((
7424 self.alloc_uninit::<i8>(n)?,
7425 self.alloc_uninit::<f32>(n / 32)?,
7426 ))
7427 }
7428
7429 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7430 self.alloc_uninit::<f32>(n)
7431 }
7432
7433 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7435 self.alloc_uninit::<i8>(n)
7436 }
7437
7438 #[allow(clippy::too_many_arguments)]
7442 pub fn rms_norm3(
7443 &self,
7444 x: &CudaSlice<f32>,
7445 w0: &CudaSlice<f32>,
7446 w1: &CudaSlice<f32>,
7447 w2: &CudaSlice<f32>,
7448 d0: &mut CudaSlice<f32>,
7449 d1: &mut CudaSlice<f32>,
7450 d2: &mut CudaSlice<f32>,
7451 ncols: usize,
7452 nrows: usize,
7453 eps: f32,
7454 ) -> Result<(), Box<dyn std::error::Error>> {
7455 let f = self.func("rms_norm3_f32");
7456 let cfg = LaunchConfig {
7457 grid_dim: (nrows as u32, 1, 1),
7458 block_dim: (rms_block(), 1, 1),
7459 shared_mem_bytes: 0,
7460 };
7461 let (nc, e) = (ncols as i32, eps);
7462 let __s_b = self.gpu.stream();
7463 let mut b = __s_b.launch_builder(&f);
7464 b.arg(x)
7465 .arg(w0)
7466 .arg(w1)
7467 .arg(w2)
7468 .arg(d0)
7469 .arg(d1)
7470 .arg(d2)
7471 .arg(&nc)
7472 .arg(&e);
7473 unsafe {
7474 b.launch(cfg)?;
7475 }
7476 Ok(())
7477 }
7478
7479 #[allow(clippy::too_many_arguments)]
7481 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
7484 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7485 *WARP_ON.get_or_init(|| {
7486 std::env::var("MEMRA_QKVNORM_W")
7487 .map(|v| v != "0")
7488 .unwrap_or(true)
7489 }) && ncols % 4 == 0
7490 && rows >= 64
7491 }
7492
7493 #[allow(clippy::too_many_arguments)]
7496 pub fn rms_norm_qkv_w4b(
7497 &self,
7498 q: &CudaSlice<f32>,
7499 k: &CudaSlice<f32>,
7500 v: &CudaSlice<f32>,
7501 wq: &CudaSlice<f32>,
7502 wk: &CudaSlice<f32>,
7503 wv: &CudaSlice<f32>,
7504 dq: &mut CudaSlice<f32>,
7505 dk: &mut CudaSlice<f32>,
7506 dv: &mut CudaSlice<f32>,
7507 dvb: &mut CudaSlice<u8>,
7508 ncols: usize,
7509 rq: usize,
7510 rk: usize,
7511 eps: f32,
7512 vf16: bool,
7513 ) -> Result<(), Box<dyn std::error::Error>> {
7514 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
7515 let f = self.func("rms_norm_qkv_w4b_f32");
7516 let rows = (rq + 2 * rk) as u32;
7517 let cfg = LaunchConfig {
7518 grid_dim: (rows.div_ceil(8), 1, 1),
7519 block_dim: (256, 1, 1),
7520 shared_mem_bytes: 0,
7521 };
7522 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7523 let vf = vf16 as i32;
7524 let __s_b = self.gpu.stream();
7525 let mut b = __s_b.launch_builder(&f);
7526 b.arg(q)
7527 .arg(k)
7528 .arg(v)
7529 .arg(wq)
7530 .arg(wk)
7531 .arg(wv)
7532 .arg(dq)
7533 .arg(dk)
7534 .arg(dv)
7535 .arg(&mut *dvb)
7536 .arg(&nc)
7537 .arg(&rqi)
7538 .arg(&rki)
7539 .arg(&rvi)
7540 .arg(&e)
7541 .arg(&vf);
7542 unsafe {
7543 b.launch(cfg)?;
7544 }
7545 Ok(())
7546 }
7547
7548 pub fn rms_norm_qkv(
7549 &self,
7550 q: &CudaSlice<f32>,
7551 k: &CudaSlice<f32>,
7552 v: &CudaSlice<f32>,
7553 wq: &CudaSlice<f32>,
7554 wk: &CudaSlice<f32>,
7555 wv: &CudaSlice<f32>,
7556 dq: &mut CudaSlice<f32>,
7557 dk: &mut CudaSlice<f32>,
7558 dv: &mut CudaSlice<f32>,
7559 ncols: usize,
7560 rq: usize,
7561 rk: usize,
7562 eps: f32,
7563 ) -> Result<(), Box<dyn std::error::Error>> {
7564 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7568 let warp_on = *WARP_ON.get_or_init(|| {
7569 std::env::var("MEMRA_QKVNORM_W")
7570 .map(|v| v != "0")
7571 .unwrap_or(true)
7572 });
7573 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
7576 let f = self.func("rms_norm_qkv_w4_f32");
7577 let rows = (rq + 2 * rk) as u32;
7578 let cfg = LaunchConfig {
7579 grid_dim: (rows.div_ceil(8), 1, 1),
7580 block_dim: (256, 1, 1),
7581 shared_mem_bytes: 0,
7582 };
7583 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
7584 let __s_b = self.gpu.stream();
7585 let mut b = __s_b.launch_builder(&f);
7586 b.arg(q)
7587 .arg(k)
7588 .arg(v)
7589 .arg(wq)
7590 .arg(wk)
7591 .arg(wv)
7592 .arg(dq)
7593 .arg(dk)
7594 .arg(dv)
7595 .arg(&nc)
7596 .arg(&rqi)
7597 .arg(&rki)
7598 .arg(&rvi)
7599 .arg(&e);
7600 unsafe {
7601 b.launch(cfg)?;
7602 }
7603 return Ok(());
7604 }
7605 let f = self.func("rms_norm_qkv_f32");
7606 let grid = (rq + 2 * rk) as u32;
7607 let cfg = LaunchConfig {
7608 grid_dim: (grid, 1, 1),
7609 block_dim: (rms_block(), 1, 1),
7610 shared_mem_bytes: 0,
7611 };
7612 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
7613 let __s_b = self.gpu.stream();
7614 let mut b = __s_b.launch_builder(&f);
7615 b.arg(q)
7616 .arg(k)
7617 .arg(v)
7618 .arg(wq)
7619 .arg(wk)
7620 .arg(wv)
7621 .arg(dq)
7622 .arg(dk)
7623 .arg(dv)
7624 .arg(&nc)
7625 .arg(&rqi)
7626 .arg(&rki)
7627 .arg(&e);
7628 unsafe {
7629 b.launch(cfg)?;
7630 }
7631 Ok(())
7632 }
7633
7634 #[allow(clippy::too_many_arguments)]
7636 pub fn rms_norm2x(
7637 &self,
7638 a: &CudaSlice<f32>,
7639 bb: &CudaSlice<f32>,
7640 wa: &CudaSlice<f32>,
7641 wb: &CudaSlice<f32>,
7642 da: &mut CudaSlice<f32>,
7643 db: &mut CudaSlice<f32>,
7644 ncols: usize,
7645 nrows: usize,
7646 eps: f32,
7647 ) -> Result<(), Box<dyn std::error::Error>> {
7648 let f = self.func("rms_norm2x_f32");
7649 let cfg = LaunchConfig {
7650 grid_dim: (2 * nrows as u32, 1, 1),
7651 block_dim: (rms_block(), 1, 1),
7652 shared_mem_bytes: 0,
7653 };
7654 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
7655 let __s_b = self.gpu.stream();
7656 let mut b = __s_b.launch_builder(&f);
7657 b.arg(a)
7658 .arg(bb)
7659 .arg(wa)
7660 .arg(wb)
7661 .arg(da)
7662 .arg(db)
7663 .arg(&nc)
7664 .arg(&nr)
7665 .arg(&e);
7666 unsafe {
7667 b.launch(cfg)?;
7668 }
7669 Ok(())
7670 }
7671
7672 pub fn softcap(
7674 &self,
7675 y: &mut CudaSlice<f32>,
7676 cap: f32,
7677 n: usize,
7678 ) -> Result<(), Box<dyn std::error::Error>> {
7679 let f = self.func("softcap_f32");
7680 let cfg = LaunchConfig::for_num_elems(n as u32);
7681 let ni = n as i32;
7682 let __s_b = self.gpu.stream();
7683 let mut b = __s_b.launch_builder(&f);
7684 b.arg(y).arg(&cap).arg(&ni);
7685 unsafe {
7686 b.launch(cfg)?;
7687 }
7688 Ok(())
7689 }
7690
7691 pub fn mask_ids_rows(
7694 &self,
7695 y: &mut CudaSlice<f32>,
7696 ids: &CudaSlice<i32>,
7697 n_ids: usize,
7698 n_vocab: usize,
7699 t: usize,
7700 ) -> Result<(), Box<dyn std::error::Error>> {
7701 let f = self.func("mask_ids_rows_f32");
7702 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
7703 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
7704 let __s_b = self.gpu.stream();
7705 let mut b = __s_b.launch_builder(&f);
7706 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
7707 unsafe {
7708 b.launch(cfg)?;
7709 }
7710 Ok(())
7711 }
7712
7713 #[allow(clippy::too_many_arguments)]
7715 pub fn add_scale_rms_norm(
7716 &self,
7717 a: &CudaSlice<f32>,
7718 b_in: &CudaSlice<f32>,
7719 c: f32,
7720 w: &CudaSlice<f32>,
7721 res: &mut CudaSlice<f32>,
7722 dst: &mut CudaSlice<f32>,
7723 ncols: usize,
7724 nrows: usize,
7725 eps: f32,
7726 ) -> Result<(), Box<dyn std::error::Error>> {
7727 let f = self.func("add_scale_rms_norm_f32");
7728 let cfg = LaunchConfig {
7729 grid_dim: (nrows as u32, 1, 1),
7730 block_dim: (rms_block(), 1, 1),
7731 shared_mem_bytes: 0,
7732 };
7733 let (nc, e2) = (ncols as i32, eps);
7734 let __s_b = self.gpu.stream();
7735 let mut b = __s_b.launch_builder(&f);
7736 b.arg(a)
7737 .arg(b_in)
7738 .arg(&c)
7739 .arg(w)
7740 .arg(res)
7741 .arg(dst)
7742 .arg(&nc)
7743 .arg(&e2);
7744 unsafe {
7745 b.launch(cfg)?;
7746 }
7747 Ok(())
7748 }
7749
7750 #[allow(clippy::too_many_arguments)]
7753 pub fn add_scale_rms_norm_q8_1(
7754 &self,
7755 a: &CudaSlice<f32>,
7756 b_in: &CudaSlice<f32>,
7757 c: f32,
7758 w: &CudaSlice<f32>,
7759 res: &mut CudaSlice<f32>,
7760 ncols: usize,
7761 nrows: usize,
7762 eps: f32,
7763 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7764 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7765 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7766 let (nc, e2) = (ncols as i32, eps);
7767 if Self::pdl_on() && Self::pdl_wb_on() {
7768 {
7769 use cudarc::driver::{DevicePtr, DevicePtrMut};
7770 let s = &self.gpu.stream();
7771 let (pa, _g0) = a.device_ptr(s);
7772 let (pb, _g1) = b_in.device_ptr(s);
7773 let (pw, _g2) = w.device_ptr(s);
7774 let (pr, _g3) = res.device_ptr_mut(s);
7775 let (pq, _g4) = out_q.device_ptr_mut(s);
7776 let (pd, _g5) = out_d.device_ptr_mut(s);
7777 let mut ps = [
7778 &pa as *const _ as *mut std::ffi::c_void,
7779 &pb as *const _ as *mut _,
7780 &c as *const _ as *mut _,
7781 &pw as *const _ as *mut _,
7782 &pr as *const _ as *mut _,
7783 &pq as *const _ as *mut _,
7784 &pd as *const _ as *mut _,
7785 &nc as *const _ as *mut _,
7786 &e2 as *const _ as *mut _,
7787 ];
7788 unsafe {
7789 self.launch_pdl(
7790 "add_scale_rms_norm_q8_1",
7791 (nrows as u32, 1, 1),
7792 (rms_block(), 1, 1),
7793 &mut ps,
7794 )?;
7795 }
7796 }
7797 return Ok((out_q, out_d));
7798 }
7799 let f = self.func("add_scale_rms_norm_q8_1");
7800 let cfg = LaunchConfig {
7801 grid_dim: (nrows as u32, 1, 1),
7802 block_dim: (rms_block(), 1, 1),
7803 shared_mem_bytes: 0,
7804 };
7805 let __s_b = self.gpu.stream();
7806 let mut b = __s_b.launch_builder(&f);
7807 b.arg(a)
7808 .arg(b_in)
7809 .arg(&c)
7810 .arg(w)
7811 .arg(res)
7812 .arg(&mut out_q)
7813 .arg(&mut out_d)
7814 .arg(&nc)
7815 .arg(&e2);
7816 unsafe {
7817 b.launch(cfg)?;
7818 }
7819 Ok((out_q, out_d))
7820 }
7821
7822 #[allow(clippy::too_many_arguments)]
7824 pub fn add_scale_rms_norm_q8_1_into(
7825 &self,
7826 a: &CudaSlice<f32>,
7827 b_in: &CudaSlice<f32>,
7828 c: f32,
7829 w: &CudaSlice<f32>,
7830 res: &mut CudaSlice<f32>,
7831 ncols: usize,
7832 nrows: usize,
7833 eps: f32,
7834 out_q: &mut CudaSlice<i8>,
7835 out_d: &mut CudaSlice<f32>,
7836 ) -> Result<(), Box<dyn std::error::Error>> {
7837 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
7838 let (nc, e2) = (ncols as i32, eps);
7839 if Self::pdl_on() && Self::pdl_wb_on() {
7840 use cudarc::driver::{DevicePtr, DevicePtrMut};
7841 let s = &self.gpu.stream();
7842 let (pa, _g0) = a.device_ptr(s);
7843 let (pb, _g1) = b_in.device_ptr(s);
7844 let (pw, _g2) = w.device_ptr(s);
7845 let (pr, _g3) = res.device_ptr_mut(s);
7846 let (pq, _g4) = out_q.device_ptr_mut(s);
7847 let (pd, _g5) = out_d.device_ptr_mut(s);
7848 let mut ps = [
7849 &pa as *const _ as *mut std::ffi::c_void,
7850 &pb as *const _ as *mut _,
7851 &c as *const _ as *mut _,
7852 &pw as *const _ as *mut _,
7853 &pr as *const _ as *mut _,
7854 &pq as *const _ as *mut _,
7855 &pd as *const _ as *mut _,
7856 &nc as *const _ as *mut _,
7857 &e2 as *const _ as *mut _,
7858 ];
7859 unsafe {
7860 self.launch_pdl(
7861 "add_scale_rms_norm_q8_1",
7862 (nrows as u32, 1, 1),
7863 (rms_block(), 1, 1),
7864 &mut ps,
7865 )?;
7866 }
7867 return Ok(());
7868 }
7869 let f = self.func("add_scale_rms_norm_q8_1");
7870 let cfg = LaunchConfig {
7871 grid_dim: (nrows as u32, 1, 1),
7872 block_dim: (rms_block(), 1, 1),
7873 shared_mem_bytes: 0,
7874 };
7875 let __s_b = self.gpu.stream();
7876 let mut b = __s_b.launch_builder(&f);
7877 b.arg(a)
7878 .arg(b_in)
7879 .arg(&c)
7880 .arg(w)
7881 .arg(res)
7882 .arg(&mut *out_q)
7883 .arg(&mut *out_d)
7884 .arg(&nc)
7885 .arg(&e2);
7886 unsafe {
7887 b.launch(cfg)?;
7888 }
7889 Ok(())
7890 }
7891
7892 #[allow(clippy::too_many_arguments)]
7895 pub fn rms_pre_add_scale_rms_norm_q8_1(
7896 &self,
7897 a: &CudaSlice<f32>,
7898 wa: &CudaSlice<f32>,
7899 b_in: &CudaSlice<f32>,
7900 c: f32,
7901 w: &CudaSlice<f32>,
7902 res: &mut CudaSlice<f32>,
7903 ncols: usize,
7904 nrows: usize,
7905 eps: f32,
7906 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7907 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7908 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7909 let (nc, e2) = (ncols as i32, eps);
7910 if Self::pdl_on() {
7911 {
7912 use cudarc::driver::{DevicePtr, DevicePtrMut};
7913 let s = &self.gpu.stream();
7914 let (pa, _g0) = a.device_ptr(s);
7915 let (pwa, _g1) = wa.device_ptr(s);
7916 let (pb, _g2) = b_in.device_ptr(s);
7917 let (pw, _g3) = w.device_ptr(s);
7918 let (pr, _g4) = res.device_ptr_mut(s);
7919 let (pq, _g5) = out_q.device_ptr_mut(s);
7920 let (pd, _g6) = out_d.device_ptr_mut(s);
7921 let mut ps = [
7922 &pa as *const _ as *mut std::ffi::c_void,
7923 &pwa as *const _ as *mut _,
7924 &pb as *const _ as *mut _,
7925 &c as *const _ as *mut _,
7926 &pw as *const _ as *mut _,
7927 &pr as *const _ as *mut _,
7928 &pq as *const _ as *mut _,
7929 &pd as *const _ as *mut _,
7930 &nc as *const _ as *mut _,
7931 &e2 as *const _ as *mut _,
7932 ];
7933 unsafe {
7934 self.launch_pdl(
7935 "rms_pre_add_scale_rms_norm_q8_1",
7936 (nrows as u32, 1, 1),
7937 (rms_block(), 1, 1),
7938 &mut ps,
7939 )?;
7940 }
7941 }
7942 return Ok((out_q, out_d));
7943 }
7944 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
7945 let cfg = LaunchConfig {
7946 grid_dim: (nrows as u32, 1, 1),
7947 block_dim: (rms_block(), 1, 1),
7948 shared_mem_bytes: 0,
7949 };
7950 let __s_b = self.gpu.stream();
7951 let mut b = __s_b.launch_builder(&f);
7952 b.arg(a)
7953 .arg(wa)
7954 .arg(b_in)
7955 .arg(&c)
7956 .arg(w)
7957 .arg(res)
7958 .arg(&mut out_q)
7959 .arg(&mut out_d)
7960 .arg(&nc)
7961 .arg(&e2);
7962 unsafe {
7963 b.launch(cfg)?;
7964 }
7965 Ok((out_q, out_d))
7966 }
7967
7968 pub fn gelu_tanh_mul_q8_1(
7971 &self,
7972 gate: &CudaSlice<f32>,
7973 up: &cudarc::driver::CudaView<f32>,
7974 act: &mut CudaSlice<f32>,
7975 ncols: usize,
7976 nrows: usize,
7977 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7978 debug_assert!(ncols % 128 == 0);
7979 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
7980 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
7981 let nc = ncols as i32;
7982 if Self::pdl_on() {
7983 {
7984 use cudarc::driver::{DevicePtr, DevicePtrMut};
7985 let s = &self.gpu.stream();
7986 let (pg, _g0) = gate.device_ptr(s);
7987 let (pu, _g1) = up.device_ptr(s);
7988 let (pact, _g2) = act.device_ptr_mut(s);
7989 let (pq, _g3) = out_q.device_ptr_mut(s);
7990 let (pd, _g4) = out_d.device_ptr_mut(s);
7991 let mut ps = [
7992 &pg as *const _ as *mut std::ffi::c_void,
7993 &pu as *const _ as *mut _,
7994 &pact as *const _ as *mut _,
7995 &pq as *const _ as *mut _,
7996 &pd as *const _ as *mut _,
7997 &nc as *const _ as *mut _,
7998 ];
7999 unsafe {
8000 self.launch_pdl(
8001 "gelu_tanh_mul_q8_1",
8002 (nrows as u32, 1, 1),
8003 (rms_block(), 1, 1),
8004 &mut ps,
8005 )?;
8006 }
8007 }
8008 return Ok((out_q, out_d));
8009 }
8010 let f = self.func("gelu_tanh_mul_q8_1");
8011 let cfg = LaunchConfig {
8012 grid_dim: (nrows as u32, 1, 1),
8013 block_dim: (rms_block(), 1, 1),
8014 shared_mem_bytes: 0,
8015 };
8016 let __s_b = self.gpu.stream();
8017 let mut b = __s_b.launch_builder(&f);
8018 b.arg(gate)
8019 .arg(up)
8020 .arg(act)
8021 .arg(&mut out_q)
8022 .arg(&mut out_d)
8023 .arg(&nc);
8024 unsafe {
8025 b.launch(cfg)?;
8026 }
8027 Ok((out_q, out_d))
8028 }
8029
8030 #[allow(clippy::too_many_arguments)]
8032 pub fn gelu_tanh_mul_q8_1_into(
8033 &self,
8034 gate: &CudaSlice<f32>,
8035 up: &cudarc::driver::CudaView<f32>,
8036 act: &mut CudaSlice<f32>,
8037 ncols: usize,
8038 nrows: usize,
8039 out_q: &mut CudaSlice<i8>,
8040 out_d: &mut CudaSlice<f32>,
8041 ) -> Result<(), Box<dyn std::error::Error>> {
8042 debug_assert!(ncols % 128 == 0);
8043 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8044 let nc = ncols as i32;
8045 if Self::pdl_on() {
8046 use cudarc::driver::{DevicePtr, DevicePtrMut};
8047 let s = &self.gpu.stream();
8048 let (pg, _g0) = gate.device_ptr(s);
8049 let (pu, _g1) = up.device_ptr(s);
8050 let (pact, _g2) = act.device_ptr_mut(s);
8051 let (pq, _g3) = out_q.device_ptr_mut(s);
8052 let (pd, _g4) = out_d.device_ptr_mut(s);
8053 let mut ps = [
8054 &pg as *const _ as *mut std::ffi::c_void,
8055 &pu as *const _ as *mut _,
8056 &pact as *const _ as *mut _,
8057 &pq as *const _ as *mut _,
8058 &pd as *const _ as *mut _,
8059 &nc as *const _ as *mut _,
8060 ];
8061 unsafe {
8062 self.launch_pdl(
8063 "gelu_tanh_mul_q8_1",
8064 (nrows as u32, 1, 1),
8065 (rms_block(), 1, 1),
8066 &mut ps,
8067 )?;
8068 }
8069 return Ok(());
8070 }
8071 let f = self.func("gelu_tanh_mul_q8_1");
8072 let cfg = LaunchConfig {
8073 grid_dim: (nrows as u32, 1, 1),
8074 block_dim: (rms_block(), 1, 1),
8075 shared_mem_bytes: 0,
8076 };
8077 let __s_b = self.gpu.stream();
8078 let mut b = __s_b.launch_builder(&f);
8079 b.arg(gate)
8080 .arg(up)
8081 .arg(&mut *act)
8082 .arg(&mut *out_q)
8083 .arg(&mut *out_d)
8084 .arg(&nc);
8085 unsafe {
8086 b.launch(cfg)?;
8087 }
8088 Ok(())
8089 }
8090
8091 #[allow(clippy::too_many_arguments)]
8093 pub fn add_rms_norm3_q8z(
8094 &self,
8095 a: &CudaSlice<f32>,
8096 b_in: &CudaSlice<f32>,
8097 w0: &CudaSlice<f32>,
8098 w1: &CudaSlice<f32>,
8099 w2: &CudaSlice<f32>,
8100 res: &mut CudaSlice<f32>,
8101 out1: &mut CudaSlice<f32>,
8102 ncols: usize,
8103 nrows: usize,
8104 eps: f32,
8105 ) -> Result<
8106 (
8107 (CudaSlice<i8>, CudaSlice<f32>),
8108 (CudaSlice<i8>, CudaSlice<f32>),
8109 ),
8110 Box<dyn std::error::Error>,
8111 > {
8112 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
8113 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8114 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
8115 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8116 let f = self.func("add_rms_norm3_q8z_f32");
8117 let cfg = LaunchConfig {
8118 grid_dim: (nrows as u32, 1, 1),
8119 block_dim: (rms_block(), 1, 1),
8120 shared_mem_bytes: 0,
8121 };
8122 let (nc, e2) = (ncols as i32, eps);
8123 let __s_b = self.gpu.stream();
8124 let mut b = __s_b.launch_builder(&f);
8125 b.arg(a)
8126 .arg(b_in)
8127 .arg(w0)
8128 .arg(w1)
8129 .arg(w2)
8130 .arg(res)
8131 .arg(&mut q0)
8132 .arg(&mut d0)
8133 .arg(out1)
8134 .arg(&mut q2)
8135 .arg(&mut d2)
8136 .arg(&nc)
8137 .arg(&e2);
8138 unsafe {
8139 b.launch(cfg)?;
8140 }
8141 Ok(((q0, d0), (q2, d2)))
8142 }
8143
8144 #[allow(clippy::too_many_arguments)]
8146 pub fn add_rms_norm3(
8147 &self,
8148 a: &CudaSlice<f32>,
8149 b_in: &CudaSlice<f32>,
8150 w0: &CudaSlice<f32>,
8151 w1: &CudaSlice<f32>,
8152 w2: &CudaSlice<f32>,
8153 res: &mut CudaSlice<f32>,
8154 d0: &mut CudaSlice<f32>,
8155 d1: &mut CudaSlice<f32>,
8156 d2: &mut CudaSlice<f32>,
8157 ncols: usize,
8158 nrows: usize,
8159 eps: f32,
8160 ) -> Result<(), Box<dyn std::error::Error>> {
8161 let f = self.func("add_rms_norm3_f32");
8162 let cfg = LaunchConfig {
8163 grid_dim: (nrows as u32, 1, 1),
8164 block_dim: (rms_block(), 1, 1),
8165 shared_mem_bytes: 0,
8166 };
8167 let (nc, e2) = (ncols as i32, eps);
8168 let __s_b = self.gpu.stream();
8169 let mut b = __s_b.launch_builder(&f);
8170 b.arg(a)
8171 .arg(b_in)
8172 .arg(w0)
8173 .arg(w1)
8174 .arg(w2)
8175 .arg(res)
8176 .arg(d0)
8177 .arg(d1)
8178 .arg(d2)
8179 .arg(&nc)
8180 .arg(&e2);
8181 unsafe {
8182 b.launch(cfg)?;
8183 }
8184 Ok(())
8185 }
8186
8187 pub fn add_scale(
8189 &self,
8190 a: &CudaSlice<f32>,
8191 b_in: &CudaSlice<f32>,
8192 c: f32,
8193 dst: &mut CudaSlice<f32>,
8194 n: usize,
8195 ) -> Result<(), Box<dyn std::error::Error>> {
8196 let f = self.func("add_scale_f32");
8197 let cfg = LaunchConfig::for_num_elems(n as u32);
8198 let ni = n as i32;
8199 let __s_b = self.gpu.stream();
8200 let mut b = __s_b.launch_builder(&f);
8201 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
8202 unsafe {
8203 b.launch(cfg)?;
8204 }
8205 Ok(())
8206 }
8207
8208 pub fn layer_norm_bias(
8210 &self,
8211 x: &CudaSlice<f32>,
8212 w: &CudaSlice<f32>,
8213 b: &CudaSlice<f32>,
8214 dst: &mut CudaSlice<f32>,
8215 ncols: usize,
8216 nrows: usize,
8217 eps: f32,
8218 ) -> Result<(), Box<dyn std::error::Error>> {
8219 let f = self.func("layer_norm_bias_f32");
8220 let (nc, e) = (ncols as i32, eps);
8221 let cfg = LaunchConfig {
8222 grid_dim: (nrows as u32, 1, 1),
8223 block_dim: (256, 1, 1),
8224 shared_mem_bytes: 0,
8225 };
8226 let __s_b = self.gpu.stream();
8227 let mut lb = __s_b.launch_builder(&f);
8228 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
8229 unsafe {
8230 lb.launch(cfg)?;
8231 }
8232 Ok(())
8233 }
8234
8235 pub fn gelu_tanh(
8237 &self,
8238 x: &CudaSlice<f32>,
8239 dst: &mut CudaSlice<f32>,
8240 n: usize,
8241 ) -> Result<(), Box<dyn std::error::Error>> {
8242 let f = self.func("gelu_tanh_f32");
8243 let ni = n as i64;
8244 let cfg = LaunchConfig {
8245 grid_dim: (n.div_ceil(256) as u32, 1, 1),
8246 block_dim: (256, 1, 1),
8247 shared_mem_bytes: 0,
8248 };
8249 let __s_b = self.gpu.stream();
8250 let mut lb = __s_b.launch_builder(&f);
8251 lb.arg(x).arg(&mut *dst).arg(&ni);
8252 unsafe {
8253 lb.launch(cfg)?;
8254 }
8255 Ok(())
8256 }
8257
8258 pub fn row_softmax(
8260 &self,
8261 x: &mut CudaSlice<f32>,
8262 ncols: usize,
8263 nrows: usize,
8264 ) -> Result<(), Box<dyn std::error::Error>> {
8265 let f = self.func("row_softmax_f32");
8266 let nc = ncols as i32;
8267 let cfg = LaunchConfig {
8268 grid_dim: (nrows as u32, 1, 1),
8269 block_dim: (256, 1, 1),
8270 shared_mem_bytes: 0,
8271 };
8272 let __s_b = self.gpu.stream();
8273 let mut lb = __s_b.launch_builder(&f);
8274 lb.arg(&mut *x).arg(&nc);
8275 unsafe {
8276 lb.launch(cfg)?;
8277 }
8278 Ok(())
8279 }
8280
8281 pub fn rms_norm(
8282 &self,
8283 x: &CudaSlice<f32>,
8284 w: &CudaSlice<f32>,
8285 dst: &mut CudaSlice<f32>,
8286 ncols: usize,
8287 nrows: usize,
8288 eps: f32,
8289 ) -> Result<(), Box<dyn std::error::Error>> {
8290 let (nc, e) = (ncols as i32, eps);
8291 if Self::pdl_on() && Self::pdl_wb_on() {
8292 use cudarc::driver::{DevicePtr, DevicePtrMut};
8293 let s = &self.gpu.stream();
8294 let (px, _g0) = x.device_ptr(s);
8295 let (pw, _g1) = w.device_ptr(s);
8296 let (pd, _g2) = dst.device_ptr_mut(s);
8297 let mut ps = [
8298 &px as *const _ as *mut std::ffi::c_void,
8299 &pw as *const _ as *mut _,
8300 &pd as *const _ as *mut _,
8301 &nc as *const _ as *mut _,
8302 &e as *const _ as *mut _,
8303 ];
8304 unsafe {
8305 self.launch_pdl(
8306 "rms_norm_f32",
8307 (nrows as u32, 1, 1),
8308 (rms_block(), 1, 1),
8309 &mut ps,
8310 )?;
8311 }
8312 return Ok(());
8313 }
8314 let f = self.func("rms_norm_f32");
8315 let cfg = LaunchConfig {
8316 grid_dim: (nrows as u32, 1, 1),
8317 block_dim: (rms_block(), 1, 1),
8318 shared_mem_bytes: 0,
8319 };
8320 let __s_b = self.gpu.stream();
8321 let mut b = __s_b.launch_builder(&f);
8322 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8323 unsafe {
8324 b.launch(cfg)?;
8325 }
8326 Ok(())
8327 }
8328
8329 pub fn rms_norm_decode(
8337 &self,
8338 x: &CudaSlice<f32>,
8339 w: &CudaSlice<f32>,
8340 dst: &mut CudaSlice<f32>,
8341 ncols: usize,
8342 nrows: usize,
8343 eps: f32,
8344 ) -> Result<(), Box<dyn std::error::Error>> {
8345 let f = self.func("rms_norm_f32");
8346 let cfg = LaunchConfig {
8347 grid_dim: (nrows as u32, 1, 1),
8348 block_dim: (1024, 1, 1),
8349 shared_mem_bytes: 0,
8350 };
8351 let (nc, e) = (ncols as i32, eps);
8352 let __s_b = self.gpu.stream();
8353 let mut b = __s_b.launch_builder(&f);
8354 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
8355 unsafe {
8356 b.launch(cfg)?;
8357 }
8358 Ok(())
8359 }
8360
8361 pub fn rms_norm_q8_1(
8365 &self,
8366 x: &CudaSlice<f32>,
8367 w: &CudaSlice<f32>,
8368 ncols: usize,
8369 nrows: usize,
8370 eps: f32,
8371 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8372 let nblk = ncols / 32;
8373 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8374 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8375 let (nc, e) = (ncols as i32, eps);
8376 if Self::pdl_on() {
8377 {
8378 use cudarc::driver::{DevicePtr, DevicePtrMut};
8379 let s = &self.gpu.stream();
8380 let (px, _g0) = x.device_ptr(s);
8381 let (pw, _g1) = w.device_ptr(s);
8382 let (pq, _g2) = q.device_ptr_mut(s);
8383 let (pd, _g3) = d.device_ptr_mut(s);
8384 let mut ps = [
8385 &px as *const _ as *mut std::ffi::c_void,
8386 &pw as *const _ as *mut _,
8387 &pq as *const _ as *mut _,
8388 &pd as *const _ as *mut _,
8389 &nc as *const _ as *mut _,
8390 &e as *const _ as *mut _,
8391 ];
8392 unsafe {
8393 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8394 }
8395 }
8396 return Ok((q, d));
8397 }
8398 let f = self.func("rms_norm_q8_1");
8399 let cfg = LaunchConfig {
8402 grid_dim: (nrows as u32, 1, 1),
8403 block_dim: (1024, 1, 1),
8404 shared_mem_bytes: 0,
8405 };
8406 let __s_b = self.gpu.stream();
8407 let mut b = __s_b.launch_builder(&f);
8408 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
8409 unsafe {
8410 b.launch(cfg)?;
8411 }
8412 Ok((q, d))
8413 }
8414
8415 pub fn rms_norm_q8_1_into(
8418 &self,
8419 x: &CudaSlice<f32>,
8420 w: &CudaSlice<f32>,
8421 ncols: usize,
8422 nrows: usize,
8423 eps: f32,
8424 q: &mut CudaSlice<i8>,
8425 d: &mut CudaSlice<f32>,
8426 ) -> Result<(), Box<dyn std::error::Error>> {
8427 let nblk = ncols / 32;
8428 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
8429 let (nc, e) = (ncols as i32, eps);
8430 if Self::pdl_on() {
8431 use cudarc::driver::{DevicePtr, DevicePtrMut};
8432 let s = &self.gpu.stream();
8433 let (px, _g0) = x.device_ptr(s);
8434 let (pw, _g1) = w.device_ptr(s);
8435 let (pq, _g2) = q.device_ptr_mut(s);
8436 let (pd, _g3) = d.device_ptr_mut(s);
8437 let mut ps = [
8438 &px as *const _ as *mut std::ffi::c_void,
8439 &pw as *const _ as *mut _,
8440 &pq as *const _ as *mut _,
8441 &pd as *const _ as *mut _,
8442 &nc as *const _ as *mut _,
8443 &e as *const _ as *mut _,
8444 ];
8445 unsafe {
8446 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
8447 }
8448 return Ok(());
8449 }
8450 let f = self.func("rms_norm_q8_1");
8451 let cfg = LaunchConfig {
8452 grid_dim: (nrows as u32, 1, 1),
8453 block_dim: (1024, 1, 1),
8454 shared_mem_bytes: 0,
8455 };
8456 let __s_b = self.gpu.stream();
8457 let mut b = __s_b.launch_builder(&f);
8458 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
8459 unsafe {
8460 b.launch(cfg)?;
8461 }
8462 Ok(())
8463 }
8464
8465 pub fn quantize_q8_1_into(
8467 &self,
8468 x: &CudaSlice<f32>,
8469 m: usize,
8470 in_f: usize,
8471 q: &mut CudaSlice<i8>,
8472 d: &mut CudaSlice<f32>,
8473 ) -> Result<(), Box<dyn std::error::Error>> {
8474 let nblk = in_f / 32;
8475 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
8476 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
8477 let (inf, mi) = (in_f as i32, m as i32);
8478 if Self::pdl_on() && Self::pdl_wb_on() {
8479 use cudarc::driver::{DevicePtr, DevicePtrMut};
8480 let s = &self.gpu.stream();
8481 let (px, _g0) = x.device_ptr(s);
8482 let (pq, _g1) = q.device_ptr_mut(s);
8483 let (pd, _g2) = d.device_ptr_mut(s);
8484 let mut ps = [
8485 &px as *const _ as *mut std::ffi::c_void,
8486 &pq as *const _ as *mut _,
8487 &pd as *const _ as *mut _,
8488 &inf as *const _ as *mut _,
8489 &mi as *const _ as *mut _,
8490 ];
8491 unsafe {
8492 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
8493 }
8494 return Ok(());
8495 }
8496 let f = self.func("quantize_q8_1");
8497 let __s_b = self.gpu.stream();
8498 let mut b = __s_b.launch_builder(&f);
8499 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
8500 unsafe {
8501 b.launch(cfg)?;
8502 }
8503 Ok(())
8504 }
8505
8506 pub fn add_rms_norm_q8_1(
8510 &self,
8511 a: &CudaSlice<f32>,
8512 b_in: &CudaSlice<f32>,
8513 w: &CudaSlice<f32>,
8514 res: &mut CudaSlice<f32>,
8515 ncols: usize,
8516 nrows: usize,
8517 eps: f32,
8518 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8519 let nblk = ncols / 32;
8520 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
8521 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
8522 let f = self.func("add_rms_norm_q8_1");
8523 let cfg = LaunchConfig {
8525 grid_dim: (nrows as u32, 1, 1),
8526 block_dim: (1024, 1, 1),
8527 shared_mem_bytes: 0,
8528 };
8529 let (nc, e) = (ncols as i32, eps);
8530 let __s_bld = self.gpu.stream();
8531 let mut bld = __s_bld.launch_builder(&f);
8532 bld.arg(a)
8533 .arg(b_in)
8534 .arg(w)
8535 .arg(res)
8536 .arg(&mut q)
8537 .arg(&mut d)
8538 .arg(&nc)
8539 .arg(&e);
8540 unsafe {
8541 bld.launch(cfg)?;
8542 }
8543 Ok((q, d))
8544 }
8545
8546 pub fn add_rms_norm(
8550 &self,
8551 a: &CudaSlice<f32>,
8552 b: &CudaSlice<f32>,
8553 w: &CudaSlice<f32>,
8554 res: &mut CudaSlice<f32>,
8555 dst: &mut CudaSlice<f32>,
8556 ncols: usize,
8557 nrows: usize,
8558 eps: f32,
8559 ) -> Result<(), Box<dyn std::error::Error>> {
8560 let (nc, e) = (ncols as i32, eps);
8561 if Self::pdl_on() && Self::pdl_wb_on() {
8562 use cudarc::driver::{DevicePtr, DevicePtrMut};
8563 let s = &self.gpu.stream();
8564 let (pa, _g0) = a.device_ptr(s);
8565 let (pb, _g1) = b.device_ptr(s);
8566 let (pw, _g2) = w.device_ptr(s);
8567 let (pr, _g3) = res.device_ptr_mut(s);
8568 let (pd, _g4) = dst.device_ptr_mut(s);
8569 let mut ps = [
8570 &pa as *const _ as *mut std::ffi::c_void,
8571 &pb as *const _ as *mut _,
8572 &pw as *const _ as *mut _,
8573 &pr as *const _ as *mut _,
8574 &pd as *const _ as *mut _,
8575 &nc as *const _ as *mut _,
8576 &e as *const _ as *mut _,
8577 ];
8578 unsafe {
8579 self.launch_pdl(
8580 "add_rms_norm_f32",
8581 (nrows as u32, 1, 1),
8582 (rms_block(), 1, 1),
8583 &mut ps,
8584 )?;
8585 }
8586 return Ok(());
8587 }
8588 let f = self.func("add_rms_norm_f32");
8589 let cfg = LaunchConfig {
8590 grid_dim: (nrows as u32, 1, 1),
8591 block_dim: (rms_block(), 1, 1),
8592 shared_mem_bytes: 0,
8593 };
8594 let __s_b2 = self.gpu.stream();
8595 let mut b2 = __s_b2.launch_builder(&f);
8596 b2.arg(a)
8597 .arg(b)
8598 .arg(w)
8599 .arg(&mut *res)
8600 .arg(&mut *dst)
8601 .arg(&nc)
8602 .arg(&e);
8603 unsafe {
8604 b2.launch(cfg)?;
8605 }
8606 Ok(())
8607 }
8608
8609 #[allow(clippy::too_many_arguments)]
8612 pub fn rms_pre_add_rms_norm(
8613 &self,
8614 a: &CudaSlice<f32>,
8615 wa: &CudaSlice<f32>,
8616 b: &CudaSlice<f32>,
8617 w: &CudaSlice<f32>,
8618 res: &mut CudaSlice<f32>,
8619 dst: &mut CudaSlice<f32>,
8620 ncols: usize,
8621 nrows: usize,
8622 eps: f32,
8623 ) -> Result<(), Box<dyn std::error::Error>> {
8624 let f = self.func("rms_pre_add_rms_norm_f32");
8625 let cfg = LaunchConfig {
8626 grid_dim: (nrows as u32, 1, 1),
8627 block_dim: (rms_block(), 1, 1),
8628 shared_mem_bytes: 0,
8629 };
8630 let (nc, e) = (ncols as i32, eps);
8631 let __s_b2 = self.gpu.stream();
8632 let mut b2 = __s_b2.launch_builder(&f);
8633 b2.arg(a)
8634 .arg(wa)
8635 .arg(b)
8636 .arg(w)
8637 .arg(&mut *res)
8638 .arg(&mut *dst)
8639 .arg(&nc)
8640 .arg(&e);
8641 unsafe {
8642 b2.launch(cfg)?;
8643 }
8644 Ok(())
8645 }
8646
8647 #[allow(clippy::too_many_arguments)]
8649 pub fn rms_pre_add_rms_norm_q8z(
8650 &self,
8651 a: &CudaSlice<f32>,
8652 wa: &CudaSlice<f32>,
8653 b: &CudaSlice<f32>,
8654 w: &CudaSlice<f32>,
8655 res: &mut CudaSlice<f32>,
8656 dst: &mut CudaSlice<f32>,
8657 ncols: usize,
8658 nrows: usize,
8659 eps: f32,
8660 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8661 debug_assert!(ncols % 128 == 0);
8662 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8663 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8664 let (nc, e) = (ncols as i32, eps);
8665 if Self::pdl_on() {
8666 {
8667 use cudarc::driver::{DevicePtr, DevicePtrMut};
8668 let s = &self.gpu.stream();
8669 let (pa, _g0) = a.device_ptr(s);
8670 let (pwa, _g1) = wa.device_ptr(s);
8671 let (pb, _g2) = b.device_ptr(s);
8672 let (pw, _g3) = w.device_ptr(s);
8673 let (pr, _g4) = res.device_ptr_mut(s);
8674 let (pdst, _g5) = dst.device_ptr_mut(s);
8675 let (pq, _g6) = out_q.device_ptr_mut(s);
8676 let (pd, _g7) = out_d.device_ptr_mut(s);
8677 let mut ps = [
8678 &pa as *const _ as *mut std::ffi::c_void,
8679 &pwa as *const _ as *mut _,
8680 &pb as *const _ as *mut _,
8681 &pw as *const _ as *mut _,
8682 &pr as *const _ as *mut _,
8683 &pdst as *const _ as *mut _,
8684 &pq as *const _ as *mut _,
8685 &pd as *const _ as *mut _,
8686 &nc as *const _ as *mut _,
8687 &e as *const _ as *mut _,
8688 ];
8689 unsafe {
8690 self.launch_pdl(
8691 "rms_pre_add_rms_norm_q8z_f32",
8692 (nrows as u32, 1, 1),
8693 (rms_block(), 1, 1),
8694 &mut ps,
8695 )?;
8696 }
8697 }
8698 return Ok((out_q, out_d));
8699 }
8700 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8701 let cfg = LaunchConfig {
8702 grid_dim: (nrows as u32, 1, 1),
8703 block_dim: (rms_block(), 1, 1),
8704 shared_mem_bytes: 0,
8705 };
8706 let __s_b2 = self.gpu.stream();
8707 let mut b2 = __s_b2.launch_builder(&f);
8708 b2.arg(a)
8709 .arg(wa)
8710 .arg(b)
8711 .arg(w)
8712 .arg(&mut *res)
8713 .arg(&mut *dst)
8714 .arg(&mut out_q)
8715 .arg(&mut out_d)
8716 .arg(&nc)
8717 .arg(&e);
8718 unsafe {
8719 b2.launch(cfg)?;
8720 }
8721 Ok((out_q, out_d))
8722 }
8723
8724 #[allow(clippy::too_many_arguments)]
8728 pub fn rms_pre_add_rms_norm_q8z_into(
8729 &self,
8730 a: &CudaSlice<f32>,
8731 wa: &CudaSlice<f32>,
8732 b: &CudaSlice<f32>,
8733 w: &CudaSlice<f32>,
8734 res: &mut CudaSlice<f32>,
8735 dst: &mut CudaSlice<f32>,
8736 ncols: usize,
8737 nrows: usize,
8738 eps: f32,
8739 out_q: &mut CudaSlice<i8>,
8740 out_d: &mut CudaSlice<f32>,
8741 ) -> Result<(), Box<dyn std::error::Error>> {
8742 debug_assert!(ncols % 128 == 0);
8743 let (nc, e) = (ncols as i32, eps);
8744 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
8745 let cfg = LaunchConfig {
8746 grid_dim: (nrows as u32, 1, 1),
8747 block_dim: (rms_block(), 1, 1),
8748 shared_mem_bytes: 0,
8749 };
8750 let __s_b = self.gpu.stream();
8751 let mut b2 = __s_b.launch_builder(&f);
8752 b2.arg(a)
8753 .arg(wa)
8754 .arg(b)
8755 .arg(w)
8756 .arg(&mut *res)
8757 .arg(&mut *dst)
8758 .arg(&mut *out_q)
8759 .arg(&mut *out_d)
8760 .arg(&nc)
8761 .arg(&e);
8762 unsafe {
8763 b2.launch(cfg)?;
8764 }
8765 Ok(())
8766 }
8767
8768 #[allow(clippy::too_many_arguments)]
8771 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
8772 &self,
8773 a: &CudaSlice<f32>,
8774 wa: &CudaSlice<f32>,
8775 b_in: &CudaSlice<f32>,
8776 c: f32,
8777 w: &CudaSlice<f32>,
8778 res: &mut CudaSlice<f32>,
8779 ncols: usize,
8780 nrows: usize,
8781 eps: f32,
8782 out_q: &mut CudaSlice<i8>,
8783 out_d: &mut CudaSlice<f32>,
8784 ) -> Result<(), Box<dyn std::error::Error>> {
8785 debug_assert!(ncols % 128 == 0);
8786 let (nc, e2) = (ncols as i32, eps);
8787 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8788 let cfg = LaunchConfig {
8789 grid_dim: (nrows as u32, 1, 1),
8790 block_dim: (rms_block(), 1, 1),
8791 shared_mem_bytes: 0,
8792 };
8793 let __s_b = self.gpu.stream();
8794 let mut b2 = __s_b.launch_builder(&f);
8795 b2.arg(a)
8796 .arg(wa)
8797 .arg(b_in)
8798 .arg(&c)
8799 .arg(w)
8800 .arg(&mut *res)
8801 .arg(&mut *out_q)
8802 .arg(&mut *out_d)
8803 .arg(&nc)
8804 .arg(&e2);
8805 unsafe {
8806 b2.launch(cfg)?;
8807 }
8808 Ok(())
8809 }
8810
8811 pub fn g4_pnfold_on() -> bool {
8819 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8820 *ON.get_or_init(|| {
8821 std::env::var("MEMRA_G4_PNFOLD")
8822 .map(|v| v != "0")
8823 .unwrap_or(true)
8824 })
8825 }
8826
8827 pub fn build_q4_out_concat3(
8831 &self,
8832 w0: &crate::model::GpuTensor,
8833 w1: &crate::model::GpuTensor,
8834 w2: &crate::model::GpuTensor,
8835 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
8836 use crate::model::GpuTensor;
8837 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
8838 match w {
8839 GpuTensor::Quant {
8840 qtype,
8841 row_bytes,
8842 rp,
8843 ..
8844 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
8845 _ => None,
8846 }
8847 };
8848 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
8849 else {
8850 return Ok(None);
8851 };
8852 if rb0 != rb1
8853 || rb0 != rb2
8854 || w0.in_features() != w1.in_features()
8855 || w0.in_features() != w2.in_features()
8856 {
8857 return Ok(None);
8858 }
8859 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
8860 match w {
8861 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
8862 _ => unreachable!(),
8863 }
8864 }
8865 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
8866 let total = rb0 * (o0 + o1 + o2);
8867 let mut cat = self.alloc_u8(total)?;
8868 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
8869 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
8870 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
8871 Ok(Some(GpuTensor::Quant {
8872 bytes: cat,
8873 qtype: QT_Q4_0,
8874 row_bytes: rb0,
8875 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
8876 scale: 1.0,
8877 rp: false,
8878 #[cfg(memra_cutlass)]
8879 cutlass: None,
8880 fp8: None,
8881 blk: None,
8882 rp4: None,
8883 f16: None,
8884 }))
8885 }
8886
8887 fn full_width_rope_only(
8905 kernel: &str,
8906 n_rot: usize,
8907 head_dim: usize,
8908 ) -> Result<(), Box<dyn std::error::Error>> {
8909 if n_rot == head_dim {
8910 return Ok(());
8911 }
8912 Err(format!(
8913 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
8914 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
8915 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
8916 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
8917 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
8918 )
8919 .into())
8920 }
8921
8922 #[allow(clippy::too_many_arguments)]
8926 pub fn rms_norm_qkv_rope_cat(
8927 &self,
8928 qkv: &CudaSlice<f32>,
8929 wq: &CudaSlice<f32>,
8930 wk: &CudaSlice<f32>,
8931 wv: &CudaSlice<f32>,
8932 q: &mut CudaSlice<f32>,
8933 k: &mut CudaSlice<f32>,
8934 v: &mut CudaSlice<f32>,
8935 head_dim: usize,
8936 n_rot: usize,
8937 rq: usize,
8938 rk: usize,
8939 pos: &CudaSlice<i32>,
8940 nh_q: usize,
8941 nh_k: usize,
8942 base: f32,
8943 freq_scale: f32,
8944 ff: Option<&CudaSlice<f32>>,
8945 eps: f32,
8946 ) -> Result<(), Box<dyn std::error::Error>> {
8947 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
8948 let rows = rq + rk + rk;
8949 let theta_scale = base.powf(-2.0 / head_dim as f32);
8950 let (nc, rqi, rki, nhq, nhk) = (
8951 head_dim as i32,
8952 rq as i32,
8953 rk as i32,
8954 nh_q as i32,
8955 nh_k as i32,
8956 );
8957 if Self::pdl_on() {
8958 use cudarc::driver::{DevicePtr, DevicePtrMut};
8959 let s = &self.gpu.stream();
8960 let (pqkv, _g0) = qkv.device_ptr(s);
8961 let (pwq, _g1) = wq.device_ptr(s);
8962 let (pwk, _g2) = wk.device_ptr(s);
8963 let (pwv, _g3) = wv.device_ptr(s);
8964 let (pq, _g4) = q.device_ptr_mut(s);
8965 let (pk, _g5) = k.device_ptr_mut(s);
8966 let (pv, _g6) = v.device_ptr_mut(s);
8967 let (ppos, _g7) = pos.device_ptr(s);
8968 let (pff, _g8) = match ff {
8969 Some(t) => {
8970 let (p, g) = t.device_ptr(s);
8971 (p, Some(g))
8972 }
8973 None => (0, None),
8974 };
8975 let mut ps = [
8976 &pqkv as *const _ as *mut std::ffi::c_void,
8977 &pwq as *const _ as *mut _,
8978 &pwk as *const _ as *mut _,
8979 &pwv as *const _ as *mut _,
8980 &pq as *const _ as *mut _,
8981 &pk as *const _ as *mut _,
8982 &pv as *const _ as *mut _,
8983 &nc as *const _ as *mut _,
8984 &rqi as *const _ as *mut _,
8985 &rki as *const _ as *mut _,
8986 &ppos as *const _ as *mut _,
8987 &nhq as *const _ as *mut _,
8988 &nhk as *const _ as *mut _,
8989 &theta_scale as *const _ as *mut _,
8990 &freq_scale as *const _ as *mut _,
8991 &pff as *const _ as *mut _,
8992 &eps as *const _ as *mut _,
8993 ];
8994 unsafe {
8995 self.launch_pdl(
8996 "rms_norm_qkv_rope_cat_f32",
8997 (rows as u32, 1, 1),
8998 (rms_block(), 1, 1),
8999 &mut ps,
9000 )?;
9001 }
9002 return Ok(());
9003 }
9004 let f = self.func("rms_norm_qkv_rope_cat_f32");
9005 let cfg = LaunchConfig {
9006 grid_dim: (rows as u32, 1, 1),
9007 block_dim: (rms_block(), 1, 1),
9008 shared_mem_bytes: 0,
9009 };
9010 let __s_b = self.gpu.stream();
9011 let mut b = __s_b.launch_builder(&f);
9012 match ff {
9013 Some(t) => {
9014 b.arg(qkv)
9015 .arg(wq)
9016 .arg(wk)
9017 .arg(wv)
9018 .arg(&mut *q)
9019 .arg(&mut *k)
9020 .arg(&mut *v)
9021 .arg(&nc)
9022 .arg(&rqi)
9023 .arg(&rki)
9024 .arg(pos)
9025 .arg(&nhq)
9026 .arg(&nhk)
9027 .arg(&theta_scale)
9028 .arg(&freq_scale)
9029 .arg(t)
9030 .arg(&eps);
9031 unsafe {
9032 b.launch(cfg)?;
9033 }
9034 }
9035 None => {
9036 let null: u64 = 0;
9037 b.arg(qkv)
9038 .arg(wq)
9039 .arg(wk)
9040 .arg(wv)
9041 .arg(&mut *q)
9042 .arg(&mut *k)
9043 .arg(&mut *v)
9044 .arg(&nc)
9045 .arg(&rqi)
9046 .arg(&rki)
9047 .arg(pos)
9048 .arg(&nhq)
9049 .arg(&nhk)
9050 .arg(&theta_scale)
9051 .arg(&freq_scale)
9052 .arg(&null)
9053 .arg(&eps);
9054 unsafe {
9055 b.launch(cfg)?;
9056 }
9057 }
9058 }
9059 Ok(())
9060 }
9061
9062 #[allow(clippy::too_many_arguments)]
9066 pub fn rms_norm_qkv_rope(
9067 &self,
9068 q0: &CudaSlice<f32>,
9069 k0: &CudaSlice<f32>,
9070 v0: &CudaSlice<f32>,
9071 wq: &CudaSlice<f32>,
9072 wk: &CudaSlice<f32>,
9073 wv: &CudaSlice<f32>,
9074 q: &mut CudaSlice<f32>,
9075 k: &mut CudaSlice<f32>,
9076 v: &mut CudaSlice<f32>,
9077 head_dim: usize,
9078 n_rot: usize,
9079 rq: usize,
9080 rk: usize,
9081 pos: &CudaSlice<i32>,
9082 nh_q: usize,
9083 nh_k: usize,
9084 base: f32,
9085 freq_scale: f32,
9086 ff: Option<&CudaSlice<f32>>,
9087 eps: f32,
9088 ) -> Result<(), Box<dyn std::error::Error>> {
9089 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
9090 let f = self.func("rms_norm_qkv_rope_f32");
9091 let rows = rq + rk + rk; let cfg = LaunchConfig {
9093 grid_dim: (rows as u32, 1, 1),
9094 block_dim: (rms_block(), 1, 1),
9095 shared_mem_bytes: 0,
9096 };
9097 let theta_scale = base.powf(-2.0 / head_dim as f32);
9098 let (nc, rqi, rki, nhq, nhk) = (
9099 head_dim as i32,
9100 rq as i32,
9101 rk as i32,
9102 nh_q as i32,
9103 nh_k as i32,
9104 );
9105 let __s_b = self.gpu.stream();
9106 let mut b = __s_b.launch_builder(&f);
9107 match ff {
9108 Some(t) => {
9109 b.arg(q0)
9110 .arg(k0)
9111 .arg(v0)
9112 .arg(wq)
9113 .arg(wk)
9114 .arg(wv)
9115 .arg(&mut *q)
9116 .arg(&mut *k)
9117 .arg(&mut *v)
9118 .arg(&nc)
9119 .arg(&rqi)
9120 .arg(&rki)
9121 .arg(pos)
9122 .arg(&nhq)
9123 .arg(&nhk)
9124 .arg(&theta_scale)
9125 .arg(&freq_scale)
9126 .arg(t)
9127 .arg(&eps);
9128 unsafe {
9129 b.launch(cfg)?;
9130 }
9131 }
9132 None => {
9133 let null: u64 = 0;
9134 b.arg(q0)
9135 .arg(k0)
9136 .arg(v0)
9137 .arg(wq)
9138 .arg(wk)
9139 .arg(wv)
9140 .arg(&mut *q)
9141 .arg(&mut *k)
9142 .arg(&mut *v)
9143 .arg(&nc)
9144 .arg(&rqi)
9145 .arg(&rki)
9146 .arg(pos)
9147 .arg(&nhq)
9148 .arg(&nhk)
9149 .arg(&theta_scale)
9150 .arg(&freq_scale)
9151 .arg(&null)
9152 .arg(&eps);
9153 unsafe {
9154 b.launch(cfg)?;
9155 }
9156 }
9157 }
9158 Ok(())
9159 }
9160
9161 #[allow(clippy::too_many_arguments)]
9167 pub fn rms_norm_qkv_rope_append_dc(
9168 &self,
9169 q0: &CudaSlice<f32>,
9170 k0: &CudaSlice<f32>,
9171 v0: &CudaSlice<f32>,
9172 wq: &CudaSlice<f32>,
9173 wk: &CudaSlice<f32>,
9174 wv: &CudaSlice<f32>,
9175 q: &mut CudaSlice<f32>,
9176 k: &mut CudaSlice<f32>,
9177 v: &mut CudaSlice<f32>,
9178 head_dim: usize,
9179 n_rot: usize,
9180 rq: usize,
9181 rk: usize,
9182 pos: &CudaSlice<i32>,
9183 nh_q: usize,
9184 nh_k: usize,
9185 base: f32,
9186 freq_scale: f32,
9187 ff: Option<&CudaSlice<f32>>,
9188 eps: f32,
9189 kc: &mut CudaSlice<u8>,
9190 vc: &mut CudaSlice<u8>,
9191 t_dev: &CudaSlice<i32>,
9192 k_tok_bytes: usize,
9193 v_tok_bytes: usize,
9194 g: bool,
9195 ) -> Result<(), Box<dyn std::error::Error>> {
9196 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
9197 let rows = rq + rk + rk;
9198 let theta_scale = base.powf(-2.0 / head_dim as f32);
9199 let (nc, rqi, rki, nhq, nhk) = (
9200 head_dim as i32,
9201 rq as i32,
9202 rk as i32,
9203 nh_q as i32,
9204 nh_k as i32,
9205 );
9206 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9207 if Self::pdl_on() && Self::pdl_wb_on() {
9208 use cudarc::driver::{DevicePtr, DevicePtrMut};
9209 let s = &self.gpu.stream();
9210 let (p0, _a0) = q0.device_ptr(s);
9211 let (p1, _a1) = k0.device_ptr(s);
9212 let (p2, _a2) = v0.device_ptr(s);
9213 let (pwq, _a3) = wq.device_ptr(s);
9214 let (pwk, _a4) = wk.device_ptr(s);
9215 let (pwv, _a5) = wv.device_ptr(s);
9216 let (pq, _a6) = q.device_ptr_mut(s);
9217 let (pk, _a7) = k.device_ptr_mut(s);
9218 let (pv, _a8) = v.device_ptr_mut(s);
9219 let (pp, _a9) = pos.device_ptr(s);
9220 let pff: u64 = match ff {
9221 Some(t) => {
9222 let (p, _gg) = t.device_ptr(s);
9223 p as u64
9224 }
9225 None => 0,
9226 };
9227 let (pkc, _a10) = kc.device_ptr_mut(s);
9228 let (pvc, _a11) = vc.device_ptr_mut(s);
9229 let (pt, _a12) = t_dev.device_ptr(s);
9230 let mut ps = [
9231 &p0 as *const _ as *mut std::ffi::c_void,
9232 &p1 as *const _ as *mut _,
9233 &p2 as *const _ as *mut _,
9234 &pwq as *const _ as *mut _,
9235 &pwk as *const _ as *mut _,
9236 &pwv as *const _ as *mut _,
9237 &pq as *const _ as *mut _,
9238 &pk as *const _ as *mut _,
9239 &pv as *const _ as *mut _,
9240 &nc as *const _ as *mut _,
9241 &rqi as *const _ as *mut _,
9242 &rki as *const _ as *mut _,
9243 &pp as *const _ as *mut _,
9244 &nhq as *const _ as *mut _,
9245 &nhk as *const _ as *mut _,
9246 &theta_scale as *const _ as *mut _,
9247 &freq_scale as *const _ as *mut _,
9248 &pff as *const _ as *mut _,
9249 &eps as *const _ as *mut _,
9250 &pkc as *const _ as *mut _,
9251 &pvc as *const _ as *mut _,
9252 &pt as *const _ as *mut _,
9253 &ktb as *const _ as *mut _,
9254 &vtb as *const _ as *mut _,
9255 ];
9256 unsafe {
9257 self.launch_pdl_flash(
9258 g,
9259 "rms_norm_qkv_rope_append_dc_f32",
9260 (rows as u32, 1, 1),
9261 (rms_block(), 1, 1),
9262 0,
9263 &mut ps,
9264 )?;
9265 }
9266 return Ok(());
9267 }
9268 let f = if g {
9269 self.func_g("rms_norm_qkv_rope_append_dc_f32")
9270 } else {
9271 self.func("rms_norm_qkv_rope_append_dc_f32")
9272 };
9273 let cfg = LaunchConfig {
9274 grid_dim: (rows as u32, 1, 1),
9275 block_dim: (rms_block(), 1, 1),
9276 shared_mem_bytes: 0,
9277 };
9278 let __s_b = self.gpu.stream();
9279 let mut b = __s_b.launch_builder(&f);
9280 match ff {
9281 Some(t) => {
9282 b.arg(q0)
9283 .arg(k0)
9284 .arg(v0)
9285 .arg(wq)
9286 .arg(wk)
9287 .arg(wv)
9288 .arg(&mut *q)
9289 .arg(&mut *k)
9290 .arg(&mut *v)
9291 .arg(&nc)
9292 .arg(&rqi)
9293 .arg(&rki)
9294 .arg(pos)
9295 .arg(&nhq)
9296 .arg(&nhk)
9297 .arg(&theta_scale)
9298 .arg(&freq_scale)
9299 .arg(t)
9300 .arg(&eps)
9301 .arg(&mut *kc)
9302 .arg(&mut *vc)
9303 .arg(t_dev)
9304 .arg(&ktb)
9305 .arg(&vtb);
9306 unsafe {
9307 b.launch(cfg)?;
9308 }
9309 }
9310 None => {
9311 let null: u64 = 0;
9312 b.arg(q0)
9313 .arg(k0)
9314 .arg(v0)
9315 .arg(wq)
9316 .arg(wk)
9317 .arg(wv)
9318 .arg(&mut *q)
9319 .arg(&mut *k)
9320 .arg(&mut *v)
9321 .arg(&nc)
9322 .arg(&rqi)
9323 .arg(&rki)
9324 .arg(pos)
9325 .arg(&nhq)
9326 .arg(&nhk)
9327 .arg(&theta_scale)
9328 .arg(&freq_scale)
9329 .arg(&null)
9330 .arg(&eps)
9331 .arg(&mut *kc)
9332 .arg(&mut *vc)
9333 .arg(t_dev)
9334 .arg(&ktb)
9335 .arg(&vtb);
9336 unsafe {
9337 b.launch(cfg)?;
9338 }
9339 }
9340 }
9341 Ok(())
9342 }
9343
9344 #[allow(clippy::too_many_arguments)]
9352 pub fn rms_norm_qkv_rope_append(
9353 &self,
9354 q0: &CudaSlice<f32>,
9355 k0: &CudaSlice<f32>,
9356 v0: &CudaSlice<f32>,
9357 wq: &CudaSlice<f32>,
9358 wk: &CudaSlice<f32>,
9359 wv: &CudaSlice<f32>,
9360 q: &mut CudaSlice<f32>,
9361 k: &mut CudaSlice<f32>,
9362 v: &mut CudaSlice<f32>,
9363 head_dim: usize,
9364 n_rot: usize,
9365 rq: usize,
9366 rk: usize,
9367 pos: &CudaSlice<i32>,
9368 nh_q: usize,
9369 nh_k: usize,
9370 base: f32,
9371 freq_scale: f32,
9372 ff: Option<&CudaSlice<f32>>,
9373 eps: f32,
9374 kc: &mut CudaSlice<u8>,
9375 vc: &mut CudaSlice<u8>,
9376 t: usize,
9377 k_tok_bytes: usize,
9378 v_tok_bytes: usize,
9379 g: bool,
9380 ) -> Result<(), Box<dyn std::error::Error>> {
9381 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
9382 let rows = rq + rk + rk;
9383 let theta_scale = base.powf(-2.0 / head_dim as f32);
9384 let (nc, rqi, rki, nhq, nhk) = (
9385 head_dim as i32,
9386 rq as i32,
9387 rk as i32,
9388 nh_q as i32,
9389 nh_k as i32,
9390 );
9391 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9392 let ti = t as i32;
9393 if Self::pdl_on() && Self::pdl_wb_on() {
9394 use cudarc::driver::{DevicePtr, DevicePtrMut};
9395 let s = &self.gpu.stream();
9396 let (p0, _a0) = q0.device_ptr(s);
9397 let (p1, _a1) = k0.device_ptr(s);
9398 let (p2, _a2) = v0.device_ptr(s);
9399 let (pwq, _a3) = wq.device_ptr(s);
9400 let (pwk, _a4) = wk.device_ptr(s);
9401 let (pwv, _a5) = wv.device_ptr(s);
9402 let (pq, _a6) = q.device_ptr_mut(s);
9403 let (pk, _a7) = k.device_ptr_mut(s);
9404 let (pv, _a8) = v.device_ptr_mut(s);
9405 let (pp, _a9) = pos.device_ptr(s);
9406 let pff: u64 = match ff {
9407 Some(t) => {
9408 let (p, _gg) = t.device_ptr(s);
9409 p as u64
9410 }
9411 None => 0,
9412 };
9413 let (pkc, _a10) = kc.device_ptr_mut(s);
9414 let (pvc, _a11) = vc.device_ptr_mut(s);
9415 let mut ps = [
9416 &p0 as *const _ as *mut std::ffi::c_void,
9417 &p1 as *const _ as *mut _,
9418 &p2 as *const _ as *mut _,
9419 &pwq as *const _ as *mut _,
9420 &pwk as *const _ as *mut _,
9421 &pwv as *const _ as *mut _,
9422 &pq as *const _ as *mut _,
9423 &pk as *const _ as *mut _,
9424 &pv as *const _ as *mut _,
9425 &nc as *const _ as *mut _,
9426 &rqi as *const _ as *mut _,
9427 &rki as *const _ as *mut _,
9428 &pp as *const _ as *mut _,
9429 &nhq as *const _ as *mut _,
9430 &nhk as *const _ as *mut _,
9431 &theta_scale as *const _ as *mut _,
9432 &freq_scale as *const _ as *mut _,
9433 &pff as *const _ as *mut _,
9434 &eps as *const _ as *mut _,
9435 &pkc as *const _ as *mut _,
9436 &pvc as *const _ as *mut _,
9437 &ti as *const _ as *mut _,
9438 &ktb as *const _ as *mut _,
9439 &vtb as *const _ as *mut _,
9440 ];
9441 unsafe {
9442 self.launch_pdl_flash(
9443 g,
9444 "rms_norm_qkv_rope_append_f32",
9445 (rows as u32, 1, 1),
9446 (rms_block(), 1, 1),
9447 0,
9448 &mut ps,
9449 )?;
9450 }
9451 return Ok(());
9452 }
9453 let f = if g {
9454 self.func_g("rms_norm_qkv_rope_append_f32")
9455 } else {
9456 self.func("rms_norm_qkv_rope_append_f32")
9457 };
9458 let cfg = LaunchConfig {
9459 grid_dim: (rows as u32, 1, 1),
9460 block_dim: (rms_block(), 1, 1),
9461 shared_mem_bytes: 0,
9462 };
9463 let __s_b = self.gpu.stream();
9464 let mut b = __s_b.launch_builder(&f);
9465 let null: u64 = 0;
9466 b.arg(q0)
9467 .arg(k0)
9468 .arg(v0)
9469 .arg(wq)
9470 .arg(wk)
9471 .arg(wv)
9472 .arg(&mut *q)
9473 .arg(&mut *k)
9474 .arg(&mut *v)
9475 .arg(&nc)
9476 .arg(&rqi)
9477 .arg(&rki)
9478 .arg(pos)
9479 .arg(&nhq)
9480 .arg(&nhk)
9481 .arg(&theta_scale)
9482 .arg(&freq_scale);
9483 match ff {
9484 Some(t) => {
9485 b.arg(t);
9486 }
9487 None => {
9488 b.arg(&null);
9489 }
9490 }
9491 b.arg(&eps)
9492 .arg(&mut *kc)
9493 .arg(&mut *vc)
9494 .arg(&ti)
9495 .arg(&ktb)
9496 .arg(&vtb);
9497 unsafe {
9498 b.launch(cfg)?;
9499 }
9500 Ok(())
9501 }
9502
9503 pub fn add_q8_1(
9504 &self,
9505 a: &CudaSlice<f32>,
9506 b: &CudaSlice<f32>,
9507 res: &mut CudaSlice<f32>,
9508 ncols: usize,
9509 nrows: usize,
9510 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9511 debug_assert!(ncols % 128 == 0);
9512 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9513 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9514 let f = self.func("add_q8_1_f32");
9515 let cfg = LaunchConfig {
9516 grid_dim: (nrows as u32, 1, 1),
9517 block_dim: (rms_block(), 1, 1),
9518 shared_mem_bytes: 0,
9519 };
9520 let nc = ncols as i32;
9521 let __s_b2 = self.gpu.stream();
9522 let mut b2 = __s_b2.launch_builder(&f);
9523 b2.arg(a)
9524 .arg(b)
9525 .arg(&mut *res)
9526 .arg(&mut out_q)
9527 .arg(&mut out_d)
9528 .arg(&nc);
9529 unsafe {
9530 b2.launch(cfg)?;
9531 }
9532 Ok((out_q, out_d))
9533 }
9534
9535 pub fn rms_pre_add_q8_1(
9539 &self,
9540 a: &CudaSlice<f32>,
9541 wa: &CudaSlice<f32>,
9542 b: &CudaSlice<f32>,
9543 res: &mut CudaSlice<f32>,
9544 ncols: usize,
9545 nrows: usize,
9546 eps: f32,
9547 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9548 debug_assert!(ncols % 128 == 0);
9549 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9550 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9551 let f = self.func("rms_pre_add_q8_1_f32");
9552 let cfg = LaunchConfig {
9553 grid_dim: (nrows as u32, 1, 1),
9554 block_dim: (rms_block(), 1, 1),
9555 shared_mem_bytes: 0,
9556 };
9557 let (nc, ep) = (ncols as i32, eps);
9558 let __s_b2 = self.gpu.stream();
9559 let mut b2 = __s_b2.launch_builder(&f);
9560 b2.arg(a)
9561 .arg(wa)
9562 .arg(b)
9563 .arg(&mut *res)
9564 .arg(&mut out_q)
9565 .arg(&mut out_d)
9566 .arg(&nc)
9567 .arg(&ep);
9568 unsafe {
9569 b2.launch(cfg)?;
9570 }
9571 Ok((out_q, out_d))
9572 }
9573
9574 pub fn l2_v2_on(ncols: usize) -> bool {
9578 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
9579 }
9580
9581 pub fn l2_norm_pp(
9582 &self,
9583 x: &CudaSlice<f32>,
9584 dst: &mut CudaSlice<f32>,
9585 dst16: Option<&mut CudaSlice<u8>>,
9586 ncols: usize,
9587 nrows: usize,
9588 eps: f32,
9589 ) -> Result<(), Box<dyn std::error::Error>> {
9590 if Self::l2_v2_on(ncols) {
9591 let f = self.func("l2_norm_pp_v2_f32");
9592 let rows_per_block = 8u32; let cfg = LaunchConfig {
9594 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
9595 block_dim: (256, 1, 1),
9596 shared_mem_bytes: 0,
9597 };
9598 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9599 let d16: u64 = match dst16 {
9601 Some(d) => self.addr_u8(d),
9602 None => 0,
9603 };
9604 let __s_b = self.gpu.stream();
9605 let mut b = __s_b.launch_builder(&f);
9606 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
9607 unsafe {
9608 b.launch(cfg)?;
9609 }
9610 return Ok(());
9611 }
9612 self.l2_norm(x, dst, ncols, nrows, eps)
9613 }
9614
9615 pub fn l2_norm(
9616 &self,
9617 x: &CudaSlice<f32>,
9618 dst: &mut CudaSlice<f32>,
9619 ncols: usize,
9620 nrows: usize,
9621 eps: f32,
9622 ) -> Result<(), Box<dyn std::error::Error>> {
9623 let f = self.func("l2_norm_f32");
9624 let cfg = LaunchConfig {
9625 grid_dim: (nrows as u32, 1, 1),
9626 block_dim: (256, 1, 1),
9627 shared_mem_bytes: 0,
9628 };
9629 let (nc, e) = (ncols as i32, eps);
9630 let __s_b = self.gpu.stream();
9631 let mut b = __s_b.launch_builder(&f);
9632 b.arg(x).arg(dst).arg(&nc).arg(&e);
9633 unsafe {
9634 b.launch(cfg)?;
9635 }
9636 Ok(())
9637 }
9638
9639 pub fn l2_norm_decode(
9645 &self,
9646 x: &CudaSlice<f32>,
9647 dst: &mut CudaSlice<f32>,
9648 ncols: usize,
9649 nrows: usize,
9650 eps: f32,
9651 ) -> Result<(), Box<dyn std::error::Error>> {
9652 let f = self.func("l2_norm_f32");
9653 let cfg = LaunchConfig {
9654 grid_dim: (nrows as u32, 1, 1),
9655 block_dim: (32, 1, 1),
9656 shared_mem_bytes: 0,
9657 };
9658 let (nc, e) = (ncols as i32, eps);
9659 let __s_b = self.gpu.stream();
9660 let mut b = __s_b.launch_builder(&f);
9661 b.arg(x).arg(dst).arg(&nc).arg(&e);
9662 unsafe {
9663 b.launch(cfg)?;
9664 }
9665 Ok(())
9666 }
9667
9668 pub fn rope_neox(
9670 &self,
9671 x: &mut CudaSlice<f32>,
9672 pos: &CudaSlice<i32>,
9673 head_dim: usize,
9674 n_dims: usize,
9675 n_heads: usize,
9676 n_tokens: usize,
9677 freq_base: f32,
9678 freq_scale: f32,
9679 ) -> Result<(), Box<dyn std::error::Error>> {
9680 let f = self.func("rope_neox_f32");
9681 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9682 let grid = (n_heads * n_tokens) as u32;
9683 let cfg = LaunchConfig {
9684 grid_dim: (grid, 1, 1),
9685 block_dim: ((head_dim / 2) as u32, 1, 1),
9686 shared_mem_bytes: 0,
9687 };
9688 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9689 let __s_b = self.gpu.stream();
9690 let mut b = __s_b.launch_builder(&f);
9691 b.arg(x)
9692 .arg(pos)
9693 .arg(&hd)
9694 .arg(&nd)
9695 .arg(&nh)
9696 .arg(&theta_scale)
9697 .arg(&freq_scale);
9698 unsafe {
9699 b.launch(cfg)?;
9700 }
9701 Ok(())
9702 }
9703
9704 pub fn rope_neox_ff(
9706 &self,
9707 x: &mut CudaSlice<f32>,
9708 pos: &CudaSlice<i32>,
9709 head_dim: usize,
9710 n_dims: usize,
9711 n_heads: usize,
9712 n_tokens: usize,
9713 freq_base: f32,
9714 freq_scale: f32,
9715 ff: &CudaSlice<f32>,
9716 ) -> Result<(), Box<dyn std::error::Error>> {
9717 let f = self.func("rope_neox_ff_f32");
9718 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9719 let grid = (n_heads * n_tokens) as u32;
9720 let cfg = LaunchConfig {
9721 grid_dim: (grid, 1, 1),
9722 block_dim: ((head_dim / 2) as u32, 1, 1),
9723 shared_mem_bytes: 0,
9724 };
9725 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
9726 let __s_b = self.gpu.stream();
9727 let mut b = __s_b.launch_builder(&f);
9728 b.arg(x)
9729 .arg(pos)
9730 .arg(&hd)
9731 .arg(&nd)
9732 .arg(&nh)
9733 .arg(&theta_scale)
9734 .arg(&freq_scale)
9735 .arg(ff);
9736 unsafe {
9737 b.launch(cfg)?;
9738 }
9739 Ok(())
9740 }
9741
9742 #[allow(clippy::too_many_arguments)]
9744 pub fn rope_neox2(
9745 &self,
9746 q: &mut CudaSlice<f32>,
9747 k: &mut CudaSlice<f32>,
9748 pos: &CudaSlice<i32>,
9749 head_dim: usize,
9750 n_dims: usize,
9751 nh_q: usize,
9752 nh_k: usize,
9753 n_tokens: usize,
9754 freq_base: f32,
9755 freq_scale: f32,
9756 ff: Option<&CudaSlice<f32>>,
9757 ) -> Result<(), Box<dyn std::error::Error>> {
9758 let f = self.func("rope_neox2_f32");
9759 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
9760 let grid = ((nh_q + nh_k) * n_tokens) as u32;
9761 let cfg = LaunchConfig {
9762 grid_dim: (grid, 1, 1),
9763 block_dim: ((head_dim / 2) as u32, 1, 1),
9764 shared_mem_bytes: 0,
9765 };
9766 let (hd, nd, nq, nk, nt) = (
9767 head_dim as i32,
9768 n_dims as i32,
9769 nh_q as i32,
9770 nh_k as i32,
9771 n_tokens as i32,
9772 );
9773 let __s_b = self.gpu.stream();
9774 let mut b = __s_b.launch_builder(&f);
9775 b.arg(q)
9776 .arg(k)
9777 .arg(pos)
9778 .arg(&hd)
9779 .arg(&nd)
9780 .arg(&nq)
9781 .arg(&nk)
9782 .arg(&nt)
9783 .arg(&theta_scale)
9784 .arg(&freq_scale);
9785 match ff {
9786 Some(ffv) => {
9787 b.arg(ffv);
9788 unsafe {
9789 b.launch(cfg)?;
9790 }
9791 }
9792 None => {
9793 let null: u64 = 0;
9794 b.arg(&null);
9795 unsafe {
9796 b.launch(cfg)?;
9797 }
9798 }
9799 }
9800 Ok(())
9801 }
9802
9803 pub fn gelu_tanh_mul(
9805 &self,
9806 gate: &CudaSlice<f32>,
9807 up: &CudaSlice<f32>,
9808 dst: &mut CudaSlice<f32>,
9809 n: usize,
9810 ) -> Result<(), Box<dyn std::error::Error>> {
9811 let f = self.func("gelu_tanh_mul_f32");
9812 let cfg = LaunchConfig::for_num_elems(n as u32);
9813 let ni = n as i32;
9814 let __s_b = self.gpu.stream();
9815 let mut b = __s_b.launch_builder(&f);
9816 b.arg(gate).arg(up).arg(dst).arg(&ni);
9817 unsafe {
9818 b.launch(cfg)?;
9819 }
9820 Ok(())
9821 }
9822
9823 pub fn silu_mul(
9824 &self,
9825 gate: &CudaSlice<f32>,
9826 up: &CudaSlice<f32>,
9827 dst: &mut CudaSlice<f32>,
9828 n: usize,
9829 ) -> Result<(), Box<dyn std::error::Error>> {
9830 let f = self.func("silu_mul_f32");
9831 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9833 let ni = n as i32;
9834 let __s_b = self.gpu.stream();
9835 let mut b = __s_b.launch_builder(&f);
9836 b.arg(gate).arg(up).arg(dst).arg(&ni);
9837 unsafe {
9838 b.launch(cfg)?;
9839 }
9840 Ok(())
9841 }
9842
9843 pub fn silu_mul_f16out(
9846 &self,
9847 gate: &CudaSlice<f32>,
9848 up: &CudaSlice<f32>,
9849 dst: &mut CudaSlice<f32>,
9850 dst16: &mut CudaSlice<u8>,
9851 n: usize,
9852 ) -> Result<(), Box<dyn std::error::Error>> {
9853 let f = self.func("silu_mul_f16out_f32");
9854 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9855 let ni = n as i32;
9856 let __s_b = self.gpu.stream();
9857 let mut b = __s_b.launch_builder(&f);
9858 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
9859 unsafe {
9860 b.launch(cfg)?;
9861 }
9862 Ok(())
9863 }
9864
9865 pub fn silu_mul_scaled(
9872 &self,
9873 gate: &CudaSlice<f32>,
9874 up: &CudaSlice<f32>,
9875 gs: f32,
9876 us: f32,
9877 dst: &mut CudaSlice<f32>,
9878 n: usize,
9879 ) -> Result<(), Box<dyn std::error::Error>> {
9880 let f = self.func("silu_mul_scaled_f32");
9881 let cfg = LaunchConfig::for_num_elems(n as u32);
9882 let ni = n as i32;
9883 let (gsf, usf) = (gs, us);
9884 let __s_b = self.gpu.stream();
9885 let mut b = __s_b.launch_builder(&f);
9886 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
9887 unsafe {
9888 b.launch(cfg)?;
9889 }
9890 Ok(())
9891 }
9892
9893 #[allow(clippy::too_many_arguments)]
9897 pub fn swigluoai_mul_scaled(
9898 &self,
9899 gate: &CudaSlice<f32>,
9900 up: &CudaSlice<f32>,
9901 gs: f32,
9902 us: f32,
9903 alpha: f32,
9904 limit: f32,
9905 dst: &mut CudaSlice<f32>,
9906 n: usize,
9907 ) -> Result<(), Box<dyn std::error::Error>> {
9908 let f = self.func("swigluoai_mul_scaled_f32");
9909 let cfg = LaunchConfig::for_num_elems(n as u32);
9910 let ni = n as i32;
9911 let __s_b = self.gpu.stream();
9912 let mut b = __s_b.launch_builder(&f);
9913 b.arg(gate)
9914 .arg(up)
9915 .arg(&gs)
9916 .arg(&us)
9917 .arg(&alpha)
9918 .arg(&limit)
9919 .arg(dst)
9920 .arg(&ni);
9921 unsafe {
9922 b.launch(cfg)?;
9923 }
9924 Ok(())
9925 }
9926
9927 pub fn silu_mul_scaled_q8_1(
9935 &self,
9936 gate: &CudaSlice<f32>,
9937 up: &CudaSlice<f32>,
9938 gs: f32,
9939 us: f32,
9940 n: usize,
9941 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9942 let f = self.func("silu_mul_scaled_q8_1");
9943 let nblk = n / 32;
9944 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);
9948 let (gsf, usf, ni) = (gs, us, n as i32);
9949 let __s_b = self.gpu.stream();
9950 let mut b = __s_b.launch_builder(&f);
9951 b.arg(gate)
9952 .arg(up)
9953 .arg(&gsf)
9954 .arg(&usf)
9955 .arg(&mut aq)
9956 .arg(&mut ad)
9957 .arg(&ni);
9958 unsafe {
9959 b.launch(cfg)?;
9960 }
9961 Ok((aq, ad))
9962 }
9963
9964 pub fn add(
9965 &self,
9966 a: &CudaSlice<f32>,
9967 b_in: &CudaSlice<f32>,
9968 dst: &mut CudaSlice<f32>,
9969 n: usize,
9970 ) -> Result<(), Box<dyn std::error::Error>> {
9971 let f = self.func("add_f32");
9972 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
9974 let ni = n as i32;
9975 let __s_bld = self.gpu.stream();
9976 let mut bld = __s_bld.launch_builder(&f);
9977 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9978 unsafe {
9979 bld.launch(cfg)?;
9980 }
9981 Ok(())
9982 }
9983
9984 pub fn mul(
9985 &self,
9986 a: &CudaSlice<f32>,
9987 b_in: &CudaSlice<f32>,
9988 dst: &mut CudaSlice<f32>,
9989 n: usize,
9990 ) -> Result<(), Box<dyn std::error::Error>> {
9991 let f = self.func("mul_f32");
9992 let cfg = LaunchConfig::for_num_elems(n as u32);
9993 let ni = n as i32;
9994 let __s_bld = self.gpu.stream();
9995 let mut bld = __s_bld.launch_builder(&f);
9996 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
9997 unsafe {
9998 bld.launch(cfg)?;
9999 }
10000 Ok(())
10001 }
10002
10003 pub fn matmul(
10006 &self,
10007 w: &crate::model::GpuTensor,
10008 x: &CudaSlice<f32>,
10009 m: usize,
10010 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10011 use crate::model::GpuTensor;
10012 let in_f = w.in_features();
10013 let out_f = w.out_features();
10014 #[allow(non_snake_case)]
10022 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
10025 usize::MAX
10026 } else {
10027 16usize
10028 };
10029
10030 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
10055 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
10056 return Ok(y);
10057 }
10058 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
10065 return Ok(y);
10066 }
10067 if let Some(y) = self.try_f16_gemm(w, x, m)? {
10070 return Ok(y);
10071 }
10072 }
10073 if let GpuTensor::Quant { qtype, .. } = w {
10088 if *qtype == QT_F8_E4M3_BLK {
10089 if m >= GEMM_M_THRESHOLD {
10090 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
10091 return Ok(y);
10092 }
10093 }
10094 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10095 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10096 return Ok(y);
10097 }
10098 }
10099 }
10100 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
10101 return self.qmatvec_mmq(w, x, m);
10102 }
10103 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
10104 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10105 return self.qmatvec_gemm(w, &aq, &ad, m);
10106 }
10107 if m >= GEMM_M_THRESHOLD {
10110 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
10111 return Ok(y);
10112 }
10113 }
10114 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
10118 if m == 1 && fast {
10123 if let GpuTensor::Quant {
10124 bytes,
10125 qtype,
10126 row_bytes,
10127 rp,
10128 rp4,
10129 scale,
10130 ..
10131 } = w
10132 {
10133 if self.mmvq_supports(*qtype) {
10134 let (bytes, rp) = match rp4 {
10138 Some(m4) => (m4, true),
10139 None => (bytes, *rp),
10140 };
10141 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10142 return self.qmatvec_mmvq(
10143 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
10144 );
10145 }
10146 }
10147 }
10148 if (2..=16).contains(&m)
10164 && fast
10165 && std::env::var("MEMRA_NO_BATCHED").is_err()
10166 && (m <= 4 || Self::b8_enabled())
10167 {
10168 let m_ok = m <= 8
10178 || matches!(w, GpuTensor::Quant { qtype, .. }
10179 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
10180 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
10181 if m_ok {
10182 if let GpuTensor::Quant {
10183 bytes,
10184 qtype,
10185 row_bytes,
10186 rp,
10187 rp4,
10188 ..
10189 } = w
10190 {
10191 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
10192 let (bytes, rp) = match rp4 {
10193 Some(m4) => (m4, true),
10194 None => (bytes, *rp),
10195 };
10196 let mcols = Self::batched_mcols(m);
10197 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10198 let mut y = self.qmatvec_mmvq_batched(
10199 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
10200 )?;
10201 if let GpuTensor::Quant { scale, .. } = w {
10202 if *scale != 1.0 {
10203 self.scale_inplace(&mut y, *scale, m * out_f)?;
10204 }
10205 }
10206 return Ok(y);
10207 }
10208 }
10209 }
10210 }
10211 if fast {
10217 if let GpuTensor::Quant {
10218 bytes,
10219 qtype,
10220 row_bytes,
10221 scale,
10222 ..
10223 } = w
10224 {
10225 if *qtype == QT_F8_E4M3 {
10226 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10227 return self.qmatvec_mmvq(
10228 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
10229 );
10230 }
10231 }
10232 }
10233 let mut y = match w {
10234 GpuTensor::Quant {
10235 bytes,
10236 qtype,
10237 row_bytes,
10238 ..
10239 } if fast && *qtype == QT_Q8_0 => {
10240 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10241 }
10242 GpuTensor::Quant {
10243 bytes,
10244 qtype,
10245 row_bytes,
10246 ..
10247 } if fast && *qtype == QT_Q4_K => {
10248 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10249 }
10250 GpuTensor::Quant {
10251 bytes,
10252 qtype,
10253 row_bytes,
10254 ..
10255 } if fast && *qtype == QT_Q6_K => {
10256 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10257 }
10258 GpuTensor::Quant {
10259 bytes,
10260 qtype,
10261 row_bytes,
10262 ..
10263 } if fast && *qtype == QT_Q5_K => {
10264 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10265 }
10266 GpuTensor::Quant {
10267 bytes,
10268 qtype,
10269 row_bytes,
10270 ..
10271 } if fast && *qtype == QT_Q3_K => {
10272 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10273 }
10274 GpuTensor::Quant {
10275 bytes,
10276 qtype,
10277 row_bytes,
10278 rp,
10279 ..
10280 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
10281 if *rp {
10282 "qmatvec_nvfp4_dp4a_rp"
10283 } else {
10284 "qmatvec_nvfp4_dp4a"
10285 },
10286 bytes,
10287 x,
10288 m,
10289 in_f,
10290 out_f,
10291 *row_bytes,
10292 )?,
10293 GpuTensor::Quant {
10297 bytes,
10298 qtype,
10299 row_bytes,
10300 ..
10301 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
10302 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
10303 }
10304 GpuTensor::Quant {
10309 bytes,
10310 qtype,
10311 row_bytes,
10312 rp,
10313 ..
10314 } =>
10315 {
10318 self.qmatvec(
10319 bytes,
10320 x,
10321 m,
10322 in_f,
10323 out_f,
10324 if *rp && *qtype == QT_NVFP4 {
10325 QT_NVFP4_RP
10326 } else {
10327 *qtype
10328 },
10329 *row_bytes,
10330 )?
10331 }
10332 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
10333 GpuTensor::FloatBf16 { data, .. } => {
10336 self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?
10337 }
10338 };
10339 if let GpuTensor::Quant { scale, .. } = w {
10341 if *scale != 1.0 {
10342 self.scale_inplace(&mut y, *scale, m * out_f)?;
10343 }
10344 }
10345 Ok(y)
10346 }
10347
10348 pub fn stage_a_raw_needed() -> bool {
10358 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10359 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
10360 }
10361
10362 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
10365 use crate::model::GpuTensor;
10366 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
10367 return false;
10368 }
10369 match w {
10370 GpuTensor::Quant { qtype, .. } => {
10377 matches!(
10378 *qtype,
10379 QT_Q8_0
10380 | QT_Q4_K
10381 | QT_Q6_K
10382 | QT_Q5_K
10383 | QT_Q3_K
10384 | QT_NVFP4
10385 | QT_F8_E4M3
10386 | QT_F8_E4M3_BLK
10387 | QT_Q4_0
10388 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
10389 }
10390 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
10391 }
10392 }
10393
10394 pub fn matmul_pre(
10399 &self,
10400 w: &crate::model::GpuTensor,
10401 aq: &CudaSlice<i8>,
10402 ad: &CudaSlice<f32>,
10403 x_fallback: &CudaSlice<f32>,
10404 m: usize,
10405 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10406 use crate::model::GpuTensor;
10407 let x_raw_ok = x_fallback.len() >= m * w.in_features();
10413 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10416 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
10417 return Ok(y);
10418 }
10419 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
10422 return Ok(y);
10423 }
10424 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
10426 return Ok(y);
10427 }
10428 }
10429 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10435 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
10436 return Ok(y);
10437 }
10438 }
10439 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10440 return Ok(y);
10441 }
10442 if m >= 16
10447 && w.out_features() >= 128
10448 && self.mmq_supports(w)
10449 && !self.verify_exact_on()
10450 && x_raw_ok
10451 {
10452 return self.qmatvec_mmq(w, x_fallback, m);
10453 }
10454 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
10457 if let Some(y) =
10458 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
10459 {
10460 return Ok(y);
10461 }
10462 }
10463 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
10466 return self.qmatvec_gemm(w, aq, ad, m);
10467 }
10468 if !self.uses_q8_1_fast(w) {
10487 if !x_raw_ok {
10488 return Err(format!(
10489 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
10490 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
10491 dropped the f32, so there is nothing to fall back to — pass the real f32 \
10492 activation (see Engine::rms_norm_decode, which is bit-identical to \
10493 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
10494 x_fallback.len(),
10495 m,
10496 w.in_features(),
10497 m * w.in_features()
10498 )
10499 .into());
10500 }
10501 return self.matmul(w, x_fallback, m);
10502 }
10503 let in_f = w.in_features();
10504 let out_f = w.out_features();
10505 let (bytes, qtype, row_bytes, scale, rp) = match w {
10506 GpuTensor::Quant {
10507 bytes,
10508 qtype,
10509 row_bytes,
10510 scale,
10511 rp,
10512 ..
10513 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10514 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
10515 };
10516 let (mbytes, mrp) = match w {
10519 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10520 _ => (bytes, rp),
10521 };
10522 if m == 1 && self.mmvq_supports(qtype) {
10526 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
10527 }
10528 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10541 && std::env::var("MEMRA_NO_BATCHED").is_err()
10542 && (m <= 4 || Self::b8_enabled())
10543 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
10547 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
10548 {
10549 let mcols = Self::batched_mcols(m);
10550 return self.qmatvec_mmvq_batched(
10551 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
10552 );
10553 }
10554 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
10560 let (b2, r2) = if qtype == QT_Q4_0 {
10561 (mbytes, mrp)
10562 } else {
10563 (bytes, rp)
10564 };
10565 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
10566 }
10567 let name = match qtype {
10568 QT_Q8_0 => "qmatvec_q8_0_dp4a",
10569 QT_Q4_K => "qmatvec_q4_K_dp4a",
10570 QT_Q6_K => "qmatvec_q6_K_dp4a",
10571 QT_Q5_K => "qmatvec_q5_K_dp4a",
10572 QT_Q3_K => "qmatvec_q3_K_dp4a",
10573 QT_NVFP4 => {
10574 if rp {
10575 "qmatvec_nvfp4_dp4a_rp"
10576 } else {
10577 "qmatvec_nvfp4_dp4a"
10578 }
10579 }
10580 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
10581 _ => unreachable!(),
10582 };
10583 let f = self.func(name);
10584 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10586 grid_dim: (out_f as u32, m as u32, 1),
10587 block_dim: (128, 1, 1),
10588 shared_mem_bytes: 0,
10589 };
10590 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10591 let __s_b = self.gpu.stream();
10592 let mut b = __s_b.launch_builder(&f);
10593 b.arg(bytes)
10594 .arg(aq)
10595 .arg(ad)
10596 .arg(&mut y)
10597 .arg(&inf)
10598 .arg(&outf)
10599 .arg(&mi)
10600 .arg(&rb);
10601 unsafe {
10602 b.launch(cfg)?;
10603 }
10604 if scale != 1.0 {
10605 self.scale_inplace(&mut y, scale, m * out_f)?;
10606 }
10607 Ok(y)
10608 }
10609
10610 pub fn matmul_decode_exact(
10618 &self,
10619 w: &crate::model::GpuTensor,
10620 x: &CudaSlice<f32>,
10621 m: usize,
10622 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10623 use crate::model::GpuTensor;
10624 if let GpuTensor::Float { data, .. } = w {
10632 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
10633 }
10634 if let GpuTensor::FloatBf16 { data, .. } = w {
10637 let (in_f, out_f) = (w.in_features(), w.out_features());
10638 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
10639 }
10640 if !self.uses_q8_1_fast(w) {
10641 return self.matmul(w, x, m);
10642 }
10643 let in_f = w.in_features();
10644 let out_f = w.out_features();
10645 let (bytes, qtype, row_bytes, scale, rp) = match w {
10646 GpuTensor::Quant {
10647 bytes,
10648 qtype,
10649 row_bytes,
10650 scale,
10651 rp,
10652 ..
10653 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10654 _ => return self.matmul(w, x, m),
10655 };
10656 let (bytes, rp) = match w {
10659 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10660 _ => (bytes, rp),
10661 };
10662 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10663 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
10667 return Ok(y);
10668 }
10669 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
10678 && std::env::var("MEMRA_NO_BATCHED").is_err()
10679 && (m <= 4 || Self::b8_enabled())
10680 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
10683 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
10684 {
10685 let mcols = Self::batched_mcols(m);
10686 return self.qmatvec_mmvq_batched(
10687 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10688 );
10689 }
10690 if self.mmvq_supports(qtype) {
10691 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10694 }
10695 self.matmul_pre(w, &aq, &ad, x, m)
10698 }
10699
10700 pub fn matmul_decode_exact_pre(
10710 &self,
10711 w: &crate::model::GpuTensor,
10712 aq: &CudaSlice<i8>,
10713 ad: &CudaSlice<f32>,
10714 m: usize,
10715 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10716 use crate::model::GpuTensor;
10717 debug_assert!(
10718 self.uses_q8_1_fast(w),
10719 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
10720 );
10721 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
10723 return Ok(y);
10724 }
10725 let in_f = w.in_features();
10726 let out_f = w.out_features();
10727 let (bytes, qtype, row_bytes, scale, rp) = match w {
10728 GpuTensor::Quant {
10729 bytes,
10730 qtype,
10731 row_bytes,
10732 scale,
10733 rp,
10734 ..
10735 } => (bytes, *qtype, *row_bytes, *scale, *rp),
10736 _ => {
10737 return Err(
10738 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
10739 );
10740 }
10741 };
10742 let (bytes, rp) = match w {
10744 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10745 _ => (bytes, rp),
10746 };
10747 if (2..=16).contains(&m)
10749 && self.batched_supports(qtype)
10750 && self.mmvq_supports(qtype)
10751 && std::env::var("MEMRA_NO_BATCHED").is_err()
10752 && (m <= 4 || Self::b8_enabled())
10753 && (m <= 8
10754 || qtype == QT_Q4_0
10755 || qtype == QT_Q6_K
10756 || qtype == QT_F8_E4M3
10757 || qtype == QT_NVFP4
10758 || qtype == QT_Q4_K
10759 || qtype == QT_Q5_K
10760 || qtype == QT_Q8_0)
10761 {
10762 let mcols = Self::batched_mcols(m);
10763 return self.qmatvec_mmvq_batched(
10764 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
10765 );
10766 }
10767 if self.mmvq_supports(qtype) {
10768 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
10769 }
10770 let x0 = self.zeros(0)?;
10773 self.matmul_pre(w, aq, ad, &x0, m)
10774 }
10775
10776 pub fn matmul_decode_exact_dual_pre(
10785 &self,
10786 w0: &crate::model::GpuTensor,
10787 w1: &crate::model::GpuTensor,
10788 aq: &CudaSlice<i8>,
10789 ad: &CudaSlice<f32>,
10790 m: usize,
10791 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
10792 {
10793 use crate::model::GpuTensor;
10794 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10795 let on = *ON.get_or_init(|| {
10796 std::env::var("MEMRA_SPEC_DUAL_T")
10797 .map(|v| v != "0")
10798 .unwrap_or(true)
10799 });
10800 if !on
10801 || !(2..=7).contains(&m)
10802 || std::env::var("MEMRA_NO_BATCHED").is_ok()
10803 || !self.uses_q8_1_fast(w0)
10804 || !self.uses_q8_1_fast(w1)
10805 {
10806 return Ok(None);
10807 }
10808 if !self.mmvq_supports(QT_NVFP4) {
10813 return Ok(None);
10814 }
10815 let (in_f, out_f) = (w0.in_features(), w0.out_features());
10816 if w1.in_features() != in_f || w1.out_features() != out_f {
10817 return Ok(None);
10818 }
10819 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10820 (
10821 GpuTensor::Quant {
10822 bytes: b0,
10823 qtype: q0,
10824 row_bytes: rb0,
10825 scale: s0,
10826 rp: rp0,
10827 rp4: None,
10828 ..
10829 },
10830 GpuTensor::Quant {
10831 bytes: b1,
10832 qtype: q1,
10833 row_bytes: rb1,
10834 scale: s1,
10835 rp: rp1,
10836 rp4: None,
10837 ..
10838 },
10839 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10840 (b0, b1, *rb0, *s0, *s1, *rp0)
10841 }
10842 _ => return Ok(None),
10843 };
10844 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
10847 {
10848 return Ok(None);
10849 }
10850 let (y0, y1) =
10851 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
10852 Ok(Some(((y0, s0), (y1, s1))))
10853 }
10854
10855 pub fn matmul_decode_exact_dual(
10871 &self,
10872 w0: &crate::model::GpuTensor,
10873 w1: &crate::model::GpuTensor,
10874 x: &CudaSlice<f32>,
10875 m: usize,
10876 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
10877 use crate::model::GpuTensor;
10878 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10879 let on = *ON.get_or_init(|| {
10880 std::env::var("MEMRA_SPEC_DUAL_T")
10881 .map(|v| v != "0")
10882 .unwrap_or(true)
10883 });
10884 if !on
10885 || !(2..=4).contains(&m)
10886 || std::env::var("MEMRA_NO_BATCHED").is_ok()
10887 || !self.uses_q8_1_fast(w0)
10888 || !self.uses_q8_1_fast(w1)
10889 {
10890 return Ok(None);
10891 }
10892 if !self.mmvq_supports(QT_NVFP4) {
10897 return Ok(None);
10898 }
10899 let (in_f, out_f) = (w0.in_features(), w0.out_features());
10900 if w1.in_features() != in_f || w1.out_features() != out_f {
10901 return Ok(None);
10902 }
10903 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
10904 (
10905 GpuTensor::Quant {
10906 bytes: b0,
10907 qtype: q0,
10908 row_bytes: rb0,
10909 scale: s0,
10910 rp: rp0,
10911 rp4: None,
10912 ..
10913 },
10914 GpuTensor::Quant {
10915 bytes: b1,
10916 qtype: q1,
10917 row_bytes: rb1,
10918 scale: s1,
10919 rp: rp1,
10920 rp4: None,
10921 ..
10922 },
10923 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
10924 (b0, b1, *rb0, *s0, *s1, *rp0)
10925 }
10926 _ => return Ok(None),
10927 };
10928 if std::env::var("MEMRA_DEBUG").is_ok() {
10931 static ONCE: std::sync::Once = std::sync::Once::new();
10932 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
10933 }
10934 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10935 let (y0, y1) =
10936 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
10937 let mut y0 = y0;
10938 let mut y1 = y1;
10939 if s0 != 1.0 {
10940 self.scale_inplace(&mut y0, s0, m * out_f)?;
10941 }
10942 if s1 != 1.0 {
10943 self.scale_inplace(&mut y1, s1, m * out_f)?;
10944 }
10945 Ok(Some((y0, y1)))
10946 }
10947
10948 #[allow(clippy::too_many_arguments)]
10953 pub fn qmatvec_batched_dual_raw(
10954 &self,
10955 b0: &CudaSlice<u8>,
10956 b1: &CudaSlice<u8>,
10957 aq: &CudaSlice<i8>,
10958 ad: &CudaSlice<f32>,
10959 m: usize,
10960 in_f: usize,
10961 out_f: usize,
10962 row_bytes: usize,
10963 rp: bool,
10964 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10965 const ROWS_PER_BLOCK: u32 = 4;
10966 let mcols = Self::batched_mcols(m);
10967 let tiny_rp1 = rp
10970 && mcols == 4
10971 && out_f <= 128
10972 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
10973 let (name, rows_per_block) = if tiny_rp1 {
10974 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
10975 } else {
10976 match (mcols, rp, m) {
10977 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
10978 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
10979 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
10980 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
10981 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
10982 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
10983 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
10984 _ => {
10985 return Err(
10986 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
10987 );
10988 }
10989 }
10990 };
10991 let f = self.func(name);
10992 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
10993 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
10994 let cfg = LaunchConfig {
10995 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
10996 block_dim: (32, ROWS_PER_BLOCK, 1),
10997 shared_mem_bytes: 0,
10998 };
10999 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11000 let __s_b = self.gpu.stream();
11001 let mut b = __s_b.launch_builder(&f);
11002 b.arg(b0)
11003 .arg(b1)
11004 .arg(aq)
11005 .arg(ad)
11006 .arg(&mut y0)
11007 .arg(&mut y1)
11008 .arg(&inf)
11009 .arg(&outf)
11010 .arg(&mi)
11011 .arg(&rb);
11012 unsafe {
11013 b.launch(cfg)?;
11014 }
11015 Ok((y0, y1))
11016 }
11017
11018 pub fn matmul_pre_dual_noscale(
11030 &self,
11031 w0: &crate::model::GpuTensor,
11032 w1: &crate::model::GpuTensor,
11033 aq: &CudaSlice<i8>,
11034 ad: &CudaSlice<f32>,
11035 m: usize,
11036 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11037 {
11038 use crate::model::GpuTensor;
11039 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11040 return Ok(None);
11041 }
11042 if !self.mmvq_supports(QT_NVFP4) {
11052 return Ok(None);
11053 }
11054 let (in_f, out_f) = (w0.in_features(), w0.out_features());
11055 if w1.in_features() != in_f || w1.out_features() != out_f {
11056 return Ok(None);
11057 }
11058 let no_mirror =
11071 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
11072 if self.q8_ffn_fuse2_on()
11073 && no_mirror(w0)
11074 && no_mirror(w1)
11075 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
11076 {
11077 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
11078 return Ok(Some(((y0, 1.0), (y1, 1.0))));
11079 }
11080 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11090 let (y0, y1) =
11091 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
11092 return Ok(Some(((y0, p0.3), (y1, p1.3))));
11093 }
11094 let (b0, q0, rb0, s0, rp0) = match w0 {
11095 GpuTensor::Quant {
11096 bytes,
11097 qtype,
11098 row_bytes,
11099 scale,
11100 rp,
11101 ..
11102 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11103 _ => return Ok(None),
11104 };
11105 let (b1, q1, rb1, s1, rp1) = match w1 {
11106 GpuTensor::Quant {
11107 bytes,
11108 qtype,
11109 row_bytes,
11110 scale,
11111 rp,
11112 ..
11113 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11114 _ => return Ok(None),
11115 };
11116 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
11117 return Ok(None);
11118 }
11119 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11121 let rows_per_block = ROWS_PER_BLOCK * RPW;
11122 let f = self.func(if rp0 {
11123 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
11124 } else {
11125 "qmatvec_nvfp4_mmvq_dual_mr2"
11126 });
11127 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
11128 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
11129 let cfg = LaunchConfig {
11130 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
11131 block_dim: (32, ROWS_PER_BLOCK, 1),
11132 shared_mem_bytes: 0,
11133 };
11134 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
11135 let one = 1.0f32;
11138 let __s_b = self.gpu.stream();
11139 let mut b = __s_b.launch_builder(&f);
11140 b.arg(b0)
11141 .arg(b1)
11142 .arg(aq)
11143 .arg(ad)
11144 .arg(&mut y0)
11145 .arg(&mut y1)
11146 .arg(&inf)
11147 .arg(&outf)
11148 .arg(&mi)
11149 .arg(&rb)
11150 .arg(&one)
11151 .arg(&one);
11152 unsafe {
11153 b.launch(cfg)?;
11154 }
11155 Ok(Some(((y0, s0), (y1, s1))))
11156 }
11157
11158 #[allow(clippy::too_many_arguments)]
11166 pub fn matmul_nvfp4_fused3(
11167 &self,
11168 w0: &crate::model::GpuTensor,
11169 w1: &crate::model::GpuTensor,
11170 w2: &crate::model::GpuTensor,
11171 aq: &CudaSlice<i8>,
11172 ad: &CudaSlice<f32>,
11173 m: usize,
11174 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11175 {
11176 use crate::model::GpuTensor;
11177 if !(1..=8).contains(&m)
11184 || !self.mmvq_supports(QT_NVFP4)
11185 || !self.uses_q8_1_fast(w0)
11186 || !self.uses_q8_1_fast(w1)
11187 || !self.uses_q8_1_fast(w2)
11188 {
11189 return Ok(None);
11190 }
11191 if m > 1 {
11192 let in_f = w0.in_features();
11193 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
11194 || !self.batched_supports(QT_NVFP4)
11195 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11196 || (m > 4 && !Self::b8_enabled())
11197 || in_f % 512 != 0
11198 || in_f / 64 > 272
11199 {
11200 return Ok(None);
11201 }
11202 }
11203 let unpack = |w: &crate::model::GpuTensor| match w {
11204 GpuTensor::Quant {
11205 bytes,
11206 qtype,
11207 scale,
11208 rp,
11209 ..
11210 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11211 _ => None,
11212 };
11213 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
11214 return Ok(None);
11215 };
11216 let in_f = w0.in_features();
11217 if w1.in_features() != in_f || w2.in_features() != in_f {
11218 return Ok(None);
11219 }
11220 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
11221 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11223 let rows_pb = ROWS_PER_BLOCK * RPW;
11224 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11225 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11226 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11227 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11228 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
11229 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
11232 if m > 1 {
11233 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
11235 return Ok(None);
11236 }
11237 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
11238 let cfg = LaunchConfig {
11239 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
11240 block_dim: (32, ROWS_PER_BLOCK, 1),
11241 shared_mem_bytes: 0,
11242 };
11243 let __s_b = self.gpu.stream();
11244 let mut b = __s_b.launch_builder(&f);
11245 b.arg(b0)
11246 .arg(b1)
11247 .arg(b2)
11248 .arg(aq)
11249 .arg(ad)
11250 .arg(&mut y0)
11251 .arg(&mut y1)
11252 .arg(&mut y2)
11253 .arg(&inf)
11254 .arg(&oi0)
11255 .arg(&oi1)
11256 .arg(&oi2)
11257 .arg(&mi);
11258 unsafe {
11259 b.launch(cfg)?;
11260 }
11261 return Ok(Some((y0, y1, y2)));
11262 }
11263 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
11264 let cfg = LaunchConfig {
11265 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
11266 block_dim: (32, ROWS_PER_BLOCK, 1),
11267 shared_mem_bytes: 0,
11268 };
11269 let __s_b = self.gpu.stream();
11270 let mut b = __s_b.launch_builder(&f);
11271 b.arg(b0)
11272 .arg(b1)
11273 .arg(b2)
11274 .arg(aq)
11275 .arg(ad)
11276 .arg(&mut y0)
11277 .arg(&mut y1)
11278 .arg(&mut y2)
11279 .arg(&inf)
11280 .arg(&oi0)
11281 .arg(&oi1)
11282 .arg(&oi2)
11283 .arg(&mi)
11284 .arg(&p0.1)
11285 .arg(&p1.1)
11286 .arg(&p2.1);
11287 unsafe {
11288 b.launch(cfg)?;
11289 }
11290 Ok(Some((y0, y1, y2)))
11291 }
11292
11293 pub fn matmul_nvfp4_fused2(
11302 &self,
11303 w0: &crate::model::GpuTensor,
11304 w1: &crate::model::GpuTensor,
11305 aq: &CudaSlice<i8>,
11306 ad: &CudaSlice<f32>,
11307 m: usize,
11308 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11309 use crate::model::GpuTensor;
11310 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11311 let off =
11312 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11313 if off
11316 || m != 1
11317 || !self.mmvq_supports(QT_NVFP4)
11318 || !self.uses_q8_1_fast(w0)
11319 || !self.uses_q8_1_fast(w1)
11320 {
11321 return Ok(None);
11322 }
11323 let unpack = |w: &crate::model::GpuTensor| match w {
11324 GpuTensor::Quant {
11325 bytes,
11326 qtype,
11327 scale,
11328 rp,
11329 ..
11330 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11331 _ => None,
11332 };
11333 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11334 return Ok(None);
11335 };
11336 let in_f = w0.in_features();
11337 if w1.in_features() != in_f {
11338 return Ok(None);
11339 }
11340 let (o0, o1) = (w0.out_features(), w1.out_features());
11341 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11343 let rows_pb = ROWS_PER_BLOCK * RPW;
11344 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11345 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11346 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11347 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11348 let cfg = LaunchConfig {
11349 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
11350 block_dim: (32, ROWS_PER_BLOCK, 1),
11351 shared_mem_bytes: 0,
11352 };
11353 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
11354 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11357 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
11360 {
11361 use cudarc::driver::{DevicePtr, DevicePtrMut};
11362 let s = &self.gpu.stream();
11363 let (pw0, _g0) = b0.device_ptr(s);
11364 let (pw1, _g1) = b1.device_ptr(s);
11365 let (paq, _g2) = aq.device_ptr(s);
11366 let (pad, _g3) = ad.device_ptr(s);
11367 let (py0, _g4) = y0.device_ptr_mut(s);
11368 let (py1, _g5) = y1.device_ptr_mut(s);
11369 let (s0, s1) = (p0.1, p1.1);
11370 let mut ps = [
11371 &pw0 as *const _ as *mut std::ffi::c_void,
11372 &pw1 as *const _ as *mut _,
11373 &paq as *const _ as *mut _,
11374 &pad as *const _ as *mut _,
11375 &py0 as *const _ as *mut _,
11376 &py1 as *const _ as *mut _,
11377 &inf as *const _ as *mut _,
11378 &oi0 as *const _ as *mut _,
11379 &oi1 as *const _ as *mut _,
11380 &mi as *const _ as *mut _,
11381 &s0 as *const _ as *mut _,
11382 &s1 as *const _ as *mut _,
11383 ];
11384 unsafe {
11385 self.launch_pdl(
11386 "qmatvec_nvfp4_mmvq_fused2_rp",
11387 cfg.grid_dim,
11388 cfg.block_dim,
11389 &mut ps,
11390 )?;
11391 }
11392 }
11393 return Ok(Some((y0, y1)));
11394 }
11395 let __s_b = self.gpu.stream();
11396 let mut b = __s_b.launch_builder(&f);
11397 b.arg(b0)
11398 .arg(b1)
11399 .arg(aq)
11400 .arg(ad)
11401 .arg(&mut y0)
11402 .arg(&mut y1)
11403 .arg(&inf)
11404 .arg(&oi0)
11405 .arg(&oi1)
11406 .arg(&mi)
11407 .arg(&p0.1)
11408 .arg(&p1.1);
11409 unsafe {
11410 b.launch(cfg)?;
11411 }
11412 Ok(Some((y0, y1)))
11413 }
11414
11415 pub fn matmul_nvfp4_fused2_into(
11420 &self,
11421 w0: &crate::model::GpuTensor,
11422 w1: &crate::model::GpuTensor,
11423 aq: &CudaSlice<i8>,
11424 ad: &CudaSlice<f32>,
11425 y0: &mut CudaSlice<f32>,
11426 y1: &mut CudaSlice<f32>,
11427 ) -> Result<bool, Box<dyn std::error::Error>> {
11428 use crate::model::GpuTensor;
11429 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11430 let off =
11431 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
11432 if off
11433 || !self.mmvq_supports(QT_NVFP4)
11434 || !self.uses_q8_1_fast(w0)
11435 || !self.uses_q8_1_fast(w1)
11436 {
11437 return Ok(false);
11438 }
11439 let unpack = |w: &crate::model::GpuTensor| match w {
11440 GpuTensor::Quant {
11441 bytes,
11442 qtype,
11443 scale,
11444 rp,
11445 ..
11446 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
11447 _ => None,
11448 };
11449 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
11450 return Ok(false);
11451 };
11452 let in_f = w0.in_features();
11453 if w1.in_features() != in_f {
11454 return Ok(false);
11455 }
11456 let (o0, o1) = (w0.out_features(), w1.out_features());
11457 if y0.len() < o0 || y1.len() < o1 {
11458 return Ok(false);
11459 }
11460 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11462 let rows_pb = ROWS_PER_BLOCK * RPW;
11463 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11464 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
11465 let cfg = LaunchConfig {
11466 grid_dim: (nb(o0) + nb(o1), 1, 1),
11467 block_dim: (32, ROWS_PER_BLOCK, 1),
11468 shared_mem_bytes: 0,
11469 };
11470 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
11471 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
11474 let __s_b = self.gpu.stream();
11475 let mut b = __s_b.launch_builder(&f);
11476 b.arg(b0)
11477 .arg(b1)
11478 .arg(aq)
11479 .arg(ad)
11480 .arg(&mut *y0)
11481 .arg(&mut *y1)
11482 .arg(&inf)
11483 .arg(&oi0)
11484 .arg(&oi1)
11485 .arg(&mi)
11486 .arg(&p0.1)
11487 .arg(&p1.1);
11488 unsafe {
11489 b.launch(cfg)?;
11490 }
11491 Ok(true)
11492 }
11493
11494 #[allow(clippy::type_complexity)]
11499 pub fn matmul_nvfp4_fused4(
11500 &self,
11501 w0: &crate::model::GpuTensor,
11502 w1: &crate::model::GpuTensor,
11503 w2: &crate::model::GpuTensor,
11504 w3: &crate::model::GpuTensor,
11505 aq: &CudaSlice<i8>,
11506 ad: &CudaSlice<f32>,
11507 m: usize,
11508 ) -> Result<
11509 Option<(
11510 CudaSlice<f32>,
11511 CudaSlice<f32>,
11512 CudaSlice<f32>,
11513 CudaSlice<f32>,
11514 )>,
11515 Box<dyn std::error::Error>,
11516 > {
11517 use crate::model::GpuTensor;
11518 if !(1..=8).contains(&m)
11525 || std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
11526 || !self.mmvq_supports(QT_NVFP4)
11527 || !self.uses_q8_1_fast(w0)
11528 || !self.uses_q8_1_fast(w1)
11529 || !self.uses_q8_1_fast(w2)
11530 || !self.uses_q8_1_fast(w3)
11531 {
11532 return Ok(None);
11533 }
11534 if m > 1 {
11535 let in_f = w0.in_features();
11538 if !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), Some(p3)) =
11558 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
11559 else {
11560 return Ok(None);
11561 };
11562 let in_f = w0.in_features();
11563 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
11564 return Ok(None);
11565 }
11566 let (o0, o1, o2, o3) = (
11567 w0.out_features(),
11568 w1.out_features(),
11569 w2.out_features(),
11570 w3.out_features(),
11571 );
11572 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
11574 let rows_pb = ROWS_PER_BLOCK * RPW;
11575 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
11576 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
11577 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
11578 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
11579 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
11580 let (inf, oi0, oi1, oi2, oi3, mi) = (
11581 in_f as i32,
11582 o0 as i32,
11583 o1 as i32,
11584 o2 as i32,
11585 o3 as i32,
11586 m as i32,
11587 );
11588 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
11591 if m > 1 {
11592 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
11595 return Ok(None);
11596 }
11597 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
11598 let cfg = LaunchConfig {
11599 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
11600 block_dim: (32, ROWS_PER_BLOCK, 1),
11601 shared_mem_bytes: 0,
11602 };
11603 let __s_b = self.gpu.stream();
11604 let mut b = __s_b.launch_builder(&f);
11605 b.arg(b0)
11606 .arg(b1)
11607 .arg(b2)
11608 .arg(b3)
11609 .arg(aq)
11610 .arg(ad)
11611 .arg(&mut y0)
11612 .arg(&mut y1)
11613 .arg(&mut y2)
11614 .arg(&mut y3)
11615 .arg(&inf)
11616 .arg(&oi0)
11617 .arg(&oi1)
11618 .arg(&oi2)
11619 .arg(&oi3)
11620 .arg(&mi);
11621 unsafe {
11622 b.launch(cfg)?;
11623 }
11624 return Ok(Some((y0, y1, y2, y3)));
11625 }
11626 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
11627 let cfg = LaunchConfig {
11628 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
11629 block_dim: (32, ROWS_PER_BLOCK, 1),
11630 shared_mem_bytes: 0,
11631 };
11632 let __s_b = self.gpu.stream();
11633 let mut b = __s_b.launch_builder(&f);
11634 b.arg(b0)
11635 .arg(b1)
11636 .arg(b2)
11637 .arg(b3)
11638 .arg(aq)
11639 .arg(ad)
11640 .arg(&mut y0)
11641 .arg(&mut y1)
11642 .arg(&mut y2)
11643 .arg(&mut y3)
11644 .arg(&inf)
11645 .arg(&oi0)
11646 .arg(&oi1)
11647 .arg(&oi2)
11648 .arg(&oi3)
11649 .arg(&mi)
11650 .arg(&p0.1)
11651 .arg(&p1.1)
11652 .arg(&p2.1)
11653 .arg(&p3.1);
11654 unsafe {
11655 b.launch(cfg)?;
11656 }
11657 Ok(Some((y0, y1, y2, y3)))
11658 }
11659
11660 pub fn matmul_q8_fused2(
11668 &self,
11669 w0: &crate::model::GpuTensor,
11670 w1: &crate::model::GpuTensor,
11671 aq: &CudaSlice<i8>,
11672 ad: &CudaSlice<f32>,
11673 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11674 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11680 return Ok(Some(self.e4m3_fused2_core(
11681 p0.0,
11682 p1.0,
11683 aq,
11684 ad,
11685 w0.in_features(),
11686 p0.1,
11687 p1.1,
11688 p0.2,
11689 p0.3,
11690 p1.3,
11691 )?));
11692 }
11693 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11694 return Ok(None);
11695 };
11696 Ok(Some(self.q8_fused2_core(
11697 p0.0,
11698 p1.0,
11699 aq,
11700 ad,
11701 w0.in_features(),
11702 p0.1,
11703 p1.1,
11704 p0.2,
11705 )?))
11706 }
11707
11708 #[allow(clippy::too_many_arguments)]
11709 fn q8_fused2_core(
11710 &self,
11711 b0: &CudaSlice<u8>,
11712 b1: &CudaSlice<u8>,
11713 aq: &CudaSlice<i8>,
11714 ad: &CudaSlice<f32>,
11715 in_f: usize,
11716 out0: usize,
11717 out1: usize,
11718 row_bytes: usize,
11719 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11720 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
11722 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
11723 let f = self.func("qmatvec_q8_0_mmvq_fused2");
11724 let mut y0 = self.alloc_uninit::<f32>(out0)?;
11725 let mut y1 = self.alloc_uninit::<f32>(out1)?;
11726 let cfg = LaunchConfig {
11727 grid_dim: (nb0 + nb1, 1, 1),
11728 block_dim: (32, ROWS_PER_BLOCK, 1),
11729 shared_mem_bytes: 0,
11730 };
11731 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
11732 let __s_b = self.gpu.stream();
11733 let mut b = __s_b.launch_builder(&f);
11734 b.arg(b0)
11735 .arg(b1)
11736 .arg(aq)
11737 .arg(ad)
11738 .arg(&mut y0)
11739 .arg(&mut y1)
11740 .arg(&inf)
11741 .arg(&o0)
11742 .arg(&o1)
11743 .arg(&rbl);
11744 unsafe {
11745 b.launch(cfg)?;
11746 }
11747 Ok((y0, y1))
11748 }
11749
11750 pub fn matmul_q8_fused2_x(
11756 &self,
11757 w0: &crate::model::GpuTensor,
11758 w1: &crate::model::GpuTensor,
11759 x: &CudaSlice<f32>,
11760 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
11761 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
11762 return Ok(None);
11763 }
11764 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
11765 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11766 return Ok(Some(self.e4m3_fused2_core(
11767 p0.0,
11768 p1.0,
11769 &aq,
11770 &ad,
11771 w0.in_features(),
11772 p0.1,
11773 p1.1,
11774 p0.2,
11775 p0.3,
11776 p1.3,
11777 )?));
11778 }
11779 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
11780 return Ok(None);
11781 };
11782 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
11783 Ok(Some(self.q8_fused2_core(
11784 p0.0,
11785 p1.0,
11786 &aq,
11787 &ad,
11788 w0.in_features(),
11789 p0.1,
11790 p1.1,
11791 p0.2,
11792 )?))
11793 }
11794
11795 #[allow(clippy::too_many_arguments)]
11798 pub fn qmatvec_q8_fused2_raw(
11799 &self,
11800 b0: &CudaSlice<u8>,
11801 b1: &CudaSlice<u8>,
11802 x: &CudaSlice<f32>,
11803 in_f: usize,
11804 out0: usize,
11805 out1: usize,
11806 row_bytes: usize,
11807 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11808 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
11809 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
11810 }
11811
11812 pub fn matmul_q4_fused3(
11818 &self,
11819 w0: &crate::model::GpuTensor,
11820 w1: &crate::model::GpuTensor,
11821 w2: &crate::model::GpuTensor,
11822 aq: &CudaSlice<i8>,
11823 ad: &CudaSlice<f32>,
11824 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
11825 {
11826 use crate::model::GpuTensor;
11827 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11828 match w {
11829 GpuTensor::Quant {
11830 qtype, row_bytes, ..
11831 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11832 _ => None,
11833 }
11834 };
11835 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11836 return Ok(None);
11837 };
11838 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11839 return Ok(None);
11840 }
11841 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11845 match w {
11846 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11847 Some(m) => (m, true),
11848 None => (bytes, *rp),
11849 },
11850 _ => unreachable!(),
11851 }
11852 }
11853 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11854 if rp0 != rp1 || rp1 != rp2 {
11855 return Ok(None);
11856 }
11857 let rp = rp0;
11858 let rpb: u32 = 4;
11859 let mr1 = rp && Self::q40_mr1_on();
11863 let nb = |o: usize| {
11864 if mr1 {
11865 (o as u32).div_ceil(rpb)
11866 } else {
11867 (o as u32).div_ceil(2).div_ceil(rpb)
11868 }
11869 };
11870 let grid = nb(o0) + nb(o1) + nb(o2);
11871 let mut y0 = self.alloc_uninit::<f32>(o0)?;
11872 let mut y1 = self.alloc_uninit::<f32>(o1)?;
11873 let mut y2 = self.alloc_uninit::<f32>(o2)?;
11874 let f = self.func(if mr1 {
11875 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
11876 } else if rp {
11877 "qmatvec_q4_0_mmvq_fused3_rp"
11878 } else {
11879 "qmatvec_q4_0_mmvq_fused3"
11880 });
11881 let cfg = LaunchConfig {
11882 grid_dim: (grid, 1, 1),
11883 block_dim: (32, rpb, 1),
11884 shared_mem_bytes: 0,
11885 };
11886 let inf = w0.in_features() as i32;
11887 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
11888 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
11889 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
11892 {
11893 use cudarc::driver::{DevicePtr, DevicePtrMut};
11894 let s = &self.gpu.stream();
11895 let (p0, _g0) = b0.device_ptr(s);
11896 let (p1, _g1) = b1.device_ptr(s);
11897 let (p2, _g2) = b2.device_ptr(s);
11898 let (paq, _g3) = aq.device_ptr(s);
11899 let (pad, _g4) = ad.device_ptr(s);
11900 let (py0, _g5) = y0.device_ptr_mut(s);
11901 let (py1, _g6) = y1.device_ptr_mut(s);
11902 let (py2, _g7) = y2.device_ptr_mut(s);
11903 let mut ps = [
11904 &p0 as *const _ as *mut std::ffi::c_void,
11905 &p1 as *const _ as *mut _,
11906 &p2 as *const _ as *mut _,
11907 &paq as *const _ as *mut _,
11908 &pad as *const _ as *mut _,
11909 &py0 as *const _ as *mut _,
11910 &py1 as *const _ as *mut _,
11911 &py2 as *const _ as *mut _,
11912 &inf as *const _ as *mut _,
11913 &oo0 as *const _ as *mut _,
11914 &oo1 as *const _ as *mut _,
11915 &oo2 as *const _ as *mut _,
11916 &r0 as *const _ as *mut _,
11917 &r1 as *const _ as *mut _,
11918 &r2 as *const _ as *mut _,
11919 ];
11920 unsafe {
11921 self.launch_pdl(
11922 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
11923 (grid, 1, 1),
11924 (32, rpb, 1),
11925 &mut ps,
11926 )?;
11927 }
11928 }
11929 return Ok(Some((y0, y1, y2)));
11930 }
11931 let __s_b = self.gpu.stream();
11932 let mut b = __s_b.launch_builder(&f);
11933 b.arg(b0)
11934 .arg(b1)
11935 .arg(b2)
11936 .arg(aq)
11937 .arg(ad)
11938 .arg(&mut y0)
11939 .arg(&mut y1)
11940 .arg(&mut y2)
11941 .arg(&inf)
11942 .arg(&oo0)
11943 .arg(&oo1)
11944 .arg(&oo2)
11945 .arg(&r0)
11946 .arg(&r1)
11947 .arg(&r2);
11948 unsafe {
11949 b.launch(cfg)?;
11950 }
11951 Ok(Some((y0, y1, y2)))
11952 }
11953
11954 #[allow(clippy::too_many_arguments)]
11957 pub fn matmul_q4_fused3_into(
11958 &self,
11959 w0: &crate::model::GpuTensor,
11960 w1: &crate::model::GpuTensor,
11961 w2: &crate::model::GpuTensor,
11962 aq: &CudaSlice<i8>,
11963 ad: &CudaSlice<f32>,
11964 y0: &mut CudaSlice<f32>,
11965 y1: &mut CudaSlice<f32>,
11966 y2: &mut CudaSlice<f32>,
11967 ) -> Result<bool, Box<dyn std::error::Error>> {
11968 use crate::model::GpuTensor;
11969 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
11970 match w {
11971 GpuTensor::Quant {
11972 qtype, row_bytes, ..
11973 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
11974 _ => None,
11975 }
11976 };
11977 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
11978 return Ok(false);
11979 };
11980 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
11981 return Ok(false);
11982 }
11983 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
11984 match w {
11985 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
11986 Some(m) => (m, true),
11987 None => (bytes, *rp),
11988 },
11989 _ => unreachable!(),
11990 }
11991 }
11992 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
11993 if rp0 != rp1 || rp1 != rp2 {
11994 return Ok(false);
11995 }
11996 let rp = rp0;
11997 let rpb: u32 = 4;
11998 let mr1 = rp && Self::q40_mr1_on();
11999 let nb = |o: usize| {
12000 if mr1 {
12001 (o as u32).div_ceil(rpb)
12002 } else {
12003 (o as u32).div_ceil(2).div_ceil(rpb)
12004 }
12005 };
12006 let grid = nb(o0) + nb(o1) + nb(o2);
12007 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
12008 let f = self.func(if mr1 {
12009 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
12010 } else if rp {
12011 "qmatvec_q4_0_mmvq_fused3_rp"
12012 } else {
12013 "qmatvec_q4_0_mmvq_fused3"
12014 });
12015 let cfg = LaunchConfig {
12016 grid_dim: (grid, 1, 1),
12017 block_dim: (32, rpb, 1),
12018 shared_mem_bytes: 0,
12019 };
12020 let inf = w0.in_features() as i32;
12021 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
12022 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
12023 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12025 use cudarc::driver::{DevicePtr, DevicePtrMut};
12026 let s = &self.gpu.stream();
12027 let (p0, _g0) = b0.device_ptr(s);
12028 let (p1, _g1) = b1.device_ptr(s);
12029 let (p2, _g2) = b2.device_ptr(s);
12030 let (paq, _g3) = aq.device_ptr(s);
12031 let (pad, _g4) = ad.device_ptr(s);
12032 let (py0, _g5) = y0.device_ptr_mut(s);
12033 let (py1, _g6) = y1.device_ptr_mut(s);
12034 let (py2, _g7) = y2.device_ptr_mut(s);
12035 let mut ps = [
12036 &p0 as *const _ as *mut std::ffi::c_void,
12037 &p1 as *const _ as *mut _,
12038 &p2 as *const _ as *mut _,
12039 &paq as *const _ as *mut _,
12040 &pad as *const _ as *mut _,
12041 &py0 as *const _ as *mut _,
12042 &py1 as *const _ as *mut _,
12043 &py2 as *const _ as *mut _,
12044 &inf as *const _ as *mut _,
12045 &oo0 as *const _ as *mut _,
12046 &oo1 as *const _ as *mut _,
12047 &oo2 as *const _ as *mut _,
12048 &r0 as *const _ as *mut _,
12049 &r1 as *const _ as *mut _,
12050 &r2 as *const _ as *mut _,
12051 ];
12052 unsafe {
12053 self.launch_pdl(
12054 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
12055 (grid, 1, 1),
12056 (32, rpb, 1),
12057 &mut ps,
12058 )?;
12059 }
12060 return Ok(true);
12061 }
12062 let __s_b = self.gpu.stream();
12063 let mut b = __s_b.launch_builder(&f);
12064 b.arg(b0)
12065 .arg(b1)
12066 .arg(b2)
12067 .arg(aq)
12068 .arg(ad)
12069 .arg(&mut *y0)
12070 .arg(&mut *y1)
12071 .arg(&mut *y2)
12072 .arg(&inf)
12073 .arg(&oo0)
12074 .arg(&oo1)
12075 .arg(&oo2)
12076 .arg(&r0)
12077 .arg(&r1)
12078 .arg(&r2);
12079 unsafe {
12080 b.launch(cfg)?;
12081 }
12082 Ok(true)
12083 }
12084
12085 pub fn matmul_q4_fused2(
12087 &self,
12088 w0: &crate::model::GpuTensor,
12089 w1: &crate::model::GpuTensor,
12090 aq: &CudaSlice<i8>,
12091 ad: &CudaSlice<f32>,
12092 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12093 use crate::model::GpuTensor;
12094 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12095 match w {
12096 GpuTensor::Quant {
12097 qtype, row_bytes, ..
12098 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12099 _ => None,
12100 }
12101 };
12102 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12103 return Ok(None);
12104 };
12105 if w0.in_features() != w1.in_features() {
12106 return Ok(None);
12107 }
12108 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12110 match w {
12111 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12112 Some(m) => (m, true),
12113 None => (bytes, *rp),
12114 },
12115 _ => unreachable!(),
12116 }
12117 }
12118 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12119 if rp0 != rp1 {
12120 return Ok(None);
12121 }
12122 let rp = rp0;
12123 let rpb: u32 = 4;
12124 let mr1 = rp && Self::q40_mr1_on();
12126 let nb = |o: usize| {
12127 if mr1 {
12128 (o as u32).div_ceil(rpb)
12129 } else {
12130 (o as u32).div_ceil(2).div_ceil(rpb)
12131 }
12132 };
12133 let grid = nb(o0) + nb(o1);
12134 let mut y0 = self.alloc_uninit::<f32>(o0)?;
12135 let mut y1 = self.alloc_uninit::<f32>(o1)?;
12136 let f = self.func(if mr1 {
12137 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12138 } else if rp {
12139 "qmatvec_q4_0_mmvq_fused2_rp"
12140 } else {
12141 "qmatvec_q4_0_mmvq_fused2"
12142 });
12143 let cfg = LaunchConfig {
12144 grid_dim: (grid, 1, 1),
12145 block_dim: (32, rpb, 1),
12146 shared_mem_bytes: 0,
12147 };
12148 let inf = w0.in_features() as i32;
12149 let (oo0, oo1) = (o0 as i32, o1 as i32);
12150 let (r0, r1) = (rb0 as i64, rb1 as i64);
12151 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12153 {
12154 use cudarc::driver::{DevicePtr, DevicePtrMut};
12155 let s = &self.gpu.stream();
12156 let (p0, _g0) = b0.device_ptr(s);
12157 let (p1, _g1) = b1.device_ptr(s);
12158 let (paq, _g2) = aq.device_ptr(s);
12159 let (pad, _g3) = ad.device_ptr(s);
12160 let (py0, _g4) = y0.device_ptr_mut(s);
12161 let (py1, _g5) = y1.device_ptr_mut(s);
12162 let mut ps = [
12163 &p0 as *const _ as *mut std::ffi::c_void,
12164 &p1 as *const _ as *mut _,
12165 &paq as *const _ as *mut _,
12166 &pad as *const _ as *mut _,
12167 &py0 as *const _ as *mut _,
12168 &py1 as *const _ as *mut _,
12169 &inf as *const _ as *mut _,
12170 &oo0 as *const _ as *mut _,
12171 &oo1 as *const _ as *mut _,
12172 &r0 as *const _ as *mut _,
12173 &r1 as *const _ as *mut _,
12174 ];
12175 unsafe {
12176 self.launch_pdl(
12177 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12178 (grid, 1, 1),
12179 (32, rpb, 1),
12180 &mut ps,
12181 )?;
12182 }
12183 }
12184 return Ok(Some((y0, y1)));
12185 }
12186 let __s_b = self.gpu.stream();
12187 let mut b = __s_b.launch_builder(&f);
12188 b.arg(b0)
12189 .arg(b1)
12190 .arg(aq)
12191 .arg(ad)
12192 .arg(&mut y0)
12193 .arg(&mut y1)
12194 .arg(&inf)
12195 .arg(&oo0)
12196 .arg(&oo1)
12197 .arg(&r0)
12198 .arg(&r1);
12199 unsafe {
12200 b.launch(cfg)?;
12201 }
12202 Ok(Some((y0, y1)))
12203 }
12204
12205 pub fn matmul_q4_fused2_into(
12207 &self,
12208 w0: &crate::model::GpuTensor,
12209 w1: &crate::model::GpuTensor,
12210 aq: &CudaSlice<i8>,
12211 ad: &CudaSlice<f32>,
12212 y0: &mut CudaSlice<f32>,
12213 y1: &mut CudaSlice<f32>,
12214 ) -> Result<bool, Box<dyn std::error::Error>> {
12215 use crate::model::GpuTensor;
12216 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12217 match w {
12218 GpuTensor::Quant {
12219 qtype, row_bytes, ..
12220 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12221 _ => None,
12222 }
12223 };
12224 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
12225 return Ok(false);
12226 };
12227 if w0.in_features() != w1.in_features() {
12228 return Ok(false);
12229 }
12230 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12231 match w {
12232 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12233 Some(m) => (m, true),
12234 None => (bytes, *rp),
12235 },
12236 _ => unreachable!(),
12237 }
12238 }
12239 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12240 if rp0 != rp1 {
12241 return Ok(false);
12242 }
12243 let rp = rp0;
12244 let rpb: u32 = 4;
12245 let mr1 = rp && Self::q40_mr1_on();
12246 let nb = |o: usize| {
12247 if mr1 {
12248 (o as u32).div_ceil(rpb)
12249 } else {
12250 (o as u32).div_ceil(2).div_ceil(rpb)
12251 }
12252 };
12253 let grid = nb(o0) + nb(o1);
12254 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
12255 let f = self.func(if mr1 {
12256 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
12257 } else if rp {
12258 "qmatvec_q4_0_mmvq_fused2_rp"
12259 } else {
12260 "qmatvec_q4_0_mmvq_fused2"
12261 });
12262 let cfg = LaunchConfig {
12263 grid_dim: (grid, 1, 1),
12264 block_dim: (32, rpb, 1),
12265 shared_mem_bytes: 0,
12266 };
12267 let inf = w0.in_features() as i32;
12268 let (oo0, oo1) = (o0 as i32, o1 as i32);
12269 let (r0, r1) = (rb0 as i64, rb1 as i64);
12270 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
12272 use cudarc::driver::{DevicePtr, DevicePtrMut};
12273 let s = &self.gpu.stream();
12274 let (p0, _g0) = b0.device_ptr(s);
12275 let (p1, _g1) = b1.device_ptr(s);
12276 let (paq, _g2) = aq.device_ptr(s);
12277 let (pad, _g3) = ad.device_ptr(s);
12278 let (py0, _g4) = y0.device_ptr_mut(s);
12279 let (py1, _g5) = y1.device_ptr_mut(s);
12280 let mut ps = [
12281 &p0 as *const _ as *mut std::ffi::c_void,
12282 &p1 as *const _ as *mut _,
12283 &paq as *const _ as *mut _,
12284 &pad as *const _ as *mut _,
12285 &py0 as *const _ as *mut _,
12286 &py1 as *const _ as *mut _,
12287 &inf as *const _ as *mut _,
12288 &oo0 as *const _ as *mut _,
12289 &oo1 as *const _ as *mut _,
12290 &r0 as *const _ as *mut _,
12291 &r1 as *const _ as *mut _,
12292 ];
12293 unsafe {
12294 self.launch_pdl(
12295 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
12296 (grid, 1, 1),
12297 (32, rpb, 1),
12298 &mut ps,
12299 )?;
12300 }
12301 return Ok(true);
12302 }
12303 let __s_b = self.gpu.stream();
12304 let mut b = __s_b.launch_builder(&f);
12305 b.arg(b0)
12306 .arg(b1)
12307 .arg(aq)
12308 .arg(ad)
12309 .arg(&mut *y0)
12310 .arg(&mut *y1)
12311 .arg(&inf)
12312 .arg(&oo0)
12313 .arg(&oo1)
12314 .arg(&r0)
12315 .arg(&r1);
12316 unsafe {
12317 b.launch(cfg)?;
12318 }
12319 Ok(true)
12320 }
12321
12322 pub fn matmul_q4_fused2_batched(
12327 &self,
12328 w0: &crate::model::GpuTensor,
12329 w1: &crate::model::GpuTensor,
12330 aq: &CudaSlice<i8>,
12331 ad: &CudaSlice<f32>,
12332 m: usize,
12333 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12334 use crate::model::GpuTensor;
12335 if m < 2 || m > 8 {
12336 return Ok(None);
12337 }
12338 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
12339 match w {
12340 GpuTensor::Quant {
12341 qtype, row_bytes, ..
12342 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
12343 _ => None,
12344 }
12345 };
12346 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
12347 return Ok(None);
12348 };
12349 if w0.in_features() != w1.in_features() {
12350 return Ok(None);
12351 }
12352 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12353 match w {
12354 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12355 Some(mr) => (mr, true),
12356 None => (bytes, *rp),
12357 },
12358 _ => unreachable!(),
12359 }
12360 }
12361 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
12362 if !rp0 || !rp1 {
12363 return Ok(None);
12364 }
12365 let mcols = Self::batched_mcols(m);
12366 let rpb: u32 = 4;
12367 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12368 let grid = nb(o0) + nb(o1);
12369 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12370 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12371 let f = self.func(match mcols {
12372 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
12373 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
12374 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
12375 });
12376 let cfg = LaunchConfig {
12377 grid_dim: (grid, 1, 1),
12378 block_dim: (32, rpb, 1),
12379 shared_mem_bytes: 0,
12380 };
12381 let inf = w0.in_features() as i32;
12382 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
12383 let rb = rb0 as i64;
12384 let __s_b = self.gpu.stream();
12385 let mut b = __s_b.launch_builder(&f);
12386 b.arg(b0)
12387 .arg(b1)
12388 .arg(aq)
12389 .arg(ad)
12390 .arg(&mut y0)
12391 .arg(&mut y1)
12392 .arg(&inf)
12393 .arg(&oo0)
12394 .arg(&oo1)
12395 .arg(&mi)
12396 .arg(&rb);
12397 unsafe {
12398 b.launch(cfg)?;
12399 }
12400 Ok(Some((y0, y1)))
12401 }
12402
12403 #[allow(clippy::too_many_arguments)]
12406 pub fn matmul_q4_fused3_batched(
12407 &self,
12408 w0: &crate::model::GpuTensor,
12409 w1: &crate::model::GpuTensor,
12410 w2: &crate::model::GpuTensor,
12411 aq: &CudaSlice<i8>,
12412 ad: &CudaSlice<f32>,
12413 m: usize,
12414 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12415 {
12416 use crate::model::GpuTensor;
12417 if m < 2 || m > 8 {
12418 return Ok(None);
12419 }
12420 let q4 = |w: &GpuTensor| -> Option<usize> {
12421 match w {
12422 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
12423 _ => None,
12424 }
12425 };
12426 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
12427 return Ok(None);
12428 };
12429 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
12430 return Ok(None);
12431 }
12432 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
12433 match w {
12434 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
12435 Some(mr) => (mr, true),
12436 None => (bytes, *rp),
12437 },
12438 _ => unreachable!(),
12439 }
12440 }
12441 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
12442 if !rp0 || !rp1 || !rp2 {
12443 return Ok(None);
12444 }
12445 let mcols = Self::batched_mcols(m);
12446 let rpb: u32 = 4;
12447 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
12448 let grid = nb(o0) + nb(o1) + nb(o2);
12449 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12450 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12451 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12452 let f = self.func(match mcols {
12453 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
12454 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
12455 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
12456 });
12457 let cfg = LaunchConfig {
12458 grid_dim: (grid, 1, 1),
12459 block_dim: (32, rpb, 1),
12460 shared_mem_bytes: 0,
12461 };
12462 let inf = w0.in_features() as i32;
12463 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
12464 let rb = 0i64;
12465 let __s_b = self.gpu.stream();
12466 let mut b = __s_b.launch_builder(&f);
12467 b.arg(b0)
12468 .arg(b1)
12469 .arg(b2)
12470 .arg(aq)
12471 .arg(ad)
12472 .arg(&mut y0)
12473 .arg(&mut y1)
12474 .arg(&mut y2)
12475 .arg(&inf)
12476 .arg(&oo0)
12477 .arg(&oo1)
12478 .arg(&oo2)
12479 .arg(&mi)
12480 .arg(&rb);
12481 unsafe {
12482 b.launch(cfg)?;
12483 }
12484 Ok(Some((y0, y1, y2)))
12485 }
12486
12487 pub fn matmul_q8_fused3(
12488 &self,
12489 w0: &crate::model::GpuTensor,
12490 w1: &crate::model::GpuTensor,
12491 w2: &crate::model::GpuTensor,
12492 aq: &CudaSlice<i8>,
12493 ad: &CudaSlice<f32>,
12494 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12495 {
12496 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12499 return Ok(Some(self.e4m3_fused3_core(
12500 p0.0,
12501 p1.0,
12502 p2.0,
12503 aq,
12504 ad,
12505 w0.in_features(),
12506 p0.1,
12507 p1.1,
12508 p2.1,
12509 p0.2,
12510 p0.3,
12511 p1.3,
12512 p2.3,
12513 )?));
12514 }
12515 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12516 return Ok(None);
12517 };
12518 Ok(Some(self.q8_fused3_core(
12519 p0.0,
12520 p1.0,
12521 p2.0,
12522 aq,
12523 ad,
12524 w0.in_features(),
12525 p0.1,
12526 p1.1,
12527 p2.1,
12528 p0.2,
12529 )?))
12530 }
12531
12532 #[allow(clippy::too_many_arguments)]
12533 fn q8_fused3_core(
12534 &self,
12535 b0: &CudaSlice<u8>,
12536 b1: &CudaSlice<u8>,
12537 b2: &CudaSlice<u8>,
12538 aq: &CudaSlice<i8>,
12539 ad: &CudaSlice<f32>,
12540 in_f: usize,
12541 out0: usize,
12542 out1: usize,
12543 out2: usize,
12544 row_bytes: usize,
12545 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12546 const ROWS_PER_BLOCK: u32 = 4;
12547 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12548 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12549 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12550 let f = self.func("qmatvec_q8_0_mmvq_fused3");
12551 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12552 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12553 let mut y2 = self.alloc_uninit::<f32>(out2)?;
12554 let cfg = LaunchConfig {
12555 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12556 block_dim: (32, ROWS_PER_BLOCK, 1),
12557 shared_mem_bytes: 0,
12558 };
12559 let (inf, o0, o1, o2, rbl) = (
12560 in_f as i32,
12561 out0 as i32,
12562 out1 as i32,
12563 out2 as i32,
12564 row_bytes as i64,
12565 );
12566 let __s_b = self.gpu.stream();
12567 let mut b = __s_b.launch_builder(&f);
12568 b.arg(b0)
12569 .arg(b1)
12570 .arg(b2)
12571 .arg(aq)
12572 .arg(ad)
12573 .arg(&mut y0)
12574 .arg(&mut y1)
12575 .arg(&mut y2)
12576 .arg(&inf)
12577 .arg(&o0)
12578 .arg(&o1)
12579 .arg(&o2)
12580 .arg(&rbl);
12581 unsafe {
12582 b.launch(cfg)?;
12583 }
12584 Ok((y0, y1, y2))
12585 }
12586
12587 #[allow(clippy::too_many_arguments)]
12589 pub fn qmatvec_q8_fused3_raw(
12590 &self,
12591 b0: &CudaSlice<u8>,
12592 b1: &CudaSlice<u8>,
12593 b2: &CudaSlice<u8>,
12594 x: &CudaSlice<f32>,
12595 in_f: usize,
12596 out0: usize,
12597 out1: usize,
12598 out2: usize,
12599 row_bytes: usize,
12600 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12601 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
12602 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
12603 }
12604
12605 pub fn matmul_q8_fused2_t(
12616 &self,
12617 w0: &crate::model::GpuTensor,
12618 w1: &crate::model::GpuTensor,
12619 aq: &CudaSlice<i8>,
12620 ad: &CudaSlice<f32>,
12621 m: usize,
12622 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12623 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12627 return Ok(None);
12628 }
12629 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12632 if m > 4 && !Self::b8_enabled() {
12633 return Ok(None);
12634 }
12635 return Ok(Some(self.e4m3_fused2_t_core(
12636 p0.0,
12637 p1.0,
12638 aq,
12639 ad,
12640 m,
12641 w0.in_features(),
12642 p0.1,
12643 p1.1,
12644 p0.2,
12645 p0.3,
12646 p1.3,
12647 )?));
12648 }
12649 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12650 return Ok(None);
12651 };
12652 Ok(Some(self.q8_fused2_t_core(
12653 p0.0,
12654 p1.0,
12655 aq,
12656 ad,
12657 m,
12658 w0.in_features(),
12659 p0.1,
12660 p1.1,
12661 p0.2,
12662 )?))
12663 }
12664
12665 #[allow(clippy::too_many_arguments)]
12666 fn q8_fused2_t_core(
12667 &self,
12668 b0: &CudaSlice<u8>,
12669 b1: &CudaSlice<u8>,
12670 aq: &CudaSlice<i8>,
12671 ad: &CudaSlice<f32>,
12672 m: usize,
12673 in_f: usize,
12674 out0: usize,
12675 out1: usize,
12676 row_bytes: usize,
12677 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12678 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12680 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12681 let f = self.func(match Self::batched_mcols(m) {
12682 2 => "qmatvec_q8_0_mmvq_fused2_b2",
12683 4 => "qmatvec_q8_0_mmvq_fused2_b4",
12684 _ => "qmatvec_q8_0_mmvq_fused2_b8",
12686 });
12687 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12688 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12689 let cfg = LaunchConfig {
12690 grid_dim: (nb0 + nb1, 1, 1),
12691 block_dim: (32, ROWS_PER_BLOCK, 1),
12692 shared_mem_bytes: 0,
12693 };
12694 let (inf, o0, o1, mi, rbl) = (
12695 in_f as i32,
12696 out0 as i32,
12697 out1 as i32,
12698 m as i32,
12699 row_bytes as i64,
12700 );
12701 let __s_b = self.gpu.stream();
12702 let mut b = __s_b.launch_builder(&f);
12703 b.arg(b0)
12704 .arg(b1)
12705 .arg(aq)
12706 .arg(ad)
12707 .arg(&mut y0)
12708 .arg(&mut y1)
12709 .arg(&inf)
12710 .arg(&o0)
12711 .arg(&o1)
12712 .arg(&mi)
12713 .arg(&rbl);
12714 unsafe {
12715 b.launch(cfg)?;
12716 }
12717 Ok((y0, y1))
12718 }
12719
12720 #[allow(clippy::too_many_arguments)]
12723 pub fn qmatvec_q8_fused2_t_raw(
12724 &self,
12725 b0: &CudaSlice<u8>,
12726 b1: &CudaSlice<u8>,
12727 x: &CudaSlice<f32>,
12728 m: usize,
12729 in_f: usize,
12730 out0: usize,
12731 out1: usize,
12732 row_bytes: usize,
12733 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12734 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12735 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
12736 }
12737
12738 #[allow(clippy::too_many_arguments)]
12741 pub fn matmul_q8_fused3_t(
12742 &self,
12743 w0: &crate::model::GpuTensor,
12744 w1: &crate::model::GpuTensor,
12745 w2: &crate::model::GpuTensor,
12746 aq: &CudaSlice<i8>,
12747 ad: &CudaSlice<f32>,
12748 m: usize,
12749 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12750 {
12751 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
12752 return Ok(None);
12753 }
12754 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
12755 return Ok(Some(self.e4m3_fused3_t_core(
12756 p0.0,
12757 p1.0,
12758 p2.0,
12759 aq,
12760 ad,
12761 m,
12762 w0.in_features(),
12763 p0.1,
12764 p1.1,
12765 p2.1,
12766 p0.2,
12767 p0.3,
12768 p1.3,
12769 p2.3,
12770 )?));
12771 }
12772 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
12773 return Ok(None);
12774 };
12775 Ok(Some(self.q8_fused3_t_core(
12776 p0.0,
12777 p1.0,
12778 p2.0,
12779 aq,
12780 ad,
12781 m,
12782 w0.in_features(),
12783 p0.1,
12784 p1.1,
12785 p2.1,
12786 p0.2,
12787 )?))
12788 }
12789
12790 #[allow(clippy::too_many_arguments)]
12791 fn q8_fused3_t_core(
12792 &self,
12793 b0: &CudaSlice<u8>,
12794 b1: &CudaSlice<u8>,
12795 b2: &CudaSlice<u8>,
12796 aq: &CudaSlice<i8>,
12797 ad: &CudaSlice<f32>,
12798 m: usize,
12799 in_f: usize,
12800 out0: usize,
12801 out1: usize,
12802 out2: usize,
12803 row_bytes: usize,
12804 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12805 const ROWS_PER_BLOCK: u32 = 4;
12806 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12807 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12808 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
12809 let f = self.func(if Self::batched_mcols(m) == 2 {
12810 "qmatvec_q8_0_mmvq_fused3_b2"
12811 } else {
12812 "qmatvec_q8_0_mmvq_fused3_b4"
12813 });
12814 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
12815 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
12816 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
12817 let cfg = LaunchConfig {
12818 grid_dim: (nb0 + nb1 + nb2, 1, 1),
12819 block_dim: (32, ROWS_PER_BLOCK, 1),
12820 shared_mem_bytes: 0,
12821 };
12822 let (inf, o0, o1, o2, mi, rbl) = (
12823 in_f as i32,
12824 out0 as i32,
12825 out1 as i32,
12826 out2 as i32,
12827 m as i32,
12828 row_bytes as i64,
12829 );
12830 let __s_b = self.gpu.stream();
12831 let mut b = __s_b.launch_builder(&f);
12832 b.arg(b0)
12833 .arg(b1)
12834 .arg(b2)
12835 .arg(aq)
12836 .arg(ad)
12837 .arg(&mut y0)
12838 .arg(&mut y1)
12839 .arg(&mut y2)
12840 .arg(&inf)
12841 .arg(&o0)
12842 .arg(&o1)
12843 .arg(&o2)
12844 .arg(&mi)
12845 .arg(&rbl);
12846 unsafe {
12847 b.launch(cfg)?;
12848 }
12849 Ok((y0, y1, y2))
12850 }
12851
12852 #[allow(clippy::too_many_arguments)]
12854 pub fn qmatvec_q8_fused3_t_raw(
12855 &self,
12856 b0: &CudaSlice<u8>,
12857 b1: &CudaSlice<u8>,
12858 b2: &CudaSlice<u8>,
12859 x: &CudaSlice<f32>,
12860 m: usize,
12861 in_f: usize,
12862 out0: usize,
12863 out1: usize,
12864 out2: usize,
12865 row_bytes: usize,
12866 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12867 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12868 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
12869 }
12870
12871 pub fn q8_ffn_fuse2_on(&self) -> bool {
12875 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12876 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
12877 }
12878
12879 #[allow(clippy::type_complexity)]
12885 fn q8_fused_params<'w, const N: usize>(
12886 &self,
12887 ws: &[&'w crate::model::GpuTensor; N],
12888 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
12889 use crate::model::GpuTensor;
12890 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
12891 return None;
12892 }
12893 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
12894 return None;
12895 }
12896 let in_f = ws[0].in_features();
12897 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
12898 for (i, w) in ws.iter().enumerate() {
12899 match w {
12900 GpuTensor::Quant {
12901 bytes,
12902 qtype,
12903 row_bytes,
12904 scale,
12905 ..
12906 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
12907 out[i] = Some((bytes, w.out_features(), *row_bytes))
12908 }
12909 _ => return None,
12910 }
12911 }
12912 Some(out.map(|o| o.unwrap()))
12913 }
12914
12915 pub fn e4m3_dual_on(&self) -> bool {
12918 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12919 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
12920 }
12921
12922 #[allow(clippy::type_complexity)]
12934 fn e4m3_fused_params<'w, const N: usize>(
12935 &self,
12936 ws: &[&'w crate::model::GpuTensor; N],
12937 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
12938 use crate::model::GpuTensor;
12939 if !self.e4m3_dual_on() {
12940 return None;
12941 }
12942 let in_f = ws[0].in_features();
12943 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
12944 for (i, w) in ws.iter().enumerate() {
12945 match w {
12946 GpuTensor::Quant {
12947 bytes,
12948 qtype,
12949 row_bytes,
12950 scale,
12951 rp,
12952 rp4,
12953 ..
12954 } if *qtype == QT_F8_E4M3
12955 && w.in_features() == in_f
12956 && *row_bytes == in_f
12957 && !*rp
12958 && rp4.is_none() =>
12959 {
12960 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
12961 }
12962 _ => return None,
12963 }
12964 }
12965 Some(out.map(|o| o.unwrap()))
12966 }
12967
12968 #[allow(clippy::too_many_arguments)]
12972 fn e4m3_fused2_core(
12973 &self,
12974 b0: &CudaSlice<u8>,
12975 b1: &CudaSlice<u8>,
12976 aq: &CudaSlice<i8>,
12977 ad: &CudaSlice<f32>,
12978 in_f: usize,
12979 out0: usize,
12980 out1: usize,
12981 row_bytes: usize,
12982 ws0: f32,
12983 ws1: f32,
12984 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12985 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
12987 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
12988 let f = self.func("qmatvec_e4m3_mmvq_fused2");
12989 let mut y0 = self.alloc_uninit::<f32>(out0)?;
12990 let mut y1 = self.alloc_uninit::<f32>(out1)?;
12991 let cfg = LaunchConfig {
12992 grid_dim: (nb0 + nb1, 1, 1),
12993 block_dim: (32, ROWS_PER_BLOCK, 1),
12994 shared_mem_bytes: 0,
12995 };
12996 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
12997 let __s_b = self.gpu.stream();
12998 let mut b = __s_b.launch_builder(&f);
12999 b.arg(b0)
13000 .arg(b1)
13001 .arg(aq)
13002 .arg(ad)
13003 .arg(&mut y0)
13004 .arg(&mut y1)
13005 .arg(&inf)
13006 .arg(&o0)
13007 .arg(&o1)
13008 .arg(&rbl)
13009 .arg(&ws0)
13010 .arg(&ws1);
13011 unsafe {
13012 b.launch(cfg)?;
13013 }
13014 Ok((y0, y1))
13015 }
13016
13017 #[allow(clippy::too_many_arguments)]
13019 fn e4m3_fused3_core(
13020 &self,
13021 b0: &CudaSlice<u8>,
13022 b1: &CudaSlice<u8>,
13023 b2: &CudaSlice<u8>,
13024 aq: &CudaSlice<i8>,
13025 ad: &CudaSlice<f32>,
13026 in_f: usize,
13027 out0: usize,
13028 out1: usize,
13029 out2: usize,
13030 row_bytes: usize,
13031 ws0: f32,
13032 ws1: f32,
13033 ws2: f32,
13034 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13035 const ROWS_PER_BLOCK: u32 = 4;
13036 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13037 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13038 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13039 let f = self.func("qmatvec_e4m3_mmvq_fused3");
13040 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13041 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13042 let mut y2 = self.alloc_uninit::<f32>(out2)?;
13043 let cfg = LaunchConfig {
13044 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13045 block_dim: (32, ROWS_PER_BLOCK, 1),
13046 shared_mem_bytes: 0,
13047 };
13048 let (inf, o0, o1, o2, rbl) = (
13049 in_f as i32,
13050 out0 as i32,
13051 out1 as i32,
13052 out2 as i32,
13053 row_bytes as i64,
13054 );
13055 let __s_b = self.gpu.stream();
13056 let mut b = __s_b.launch_builder(&f);
13057 b.arg(b0)
13058 .arg(b1)
13059 .arg(b2)
13060 .arg(aq)
13061 .arg(ad)
13062 .arg(&mut y0)
13063 .arg(&mut y1)
13064 .arg(&mut y2)
13065 .arg(&inf)
13066 .arg(&o0)
13067 .arg(&o1)
13068 .arg(&o2)
13069 .arg(&rbl)
13070 .arg(&ws0)
13071 .arg(&ws1)
13072 .arg(&ws2);
13073 unsafe {
13074 b.launch(cfg)?;
13075 }
13076 Ok((y0, y1, y2))
13077 }
13078
13079 #[allow(clippy::too_many_arguments)]
13083 fn e4m3_fused2_t_core(
13084 &self,
13085 b0: &CudaSlice<u8>,
13086 b1: &CudaSlice<u8>,
13087 aq: &CudaSlice<i8>,
13088 ad: &CudaSlice<f32>,
13089 m: usize,
13090 in_f: usize,
13091 out0: usize,
13092 out1: usize,
13093 row_bytes: usize,
13094 ws0: f32,
13095 ws1: f32,
13096 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13097 const ROWS_PER_BLOCK: u32 = 4;
13098 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13099 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13100 let f = self.func(match Self::batched_mcols(m) {
13101 2 => "qmatvec_e4m3_mmvq_fused2_b2",
13102 4 => "qmatvec_e4m3_mmvq_fused2_b4",
13103 _ => "qmatvec_e4m3_mmvq_fused2_b8",
13104 });
13105 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13106 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13107 let cfg = LaunchConfig {
13108 grid_dim: (nb0 + nb1, 1, 1),
13109 block_dim: (32, ROWS_PER_BLOCK, 1),
13110 shared_mem_bytes: 0,
13111 };
13112 let (inf, o0, o1, mi, rbl) = (
13113 in_f as i32,
13114 out0 as i32,
13115 out1 as i32,
13116 m as i32,
13117 row_bytes as i64,
13118 );
13119 let __s_b = self.gpu.stream();
13120 let mut b = __s_b.launch_builder(&f);
13121 b.arg(b0)
13122 .arg(b1)
13123 .arg(aq)
13124 .arg(ad)
13125 .arg(&mut y0)
13126 .arg(&mut y1)
13127 .arg(&inf)
13128 .arg(&o0)
13129 .arg(&o1)
13130 .arg(&mi)
13131 .arg(&rbl);
13132 unsafe {
13133 b.launch(cfg)?;
13134 }
13135 if ws0 != 1.0 {
13136 self.scale_inplace(&mut y0, ws0, m * out0)?;
13137 }
13138 if ws1 != 1.0 {
13139 self.scale_inplace(&mut y1, ws1, m * out1)?;
13140 }
13141 Ok((y0, y1))
13142 }
13143
13144 #[allow(clippy::too_many_arguments)]
13146 fn e4m3_fused3_t_core(
13147 &self,
13148 b0: &CudaSlice<u8>,
13149 b1: &CudaSlice<u8>,
13150 b2: &CudaSlice<u8>,
13151 aq: &CudaSlice<i8>,
13152 ad: &CudaSlice<f32>,
13153 m: usize,
13154 in_f: usize,
13155 out0: usize,
13156 out1: usize,
13157 out2: usize,
13158 row_bytes: usize,
13159 ws0: f32,
13160 ws1: f32,
13161 ws2: f32,
13162 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13163 const ROWS_PER_BLOCK: u32 = 4;
13164 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13165 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13166 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13167 let f = self.func(if Self::batched_mcols(m) == 2 {
13168 "qmatvec_e4m3_mmvq_fused3_b2"
13169 } else {
13170 "qmatvec_e4m3_mmvq_fused3_b4"
13171 });
13172 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13173 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13174 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
13175 let cfg = LaunchConfig {
13176 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13177 block_dim: (32, ROWS_PER_BLOCK, 1),
13178 shared_mem_bytes: 0,
13179 };
13180 let (inf, o0, o1, o2, mi, rbl) = (
13181 in_f as i32,
13182 out0 as i32,
13183 out1 as i32,
13184 out2 as i32,
13185 m as i32,
13186 row_bytes as i64,
13187 );
13188 let __s_b = self.gpu.stream();
13189 let mut b = __s_b.launch_builder(&f);
13190 b.arg(b0)
13191 .arg(b1)
13192 .arg(b2)
13193 .arg(aq)
13194 .arg(ad)
13195 .arg(&mut y0)
13196 .arg(&mut y1)
13197 .arg(&mut y2)
13198 .arg(&inf)
13199 .arg(&o0)
13200 .arg(&o1)
13201 .arg(&o2)
13202 .arg(&mi)
13203 .arg(&rbl);
13204 unsafe {
13205 b.launch(cfg)?;
13206 }
13207 if ws0 != 1.0 {
13208 self.scale_inplace(&mut y0, ws0, m * out0)?;
13209 }
13210 if ws1 != 1.0 {
13211 self.scale_inplace(&mut y1, ws1, m * out1)?;
13212 }
13213 if ws2 != 1.0 {
13214 self.scale_inplace(&mut y2, ws2, m * out2)?;
13215 }
13216 Ok((y0, y1, y2))
13217 }
13218
13219 pub fn qmatvec_e4m3_blk_mmvq(
13229 &self,
13230 bytes: &CudaSlice<u8>,
13231 aq: &CudaSlice<i8>,
13232 ad: &CudaSlice<f32>,
13233 scales: &CudaSlice<f32>,
13234 m: usize,
13235 in_f: usize,
13236 out_f: usize,
13237 row_bytes: usize,
13238 scale_cols: usize,
13239 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13240 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
13242 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
13243 )?;
13244 Ok(y)
13245 }
13246
13247 #[allow(clippy::too_many_arguments)]
13249 pub fn qmatvec_e4m3_blk_mmvq_into(
13250 &self,
13251 bytes: &CudaSlice<u8>,
13252 aq: &CudaSlice<i8>,
13253 ad: &CudaSlice<f32>,
13254 scales: &CudaSlice<f32>,
13255 m: usize,
13256 in_f: usize,
13257 out_f: usize,
13258 row_bytes: usize,
13259 scale_cols: usize,
13260 y: &mut CudaSlice<f32>,
13261 ) -> Result<(), Box<dyn std::error::Error>> {
13262 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
13264 let cfg = LaunchConfig {
13265 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
13266 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
13269 let (inf, outf, mi, rb, sc) = (
13270 in_f as i32,
13271 out_f as i32,
13272 m as i32,
13273 row_bytes as i64,
13274 scale_cols as i32,
13275 );
13276 let __s_b = self.gpu.stream();
13277 let mut b = __s_b.launch_builder(&f);
13278 b.arg(bytes)
13279 .arg(aq)
13280 .arg(ad)
13281 .arg(scales)
13282 .arg(&mut *y)
13283 .arg(&inf)
13284 .arg(&outf)
13285 .arg(&mi)
13286 .arg(&rb)
13287 .arg(&sc);
13288 unsafe {
13289 b.launch(cfg)?;
13290 }
13291 Ok(())
13292 }
13293
13294 #[allow(clippy::too_many_arguments)]
13300 pub fn qmatvec_e4m3_blk_mmvq_batched(
13301 &self,
13302 bytes: &CudaSlice<u8>,
13303 aq: &CudaSlice<i8>,
13304 ad: &CudaSlice<f32>,
13305 scales: &CudaSlice<f32>,
13306 m: usize,
13307 in_f: usize,
13308 out_f: usize,
13309 row_bytes: usize,
13310 scale_cols: usize,
13311 mcols: usize,
13312 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13313 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
13315 let name = match mcols {
13316 2 => "qmatvec_e4m3_blk_mmvq_b2",
13317 4 => "qmatvec_e4m3_blk_mmvq_b4",
13318 8 => "qmatvec_e4m3_blk_mmvq_b8",
13319 16 => "qmatvec_e4m3_blk_mmvq_b16",
13320 _ => {
13321 return Err(
13322 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
13323 );
13324 }
13325 };
13326 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13327 let f = self.func(name);
13328 let cfg = LaunchConfig {
13329 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
13330 block_dim: (32, ROWS_PER_BLOCK, 1),
13331 shared_mem_bytes: 0,
13332 };
13333 let (inf, outf, mi, rb, sc) = (
13334 in_f as i32,
13335 out_f as i32,
13336 m as i32,
13337 row_bytes as i64,
13338 scale_cols as i32,
13339 );
13340 let __s_b = self.gpu.stream();
13341 let mut b = __s_b.launch_builder(&f);
13342 b.arg(bytes)
13343 .arg(aq)
13344 .arg(ad)
13345 .arg(scales)
13346 .arg(&mut y)
13347 .arg(&inf)
13348 .arg(&outf)
13349 .arg(&mi)
13350 .arg(&rb)
13351 .arg(&sc);
13352 unsafe {
13353 b.launch(cfg)?;
13354 }
13355 Ok(y)
13356 }
13357
13358 #[allow(clippy::too_many_arguments)]
13361 pub fn qmatvec_e4m3_blk_batched_raw(
13362 &self,
13363 bytes: &CudaSlice<u8>,
13364 x: &CudaSlice<f32>,
13365 scales: &CudaSlice<f32>,
13366 m: usize,
13367 in_f: usize,
13368 out_f: usize,
13369 row_bytes: usize,
13370 scale_cols: usize,
13371 mcols: usize,
13372 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13373 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13374 self.qmatvec_e4m3_blk_mmvq_batched(
13375 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
13376 )
13377 }
13378
13379 #[allow(clippy::too_many_arguments)]
13382 pub fn qmatvec_e4m3_blk_mmvq_raw(
13383 &self,
13384 bytes: &CudaSlice<u8>,
13385 x: &CudaSlice<f32>,
13386 scales: &CudaSlice<f32>,
13387 m: usize,
13388 in_f: usize,
13389 out_f: usize,
13390 row_bytes: usize,
13391 scale_cols: usize,
13392 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13393 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13394 self.qmatvec_e4m3_blk_mmvq(
13395 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
13396 )
13397 }
13398
13399 #[allow(clippy::too_many_arguments)]
13402 pub fn qmatvec_e4m3_fused2_raw(
13403 &self,
13404 b0: &CudaSlice<u8>,
13405 b1: &CudaSlice<u8>,
13406 x: &CudaSlice<f32>,
13407 in_f: usize,
13408 out0: usize,
13409 out1: usize,
13410 row_bytes: usize,
13411 ws0: f32,
13412 ws1: f32,
13413 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13414 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13415 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
13416 }
13417
13418 #[allow(clippy::too_many_arguments)]
13419 pub fn qmatvec_e4m3_fused3_raw(
13420 &self,
13421 b0: &CudaSlice<u8>,
13422 b1: &CudaSlice<u8>,
13423 b2: &CudaSlice<u8>,
13424 x: &CudaSlice<f32>,
13425 in_f: usize,
13426 out0: usize,
13427 out1: usize,
13428 out2: usize,
13429 row_bytes: usize,
13430 ws0: f32,
13431 ws1: f32,
13432 ws2: f32,
13433 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13434 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13435 self.e4m3_fused3_core(
13436 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13437 )
13438 }
13439
13440 #[allow(clippy::too_many_arguments)]
13441 pub fn qmatvec_e4m3_fused2_t_raw(
13442 &self,
13443 b0: &CudaSlice<u8>,
13444 b1: &CudaSlice<u8>,
13445 x: &CudaSlice<f32>,
13446 m: usize,
13447 in_f: usize,
13448 out0: usize,
13449 out1: usize,
13450 row_bytes: usize,
13451 ws0: f32,
13452 ws1: f32,
13453 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13454 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13455 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
13456 }
13457
13458 #[allow(clippy::too_many_arguments)]
13459 pub fn qmatvec_e4m3_fused3_t_raw(
13460 &self,
13461 b0: &CudaSlice<u8>,
13462 b1: &CudaSlice<u8>,
13463 b2: &CudaSlice<u8>,
13464 x: &CudaSlice<f32>,
13465 m: usize,
13466 in_f: usize,
13467 out0: usize,
13468 out1: usize,
13469 out2: usize,
13470 row_bytes: usize,
13471 ws0: f32,
13472 ws1: f32,
13473 ws2: f32,
13474 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13475 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13476 self.e4m3_fused3_t_core(
13477 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
13478 )
13479 }
13480
13481 fn try_e4m3_blk_pre(
13492 &self,
13493 w: &crate::model::GpuTensor,
13494 aq: &CudaSlice<i8>,
13495 ad: &CudaSlice<f32>,
13496 m: usize,
13497 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13498 use crate::model::GpuTensor;
13499 if let GpuTensor::Quant {
13500 bytes,
13501 qtype,
13502 row_bytes,
13503 blk: Some(g),
13504 ..
13505 } = w
13506 {
13507 if *qtype == QT_F8_E4M3_BLK {
13508 if (2..=16).contains(&m)
13514 && std::env::var("MEMRA_NO_BATCHED").is_err()
13515 && (m <= 4 || Self::b8_enabled())
13516 {
13517 let mcols = Self::batched_mcols(m);
13518 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
13519 bytes,
13520 aq,
13521 ad,
13522 &g.scales,
13523 m,
13524 w.in_features(),
13525 w.out_features(),
13526 *row_bytes,
13527 g.cols,
13528 mcols,
13529 )?));
13530 }
13531 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
13532 bytes,
13533 aq,
13534 ad,
13535 &g.scales,
13536 m,
13537 w.in_features(),
13538 w.out_features(),
13539 *row_bytes,
13540 g.cols,
13541 )?));
13542 }
13543 }
13544 Ok(None)
13545 }
13546
13547 fn try_e4m3_blk_prefill(
13594 &self,
13595 w: &crate::model::GpuTensor,
13596 x: &CudaSlice<f32>,
13597 m: usize,
13598 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13599 use crate::model::GpuTensor;
13600 let GpuTensor::Quant {
13601 bytes,
13602 qtype,
13603 blk: Some(g),
13604 ..
13605 } = w
13606 else {
13607 return Ok(None);
13608 };
13609 if *qtype != QT_F8_E4M3_BLK {
13610 return Ok(None);
13611 }
13612 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
13617 return Ok(Some(y));
13618 }
13619 let (in_f, out_f) = (w.in_features(), w.out_features());
13620 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
13621 let tmp = GpuTensor::Quant {
13622 bytes: slab,
13623 qtype: QT_Q8_0,
13624 row_bytes: in_f / 32 * 34,
13625 ne: vec![in_f as u64, out_f as u64],
13626 scale: 1.0,
13627 rp: false,
13628 #[cfg(memra_cutlass)]
13629 cutlass: None,
13630 fp8: None,
13631 blk: None,
13632 f16: None,
13633 rp4: None,
13634 };
13635 Ok(Some(self.matmul(&tmp, x, m)?))
13637 }
13638
13639 pub fn matmul_pre_noscale(
13640 &self,
13641 w: &crate::model::GpuTensor,
13642 aq: &CudaSlice<i8>,
13643 ad: &CudaSlice<f32>,
13644 m: usize,
13645 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
13646 use crate::model::GpuTensor;
13647 if m == 1 {
13651 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
13652 return Ok(Some((y, 1.0)));
13653 }
13654 }
13655 if m != 1 || !self.uses_q8_1_fast(w) {
13657 return Ok(None);
13658 }
13659 let in_f = w.in_features();
13660 let out_f = w.out_features();
13661 let (bytes, qtype, row_bytes, scale, rp) = match w {
13662 GpuTensor::Quant {
13663 bytes,
13664 qtype,
13665 row_bytes,
13666 scale,
13667 rp,
13668 ..
13669 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13670 _ => return Ok(None),
13671 };
13672 if self.mmvq_supports(qtype) {
13674 let (mbytes, mrp) = match w {
13676 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
13677 _ => (bytes, rp),
13678 };
13679 let y = self.qmatvec_mmvq(
13680 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
13681 )?;
13682 return Ok(Some((y, scale)));
13683 }
13684 let name = match qtype {
13686 QT_Q8_0 => "qmatvec_q8_0_dp4a",
13687 QT_Q4_K => "qmatvec_q4_K_dp4a",
13688 QT_Q6_K => "qmatvec_q6_K_dp4a",
13689 QT_Q5_K => "qmatvec_q5_K_dp4a",
13690 QT_Q3_K => "qmatvec_q3_K_dp4a",
13691 QT_NVFP4 => {
13692 if rp {
13693 "qmatvec_nvfp4_dp4a_rp"
13694 } else {
13695 "qmatvec_nvfp4_dp4a"
13696 }
13697 }
13698 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
13699 _ => return Ok(None),
13700 };
13701 let f = self.func(name);
13702 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
13703 let cfg = LaunchConfig {
13704 grid_dim: (out_f as u32, m as u32, 1),
13705 block_dim: (128, 1, 1),
13706 shared_mem_bytes: 0,
13707 };
13708 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13709 let __s_b = self.gpu.stream();
13710 let mut b = __s_b.launch_builder(&f);
13711 b.arg(bytes)
13712 .arg(aq)
13713 .arg(ad)
13714 .arg(&mut y)
13715 .arg(&inf)
13716 .arg(&outf)
13717 .arg(&mi)
13718 .arg(&rb);
13719 unsafe {
13720 b.launch(cfg)?;
13721 }
13722 Ok(Some((y, scale)))
13723 }
13724
13725 pub fn mmvq_supports(&self, qtype: i32) -> bool {
13728 if qtype == QT_F8_E4M3 {
13733 return true;
13734 }
13735 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
13736 return false;
13737 }
13738 matches!(
13739 qtype,
13740 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
13741 )
13742 }
13743
13744 pub fn qmatvec_mmvq(
13749 &self,
13750 bytes: &CudaSlice<u8>,
13751 aq: &CudaSlice<i8>,
13752 ad: &CudaSlice<f32>,
13753 m: usize,
13754 in_f: usize,
13755 out_f: usize,
13756 qtype: i32,
13757 row_bytes: usize,
13758 scale: f32,
13759 rp: bool,
13760 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13761 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
13763 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
13764 )?;
13765 Ok(y)
13766 }
13767
13768 #[allow(clippy::too_many_arguments)]
13770 pub fn qmatvec_mmvq_into(
13771 &self,
13772 bytes: &CudaSlice<u8>,
13773 aq: &CudaSlice<i8>,
13774 ad: &CudaSlice<f32>,
13775 m: usize,
13776 in_f: usize,
13777 out_f: usize,
13778 qtype: i32,
13779 row_bytes: usize,
13780 scale: f32,
13781 rp: bool,
13782 y: &mut CudaSlice<f32>,
13783 ) -> Result<(), Box<dyn std::error::Error>> {
13784 debug_assert!(y.len() >= m * out_f);
13785 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
13791 && rp
13792 && m == 1
13793 && out_f >= 64
13794 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
13795 && {
13796 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13797 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
13798 }
13799 {
13800 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
13801 let cfg = LaunchConfig {
13802 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
13803 block_dim: (32, 2, 1),
13804 shared_mem_bytes: 0,
13805 };
13806 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
13807 let __s_b = self.gpu.stream();
13808 let mut b = __s_b.launch_builder(&f);
13809 b.arg(bytes)
13810 .arg(aq)
13811 .arg(ad)
13812 .arg(&mut *y)
13813 .arg(&inf)
13814 .arg(&outf)
13815 .arg(&mi)
13816 .arg(&rb);
13817 unsafe {
13818 b.launch(cfg)?;
13819 }
13820 if scale != 1.0 {
13821 self.scale_inplace(y, scale, out_f)?;
13822 }
13823 return Ok(());
13824 }
13825 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
13834 2
13835 } else {
13836 1
13837 };
13838 if m == 1 && qtype == QT_Q4_0 {
13843 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13844 mr = *Q40MR.get_or_init(|| {
13847 std::env::var("MEMRA_Q40_MR")
13848 .ok()
13849 .and_then(|v| v.parse().ok())
13850 .unwrap_or(1)
13851 });
13852 }
13853 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
13864 let q5_force = q5_mode.as_deref() == Some("2");
13865 let q5_il = qtype == QT_Q5_K
13868 && m == 1
13869 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
13870 if q5_il && !q5_force && out_f > 65536 {
13871 mr = 1;
13872 }
13873 if qtype == QT_Q4_0 && rp && mr != 1 {
13876 mr = 2;
13877 }
13878 if qtype == QT_Q8_0 && rp {
13882 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
13883 mr = *Q80MR.get_or_init(|| {
13884 std::env::var("MEMRA_Q80_MR")
13885 .ok()
13886 .and_then(|v| v.parse().ok())
13887 .unwrap_or(1)
13888 });
13889 }
13890 let name = match (qtype, mr, rp) {
13891 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
13892 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
13893 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
13894 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
13895 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
13896 (QT_Q5_K, 2, _) => {
13897 if q5_il {
13898 "qmatvec_q5_K_mmvq_mr2_il"
13899 } else {
13900 "qmatvec_q5_K_mmvq_mr2"
13901 }
13902 }
13903 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
13904 (QT_Q8_0, _, true)
13909 if in_f % 1024 == 0 && {
13910 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13911 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
13912 } =>
13913 {
13914 "qmatvec_q8_0_mmvq_rpca"
13915 }
13916 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
13917 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
13918 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
13922 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
13923 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
13924 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
13925 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
13926 (QT_Q5_K, _, _) => {
13927 if q5_il {
13928 "qmatvec_q5_K_mmvq_il"
13929 } else {
13930 "qmatvec_q5_K_mmvq"
13931 }
13932 }
13933 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
13934 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
13935 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
13936 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
13937 };
13938 let f = self.func(name);
13939 let rows_per_block = ROWS_PER_BLOCK * mr;
13941 let cfg = LaunchConfig {
13942 grid_dim: (
13943 (out_f as u32 + rows_per_block - 1) / rows_per_block,
13944 m as u32,
13945 1,
13946 ),
13947 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
13950 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13951 let __s_b = self.gpu.stream();
13952 let mut b = __s_b.launch_builder(&f);
13953 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
13958 if Self::pdl_on()
13961 && Self::pdl_mmvq_on()
13962 && Self::pdl_nvfp4q8_on()
13963 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
13964 {
13965 use cudarc::driver::{DevicePtr, DevicePtrMut};
13966 let s = &self.gpu.stream();
13967 let (pw, _g0) = bytes.device_ptr(s);
13968 let (paq, _g1) = aq.device_ptr(s);
13969 let (pad, _g2) = ad.device_ptr(s);
13970 let (py, _g3) = y.device_ptr_mut(s);
13971 let mut ps = [
13972 &pw as *const _ as *mut std::ffi::c_void,
13973 &paq as *const _ as *mut _,
13974 &pad as *const _ as *mut _,
13975 &py as *const _ as *mut _,
13976 &inf as *const _ as *mut _,
13977 &outf as *const _ as *mut _,
13978 &mi as *const _ as *mut _,
13979 &rb as *const _ as *mut _,
13980 &scale as *const _ as *mut _,
13981 ];
13982 unsafe {
13983 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
13984 }
13985 return Ok(());
13986 }
13987 b.arg(bytes)
13988 .arg(aq)
13989 .arg(ad)
13990 .arg(&mut *y)
13991 .arg(&inf)
13992 .arg(&outf)
13993 .arg(&mi)
13994 .arg(&rb)
13995 .arg(&scale);
13996 unsafe {
13997 b.launch(cfg)?;
13998 }
13999 } else if Self::pdl_on()
14000 && Self::pdl_mmvq_on()
14001 && (matches!(
14002 name,
14003 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
14004 ) || (Self::pdl_nvfp4q8_on()
14005 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
14006 {
14007 {
14011 use cudarc::driver::{DevicePtr, DevicePtrMut};
14012 let s = &self.gpu.stream();
14013 let (pw, _g0) = bytes.device_ptr(s);
14014 let (paq, _g1) = aq.device_ptr(s);
14015 let (pad, _g2) = ad.device_ptr(s);
14016 let (py, _g3) = y.device_ptr_mut(s);
14017 let mut ps = [
14018 &pw as *const _ as *mut std::ffi::c_void,
14019 &paq as *const _ as *mut _,
14020 &pad as *const _ as *mut _,
14021 &py as *const _ as *mut _,
14022 &inf as *const _ as *mut _,
14023 &outf as *const _ as *mut _,
14024 &mi as *const _ as *mut _,
14025 &rb as *const _ as *mut _,
14026 ];
14027 unsafe {
14028 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
14029 }
14030 }
14031 if scale != 1.0 {
14032 self.scale_inplace(y, scale, m * out_f)?;
14033 }
14034 } else {
14035 b.arg(bytes)
14036 .arg(aq)
14037 .arg(ad)
14038 .arg(&mut *y)
14039 .arg(&inf)
14040 .arg(&outf)
14041 .arg(&mi)
14042 .arg(&rb);
14043 unsafe {
14044 b.launch(cfg)?;
14045 }
14046 if scale != 1.0 {
14047 self.scale_inplace(y, scale, m * out_f)?;
14048 }
14049 }
14050 Ok(())
14051 }
14052
14053 pub fn qmatvec_mmvq_raw(
14057 &self,
14058 bytes: &CudaSlice<u8>,
14059 x: &CudaSlice<f32>,
14060 m: usize,
14061 in_f: usize,
14062 out_f: usize,
14063 qtype: i32,
14064 row_bytes: usize,
14065 rp: bool,
14066 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14067 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14068 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
14069 }
14070
14071 pub fn batched_supports(&self, qtype: i32) -> bool {
14075 matches!(
14076 qtype,
14077 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
14078 )
14079 }
14080
14081 pub fn iq_fast_enabled() -> bool {
14089 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14090 *ON.get_or_init(|| {
14091 std::env::var("MEMRA_IQ_FAST")
14092 .map(|v| v != "0")
14093 .unwrap_or(true)
14094 })
14095 }
14096
14097 pub fn b8_enabled() -> bool {
14100 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14101 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
14102 }
14103
14104 pub fn batched_mcols(m: usize) -> usize {
14106 if m == 2 {
14107 2
14108 } else if m <= 4 {
14109 4
14110 } else if m <= 8 {
14111 8
14112 } else {
14113 16
14114 }
14115 }
14116
14117 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
14122 Some(match (qtype, mcols) {
14123 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
14124 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
14125 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
14126 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
14132 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
14133 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
14134 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
14135 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
14138 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
14139 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
14140 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
14141 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
14144 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
14145 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
14146 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
14147 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
14148 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
14149 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
14150 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
14151 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
14155 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
14156 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
14157 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
14158 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
14162 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
14163 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
14164 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
14165 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
14166 _ => return None,
14167 })
14168 }
14169
14170 pub fn sm_count(&self) -> i32 {
14205 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14206 *SMS.get_or_init(|| {
14207 use cudarc::driver::sys::CUdevice_attribute_enum as A;
14208 self.gpu
14209 .ctx
14210 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14211 .unwrap_or(82)
14212 })
14213 }
14214
14215 pub fn batched_variant(
14216 &self,
14217 _m: usize,
14218 in_f: usize,
14219 out_f: usize,
14220 qtype: i32,
14221 row_bytes: usize,
14222 mcols: usize,
14223 rp: bool,
14224 ) -> &'static str {
14225 if qtype == QT_Q8_0 {
14230 return if rp { "rp" } else { "base" };
14231 }
14232 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14233 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
14234 Ok("base") => "base",
14235 Ok("pf") => "pf",
14236 Ok("r2") => "r2",
14237 Ok("r2w8") => "r2w8",
14238 Ok("pfr2") => "pfr2",
14239 Ok("ca") => "ca",
14240 Ok("car2") => "car2",
14241 Ok("rp") => "rp",
14244 Ok("rpr2") => "rpr2",
14245 Ok("rpr2w8") => "rpr2w8",
14246 Ok("rpca") => "rpca",
14249 Ok("rpcar2") => "rpcar2",
14250 Ok("rpsc") => "rpsc",
14257 Ok("rpms") => "rpms",
14258 Ok("rpmsc") => "rpmsc",
14259 Ok("rpks") => "rpks",
14260 Ok("rpksc") => "rpksc",
14261 _ => "auto",
14262 });
14263 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
14267 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14272 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
14273 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
14274 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
14275 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
14276 let sms = *SMS.get_or_init(|| {
14277 use cudarc::driver::sys::CUdevice_attribute_enum as A;
14278 self.gpu
14279 .ctx
14280 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
14281 .unwrap_or(82)
14282 });
14283 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
14303 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14306 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
14307 Ok("base") => "base",
14308 Ok("r2") => "r2",
14309 Ok("r2w8") => "r2w8",
14310 _ => "auto",
14311 });
14312 let variant: &'static str = if qtype == QT_Q4_0 {
14313 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
14317 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
14318 Ok("base") => "base",
14324 Ok("r2") => "r2",
14325 Ok("ms") => "ms",
14326 Ok("sm") => "sm",
14327 Ok("la") => "la",
14328 _ => "auto",
14329 });
14330 let v = if q40 != "auto" {
14331 q40
14332 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
14333 "r2"
14334 } else {
14335 "base"
14336 };
14337 if rp {
14342 match v {
14343 "ms" => "r2ms_rp",
14344 "sm" => "r2sm_rp",
14345 "la" => "r2la_rp",
14346 "r2" => "r2_rp",
14347 _ => "rp",
14348 }
14349 } else if matches!(v, "ms" | "sm" | "la") {
14350 "r2"
14351 } else {
14352 v
14353 }
14354 } else if qtype != QT_NVFP4 && !kq_r2 {
14355 "base"
14356 } else if kq_r2 && rp {
14357 "rp"
14361 } else if kq_r2 {
14362 if kq_bv != "auto" {
14365 if kq_bv == "r2w8" && mcols != 4 {
14366 "r2"
14367 } else {
14368 kq_bv
14369 }
14370 } else if bv != "auto" {
14371 match bv {
14372 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
14373 "r2w8" | "rpr2w8" => {
14374 if mcols != 4 {
14375 "r2"
14376 } else {
14377 "r2w8"
14378 }
14379 }
14380 _ => "base", }
14382 } else {
14383 let blocks = (out_f + 7) / 8;
14384 let waves = blocks as f64 / (7 * sms as usize) as f64;
14385 let filled = blocks >= 4 * sms as usize;
14386 let use_r2 = if qtype == QT_Q4_K {
14387 filled
14388 } else {
14389 waves >= 2.0
14390 };
14391 if use_r2 { "r2" } else { "base" }
14392 }
14393 } else if bv != "auto" {
14394 let v = if bv == "r2w8" && mcols == 2 {
14399 "r2"
14400 } else if bv == "ca" && (!ca_ok || mcols == 8) {
14401 "pf"
14402 } else if bv == "car2" && (!ca_ok || mcols == 8) {
14403 "r2"
14404 } else if bv == "pfr2" && mcols == 8 {
14405 "r2"
14406 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
14407 "rpr2"
14408 }
14409 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
14411 if mcols == 8 { "rpr2w8" } else { "rpr2" }
14412 } else if bv == "rpcar2" && mcols == 2 {
14413 "rpca"
14414 }
14415 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
14418 "rpr2"
14419 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
14420 "rpr2"
14421 } else {
14422 bv
14423 };
14424 if rp {
14425 match v {
14426 "base" | "pf" | "ca" | "rp" => "rp",
14427 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
14428 "r2w8" | "rpr2w8" => {
14429 if mcols == 2 {
14430 "rpr2"
14431 } else {
14432 "rpr2w8"
14433 }
14434 }
14435 other => other, }
14437 } else {
14438 v
14439 }
14440 } else if mcols == 8 {
14441 if rp {
14452 if sc_ok { "rpsc" } else { "rpr2w8" }
14453 } else {
14454 "r2w8"
14455 }
14456 } else if mcols >= 4 {
14457 let blocks = (out_f + 7) / 8;
14461 let r7 = 7 * sms as usize;
14462 let r8 = 8 * sms as usize;
14463 let waves = blocks as f64 / r7 as f64;
14464 let filled = blocks >= 4 * sms as usize;
14465 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
14469 if rp { "rpr2w8" } else { "r2w8" }
14473 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
14474 if rp { "rpr2" } else { "r2" }
14477 } else {
14478 if rp { "rp" } else { "pf" }
14482 }
14483 } else if in_f >= 6144 {
14484 if rp { "rpr2" } else { "r2" }
14488 } else if rp {
14489 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
14494 if sc_ok && waves >= 0.9 && waves <= 1.1 {
14495 "rpsc"
14496 } else {
14497 "rp"
14498 }
14499 } else {
14500 "base"
14501 };
14502 variant
14503 }
14504
14505 pub fn qmatvec_mmvq_batched(
14506 &self,
14507 bytes: &CudaSlice<u8>,
14508 aq: &CudaSlice<i8>,
14509 ad: &CudaSlice<f32>,
14510 m: usize,
14511 in_f: usize,
14512 out_f: usize,
14513 qtype: i32,
14514 row_bytes: usize,
14515 mcols: usize,
14516 scale: f32,
14517 rp: bool,
14518 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14519 const ROWS_PER_BLOCK: u32 = 4;
14520 let forced: Option<&'static str> = {
14525 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
14526 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
14527 .as_deref()
14528 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
14529 };
14530 let variant = match forced {
14531 Some(v) if !rp || v.contains("rp") => v,
14532 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
14533 };
14534 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
14535 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
14536 })?;
14537 let variant = if mcols == 16 {
14541 if rp { "rp" } else { "base" }
14542 } else {
14543 variant
14544 };
14545 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14552 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
14553 if b567
14554 && qtype == QT_NVFP4
14555 && rp
14556 && mcols == 8
14557 && (5..=7).contains(&m)
14558 && matches!(variant, "rpsc" | "rpr2w8")
14559 {
14560 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
14561 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14563 let cfg = LaunchConfig {
14564 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14565 block_dim: (32, ROWS_PER_BLOCK, 1),
14566 shared_mem_bytes: 0,
14567 };
14568 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14569 let __s_b = self.gpu.stream();
14570 let mut b = __s_b.launch_builder(&f);
14571 b.arg(bytes)
14572 .arg(aq)
14573 .arg(ad)
14574 .arg(&mut y)
14575 .arg(&inf)
14576 .arg(&outf)
14577 .arg(&mi)
14578 .arg(&rb);
14579 unsafe {
14580 b.launch(cfg)?;
14581 }
14582 if scale != 1.0 {
14583 self.scale_inplace(&mut y, scale, m * out_f)?;
14584 }
14585 return Ok(y);
14586 }
14587 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
14588 "base" => (base_name.into(), ROWS_PER_BLOCK),
14589 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
14590 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
14591 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
14592 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
14596 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
14597 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
14598 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
14599 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
14600 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
14601 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
14602 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
14604 debug_assert!(
14605 !rp || name.contains("_rp"),
14606 "rp weight dispatched to a GGUF-layout kernel"
14607 );
14608 let f = self.func(&name);
14609 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14610 let smem = if name.contains("_r2sm_rp") {
14612 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
14613 } else {
14614 0
14615 };
14616 let cfg = LaunchConfig {
14617 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
14618 block_dim: (32, ROWS_PER_BLOCK, 1),
14619 shared_mem_bytes: smem,
14620 };
14621 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14622 let __s_b = self.gpu.stream();
14623 let mut b = __s_b.launch_builder(&f);
14624 b.arg(bytes)
14625 .arg(aq)
14626 .arg(ad)
14627 .arg(&mut y)
14628 .arg(&inf)
14629 .arg(&outf)
14630 .arg(&mi)
14631 .arg(&rb);
14632 unsafe {
14633 b.launch(cfg)?;
14634 }
14635 if scale != 1.0 {
14636 self.scale_inplace(&mut y, scale, m * out_f)?;
14637 }
14638 Ok(y)
14639 }
14640
14641 pub fn qmatvec_batched_raw(
14645 &self,
14646 bytes: &CudaSlice<u8>,
14647 x: &CudaSlice<f32>,
14648 m: usize,
14649 in_f: usize,
14650 out_f: usize,
14651 qtype: i32,
14652 row_bytes: usize,
14653 mcols: usize,
14654 rp: bool,
14655 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14656 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14657 self.qmatvec_mmvq_batched(
14658 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
14659 )
14660 }
14661
14662 pub fn qmatvec_nvfp4_batched_raw(
14664 &self,
14665 bytes: &CudaSlice<u8>,
14666 x: &CudaSlice<f32>,
14667 m: usize,
14668 in_f: usize,
14669 out_f: usize,
14670 row_bytes: usize,
14671 mcols: usize,
14672 rp: bool,
14673 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14674 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
14675 }
14676
14677 fn try_fp4_gemm(
14681 &self,
14682 w: &crate::model::GpuTensor,
14683 x: &CudaSlice<f32>,
14684 m: usize,
14685 in_f: usize,
14686 out_f: usize,
14687 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14688 use crate::model::GpuTensor;
14689 if cfg!(memra_portable_cuda) {
14690 return Ok(None);
14691 }
14692 if std::env::var("MEMRA_FP4").is_err() {
14693 return Ok(None);
14694 }
14695 #[cfg(memra_cutlass)]
14704 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
14705 if let GpuTensor::Quant {
14706 bytes,
14707 qtype,
14708 scale,
14709 row_bytes,
14710 cutlass,
14711 ..
14712 } = w
14713 {
14714 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
14715 if let Some(cw) = cutlass {
14716 let y = self.cutlass_fp4_gemm(
14718 &cw.b_packed,
14719 &cw.sfb_swizzled,
14720 x,
14721 *scale,
14722 m,
14723 out_f,
14724 in_f,
14725 )?;
14726 return Ok(Some(y));
14727 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
14728 let (b_packed, sfb_sw) =
14733 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
14734 let y =
14735 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
14736 return Ok(Some(y));
14737 }
14738 }
14739 }
14740 }
14741 if let GpuTensor::Quant {
14742 bytes,
14743 qtype,
14744 row_bytes,
14745 scale,
14746 rp,
14747 ..
14748 } = w
14749 {
14750 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
14753 let y =
14754 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
14755 return Ok(Some(y));
14756 }
14757 }
14758 Ok(None)
14759 }
14760
14761 pub fn rms_norm_f16out(
14765 &self,
14766 x: &CudaSlice<f32>,
14767 w: &CudaSlice<f32>,
14768 dst: &mut CudaSlice<f32>,
14769 dst16: &mut CudaSlice<u8>,
14770 ncols: usize,
14771 nrows: usize,
14772 eps: f32,
14773 ) -> Result<(), Box<dyn std::error::Error>> {
14774 let f = self.func("rms_norm_f16out_f32");
14775 let cfg = LaunchConfig {
14776 grid_dim: (nrows as u32, 1, 1),
14777 block_dim: (rms_block(), 1, 1),
14778 shared_mem_bytes: 0,
14779 };
14780 let (nc, e) = (ncols as i32, eps);
14781 let __s_b = self.gpu.stream();
14782 let mut b = __s_b.launch_builder(&f);
14783 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
14784 unsafe {
14785 b.launch(cfg)?;
14786 }
14787 Ok(())
14788 }
14789
14790 #[allow(clippy::too_many_arguments)]
14793 pub fn add_rms_norm_f16out(
14794 &self,
14795 a: &CudaSlice<f32>,
14796 b: &CudaSlice<f32>,
14797 w: &CudaSlice<f32>,
14798 res: &mut CudaSlice<f32>,
14799 dst: &mut CudaSlice<f32>,
14800 dst16: &mut CudaSlice<u8>,
14801 ncols: usize,
14802 nrows: usize,
14803 eps: f32,
14804 ) -> Result<(), Box<dyn std::error::Error>> {
14805 let f = self.func("add_rms_norm_f16out_f32");
14806 let cfg = LaunchConfig {
14807 grid_dim: (nrows as u32, 1, 1),
14808 block_dim: (rms_block(), 1, 1),
14809 shared_mem_bytes: 0,
14810 };
14811 let (nc, e) = (ncols as i32, eps);
14812 let __s_lb = self.gpu.stream();
14813 let mut lb = __s_lb.launch_builder(&f);
14814 lb.arg(a)
14815 .arg(b)
14816 .arg(w)
14817 .arg(res)
14818 .arg(dst)
14819 .arg(dst16)
14820 .arg(&nc)
14821 .arg(&e);
14822 unsafe {
14823 lb.launch(cfg)?;
14824 }
14825 Ok(())
14826 }
14827
14828 pub fn matmul_group_xh(
14831 &self,
14832 ws: &[&crate::model::GpuTensor],
14833 x: &CudaSlice<f32>,
14834 xh: &CudaSlice<u8>,
14835 m: usize,
14836 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14837 let mut out = Vec::with_capacity(ws.len());
14838 let in_f = ws[0].in_features();
14839 for w in ws {
14840 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
14841 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
14842 out.push(y);
14843 continue;
14844 }
14845 }
14846 out.push(self.matmul(w, x, m)?);
14847 }
14848 Ok(out)
14849 }
14850
14851 pub fn gdn_pad_mask(
14854 &self,
14855 beta: &mut CudaSlice<f32>,
14856 g_log: &mut CudaSlice<f32>,
14857 len_d: &CudaSlice<i32>,
14858 h: usize,
14859 t: usize,
14860 ) -> Result<(), Box<dyn std::error::Error>> {
14861 let f = self.func("gdn_pad_mask_f32");
14862 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
14863 let (hi, ti) = (h as i32, t as i32);
14864 let __s_b = self.gpu.stream();
14865 let mut b = __s_b.launch_builder(&f);
14866 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
14867 unsafe {
14868 b.launch(cfg)?;
14869 }
14870 Ok(())
14871 }
14872
14873 pub fn row_gather_dev(
14876 &self,
14877 src: &CudaSlice<f32>,
14878 dst: &mut CudaSlice<f32>,
14879 len_d: &CudaSlice<i32>,
14880 ncols: usize,
14881 ) -> Result<(), Box<dyn std::error::Error>> {
14882 let f = self.func("row_gather_dev_f32");
14883 let cfg = LaunchConfig::for_num_elems(ncols as u32);
14884 let nc = ncols as i32;
14885 let __s_b = self.gpu.stream();
14886 let mut b = __s_b.launch_builder(&f);
14887 b.arg(src).arg(dst).arg(len_d).arg(&nc);
14888 unsafe {
14889 b.launch(cfg)?;
14890 }
14891 Ok(())
14892 }
14893
14894 pub fn matmul_group(
14901 &self,
14902 ws: &[&crate::model::GpuTensor],
14903 x: &CudaSlice<f32>,
14904 m: usize,
14905 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14906 use crate::model::GpuTensor;
14907 let mut out = Vec::with_capacity(ws.len());
14908 let any_mirror = ws
14909 .iter()
14910 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
14911 if m >= 16 && any_mirror && !self.verify_exact_on() {
14912 let in_f = ws[0].in_features();
14913 let xh = self.f16_act(x, m * in_f, in_f)?;
14914 for w in ws {
14915 if w.in_features() == in_f {
14916 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
14917 out.push(y);
14918 continue;
14919 }
14920 }
14921 out.push(self.matmul(w, x, m)?);
14922 }
14923 return Ok(out);
14924 }
14925 for w in ws {
14926 out.push(self.matmul(w, x, m)?);
14927 }
14928 Ok(out)
14929 }
14930
14931 pub fn matmul_group_multi(
14938 &self,
14939 ws: &[&crate::model::GpuTensor],
14940 xs: &[&CudaSlice<f32>],
14941 ms: &[usize],
14942 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
14943 assert_eq!(xs.len(), ms.len());
14944 let in_f = ws[0].in_features();
14945 let total: usize = ms.iter().sum();
14946 let mut xcat = self.uninit(total * in_f)?;
14947 let mut off = 0usize;
14948 for (x, &m) in xs.iter().zip(ms) {
14949 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
14950 off += m;
14951 }
14952 let ys = self.matmul_group(ws, &xcat, total)?;
14953 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
14954 for (w, y) in ws.iter().zip(ys) {
14955 let out_f = w.out_features();
14956 let mut off = 0usize;
14957 for (s, &m) in ms.iter().enumerate() {
14958 let mut ys_s = self.uninit(m * out_f)?;
14959 let src = y.slice(off * out_f..(off + m) * out_f);
14960 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
14961 out[s].push(ys_s);
14962 off += m;
14963 }
14964 }
14965 Ok(out)
14966 }
14967
14968 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
14978 use crate::model::GpuTensor;
14979 if !legacy_quant_gemm_allowed(
14980 cfg!(memra_portable_cuda),
14981 cfg!(memra_hopper_mma),
14982 std::env::var_os("MEMRA_NO_GEMM").is_some(),
14983 ) {
14984 return false;
14985 }
14986 match w {
14987 GpuTensor::Quant { qtype, .. } => {
14988 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
14989 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
14990 }
14991 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
14992 }
14993 }
14994
14995 pub fn qmatvec_gemm(
15002 &self,
15003 w: &crate::model::GpuTensor,
15004 aq: &CudaSlice<i8>,
15005 ad: &CudaSlice<f32>,
15006 m: usize,
15007 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15008 use crate::model::GpuTensor;
15009 let in_f = w.in_features();
15010 let out_f = w.out_features();
15011 let (bytes, qtype, row_bytes, scale, rp) = match w {
15012 GpuTensor::Quant {
15013 bytes,
15014 qtype,
15015 row_bytes,
15016 scale,
15017 rp,
15018 ..
15019 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15020 _ => unreachable!("gemm_supports guaranteed Quant"),
15021 };
15022 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
15028 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
15029 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
15030 if scale != 1.0 {
15031 self.scale_inplace(&mut y, scale, m * out_f)?;
15032 }
15033 return Ok(y);
15034 }
15035 }
15036 let name = match qtype {
15037 QT_Q8_0 => "qmatvec_gemm_q8_0",
15038 QT_Q4_K => "qmatvec_gemm_q4_K",
15039 QT_Q4_0 => {
15040 if rp {
15041 "qmatvec_gemm_q4_0_rp"
15042 } else {
15043 "qmatvec_gemm_q4_0"
15044 }
15045 }
15046 QT_Q5_K => "qmatvec_gemm_q5_K",
15047 QT_Q6_K => "qmatvec_gemm_q6_K",
15048 QT_NVFP4 => {
15049 if rp {
15050 "qmatvec_gemm_nvfp4_rp"
15051 } else {
15052 "qmatvec_gemm_nvfp4"
15053 }
15054 }
15055 _ => unreachable!(),
15056 };
15057 let f = self.func(name);
15058 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);
15063 let k1_tile = if is_k1 {
15065 k1_launch_override().unwrap_or((128, 128, 8))
15066 } else {
15067 (128, 128, 8)
15068 };
15069 let (bm, bn): (u32, u32) = if is_k1 {
15070 (k1_tile.0, k1_tile.1)
15071 } else {
15072 (64, 256)
15073 };
15074 let warps: u32 = if is_k1 {
15075 k1_tile.2
15076 } else {
15077 match qtype {
15078 QT_NVFP4 => 8,
15079 _ => 4,
15080 }
15081 };
15082 let cfg = LaunchConfig {
15083 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15084 block_dim: (32, warps, 1),
15085 shared_mem_bytes: 0,
15086 };
15087 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15088 let __s_b = self.gpu.stream();
15089 let mut b = __s_b.launch_builder(&f);
15090 b.arg(bytes)
15091 .arg(aq)
15092 .arg(ad)
15093 .arg(&mut y)
15094 .arg(&inf)
15095 .arg(&outf)
15096 .arg(&mi)
15097 .arg(&rb);
15098 unsafe {
15099 b.launch(cfg)?;
15100 }
15101 if scale != 1.0 {
15102 self.scale_inplace(&mut y, scale, m * out_f)?;
15103 }
15104 Ok(y)
15105 }
15106
15107 pub fn qmatvec_gemm_raw(
15112 &self,
15113 bytes: &CudaSlice<u8>,
15114 x: &CudaSlice<f32>,
15115 m: usize,
15116 in_f: usize,
15117 out_f: usize,
15118 qtype: i32,
15119 row_bytes: usize,
15120 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15121 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15122 let name = match qtype {
15123 QT_Q8_0 => "qmatvec_gemm_q8_0",
15124 QT_Q4_K => "qmatvec_gemm_q4_K",
15125 QT_Q4_0 => "qmatvec_gemm_q4_0",
15126 QT_Q5_K => "qmatvec_gemm_q5_K",
15127 QT_Q6_K => "qmatvec_gemm_q6_K",
15128 QT_NVFP4 => "qmatvec_gemm_nvfp4",
15129 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
15130 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
15131 };
15132 let f = self.func(name);
15133 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);
15137 let k1_tile = if is_k1 {
15139 k1_launch_override().unwrap_or((128, 128, 8))
15140 } else {
15141 (128, 128, 8)
15142 };
15143 let (bm, bn): (u32, u32) = if is_k1 {
15144 (k1_tile.0, k1_tile.1)
15145 } else {
15146 (64, 256)
15147 };
15148 let warps: u32 = if is_k1 {
15149 k1_tile.2
15150 } else {
15151 match qtype {
15152 QT_NVFP4 | QT_NVFP4_RP => 8,
15153 _ => 4,
15154 }
15155 };
15156 let cfg = LaunchConfig {
15157 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
15158 block_dim: (32, warps, 1),
15159 shared_mem_bytes: 0,
15160 };
15161 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15162 let __s_b = self.gpu.stream();
15163 let mut b = __s_b.launch_builder(&f);
15164 b.arg(bytes)
15165 .arg(&aq)
15166 .arg(&ad)
15167 .arg(&mut y)
15168 .arg(&inf)
15169 .arg(&outf)
15170 .arg(&mi)
15171 .arg(&rb);
15172 unsafe {
15173 b.launch(cfg)?;
15174 }
15175 Ok(y)
15176 }
15177
15178 pub fn qmatvec_gemm_q8_0_wgmma_raw(
15185 &self,
15186 rp4: &CudaSlice<u8>,
15187 aq: &CudaSlice<i8>,
15188 ad: &CudaSlice<f32>,
15189 m: usize,
15190 in_f: usize,
15191 out_f: usize,
15192 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15193 assert!(
15194 out_f % 64 == 0 && in_f % 32 == 0,
15195 "wgmma GEMM needs out_f%64==0, in_f%32==0"
15196 );
15197 let f = self.func("qmatvec_gemm_q8_0_wgmma");
15198 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
15200 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
15201 block_dim: (128, 1, 1),
15202 shared_mem_bytes: 0,
15203 };
15204 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
15205 let __s_b = self.gpu.stream();
15206 let mut b = __s_b.launch_builder(&f);
15207 b.arg(rp4)
15208 .arg(aq)
15209 .arg(ad)
15210 .arg(&mut y)
15211 .arg(&inf)
15212 .arg(&outf)
15213 .arg(&mi);
15214 unsafe {
15215 b.launch(cfg)?;
15216 }
15217 Ok(y)
15218 }
15219
15220 pub fn scale_inplace(
15222 &self,
15223 y: &mut CudaSlice<f32>,
15224 s: f32,
15225 n: usize,
15226 ) -> Result<(), Box<dyn std::error::Error>> {
15227 let f = self.func("scale_f32");
15228 let cfg = LaunchConfig::for_num_elems(n as u32);
15229 let (sf, ni) = (s, n as i32);
15230 let __s_b = self.gpu.stream();
15231 let mut b = __s_b.launch_builder(&f);
15232 b.arg(y).arg(&sf).arg(&ni);
15233 unsafe {
15234 b.launch(cfg)?;
15235 }
15236 Ok(())
15237 }
15238
15239 pub fn bf16_to_f32(
15244 &self,
15245 data: &cudarc::driver::CudaView<'_, u8>,
15246 n: usize,
15247 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15248 let mut out = self.alloc_uninit::<f32>(n)?;
15249 let f = self.func("bf16_to_f32");
15250 let cfg = LaunchConfig::for_num_elems(n as u32);
15251 let ni = n as i32;
15252 let __s_b = self.gpu.stream();
15253 let mut b = __s_b.launch_builder(&f);
15254 b.arg(data).arg(&mut out).arg(&ni);
15255 unsafe {
15256 b.launch(cfg)?;
15257 }
15258 Ok(out)
15259 }
15260
15261 fn linear_bf16_chunked(
15268 &self,
15269 x: &CudaSlice<f32>,
15270 data: &CudaSlice<u8>,
15271 m: usize,
15272 in_f: usize,
15273 out_f: usize,
15274 exact: bool,
15275 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15276 const CHUNK_BYTES: usize = 256 << 20;
15277 let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
15278 if chunk_rows >= out_f {
15279 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
15280 return if exact {
15281 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
15282 } else {
15283 self.linear(x, &wf32, m, in_f, out_f)
15284 };
15285 }
15286 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15287 let mut r0 = 0usize;
15288 while r0 < out_f {
15289 let rows = chunk_rows.min(out_f - r0);
15290 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
15291 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
15292 let yc = if exact {
15293 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
15294 } else {
15295 self.linear(x, &wf32, m, in_f, rows)?
15296 };
15297 for mi in 0..m {
15299 let src = yc.slice(mi * rows..(mi + 1) * rows);
15300 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
15301 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
15302 }
15303 r0 += rows;
15304 }
15305 Ok(y)
15306 }
15307
15308 pub fn linear_decode_exact(
15315 &self,
15316 x: &CudaSlice<f32>,
15317 w: &CudaSlice<f32>,
15318 m_tokens: usize,
15319 in_f: usize,
15320 out_f: usize,
15321 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15322 if m_tokens == 1 {
15323 return self.linear(x, w, 1, in_f, out_f);
15324 }
15325 let xv = self.view(x, m_tokens * in_f);
15326 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
15327 for t in 0..m_tokens {
15328 let row = xv.slice(t * in_f..(t + 1) * in_f);
15329 let mut xr = self.alloc_uninit::<f32>(in_f)?;
15330 self.copy_view_into(&mut xr, 0, &row, in_f)?;
15331 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
15332 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
15333 }
15334 Ok(y)
15335 }
15336
15337 pub fn linear(
15338 &self,
15339 x: &CudaSlice<f32>,
15340 w: &CudaSlice<f32>,
15341 m_tokens: usize,
15342 in_f: usize,
15343 out_f: usize,
15344 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15345 use cudarc::cublaslt::{Matmul, MatmulConfig};
15346 let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; let cfg = MatmulConfig {
15348 transa: true,
15349 transb: false,
15350 transc: false,
15351 m: out_f as u64,
15352 n: m_tokens as u64,
15353 k: in_f as u64,
15354 alpha: 1.0,
15355 lda: in_f as i64,
15356 ldb: in_f as i64,
15357 beta: 0.0,
15358 ldc: out_f as i64,
15359 stride_a: None,
15360 stride_b: None,
15361 stride_c: None,
15362 stride_bias: None,
15363 batch_size: None,
15364 };
15365 unsafe {
15366 self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?;
15367 }
15368 Ok(c)
15369 }
15370
15371 pub fn sdpa_naive(
15373 &self,
15374 q: &CudaSlice<f32>,
15375 k: &CudaSlice<f32>,
15376 v: &CudaSlice<f32>,
15377 o: &mut CudaSlice<f32>,
15378 head_dim: usize,
15379 n_head: usize,
15380 n_head_kv: usize,
15381 t: usize,
15382 t_kv: usize,
15383 scale: f32,
15384 causal: bool,
15385 ) -> Result<(), Box<dyn std::error::Error>> {
15386 let f = self.func("sdpa_naive_f32");
15387 let cfg = LaunchConfig {
15388 grid_dim: (n_head as u32, t as u32, 1),
15389 block_dim: (128, 1, 1),
15390 shared_mem_bytes: (t_kv * 4) as u32,
15391 };
15392 let (hd, nh, nhkv, ti, tkvi, cz) = (
15393 head_dim as i32,
15394 n_head as i32,
15395 n_head_kv as i32,
15396 t as i32,
15397 t_kv as i32,
15398 causal as i32,
15399 );
15400 let __s_b = self.gpu.stream();
15401 let mut b = __s_b.launch_builder(&f);
15402 b.arg(q)
15403 .arg(k)
15404 .arg(v)
15405 .arg(o)
15406 .arg(&hd)
15407 .arg(&nh)
15408 .arg(&nhkv)
15409 .arg(&ti)
15410 .arg(&tkvi)
15411 .arg(&scale)
15412 .arg(&cz);
15413 unsafe {
15414 b.launch(cfg)?;
15415 }
15416 Ok(())
15417 }
15418
15419 #[allow(clippy::too_many_arguments)]
15424 pub fn sdpa_naive_island(
15425 &self,
15426 q: &CudaSlice<f32>,
15427 k: &CudaSlice<f32>,
15428 v: &CudaSlice<f32>,
15429 o: &mut CudaSlice<f32>,
15430 span_id: &CudaSlice<i32>,
15431 head_dim: usize,
15432 n_head: usize,
15433 n_head_kv: usize,
15434 t: usize,
15435 t_kv: usize,
15436 scale: f32,
15437 window: usize,
15438 ) -> Result<(), Box<dyn std::error::Error>> {
15439 let f = self.func("sdpa_naive_island_f32");
15440 let cfg = LaunchConfig {
15441 grid_dim: (n_head as u32, t as u32, 1),
15442 block_dim: (128, 1, 1),
15443 shared_mem_bytes: (t_kv * 4) as u32,
15444 };
15445 let (hd, nh, nhkv, ti, tkvi, wi) = (
15446 head_dim as i32,
15447 n_head as i32,
15448 n_head_kv as i32,
15449 t as i32,
15450 t_kv as i32,
15451 window as i32,
15452 );
15453 let __s_b = self.gpu.stream();
15454 let mut b = __s_b.launch_builder(&f);
15455 b.arg(q)
15456 .arg(k)
15457 .arg(v)
15458 .arg(o)
15459 .arg(span_id)
15460 .arg(&hd)
15461 .arg(&nh)
15462 .arg(&nhkv)
15463 .arg(&ti)
15464 .arg(&tkvi)
15465 .arg(&scale)
15466 .arg(&wi);
15467 unsafe {
15468 b.launch(cfg)?;
15469 }
15470 Ok(())
15471 }
15472
15473 #[allow(clippy::too_many_arguments)]
15475 pub fn sdpa_naive_w(
15476 &self,
15477 q: &CudaSlice<f32>,
15478 k: &CudaSlice<f32>,
15479 v: &CudaSlice<f32>,
15480 o: &mut CudaSlice<f32>,
15481 head_dim: usize,
15482 n_head: usize,
15483 n_head_kv: usize,
15484 t: usize,
15485 t_kv: usize,
15486 scale: f32,
15487 causal: bool,
15488 window: usize,
15489 ) -> Result<(), Box<dyn std::error::Error>> {
15490 let f = self.func("sdpa_naive_w_f32");
15491 let cfg = LaunchConfig {
15492 grid_dim: (n_head as u32, t as u32, 1),
15493 block_dim: (128, 1, 1),
15494 shared_mem_bytes: (t_kv * 4) as u32,
15495 };
15496 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
15497 head_dim as i32,
15498 n_head as i32,
15499 n_head_kv as i32,
15500 t as i32,
15501 t_kv as i32,
15502 causal as i32,
15503 window as i32,
15504 );
15505 let __s_b = self.gpu.stream();
15506 let mut b = __s_b.launch_builder(&f);
15507 b.arg(q)
15508 .arg(k)
15509 .arg(v)
15510 .arg(o)
15511 .arg(&hd)
15512 .arg(&nh)
15513 .arg(&nhkv)
15514 .arg(&ti)
15515 .arg(&tkvi)
15516 .arg(&scale)
15517 .arg(&cz)
15518 .arg(&wi);
15519 unsafe {
15520 b.launch(cfg)?;
15521 }
15522 Ok(())
15523 }
15524
15525 pub fn sdpa_naive_view(
15527 &self,
15528 q: &CudaSlice<f32>,
15529 k: &cudarc::driver::CudaView<f32>,
15530 v: &cudarc::driver::CudaView<f32>,
15531 o: &mut CudaSlice<f32>,
15532 head_dim: usize,
15533 n_head: usize,
15534 n_head_kv: usize,
15535 t: usize,
15536 t_kv: usize,
15537 scale: f32,
15538 causal: bool,
15539 ) -> Result<(), Box<dyn std::error::Error>> {
15540 let f = self.func("sdpa_naive_f32");
15541 let cfg = LaunchConfig {
15542 grid_dim: (n_head as u32, t as u32, 1),
15543 block_dim: (128, 1, 1),
15544 shared_mem_bytes: (t_kv * 4) as u32,
15545 };
15546 let (hd, nh, nhkv, ti, tkvi, cz) = (
15547 head_dim as i32,
15548 n_head as i32,
15549 n_head_kv as i32,
15550 t as i32,
15551 t_kv as i32,
15552 causal as i32,
15553 );
15554 let __s_b = self.gpu.stream();
15555 let mut b = __s_b.launch_builder(&f);
15556 b.arg(q)
15557 .arg(k)
15558 .arg(v)
15559 .arg(o)
15560 .arg(&hd)
15561 .arg(&nh)
15562 .arg(&nhkv)
15563 .arg(&ti)
15564 .arg(&tkvi)
15565 .arg(&scale)
15566 .arg(&cz);
15567 unsafe {
15568 b.launch(cfg)?;
15569 }
15570 Ok(())
15571 }
15572
15573 #[allow(clippy::too_many_arguments)]
15581 pub fn fa_dequant_kv_view_f32(
15582 &self,
15583 k: &cudarc::driver::CudaView<u8>,
15584 v: &cudarc::driver::CudaView<u8>,
15585 kf: &mut CudaSlice<f32>,
15586 vf: &mut CudaSlice<f32>,
15587 kv_dim_k: usize,
15588 kv_dim_v: usize,
15589 t_kv: usize,
15590 k_tok_bytes: usize,
15591 v_tok_bytes: usize,
15592 g: bool,
15593 ) -> Result<(), Box<dyn std::error::Error>> {
15594 let f = if g {
15595 self.func_g("fa_dequant_kv_ws_f32")
15596 } else {
15597 self.func("fa_dequant_kv_ws_f32")
15598 };
15599 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
15600 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15601 let cfg = LaunchConfig {
15602 grid_dim: (nblk.max(1), 1, 1),
15603 block_dim: (256, 1, 1),
15604 shared_mem_bytes: 0,
15605 };
15606 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
15607 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15608 let __s_b = self.gpu.stream();
15609 let mut b = __s_b.launch_builder(&f);
15610 b.arg(k)
15611 .arg(v)
15612 .arg(&mut *kf)
15613 .arg(&mut *vf)
15614 .arg(&kdk)
15615 .arg(&kdv)
15616 .arg(&tkvi)
15617 .arg(&ktb)
15618 .arg(&vtb);
15619 unsafe {
15620 b.launch(cfg)?;
15621 }
15622 Ok(())
15623 }
15624
15625 #[allow(clippy::too_many_arguments)]
15626 pub fn sdpa_naive_quantized_view(
15627 &self,
15628 q: &CudaSlice<f32>,
15629 k: &cudarc::driver::CudaView<u8>,
15630 v: &cudarc::driver::CudaView<u8>,
15631 o: &mut CudaSlice<f32>,
15632 head_dim: usize,
15633 n_head: usize,
15634 n_head_kv: usize,
15635 t: usize,
15636 t_kv: usize,
15637 scale: f32,
15638 causal: bool,
15639 k_tok_bytes: usize,
15640 v_tok_bytes: usize,
15641 ) -> Result<(), Box<dyn std::error::Error>> {
15642 let kv_dim = n_head_kv * head_dim;
15643 let mut kf = self.uninit(t_kv * kv_dim)?;
15644 let mut vf = self.uninit(t_kv * kv_dim)?;
15645 let f = self.func("fa_dequant_kv_ws_f32");
15646 let total = (2 * t_kv * kv_dim) as u64;
15647 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15648 let cfg = LaunchConfig {
15649 grid_dim: (nblk.max(1), 1, 1),
15650 block_dim: (256, 1, 1),
15651 shared_mem_bytes: 0,
15652 };
15653 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15654 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15655 let __s_b = self.gpu.stream();
15656 let mut b = __s_b.launch_builder(&f);
15657 b.arg(k)
15658 .arg(v)
15659 .arg(&mut kf)
15660 .arg(&mut vf)
15661 .arg(&kv_dim_i)
15662 .arg(&kv_dim_i)
15663 .arg(&t_kv_i)
15664 .arg(&k_tok_bytes_i)
15665 .arg(&v_tok_bytes_i);
15666 unsafe { b.launch(cfg)? };
15667 self.sdpa_naive(
15668 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15669 )
15670 }
15671
15672 #[allow(clippy::too_many_arguments)]
15684 pub fn sdpa_naive_w_quantized_view(
15685 &self,
15686 q: &CudaSlice<f32>,
15687 k: &cudarc::driver::CudaView<u8>,
15688 v: &cudarc::driver::CudaView<u8>,
15689 o: &mut CudaSlice<f32>,
15690 head_dim: usize,
15691 n_head: usize,
15692 n_head_kv: usize,
15693 t: usize,
15694 t_kv: usize,
15695 scale: f32,
15696 causal: bool,
15697 window: usize,
15698 k_tok_bytes: usize,
15699 v_tok_bytes: usize,
15700 ) -> Result<(), Box<dyn std::error::Error>> {
15701 let kv_dim = n_head_kv * head_dim;
15702 let mut kf = self.uninit(t_kv * kv_dim)?;
15703 let mut vf = self.uninit(t_kv * kv_dim)?;
15704 let f = self.func("fa_dequant_kv_ws_f32");
15705 let total = (2 * t_kv * kv_dim) as u64;
15706 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
15707 let cfg = LaunchConfig {
15708 grid_dim: (nblk.max(1), 1, 1),
15709 block_dim: (256, 1, 1),
15710 shared_mem_bytes: 0,
15711 };
15712 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
15713 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
15714 let __s_b = self.gpu.stream();
15715 let mut b = __s_b.launch_builder(&f);
15716 b.arg(k)
15717 .arg(v)
15718 .arg(&mut kf)
15719 .arg(&mut vf)
15720 .arg(&kv_dim_i)
15721 .arg(&kv_dim_i)
15722 .arg(&t_kv_i)
15723 .arg(&k_tok_bytes_i)
15724 .arg(&v_tok_bytes_i);
15725 unsafe { b.launch(cfg)? };
15726 self.sdpa_naive_w(
15727 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15728 )
15729 }
15730
15731 pub fn fa_prefill(
15735 &self,
15736 q: &CudaSlice<f32>,
15737 k: &CudaSlice<f32>,
15738 v: &CudaSlice<f32>,
15739 o: &mut CudaSlice<f32>,
15740 head_dim: usize,
15741 n_head: usize,
15742 n_head_kv: usize,
15743 t: usize,
15744 t_kv: usize,
15745 scale: f32,
15746 causal: bool,
15747 ) -> Result<(), Box<dyn std::error::Error>> {
15748 if portable_mma_gated() {
15749 return self.sdpa_naive(
15750 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
15751 );
15752 }
15753 let fa3_on = head_dim == 256
15761 && causal
15762 && t == t_kv
15763 && match std::env::var("MEMRA_FA3").as_deref() {
15764 Ok("0") => false,
15765 Ok("1") => true,
15766 _ => cfg!(memra_hopper_mma),
15767 };
15768 if fa3_on {
15769 let n = t * n_head * head_dim;
15770 let nkv = t * n_head_kv * head_dim;
15771 let mut q16 = self.alloc_u8_uninit(n * 2)?;
15772 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
15773 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
15774 self.f32_to_bf16_into(q, &mut q16, n)?;
15775 self.f32_to_bf16_into(k, &mut k16, nkv)?;
15776 self.f32_to_bf16_into(v, &mut v16, nkv)?;
15777 let rc = {
15778 use cudarc::driver::{DevicePtr, DevicePtrMut};
15779 let stream = self.gpu.stream();
15780 let (qp, _g1) = q16.device_ptr(&stream);
15781 let (kp, _g2) = k16.device_ptr(&stream);
15782 let (vp, _g3) = v16.device_ptr(&stream);
15783 let (op, _g4) = o.device_ptr_mut(&stream);
15784 unsafe {
15785 memra_fa3_prefill(
15786 qp as *const core::ffi::c_void,
15787 kp as *const core::ffi::c_void,
15788 vp as *const core::ffi::c_void,
15789 op as *mut f32,
15790 t as i32,
15791 n_head as i32,
15792 n_head_kv as i32,
15793 head_dim as i32,
15794 scale,
15795 stream.cu_stream() as *mut core::ffi::c_void,
15796 )
15797 }
15798 };
15799 if rc != 0 {
15800 return Err(format!("memra_fa3_prefill rc={rc}").into());
15801 }
15802 return Ok(());
15803 }
15804 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15809 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
15810 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
15811 const BLOCK_Q: usize = 64;
15812 const BKX: usize = 32;
15813 let f = self.func("fa_prefill_bf16_p1");
15814 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
15815 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
15816 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15817 f.set_attribute(
15818 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15819 shmem as i32,
15820 )?;
15821 let cfg = LaunchConfig {
15822 grid_dim: (
15823 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
15824 n_head as u32,
15825 1,
15826 ),
15827 block_dim: (32, 4, 1),
15828 shared_mem_bytes: shmem,
15829 };
15830 let (hd, nh, nhkv, ti, tkvi, cz) = (
15831 head_dim as i32,
15832 n_head as i32,
15833 n_head_kv as i32,
15834 t as i32,
15835 t_kv as i32,
15836 causal as i32,
15837 );
15838 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
15839 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
15840 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
15841 let __s_b = self.gpu.stream();
15842 let mut b = __s_b.launch_builder(&f);
15843 b.arg(&qb)
15844 .arg(&kb)
15845 .arg(&vb)
15846 .arg(o)
15847 .arg(&hd)
15848 .arg(&nh)
15849 .arg(&nhkv)
15850 .arg(&ti)
15851 .arg(&tkvi)
15852 .arg(&scale)
15853 .arg(&cz);
15854 unsafe {
15855 b.launch(cfg)?;
15856 }
15857 return Ok(());
15858 }
15859 const BK: usize = 32;
15865 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
15868 let (block_q, warps, w2_sfx): (usize, u32, &str) =
15869 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
15870 let hd_sfx = fa_hd_suffix(head_dim)?;
15874 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15875 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
15880 let (kb16, vb16) = if bf16kv {
15881 let n = t_kv * n_head_kv * head_dim;
15882 let mut kb = self.alloc_u8_uninit(n * 2)?;
15883 let mut vb = self.alloc_u8_uninit(n * 2)?;
15884 let fcv = self.func("f32_to_bf16_bulk");
15885 let ni = n as i64;
15886 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
15887 let __s_b = self.gpu.stream();
15888 let mut b = __s_b.launch_builder(&fcv);
15889 b.arg(k).arg(&mut kb).arg(&ni);
15890 unsafe {
15891 b.launch(cfgc)?;
15892 }
15893 let __s_b = self.gpu.stream();
15894 let mut b = __s_b.launch_builder(&fcv);
15895 b.arg(v).arg(&mut vb).arg(&ni);
15896 unsafe {
15897 b.launch(cfgc)?;
15898 }
15899 (Some(kb), Some(vb))
15900 } else {
15901 (None, None)
15902 };
15903 let f = self.func(&if bf16kv {
15904 format!("fa_prefill_bf16kv_pp{hd_sfx}")
15905 } else {
15906 format!(
15907 "fa_prefill_f32{}{}{hd_sfx}",
15908 if floor { "" } else { "_pp" },
15909 if floor { "" } else { w2_sfx }
15910 )
15911 });
15912 let kv_stages = if bf16kv { 2 } else { 1 };
15915 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
15916 + 4 * (block_q * BK + 2 * block_q)) as u32;
15917 use cudarc::driver::sys::CUfunction_attribute_enum as A;
15918 f.set_attribute(
15919 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
15920 shmem as i32,
15921 )?;
15922 let cfg = LaunchConfig {
15923 grid_dim: (
15924 (t as u32 + block_q as u32 - 1) / block_q as u32,
15925 n_head as u32,
15926 1,
15927 ),
15928 block_dim: (32, warps, 1),
15929 shared_mem_bytes: shmem,
15930 };
15931 let (hd, nh, nhkv, ti, tkvi, cz) = (
15932 head_dim as i32,
15933 n_head as i32,
15934 n_head_kv as i32,
15935 t as i32,
15936 t_kv as i32,
15937 causal as i32,
15938 );
15939 let __s_b = self.gpu.stream();
15940 let mut b = __s_b.launch_builder(&f);
15941 b.arg(q);
15942 match (&kb16, &vb16) {
15943 (Some(kb), Some(vb)) => {
15944 b.arg(kb).arg(vb);
15945 }
15946 _ => {
15947 b.arg(k).arg(v);
15948 }
15949 }
15950 b.arg(o)
15951 .arg(&hd)
15952 .arg(&nh)
15953 .arg(&nhkv)
15954 .arg(&ti)
15955 .arg(&tkvi)
15956 .arg(&scale)
15957 .arg(&cz);
15958 unsafe {
15959 b.launch(cfg)?;
15960 }
15961 Ok(())
15962 }
15963
15964 #[allow(clippy::too_many_arguments)]
15968 pub fn fa_prefill_w(
15969 &self,
15970 q: &CudaSlice<f32>,
15971 k: &CudaSlice<f32>,
15972 v: &CudaSlice<f32>,
15973 o: &mut CudaSlice<f32>,
15974 head_dim: usize,
15975 n_head: usize,
15976 n_head_kv: usize,
15977 t: usize,
15978 t_kv: usize,
15979 scale: f32,
15980 causal: bool,
15981 window: usize,
15982 ) -> Result<(), Box<dyn std::error::Error>> {
15983 if portable_mma_gated() {
15986 return self.sdpa_naive_w(
15987 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
15988 );
15989 }
15990 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15994 let faw_f32 =
15995 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
15996 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
15997 self.fa_prefill_w_arm(
15998 q,
15999 k,
16000 v,
16001 o,
16002 head_dim,
16003 n_head,
16004 n_head_kv,
16005 t,
16006 t_kv,
16007 scale,
16008 causal,
16009 window,
16010 floor || faw_f32,
16011 floor,
16012 )
16013 }
16014
16015 #[allow(clippy::too_many_arguments)]
16018 pub fn fa_prefill_w_pre(
16019 &self,
16020 qb: &CudaSlice<u8>,
16021 kb: &CudaSlice<u8>,
16022 vb: &CudaSlice<u8>,
16023 o: &mut CudaSlice<f32>,
16024 head_dim: usize,
16025 n_head: usize,
16026 n_head_kv: usize,
16027 t: usize,
16028 t_kv: usize,
16029 scale: f32,
16030 causal: bool,
16031 window: usize,
16032 v_f16: bool,
16033 ) -> Result<(), Box<dyn std::error::Error>> {
16034 const BLOCK_Q: usize = 64;
16035 const BK: usize = 32;
16036 debug_assert_eq!(head_dim, 256);
16037 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16038 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
16039 if hp {
16040 const BLOCK_QH: usize = 32;
16041 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16044 let vh: &CudaSlice<u8> = if v_f16 {
16045 vb
16046 } else {
16047 let n = t_kv * n_head_kv * head_dim;
16048 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
16049 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
16050 }
16051 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
16052 vguard.as_ref().unwrap()
16053 };
16054 let f = self.func("fa_prefill_w_bf16_p1h2");
16055 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16056 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16057 f.set_attribute(
16058 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16059 shmem as i32,
16060 )?;
16061 let cfg = LaunchConfig {
16062 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16063 block_dim: (32, 4, 1),
16064 shared_mem_bytes: shmem,
16065 };
16066 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16067 head_dim as i32,
16068 n_head as i32,
16069 n_head_kv as i32,
16070 t as i32,
16071 t_kv as i32,
16072 causal as i32,
16073 window as i32,
16074 );
16075 let __s_b = self.gpu.stream();
16076 let mut b = __s_b.launch_builder(&f);
16077 b.arg(qb)
16078 .arg(kb)
16079 .arg(vh)
16080 .arg(o)
16081 .arg(&hd)
16082 .arg(&nh)
16083 .arg(&nhkv)
16084 .arg(&ti)
16085 .arg(&tkvi)
16086 .arg(&scale)
16087 .arg(&cz)
16088 .arg(&wi);
16089 unsafe {
16090 b.launch(cfg)?;
16091 }
16092 return Ok(());
16093 }
16094 let f = self.func("fa_prefill_w_bf16_p1");
16095 let shmem =
16096 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16097 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16098 f.set_attribute(
16099 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16100 shmem as i32,
16101 )?;
16102 let cfg = LaunchConfig {
16103 grid_dim: (
16104 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16105 n_head as u32,
16106 1,
16107 ),
16108 block_dim: (32, 4, 1),
16109 shared_mem_bytes: shmem,
16110 };
16111 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16112 head_dim as i32,
16113 n_head as i32,
16114 n_head_kv as i32,
16115 t as i32,
16116 t_kv as i32,
16117 causal as i32,
16118 window as i32,
16119 );
16120 let __s_b = self.gpu.stream();
16121 let mut b = __s_b.launch_builder(&f);
16122 b.arg(qb)
16123 .arg(kb)
16124 .arg(vb)
16125 .arg(o)
16126 .arg(&hd)
16127 .arg(&nh)
16128 .arg(&nhkv)
16129 .arg(&ti)
16130 .arg(&tkvi)
16131 .arg(&scale)
16132 .arg(&cz)
16133 .arg(&wi);
16134 unsafe {
16135 b.launch(cfg)?;
16136 }
16137 Ok(())
16138 }
16139
16140 #[allow(clippy::too_many_arguments)]
16142 pub fn fa_prefill_w_arm(
16143 &self,
16144 q: &CudaSlice<f32>,
16145 k: &CudaSlice<f32>,
16146 v: &CudaSlice<f32>,
16147 o: &mut CudaSlice<f32>,
16148 head_dim: usize,
16149 n_head: usize,
16150 n_head_kv: usize,
16151 t: usize,
16152 t_kv: usize,
16153 scale: f32,
16154 causal: bool,
16155 window: usize,
16156 f32_stage: bool,
16157 floor: bool,
16158 ) -> Result<(), Box<dyn std::error::Error>> {
16159 const BLOCK_Q: usize = 64;
16160 const BK: usize = 32;
16161 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
16162 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16166 let p1 = !floor
16167 && !f32_stage
16168 && *P1_ON.get_or_init(|| {
16169 std::env::var("MEMRA_FAW_P1")
16170 .map(|v| v != "0")
16171 .unwrap_or(true)
16172 });
16173 let hp =
16174 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16175 if hp {
16176 const BLOCK_QH: usize = 32;
16177 let f = self.func("fa_prefill_w_bf16_p1h2");
16178 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
16179 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16180 f.set_attribute(
16181 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16182 shmem as i32,
16183 )?;
16184 let cfg = LaunchConfig {
16185 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
16186 block_dim: (32, 4, 1),
16187 shared_mem_bytes: shmem,
16188 };
16189 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16190 head_dim as i32,
16191 n_head as i32,
16192 n_head_kv as i32,
16193 t as i32,
16194 t_kv as i32,
16195 causal as i32,
16196 window as i32,
16197 );
16198 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16199 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16200 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
16201 let __s_b = self.gpu.stream();
16202 let mut b = __s_b.launch_builder(&f);
16203 b.arg(&qb)
16204 .arg(&kb)
16205 .arg(&vh)
16206 .arg(o)
16207 .arg(&hd)
16208 .arg(&nh)
16209 .arg(&nhkv)
16210 .arg(&ti)
16211 .arg(&tkvi)
16212 .arg(&scale)
16213 .arg(&cz)
16214 .arg(&wi);
16215 unsafe {
16216 b.launch(cfg)?;
16217 }
16218 return Ok(());
16219 }
16220 if p1 {
16221 let f = self.func("fa_prefill_w_bf16_p1");
16222 let shmem =
16223 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16224 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16225 f.set_attribute(
16226 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16227 shmem as i32,
16228 )?;
16229 let cfg = LaunchConfig {
16230 grid_dim: (
16231 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16232 n_head as u32,
16233 1,
16234 ),
16235 block_dim: (32, 4, 1),
16236 shared_mem_bytes: shmem,
16237 };
16238 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16239 head_dim as i32,
16240 n_head as i32,
16241 n_head_kv as i32,
16242 t as i32,
16243 t_kv as i32,
16244 causal as i32,
16245 window as i32,
16246 );
16247 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16248 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16249 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16250 let __s_b = self.gpu.stream();
16251 let mut b = __s_b.launch_builder(&f);
16252 b.arg(&qb)
16253 .arg(&kb)
16254 .arg(&vb)
16255 .arg(o)
16256 .arg(&hd)
16257 .arg(&nh)
16258 .arg(&nhkv)
16259 .arg(&ti)
16260 .arg(&tkvi)
16261 .arg(&scale)
16262 .arg(&cz)
16263 .arg(&wi);
16264 unsafe {
16265 b.launch(cfg)?;
16266 }
16267 return Ok(());
16268 }
16269 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16272 let g4 = !floor
16273 && !f32_stage
16274 && n_head_kv == 1
16275 && n_head % 4 == 0
16276 && *G4_ON.get_or_init(|| {
16277 std::env::var("MEMRA_FAW_G4")
16278 .map(|v| v != "0")
16279 .unwrap_or(true)
16280 });
16281 if g4 {
16282 const SP_M: usize = 16;
16283 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16286 let o2 = *O2_ON.get_or_init(|| {
16287 std::env::var("MEMRA_FAW_O2")
16288 .map(|v| v != "0")
16289 .unwrap_or(true)
16290 });
16291 let f = self.func(if o2 {
16292 "fa_prefill_w_bf16_g4o2"
16293 } else {
16294 "fa_prefill_w_bf16_g4"
16295 });
16296 let shmem = if o2 {
16297 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
16298 } else {
16299 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
16300 as u32
16301 };
16302 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16303 f.set_attribute(
16304 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16305 shmem as i32,
16306 )?;
16307 let cfg = LaunchConfig {
16308 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
16309 block_dim: (32, 4, 1),
16310 shared_mem_bytes: shmem,
16311 };
16312 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16313 head_dim as i32,
16314 n_head as i32,
16315 n_head_kv as i32,
16316 t as i32,
16317 t_kv as i32,
16318 causal as i32,
16319 window as i32,
16320 );
16321 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16322 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16323 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16324 let __s_b = self.gpu.stream();
16325 let mut b = __s_b.launch_builder(&f);
16326 b.arg(&qb)
16327 .arg(&kb)
16328 .arg(&vb)
16329 .arg(o)
16330 .arg(&hd)
16331 .arg(&nh)
16332 .arg(&nhkv)
16333 .arg(&ti)
16334 .arg(&tkvi)
16335 .arg(&scale)
16336 .arg(&cz)
16337 .arg(&wi);
16338 unsafe {
16339 b.launch(cfg)?;
16340 }
16341 return Ok(());
16342 }
16343 let f = self.func(if floor {
16344 "fa_prefill_w_f32"
16345 } else if f32_stage {
16346 "fa_prefill_w_f32_pp"
16347 } else {
16348 "fa_prefill_w_bf16_pp"
16349 });
16350 let shmem =
16351 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
16352 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16353 f.set_attribute(
16354 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16355 shmem as i32,
16356 )?;
16357 let cfg = LaunchConfig {
16358 grid_dim: (
16359 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16360 n_head as u32,
16361 1,
16362 ),
16363 block_dim: (32, 4, 1),
16364 shared_mem_bytes: shmem,
16365 };
16366 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
16367 head_dim as i32,
16368 n_head as i32,
16369 n_head_kv as i32,
16370 t as i32,
16371 t_kv as i32,
16372 causal as i32,
16373 window as i32,
16374 );
16375 if f32_stage {
16376 let __s_b = self.gpu.stream();
16377 let mut b = __s_b.launch_builder(&f);
16378 b.arg(q)
16379 .arg(k)
16380 .arg(v)
16381 .arg(o)
16382 .arg(&hd)
16383 .arg(&nh)
16384 .arg(&nhkv)
16385 .arg(&ti)
16386 .arg(&tkvi)
16387 .arg(&scale)
16388 .arg(&cz)
16389 .arg(&wi);
16390 unsafe {
16391 b.launch(cfg)?;
16392 }
16393 } else {
16394 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16395 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16396 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16397 let __s_b = self.gpu.stream();
16398 let mut b = __s_b.launch_builder(&f);
16399 b.arg(&qb)
16400 .arg(&kb)
16401 .arg(&vb)
16402 .arg(o)
16403 .arg(&hd)
16404 .arg(&nh)
16405 .arg(&nhkv)
16406 .arg(&ti)
16407 .arg(&tkvi)
16408 .arg(&scale)
16409 .arg(&cz)
16410 .arg(&wi);
16411 unsafe {
16412 b.launch(cfg)?;
16413 }
16414 }
16415 Ok(())
16416 }
16417
16418 #[allow(clippy::too_many_arguments)]
16422 pub fn fa_prefill_hd512(
16423 &self,
16424 q: &CudaSlice<f32>,
16425 k: &CudaSlice<f32>,
16426 v: &CudaSlice<f32>,
16427 o: &mut CudaSlice<f32>,
16428 head_dim: usize,
16429 n_head: usize,
16430 n_head_kv: usize,
16431 t: usize,
16432 t_kv: usize,
16433 scale: f32,
16434 causal: bool,
16435 ) -> Result<(), Box<dyn std::error::Error>> {
16436 if portable_mma_gated() {
16438 return self.sdpa_naive(
16439 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
16440 );
16441 }
16442 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16448 let f32_stage =
16449 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
16450 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16454 let sp = !f32_stage
16455 && *SP_ON.get_or_init(|| {
16456 std::env::var("MEMRA_FA512_SP")
16457 .map(|v| v != "0")
16458 .unwrap_or(true)
16459 });
16460 self.fa_prefill_hd512_arm(
16461 q,
16462 k,
16463 v,
16464 o,
16465 head_dim,
16466 n_head,
16467 n_head_kv,
16468 t,
16469 t_kv,
16470 scale,
16471 causal,
16472 f32_stage,
16473 sp,
16474 sp && fa_f16pv_on(),
16475 )
16476 }
16477
16478 #[allow(clippy::too_many_arguments)]
16480 pub fn fa_prefill_hd512_pre(
16481 &self,
16482 qb: &CudaSlice<u8>,
16483 kb: &CudaSlice<u8>,
16484 vb: &CudaSlice<u8>,
16485 o: &mut CudaSlice<f32>,
16486 head_dim: usize,
16487 n_head: usize,
16488 n_head_kv: usize,
16489 t: usize,
16490 t_kv: usize,
16491 scale: f32,
16492 causal: bool,
16493 v_f16: bool,
16494 ) -> Result<(), Box<dyn std::error::Error>> {
16495 debug_assert_eq!(head_dim, 512);
16496 const SP_M: usize = 16;
16497 const BKS: usize = 32;
16498 let f16pv = fa_f16pv_on();
16502 let nw = if f16pv { fa512_wide_warps() } else { 2 };
16503 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16504 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
16505 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
16506 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
16507 let n = t_kv * n_head_kv * head_dim;
16509 let need = n * 2;
16510 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
16511 *vguard = Some(self.alloc_uninit::<u8>(need)?);
16512 }
16513 let dst = vguard.as_mut().unwrap();
16514 self.bf16_to_f16_into(vb, n, dst)?;
16515 vguard.as_ref().unwrap()
16516 } else {
16517 vb
16518 };
16519 let f = self.func(if hp {
16520 "fa_prefill_bf16_hd512_sp16h2"
16521 } else {
16522 match (f16pv, nw) {
16523 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16524 (true, _) => "fa_prefill_bf16_hd512_sp16",
16525 _ => "fa_prefill_bf16_hd512_sp",
16526 }
16527 });
16528 let (nwarp, npart) = if hp {
16529 (4usize, 4usize)
16530 } else if nw > 2 {
16531 (nw, nw)
16532 } else {
16533 (2, 1)
16534 };
16535 let shmem = if hp {
16537 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
16538 as u32
16539 } else {
16540 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16541 + 4 * (npart * SP_M * BKS + SP_M)) as u32
16542 };
16543 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16544 f.set_attribute(
16545 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16546 shmem as i32,
16547 )?;
16548 let grid_y = if hp {
16549 (n_head / 2) as u32
16550 } else {
16551 n_head as u32
16552 };
16553 let cfg = LaunchConfig {
16554 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16555 block_dim: (32, nwarp as u32, 1),
16556 shared_mem_bytes: shmem,
16557 };
16558 let (hd, nh, nhkv, ti, tkvi, cz) = (
16559 head_dim as i32,
16560 n_head as i32,
16561 n_head_kv as i32,
16562 t as i32,
16563 t_kv as i32,
16564 causal as i32,
16565 );
16566 let __s_b = self.gpu.stream();
16567 let mut b = __s_b.launch_builder(&f);
16568 b.arg(qb)
16569 .arg(kb)
16570 .arg(vref)
16571 .arg(o)
16572 .arg(&hd)
16573 .arg(&nh)
16574 .arg(&nhkv)
16575 .arg(&ti)
16576 .arg(&tkvi)
16577 .arg(&scale)
16578 .arg(&cz);
16579 unsafe {
16580 b.launch(cfg)?;
16581 }
16582 Ok(())
16583 }
16584
16585 #[allow(clippy::too_many_arguments)]
16588 pub fn fa_prefill_hd512_arm(
16589 &self,
16590 q: &CudaSlice<f32>,
16591 k: &CudaSlice<f32>,
16592 v: &CudaSlice<f32>,
16593 o: &mut CudaSlice<f32>,
16594 head_dim: usize,
16595 n_head: usize,
16596 n_head_kv: usize,
16597 t: usize,
16598 t_kv: usize,
16599 scale: f32,
16600 causal: bool,
16601 f32_stage: bool,
16602 sp: bool,
16603 f16pv: bool,
16604 ) -> Result<(), Box<dyn std::error::Error>> {
16605 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
16606 if sp && !f32_stage {
16607 const SP_M: usize = 16;
16611 const BKS: usize = 32;
16612 let nw = if f16pv { fa512_wide_warps() } else { 2 };
16613 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
16614 let f = self.func(if hp {
16615 "fa_prefill_bf16_hd512_sp16h2"
16616 } else {
16617 match (f16pv, nw) {
16618 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
16619 (true, _) => "fa_prefill_bf16_hd512_sp16",
16620 _ => "fa_prefill_bf16_hd512_sp",
16621 }
16622 });
16623 let (nwarp, npart) = if hp {
16624 (4usize, 4usize)
16625 } else if nw > 2 {
16626 (nw, nw)
16627 } else {
16628 (2, 1)
16629 };
16630 let shmem = if hp {
16631 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
16632 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
16633 } else {
16634 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
16635 + 4 * (npart * SP_M * BKS + SP_M)) as u32
16636 };
16637 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16638 f.set_attribute(
16639 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16640 shmem as i32,
16641 )?;
16642 let grid_y = if hp {
16643 (n_head / 2) as u32
16644 } else {
16645 n_head as u32
16646 };
16647 let cfg = LaunchConfig {
16648 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
16649 block_dim: (32, nwarp as u32, 1),
16650 shared_mem_bytes: shmem,
16651 };
16652 let (hd, nh, nhkv, ti, tkvi, cz) = (
16653 head_dim as i32,
16654 n_head as i32,
16655 n_head_kv as i32,
16656 t as i32,
16657 t_kv as i32,
16658 causal as i32,
16659 );
16660 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16661 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16662 let vb = if f16pv {
16663 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
16664 } else {
16665 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
16666 };
16667 let __s_b = self.gpu.stream();
16668 let mut b = __s_b.launch_builder(&f);
16669 b.arg(&qb)
16670 .arg(&kb)
16671 .arg(&vb)
16672 .arg(o)
16673 .arg(&hd)
16674 .arg(&nh)
16675 .arg(&nhkv)
16676 .arg(&ti)
16677 .arg(&tkvi)
16678 .arg(&scale)
16679 .arg(&cz);
16680 unsafe {
16681 b.launch(cfg)?;
16682 }
16683 return Ok(());
16684 }
16685 const BLOCK_Q: usize = 32;
16686 const BK: usize = 32;
16687 const HALF: usize = 256;
16688 let f = self.func(if f32_stage {
16689 "fa_prefill_f32_hd512"
16690 } else {
16691 "fa_prefill_bf16_hd512"
16692 });
16693 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
16695 + 4 * BLOCK_Q) as u32;
16696 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16697 f.set_attribute(
16698 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16699 shmem as i32,
16700 )?;
16701 let cfg = LaunchConfig {
16702 grid_dim: (
16703 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
16704 n_head as u32,
16705 2,
16706 ),
16707 block_dim: (32, 2, 1),
16708 shared_mem_bytes: shmem,
16709 };
16710 let (hd, nh, nhkv, ti, tkvi, cz) = (
16711 head_dim as i32,
16712 n_head as i32,
16713 n_head_kv as i32,
16714 t as i32,
16715 t_kv as i32,
16716 causal as i32,
16717 );
16718 if f32_stage {
16719 let __s_b = self.gpu.stream();
16720 let mut b = __s_b.launch_builder(&f);
16721 b.arg(q)
16722 .arg(k)
16723 .arg(v)
16724 .arg(o)
16725 .arg(&hd)
16726 .arg(&nh)
16727 .arg(&nhkv)
16728 .arg(&ti)
16729 .arg(&tkvi)
16730 .arg(&scale)
16731 .arg(&cz);
16732 unsafe {
16733 b.launch(cfg)?;
16734 }
16735 } else {
16736 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
16737 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
16738 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
16739 let __s_b = self.gpu.stream();
16740 let mut b = __s_b.launch_builder(&f);
16741 b.arg(&qb)
16742 .arg(&kb)
16743 .arg(&vb)
16744 .arg(o)
16745 .arg(&hd)
16746 .arg(&nh)
16747 .arg(&nhkv)
16748 .arg(&ti)
16749 .arg(&tkvi)
16750 .arg(&scale)
16751 .arg(&cz);
16752 unsafe {
16753 b.launch(cfg)?;
16754 }
16755 }
16756 Ok(())
16757 }
16758
16759 #[allow(clippy::too_many_arguments)]
16763 pub fn rope_neox2_bf16e(
16764 &self,
16765 q: &mut CudaSlice<f32>,
16766 k: &mut CudaSlice<f32>,
16767 qb: &mut CudaSlice<u8>,
16768 kb: &mut CudaSlice<u8>,
16769 pos: &CudaSlice<i32>,
16770 head_dim: usize,
16771 n_dims: usize,
16772 nh_q: usize,
16773 nh_k: usize,
16774 n_tokens: usize,
16775 base: f32,
16776 freq_scale: f32,
16777 ff: Option<&CudaSlice<f32>>,
16778 ) -> Result<(), Box<dyn std::error::Error>> {
16779 let f = self.func("rope_neox2_bf16e_f32");
16780 let rows = ((nh_q + nh_k) * n_tokens) as u32;
16781 let cfg = LaunchConfig {
16782 grid_dim: (rows, 1, 1),
16783 block_dim: ((head_dim / 2) as u32, 1, 1),
16784 shared_mem_bytes: 0,
16785 };
16786 let theta_scale = base.powf(-2.0 / n_dims as f32);
16787 let (hd, nd, nhq, nhk, nt) = (
16788 head_dim as i32,
16789 n_dims as i32,
16790 nh_q as i32,
16791 nh_k as i32,
16792 n_tokens as i32,
16793 );
16794 let __s_b = self.gpu.stream();
16795 let mut b = __s_b.launch_builder(&f);
16796 match ff {
16797 Some(t) => {
16798 b.arg(&mut *q)
16799 .arg(&mut *k)
16800 .arg(&mut *qb)
16801 .arg(&mut *kb)
16802 .arg(pos)
16803 .arg(&hd)
16804 .arg(&nd)
16805 .arg(&nhq)
16806 .arg(&nhk)
16807 .arg(&nt)
16808 .arg(&theta_scale)
16809 .arg(&freq_scale)
16810 .arg(t);
16811 unsafe {
16812 b.launch(cfg)?;
16813 }
16814 }
16815 None => {
16816 let null: u64 = 0;
16817 b.arg(&mut *q)
16818 .arg(&mut *k)
16819 .arg(&mut *qb)
16820 .arg(&mut *kb)
16821 .arg(pos)
16822 .arg(&hd)
16823 .arg(&nd)
16824 .arg(&nhq)
16825 .arg(&nhk)
16826 .arg(&nt)
16827 .arg(&theta_scale)
16828 .arg(&freq_scale)
16829 .arg(&null);
16830 unsafe {
16831 b.launch(cfg)?;
16832 }
16833 }
16834 }
16835 Ok(())
16836 }
16837
16838 pub fn f32_to_bf16(
16841 &self,
16842 x: &CudaSlice<f32>,
16843 n: usize,
16844 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16845 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
16846 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16847 let f = self.func("f32_to_bf16_flat");
16848 let n_i = n as i64;
16849 let cfg = LaunchConfig {
16850 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16851 block_dim: (256, 1, 1),
16852 shared_mem_bytes: 0,
16853 };
16854 let __s_b = self.gpu.stream();
16855 let mut b = __s_b.launch_builder(&f);
16856 b.arg(x).arg(&mut y).arg(&n_i);
16857 unsafe {
16858 b.launch(cfg)?;
16859 }
16860 Ok(y)
16861 }
16862
16863 pub fn f32_to_f16(
16864 &self,
16865 x: &CudaSlice<f32>,
16866 n: usize,
16867 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16868 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
16869 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16870 let f = self.func("f32_to_f16_flat");
16871 let n_i = n as i64;
16872 let cfg = LaunchConfig {
16873 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
16874 block_dim: (256, 1, 1),
16875 shared_mem_bytes: 0,
16876 };
16877 let __s_b = self.gpu.stream();
16878 let mut b = __s_b.launch_builder(&f);
16879 b.arg(x).arg(&mut y).arg(&n_i);
16880 unsafe {
16881 b.launch(cfg)?;
16882 }
16883 Ok(y)
16884 }
16885
16886 pub fn bf16_to_f16(
16888 &self,
16889 xb: &CudaSlice<u8>,
16890 n: usize,
16891 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
16892 let mut y = self.alloc_uninit::<u8>(n * 2)?;
16893 self.bf16_to_f16_into(xb, n, &mut y)?;
16894 Ok(y)
16895 }
16896
16897 pub fn bf16_to_f16_into(
16899 &self,
16900 xb: &CudaSlice<u8>,
16901 n: usize,
16902 y: &mut CudaSlice<u8>,
16903 ) -> Result<(), Box<dyn std::error::Error>> {
16904 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
16905 assert!(y.len() >= n * 2);
16906 let f = self.func("bf16_to_f16_flat");
16907 let n2 = (n / 2) as i64;
16908 let cfg = LaunchConfig {
16909 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
16910 block_dim: (256, 1, 1),
16911 shared_mem_bytes: 0,
16912 };
16913 let __s_b = self.gpu.stream();
16914 let mut b = __s_b.launch_builder(&f);
16915 b.arg(xb).arg(y).arg(&n2);
16916 unsafe {
16917 b.launch(cfg)?;
16918 }
16919 Ok(())
16920 }
16921
16922 #[allow(clippy::too_many_arguments)]
16927 pub fn fa_prefill_vl8(
16928 &self,
16929 seqs: &[FaSeqVl],
16930 head_dim: usize,
16931 n_head: usize,
16932 n_head_kv: usize,
16933 scale: f32,
16934 ) -> Result<(), Box<dyn std::error::Error>> {
16935 const BK: usize = 32;
16936 let b = seqs.len();
16937 assert!(b >= 1 && b <= 8);
16938 let mut packed = [FaSeqVl::default(); 8];
16939 packed[..b].copy_from_slice(seqs);
16940 let v = FaVl8(packed);
16941 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
16942 let ept = (n_head_kv * head_dim) as i32;
16943 {
16944 let f = self.func("fa_mirror_vl");
16945 let max_n = (max_t as i64) * ept as i64;
16946 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
16947 for which in 0..2i32 {
16948 let cfg = LaunchConfig {
16949 grid_dim: (blocks, 1, b as u32),
16950 block_dim: (256, 1, 1),
16951 shared_mem_bytes: 0,
16952 };
16953 let __s_lb = self.gpu.stream();
16954 let mut lb = __s_lb.launch_builder(&f);
16955 lb.arg(&v).arg(&ept).arg(&which);
16956 unsafe {
16957 lb.launch(cfg)?;
16958 }
16959 }
16960 }
16961 let hd_sfx = fa_hd_suffix(head_dim)?;
16962 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
16963 let block_q = 64usize;
16964 let kv_stages = 2usize;
16965 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
16966 + 4 * (block_q * BK + 2 * block_q)) as u32;
16967 use cudarc::driver::sys::CUfunction_attribute_enum as A;
16968 f.set_attribute(
16969 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
16970 shmem as i32,
16971 )?;
16972 let cfg = LaunchConfig {
16973 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
16974 block_dim: (32, 4, 1),
16975 shared_mem_bytes: shmem,
16976 };
16977 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
16978 let __s_lb = self.gpu.stream();
16979 let mut lb = __s_lb.launch_builder(&f);
16980 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
16981 unsafe {
16982 lb.launch(cfg)?;
16983 }
16984 Ok(())
16985 }
16986
16987 #[allow(clippy::too_many_arguments)]
16991 pub fn attn_pre_vl8(
16992 &self,
16993 seqs: &[AttnPreVl],
16994 wq: &CudaSlice<f32>,
16995 wk: &CudaSlice<f32>,
16996 head_dim: usize,
16997 rope_dims: usize,
16998 n_head: usize,
16999 n_head_kv: usize,
17000 eps: f32,
17001 freq_base: f32,
17002 freq_scale: f32,
17003 kv_dim_k: usize,
17004 kv_dim_v: usize,
17005 k_tok_bytes: usize,
17006 v_tok_bytes: usize,
17007 ) -> Result<(), Box<dyn std::error::Error>> {
17008 let b = seqs.len();
17009 assert!(b >= 1 && b <= 8);
17010 let mut packed = [AttnPreVl::default(); 8];
17011 packed[..b].copy_from_slice(seqs);
17012 let v = AttnPreVl8(packed);
17013 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
17014 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
17015 {
17016 let f = self.func("q_gate_split_vl");
17017 let n = max_t * (n_head * head_dim) as u32;
17018 let cfg = LaunchConfig {
17019 grid_dim: (n.div_ceil(256), 1, b as u32),
17020 block_dim: (256, 1, 1),
17021 shared_mem_bytes: 0,
17022 };
17023 let __s_lb = self.gpu.stream();
17024 let mut lb = __s_lb.launch_builder(&f);
17025 lb.arg(&v).arg(&hd).arg(&nh);
17026 unsafe {
17027 lb.launch(cfg)?;
17028 }
17029 }
17030 {
17031 let f = self.func("attn_rms_vl");
17032 let cfg = LaunchConfig {
17033 grid_dim: (max_t * n_head as u32, 2, b as u32),
17034 block_dim: (rms_block(), 1, 1),
17035 shared_mem_bytes: 0,
17036 };
17037 let __s_lb = self.gpu.stream();
17038 let mut lb = __s_lb.launch_builder(&f);
17039 lb.arg(&v)
17040 .arg(wq)
17041 .arg(wk)
17042 .arg(&hd)
17043 .arg(&nh)
17044 .arg(&nhkv)
17045 .arg(&eps);
17046 unsafe {
17047 lb.launch(cfg)?;
17048 }
17049 }
17050 {
17051 let f = self.func("attn_rope_vl");
17052 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
17053 let nd = rope_dims as i32;
17054 let cfg = LaunchConfig {
17055 grid_dim: (max_t * n_head as u32, 2, b as u32),
17056 block_dim: ((head_dim / 2) as u32, 1, 1),
17057 shared_mem_bytes: 0,
17058 };
17059 let __s_lb = self.gpu.stream();
17060 let mut lb = __s_lb.launch_builder(&f);
17061 lb.arg(&v)
17062 .arg(&hd)
17063 .arg(&nd)
17064 .arg(&nh)
17065 .arg(&nhkv)
17066 .arg(&theta_scale)
17067 .arg(&freq_scale);
17068 unsafe {
17069 lb.launch(cfg)?;
17070 }
17071 }
17072 {
17073 let f = self.func("append_kv_vl");
17074 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
17075 let cfg = LaunchConfig {
17076 grid_dim: (nblk, max_t, b as u32),
17077 block_dim: (32, 1, 1),
17078 shared_mem_bytes: 0,
17079 };
17080 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17081 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17082 let __s_lb = self.gpu.stream();
17083 let mut lb = __s_lb.launch_builder(&f);
17084 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
17085 unsafe {
17086 lb.launch(cfg)?;
17087 }
17088 }
17089 Ok(())
17090 }
17091
17092 pub fn fa_prefill_view(
17097 &self,
17098 q: &CudaSlice<f32>,
17099 k: &cudarc::driver::CudaView<u8>,
17100 v: &cudarc::driver::CudaView<u8>,
17101 o: &mut CudaSlice<f32>,
17102 head_dim: usize,
17103 n_head: usize,
17104 n_head_kv: usize,
17105 t: usize,
17106 t_kv: usize,
17107 scale: f32,
17108 causal: bool,
17109 k_tok_bytes: usize,
17110 v_tok_bytes: usize,
17111 g: bool,
17112 ) -> Result<(), Box<dyn std::error::Error>> {
17113 if portable_mma_gated() {
17114 return self.sdpa_naive_quantized_view(
17115 q,
17116 k,
17117 v,
17118 o,
17119 head_dim,
17120 n_head,
17121 n_head_kv,
17122 t,
17123 t_kv,
17124 scale,
17125 causal,
17126 k_tok_bytes,
17127 v_tok_bytes,
17128 );
17129 }
17130 const BLOCK_Q: usize = 64;
17131 const BK: usize = 32;
17132 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
17135 let f = if g {
17136 self.func_g(&name)
17137 } else {
17138 self.func(&name)
17139 };
17140 let shmem =
17141 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
17142 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17143 f.set_attribute(
17144 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17145 shmem as i32,
17146 )?;
17147 let cfg = LaunchConfig {
17148 grid_dim: (
17149 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17150 n_head as u32,
17151 1,
17152 ),
17153 block_dim: (32, 4, 1),
17154 shared_mem_bytes: shmem,
17155 };
17156 let (hd, nh, nhkv, ti, tkvi, cz) = (
17157 head_dim as i32,
17158 n_head as i32,
17159 n_head_kv as i32,
17160 t as i32,
17161 t_kv as i32,
17162 causal as i32,
17163 );
17164 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17165 let __s_b = self.gpu.stream();
17166 let mut b = __s_b.launch_builder(&f);
17167 b.arg(q)
17168 .arg(k)
17169 .arg(v)
17170 .arg(o)
17171 .arg(&hd)
17172 .arg(&nh)
17173 .arg(&nhkv)
17174 .arg(&ti)
17175 .arg(&tkvi)
17176 .arg(&scale)
17177 .arg(&cz)
17178 .arg(&ktb)
17179 .arg(&vtb);
17180 unsafe {
17181 b.launch(cfg)?;
17182 }
17183 Ok(())
17184 }
17185
17186 #[allow(clippy::too_many_arguments)]
17196 pub fn fa_prefill_view_ws(
17197 &self,
17198 q: &CudaSlice<f32>,
17199 k: &cudarc::driver::CudaView<u8>,
17200 v: &cudarc::driver::CudaView<u8>,
17201 o: &mut CudaSlice<f32>,
17202 head_dim: usize,
17203 n_head: usize,
17204 n_head_kv: usize,
17205 t: usize,
17206 t_kv: usize,
17207 scale: f32,
17208 causal: bool,
17209 k_tok_bytes: usize,
17210 v_tok_bytes: usize,
17211 g: bool,
17212 ) -> Result<(), Box<dyn std::error::Error>> {
17213 if portable_mma_gated() {
17214 return self.sdpa_naive_quantized_view(
17215 q,
17216 k,
17217 v,
17218 o,
17219 head_dim,
17220 n_head,
17221 n_head_kv,
17222 t,
17223 t_kv,
17224 scale,
17225 causal,
17226 k_tok_bytes,
17227 v_tok_bytes,
17228 );
17229 }
17230 const BLOCK_Q: usize = 64;
17231 const BK: usize = 32;
17232 let kv_dim_k = n_head_kv * head_dim;
17233 let kv_dim_v = n_head_kv * head_dim;
17234 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17236 let mut guard = self.prime_deqw_ws.lock().unwrap();
17238 let need_grow = match guard.as_ref() {
17239 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17240 None => true,
17241 };
17242 if need_grow {
17243 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17244 let (ck, cv) = guard
17245 .as_ref()
17246 .map(|(a, b)| (a.len(), b.len()))
17247 .unwrap_or((0, 0));
17248 *guard = Some((
17249 self.alloc_u8(grow(ck, k_ws_bytes))?,
17250 self.alloc_u8(grow(cv, v_ws_bytes))?,
17251 ));
17252 }
17253 let (kw, vw) = guard.as_mut().unwrap();
17254 {
17256 let f = if g {
17258 self.func_g("fa_dequant_kv_ws_bf16")
17259 } else {
17260 self.func("fa_dequant_kv_ws_bf16")
17261 };
17262 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17263 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17264 let cfg = LaunchConfig {
17265 grid_dim: (nblk.max(1), 1, 1),
17266 block_dim: (256, 1, 1),
17267 shared_mem_bytes: 0,
17268 };
17269 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17270 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17271 let __s_b = self.gpu.stream();
17272 let mut b = __s_b.launch_builder(&f);
17273 b.arg(k)
17274 .arg(v)
17275 .arg(&mut *kw)
17276 .arg(&mut *vw)
17277 .arg(&kdk)
17278 .arg(&kdv)
17279 .arg(&tkvi)
17280 .arg(&ktb)
17281 .arg(&vtb);
17282 unsafe {
17283 b.launch(cfg)?;
17284 }
17285 }
17286 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17294 .map(|v| v != "0")
17295 .unwrap_or(true);
17296 {
17297 let hd_sfx = fa_hd_suffix(head_dim)?;
17298 let f = self.func(&format!(
17299 "fa_prefill_qw{}{hd_sfx}",
17300 if db { "_db" } else { "" }
17301 ));
17302 let shmem = if db {
17303 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17305 } else {
17306 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17307 };
17308 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17309 f.set_attribute(
17310 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17311 shmem as i32,
17312 )?;
17313 let cfg = LaunchConfig {
17314 grid_dim: (
17315 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17316 n_head as u32,
17317 1,
17318 ),
17319 block_dim: (32, 4, 1),
17320 shared_mem_bytes: shmem,
17321 };
17322 let (hd, nh, nhkv, ti, tkvi, cz) = (
17323 head_dim as i32,
17324 n_head as i32,
17325 n_head_kv as i32,
17326 t as i32,
17327 t_kv as i32,
17328 causal as i32,
17329 );
17330 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
17331 let __s_b = self.gpu.stream();
17332 let mut b = __s_b.launch_builder(&f);
17333 b.arg(q)
17334 .arg(&*kw)
17335 .arg(&*vw)
17336 .arg(o)
17337 .arg(&hd)
17338 .arg(&nh)
17339 .arg(&nhkv)
17340 .arg(&ti)
17341 .arg(&tkvi)
17342 .arg(&scale)
17343 .arg(&cz)
17344 .arg(&kdk)
17345 .arg(&kdv);
17346 unsafe {
17347 b.launch(cfg)?;
17348 }
17349 }
17350 Ok(())
17351 }
17352
17353 #[allow(clippy::too_many_arguments)]
17369 pub fn fa_prefill_view_ws_w_hd128(
17370 &self,
17371 q: &CudaSlice<f32>,
17372 k: &cudarc::driver::CudaView<u8>,
17373 v: &cudarc::driver::CudaView<u8>,
17374 o: &mut CudaSlice<f32>,
17375 head_dim: usize,
17376 n_head: usize,
17377 n_head_kv: usize,
17378 t: usize,
17379 t_kv: usize,
17380 scale: f32,
17381 causal: bool,
17382 window: usize,
17383 k_tok_bytes: usize,
17384 v_tok_bytes: usize,
17385 ) -> Result<(), Box<dyn std::error::Error>> {
17386 assert_eq!(
17387 head_dim, 128,
17388 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
17389 );
17390 if portable_mma_gated() {
17391 return self.sdpa_naive_w_quantized_view(
17392 q,
17393 k,
17394 v,
17395 o,
17396 head_dim,
17397 n_head,
17398 n_head_kv,
17399 t,
17400 t_kv,
17401 scale,
17402 causal,
17403 window,
17404 k_tok_bytes,
17405 v_tok_bytes,
17406 );
17407 }
17408 const BLOCK_Q: usize = 64;
17409 const BK: usize = 32;
17410 let kv_dim_k = n_head_kv * head_dim;
17411 let kv_dim_v = n_head_kv * head_dim;
17412 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
17414 let mut guard = self.prime_deqw_ws.lock().unwrap();
17415 let need_grow = match guard.as_ref() {
17416 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
17417 None => true,
17418 };
17419 if need_grow {
17420 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
17421 let (ck, cv) = guard
17422 .as_ref()
17423 .map(|(a, b)| (a.len(), b.len()))
17424 .unwrap_or((0, 0));
17425 *guard = Some((
17426 self.alloc_u8(grow(ck, k_ws_bytes))?,
17427 self.alloc_u8(grow(cv, v_ws_bytes))?,
17428 ));
17429 }
17430 let (kw, vw) = guard.as_mut().unwrap();
17431 {
17434 let f = self.func("fa_dequant_kv_ws_bf16");
17435 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
17436 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
17437 let cfg = LaunchConfig {
17438 grid_dim: (nblk.max(1), 1, 1),
17439 block_dim: (256, 1, 1),
17440 shared_mem_bytes: 0,
17441 };
17442 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
17443 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17444 let __s_b = self.gpu.stream();
17445 let mut b = __s_b.launch_builder(&f);
17446 b.arg(k)
17447 .arg(v)
17448 .arg(&mut *kw)
17449 .arg(&mut *vw)
17450 .arg(&kdk)
17451 .arg(&kdv)
17452 .arg(&tkvi)
17453 .arg(&ktb)
17454 .arg(&vtb);
17455 unsafe {
17456 b.launch(cfg)?;
17457 }
17458 }
17459 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
17461 .map(|v| v != "0")
17462 .unwrap_or(true);
17463 {
17464 let f = self.func(if db {
17465 "fa_prefill_qw_db_w_hd128"
17466 } else {
17467 "fa_prefill_qw_w_hd128"
17468 });
17469 let shmem = if db {
17470 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
17471 } else {
17472 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
17473 };
17474 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17475 f.set_attribute(
17476 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17477 shmem as i32,
17478 )?;
17479 let cfg = LaunchConfig {
17480 grid_dim: (
17481 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
17482 n_head as u32,
17483 1,
17484 ),
17485 block_dim: (32, 4, 1),
17486 shared_mem_bytes: shmem,
17487 };
17488 let (hd, nh, nhkv, ti, tkvi, cz) = (
17489 head_dim as i32,
17490 n_head as i32,
17491 n_head_kv as i32,
17492 t as i32,
17493 t_kv as i32,
17494 causal as i32,
17495 );
17496 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
17497 let __s_b = self.gpu.stream();
17498 let mut b = __s_b.launch_builder(&f);
17499 b.arg(q)
17500 .arg(&*kw)
17501 .arg(&*vw)
17502 .arg(o)
17503 .arg(&hd)
17504 .arg(&nh)
17505 .arg(&nhkv)
17506 .arg(&ti)
17507 .arg(&tkvi)
17508 .arg(&scale)
17509 .arg(&cz)
17510 .arg(&kdk)
17511 .arg(&kdv)
17512 .arg(&wnd);
17513 unsafe {
17514 b.launch(cfg)?;
17515 }
17516 }
17517 Ok(())
17518 }
17519
17520 pub fn fa_decode(
17524 &self,
17525 q: &CudaSlice<f32>,
17526 k: &cudarc::driver::CudaView<u8>,
17527 v: &cudarc::driver::CudaView<u8>,
17528 o: &mut CudaSlice<f32>,
17529 head_dim: usize,
17530 n_head: usize,
17531 n_head_kv: usize,
17532 t_kv: usize,
17533 scale: f32,
17534 k_tok_bytes: usize,
17535 v_tok_bytes: usize,
17536 ) -> Result<(), Box<dyn std::error::Error>> {
17537 self.fa_decode_kvmod(
17538 q,
17539 k,
17540 v,
17541 o,
17542 head_dim,
17543 n_head,
17544 n_head_kv,
17545 t_kv,
17546 scale,
17547 k_tok_bytes,
17548 v_tok_bytes,
17549 false,
17550 )
17551 }
17552
17553 #[allow(clippy::too_many_arguments)]
17557 #[allow(clippy::too_many_arguments)]
17561 #[allow(clippy::too_many_arguments)]
17562 fn fa_decode_scalar_unified(
17563 &self,
17564 q: &cudarc::driver::CudaView<f32>,
17565 k: &cudarc::driver::CudaView<u8>,
17566 v: &cudarc::driver::CudaView<u8>,
17567 o: &mut cudarc::driver::CudaViewMut<f32>,
17568 head_dim: usize,
17569 n_head: usize,
17570 n_head_kv: usize,
17571 t_kv_host: usize,
17572 t_kv_dev: Option<&CudaSlice<i32>>,
17573 scale: f32,
17574 n_splits: usize,
17575 split_keys: usize,
17576 k_tok_bytes: usize,
17577 v_tok_bytes: usize,
17578 g: bool,
17579 part_o: &mut CudaSlice<f32>,
17580 part_m: &mut CudaSlice<f32>,
17581 part_l: &mut CudaSlice<f32>,
17582 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
17583 ) -> Result<(), Box<dyn std::error::Error>> {
17584 let f = if g {
17585 self.func_g("fa_decode_f32")
17586 } else {
17587 self.fa_func("fa_decode_f32", head_dim)
17588 };
17589 let cfg = LaunchConfig {
17590 grid_dim: (n_head as u32, n_splits as u32, 1),
17591 block_dim: (head_dim as u32, 1, 1),
17592 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
17593 };
17594 let (hd, nh, nhkv, nsp) = (
17595 head_dim as i32,
17596 n_head as i32,
17597 n_head_kv as i32,
17598 n_splits as i32,
17599 );
17600 let (ktb, vtb, tkvi, ski) = (
17601 k_tok_bytes as i64,
17602 v_tok_bytes as i64,
17603 t_kv_host as i32,
17604 split_keys as i32,
17605 );
17606 let __s_b = self.gpu.stream();
17607 let mut b = __s_b.launch_builder(&f);
17608 match t_kv_dev {
17609 Some(d) => {
17610 b.arg(q)
17611 .arg(k)
17612 .arg(v)
17613 .arg(&mut *part_o)
17614 .arg(&mut *part_m)
17615 .arg(&mut *part_l)
17616 .arg(&hd)
17617 .arg(&nh)
17618 .arg(&nhkv)
17619 .arg(&tkvi)
17620 .arg(d)
17621 .arg(&scale)
17622 .arg(&nsp)
17623 .arg(&ski)
17624 .arg(&ktb)
17625 .arg(&vtb);
17626 unsafe {
17627 b.launch(cfg)?;
17628 }
17629 }
17630 None => {
17631 let null: u64 = 0;
17632 b.arg(q)
17633 .arg(k)
17634 .arg(v)
17635 .arg(&mut *part_o)
17636 .arg(&mut *part_m)
17637 .arg(&mut *part_l)
17638 .arg(&hd)
17639 .arg(&nh)
17640 .arg(&nhkv)
17641 .arg(&tkvi)
17642 .arg(&null)
17643 .arg(&scale)
17644 .arg(&nsp)
17645 .arg(&ski)
17646 .arg(&ktb)
17647 .arg(&vtb);
17648 unsafe {
17649 b.launch(cfg)?;
17650 }
17651 }
17652 }
17653 let cfg2 = LaunchConfig {
17654 grid_dim: (n_head as u32, 1, 1),
17655 block_dim: (head_dim as u32, 1, 1),
17656 shared_mem_bytes: 0,
17657 };
17658 if let Some((oq, od)) = q8_out {
17659 let fc = if g {
17661 self.func_g("fa_decode_combine_q8_1")
17662 } else {
17663 self.fa_func("fa_decode_combine_q8_1", head_dim)
17664 };
17665 let __s_b2 = self.gpu.stream();
17666 let mut b2 = __s_b2.launch_builder(&fc);
17667 b2.arg(&*part_o)
17668 .arg(&*part_m)
17669 .arg(&*part_l)
17670 .arg(oq)
17671 .arg(od)
17672 .arg(&hd)
17673 .arg(&nh)
17674 .arg(&nsp);
17675 unsafe {
17676 b2.launch(cfg2)?;
17677 }
17678 return Ok(());
17679 }
17680 let fc = if g {
17681 self.func_g("fa_decode_combine_f32")
17682 } else {
17683 self.fa_func("fa_decode_combine_f32", head_dim)
17684 };
17685 let __s_b2 = self.gpu.stream();
17686 let mut b2 = __s_b2.launch_builder(&fc);
17687 b2.arg(&*part_o)
17688 .arg(&*part_m)
17689 .arg(&*part_l)
17690 .arg(o)
17691 .arg(&hd)
17692 .arg(&nh)
17693 .arg(&nsp);
17694 unsafe {
17695 b2.launch(cfg2)?;
17696 }
17697 Ok(())
17698 }
17699
17700 pub fn fa_decode_kvmod(
17701 &self,
17702 q: &CudaSlice<f32>,
17703 k: &cudarc::driver::CudaView<u8>,
17704 v: &cudarc::driver::CudaView<u8>,
17705 o: &mut CudaSlice<f32>,
17706 head_dim: usize,
17707 n_head: usize,
17708 n_head_kv: usize,
17709 t_kv: usize,
17710 scale: f32,
17711 k_tok_bytes: usize,
17712 v_tok_bytes: usize,
17713 g: bool,
17714 ) -> Result<(), Box<dyn std::error::Error>> {
17715 let q_view = q.as_view();
17716 let mut o_view = o.as_view_mut();
17717 self.fa_decode_kvmod_view(
17718 &q_view,
17719 k,
17720 v,
17721 &mut o_view,
17722 head_dim,
17723 n_head,
17724 n_head_kv,
17725 t_kv,
17726 scale,
17727 k_tok_bytes,
17728 v_tok_bytes,
17729 g,
17730 )
17731 }
17732
17733 #[allow(clippy::too_many_arguments)]
17738 pub fn fa_decode_kvmod_view(
17739 &self,
17740 q: &cudarc::driver::CudaView<f32>,
17741 k: &cudarc::driver::CudaView<u8>,
17742 v: &cudarc::driver::CudaView<u8>,
17743 o: &mut cudarc::driver::CudaViewMut<f32>,
17744 head_dim: usize,
17745 n_head: usize,
17746 n_head_kv: usize,
17747 t_kv: usize,
17748 scale: f32,
17749 k_tok_bytes: usize,
17750 v_tok_bytes: usize,
17751 g: bool,
17752 ) -> Result<(), Box<dyn std::error::Error>> {
17753 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
17774 if g && head_dim == 256 && !fa_v4_at(t_kv) {
17778 fa_vec = false;
17779 }
17780 let sp = fa_split_keys(t_kv, n_head_kv);
17781 let n_splits = if fa_vec {
17782 ((t_kv + sp - 1) / sp).max(1)
17783 } else {
17784 ((t_kv + 255) / 256).max(1)
17785 };
17786 let o_len = n_head * n_splits * head_dim;
17787 let ml_len = n_head * n_splits;
17788 let mut part_guard = self.fa_part_pool.lock().unwrap();
17789 if part_guard
17790 .as_ref()
17791 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
17792 .unwrap_or(true)
17793 {
17794 let old = part_guard.take();
17805 let (co, cm) = old
17806 .as_ref()
17807 .map(|pp| (pp.0.len(), pp.1.len()))
17808 .unwrap_or((0, 0));
17809 if let Some(old) = old {
17810 self.fa_part_retired.lock().unwrap().push(old);
17811 }
17812 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
17813 eprintln!(
17814 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
17815 co, o_len, cm, ml_len
17816 );
17817 }
17818 *part_guard = Some((
17819 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
17820 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17821 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
17822 ));
17823 }
17824 let pg = part_guard.as_mut().unwrap();
17825 self.gpu
17826 .stream()
17827 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
17828 self.gpu
17829 .stream()
17830 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
17831 self.gpu
17832 .stream()
17833 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
17834 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
17835 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
17836 let (hd, nh, nhkv, tkvi, nsp) = (
17837 head_dim as i32,
17838 n_head as i32,
17839 n_head_kv as i32,
17840 t_kv as i32,
17841 n_splits as i32,
17842 );
17843 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17844 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
17848 let fa512_min = fa512_min_tkv();
17853 let deep = fa_vec
17856 && head_dim == 256
17857 && fa_v4_at(t_kv)
17858 && !g
17859 && fa_deep_at(t_kv)
17860 && !matches!(fa_v4_mode(), "noB3" | "stage");
17861 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
17862 let gqa = (n_head / n_head_kv).max(1) as u32;
17865 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
17866 (
17867 fv,
17868 LaunchConfig {
17869 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17870 block_dim: (32, gqa, 1),
17871 shared_mem_bytes: 0,
17872 },
17873 )
17874 } else if fa_vec && head_dim <= 256 {
17875 let gqa = (n_head / n_head_kv).max(1) as u32;
17876 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17887 let smem_tkv = *SMEM_TKV.get_or_init(|| {
17888 std::env::var("MEMRA_FA_SMEM_TKV")
17889 .ok()
17890 .and_then(|v| v.parse().ok())
17891 .unwrap_or_else(|| {
17892 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
17893 })
17894 });
17895 if fa_v4_at(t_kv) && head_dim == 256 {
17896 let v4name = match fa_v4_mode() {
17900 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
17903 _ => "fa_decode_vec_q_v4",
17904 };
17905 let fv = if g {
17906 self.func_g(v4name)
17907 } else {
17908 self.func(v4name)
17909 };
17910 let shmem = (if deep { 12160 } else { 11520 }
17913 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
17914 use cudarc::driver::sys::CUfunction_attribute_enum as A;
17915 fv.set_attribute(
17916 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17917 shmem as i32,
17918 )?;
17919 (
17920 fv,
17921 LaunchConfig {
17922 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17923 block_dim: (32, gqa, 1),
17924 shared_mem_bytes: shmem,
17925 },
17926 )
17927 } else if fa_v3_active(head_dim) {
17928 let fv = if g {
17931 self.func_g("fa_decode_vec_q_v3")
17932 } else {
17933 self.func("fa_decode_vec_q_v3")
17934 };
17935 let shmem = (32 * head_dim * 2) as u32; (
17937 fv,
17938 LaunchConfig {
17939 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17940 block_dim: (32, gqa, 1),
17941 shared_mem_bytes: shmem,
17942 },
17943 )
17944 } else if fa_v2_on() {
17945 let fv = if g {
17949 self.func_g("fa_decode_vec_q_v2")
17950 } else {
17951 self.func("fa_decode_vec_q_v2")
17952 };
17953 let shmem = (2 * 32 * head_dim * 2) as u32; (
17955 fv,
17956 LaunchConfig {
17957 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17958 block_dim: (32, gqa, 1),
17959 shared_mem_bytes: shmem,
17960 },
17961 )
17962 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
17963 {
17964 let fv = if g {
17968 self.func_g("fa_decode_vec_q_smem")
17969 } else {
17970 self.func("fa_decode_vec_q_smem")
17971 };
17972 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
17974 fv.set_attribute(
17975 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
17976 shmem as i32,
17977 )?;
17978 (
17979 fv,
17980 LaunchConfig {
17981 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17982 block_dim: (32, gqa, 1),
17983 shared_mem_bytes: shmem,
17984 },
17985 )
17986 } else {
17987 let fv = if g {
17990 self.func_g("fa_decode_vec_q")
17991 } else {
17992 self.func("fa_decode_vec_q")
17993 };
17994 (
17995 fv,
17996 LaunchConfig {
17997 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
17998 block_dim: (32, gqa, 1),
17999 shared_mem_bytes: 0,
18000 },
18001 )
18002 }
18003 } else {
18004 return self.fa_decode_scalar_unified(
18007 q,
18008 k,
18009 v,
18010 o,
18011 head_dim,
18012 n_head,
18013 n_head_kv,
18014 t_kv,
18015 None,
18016 scale,
18017 n_splits,
18018 if fa_vec { sp } else { 256 },
18019 k_tok_bytes,
18020 v_tok_bytes,
18021 g,
18022 part_o,
18023 part_m,
18024 part_l,
18025 None,
18026 );
18027 };
18028 let __s_b = self.gpu.stream();
18029 let mut b = __s_b.launch_builder(&f);
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(&scale)
18041 .arg(&nsp)
18042 .arg(&ktb)
18043 .arg(&vtb);
18044 unsafe {
18045 b.launch(cfg)?;
18046 }
18047 let (fc, cfg2) = (
18050 if g {
18051 self.func_g("fa_decode_combine_f32")
18052 } else {
18053 self.fa_func("fa_decode_combine_f32", head_dim)
18054 },
18055 LaunchConfig {
18056 grid_dim: (n_head as u32, 1, 1),
18057 block_dim: (head_dim as u32, 1, 1),
18058 shared_mem_bytes: 0,
18059 },
18060 );
18061 let __s_b2 = self.gpu.stream();
18062 let mut b2 = __s_b2.launch_builder(&fc);
18063 b2.arg(&*part_o)
18064 .arg(&*part_m)
18065 .arg(&*part_l)
18066 .arg(o)
18067 .arg(&hd)
18068 .arg(&nh)
18069 .arg(&nsp);
18070 unsafe {
18071 b2.launch(cfg2)?;
18072 }
18073 Ok(())
18074 }
18075
18076 #[allow(clippy::too_many_arguments)]
18087 pub fn fa_decode_batch_seqs_v4(
18088 &self,
18089 q: &CudaSlice<f32>,
18090 kv_ptrs: &cudarc::driver::CudaView<u64>,
18091 pos_seq: &CudaSlice<i32>,
18092 o: &mut CudaSlice<f32>,
18093 head_dim: usize,
18094 n_head: usize,
18095 n_head_kv: usize,
18096 b_n: usize,
18097 t_kv_max: usize,
18098 scale: f32,
18099 split_keys: usize,
18100 k_tok_bytes: usize,
18101 v_tok_bytes: usize,
18102 ) -> Result<(), Box<dyn std::error::Error>> {
18103 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
18104 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
18105 let o_len = b_n * n_head * n_splits_max * head_dim;
18106 let ml_len = b_n * n_head * n_splits_max;
18107 let mut part_guard = self.fa_part_pool.lock().unwrap();
18108 if part_guard
18109 .as_ref()
18110 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18111 .unwrap_or(true)
18112 {
18113 let old = part_guard.take();
18124 let (co, cm) = old
18125 .as_ref()
18126 .map(|pp| (pp.0.len(), pp.1.len()))
18127 .unwrap_or((0, 0));
18128 if let Some(old) = old {
18129 self.fa_part_retired.lock().unwrap().push(old);
18130 }
18131 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18132 eprintln!(
18133 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18134 co, o_len, cm, ml_len
18135 );
18136 }
18137 *part_guard = Some((
18138 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18139 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18140 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18141 ));
18142 }
18143 let pg = part_guard.as_mut().unwrap();
18144 self.gpu
18145 .stream()
18146 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18147 self.gpu
18148 .stream()
18149 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18150 self.gpu
18151 .stream()
18152 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18153 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18154 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18155 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
18156 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18157 let gqa = (n_head / n_head_kv).max(1) as u32;
18158 let f = self.func("fa_decode_vec_q_seqs_v4");
18159 let shmem = (11520 + 32 * head_dim * 2) as u32;
18161 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18162 f.set_attribute(
18163 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18164 shmem as i32,
18165 )?;
18166 let cfg = LaunchConfig {
18167 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
18168 block_dim: (32, gqa, 1),
18169 shared_mem_bytes: shmem,
18170 };
18171 {
18172 let __s_b = self.gpu.stream();
18173 let mut b = __s_b.launch_builder(&f);
18174 b.arg(q)
18175 .arg(kv_ptrs)
18176 .arg(pos_seq)
18177 .arg(&mut *part_o)
18178 .arg(&mut *part_m)
18179 .arg(&mut *part_l)
18180 .arg(&hd)
18181 .arg(&nh)
18182 .arg(&nhkv)
18183 .arg(&scale)
18184 .arg(&nspm)
18185 .arg(&spk)
18186 .arg(&ktb)
18187 .arg(&vtb);
18188 unsafe {
18189 b.launch(cfg)?;
18190 }
18191 }
18192 let fc = self.func("fa_decode_combine_seqs");
18193 let cfg2 = LaunchConfig {
18194 grid_dim: (n_head as u32, b_n as u32, 1),
18195 block_dim: (head_dim as u32, 1, 1),
18196 shared_mem_bytes: 0,
18197 };
18198 let __s_b2 = self.gpu.stream();
18199 let mut b2 = __s_b2.launch_builder(&fc);
18200 b2.arg(&*part_o)
18201 .arg(&*part_m)
18202 .arg(&*part_l)
18203 .arg(o)
18204 .arg(&hd)
18205 .arg(&nh)
18206 .arg(pos_seq)
18207 .arg(&nspm)
18208 .arg(&spk);
18209 unsafe {
18210 b2.launch(cfg2)?;
18211 }
18212 Ok(())
18213 }
18214
18215 #[allow(clippy::too_many_arguments)]
18222 pub fn append_kv_quantized_seqs(
18223 &self,
18224 k_rows: &CudaSlice<f32>,
18225 v_rows: &CudaSlice<f32>,
18226 kv_ptrs: &cudarc::driver::CudaView<u64>,
18227 pos_seq: &CudaSlice<i32>,
18228 b_n: usize,
18229 kv_dim_k: usize,
18230 kv_dim_v: usize,
18231 k_tok_bytes: usize,
18232 v_tok_bytes: usize,
18233 ) -> Result<(), Box<dyn std::error::Error>> {
18234 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
18235 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
18236 let cfg = LaunchConfig {
18237 grid_dim: (nblk, b_n as u32, 1),
18238 block_dim: (32, 1, 1),
18239 shared_mem_bytes: 0,
18240 };
18241 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
18242 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18243 let __s_b = self.gpu.stream();
18244 let mut b = __s_b.launch_builder(&f);
18245 b.arg(k_rows)
18246 .arg(v_rows)
18247 .arg(kv_ptrs)
18248 .arg(pos_seq)
18249 .arg(&kdk)
18250 .arg(&kdv)
18251 .arg(&ktb)
18252 .arg(&vtb);
18253 unsafe {
18254 b.launch(cfg)?;
18255 }
18256 Ok(())
18257 }
18258
18259 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
18265 std::env::var("MEMRA_NO_FA_VEC").is_err()
18266 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
18267 && base_len + 1 >= fa_vec_min_tkv()
18268 && head_dim <= 256
18269 && head_dim % 32 == 0
18270 }
18271
18272 #[allow(clippy::too_many_arguments)]
18281 pub fn fa_decode_rows(
18282 &self,
18283 q: &CudaSlice<f32>,
18284 k: &cudarc::driver::CudaView<u8>,
18285 v: &cudarc::driver::CudaView<u8>,
18286 o: &mut CudaSlice<f32>,
18287 head_dim: usize,
18288 n_head: usize,
18289 n_head_kv: usize,
18290 base_len: usize,
18291 t: usize,
18292 scale: f32,
18293 k_tok_bytes: usize,
18294 v_tok_bytes: usize,
18295 base_dev: Option<(&CudaSlice<i32>, i32)>,
18299 kv_shared: bool,
18302 g: bool,
18306 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18309 ) -> Result<(), Box<dyn std::error::Error>> {
18310 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
18311 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
18318 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18319 let v = *SP512.get_or_init(|| {
18322 std::env::var("MEMRA_FA_SP512")
18323 .ok()
18324 .and_then(|x| x.parse().ok())
18325 .unwrap_or(0)
18326 });
18327 sp = if v >= 8 {
18328 v
18329 } else {
18330 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18331 };
18332 }
18333 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18334 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18335 let gqa = (n_head / n_head_kv).max(1) as u32;
18336 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
18347 groups.push((0, t, sp));
18348 } else {
18349 let mut r0 = 0usize;
18350 while r0 < t {
18351 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
18352 let mut r1 = r0 + 1;
18353 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
18354 r1 += 1;
18355 }
18356 groups.push((r0, r1 - r0, sp_g));
18357 r0 = r1;
18358 }
18359 }
18360 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18364 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
18365 std::env::var("MEMRA_FA_SMEM_TKV")
18366 .ok()
18367 .and_then(|v| v.parse().ok())
18368 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18369 });
18370 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
18371 let v3 = fa_v3_active(head_dim);
18372 let smem_rows =
18373 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
18374 let _ = kv_shared;
18379 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
18382 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18396 let tb512 = head_dim == 512
18398 && sp <= 32
18399 && n_head / n_head_kv.max(1) <= 16
18400 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
18401 let fname = if tb512 {
18402 "fa_decode_vec_q_rows_v4_512_tb"
18403 } else if i2 {
18404 "fa_decode_vec_q_rows_dpl16_i2"
18405 } else if head_dim == 512 {
18406 "fa_decode_vec_q_rows_dpl16"
18407 }
18408 else if v4 {
18410 "fa_decode_vec_q_rows_v4"
18411 } else if v3 {
18412 "fa_decode_vec_q_rows_v3"
18413 } else if fa_v2_on() {
18414 "fa_decode_vec_q_rows_v2"
18415 } else if smem_rows {
18416 "fa_decode_vec_q_rows_smem"
18417 } else {
18418 "fa_decode_vec_q_rows"
18419 };
18420 let f = if head_dim == 512 {
18421 self.fa_func(fname, head_dim)
18422 } else if g {
18423 self.func_g(if smem_rows {
18431 "fa_decode_vec_q_rows"
18432 } else {
18433 fname
18434 })
18435 } else {
18436 self.func(fname)
18437 };
18438 let shmem = if tb512 {
18439 let gk = Self::gkv_on();
18441 let sh =
18442 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
18443 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18444 f.set_attribute(
18445 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18446 sh as i32,
18447 )?;
18448 sh
18449 } else if v4 || v3 || smem_rows || fa_v2_on() {
18450 let sh = (if v4 {
18452 11520 + 32 * head_dim * if g { 1 } else { 2 }
18453 } else if v3 {
18454 32 * head_dim * 2
18455 } else {
18456 2 * 32 * head_dim * 2
18457 }) as u32;
18458 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18459 f.set_attribute(
18460 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18461 sh as i32,
18462 )?;
18463 sh
18464 } else {
18465 0
18466 };
18467 for &(r0, t_g, sp_g) in &groups {
18471 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
18472 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
18473 let base_i = (base_len + r0) as i32;
18474 let o_len = t_g * n_head * n_splits_g * head_dim;
18475 let ml_len = t_g * n_head * n_splits_g;
18476 let mut part_guard = self.fa_part_pool.lock().unwrap();
18477 if part_guard
18478 .as_ref()
18479 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18480 .unwrap_or(true)
18481 {
18482 let old = part_guard.take();
18493 let (co, cm) = old
18494 .as_ref()
18495 .map(|pp| (pp.0.len(), pp.1.len()))
18496 .unwrap_or((0, 0));
18497 if let Some(old) = old {
18498 self.fa_part_retired.lock().unwrap().push(old);
18499 }
18500 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18501 eprintln!(
18502 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18503 co, o_len, cm, ml_len
18504 );
18505 }
18506 *part_guard = Some((
18507 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18508 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18509 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18510 ));
18511 }
18512 let pg = part_guard.as_mut().unwrap();
18513 self.gpu
18514 .stream()
18515 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18516 self.gpu
18517 .stream()
18518 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18519 self.gpu
18520 .stream()
18521 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18522 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18523 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
18524 let qv = self.view(q, t * n_head * head_dim);
18525 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18526 let cfg = LaunchConfig {
18527 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
18528 block_dim: (32, gqa, 1),
18529 shared_mem_bytes: shmem,
18530 };
18531 {
18532 let __s_b = self.gpu.stream();
18533 let mut b = __s_b.launch_builder(&f);
18534 if tb512 {
18535 let (bd, plus) =
18537 base_dev.expect("hd512 rows twin requires a device base counter");
18538 let plus_g = plus + r0 as i32;
18539 let nr = t_g as i32;
18540 if Self::pdl_on() && Self::pdl_wb_on() {
18541 use cudarc::driver::{DevicePtr, DevicePtrMut};
18543 let s = &self.gpu.stream();
18544 let (pq, _b0) = q_g.device_ptr(s);
18545 let (pk, _b1) = k.device_ptr(s);
18546 let (pv, _b2) = v.device_ptr(s);
18547 let (po, _b3) = part_o.device_ptr_mut(s);
18548 let (pm, _b4) = part_m.device_ptr_mut(s);
18549 let (pl, _b5) = part_l.device_ptr_mut(s);
18550 let (pb, _b6) = bd.device_ptr(s);
18551 let mut ps = [
18552 &pq as *const _ as *mut std::ffi::c_void,
18553 &pk as *const _ as *mut _,
18554 &pv as *const _ as *mut _,
18555 &po as *const _ as *mut _,
18556 &pm as *const _ as *mut _,
18557 &pl as *const _ as *mut _,
18558 &hd as *const _ as *mut _,
18559 &nh as *const _ as *mut _,
18560 &nhkv as *const _ as *mut _,
18561 &pb as *const _ as *mut _,
18562 &plus_g as *const _ as *mut _,
18563 &scale as *const _ as *mut _,
18564 &nspm as *const _ as *mut _,
18565 &spk as *const _ as *mut _,
18566 &ktb as *const _ as *mut _,
18567 &vtb as *const _ as *mut _,
18568 &nr as *const _ as *mut _,
18569 ];
18570 unsafe {
18571 self.launch_pdl_flash(
18572 Self::gkv_on(),
18573 "fa_decode_vec_q_rows_v4_512_tb",
18574 (n_head_kv as u32, n_splits_g as u32, 1),
18575 (32, gqa, 1),
18576 shmem,
18577 &mut ps,
18578 )?;
18579 }
18580 } else {
18581 let cfg_tb = LaunchConfig {
18582 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
18583 block_dim: (32, gqa, 1),
18584 shared_mem_bytes: shmem,
18585 };
18586 b.arg(&q_g)
18587 .arg(k)
18588 .arg(v)
18589 .arg(&mut *part_o)
18590 .arg(&mut *part_m)
18591 .arg(&mut *part_l)
18592 .arg(&hd)
18593 .arg(&nh)
18594 .arg(&nhkv)
18595 .arg(bd)
18596 .arg(&plus_g)
18597 .arg(&scale)
18598 .arg(&nspm)
18599 .arg(&spk)
18600 .arg(&ktb)
18601 .arg(&vtb)
18602 .arg(&nr);
18603 unsafe {
18604 b.launch(cfg_tb)?;
18605 }
18606 }
18607 } else if head_dim == 512 {
18608 let (bd, plus) =
18609 base_dev.expect("hd512 rows twin requires a device base counter");
18610 let plus_g = plus + r0 as i32;
18611 b.arg(&q_g)
18612 .arg(k)
18613 .arg(v)
18614 .arg(&mut *part_o)
18615 .arg(&mut *part_m)
18616 .arg(&mut *part_l)
18617 .arg(&hd)
18618 .arg(&nh)
18619 .arg(&nhkv)
18620 .arg(bd)
18621 .arg(&plus_g)
18622 .arg(&scale)
18623 .arg(&nspm)
18624 .arg(&spk)
18625 .arg(&ktb)
18626 .arg(&vtb);
18627 unsafe {
18628 b.launch(cfg)?;
18629 }
18630 } else {
18631 b.arg(&q_g)
18632 .arg(k)
18633 .arg(v)
18634 .arg(&mut *part_o)
18635 .arg(&mut *part_m)
18636 .arg(&mut *part_l)
18637 .arg(&hd)
18638 .arg(&nh)
18639 .arg(&nhkv)
18640 .arg(&base_i)
18641 .arg(&scale)
18642 .arg(&nspm)
18643 .arg(&spk)
18644 .arg(&ktb)
18645 .arg(&vtb);
18646 unsafe {
18647 b.launch(cfg)?;
18648 }
18649 }
18650 }
18651 let cfg2 = LaunchConfig {
18652 grid_dim: (n_head as u32, t_g as u32, 1),
18653 block_dim: (head_dim as u32, 1, 1),
18654 shared_mem_bytes: 0,
18655 };
18656 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
18657 if head_dim == 512 {
18658 let (bd, plus) = base_dev.unwrap();
18661 let plus_g = plus + r0 as i32;
18662 if let Some((oq, od)) = q8_out.as_mut() {
18663 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
18665 if Self::pdl_on() && Self::pdl_wb_on() {
18666 use cudarc::driver::{DevicePtr, DevicePtrMut};
18668 let s = &self.gpu.stream();
18669 let (po, _g0) = part_o.device_ptr(s);
18670 let (pm, _g1) = part_m.device_ptr(s);
18671 let (pl, _g2) = part_l.device_ptr(s);
18672 let (pq, _g3) = oq.device_ptr_mut(s);
18673 let (pd, _g4) = od.device_ptr_mut(s);
18674 let (pb, _g5) = bd.device_ptr(s);
18675 let mut ps = [
18676 &po as *const _ as *mut std::ffi::c_void,
18677 &pm as *const _ as *mut _,
18678 &pl as *const _ as *mut _,
18679 &pq as *const _ as *mut _,
18680 &pd as *const _ as *mut _,
18681 &hd as *const _ as *mut _,
18682 &nh as *const _ as *mut _,
18683 &pb as *const _ as *mut _,
18684 &plus_g as *const _ as *mut _,
18685 &nspm as *const _ as *mut _,
18686 &spk as *const _ as *mut _,
18687 ];
18688 unsafe {
18689 self.launch_pdl_flash(
18690 Self::gkv_on(),
18691 "fa_decode_combine_rows_dc_q8_1",
18692 cfg2.grid_dim,
18693 cfg2.block_dim,
18694 0,
18695 &mut ps,
18696 )?;
18697 }
18698 continue;
18699 }
18700 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
18701 let __s_b2 = self.gpu.stream();
18702 let mut b2 = __s_b2.launch_builder(&fc);
18703 b2.arg(&*part_o)
18704 .arg(&*part_m)
18705 .arg(&*part_l)
18706 .arg(&mut **oq)
18707 .arg(&mut **od)
18708 .arg(&hd)
18709 .arg(&nh)
18710 .arg(bd)
18711 .arg(&plus_g)
18712 .arg(&nspm)
18713 .arg(&spk);
18714 unsafe {
18715 b2.launch(cfg2)?;
18716 }
18717 continue;
18718 }
18719 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
18720 let __s_b2 = self.gpu.stream();
18721 let mut b2 = __s_b2.launch_builder(&fc);
18722 b2.arg(&*part_o)
18723 .arg(&*part_m)
18724 .arg(&*part_l)
18725 .arg(&mut o_g)
18726 .arg(&hd)
18727 .arg(&nh)
18728 .arg(bd)
18729 .arg(&plus_g)
18730 .arg(&nspm)
18731 .arg(&spk);
18732 unsafe {
18733 b2.launch(cfg2)?;
18734 }
18735 } else {
18736 assert!(
18739 q8_out.is_none(),
18740 "rows q8 emit requires the hd512 dc combine"
18741 );
18742 let fc = self.func("fa_decode_combine_rows");
18743 let __s_b2 = self.gpu.stream();
18744 let mut b2 = __s_b2.launch_builder(&fc);
18745 b2.arg(&*part_o)
18746 .arg(&*part_m)
18747 .arg(&*part_l)
18748 .arg(&mut o_g)
18749 .arg(&hd)
18750 .arg(&nh)
18751 .arg(&base_i)
18752 .arg(&nspm)
18753 .arg(&spk);
18754 unsafe {
18755 b2.launch(cfg2)?;
18756 }
18757 }
18758 }
18759 Ok(())
18760 }
18761
18762 #[allow(clippy::too_many_arguments)]
18766 pub fn fa_decode_rows_w(
18767 &self,
18768 q: &CudaSlice<f32>,
18769 k: &cudarc::driver::CudaView<u8>,
18770 v: &cudarc::driver::CudaView<u8>,
18771 o: &mut CudaSlice<f32>,
18772 head_dim: usize,
18773 n_head: usize,
18774 n_head_kv: usize,
18775 base_dev: &CudaSlice<i32>,
18776 base_plus: i32,
18777 t: usize,
18778 scale: f32,
18779 window: usize,
18780 k_tok_bytes: usize,
18781 v_tok_bytes: usize,
18782 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
18783 ) -> Result<(), Box<dyn std::error::Error>> {
18784 debug_assert!(head_dim == 256);
18789 let sp = {
18797 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18798 let v = *SPW.get_or_init(|| {
18799 std::env::var("MEMRA_FA_SPW")
18800 .ok()
18801 .and_then(|x| x.parse().ok())
18802 .unwrap_or(0)
18803 });
18804 if v >= 8 {
18805 v
18806 } else {
18807 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
18808 }
18809 };
18810 let n_splits_max = (window + sp - 1) / sp;
18811 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
18812 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
18813 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18814 let gqa = (n_head / n_head_kv).max(1) as u32;
18815 let o_len = t * n_head * n_splits_max * head_dim;
18816 let ml_len = t * n_head * n_splits_max;
18817 let mut part_guard = self.fa_part_pool.lock().unwrap();
18818 if part_guard
18819 .as_ref()
18820 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
18821 .unwrap_or(true)
18822 {
18823 let old = part_guard.take();
18834 let (co, cm) = old
18835 .as_ref()
18836 .map(|pp| (pp.0.len(), pp.1.len()))
18837 .unwrap_or((0, 0));
18838 if let Some(old) = old {
18839 self.fa_part_retired.lock().unwrap().push(old);
18840 }
18841 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
18842 eprintln!(
18843 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
18844 co, o_len, cm, ml_len
18845 );
18846 }
18847 *part_guard = Some((
18848 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
18849 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18850 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
18851 ));
18852 }
18853 let pg = part_guard.as_mut().unwrap();
18854 self.gpu
18855 .stream()
18856 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
18857 self.gpu
18858 .stream()
18859 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
18860 self.gpu
18861 .stream()
18862 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
18863 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
18864 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
18870 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
18871 std::env::var("MEMRA_FA_SMEM_TKV")
18872 .ok()
18873 .and_then(|v| v.parse().ok())
18874 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
18875 });
18876 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18882 let wg = Self::wkv_on();
18887 let sp2 =
18890 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
18891 if sp2 {
18892 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18893 if Self::pdl_on() && Self::pdl_wb_on() {
18894 use cudarc::driver::{DevicePtr, DevicePtrMut};
18896 let s = &self.gpu.stream();
18897 let (pq, _b0) = q.device_ptr(s);
18898 let (pk, _b1) = k.device_ptr(s);
18899 let (pv, _b2) = v.device_ptr(s);
18900 let (po, _b3) = part_o.device_ptr_mut(s);
18901 let (pm, _b4) = part_m.device_ptr_mut(s);
18902 let (pl, _b5) = part_l.device_ptr_mut(s);
18903 let (pb, _b6) = base_dev.device_ptr(s);
18904 let mut ps = [
18905 &pq as *const _ as *mut std::ffi::c_void,
18906 &pk as *const _ as *mut _,
18907 &pv as *const _ as *mut _,
18908 &po as *const _ as *mut _,
18909 &pm as *const _ as *mut _,
18910 &pl as *const _ as *mut _,
18911 &hd as *const _ as *mut _,
18912 &nh as *const _ as *mut _,
18913 &nhkv as *const _ as *mut _,
18914 &pb as *const _ as *mut _,
18915 &base_plus as *const _ as *mut _,
18916 &scale as *const _ as *mut _,
18917 &nspm as *const _ as *mut _,
18918 &spk as *const _ as *mut _,
18919 &ktb as *const _ as *mut _,
18920 &vtb as *const _ as *mut _,
18921 &wini as *const _ as *mut _,
18922 ];
18923 unsafe {
18924 self.launch_pdl_flash(
18925 wg,
18926 "fa_decode_vec_q_rows_v4_w_sp",
18927 (n_head_kv as u32, n_splits_max as u32, t as u32),
18928 (32, gqa + 1, 1),
18929 sh,
18930 &mut ps,
18931 )?;
18932 }
18933 } else {
18934 let f = if wg {
18935 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
18936 } else {
18937 self.func("fa_decode_vec_q_rows_v4_w_sp")
18938 };
18939 f.set_attribute(
18940 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18941 sh as i32,
18942 )?;
18943 let cfg = LaunchConfig {
18944 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
18945 block_dim: (32, gqa + 1, 1),
18946 shared_mem_bytes: sh,
18947 };
18948 let __s_b = self.gpu.stream();
18949 let mut b = __s_b.launch_builder(&f);
18950 b.arg(q)
18951 .arg(k)
18952 .arg(v)
18953 .arg(&mut *part_o)
18954 .arg(&mut *part_m)
18955 .arg(&mut *part_l)
18956 .arg(&hd)
18957 .arg(&nh)
18958 .arg(&nhkv)
18959 .arg(base_dev)
18960 .arg(&base_plus)
18961 .arg(&scale)
18962 .arg(&nspm)
18963 .arg(&spk)
18964 .arg(&ktb)
18965 .arg(&vtb)
18966 .arg(&wini);
18967 unsafe {
18968 b.launch(cfg)?;
18969 }
18970 }
18971 } else {
18972 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
18973 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
18975 use cudarc::driver::{DevicePtr, DevicePtrMut};
18976 let s = &self.gpu.stream();
18977 let (pq, _b0) = q.device_ptr(s);
18978 let (pk, _b1) = k.device_ptr(s);
18979 let (pv, _b2) = v.device_ptr(s);
18980 let (po, _b3) = part_o.device_ptr_mut(s);
18981 let (pm, _b4) = part_m.device_ptr_mut(s);
18982 let (pl, _b5) = part_l.device_ptr_mut(s);
18983 let (pb, _b6) = base_dev.device_ptr(s);
18984 let mut ps = [
18985 &pq as *const _ as *mut std::ffi::c_void,
18986 &pk as *const _ as *mut _,
18987 &pv as *const _ as *mut _,
18988 &po as *const _ as *mut _,
18989 &pm as *const _ as *mut _,
18990 &pl as *const _ as *mut _,
18991 &hd as *const _ as *mut _,
18992 &nh as *const _ as *mut _,
18993 &nhkv as *const _ as *mut _,
18994 &pb as *const _ as *mut _,
18995 &base_plus as *const _ as *mut _,
18996 &scale as *const _ as *mut _,
18997 &nspm as *const _ as *mut _,
18998 &spk as *const _ as *mut _,
18999 &ktb as *const _ as *mut _,
19000 &vtb as *const _ as *mut _,
19001 &wini as *const _ as *mut _,
19002 ];
19003 unsafe {
19004 self.launch_pdl_flash(
19005 wg,
19006 "fa_decode_vec_q_rows_v4_w",
19007 (n_head_kv as u32, n_splits_max as u32, t as u32),
19008 (32, gqa, 1),
19009 sh,
19010 &mut ps,
19011 )?;
19012 }
19013 } else {
19014 let pick = |name: &str| {
19015 if wg {
19016 self.func_g(name)
19017 } else {
19018 self.func(name)
19019 }
19020 };
19021 let (f, sh) = if fa_v4_at(window) {
19022 let f = pick("fa_decode_vec_q_rows_v4_w");
19023 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
19024 } else if smem_tkv > 0 && window >= smem_tkv {
19025 (
19028 pick("fa_decode_vec_q_rows_smem_w"),
19029 (2 * 32 * head_dim * 2) as u32,
19030 )
19031 } else {
19032 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
19033 };
19034 f.set_attribute(
19035 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19036 sh as i32,
19037 )?;
19038 let cfg = LaunchConfig {
19039 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19040 block_dim: (32, gqa, 1),
19041 shared_mem_bytes: sh,
19042 };
19043 let __s_b = self.gpu.stream();
19044 let mut b = __s_b.launch_builder(&f);
19045 b.arg(q)
19046 .arg(k)
19047 .arg(v)
19048 .arg(&mut *part_o)
19049 .arg(&mut *part_m)
19050 .arg(&mut *part_l)
19051 .arg(&hd)
19052 .arg(&nh)
19053 .arg(&nhkv)
19054 .arg(base_dev)
19055 .arg(&base_plus)
19056 .arg(&scale)
19057 .arg(&nspm)
19058 .arg(&spk)
19059 .arg(&ktb)
19060 .arg(&vtb)
19061 .arg(&wini);
19062 unsafe {
19063 b.launch(cfg)?;
19064 }
19065 }
19066 }
19067 let cfg2 = LaunchConfig {
19068 grid_dim: (n_head as u32, t as u32, 1),
19069 block_dim: (head_dim as u32, 1, 1),
19070 shared_mem_bytes: 0,
19071 };
19072 if let Some((oq, od)) = q8_out {
19073 if Self::pdl_on() && Self::pdl_wb_on() {
19076 use cudarc::driver::{DevicePtr, DevicePtrMut};
19078 let s = &self.gpu.stream();
19079 let (po, _g0) = part_o.device_ptr(s);
19080 let (pm, _g1) = part_m.device_ptr(s);
19081 let (pl, _g2) = part_l.device_ptr(s);
19082 let (pq, _g3) = oq.device_ptr_mut(s);
19083 let (pd, _g4) = od.device_ptr_mut(s);
19084 let mut ps = [
19085 &po as *const _ as *mut std::ffi::c_void,
19086 &pm as *const _ as *mut _,
19087 &pl as *const _ as *mut _,
19088 &pq as *const _ as *mut _,
19089 &pd as *const _ as *mut _,
19090 &hd as *const _ as *mut _,
19091 &nh as *const _ as *mut _,
19092 &nspm as *const _ as *mut _,
19093 &spk as *const _ as *mut _,
19094 &wini as *const _ as *mut _,
19095 ];
19096 unsafe {
19097 self.launch_pdl_flash(
19098 wg,
19099 "fa_decode_combine_rows_w_q8_1",
19100 cfg2.grid_dim,
19101 cfg2.block_dim,
19102 0,
19103 &mut ps,
19104 )?;
19105 }
19106 return Ok(());
19107 }
19108 let fc = if wg {
19109 self.func_g("fa_decode_combine_rows_w_q8_1")
19110 } else {
19111 self.func("fa_decode_combine_rows_w_q8_1")
19112 };
19113 let __s_b2 = self.gpu.stream();
19114 let mut b2 = __s_b2.launch_builder(&fc);
19115 b2.arg(&*part_o)
19116 .arg(&*part_m)
19117 .arg(&*part_l)
19118 .arg(oq)
19119 .arg(od)
19120 .arg(&hd)
19121 .arg(&nh)
19122 .arg(&nspm)
19123 .arg(&spk)
19124 .arg(&wini);
19125 unsafe {
19126 b2.launch(cfg2)?;
19127 }
19128 return Ok(());
19129 }
19130 let fc = if wg {
19131 self.func_g("fa_decode_combine_rows_w")
19132 } else {
19133 self.func("fa_decode_combine_rows_w")
19134 };
19135 let __s_b2 = self.gpu.stream();
19136 let mut b2 = __s_b2.launch_builder(&fc);
19137 b2.arg(&*part_o)
19138 .arg(&*part_m)
19139 .arg(&*part_l)
19140 .arg(o)
19141 .arg(&hd)
19142 .arg(&nh)
19143 .arg(&nspm)
19144 .arg(&spk)
19145 .arg(&wini);
19146 unsafe {
19147 b2.launch(cfg2)?;
19148 }
19149 Ok(())
19150 }
19151
19152 #[allow(clippy::too_many_arguments)]
19158 pub fn fa_decode_rows_dc(
19159 &self,
19160 q: &CudaSlice<f32>,
19161 k: &cudarc::driver::CudaView<u8>,
19162 v: &cudarc::driver::CudaView<u8>,
19163 o: &mut CudaSlice<f32>,
19164 head_dim: usize,
19165 n_head: usize,
19166 n_head_kv: usize,
19167 base_dev: &CudaSlice<i32>,
19168 t_kv_upper: usize,
19169 t: usize,
19170 scale: f32,
19171 k_tok_bytes: usize,
19172 v_tok_bytes: usize,
19173 base_plus: i32,
19174 g: bool,
19175 ) -> Result<(), Box<dyn std::error::Error>> {
19176 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
19177 assert!(
19178 v4 || fa_v3_active(head_dim),
19179 "stream fa rows requires the v3 or v4 lane"
19180 );
19181 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
19182 if v4 {
19183 let sp = fa_split_keys(t_kv_upper, n_head_kv);
19184 let n_splits_max = (t_kv_upper + sp - 1) / sp;
19185 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19186 let (nspm, spk) = (n_splits_max as i32, sp as i32);
19187 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19188 let gqa = (n_head / n_head_kv).max(1) as u32;
19189 let o_len = t * n_head * n_splits_max * head_dim;
19190 let ml_len = t * n_head * n_splits_max;
19191 let mut part_guard = self.fa_part_pool.lock().unwrap();
19192 if part_guard
19193 .as_ref()
19194 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19195 .unwrap_or(true)
19196 {
19197 let old = part_guard.take();
19208 let (co, cm) = old
19209 .as_ref()
19210 .map(|pp| (pp.0.len(), pp.1.len()))
19211 .unwrap_or((0, 0));
19212 if let Some(old) = old {
19213 self.fa_part_retired.lock().unwrap().push(old);
19214 }
19215 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19216 eprintln!(
19217 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19218 co, o_len, cm, ml_len
19219 );
19220 }
19221 *part_guard = Some((
19222 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19223 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19224 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19225 ));
19226 }
19227 let pg = part_guard.as_mut().unwrap();
19228 self.gpu
19229 .stream()
19230 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19231 self.gpu
19232 .stream()
19233 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19234 self.gpu
19235 .stream()
19236 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19237 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19238 let f = if g {
19239 self.func_g("fa_decode_vec_q_rows_v4_dc")
19240 } else {
19241 self.func("fa_decode_vec_q_rows_v4_dc")
19242 };
19243 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19244 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19245 f.set_attribute(
19246 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19247 sh as i32,
19248 )?;
19249 let cfg = LaunchConfig {
19250 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19251 block_dim: (32, gqa, 1),
19252 shared_mem_bytes: sh,
19253 };
19254 let __s_b = self.gpu.stream();
19255 let mut b = __s_b.launch_builder(&f);
19256 b.arg(q)
19257 .arg(k)
19258 .arg(v)
19259 .arg(&mut *part_o)
19260 .arg(&mut *part_m)
19261 .arg(&mut *part_l)
19262 .arg(&hd)
19263 .arg(&nh)
19264 .arg(&nhkv)
19265 .arg(base_dev)
19266 .arg(&base_plus)
19267 .arg(&scale)
19268 .arg(&nspm)
19269 .arg(&spk)
19270 .arg(&ktb)
19271 .arg(&vtb);
19272 unsafe {
19273 b.launch(cfg)?;
19274 }
19275 let fc = self.func("fa_decode_combine_rows_dc");
19276 let cfg2 = LaunchConfig {
19277 grid_dim: (n_head as u32, t as u32, 1),
19278 block_dim: (head_dim as u32, 1, 1),
19279 shared_mem_bytes: 0,
19280 };
19281 let __s_b2 = self.gpu.stream();
19282 let mut b2 = __s_b2.launch_builder(&fc);
19283 b2.arg(&*part_o)
19284 .arg(&*part_m)
19285 .arg(&*part_l)
19286 .arg(o)
19287 .arg(&hd)
19288 .arg(&nh)
19289 .arg(base_dev)
19290 .arg(&base_plus)
19291 .arg(&nspm)
19292 .arg(&spk);
19293 unsafe {
19294 b2.launch(cfg2)?;
19295 }
19296 return Ok(());
19297 }
19298 let sp = fa_split_keys(t_kv_upper, n_head_kv);
19299 let n_splits_max = (t_kv_upper + sp - 1) / sp;
19300 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19301 let (nspm, spk) = (n_splits_max as i32, sp as i32);
19302 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19303 let gqa = (n_head / n_head_kv).max(1) as u32;
19304 let o_len = t * n_head * n_splits_max * head_dim;
19305 let ml_len = t * n_head * n_splits_max;
19306 let mut part_guard = self.fa_part_pool.lock().unwrap();
19307 if part_guard
19308 .as_ref()
19309 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19310 .unwrap_or(true)
19311 {
19312 let old = part_guard.take();
19323 let (co, cm) = old
19324 .as_ref()
19325 .map(|pp| (pp.0.len(), pp.1.len()))
19326 .unwrap_or((0, 0));
19327 if let Some(old) = old {
19328 self.fa_part_retired.lock().unwrap().push(old);
19329 }
19330 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19331 eprintln!(
19332 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19333 co, o_len, cm, ml_len
19334 );
19335 }
19336 *part_guard = Some((
19337 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19338 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19339 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19340 ));
19341 }
19342 let pg = part_guard.as_mut().unwrap();
19343 self.gpu
19344 .stream()
19345 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19346 self.gpu
19347 .stream()
19348 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19349 self.gpu
19350 .stream()
19351 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19352 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19353 let f = self.func("fa_decode_vec_q_rows_v3_dc");
19354 let sh = (32 * head_dim * 2) as u32;
19355 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19356 f.set_attribute(
19357 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19358 sh as i32,
19359 )?;
19360 let cfg = LaunchConfig {
19361 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
19362 block_dim: (32, gqa, 1),
19363 shared_mem_bytes: sh,
19364 };
19365 let __s_b = self.gpu.stream();
19366 let mut b = __s_b.launch_builder(&f);
19367 b.arg(q)
19368 .arg(k)
19369 .arg(v)
19370 .arg(&mut *part_o)
19371 .arg(&mut *part_m)
19372 .arg(&mut *part_l)
19373 .arg(&hd)
19374 .arg(&nh)
19375 .arg(&nhkv)
19376 .arg(base_dev)
19377 .arg(&scale)
19378 .arg(&nspm)
19379 .arg(&spk)
19380 .arg(&ktb)
19381 .arg(&vtb);
19382 unsafe {
19383 b.launch(cfg)?;
19384 }
19385 let fc = self.func("fa_decode_combine_rows_dc");
19386 let cfg2 = LaunchConfig {
19387 grid_dim: (n_head as u32, t as u32, 1),
19388 block_dim: (head_dim as u32, 1, 1),
19389 shared_mem_bytes: 0,
19390 };
19391 let plus0 = 0i32;
19392 let __s_b2 = self.gpu.stream();
19393 let mut b2 = __s_b2.launch_builder(&fc);
19394 b2.arg(&*part_o)
19395 .arg(&*part_m)
19396 .arg(&*part_l)
19397 .arg(o)
19398 .arg(&hd)
19399 .arg(&nh)
19400 .arg(base_dev)
19401 .arg(&plus0)
19402 .arg(&nspm)
19403 .arg(&spk);
19404 unsafe {
19405 b2.launch(cfg2)?;
19406 }
19407 Ok(())
19408 }
19409
19410 pub fn fa_decode_dc(
19421 &self,
19422 q: &CudaSlice<f32>,
19423 k: &cudarc::driver::CudaView<u8>,
19424 v: &cudarc::driver::CudaView<u8>,
19425 o: &mut CudaSlice<f32>,
19426 head_dim: usize,
19427 n_head: usize,
19428 n_head_kv: usize,
19429 t_kv_dev: &CudaSlice<i32>,
19430 bucket_max: usize,
19431 scale: f32,
19432 k_tok_bytes: usize,
19433 v_tok_bytes: usize,
19434 g: bool,
19435 ) -> Result<(), Box<dyn std::error::Error>> {
19436 self.fa_decode_dc_q8(
19437 q,
19438 k,
19439 v,
19440 o,
19441 head_dim,
19442 n_head,
19443 n_head_kv,
19444 t_kv_dev,
19445 bucket_max,
19446 scale,
19447 k_tok_bytes,
19448 v_tok_bytes,
19449 g,
19450 None,
19451 )
19452 }
19453
19454 #[allow(clippy::too_many_arguments)]
19457 pub fn fa_decode_dc_q8(
19458 &self,
19459 q: &CudaSlice<f32>,
19460 k: &cudarc::driver::CudaView<u8>,
19461 v: &cudarc::driver::CudaView<u8>,
19462 o: &mut CudaSlice<f32>,
19463 head_dim: usize,
19464 n_head: usize,
19465 n_head_kv: usize,
19466 t_kv_dev: &CudaSlice<i32>,
19467 bucket_max: usize,
19468 scale: f32,
19469 k_tok_bytes: usize,
19470 v_tok_bytes: usize,
19471 g: bool,
19472 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
19473 ) -> Result<(), Box<dyn std::error::Error>> {
19474 let mut fa_vec =
19482 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
19483 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
19484 fa_vec = false;
19485 } let sp = fa_split_keys(bucket_max, n_head_kv);
19487 let n_splits = if fa_vec {
19488 ((bucket_max + sp - 1) / sp).max(1)
19489 } else {
19490 ((bucket_max + 255) / 256).max(1)
19491 };
19492 let o_len = n_head * n_splits * head_dim;
19493 let ml_len = n_head * n_splits;
19494 let mut part_guard = self.fa_part_pool.lock().unwrap();
19495 if part_guard
19496 .as_ref()
19497 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
19498 .unwrap_or(true)
19499 {
19500 let old = part_guard.take();
19511 let (co, cm) = old
19512 .as_ref()
19513 .map(|pp| (pp.0.len(), pp.1.len()))
19514 .unwrap_or((0, 0));
19515 if let Some(old) = old {
19516 self.fa_part_retired.lock().unwrap().push(old);
19517 }
19518 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
19519 eprintln!(
19520 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
19521 co, o_len, cm, ml_len
19522 );
19523 }
19524 *part_guard = Some((
19525 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
19526 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19527 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
19528 ));
19529 }
19530 let pg = part_guard.as_mut().unwrap();
19531 self.gpu
19532 .stream()
19533 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
19534 self.gpu
19535 .stream()
19536 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
19537 self.gpu
19538 .stream()
19539 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
19540 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
19541 let (hd, nh, nhkv, nsp) = (
19542 head_dim as i32,
19543 n_head as i32,
19544 n_head_kv as i32,
19545 n_splits as i32,
19546 );
19547 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19548 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
19549 let deep = fa_vec
19552 && head_dim == 256
19553 && fa_v4_at(bucket_max)
19554 && !g
19555 && fa_deep_at(bucket_max)
19556 && !matches!(fa_v4_mode(), "noB3" | "stage");
19557 let (f, cfg) = if fa_vec
19558 && head_dim == 512
19559 && bucket_max >= {
19560 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
19561 *FA512_MIN_DC.get_or_init(|| {
19562 std::env::var("MEMRA_FA512_MIN")
19563 .ok()
19564 .and_then(|v| v.parse().ok())
19565 .unwrap_or(512)
19566 })
19567 } {
19568 let gqa = (n_head / n_head_kv).max(1) as u32;
19570 (
19571 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
19572 LaunchConfig {
19573 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19574 block_dim: (32, gqa, 1),
19575 shared_mem_bytes: 0,
19576 },
19577 )
19578 } else if fa_vec && head_dim == 512 {
19579 let q_view = q.as_view();
19582 let mut o_view = o.as_view_mut();
19583 return self.fa_decode_scalar_unified(
19584 &q_view,
19585 k,
19586 v,
19587 &mut o_view,
19588 head_dim,
19589 n_head,
19590 n_head_kv,
19591 0,
19592 Some(t_kv_dev),
19593 scale,
19594 n_splits,
19595 sp,
19596 k_tok_bytes,
19597 v_tok_bytes,
19598 g,
19599 &mut *part_o,
19600 &mut *part_m,
19601 &mut *part_l,
19602 q8_out,
19603 );
19604 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
19605 let gqa = (n_head / n_head_kv).max(1) as u32;
19608 let fv = if g {
19609 self.func_g("fa_decode_vec_q_v4_dc")
19610 } else if deep {
19611 self.func("fa_decode_vec_q_v4_deep_dc")
19612 } else {
19613 self.func("fa_decode_vec_q_v4_dc")
19614 };
19615 let shmem =
19616 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
19617 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19618 fv.set_attribute(
19619 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19620 shmem as i32,
19621 )?;
19622 (
19623 fv,
19624 LaunchConfig {
19625 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19626 block_dim: (32, gqa, 1),
19627 shared_mem_bytes: shmem,
19628 },
19629 )
19630 } else if fa_vec && fa_v3_active(head_dim) {
19631 let gqa = (n_head / n_head_kv).max(1) as u32;
19634 let fv = if g {
19635 self.func_g("fa_decode_vec_q_v3_dc")
19636 } else {
19637 self.func("fa_decode_vec_q_v3_dc")
19638 };
19639 let shmem = (32 * head_dim * 2) as u32; (
19641 fv,
19642 LaunchConfig {
19643 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19644 block_dim: (32, gqa, 1),
19645 shared_mem_bytes: shmem,
19646 },
19647 )
19648 } else if fa_vec && fa_v2_on() {
19649 let gqa = (n_head / n_head_kv).max(1) as u32;
19653 let fv = if g {
19654 self.func_g("fa_decode_vec_q_v2_dc")
19655 } else {
19656 self.func("fa_decode_vec_q_v2_dc")
19657 };
19658 let shmem = (2 * 32 * head_dim * 2) as u32; (
19660 fv,
19661 LaunchConfig {
19662 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19663 block_dim: (32, gqa, 1),
19664 shared_mem_bytes: shmem,
19665 },
19666 )
19667 } else if fa_vec {
19668 let gqa = (n_head / n_head_kv).max(1) as u32;
19669 let fv = if g {
19671 self.func_g("fa_decode_vec_q_dc")
19672 } else {
19673 self.func("fa_decode_vec_q_dc")
19674 };
19675 (
19676 fv,
19677 LaunchConfig {
19678 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
19679 block_dim: (32, gqa, 1),
19680 shared_mem_bytes: 0,
19681 },
19682 )
19683 } else {
19684 let q_view = q.as_view();
19685 let mut o_view = o.as_view_mut();
19686 return self.fa_decode_scalar_unified(
19687 &q_view,
19688 k,
19689 v,
19690 &mut o_view,
19691 head_dim,
19692 n_head,
19693 n_head_kv,
19694 0,
19695 Some(t_kv_dev),
19696 scale,
19697 n_splits,
19698 if fa_vec { sp } else { 256 },
19699 k_tok_bytes,
19700 v_tok_bytes,
19701 g,
19702 &mut *part_o,
19703 &mut *part_m,
19704 &mut *part_l,
19705 q8_out,
19706 );
19707 };
19708 let ski = sp as i32; let __s_b = self.gpu.stream();
19710 let mut b = __s_b.launch_builder(&f);
19711 b.arg(q)
19712 .arg(k)
19713 .arg(v)
19714 .arg(&mut *part_o)
19715 .arg(&mut *part_m)
19716 .arg(&mut *part_l)
19717 .arg(&hd)
19718 .arg(&nh)
19719 .arg(&nhkv)
19720 .arg(t_kv_dev)
19721 .arg(&scale)
19722 .arg(&nsp)
19723 .arg(&ski)
19724 .arg(&ktb)
19725 .arg(&vtb);
19726 unsafe {
19727 b.launch(cfg)?;
19728 }
19729 let cfg2 = LaunchConfig {
19730 grid_dim: (n_head as u32, 1, 1),
19731 block_dim: (head_dim as u32, 1, 1),
19732 shared_mem_bytes: 0,
19733 };
19734 if let Some((oq, od)) = q8_out {
19735 let fc = if g {
19736 self.func_g("fa_decode_combine_q8_1")
19737 } else {
19738 self.fa_func("fa_decode_combine_q8_1", head_dim)
19739 };
19740 let __s_b2 = self.gpu.stream();
19741 let mut b2 = __s_b2.launch_builder(&fc);
19742 b2.arg(&*part_o)
19743 .arg(&*part_m)
19744 .arg(&*part_l)
19745 .arg(oq)
19746 .arg(od)
19747 .arg(&hd)
19748 .arg(&nh)
19749 .arg(&nsp);
19750 unsafe {
19751 b2.launch(cfg2)?;
19752 }
19753 return Ok(());
19754 }
19755 let fc = if g {
19756 self.func_g("fa_decode_combine_f32")
19757 } else {
19758 self.fa_func("fa_decode_combine_f32", head_dim)
19759 };
19760 let __s_b2 = self.gpu.stream();
19761 let mut b2 = __s_b2.launch_builder(&fc);
19762 b2.arg(&*part_o)
19763 .arg(&*part_m)
19764 .arg(&*part_l)
19765 .arg(o)
19766 .arg(&hd)
19767 .arg(&nh)
19768 .arg(&nsp);
19769 unsafe {
19770 b2.launch(cfg2)?;
19771 }
19772 Ok(())
19773 }
19774
19775 pub fn fa_geom_eager(
19781 &self,
19782 t_kv: usize,
19783 head_dim: usize,
19784 n_head_kv: usize,
19785 g: bool,
19786 ) -> (bool, usize) {
19787 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
19791 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
19797 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
19798 if g && head_dim == 256 && !fa_v4_at(t_kv) {
19804 fa_vec = false;
19805 }
19806 let sp = fa_split_keys(t_kv, n_head_kv);
19807 let n_splits = if fa_vec {
19808 ((t_kv + sp - 1) / sp).max(1)
19809 } else {
19810 ((t_kv + 255) / 256).max(1)
19811 };
19812 (fa_vec, n_splits)
19813 }
19814
19815 pub fn fa_bucket_key(
19821 &self,
19822 t_kv: usize,
19823 head_dim: usize,
19824 n_head_kv: usize,
19825 g: bool,
19826 ) -> (bool, usize) {
19827 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
19828 }
19829
19830 pub fn capture_graph_retained<F>(
19842 &self,
19843 step: F,
19844 ) -> Result<
19845 (
19846 cudarc::driver::CudaGraph,
19847 Vec<Box<dyn std::any::Any + Send>>,
19848 ),
19849 Box<dyn std::error::Error>,
19850 >
19851 where
19852 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19853 {
19854 use cudarc::driver::sys::CUgraphInstantiate_flags;
19855 self.capture_graph_retained_flags(
19856 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19857 step,
19858 )
19859 }
19860
19861 pub fn capture_graph_retained_flags<F>(
19866 &self,
19867 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
19868 mut step: F,
19869 ) -> Result<
19870 (
19871 cudarc::driver::CudaGraph,
19872 Vec<Box<dyn std::any::Any + Send>>,
19873 ),
19874 Box<dyn std::error::Error>,
19875 >
19876 where
19877 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19878 {
19879 use cudarc::driver::sys::CUstreamCaptureMode;
19880 self.capture_keep.lock().unwrap().clear();
19888 let was_tracking = self.gpu.ctx.is_event_tracking();
19889 if was_tracking {
19890 unsafe {
19891 self.gpu.ctx.disable_event_tracking();
19892 }
19893 }
19894 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
19895 self.capture_keep_on
19896 .store(true, std::sync::atomic::Ordering::Relaxed);
19897 let w = (|| {
19898 step(self)?;
19899 step(self)
19900 })();
19901 self.capture_keep_on
19902 .store(false, std::sync::atomic::Ordering::Relaxed);
19903 w?;
19904 self.gpu.stream().synchronize()?;
19905 self.gpu
19906 .stream()
19907 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
19908 let r = step(self);
19909 let g = self.gpu.stream().end_capture(flags);
19910 r?;
19911 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
19912 graph.upload()?;
19913 Ok(graph)
19914 };
19915 let result = run();
19916 self.capture_keep_on
19917 .store(false, std::sync::atomic::Ordering::Relaxed);
19918 if was_tracking {
19919 unsafe {
19920 self.gpu.ctx.enable_event_tracking();
19921 }
19922 }
19923 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
19924 Ok((result?, keeper))
19925 }
19926
19927 pub fn capture_graph<F>(
19928 &self,
19929 mut step: F,
19930 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
19931 where
19932 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
19933 {
19934 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
19935 let was_tracking = self.gpu.ctx.is_event_tracking();
19943 if was_tracking {
19944 unsafe {
19945 self.gpu.ctx.disable_event_tracking();
19946 }
19947 }
19948 let iflag = {
19955 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
19956 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
19957 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
19960 Ok("priority") => {
19961 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
19962 }
19963 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
19964 })
19965 };
19966 let ct = {
19973 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19974 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
19975 };
19976 let warmups = {
19999 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20000 *W.get_or_init(|| {
20001 std::env::var("MEMRA_GRAPH_WARMUPS")
20002 .ok()
20003 .and_then(|v| v.parse().ok())
20004 .filter(|n| *n >= 1)
20005 .unwrap_or(1)
20006 })
20007 };
20008 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
20009 let t_w = std::time::Instant::now();
20010 for _ in 0..warmups {
20012 step(self)?;
20013 }
20014 self.gpu.stream().synchronize()?;
20015 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
20016 let t_c = std::time::Instant::now();
20018 self.gpu
20019 .stream()
20020 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
20021 let r = step(self);
20024 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
20025 let t_i = std::time::Instant::now();
20026 let g = self.gpu.stream().end_capture(iflag);
20027 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
20028 r?;
20029 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
20030 let t_u = std::time::Instant::now();
20031 graph.upload()?;
20032 if ct {
20033 println!(
20034 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
20035 instantiate {ms_inst:.2} ms upload {:.2} ms",
20036 t_u.elapsed().as_secs_f64() * 1e3
20037 );
20038 }
20039 Ok(graph)
20040 };
20041 let result = run();
20042 if was_tracking {
20043 unsafe {
20044 self.gpu.ctx.enable_event_tracking();
20045 }
20046 }
20047 result
20048 }
20049
20050 pub fn gdn_scan_s128_view(
20052 &self,
20053 q: &CudaSlice<f32>,
20054 k: &CudaSlice<f32>,
20055 v: &CudaSlice<f32>,
20056 g: &CudaSlice<f32>,
20057 beta: &CudaSlice<f32>,
20058 state_in: &cudarc::driver::CudaView<f32>,
20059 state_out: &mut cudarc::driver::CudaViewMut<f32>,
20060 o: &mut CudaSlice<f32>,
20061 n_head: usize,
20062 t: usize,
20063 scale: f32,
20064 ) -> Result<(), Box<dyn std::error::Error>> {
20065 let f = self.func("gdn_scan_s128");
20066 const S_V: u32 = 128;
20067 const WARP: u32 = 32;
20068 const COLS: u32 = 4;
20069 let cfg = LaunchConfig {
20070 grid_dim: (n_head as u32, 1, S_V / COLS),
20071 block_dim: (WARP, COLS, 1),
20072 shared_mem_bytes: 0,
20073 };
20074 let (h, ti) = (n_head as i32, t as i32);
20075 let __s_b = self.gpu.stream();
20076 let mut b = __s_b.launch_builder(&f);
20077 b.arg(q)
20078 .arg(k)
20079 .arg(v)
20080 .arg(g)
20081 .arg(beta)
20082 .arg(state_in)
20083 .arg(state_out)
20084 .arg(o)
20085 .arg(&h)
20086 .arg(&ti)
20087 .arg(&scale);
20088 unsafe {
20089 b.launch(cfg)?;
20090 }
20091 Ok(())
20092 }
20093
20094 pub fn ssm_conv1d_view(
20096 &self,
20097 x: &cudarc::driver::CudaView<f32>,
20098 w: &CudaSlice<f32>,
20099 y: &mut CudaSlice<f32>,
20100 conv_dim: usize,
20101 t: usize,
20102 d_conv: usize,
20103 silu: bool,
20104 ) -> Result<(), Box<dyn std::error::Error>> {
20105 let f = self.func("ssm_conv1d_silu_f32");
20106 let cfg = LaunchConfig {
20108 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20109 block_dim: (256, 1, 1),
20110 shared_mem_bytes: 0,
20111 };
20112 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20113 let __s_b = self.gpu.stream();
20114 let mut b = __s_b.launch_builder(&f);
20115 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20116 unsafe {
20117 b.launch(cfg)?;
20118 }
20119 Ok(())
20120 }
20121
20122 pub fn ssm_conv1d_tm(
20129 &self,
20130 qkv_tm: &CudaSlice<f32>,
20131 w: &CudaSlice<f32>,
20132 y: &mut CudaSlice<f32>,
20133 conv_dim: usize,
20134 t: usize,
20135 d_conv: usize,
20136 ) -> Result<(), Box<dyn std::error::Error>> {
20137 let f = self.func("ssm_conv1d_tm_f32");
20138 let cfg = LaunchConfig {
20139 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20140 block_dim: (256, 1, 1),
20141 shared_mem_bytes: 0,
20142 };
20143 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20144 let __s_b = self.gpu.stream();
20145 let mut b = __s_b.launch_builder(&f);
20146 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
20147 unsafe {
20148 b.launch(cfg)?;
20149 }
20150 Ok(())
20151 }
20152
20153 pub fn ssm_conv1d_tm_state(
20161 &self,
20162 qkv_tm: &CudaSlice<f32>,
20163 conv_state: &mut CudaSlice<f32>,
20164 w: &CudaSlice<f32>,
20165 y: &mut CudaSlice<f32>,
20166 conv_dim: usize,
20167 t: usize,
20168 d_conv: usize,
20169 ) -> Result<(), Box<dyn std::error::Error>> {
20170 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
20171 }
20172
20173 #[allow(clippy::too_many_arguments)]
20176 pub fn ssm_conv1d_tm_state_pad(
20177 &self,
20178 qkv_tm: &CudaSlice<f32>,
20179 conv_state: &mut CudaSlice<f32>,
20180 w: &CudaSlice<f32>,
20181 y: &mut CudaSlice<f32>,
20182 conv_dim: usize,
20183 t: usize,
20184 d_conv: usize,
20185 pad_len: Option<&CudaSlice<i32>>,
20186 ) -> Result<(), Box<dyn std::error::Error>> {
20187 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20188 let ring_old = if t < d_conv - 1 {
20192 Some(self.clone_dtod(conv_state)?)
20193 } else {
20194 None
20195 };
20196 {
20197 let f = self.func("ssm_conv1d_tm_state_f32");
20198 let cfg = LaunchConfig {
20199 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20200 block_dim: (256, 1, 1),
20201 shared_mem_bytes: 0,
20202 };
20203 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20204 let __s_b = self.gpu.stream();
20205 let mut b = __s_b.launch_builder(&f);
20206 b.arg(qkv_tm)
20207 .arg(&*conv_state)
20208 .arg(w)
20209 .arg(y)
20210 .arg(&cd)
20211 .arg(&ti)
20212 .arg(&dc);
20213 unsafe {
20214 b.launch(cfg)?;
20215 }
20216 }
20217 match (ring_old, pad_len) {
20218 (None, Some(len_d)) => {
20219 let f = self.func("ssm_conv_ring_update_dev_f32");
20220 let n = conv_dim * (d_conv - 1);
20221 let cfg = LaunchConfig::for_num_elems(n as u32);
20222 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20223 let __s_b = self.gpu.stream();
20224 let mut b = __s_b.launch_builder(&f);
20225 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20226 unsafe {
20227 b.launch(cfg)?;
20228 }
20229 }
20230 (None, None) => {
20231 let f = self.func("ssm_conv_ring_update_f32");
20232 let n = conv_dim * (d_conv - 1);
20233 let cfg = LaunchConfig::for_num_elems(n as u32);
20234 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20235 let __s_b = self.gpu.stream();
20236 let mut b = __s_b.launch_builder(&f);
20237 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20238 unsafe {
20239 b.launch(cfg)?;
20240 }
20241 }
20242 (Some(old), _) => {
20243 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
20244 }
20245 }
20246 Ok(())
20247 }
20248
20249 pub fn ssm_conv1d_tm_state_pad_v(
20251 &self,
20252 qkv_tm: &cudarc::driver::CudaView<f32>,
20253 conv_state: &mut CudaSlice<f32>,
20254 w: &CudaSlice<f32>,
20255 y: &mut CudaSlice<f32>,
20256 conv_dim: usize,
20257 t: usize,
20258 d_conv: usize,
20259 pad_len: Option<&CudaSlice<i32>>,
20260 ) -> Result<(), Box<dyn std::error::Error>> {
20261 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
20262 let ring_old = if t < d_conv - 1 {
20266 Some(self.clone_dtod(conv_state)?)
20267 } else {
20268 None
20269 };
20270 {
20271 let f = self.func("ssm_conv1d_tm_state_f32");
20272 let cfg = LaunchConfig {
20273 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20274 block_dim: (256, 1, 1),
20275 shared_mem_bytes: 0,
20276 };
20277 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20278 let __s_b = self.gpu.stream();
20279 let mut b = __s_b.launch_builder(&f);
20280 b.arg(qkv_tm)
20281 .arg(&*conv_state)
20282 .arg(w)
20283 .arg(y)
20284 .arg(&cd)
20285 .arg(&ti)
20286 .arg(&dc);
20287 unsafe {
20288 b.launch(cfg)?;
20289 }
20290 }
20291 match (ring_old, pad_len) {
20292 (None, Some(len_d)) => {
20293 let f = self.func("ssm_conv_ring_update_dev_f32");
20294 let n = conv_dim * (d_conv - 1);
20295 let cfg = LaunchConfig::for_num_elems(n as u32);
20296 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20297 let __s_b = self.gpu.stream();
20298 let mut b = __s_b.launch_builder(&f);
20299 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
20300 unsafe {
20301 b.launch(cfg)?;
20302 }
20303 }
20304 (None, None) => {
20305 let f = self.func("ssm_conv_ring_update_f32");
20306 let n = conv_dim * (d_conv - 1);
20307 let cfg = LaunchConfig::for_num_elems(n as u32);
20308 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20309 let __s_b = self.gpu.stream();
20310 let mut b = __s_b.launch_builder(&f);
20311 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
20312 unsafe {
20313 b.launch(cfg)?;
20314 }
20315 }
20316 (Some(_), _) => unreachable!(
20317 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
20318 ),
20319 }
20320 Ok(())
20321 }
20322
20323 pub fn ssm_conv_ring_rebuild(
20328 &self,
20329 qkv_tm: &CudaSlice<f32>,
20330 ring_old: &CudaSlice<f32>,
20331 conv_state: &mut CudaSlice<f32>,
20332 conv_dim: usize,
20333 tc: usize,
20334 d_conv: usize,
20335 ) -> Result<(), Box<dyn std::error::Error>> {
20336 let f = self.func("ssm_conv_ring_rebuild_f32");
20337 let n = conv_dim * (d_conv - 1);
20338 let cfg = LaunchConfig::for_num_elems(n as u32);
20339 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
20340 let __s_b = self.gpu.stream();
20341 let mut b = __s_b.launch_builder(&f);
20342 b.arg(qkv_tm)
20343 .arg(ring_old)
20344 .arg(conv_state)
20345 .arg(&cd)
20346 .arg(&ti)
20347 .arg(&dc);
20348 unsafe {
20349 b.launch(cfg)?;
20350 }
20351 Ok(())
20352 }
20353
20354 #[allow(clippy::too_many_arguments)]
20359 pub fn gdn_prep_decode(
20360 &self,
20361 conv_out: &CudaSlice<f32>,
20362 beta_raw: &CudaSlice<f32>,
20363 alpha: &CudaSlice<f32>,
20364 dt_bias: &CudaSlice<f32>,
20365 a: &CudaSlice<f32>,
20366 q_l2: &mut CudaSlice<f32>,
20367 k_l2: &mut CudaSlice<f32>,
20368 v_g: &mut CudaSlice<f32>,
20369 beta: &mut CudaSlice<f32>,
20370 g_log: &mut CudaSlice<f32>,
20371 d_state: usize,
20372 num_v: usize,
20373 num_k: usize,
20374 key_dim: usize,
20375 eps: f32,
20376 ) -> Result<(), Box<dyn std::error::Error>> {
20377 let f = self.func("gdn_prep_decode_f32");
20378 let cfg = LaunchConfig {
20379 grid_dim: (num_v as u32, 1, 1),
20380 block_dim: (32, 4, 1),
20381 shared_mem_bytes: 0,
20382 };
20383 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20384 let __s_b = self.gpu.stream();
20385 let mut b = __s_b.launch_builder(&f);
20386 b.arg(conv_out)
20387 .arg(beta_raw)
20388 .arg(alpha)
20389 .arg(dt_bias)
20390 .arg(a)
20391 .arg(q_l2)
20392 .arg(k_l2)
20393 .arg(v_g)
20394 .arg(beta)
20395 .arg(g_log)
20396 .arg(&ds)
20397 .arg(&nv)
20398 .arg(&nk)
20399 .arg(&kd)
20400 .arg(&eps);
20401 unsafe {
20402 b.launch(cfg)?;
20403 }
20404 Ok(())
20405 }
20406
20407 #[allow(clippy::too_many_arguments)]
20411 pub fn ssm_conv1d_gdn(
20412 &self,
20413 qkv_tm: &CudaSlice<f32>,
20414 w: &CudaSlice<f32>,
20415 q_g: &mut CudaSlice<f32>,
20416 k_g: &mut CudaSlice<f32>,
20417 v_g: &mut CudaSlice<f32>,
20418 conv_dim: usize,
20419 t: usize,
20420 d_conv: usize,
20421 d_state: usize,
20422 num_v: usize,
20423 num_k: usize,
20424 key_dim: usize,
20425 ) -> Result<(), Box<dyn std::error::Error>> {
20426 let f = self.func("ssm_conv1d_gdn_f32");
20427 let cfg = LaunchConfig {
20428 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
20429 block_dim: (256, 1, 1),
20430 shared_mem_bytes: 0,
20431 };
20432 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
20433 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
20434 let __s_b = self.gpu.stream();
20435 let mut b = __s_b.launch_builder(&f);
20436 b.arg(qkv_tm)
20437 .arg(w)
20438 .arg(q_g)
20439 .arg(k_g)
20440 .arg(v_g)
20441 .arg(&cd)
20442 .arg(&ti)
20443 .arg(&dc)
20444 .arg(&ds)
20445 .arg(&nv)
20446 .arg(&nk)
20447 .arg(&kd);
20448 unsafe {
20449 b.launch(cfg)?;
20450 }
20451 Ok(())
20452 }
20453
20454 pub fn ssm_conv1d(
20455 &self,
20456 x: &CudaSlice<f32>,
20457 w: &CudaSlice<f32>,
20458 y: &mut CudaSlice<f32>,
20459 conv_dim: usize,
20460 t: usize,
20461 d_conv: usize,
20462 silu: bool,
20463 ) -> Result<(), Box<dyn std::error::Error>> {
20464 let f = self.func("ssm_conv1d_silu_f32");
20465 let cfg = LaunchConfig {
20466 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
20467 block_dim: (256, 1, 1),
20468 shared_mem_bytes: 0,
20469 };
20470 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
20471 let __s_b = self.gpu.stream();
20472 let mut b = __s_b.launch_builder(&f);
20473 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
20474 unsafe {
20475 b.launch(cfg)?;
20476 }
20477 Ok(())
20478 }
20479
20480 pub fn gdn_scan_s128(
20483 &self,
20484 q: &CudaSlice<f32>,
20485 k: &CudaSlice<f32>,
20486 v: &CudaSlice<f32>,
20487 g: &CudaSlice<f32>,
20488 beta: &CudaSlice<f32>,
20489 state_in: &CudaSlice<f32>,
20490 state_out: &mut CudaSlice<f32>,
20491 o: &mut CudaSlice<f32>,
20492 n_head: usize,
20493 t: usize,
20494 scale: f32,
20495 ) -> Result<(), Box<dyn std::error::Error>> {
20496 let f = self.func("gdn_scan_s128");
20497 const S_V: u32 = 128;
20498 const WARP: u32 = 32;
20499 const COLS_PER_BLOCK: u32 = 4;
20500 let cfg = LaunchConfig {
20501 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
20502 block_dim: (WARP, COLS_PER_BLOCK, 1),
20503 shared_mem_bytes: 0,
20504 };
20505 let (h, ti) = (n_head as i32, t as i32);
20506 let __s_b = self.gpu.stream();
20507 let mut b = __s_b.launch_builder(&f);
20508 b.arg(q)
20509 .arg(k)
20510 .arg(v)
20511 .arg(g)
20512 .arg(beta)
20513 .arg(state_in)
20514 .arg(state_out)
20515 .arg(o)
20516 .arg(&h)
20517 .arg(&ti)
20518 .arg(&scale);
20519 unsafe {
20520 b.launch(cfg)?;
20521 }
20522 Ok(())
20523 }
20524
20525 #[allow(clippy::too_many_arguments)]
20530 pub fn ssm_conv1d_fused_decode_b(
20531 &self,
20532 qkv_cols: &CudaSlice<f32>,
20533 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20534 w: &CudaSlice<f32>,
20535 conv_outs: &mut CudaSlice<f32>,
20536 conv_dim: usize,
20537 d_conv: usize,
20538 b_n: usize,
20539 ) -> Result<(), Box<dyn std::error::Error>> {
20540 let f = self.func("ssm_conv1d_fused_decode_b_f32");
20541 let cfg = LaunchConfig {
20542 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20543 block_dim: (256, 1, 1),
20544 shared_mem_bytes: 0,
20545 };
20546 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20547 let __s_b = self.gpu.stream();
20548 let mut b = __s_b.launch_builder(&f);
20549 b.arg(qkv_cols)
20550 .arg(conv_state_ptrs)
20551 .arg(w)
20552 .arg(conv_outs)
20553 .arg(&cd)
20554 .arg(&dc);
20555 unsafe {
20556 b.launch(cfg)?;
20557 }
20558 Ok(())
20559 }
20560
20561 #[allow(clippy::too_many_arguments)]
20562 pub fn gdn_prep_decode_b(
20563 &self,
20564 conv_outs: &CudaSlice<f32>,
20565 beta_raws: &CudaSlice<f32>,
20566 alphas: &CudaSlice<f32>,
20567 dt_bias: &CudaSlice<f32>,
20568 a: &CudaSlice<f32>,
20569 q_l2: &mut CudaSlice<f32>,
20570 k_l2: &mut CudaSlice<f32>,
20571 v_g: &mut CudaSlice<f32>,
20572 beta: &mut CudaSlice<f32>,
20573 g_log: &mut CudaSlice<f32>,
20574 d_state: usize,
20575 num_v: usize,
20576 num_k: usize,
20577 key_dim: usize,
20578 eps: f32,
20579 conv_dim: usize,
20580 b_n: usize,
20581 ) -> Result<(), Box<dyn std::error::Error>> {
20582 let f = self.func("gdn_prep_decode_b_f32");
20583 let cfg = LaunchConfig {
20584 grid_dim: (num_v as u32, 1, b_n as u32),
20585 block_dim: (32, 4, 1),
20586 shared_mem_bytes: 0,
20587 };
20588 let (ds, nv, nk, kd, cd) = (
20589 d_state as i32,
20590 num_v as i32,
20591 num_k as i32,
20592 key_dim as i32,
20593 conv_dim as i32,
20594 );
20595 let __s_b = self.gpu.stream();
20596 let mut b = __s_b.launch_builder(&f);
20597 b.arg(conv_outs)
20598 .arg(beta_raws)
20599 .arg(alphas)
20600 .arg(dt_bias)
20601 .arg(a)
20602 .arg(q_l2)
20603 .arg(k_l2)
20604 .arg(v_g)
20605 .arg(beta)
20606 .arg(g_log)
20607 .arg(&ds)
20608 .arg(&nv)
20609 .arg(&nk)
20610 .arg(&kd)
20611 .arg(&eps)
20612 .arg(&cd);
20613 unsafe {
20614 b.launch(cfg)?;
20615 }
20616 Ok(())
20617 }
20618
20619 #[allow(clippy::too_many_arguments)]
20620 pub fn gdn_scan_s128_batched(
20621 &self,
20622 q: &CudaSlice<f32>,
20623 k: &CudaSlice<f32>,
20624 v: &CudaSlice<f32>,
20625 g: &CudaSlice<f32>,
20626 beta: &CudaSlice<f32>,
20627 state_in_ptrs: &cudarc::driver::CudaView<u64>,
20628 state_out_ptrs: &cudarc::driver::CudaView<u64>,
20629 o: &mut CudaSlice<f32>,
20630 n_head: usize,
20631 b_n: usize,
20632 scale: f32,
20633 ) -> Result<(), Box<dyn std::error::Error>> {
20634 let f = self.func("gdn_scan_s128_b");
20635 const S_V: u32 = 128;
20636 const WARP: u32 = 32;
20637 const COLS_PER_BLOCK: u32 = 4;
20638 let cfg = LaunchConfig {
20639 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20640 block_dim: (WARP, COLS_PER_BLOCK, 1),
20641 shared_mem_bytes: 0,
20642 };
20643 let h = n_head as i32;
20644 let __s_b = self.gpu.stream();
20645 let mut b = __s_b.launch_builder(&f);
20646 b.arg(q)
20647 .arg(k)
20648 .arg(v)
20649 .arg(g)
20650 .arg(beta)
20651 .arg(state_in_ptrs)
20652 .arg(state_out_ptrs)
20653 .arg(o)
20654 .arg(&h)
20655 .arg(&scale);
20656 unsafe {
20657 b.launch(cfg)?;
20658 }
20659 Ok(())
20660 }
20661
20662 #[allow(clippy::too_many_arguments)]
20668 pub fn ssm_conv1d_fused_decode_b_view(
20669 &self,
20670 qkv_cols: &cudarc::driver::CudaView<f32>,
20671 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
20672 w: &CudaSlice<f32>,
20673 conv_outs: &mut CudaSlice<f32>,
20674 conv_dim: usize,
20675 d_conv: usize,
20676 b_n: usize,
20677 ) -> Result<(), Box<dyn std::error::Error>> {
20678 let f = self.func("ssm_conv1d_fused_decode_b_f32");
20679 let cfg = LaunchConfig {
20680 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
20681 block_dim: (256, 1, 1),
20682 shared_mem_bytes: 0,
20683 };
20684 let (cd, dc) = (conv_dim as i32, d_conv as i32);
20685 let __s_b = self.gpu.stream();
20686 let mut b = __s_b.launch_builder(&f);
20687 b.arg(qkv_cols)
20688 .arg(conv_state_ptrs)
20689 .arg(w)
20690 .arg(conv_outs)
20691 .arg(&cd)
20692 .arg(&dc);
20693 unsafe {
20694 b.launch(cfg)?;
20695 }
20696 Ok(())
20697 }
20698
20699 #[allow(clippy::too_many_arguments)]
20700 pub fn gdn_prep_decode_b_view(
20701 &self,
20702 conv_outs: &CudaSlice<f32>,
20703 beta_raws: &cudarc::driver::CudaView<f32>,
20704 alphas: &cudarc::driver::CudaView<f32>,
20705 dt_bias: &CudaSlice<f32>,
20706 a: &CudaSlice<f32>,
20707 q_l2: &mut CudaSlice<f32>,
20708 k_l2: &mut CudaSlice<f32>,
20709 v_g: &mut CudaSlice<f32>,
20710 beta: &mut CudaSlice<f32>,
20711 g_log: &mut CudaSlice<f32>,
20712 d_state: usize,
20713 num_v: usize,
20714 num_k: usize,
20715 key_dim: usize,
20716 eps: f32,
20717 conv_dim: usize,
20718 b_n: usize,
20719 ) -> Result<(), Box<dyn std::error::Error>> {
20720 let f = self.func("gdn_prep_decode_b_f32");
20721 let cfg = LaunchConfig {
20722 grid_dim: (num_v as u32, 1, b_n as u32),
20723 block_dim: (32, 4, 1),
20724 shared_mem_bytes: 0,
20725 };
20726 let (ds, nv, nk, kd, cd) = (
20727 d_state as i32,
20728 num_v as i32,
20729 num_k as i32,
20730 key_dim as i32,
20731 conv_dim as i32,
20732 );
20733 let __s_b = self.gpu.stream();
20734 let mut b = __s_b.launch_builder(&f);
20735 b.arg(conv_outs)
20736 .arg(beta_raws)
20737 .arg(alphas)
20738 .arg(dt_bias)
20739 .arg(a)
20740 .arg(q_l2)
20741 .arg(k_l2)
20742 .arg(v_g)
20743 .arg(beta)
20744 .arg(g_log)
20745 .arg(&ds)
20746 .arg(&nv)
20747 .arg(&nk)
20748 .arg(&kd)
20749 .arg(&eps)
20750 .arg(&cd);
20751 unsafe {
20752 b.launch(cfg)?;
20753 }
20754 Ok(())
20755 }
20756
20757 #[allow(clippy::too_many_arguments)]
20758 pub fn gdn_scan_s128_batched_view(
20759 &self,
20760 q: &CudaSlice<f32>,
20761 k: &CudaSlice<f32>,
20762 v: &CudaSlice<f32>,
20763 g: &CudaSlice<f32>,
20764 beta: &CudaSlice<f32>,
20765 state_in_ptrs: &cudarc::driver::CudaView<u64>,
20766 state_out_ptrs: &cudarc::driver::CudaView<u64>,
20767 o: &mut cudarc::driver::CudaViewMut<f32>,
20768 n_head: usize,
20769 b_n: usize,
20770 scale: f32,
20771 ) -> Result<(), Box<dyn std::error::Error>> {
20772 let f = self.func("gdn_scan_s128_b");
20773 const S_V: u32 = 128;
20774 const WARP: u32 = 32;
20775 const COLS_PER_BLOCK: u32 = 4;
20776 let cfg = LaunchConfig {
20777 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
20778 block_dim: (WARP, COLS_PER_BLOCK, 1),
20779 shared_mem_bytes: 0,
20780 };
20781 let h = n_head as i32;
20782 let __s_b = self.gpu.stream();
20783 let mut b = __s_b.launch_builder(&f);
20784 b.arg(q)
20785 .arg(k)
20786 .arg(v)
20787 .arg(g)
20788 .arg(beta)
20789 .arg(state_in_ptrs)
20790 .arg(state_out_ptrs)
20791 .arg(o)
20792 .arg(&h)
20793 .arg(&scale);
20794 unsafe {
20795 b.launch(cfg)?;
20796 }
20797 Ok(())
20798 }
20799
20800 pub fn gdn_chunked_enabled() -> bool {
20809 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20810 *E.get_or_init(|| {
20811 std::env::var("MEMRA_GDN_CHUNKED")
20812 .map(|v| v != "0")
20813 .unwrap_or(true)
20814 })
20815 }
20816
20817 pub fn gdn_chunk_size() -> usize {
20822 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20823 *C.get_or_init(|| {
20824 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
20825 .ok()
20826 .and_then(|v| v.parse().ok())
20827 .unwrap_or(32);
20828 c.clamp(32, 128) / 32 * 32
20829 })
20830 }
20831
20832 #[allow(clippy::too_many_arguments)]
20837 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
20840 #[allow(clippy::too_many_arguments)]
20841 pub fn gdn_chunk_k123(
20842 &self,
20843 q: &CudaSlice<f32>,
20844 k: &CudaSlice<f32>,
20845 v: &CudaSlice<f32>,
20846 g: &CudaSlice<f32>,
20847 beta: &CudaSlice<f32>,
20848 wb16: Option<&mut CudaSlice<u8>>,
20849 n_head: usize,
20850 t: usize,
20851 c: usize,
20852 hk: usize,
20853 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
20854 ) -> Result<
20855 (
20856 CudaSlice<f32>,
20857 CudaSlice<f32>,
20858 CudaSlice<f32>,
20859 CudaSlice<f32>,
20860 ),
20861 Box<dyn std::error::Error>,
20862 > {
20863 const D: usize = 128;
20864 let h = n_head;
20865 let nc = (t + c - 1) / c;
20866 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
20867 let mut gcum = self.uninit(t * h)?;
20868 let mut a = self.uninit(nc * h * c * c)?;
20869 let mut p = self.uninit(nc * h * c * c)?;
20870 let mut u = self.uninit(nc * h * c * D)?;
20871 let mut w = self.uninit(nc * h * c * D)?;
20872 {
20873 let f = self.func("gdn_chunk_cumgate_f32");
20875 let cfg = LaunchConfig {
20876 grid_dim: (nc as u32, h as u32, 1),
20877 block_dim: (32, 1, 1),
20878 shared_mem_bytes: 0,
20879 };
20880 let __s_b = self.gpu.stream();
20881 let mut b = __s_b.launch_builder(&f);
20882 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
20883 unsafe {
20884 b.launch(cfg)?;
20885 }
20886 }
20887 if let Some((qb, kb, pb)) = k2w {
20888 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
20891 let f = self.func("gdn_k2_wgmma");
20892 let cfg = LaunchConfig {
20893 grid_dim: (nc as u32, h as u32, 1),
20894 block_dim: (128, 1, 1),
20895 shared_mem_bytes: 0,
20896 };
20897 let hki = hk as i32;
20898 let __s_b = self.gpu.stream();
20899 let mut b = __s_b.launch_builder(&f);
20900 b.arg(qb)
20901 .arg(kb)
20902 .arg(&gcum)
20903 .arg(beta)
20904 .arg(&mut a)
20905 .arg(&mut *pb)
20906 .arg(&hi)
20907 .arg(&ti)
20908 .arg(&ci)
20909 .arg(&hki);
20910 unsafe {
20911 b.launch(cfg)?;
20912 }
20913 } else if c <= 64 && !portable_mma_gated() {
20914 let f = self.func("gdn_chunk_attn_f32");
20916 let jt = ((c + 31) / 32) as u32;
20917 let cfg = LaunchConfig {
20918 grid_dim: (nc as u32, h as u32, jt),
20919 block_dim: (256, 1, 1),
20920 shared_mem_bytes: 0,
20921 };
20922 let hki = hk as i32;
20923 let __s_b = self.gpu.stream();
20924 let mut b = __s_b.launch_builder(&f);
20925 b.arg(q)
20926 .arg(k)
20927 .arg(&gcum)
20928 .arg(beta)
20929 .arg(&mut a)
20930 .arg(&mut p)
20931 .arg(&hi)
20932 .arg(&ti)
20933 .arg(&ci)
20934 .arg(&hki);
20935 unsafe {
20936 b.launch(cfg)?;
20937 }
20938 } else {
20939 assert!(
20941 hk == h,
20942 "generic K2 is broadcast-only (de-broadcast rides C==32)"
20943 );
20944 let f = self.func("gdn_chunk_attn_g_f32");
20945 let cfg = LaunchConfig {
20946 grid_dim: (nc as u32, h as u32, 1),
20947 block_dim: (32, 8, 1),
20948 shared_mem_bytes: 0,
20949 };
20950 let __s_b = self.gpu.stream();
20951 let mut b = __s_b.launch_builder(&f);
20952 b.arg(q)
20953 .arg(k)
20954 .arg(&gcum)
20955 .arg(beta)
20956 .arg(&mut a)
20957 .arg(&mut p)
20958 .arg(&hi)
20959 .arg(&ti)
20960 .arg(&ci);
20961 unsafe {
20962 b.launch(cfg)?;
20963 }
20964 }
20965 {
20966 let cfg = LaunchConfig {
20968 grid_dim: (nc as u32, h as u32, 1),
20969 block_dim: (256, 1, 1),
20970 shared_mem_bytes: 0,
20971 };
20972 match c {
20973 32 | 64 => {
20974 let f = self.func(if c == 32 {
20975 "gdn_chunk_solve32_f32"
20976 } else {
20977 "gdn_chunk_solve64_f32"
20978 });
20979 let wb: u64 = match wb16 {
20981 Some(d) => self.addr_u8(d),
20982 None => 0,
20983 };
20984 let hki = hk as i32;
20985 let __s_b = self.gpu.stream();
20986 let mut b = __s_b.launch_builder(&f);
20987 b.arg(v)
20988 .arg(k)
20989 .arg(&a)
20990 .arg(&gcum)
20991 .arg(&mut u)
20992 .arg(&mut w)
20993 .arg(&wb)
20994 .arg(&hi)
20995 .arg(&ti)
20996 .arg(&hki);
20997 unsafe {
20998 b.launch(cfg)?;
20999 }
21000 }
21001 _ => {
21002 assert!(hk == h, "generic K3 is broadcast-only");
21003 let f = self.func("gdn_chunk_solve_f32");
21004 let __s_b = self.gpu.stream();
21005 let mut b = __s_b.launch_builder(&f);
21006 b.arg(v)
21007 .arg(k)
21008 .arg(&a)
21009 .arg(&gcum)
21010 .arg(&mut u)
21011 .arg(&mut w)
21012 .arg(&hi)
21013 .arg(&ti)
21014 .arg(&ci);
21015 unsafe {
21016 b.launch(cfg)?;
21017 }
21018 }
21019 }
21020 }
21021 Ok((gcum, p, u, w))
21022 }
21023
21024 pub fn gdn_db_on() -> bool {
21028 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
21029 }
21030
21031 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
21040 !portable_mma_gated()
21041 && c == 32
21042 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21043 Ok("1") => true,
21044 Ok("0") => false,
21045 _ => gdn_mma_default_on(),
21046 }
21047 }
21048
21049 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
21056 cfg!(memra_hopper_mma)
21057 && self.gdn_mma_enabled(c)
21058 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
21059 }
21060
21061 #[allow(clippy::too_many_arguments)]
21066 pub fn ssm_conv1d_gdn_state_pad(
21067 &self,
21068 qkv_tm: &cudarc::driver::CudaView<f32>,
21069 conv_state: &mut CudaSlice<f32>,
21070 w: &CudaSlice<f32>,
21071 q_g: &mut CudaSlice<f32>,
21072 k_g: &mut CudaSlice<f32>,
21073 v_g: &mut CudaSlice<f32>,
21074 conv_dim: usize,
21075 t: usize,
21076 d_conv: usize,
21077 d_state: usize,
21078 num_v: usize,
21079 num_k: usize,
21080 key_dim: usize,
21081 hk: usize,
21082 pad_len: Option<&CudaSlice<i32>>,
21083 ) -> Result<(), Box<dyn std::error::Error>> {
21084 assert!(
21085 t >= d_conv - 1,
21086 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
21087 );
21088 {
21089 let f = self.func("ssm_conv1d_gdn_state_f32");
21090 let cfg = LaunchConfig {
21091 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
21092 block_dim: (256, 1, 1),
21093 shared_mem_bytes: 0,
21094 };
21095 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21096 let (ds, nv, nk, kd, hki) = (
21097 d_state as i32,
21098 num_v as i32,
21099 num_k as i32,
21100 key_dim as i32,
21101 hk as i32,
21102 );
21103 let __s_b = self.gpu.stream();
21104 let mut b = __s_b.launch_builder(&f);
21105 b.arg(qkv_tm)
21106 .arg(&*conv_state)
21107 .arg(w)
21108 .arg(q_g)
21109 .arg(k_g)
21110 .arg(v_g)
21111 .arg(&cd)
21112 .arg(&ti)
21113 .arg(&dc)
21114 .arg(&ds)
21115 .arg(&nv)
21116 .arg(&nk)
21117 .arg(&kd)
21118 .arg(&hki);
21119 unsafe {
21120 b.launch(cfg)?;
21121 }
21122 }
21123 match pad_len {
21124 Some(len_d) => {
21125 let f = self.func("ssm_conv_ring_update_dev_f32");
21126 let n = conv_dim * (d_conv - 1);
21127 let cfg = LaunchConfig::for_num_elems(n as u32);
21128 let (cd, dc) = (conv_dim as i32, d_conv as i32);
21129 let __s_b = self.gpu.stream();
21130 let mut b = __s_b.launch_builder(&f);
21131 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
21132 unsafe {
21133 b.launch(cfg)?;
21134 }
21135 }
21136 None => {
21137 let f = self.func("ssm_conv_ring_update_f32");
21138 let n = conv_dim * (d_conv - 1);
21139 let cfg = LaunchConfig::for_num_elems(n as u32);
21140 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
21141 let __s_b = self.gpu.stream();
21142 let mut b = __s_b.launch_builder(&f);
21143 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
21144 unsafe {
21145 b.launch(cfg)?;
21146 }
21147 }
21148 }
21149 Ok(())
21150 }
21151
21152 pub fn gdn_chunk_alloc(
21156 &self,
21157 n_head: usize,
21158 t: usize,
21159 c: usize,
21160 hk: usize,
21161 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
21162 const D: usize = 128;
21163 assert!(
21164 c == 32,
21165 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
21166 );
21167 let h = n_head;
21168 let nc = (t + c - 1) / c;
21169 Ok(GdnChunkBufs {
21170 gcum: self.uninit(t * h)?,
21171 a: self.uninit(nc * h * c * c)?,
21172 p: self.uninit(nc * h * c * c)?,
21173 u: self.uninit(nc * h * c * D)?,
21174 w: self.uninit(nc * h * c * D)?,
21175 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21176 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21177 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
21178 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
21179 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
21180 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
21181 o: self.uninit(D * h * t)?,
21182 t,
21183 nc,
21184 })
21185 }
21186
21187 pub fn f32_to_bf16_v(
21189 &self,
21190 x: &cudarc::driver::CudaView<f32>,
21191 dst: &mut CudaSlice<u8>,
21192 n: usize,
21193 ) -> Result<(), Box<dyn std::error::Error>> {
21194 let f = self.func("f32_to_bf16_bulk");
21195 let ni = n as i64;
21196 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21197 let __s_b = self.gpu.stream();
21198 let mut b = __s_b.launch_builder(&f);
21199 b.arg(x).arg(dst).arg(&ni);
21200 unsafe {
21201 b.launch(cfg)?;
21202 }
21203 Ok(())
21204 }
21205
21206 pub fn f32_to_bf16_into(
21208 &self,
21209 x: &CudaSlice<f32>,
21210 dst: &mut CudaSlice<u8>,
21211 n: usize,
21212 ) -> Result<(), Box<dyn std::error::Error>> {
21213 let f = self.func("f32_to_bf16_bulk");
21214 let ni = n as i64;
21215 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
21216 let __s_b = self.gpu.stream();
21217 let mut b = __s_b.launch_builder(&f);
21218 b.arg(x).arg(dst).arg(&ni);
21219 unsafe {
21220 b.launch(cfg)?;
21221 }
21222 Ok(())
21223 }
21224
21225 pub fn gdn_chunk_k123_vl8(
21228 &self,
21229 seqs: &[GdnSeqVl],
21230 n_head: usize,
21231 hk: usize,
21232 wq: Option<&GdnWVl8>,
21233 ) -> Result<(), Box<dyn std::error::Error>> {
21234 let b = seqs.len();
21235 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
21236 let mut packed = [GdnSeqVl::default(); 8];
21237 packed[..b].copy_from_slice(seqs);
21238 let v = GdnVl8(packed);
21239 let (hi, ci) = (n_head as i32, 32i32);
21240 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21241 {
21242 let f = self.func("gdn_chunk_cumgate_vl");
21243 let cfg = LaunchConfig {
21244 grid_dim: (max_nc, n_head as u32, b as u32),
21245 block_dim: (32, 1, 1),
21246 shared_mem_bytes: 0,
21247 };
21248 let __s_lb = self.gpu.stream();
21249 let mut lb = __s_lb.launch_builder(&f);
21250 lb.arg(&v).arg(&hi).arg(&ci);
21251 unsafe {
21252 lb.launch(cfg)?;
21253 }
21254 }
21255 let hki = hk as i32;
21256 if let Some(w) = wq {
21257 let f = self.func("gdn_k2_wgmma_vl");
21259 let cfg = LaunchConfig {
21260 grid_dim: (max_nc, n_head as u32, b as u32),
21261 block_dim: (128, 1, 1),
21262 shared_mem_bytes: 0,
21263 };
21264 let __s_lb = self.gpu.stream();
21265 let mut lb = __s_lb.launch_builder(&f);
21266 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
21267 unsafe {
21268 lb.launch(cfg)?;
21269 }
21270 } else {
21271 let f = self.func("gdn_chunk_attn_vl");
21272 let cfg = LaunchConfig {
21273 grid_dim: (max_nc, n_head as u32, b as u32),
21274 block_dim: (256, 1, 1),
21275 shared_mem_bytes: 0,
21276 };
21277 let __s_lb = self.gpu.stream();
21278 let mut lb = __s_lb.launch_builder(&f);
21279 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21280 unsafe {
21281 lb.launch(cfg)?;
21282 }
21283 }
21284 {
21285 let f = self.func("gdn_chunk_solve32_vl");
21286 let cfg = LaunchConfig {
21287 grid_dim: (max_nc, n_head as u32, b as u32),
21288 block_dim: (256, 1, 1),
21289 shared_mem_bytes: 0,
21290 };
21291 let __s_lb = self.gpu.stream();
21292 let mut lb = __s_lb.launch_builder(&f);
21293 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21294 unsafe {
21295 lb.launch(cfg)?;
21296 }
21297 }
21298 Ok(())
21299 }
21300
21301 #[allow(clippy::too_many_arguments)]
21305 pub fn gdn_prep_vl8(
21306 &self,
21307 seqs: &[GdnPrepVl],
21308 conv_w: &CudaSlice<f32>,
21309 dt_bias: &CudaSlice<f32>,
21310 a: &CudaSlice<f32>,
21311 conv_dim: usize,
21312 d_conv: usize,
21313 d_state: usize,
21314 num_v: usize,
21315 num_k: usize,
21316 key_dim: usize,
21317 hk: usize,
21318 eps: f32,
21319 ) -> Result<(), Box<dyn std::error::Error>> {
21320 let b = seqs.len();
21321 assert!(b >= 1 && b <= 8);
21322 let mut packed = [GdnPrepVl::default(); 8];
21323 packed[..b].copy_from_slice(seqs);
21324 let v = GdnPrepVl8(packed);
21325 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21326 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
21327 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
21328 assert!(
21329 conv_fuse || hk == num_v,
21330 "de-broadcast requires the fused conv"
21331 );
21332 if conv_fuse {
21333 let f = self.func("ssm_conv1d_gdn_state_vl");
21334 let cfg = LaunchConfig {
21335 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21336 block_dim: (256, 1, 1),
21337 shared_mem_bytes: 0,
21338 };
21339 let (dsi, nvi, nki, kdi, hki) = (
21340 d_state as i32,
21341 num_v as i32,
21342 num_k as i32,
21343 key_dim as i32,
21344 hk as i32,
21345 );
21346 let __s_lb = self.gpu.stream();
21347 let mut lb = __s_lb.launch_builder(&f);
21348 lb.arg(&v)
21349 .arg(conv_w)
21350 .arg(&cdi)
21351 .arg(&dci)
21352 .arg(&dsi)
21353 .arg(&nvi)
21354 .arg(&nki)
21355 .arg(&kdi)
21356 .arg(&hki);
21357 unsafe {
21358 lb.launch(cfg)?;
21359 }
21360 } else {
21361 let f = self.func("ssm_conv1d_tm_state_vl");
21362 let cfg = LaunchConfig {
21363 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
21364 block_dim: (256, 1, 1),
21365 shared_mem_bytes: 0,
21366 };
21367 let __s_lb = self.gpu.stream();
21368 let mut lb = __s_lb.launch_builder(&f);
21369 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
21370 unsafe {
21371 lb.launch(cfg)?;
21372 }
21373 }
21374 {
21375 let f = self.func("ssm_conv_ring_update_vl");
21376 let n = (conv_dim * (d_conv - 1)) as u32;
21377 let cfg = LaunchConfig {
21378 grid_dim: (n.div_ceil(256), 1, b as u32),
21379 block_dim: (256, 1, 1),
21380 shared_mem_bytes: 0,
21381 };
21382 let __s_lb = self.gpu.stream();
21383 let mut lb = __s_lb.launch_builder(&f);
21384 lb.arg(&v).arg(&cdi).arg(&dci);
21385 unsafe {
21386 lb.launch(cfg)?;
21387 }
21388 }
21389 if !conv_fuse {
21390 let f = self.func("qkv_to_gdn_repack_vl");
21391 let n = max_t * (num_v * d_state) as u32;
21392 let cfg = LaunchConfig {
21393 grid_dim: (n.div_ceil(256), 1, b as u32),
21394 block_dim: (256, 1, 1),
21395 shared_mem_bytes: 0,
21396 };
21397 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
21398 let __s_lb = self.gpu.stream();
21399 let mut lb = __s_lb.launch_builder(&f);
21400 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
21401 unsafe {
21402 lb.launch(cfg)?;
21403 }
21404 }
21405 if Self::l2_v2_on(d_state) {
21406 let f = self.func("gdn_l2_v2_vl");
21407 let cfg = LaunchConfig {
21408 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
21409 block_dim: (256, 1, 1),
21410 shared_mem_bytes: 0,
21411 };
21412 let (dsi, nvi) = (d_state as i32, hk as i32);
21413 let __s_lb = self.gpu.stream();
21414 let mut lb = __s_lb.launch_builder(&f);
21415 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21416 unsafe {
21417 lb.launch(cfg)?;
21418 }
21419 } else {
21420 let f = self.func("gdn_l2_vl");
21421 let cfg = LaunchConfig {
21422 grid_dim: (max_t * hk as u32, 2, b as u32),
21423 block_dim: (256, 1, 1),
21424 shared_mem_bytes: 0,
21425 };
21426 let (dsi, nvi) = (d_state as i32, hk as i32);
21427 let __s_lb = self.gpu.stream();
21428 let mut lb = __s_lb.launch_builder(&f);
21429 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
21430 unsafe {
21431 lb.launch(cfg)?;
21432 }
21433 }
21434 {
21435 let f = self.func("gdn_gate_prep_vl");
21436 let n = max_t * num_v as u32;
21437 let cfg = LaunchConfig {
21438 grid_dim: (n.div_ceil(256), 1, b as u32),
21439 block_dim: (256, 1, 1),
21440 shared_mem_bytes: 0,
21441 };
21442 let nvi = num_v as i32;
21443 let __s_lb = self.gpu.stream();
21444 let mut lb = __s_lb.launch_builder(&f);
21445 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
21446 unsafe {
21447 lb.launch(cfg)?;
21448 }
21449 }
21450 Ok(())
21451 }
21452
21453 pub fn gdn_mirror_vl8(
21455 &self,
21456 seqs: &[GdnSeqVl],
21457 n_head: usize,
21458 which: i32,
21459 hk: usize,
21460 ) -> Result<(), Box<dyn std::error::Error>> {
21461 let b = seqs.len();
21462 assert!(b >= 1 && b <= 8);
21463 let mut packed = [GdnSeqVl::default(); 8];
21464 packed[..b].copy_from_slice(seqs);
21465 let v = GdnVl8(packed);
21466 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
21467 let max_n = seqs
21468 .iter()
21469 .map(|s| {
21470 if which == 0 {
21471 s.t as i64 * ept as i64
21472 } else {
21473 s.nc as i64 * ept as i64 * 32
21474 }
21475 })
21476 .max()
21477 .unwrap();
21478 let f = self.func("gdn_mirror_vl");
21479 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21480 let cfg = LaunchConfig {
21481 grid_dim: (blocks, 1, b as u32),
21482 block_dim: (256, 1, 1),
21483 shared_mem_bytes: 0,
21484 };
21485 let __s_lb = self.gpu.stream();
21486 let mut lb = __s_lb.launch_builder(&f);
21487 lb.arg(&v).arg(&ept).arg(&which);
21488 unsafe {
21489 lb.launch(cfg)?;
21490 }
21491 Ok(())
21492 }
21493
21494 pub fn gdn_tail_vl8(
21496 &self,
21497 seqs: &[GdnPrepVl],
21498 norm_w: &CudaSlice<f32>,
21499 d_state: usize,
21500 num_v: usize,
21501 eps: f32,
21502 ) -> Result<(), Box<dyn std::error::Error>> {
21503 let b = seqs.len();
21504 assert!(b >= 1 && b <= 8);
21505 let mut packed = [GdnPrepVl::default(); 8];
21506 packed[..b].copy_from_slice(seqs);
21507 let v = GdnPrepVl8(packed);
21508 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21509 let f = self.func("gated_rmsnorm_f16out_vl");
21510 let cfg = LaunchConfig {
21512 grid_dim: (max_t * num_v as u32, 1, b as u32),
21513 block_dim: (128, 1, 1),
21514 shared_mem_bytes: 0,
21515 };
21516 let (dsi, nvi) = (d_state as i32, num_v as i32);
21517 let __s_lb = self.gpu.stream();
21518 let mut lb = __s_lb.launch_builder(&f);
21519 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
21520 unsafe {
21521 lb.launch(cfg)?;
21522 }
21523 Ok(())
21524 }
21525
21526 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
21529 use cudarc::driver::DevicePtr;
21530 let s = self.gpu.stream();
21531 let (p, _g) = x.device_ptr(&s);
21532 p as u64
21533 }
21534 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
21535 use cudarc::driver::DevicePtrMut;
21536 let s = self.gpu.stream();
21537 let (p, _g) = x.device_ptr_mut(&s);
21538 p as u64
21539 }
21540 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
21541 use cudarc::driver::DevicePtr;
21542 let s = self.gpu.stream();
21543 let (p, _g) = x.device_ptr(&s);
21544 p as u64
21545 }
21546 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
21547 use cudarc::driver::DevicePtr;
21548 let s = self.gpu.stream();
21549 let (p, _g) = x.device_ptr(&s);
21550 p as u64
21551 }
21552
21553 pub fn gdn_chunk_vl8(
21557 &self,
21558 seqs: &[GdnSeqVl],
21559 n_head: usize,
21560 scale: f32,
21561 hk: usize,
21562 wq: Option<&GdnWVl8>,
21563 ) -> Result<(), Box<dyn std::error::Error>> {
21564 const NSPLIT: u32 = 4;
21565 let b = seqs.len();
21566 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
21567 let mut packed = [GdnSeqVl::default(); 8];
21568 packed[..b].copy_from_slice(seqs);
21569 let v = GdnVl8(packed);
21570 let (hi, ci) = (n_head as i32, 32i32);
21571 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
21572 let hki = hk as i32;
21573 if let Some(w) = wq {
21574 let f = self.func("gdn_k45_wgmma_vl");
21576 let cfg = LaunchConfig {
21577 grid_dim: (n_head as u32, NSPLIT, b as u32),
21578 block_dim: (256, 1, 1),
21579 shared_mem_bytes: 0,
21580 };
21581 let __s_lb = self.gpu.stream();
21582 let mut lb = __s_lb.launch_builder(&f);
21583 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
21584 unsafe {
21585 lb.launch(cfg)?;
21586 }
21587 let _ = max_nc;
21588 return Ok(());
21589 }
21590 {
21591 let f = self.func("gdn_chunk_state_mma_vl");
21592 let cfg = LaunchConfig {
21593 grid_dim: (n_head as u32, NSPLIT, b as u32),
21594 block_dim: (256, 1, 1),
21595 shared_mem_bytes: 0,
21596 };
21597 let __s_lb = self.gpu.stream();
21598 let mut lb = __s_lb.launch_builder(&f);
21599 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
21600 unsafe {
21601 lb.launch(cfg)?;
21602 }
21603 }
21604 {
21605 let f = self.func("gdn_chunk_output_mma_vl");
21606 let cfg = LaunchConfig {
21607 grid_dim: (max_nc, n_head as u32, b as u32),
21608 block_dim: (256, 1, 1),
21609 shared_mem_bytes: 0,
21610 };
21611 let __s_lb = self.gpu.stream();
21612 let mut lb = __s_lb.launch_builder(&f);
21613 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
21614 unsafe {
21615 lb.launch(cfg)?;
21616 }
21617 }
21618 Ok(())
21619 }
21620 pub fn gdn_scan_chunked(
21621 &self,
21622 q: &CudaSlice<f32>,
21623 k: &CudaSlice<f32>,
21624 v: &CudaSlice<f32>,
21625 g: &CudaSlice<f32>,
21626 beta: &CudaSlice<f32>,
21627 kb16_pre: Option<&CudaSlice<u8>>,
21628 qb16_pre: Option<&CudaSlice<u8>>,
21629 state_in: &CudaSlice<f32>,
21630 state_out: &mut CudaSlice<f32>,
21631 o: &mut CudaSlice<f32>,
21632 n_head: usize,
21633 t: usize,
21634 scale: f32,
21635 c: usize,
21636 hk: usize,
21637 ) -> Result<(), Box<dyn std::error::Error>> {
21638 const D: usize = 128;
21639 const NSPLIT: u32 = 4;
21640 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
21641 let h = n_head;
21642 let nc = (t + c - 1) / c;
21643 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
21644 let gdn_mma_pre = !portable_mma_gated()
21649 && c == 32
21650 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21651 Ok("1") => true,
21652 Ok("0") => false,
21653 _ => gdn_mma_default_on(),
21654 };
21655 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
21656 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
21657 } else {
21658 None
21659 };
21660 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
21665 && gdn_mma_pre
21666 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
21667 let nk = t * hk * D;
21668 let mut kb16_local: Option<CudaSlice<u8>> = None;
21669 if gdn_mma_pre && kb16_pre.is_none() {
21670 let mut kb = self.alloc_u8_uninit(nk * 2)?;
21671 let f = self.func("f32_to_bf16_bulk");
21672 let n2 = nk as i64;
21673 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21674 let __s_b = self.gpu.stream();
21675 let mut b = __s_b.launch_builder(&f);
21676 b.arg(k).arg(&mut kb).arg(&n2);
21677 unsafe {
21678 b.launch(cfg2)?;
21679 }
21680 kb16_local = Some(kb);
21681 }
21682 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
21683 if let Some(kb) = kb16_pre {
21684 assert!(kb.len() >= nk * 2, "kb16_pre too small");
21685 }
21686 let mut qb16: Option<CudaSlice<u8>> = None;
21687 let mut pb16: Option<CudaSlice<u8>> = None;
21688 if gdn_wgmma_pre {
21689 if qb16_pre.is_none() {
21692 let mut qb = self.alloc_u8_uninit(nk * 2)?;
21693 let f = self.func("f32_to_bf16_bulk");
21694 let n2 = nk as i64;
21695 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
21696 let __s_b = self.gpu.stream();
21697 let mut b = __s_b.launch_builder(&f);
21698 b.arg(q).arg(&mut qb).arg(&n2);
21699 unsafe {
21700 b.launch(cfg2)?;
21701 }
21702 qb16 = Some(qb);
21703 } else if let Some(qb) = qb16_pre {
21704 assert!(qb.len() >= nk * 2, "qb16_pre too small");
21705 }
21706 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
21707 }
21708 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
21709 let k2w = if gdn_wgmma_pre {
21710 Some((
21711 *qb16_ref0.as_ref().unwrap(),
21712 *kb16_ref0.as_ref().unwrap(),
21713 pb16.as_mut().unwrap(),
21714 ))
21715 } else {
21716 None
21717 };
21718 let (gcum, p, u, w) =
21719 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
21720 let _ = &w;
21721 let mut y = self.uninit(nc * h * c * D)?;
21722 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
21736 && c == 32
21737 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
21738 Ok("1") => true,
21739 Ok("0") => false,
21740 _ => gdn_mma_default_on(),
21741 };
21742 if gdn_mma {
21743 let wb16 = wb16_pre
21744 .take()
21745 .expect("mma path pre-allocates wb16 (K3 store fold)");
21746 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
21747 if gdn_wgmma_pre {
21759 let qb16 = qb16_ref0.unwrap();
21761 let pb16 = pb16.as_ref().unwrap();
21762 {
21763 let f = self.func("gdn_k45_wgmma");
21764 let cfg = LaunchConfig {
21765 grid_dim: (h as u32, 4, 1),
21766 block_dim: (256, 1, 1),
21767 shared_mem_bytes: 0,
21768 };
21769 let hki = hk as i32;
21770 let __s_b = self.gpu.stream();
21771 let mut b = __s_b.launch_builder(&f);
21772 b.arg(kb16_ref)
21773 .arg(&gcum)
21774 .arg(beta)
21775 .arg(&u)
21776 .arg(&wb16)
21777 .arg(qb16)
21778 .arg(pb16)
21779 .arg(o)
21780 .arg(&scale)
21781 .arg(state_in)
21782 .arg(&mut *state_out)
21783 .arg(&hi)
21784 .arg(&ti)
21785 .arg(&ci)
21786 .arg(&hki);
21787 unsafe {
21788 b.launch(cfg)?;
21789 }
21790 }
21791 return Ok(());
21792 }
21793 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
21797 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
21798 {
21799 let f = self.func("gdn_chunk_state_mma");
21800 let cfg = LaunchConfig {
21801 grid_dim: (h as u32, NSPLIT, 1),
21802 block_dim: (256, 1, 1),
21803 shared_mem_bytes: 0,
21804 };
21805 let hki = hk as i32;
21806 let __s_b = self.gpu.stream();
21807 let mut b = __s_b.launch_builder(&f);
21808 b.arg(kb16_ref)
21809 .arg(&gcum)
21810 .arg(beta)
21811 .arg(&u)
21812 .arg(&wb16)
21813 .arg(&mut y16)
21814 .arg(&mut ssnap16)
21815 .arg(state_in)
21816 .arg(&mut *state_out)
21817 .arg(&hi)
21818 .arg(&ti)
21819 .arg(&ci)
21820 .arg(&hki);
21821 unsafe {
21822 b.launch(cfg)?;
21823 }
21824 }
21825 {
21826 let f = self.func("gdn_chunk_output_mma");
21828 let jt = ((c + 31) / 32) as u32;
21829 let cfg = LaunchConfig {
21830 grid_dim: (nc as u32, h as u32, jt),
21831 block_dim: (256, 1, 1),
21832 shared_mem_bytes: 0,
21833 };
21834 let hki = hk as i32;
21835 let __s_b = self.gpu.stream();
21836 let mut b = __s_b.launch_builder(&f);
21837 b.arg(q)
21838 .arg(&gcum)
21839 .arg(&p)
21840 .arg(&y16)
21841 .arg(&ssnap16)
21842 .arg(o)
21843 .arg(&hi)
21844 .arg(&ti)
21845 .arg(&ci)
21846 .arg(&scale)
21847 .arg(&hki);
21848 unsafe {
21849 b.launch(cfg)?;
21850 }
21851 }
21852 return Ok(());
21853 }
21854 {
21855 let f = self.func("gdn_chunk_state_f32");
21857 let cfg = LaunchConfig {
21858 grid_dim: (h as u32, NSPLIT, 1),
21859 block_dim: (256, 1, 1),
21860 shared_mem_bytes: 0,
21861 };
21862 let __s_b = self.gpu.stream();
21863 let mut b = __s_b.launch_builder(&f);
21864 b.arg(k)
21865 .arg(&gcum)
21866 .arg(beta)
21867 .arg(&u)
21868 .arg(&w)
21869 .arg(&mut y)
21870 .arg(&mut ssnap)
21871 .arg(state_in)
21872 .arg(&mut *state_out)
21873 .arg(&hi)
21874 .arg(&ti)
21875 .arg(&ci);
21876 unsafe {
21877 b.launch(cfg)?;
21878 }
21879 }
21880 {
21881 let f = self.func("gdn_chunk_output_f32");
21883 let jt = ((c + 31) / 32) as u32;
21884 let cfg = LaunchConfig {
21885 grid_dim: (nc as u32, h as u32, jt),
21886 block_dim: (256, 1, 1),
21887 shared_mem_bytes: 0,
21888 };
21889 let __s_b = self.gpu.stream();
21890 let mut b = __s_b.launch_builder(&f);
21891 b.arg(q)
21892 .arg(&gcum)
21893 .arg(&p)
21894 .arg(&y)
21895 .arg(&ssnap)
21896 .arg(o)
21897 .arg(&hi)
21898 .arg(&ti)
21899 .arg(&ci)
21900 .arg(&scale);
21901 unsafe {
21902 b.launch(cfg)?;
21903 }
21904 }
21905 Ok(())
21906 }
21907
21908 #[allow(clippy::too_many_arguments)]
21917 #[allow(clippy::too_many_arguments)]
21918 pub fn gdn_scan_prefill(
21919 &self,
21920 q: &CudaSlice<f32>,
21921 k: &CudaSlice<f32>,
21922 v: &CudaSlice<f32>,
21923 g: &CudaSlice<f32>,
21924 beta: &CudaSlice<f32>,
21925 kb16_pre: Option<&CudaSlice<u8>>,
21926 qb16_pre: Option<&CudaSlice<u8>>,
21927 state_in: &CudaSlice<f32>,
21928 state_out: &mut CudaSlice<f32>,
21929 o: &mut CudaSlice<f32>,
21930 n_head: usize,
21931 t: usize,
21932 scale: f32,
21933 hk: usize,
21934 ) -> Result<(), Box<dyn std::error::Error>> {
21935 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
21936 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
21937 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
21938 }
21939 if Self::gdn_chunked_enabled() && t >= 16 {
21940 self.gdn_scan_chunked(
21941 q,
21942 k,
21943 v,
21944 g,
21945 beta,
21946 kb16_pre,
21947 qb16_pre,
21948 state_in,
21949 state_out,
21950 o,
21951 n_head,
21952 t,
21953 scale,
21954 Self::gdn_chunk_size(),
21955 hk,
21956 )
21957 } else {
21958 assert!(
21959 hk == n_head,
21960 "s128 scan is broadcast-only (prep guarantees by predicate)"
21961 );
21962 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
21963 }
21964 }
21965
21966 #[allow(clippy::too_many_arguments)]
21968 fn gdn_scan_diff(
21969 &self,
21970 q: &CudaSlice<f32>,
21971 k: &CudaSlice<f32>,
21972 v: &CudaSlice<f32>,
21973 g: &CudaSlice<f32>,
21974 beta: &CudaSlice<f32>,
21975 state_in: &CudaSlice<f32>,
21976 state_out: &mut CudaSlice<f32>,
21977 o: &mut CudaSlice<f32>,
21978 n_head: usize,
21979 t: usize,
21980 scale: f32,
21981 ) -> Result<(), Box<dyn std::error::Error>> {
21982 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
21983 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
21984 let mut o_c = self.uninit(o.len())?;
21985 let mut st_c = self.uninit(state_out.len())?;
21986 self.gdn_scan_chunked(
21987 q,
21988 k,
21989 v,
21990 g,
21991 beta,
21992 None,
21993 None,
21994 state_in,
21995 &mut st_c,
21996 &mut o_c,
21997 n_head,
21998 t,
21999 scale,
22000 Self::gdn_chunk_size(),
22001 n_head,
22002 )?;
22003 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
22004 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
22005 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
22006 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
22007 let mut max_abs = 0f32;
22008 let mut max_rel = 0f32;
22009 let mut sum_rel = 0f64;
22010 for (x, y) in a.iter().zip(b) {
22011 let ad = (x - y).abs();
22012 let rel = ad / x.abs().max(y.abs()).max(1e-3);
22013 if ad > max_abs {
22014 max_abs = ad;
22015 }
22016 if rel > max_rel {
22017 max_rel = rel;
22018 }
22019 sum_rel += rel as f64;
22020 }
22021 (max_abs, max_rel, sum_rel / a.len() as f64)
22022 };
22023 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
22024 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
22025 println!(
22026 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
22027 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
22028 Self::gdn_chunk_size()
22029 );
22030 Ok(())
22031 }
22032
22033 pub fn gdn_glog(
22035 &self,
22036 alpha: &CudaSlice<f32>,
22037 dt_bias: &CudaSlice<f32>,
22038 a: &CudaSlice<f32>,
22039 g_log: &mut CudaSlice<f32>,
22040 n_head: usize,
22041 t: usize,
22042 ) -> Result<(), Box<dyn std::error::Error>> {
22043 let f = self.func("gdn_glog_f32");
22044 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22045 let (h, ti) = (n_head as i32, t as i32);
22046 let __s_b = self.gpu.stream();
22047 let mut b = __s_b.launch_builder(&f);
22048 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22049 unsafe {
22050 b.launch(cfg)?;
22051 }
22052 Ok(())
22053 }
22054
22055 pub fn sigmoid_v(
22058 &self,
22059 x: &cudarc::driver::CudaView<f32>,
22060 y: &mut CudaSlice<f32>,
22061 n: usize,
22062 ) -> Result<(), Box<dyn std::error::Error>> {
22063 let f = self.func("sigmoid_f32");
22064 let cfg = LaunchConfig::for_num_elems(n as u32);
22065 let ni = n as i32;
22066 let __s_b = self.gpu.stream();
22067 let mut b = __s_b.launch_builder(&f);
22068 b.arg(x).arg(y).arg(&ni);
22069 unsafe {
22070 b.launch(cfg)?;
22071 }
22072 Ok(())
22073 }
22074
22075 pub fn gdn_glog_v(
22076 &self,
22077 alpha: &cudarc::driver::CudaView<f32>,
22078 dt_bias: &CudaSlice<f32>,
22079 a: &CudaSlice<f32>,
22080 g_log: &mut CudaSlice<f32>,
22081 n_head: usize,
22082 t: usize,
22083 ) -> Result<(), Box<dyn std::error::Error>> {
22084 let f = self.func("gdn_glog_f32");
22085 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
22086 let (h, ti) = (n_head as i32, t as i32);
22087 let __s_b = self.gpu.stream();
22088 let mut b = __s_b.launch_builder(&f);
22089 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
22090 unsafe {
22091 b.launch(cfg)?;
22092 }
22093 Ok(())
22094 }
22095
22096 pub fn sigmoid(
22097 &self,
22098 x: &CudaSlice<f32>,
22099 y: &mut CudaSlice<f32>,
22100 n: usize,
22101 ) -> Result<(), Box<dyn std::error::Error>> {
22102 let f = self.func("sigmoid_f32");
22103 let cfg = LaunchConfig::for_num_elems(n as u32);
22104 let ni = n as i32;
22105 let __s_b = self.gpu.stream();
22106 let mut b = __s_b.launch_builder(&f);
22107 b.arg(x).arg(y).arg(&ni);
22108 unsafe {
22109 b.launch(cfg)?;
22110 }
22111 Ok(())
22112 }
22113
22114 pub fn sig_mul_f16out(
22117 &self,
22118 a: &CudaSlice<f32>,
22119 g: &CudaSlice<f32>,
22120 dst: &mut CudaSlice<f32>,
22121 dst16: &mut CudaSlice<u8>,
22122 n: usize,
22123 ) -> Result<(), Box<dyn std::error::Error>> {
22124 let f = self.func("sig_mul_f16out_f32");
22125 let cfg = LaunchConfig::for_num_elems(n as u32);
22126 let ni = n as i32;
22127 let __s_b = self.gpu.stream();
22128 let mut b = __s_b.launch_builder(&f);
22129 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
22130 unsafe {
22131 b.launch(cfg)?;
22132 }
22133 Ok(())
22134 }
22135
22136 #[allow(clippy::too_many_arguments)]
22145 pub fn attn_head_gate(
22146 &self,
22147 a: &CudaSlice<f32>,
22148 g: &CudaSlice<f32>,
22149 dst: &mut CudaSlice<f32>,
22150 dst16: Option<&mut CudaSlice<u8>>,
22151 head_dim: usize,
22152 n_head: usize,
22153 t: usize,
22154 ) -> Result<(), Box<dyn std::error::Error>> {
22155 let f = self.func("attn_head_gate_f32");
22156 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22157 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22158 let d16: u64 = match dst16 {
22160 Some(d) => self.addr_u8(d),
22161 None => 0,
22162 };
22163 let __s_b = self.gpu.stream();
22164 let mut b = __s_b.launch_builder(&f);
22165 b.arg(a)
22166 .arg(g)
22167 .arg(dst)
22168 .arg(&d16)
22169 .arg(&hd)
22170 .arg(&nh)
22171 .arg(&ti);
22172 unsafe {
22173 b.launch(cfg)?;
22174 }
22175 Ok(())
22176 }
22177
22178 #[allow(clippy::too_many_arguments)]
22187 pub fn swiglu_clamped_mul_scaled(
22188 &self,
22189 gate: &CudaSlice<f32>,
22190 up: &CudaSlice<f32>,
22191 gs: f32,
22192 us: f32,
22193 limit: f32,
22194 dst: &mut CudaSlice<f32>,
22195 n: usize,
22196 ) -> Result<(), Box<dyn std::error::Error>> {
22197 debug_assert!(
22198 limit > 1e-6,
22199 "swiglu_clamped needs a live limit; use silu_mul_scaled"
22200 );
22201 let f = self.func("swiglu_clamped_mul_scaled_f32");
22202 let cfg = LaunchConfig::for_num_elems(n as u32);
22203 let ni = n as i32;
22204 let __s_b = self.gpu.stream();
22205 let mut b = __s_b.launch_builder(&f);
22206 b.arg(gate)
22207 .arg(up)
22208 .arg(&gs)
22209 .arg(&us)
22210 .arg(&limit)
22211 .arg(dst)
22212 .arg(&ni);
22213 unsafe {
22214 b.launch(cfg)?;
22215 }
22216 Ok(())
22217 }
22218
22219 pub fn gated_rmsnorm(
22221 &self,
22222 o: &CudaSlice<f32>,
22223 w: &CudaSlice<f32>,
22224 z: &CudaSlice<f32>,
22225 dst: &mut CudaSlice<f32>,
22226 ncols: usize,
22227 nrows: usize,
22228 eps: f32,
22229 ) -> Result<(), Box<dyn std::error::Error>> {
22230 let f = self.func("gated_rmsnorm_f32");
22231 let cfg = LaunchConfig {
22232 grid_dim: (nrows as u32, 1, 1),
22233 block_dim: (128, 1, 1),
22234 shared_mem_bytes: 0,
22235 };
22236 let (nc, e) = (ncols as i32, eps);
22237 let __s_b = self.gpu.stream();
22238 let mut b = __s_b.launch_builder(&f);
22239 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22240 unsafe {
22241 b.launch(cfg)?;
22242 }
22243 Ok(())
22244 }
22245
22246 pub fn gated_rmsnorm_f16out(
22249 &self,
22250 o: &CudaSlice<f32>,
22251 w: &CudaSlice<f32>,
22252 z: &CudaSlice<f32>,
22253 dst: &mut CudaSlice<f32>,
22254 dst16: &mut CudaSlice<u8>,
22255 ncols: usize,
22256 nrows: usize,
22257 eps: f32,
22258 ) -> Result<(), Box<dyn std::error::Error>> {
22259 let f = self.func("gated_rmsnorm_f16out_f32");
22260 let cfg = LaunchConfig {
22262 grid_dim: (nrows as u32, 1, 1),
22263 block_dim: (128, 1, 1),
22264 shared_mem_bytes: 0,
22265 };
22266 let (nc, e) = (ncols as i32, eps);
22267 let __s_b = self.gpu.stream();
22268 let mut b = __s_b.launch_builder(&f);
22269 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22270 unsafe {
22271 b.launch(cfg)?;
22272 }
22273 Ok(())
22274 }
22275
22276 #[allow(clippy::too_many_arguments)]
22280 pub fn add_rms_norm_zq8(
22281 &self,
22282 a: &CudaSlice<f32>,
22283 b_in: &CudaSlice<f32>,
22284 w: &CudaSlice<f32>,
22285 res: &mut CudaSlice<f32>,
22286 z: &mut CudaSlice<f32>,
22287 ncols: usize,
22288 nrows: usize,
22289 eps: f32,
22290 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22291 assert!(ncols % 32 == 0);
22292 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
22293 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22294 let f = self.func("add_rms_norm_zq8");
22295 let cfg = LaunchConfig {
22296 grid_dim: (nrows as u32, 1, 1),
22297 block_dim: (1024, 1, 1),
22298 shared_mem_bytes: 0,
22299 };
22300 let (nc, ep) = (ncols as i32, eps);
22301 let __s_b = self.gpu.stream();
22302 let mut b = __s_b.launch_builder(&f);
22303 b.arg(a)
22304 .arg(b_in)
22305 .arg(w)
22306 .arg(res)
22307 .arg(z)
22308 .arg(&mut q)
22309 .arg(&mut d)
22310 .arg(&nc)
22311 .arg(&ep);
22312 unsafe {
22313 b.launch(cfg)?;
22314 }
22315 Ok((q, d))
22316 }
22317
22318 pub fn gated_rmsnorm_zv(
22323 &self,
22324 o: &CudaSlice<f32>,
22325 w: &CudaSlice<f32>,
22326 z: &cudarc::driver::CudaView<f32>,
22327 dst: &mut CudaSlice<f32>,
22328 ncols: usize,
22329 nrows: usize,
22330 eps: f32,
22331 ) -> Result<(), Box<dyn std::error::Error>> {
22332 let f = self.func("gated_rmsnorm_f32");
22333 let cfg = LaunchConfig {
22334 grid_dim: (nrows as u32, 1, 1),
22335 block_dim: (128, 1, 1),
22336 shared_mem_bytes: 0,
22337 };
22338 let (nc, e) = (ncols as i32, eps);
22339 let __s_b = self.gpu.stream();
22340 let mut b = __s_b.launch_builder(&f);
22341 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
22342 unsafe {
22343 b.launch(cfg)?;
22344 }
22345 Ok(())
22346 }
22347
22348 pub fn gated_rmsnorm_f16out_zv(
22349 &self,
22350 o: &CudaSlice<f32>,
22351 w: &CudaSlice<f32>,
22352 z: &cudarc::driver::CudaView<f32>,
22353 dst: &mut CudaSlice<f32>,
22354 dst16: &mut CudaSlice<u8>,
22355 ncols: usize,
22356 nrows: usize,
22357 eps: f32,
22358 ) -> Result<(), Box<dyn std::error::Error>> {
22359 let f = self.func("gated_rmsnorm_f16out_f32");
22360 let cfg = LaunchConfig {
22362 grid_dim: (nrows as u32, 1, 1),
22363 block_dim: (128, 1, 1),
22364 shared_mem_bytes: 0,
22365 };
22366 let (nc, e) = (ncols as i32, eps);
22367 let __s_b = self.gpu.stream();
22368 let mut b = __s_b.launch_builder(&f);
22369 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
22370 unsafe {
22371 b.launch(cfg)?;
22372 }
22373 Ok(())
22374 }
22375
22376 pub fn gated_rmsnorm_q8_1(
22377 &self,
22378 o: &CudaSlice<f32>,
22379 w: &CudaSlice<f32>,
22380 z: &CudaSlice<f32>,
22381 ncols: usize,
22382 nrows: usize,
22383 eps: f32,
22384 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22385 assert!(ncols % 32 == 0);
22386 let f = self.func("gated_rmsnorm_q8_1");
22387 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
22388 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
22389 let cfg = LaunchConfig {
22390 grid_dim: (nrows as u32, 1, 1),
22391 block_dim: (128, 1, 1),
22392 shared_mem_bytes: 0,
22393 };
22394 let (nc, ep) = (ncols as i32, eps);
22395 let __s_b = self.gpu.stream();
22396 let mut b = __s_b.launch_builder(&f);
22397 b.arg(o)
22398 .arg(w)
22399 .arg(z)
22400 .arg(&mut out_q)
22401 .arg(&mut out_d)
22402 .arg(&nc)
22403 .arg(&ep);
22404 unsafe {
22405 b.launch(cfg)?;
22406 }
22407 Ok((out_q, out_d))
22408 }
22409
22410 pub fn transpose(
22412 &self,
22413 inp: &CudaSlice<f32>,
22414 rows: usize,
22415 cols: usize,
22416 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22417 let f = self.func("transpose_f32");
22418 let mut out = self.zeros(rows * cols)?;
22419 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
22420 let (r, c) = (rows as i32, cols as i32);
22421 let __s_b = self.gpu.stream();
22422 let mut b = __s_b.launch_builder(&f);
22423 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
22424 unsafe {
22425 b.launch(cfg)?;
22426 }
22427 Ok(out)
22428 }
22429
22430 pub fn repeat_heads(
22432 &self,
22433 inp: &CudaSlice<f32>,
22434 out: &mut CudaSlice<f32>,
22435 head_dim: usize,
22436 n_in: usize,
22437 n_out: usize,
22438 t: usize,
22439 ) -> Result<(), Box<dyn std::error::Error>> {
22440 let f = self.func("repeat_heads_f32");
22441 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
22442 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
22443 let __s_b = self.gpu.stream();
22444 let mut b = __s_b.launch_builder(&f);
22445 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
22446 unsafe {
22447 b.launch(cfg)?;
22448 }
22449 Ok(())
22450 }
22451
22452 pub fn q_gate_split(
22459 &self,
22460 qf: &CudaSlice<f32>,
22461 q_out: &mut CudaSlice<f32>,
22462 gate_out: &mut CudaSlice<f32>,
22463 head_dim: usize,
22464 n_head: usize,
22465 t: usize,
22466 ) -> Result<(), Box<dyn std::error::Error>> {
22467 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
22468 let out_need = head_dim * n_head * t;
22469 if q_out.len() < out_need || gate_out.len() < out_need {
22470 return Err(format!(
22471 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
22472 q_out.len(),
22473 gate_out.len()
22474 )
22475 .into());
22476 }
22477 let f = self.func("q_gate_split_f32");
22478 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
22479 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
22480 let __s_b = self.gpu.stream();
22481 let mut b = __s_b.launch_builder(&f);
22482 b.arg(qf)
22483 .arg(q_out)
22484 .arg(gate_out)
22485 .arg(&hd)
22486 .arg(&nh)
22487 .arg(&ti);
22488 unsafe {
22489 b.launch(cfg)?;
22490 }
22491 Ok(())
22492 }
22493
22494 pub fn qkv_to_gdn_repack(
22498 &self,
22499 conv_out: &CudaSlice<f32>,
22500 q_g: &mut CudaSlice<f32>,
22501 k_g: &mut CudaSlice<f32>,
22502 v_g: &mut CudaSlice<f32>,
22503 d_state: usize,
22504 num_v: usize,
22505 num_k: usize,
22506 key_dim: usize,
22507 t: usize,
22508 ) -> Result<(), Box<dyn std::error::Error>> {
22509 let f = self.func("qkv_to_gdn_repack_f32");
22510 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
22511 let (ds, nv, nk, kd, ti) = (
22512 d_state as i32,
22513 num_v as i32,
22514 num_k as i32,
22515 key_dim as i32,
22516 t as i32,
22517 );
22518 let __s_b = self.gpu.stream();
22519 let mut b = __s_b.launch_builder(&f);
22520 b.arg(conv_out)
22521 .arg(q_g)
22522 .arg(k_g)
22523 .arg(v_g)
22524 .arg(&ds)
22525 .arg(&nv)
22526 .arg(&nk)
22527 .arg(&kd)
22528 .arg(&ti);
22529 unsafe {
22530 b.launch(cfg)?;
22531 }
22532 Ok(())
22533 }
22534
22535 pub fn conv_left_pad(
22538 &self,
22539 src: &CudaSlice<f32>,
22540 dst: &mut CudaSlice<f32>,
22541 conv_dim: usize,
22542 t: usize,
22543 pad: usize,
22544 ) -> Result<(), Box<dyn std::error::Error>> {
22545 let f = self.func("conv_left_pad_f32");
22546 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
22547 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
22548 let __s_b = self.gpu.stream();
22549 let mut b = __s_b.launch_builder(&f);
22550 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
22551 unsafe {
22552 b.launch(cfg)?;
22553 }
22554 Ok(())
22555 }
22556
22557 pub fn conv_assemble_and_roll(
22561 &self,
22562 qkv_col: &CudaSlice<f32>,
22563 conv_state: &mut CudaSlice<f32>,
22564 conv_in: &mut CudaSlice<f32>,
22565 conv_dim: usize,
22566 pad: usize,
22567 ) -> Result<(), Box<dyn std::error::Error>> {
22568 let f = self.func("conv_assemble_and_roll_f32");
22569 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22570 let (cd, p) = (conv_dim as i32, pad as i32);
22571 let __s_b = self.gpu.stream();
22572 let mut b = __s_b.launch_builder(&f);
22573 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
22574 unsafe {
22575 b.launch(cfg)?;
22576 }
22577 Ok(())
22578 }
22579
22580 pub fn ssm_conv1d_fused_decode(
22586 &self,
22587 qkv_col: &CudaSlice<f32>,
22588 conv_state: &mut CudaSlice<f32>,
22589 w: &CudaSlice<f32>,
22590 conv_out: &mut CudaSlice<f32>,
22591 conv_dim: usize,
22592 d_conv: usize,
22593 ) -> Result<(), Box<dyn std::error::Error>> {
22594 let f = self.func("ssm_conv1d_fused_decode_f32");
22595 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
22596 let (cd, dc) = (conv_dim as i32, d_conv as i32);
22597 let __s_b = self.gpu.stream();
22598 let mut b = __s_b.launch_builder(&f);
22599 b.arg(qkv_col)
22600 .arg(conv_state)
22601 .arg(w)
22602 .arg(conv_out)
22603 .arg(&cd)
22604 .arg(&dc);
22605 unsafe {
22606 b.launch(cfg)?;
22607 }
22608 Ok(())
22609 }
22610
22611 pub fn slice_range(
22614 &self,
22615 src: &CudaSlice<f32>,
22616 start: usize,
22617 len: usize,
22618 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22619 let host = self.gpu.stream().clone_dtoh(src)?;
22620 self.gpu.stream().synchronize()?;
22621 Ok(self.htod(&host[start..start + len])?)
22622 }
22623}
22624
22625#[cfg(test)]
22626mod target_dispatch_tests {
22627 use super::legacy_quant_gemm_allowed;
22628
22629 #[test]
22630 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
22631 assert!(legacy_quant_gemm_allowed(false, false, false));
22633 assert!(!legacy_quant_gemm_allowed(false, false, true));
22634 assert!(!legacy_quant_gemm_allowed(true, false, false));
22636 assert!(!legacy_quant_gemm_allowed(true, false, true));
22637 assert!(legacy_quant_gemm_allowed(true, true, false));
22639 assert!(!legacy_quant_gemm_allowed(true, true, true));
22640 }
22641
22642 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
22643 #[test]
22644 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
22645 assert!(!legacy_quant_gemm_allowed(
22646 cfg!(memra_portable_cuda),
22647 cfg!(memra_hopper_mma),
22648 false
22649 ));
22650 }
22651
22652 #[cfg(memra_hopper_mma)]
22653 #[test]
22654 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
22655 assert!(legacy_quant_gemm_allowed(
22656 cfg!(memra_portable_cuda),
22657 cfg!(memra_hopper_mma),
22658 false
22659 ));
22660 assert!(super::portable_mma_gated() == false);
22661 }
22662}
22663
22664impl memra_kv::KvDev for Engine {
22667 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22668 Engine::zeros(self, n)
22669 }
22670 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22671 Engine::uninit(self, n)
22672 }
22673 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
22674 Engine::alloc_u8(self, n)
22675 }
22676 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
22677 Engine::htod_i32(self, v)
22678 }
22679 fn clone_dtod(
22680 &self,
22681 src: &CudaSlice<f32>,
22682 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22683 Engine::clone_dtod(self, src)
22684 }
22685 fn copy_into(
22686 &self,
22687 dst: &mut CudaSlice<f32>,
22688 off: usize,
22689 src: &CudaSlice<f32>,
22690 len: usize,
22691 ) -> Result<(), Box<dyn std::error::Error>> {
22692 Engine::copy_into(self, dst, off, src, len)
22693 }
22694 fn set_i32_one(
22695 &self,
22696 d: &mut CudaSlice<i32>,
22697 v: i32,
22698 ) -> Result<(), Box<dyn std::error::Error>> {
22699 Engine::set_i32_one(self, d, v)
22700 }
22701}
22702
22703#[cfg(test)]
22704mod fused_gate_bounds_tests {
22705 use super::*;
22706
22707 #[test]
22720 #[ignore = "requires a CUDA GPU"]
22721 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
22722 let e = Engine::new(0).unwrap();
22723 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
22724 let fused = 2 * head_dim * n_head * t;
22725 let out_n = head_dim * n_head * t;
22726
22727 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
22729 let mut q = e.uninit(out_n).unwrap();
22730 let mut gate = e.uninit(out_n).unwrap();
22731 let err = e
22732 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
22733 .expect_err("half-width wq must be refused, not read past")
22734 .to_string();
22735 assert!(err.contains("NO fused gate"), "{err}");
22736 assert!(err.contains(&format!("{fused}")), "{err}");
22737
22738 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
22741 let wide = e.htod(&host).unwrap();
22742 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
22743 .expect("full-width wq splits");
22744 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
22745 for tok in 0..t {
22746 for hh in 0..n_head {
22747 for d in 0..head_dim {
22748 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
22749 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
22750 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
22751 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
22752 }
22753 }
22754 }
22755
22756 let mut small = e.uninit(out_n - 1).unwrap();
22758 assert!(
22759 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
22760 .is_err()
22761 );
22762 }
22763}
22764
22765#[cfg(test)]
22769mod fused_rope_width_tests {
22770 use super::Engine;
22771
22772 #[test]
22775 fn full_width_is_accepted() {
22776 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
22777 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
22778 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
22779 }
22780
22781 #[test]
22794 fn gemma4_official_artifact_widths_pass() {
22795 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
22796 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
22797 }
22798
22799 #[test]
22802 fn partial_rotary_is_refused_with_the_geometry_named() {
22803 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
22805 .expect_err("partial rotary must refuse");
22806 let msg = err.to_string();
22807 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
22808 assert!(msg.contains("n_rot 64"), "{msg}");
22809 assert!(msg.contains("head_dim 256"), "{msg}");
22810 assert!(
22811 msg.contains("64..256"),
22812 "names the band it would corrupt: {msg}"
22813 );
22814 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
22816 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
22818 }
22819}