1use cudarc::driver::sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES;
4use cudarc::driver::{
5 CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DeviceSlice, LaunchConfig,
6 PushKernelArg,
7};
8use cudarc::nvrtc::Ptx;
9use std::sync::{Arc, Mutex};
10
11const GDN_K2_DYNAMIC_SHARED_BYTES: u32 = 67_072;
12
13const SDPA_NAIVE_SMEM_MAX: usize = 48 * 1024;
18
19const SDPA_NAIVE_GMEM_WS_MAX: usize = 1 << 30;
23
24#[cfg(debug_assertions)]
25pub(crate) fn debug_assert_tensor_stream_device<T>(
26 tensor: &CudaSlice<T>,
27 stream: &CudaStream,
28 site: &str,
29) {
30 let tensor_dev = tensor.ordinal();
31 let stream_dev = stream.context().ordinal();
32 assert_eq!(
33 tensor_dev, stream_dev,
34 "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
35 );
36}
37
38fn ensure_tensor_stream_device<T>(
39 tensor: &impl DeviceSlice<T>,
40 stream: &CudaStream,
41 site: &str,
42) -> Result<(), Box<dyn std::error::Error>> {
43 let tensor_dev = tensor.stream().context().ordinal();
44 let stream_dev = stream.context().ordinal();
45 if tensor_dev != stream_dev {
46 return Err(format!(
47 "PP cross-device tensor access at {site}: tensor on dev{tensor_dev}, \
48 stream on dev{stream_dev}"
49 )
50 .into());
51 }
52 Ok(())
53}
54
55pub use memra_gguf;
56pub use memra_runtime;
57
58pub mod forward;
59pub mod hybrid;
60pub mod hybrid_forward;
61pub mod model;
62pub mod sigrouter_contract;
63pub mod vision;
64pub mod vision_gemma;
65pub mod vision_pre;
66pub mod vision_step;
67pub mod cache {
70 pub use memra_kv::*;
71}
72pub mod decode;
73pub mod decode_batch;
74pub mod dflash;
75pub mod eagle;
76pub mod gemma_spec;
77pub mod graph_update;
78pub mod mla;
82pub mod moesd;
83pub mod parallel;
84pub mod plan_backend;
85pub mod pp;
86pub mod round_stream;
87pub mod spec;
88pub mod tp;
89pub use memra_sampling as sampler;
90
91pub fn moe_f16g_mode() -> u8 {
135 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
136 *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
137 Ok("0") => 0,
138 Ok("2") => 2,
139 Ok("3") => 3,
140 Ok(_) => 1,
141 Err(_) => 2,
144 })
145}
146pub fn moe_f16g_sk_params() -> (i32, i32) {
160 static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
161 *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
162 Ok("0") => (-1, 0),
163 Ok("32") => (0, i32::MAX),
164 Ok("128") => (0, 1),
165 _ => {
166 let cross = std::env::var("MEMRA_F16G_SK_CROSS")
167 .ok()
168 .and_then(|v| v.parse().ok())
169 .unwrap_or(64);
170 (0, cross)
171 }
172 })
173}
174pub fn moe_f16g_direct_on(qtype: i32) -> bool {
185 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
186 let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
187 Ok("0") => 0,
188 Ok("kq") => 1,
189 _ => 2,
190 });
191 match m {
192 0 => false,
193 1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
194 _ => true,
195 }
196}
197pub fn moe_f16g_tail_on() -> bool {
206 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
207 *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
208}
209
210pub fn moe_f16g_gemma_on() -> bool {
217 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
218 *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
219}
220
221pub fn moe_fuse_actq_on() -> bool {
225 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
226 *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
227}
228
229pub fn router_prefill_exact_on() -> bool {
239 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
240 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
241}
242
243pub fn router_kernel_on() -> bool {
244 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
245 *ON.get_or_init(|| {
246 let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
247 if !on {
248 eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
249 }
250 on
251 })
252}
253
254pub const ROUTER_BATCH_MIN_T: usize = 8;
269pub fn router_batch_on() -> bool {
270 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
271 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
272}
273mod cpu_experts;
274#[cfg(memra_cutlass)]
275pub mod cutlass_ffi;
276pub mod dsv4_ffi;
277pub mod dsv4_gpu;
278pub mod f16_ffi;
279pub mod fp8_ffi;
280pub mod mmq_ffi;
281pub mod moe_cache;
282pub mod prime_graph;
283pub mod spill;
284mod spill_pread;
285
286const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
293const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
294const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
295const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
296const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
297const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
298const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
300
301fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
307 assert!(
308 !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
309 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
310 );
311 match std::env::var("MEMRA_GEMM_FATBIN") {
312 Ok(path) => std::borrow::Cow::Owned(
313 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
314 ),
315 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
316 }
317}
318
319pub(crate) const fn portable_mma_gated() -> bool {
326 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
327}
328
329#[track_caller]
346pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
347 assert!(
348 !portable_mma_gated(),
349 "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
350 portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
351 );
352}
353
354pub(crate) const fn gdn_mma_default_on() -> bool {
363 cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
364}
365
366const fn konst_eq(a: &str, b: &str) -> bool {
368 let (a, b) = (a.as_bytes(), b.as_bytes());
369 if a.len() != b.len() {
370 return false;
371 }
372 let mut i = 0;
373 while i < a.len() {
374 if a[i] != b[i] {
375 return false;
376 }
377 i += 1;
378 }
379 true
380}
381
382const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
387 (!portable_cuda || hopper_mma) && !no_gemm
388}
389
390const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
398const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
399const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
400const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
401const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
402
403pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
406
407fn flash_fatbin_bytes() -> &'static [u8] {
409 match kv_cache_formats() {
410 ("q8_0", "q5_1") => FLASH_FATBIN,
411 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
412 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
413 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
414 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
415 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
416 other => unreachable!("kv_cache_formats returned {other:?}"),
417 }
418}
419
420fn k1_launch_override() -> Option<(u32, u32, u32)> {
427 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
428 *K1.get_or_init(|| {
429 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
430 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
431 match p.as_slice() {
432 [bm, bn, w] => Some((*bm, *bn, *w)),
433 _ => None,
434 }
435 })
436}
437
438pub(crate) fn wgmma_gemm_enabled() -> bool {
445 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
446 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
447}
448
449pub const FA_VEC_MIN_TKV: usize = 96;
464pub fn fa_vec_min_tkv() -> usize {
468 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
469 *V.get_or_init(|| {
470 std::env::var("MEMRA_FA_VEC_MIN")
471 .ok()
472 .and_then(|v| v.parse().ok())
473 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
474 })
475}
476
477pub fn fa_f16pv_on() -> bool {
488 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
489 *ON.get_or_init(|| {
490 std::env::var("MEMRA_FA_F16PV")
491 .map(|v| v != "0")
492 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
493 })
494}
495
496pub fn fa512_hp_on() -> bool {
500 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
501 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
502}
503
504pub fn faw_hp_on() -> bool {
508 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
509 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
510}
511
512pub fn fa512_wide_warps() -> usize {
516 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
517 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
518 Ok("1") => 4,
519 _ => 2,
520 })
521}
522
523pub fn fa512_min_tkv() -> usize {
526 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
527 *FA512_MIN.get_or_init(|| {
528 std::env::var("MEMRA_FA512_MIN")
529 .ok()
530 .and_then(|v| v.parse().ok())
531 .unwrap_or(512)
532 })
533}
534pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
538 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
539pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
543pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
549 std::sync::atomic::AtomicBool::new(false);
550pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
557 std::sync::atomic::AtomicBool::new(true);
558pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
559 std::sync::atomic::AtomicUsize::new(16);
560pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
565pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
567pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
572pub use memra_kv::KV_FP8_FORCE;
575pub(crate) fn mmv_block() -> u32 {
580 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
581 *V.get_or_init(|| {
582 std::env::var("MEMRA_MMV_BLOCK")
583 .ok()
584 .and_then(|v| v.parse().ok())
585 .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
586 .unwrap_or(128)
587 })
588}
589
590static STEP37_SERVING_DEFAULTS: std::sync::atomic::AtomicBool =
616 std::sync::atomic::AtomicBool::new(false);
617
618pub fn arm_step37_serving_defaults() {
619 STEP37_SERVING_DEFAULTS.store(true, std::sync::atomic::Ordering::Relaxed);
620 crate::cache::set_swa_ring_default(true);
621 eprintln!(
622 "[step37-defaults] serving doors armed ON for the SlidingGatedMoe program \
623 (per-flag =0 kills, =1 forces; owner flip 2026-08-27)"
624 );
625}
626
627pub(crate) fn step37_defaults_armed() -> bool {
628 STEP37_SERVING_DEFAULTS.load(std::sync::atomic::Ordering::Relaxed)
629}
630
631pub(crate) fn step37_door(cell: &'static std::sync::OnceLock<Option<bool>>, name: &str) -> bool {
635 match *cell.get_or_init(|| match std::env::var(name).ok().as_deref() {
636 Some("1") => Some(true),
637 Some("0") => Some(false),
638 _ => None,
639 }) {
640 Some(forced) => forced,
641 None => step37_defaults_armed(),
642 }
643}
644
645pub(crate) fn w8_hybrid_on() -> bool {
646 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
647 step37_door(&ENV, "MEMRA_W8_HYBRID")
648}
649
650pub(crate) fn step_tp_w8_on() -> bool {
651 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
652 step37_door(&ENV, "MEMRA_STEP_TP_W8")
653}
654
655pub(crate) fn w8_view_on() -> bool {
660 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
661 *ON.get_or_init(|| std::env::var("MEMRA_W8_VIEW").as_deref() == Ok("1"))
662}
663
664pub(crate) fn step_gemm_prime_on() -> bool {
678 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
679 step37_door(&ENV, "MEMRA_STEP_GEMM_PRIME")
680}
681
682pub(crate) fn q8t_wonce_on() -> bool {
683 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
684 step37_door(&ENV, "MEMRA_Q8T_WONCE")
685}
686
687pub(crate) fn sig_expf_dev_on() -> bool {
692 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
693 *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
694}
695
696pub(crate) fn topk_fast_on() -> bool {
697 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
698 *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
699}
700
701fn sigmoid_topk_kernel(sig_expf: bool, fast: bool, n_used: usize) -> &'static str {
706 match (sig_expf, fast && n_used <= 8) {
707 (true, true) => "moe_router_sigmoid_topk_f32_dexp_fast",
708 (true, false) => "moe_router_sigmoid_topk_f32_dexp",
709 (false, true) => "moe_router_sigmoid_topk_f32_fast",
710 (false, false) => "moe_router_sigmoid_topk_f32",
711 }
712}
713
714#[cfg(test)]
715mod sigmoid_topk_dispatch_tests {
716 #[test]
717 fn fast_kernel_refuses_wide_topk_and_composes_with_dexp() {
718 use super::sigmoid_topk_kernel;
719
720 assert_eq!(
721 sigmoid_topk_kernel(false, true, 8),
722 "moe_router_sigmoid_topk_f32_fast"
723 );
724 assert_eq!(
725 sigmoid_topk_kernel(true, true, 8),
726 "moe_router_sigmoid_topk_f32_dexp_fast"
727 );
728 assert_eq!(
729 sigmoid_topk_kernel(false, true, 9),
730 "moe_router_sigmoid_topk_f32"
731 );
732 assert_eq!(
733 sigmoid_topk_kernel(true, true, 9),
734 "moe_router_sigmoid_topk_f32_dexp"
735 );
736 }
737}
738
739pub(crate) fn rms_block() -> u32 {
740 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
741 *V.get_or_init(|| {
742 std::env::var("MEMRA_RMS_BLOCK")
743 .ok()
744 .and_then(|v| v.parse().ok())
745 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
746 })
747}
748
749pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
750 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
751 if let Some(forced) = *S.get_or_init(|| {
752 std::env::var("MEMRA_FA_SPLIT")
753 .ok()
754 .and_then(|v| v.parse().ok())
755 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
756 }) {
757 return forced;
758 }
759 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
777 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
778 {
779 return if t_kv <= 8192 {
780 16
781 } else if t_kv <= 16384 {
782 64
783 } else {
784 128
785 };
786 }
787 let big_rig = fa_sm_count() >= 128;
788 if big_rig {
789 let _ = n_head_kv;
790 if t_kv <= 2048 {
791 static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
800 if let Some(sp) = *SHORT.get_or_init(|| {
801 std::env::var("MEMRA_FA_SP_SHORT")
802 .ok()
803 .and_then(|v| v.parse().ok())
804 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
805 }) {
806 return sp;
807 }
808 16
809 } else if t_kv <= 16384 {
810 64
811 } else {
812 128
813 }
814 } else if n_head_kv <= 4 {
815 if t_kv <= 512 {
836 8
837 } else if t_kv <= 16384 {
838 64
839 } else {
840 128
841 }
842 } else {
843 if t_kv <= 8192 {
844 32
845 } else if t_kv <= 16384 {
846 64
847 } else {
848 128
849 }
850 }
851}
852
853pub(crate) fn fa_sm_count() -> i32 {
856 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
857 *N.get_or_init(|| {
858 cudarc::driver::result::init().ok();
859 cudarc::driver::result::device::get(0)
860 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
861 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
862 .unwrap_or(82)
863 })
864}
865
866fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
870 match head_dim {
871 256 => Ok(""),
872 128 => Ok("_hd128"),
873 d => Err(format!(
874 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
875 callers must gate to sdpa_naive"
876 )
877 .into()),
878 }
879}
880
881pub const QT_Q8_0: i32 = 0;
883pub const QT_Q4_K: i32 = 1;
884pub const QT_Q6_K: i32 = 2;
885pub const QT_Q5_K: i32 = 3;
886pub const QT_Q3_K: i32 = 4;
887pub const QT_IQ4_XS: i32 = 5;
888pub const QT_IQ3_S: i32 = 6;
889pub const QT_NVFP4: i32 = 7;
890pub const QT_NVFP4_V2: i32 = 107;
893pub const QT_F8_E4M3: i32 = 10;
899pub const QT_NVFP4_RP: i32 = 9;
902pub const QT_F32: i32 = 8;
904pub const QT_BF16: i32 = 11;
905pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
910pub const QT_F8_E4M3_BLK: i32 = 14;
926
927pub struct Engine {
929 pub gpu: memra_runtime::Gpu,
930 module: Arc<CudaModule>,
931 hybrid: Arc<CudaModule>,
932 qmatvec: Arc<CudaModule>,
933 flash: Arc<CudaModule>,
934 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
938 gemm: Arc<CudaModule>,
939 router: Arc<CudaModule>,
940 sample: Arc<CudaModule>,
942 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
946 w8_mirrors: Mutex<std::collections::HashMap<(u64, u32, u32), CudaSlice<u8>>>,
955 w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
958 moe_cache_layout: Mutex<Option<Vec<usize>>>,
962 capture_keep_on: std::sync::atomic::AtomicBool,
968 verify_exact: std::sync::atomic::AtomicBool,
973 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
974 pub copy_stream: Arc<CudaStream>,
976 #[cfg(memra_cutlass)]
983 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
984 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
988 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
991 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
995 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
999 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
1001 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
1002 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
1007 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
1012 router_stage: Mutex<Option<PinnedStage>>,
1016}
1017
1018fn fa_v2_on() -> bool {
1028 std::env::var("MEMRA_FA_V2")
1034 .map(|v| v != "0")
1035 .unwrap_or(true)
1036}
1037
1038pub(crate) fn fa_part_zero_on() -> bool {
1048 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1049 *ON.get_or_init(|| std::env::var("MEMRA_FA_PART_ZERO").as_deref() == Ok("1"))
1050}
1051
1052pub(crate) fn fa_v3_on() -> bool {
1053 std::env::var("MEMRA_FA_V3")
1057 .map(|v| v != "0")
1058 .unwrap_or(true)
1059}
1060
1061fn fa_v4_mode() -> &'static str {
1066 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
1067 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
1068}
1069fn fa_v4_on() -> bool {
1070 fa_v4_mode() != "0"
1071} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
1080 std::sync::atomic::AtomicUsize::new(1024);
1081pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
1082 std::sync::atomic::AtomicUsize::new(usize::MAX);
1083pub fn fa_v4_at_pub(t_kv: usize) -> bool {
1084 fa_v4_at(t_kv)
1085}
1086fn fa_v4_at(t_kv: usize) -> bool {
1087 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
1088 let mx = *M.get_or_init(|| {
1089 std::env::var("MEMRA_FA_V4_MAX")
1090 .ok()
1091 .and_then(|v| v.parse().ok())
1092 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1093 });
1094 fa_v4_on() && t_kv < mx
1095}
1096pub const FA_DEEP_MIN_DEFAULT: usize = 0;
1110fn fa_deep_at(t_kv: usize) -> bool {
1111 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
1112 return false;
1113 }
1114 let min = std::env::var("MEMRA_FA_DEEP_MIN")
1115 .ok()
1116 .and_then(|v| v.parse().ok())
1117 .unwrap_or(FA_DEEP_MIN_DEFAULT);
1118 t_kv >= min
1119}
1120pub fn fa_deep_at_pub(t_kv: usize) -> bool {
1122 fa_deep_at(t_kv)
1123}
1124
1125fn fa_v3_active(head_dim: usize) -> bool {
1126 fa_v3_on()
1129 && head_dim % 128 == 0
1130 && kv_cache_formats() == ("q8_0", "q5_1")
1131 && !Engine::kv_fp8_on()
1132}
1133
1134pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1142 std::env::var("MEMRA_NO_FA_VEC").is_err()
1143 && t_kv >= fa_vec_min_tkv()
1144 && head_dim == 256
1145 && fa_v4_at(t_kv)
1146 && !matches!(fa_v4_mode(), "noB3" | "stage")
1147 && !Engine::kv_fp8_on()
1148}
1149pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1151 fa_split_keys(t_kv, n_head_kv)
1152}
1153
1154struct PinnedStage {
1159 ptr: *mut u8,
1160 cap: usize,
1161}
1162unsafe impl Send for PinnedStage {}
1163impl PinnedStage {
1164 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1165 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1166 Ok(PinnedStage { ptr, cap })
1167 }
1168}
1169impl Drop for PinnedStage {
1170 fn drop(&mut self) {
1171 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1172 }
1173}
1174
1175pub struct PinnedHostBuf {
1182 ptr: *mut u8,
1183 len: usize,
1184}
1185unsafe impl Send for PinnedHostBuf {}
1188impl PinnedHostBuf {
1189 pub fn new(len: usize) -> Result<Self, Box<dyn std::error::Error>> {
1192 let ptr = unsafe { cudarc::driver::result::malloc_host(len.max(1), 0)? } as *mut u8;
1193 Ok(PinnedHostBuf { ptr, len })
1194 }
1195 pub fn len(&self) -> usize {
1196 self.len
1197 }
1198 pub fn is_empty(&self) -> bool {
1199 self.len == 0
1200 }
1201 pub fn as_slice(&self) -> &[u8] {
1202 unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
1203 }
1204 pub fn as_mut_slice(&mut self) -> &mut [u8] {
1205 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
1206 }
1207}
1208impl Drop for PinnedHostBuf {
1209 fn drop(&mut self) {
1210 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1211 }
1212}
1213
1214pub const ARGMAX_NB: usize = 256;
1217
1218pub(crate) use memra_fa3_vl as fa3_vl_raw;
1220
1221unsafe extern "C" {
1222 fn memra_fa3_prefill(
1224 q16: *const core::ffi::c_void,
1225 k16: *const core::ffi::c_void,
1226 v16: *const core::ffi::c_void,
1227 o: *mut f32,
1228 t: i32,
1229 h: i32,
1230 hkv: i32,
1231 d: i32,
1232 scale: f32,
1233 stream: *mut core::ffi::c_void,
1234 ) -> i32;
1235 pub(crate) fn memra_fa3_vl(
1237 q16s: *const *const core::ffi::c_void,
1238 k16s: *const *const core::ffi::c_void,
1239 v16s: *const *const core::ffi::c_void,
1240 os: *const *mut f32,
1241 ts: *const i32,
1242 b: i32,
1243 h: i32,
1244 hkv: i32,
1245 d: i32,
1246 scale: f32,
1247 stream: *mut core::ffi::c_void,
1248 ) -> i32;
1249}
1250
1251#[repr(C)]
1256#[derive(Clone, Copy)]
1257pub struct WPtr8(pub [u64; 8]);
1258unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1259
1260#[repr(C)]
1265#[derive(Clone, Copy, Default)]
1266pub struct GdnSeqVl {
1267 pub kb16: u64,
1268 pub gcum: u64,
1269 pub beta: u64,
1270 pub u: u64,
1271 pub wb16: u64,
1272 pub y: u64,
1273 pub ssnap: u64,
1274 pub state_in: u64,
1275 pub state_out: u64,
1276 pub q: u64,
1277 pub p: u64,
1278 pub o: u64,
1279 pub k: u64,
1280 pub v: u64,
1281 pub g: u64,
1282 pub a: u64,
1283 pub w: u64,
1284 pub t: i32,
1285 pub nc: i32,
1286}
1287unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1288#[repr(C)]
1289#[derive(Clone, Copy)]
1290pub struct GdnVl8(pub [GdnSeqVl; 8]);
1291unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1292
1293#[repr(C)]
1296#[derive(Clone, Copy, Default)]
1297pub struct GdnWVl {
1298 pub qb16: u64,
1299 pub pb16: u64,
1300}
1301unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1302#[repr(C)]
1303#[derive(Clone, Copy)]
1304pub struct GdnWVl8(pub [GdnWVl; 8]);
1305unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1306
1307#[repr(C)]
1309#[derive(Clone, Copy, Default)]
1310pub struct GdnPrepVl {
1311 pub qkv: u64,
1312 pub conv_state: u64,
1313 pub conv_out: u64,
1314 pub q_g: u64,
1315 pub k_g: u64,
1316 pub v_g: u64,
1317 pub q_l2: u64,
1318 pub k_l2: u64,
1319 pub beta_raw: u64,
1320 pub alpha: u64,
1321 pub beta: u64,
1322 pub g_log: u64,
1323 pub o: u64,
1324 pub z: u64,
1325 pub gn: u64,
1326 pub gn16: u64,
1327 pub kb16: u64,
1328 pub qb16: u64,
1329 pub t: i32,
1330 pub pad: i32,
1331}
1332unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1333#[repr(C)]
1334#[derive(Clone, Copy)]
1335pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1336unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1337
1338#[repr(C)]
1340#[derive(Clone, Copy, Default)]
1341pub struct FaSeqVl {
1342 pub q: u64,
1343 pub k16: u64,
1344 pub v16: u64,
1345 pub o: u64,
1346 pub kf: u64,
1347 pub vf: u64,
1348 pub t: i32,
1349 pub pad: i32,
1350}
1351unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1352#[repr(C)]
1353#[derive(Clone, Copy)]
1354pub struct FaVl8(pub [FaSeqVl; 8]);
1355unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1356
1357#[repr(C)]
1359#[derive(Clone, Copy, Default)]
1360pub struct AttnPreVl {
1361 pub qf: u64,
1362 pub kf: u64,
1363 pub vf: u64,
1364 pub q: u64,
1365 pub gate: u64,
1366 pub qn: u64,
1367 pub kn: u64,
1368 pub kc: u64,
1369 pub vc: u64,
1370 pub t: i32,
1371 pub pad: i32,
1372}
1373unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1374#[repr(C)]
1375#[derive(Clone, Copy)]
1376pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1377unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1378
1379pub struct GdnChunkBufs {
1382 pub gcum: CudaSlice<f32>,
1383 pub a: CudaSlice<f32>,
1384 pub p: CudaSlice<f32>,
1385 pub u: CudaSlice<f32>,
1386 pub w: CudaSlice<f32>,
1387 pub kb16: CudaSlice<u8>,
1388 pub wb16: CudaSlice<u8>,
1389 pub y16: CudaSlice<u8>,
1390 pub ssnap16: CudaSlice<u8>,
1391 pub qb16: CudaSlice<u8>,
1392 pub pb16: CudaSlice<u8>,
1393 pub o: CudaSlice<f32>,
1394 pub t: usize,
1395 pub nc: usize,
1396}
1397
1398#[repr(C)]
1400#[derive(Clone, Copy)]
1401pub struct F32x8(pub [f32; 8]);
1402unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1403
1404pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1408
1409#[must_use = "dropping immediately ends the exact scope"]
1414pub struct ExactScope<'a> {
1415 flag: &'a std::sync::atomic::AtomicBool,
1416 prev: bool,
1417}
1418
1419impl<'a> ExactScope<'a> {
1420 pub(crate) fn set(flag: &'a std::sync::atomic::AtomicBool, on: bool) -> Self {
1421 let prev = flag.load(std::sync::atomic::Ordering::Relaxed);
1422 flag.store(on, std::sync::atomic::Ordering::Relaxed);
1423 ExactScope { flag, prev }
1424 }
1425}
1426
1427impl Drop for ExactScope<'_> {
1428 fn drop(&mut self) {
1429 self.flag
1430 .store(self.prev, std::sync::atomic::Ordering::Relaxed);
1431 }
1432}
1433
1434#[cfg(test)]
1435mod exact_scope_tests {
1436 use std::sync::atomic::{AtomicBool, Ordering};
1437
1438 #[test]
1439 fn error_path_restores_verify_exact() {
1440 let flag = AtomicBool::new(false);
1445 let failing = |flag: &AtomicBool| -> Result<(), &'static str> {
1446 let _scope = super::ExactScope::set(flag, true);
1447 assert!(flag.load(Ordering::Relaxed), "scope arms the flag");
1448 Err("draft forward failed")? };
1450 assert!(failing(&flag).is_err());
1451 assert!(
1452 !flag.load(Ordering::Relaxed),
1453 "error propagation must restore the pre-scope value"
1454 );
1455 let flag = AtomicBool::new(true);
1457 {
1458 let _scope = super::ExactScope::set(&flag, true);
1459 }
1460 assert!(flag.load(Ordering::Relaxed));
1461 let flag = AtomicBool::new(false);
1463 let scope = super::ExactScope::set(&flag, true);
1464 drop(scope);
1465 assert!(!flag.load(Ordering::Relaxed));
1466 }
1467}
1468
1469impl Engine {
1470 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1471 let gpu = memra_runtime::Gpu::new(ordinal)?;
1472 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1476 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1477 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1478 .and_then(|d| unsafe {
1479 Ok((
1480 cudarc::driver::result::device::get_attribute(
1481 d,
1482 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1483 )?,
1484 cudarc::driver::result::device::get_attribute(
1485 d,
1486 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1487 )?,
1488 ))
1489 })
1490 .unwrap_or((0, 0));
1491 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1492 let ok = matches!(
1493 (built, maj, min),
1494 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1495 );
1496 if !ok {
1497 return Err(format!(
1498 "memra was built for sm_{built} but device {ordinal} reports compute \
1499 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1500 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1501 )
1502 .into());
1503 }
1504 }
1505 unsafe {
1510 use cudarc::driver::sys;
1511 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1512 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1513 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1514 let mut thresh: u64 = u64::MAX;
1515 let _ = sys::cuMemPoolSetAttribute(
1516 pool,
1517 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1518 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1519 );
1520 }
1521 }
1522 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1523 let hybrid = gpu
1524 .ctx
1525 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1526 let qmatvec = gpu
1527 .ctx
1528 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1529 let flash = gpu
1530 .ctx
1531 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1532 let gemm = gpu
1533 .ctx
1534 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1535 let router = gpu
1536 .ctx
1537 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1538 let sample = gpu
1539 .ctx
1540 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1541 let copy_stream = gpu.ctx.new_stream()?;
1542 if std::env::var("MEMRA_EVT")
1558 .map(|v| v == "1")
1559 .unwrap_or(false)
1560 {
1561 } else {
1563 unsafe {
1564 gpu.ctx.disable_event_tracking();
1565 }
1566 }
1567 Ok(Self {
1568 gpu,
1569 module,
1570 hybrid,
1571 qmatvec,
1572 flash,
1573 flash_g: std::sync::OnceLock::new(),
1574 gemm,
1575 router,
1576 sample,
1577 moe_cache: Mutex::new(None),
1578 w8_mirrors: Mutex::new(std::collections::HashMap::new()),
1579 w8_act: Mutex::new(std::collections::HashMap::new()),
1580 moe_cache_layout: Mutex::new(None),
1581 copy_stream,
1582 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1583 verify_exact: std::sync::atomic::AtomicBool::new(false),
1584 capture_keep: Mutex::new(Vec::new()),
1585 argmax_partials: Mutex::new(None),
1586 prime_deqw_ws: Mutex::new(None),
1587 router_stage: Mutex::new(None),
1588 fp8_scratch: Mutex::new(None),
1589 fa_vf16_scratch: Mutex::new(None),
1590 fa_part_pool: Mutex::new(None),
1591 fa_part_retired: Mutex::new(Vec::new()),
1592 fn_cache: Mutex::new(Default::default()),
1593 f16_scratch: Mutex::new(None),
1594 #[cfg(memra_cutlass)]
1595 cutlass_scratch: Mutex::new(None),
1596 })
1597 }
1598
1599 pub fn ctx(&self) -> &Arc<CudaContext> {
1600 &self.gpu.ctx
1601 }
1602
1603 pub fn pool_cached_bytes(&self) -> usize {
1621 let (reserved, used) = self.pool_reserved_used();
1622 reserved.saturating_sub(used)
1623 }
1624
1625 pub fn device_graph_mem_reserved(&self) -> usize {
1635 use cudarc::driver::sys as cus;
1636 let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
1637 return 0;
1638 };
1639 let mut bytes: u64 = 0;
1640 let rc = unsafe {
1641 cus::cuDeviceGetGraphMemAttribute(
1642 dev,
1643 cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
1644 &mut bytes as *mut u64 as *mut std::ffi::c_void,
1645 )
1646 };
1647 if rc == cus::cudaError_enum::CUDA_SUCCESS {
1648 bytes as usize
1649 } else {
1650 0
1651 }
1652 }
1653
1654 pub fn pool_trim_to_zero(&self) -> usize {
1668 use cudarc::driver::sys;
1669 let (before, _) = self.pool_reserved_used();
1670 unsafe {
1671 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1672 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1673 != sys::CUresult::CUDA_SUCCESS
1674 {
1675 return 0;
1676 }
1677 let _ = sys::cuMemPoolTrimTo(pool, 0);
1678 }
1679 let (after, _) = self.pool_reserved_used();
1680 before.saturating_sub(after)
1681 }
1682
1683 pub fn pool_reserved_used(&self) -> (usize, usize) {
1684 use cudarc::driver::sys;
1685 unsafe {
1686 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1687 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1688 != sys::CUresult::CUDA_SUCCESS
1689 {
1690 return (0, 0);
1691 }
1692 let (mut reserved, mut used) = (0u64, 0u64);
1693 if sys::cuMemPoolGetAttribute(
1694 pool,
1695 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1696 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1697 ) != sys::CUresult::CUDA_SUCCESS
1698 {
1699 return (0, 0);
1700 }
1701 if sys::cuMemPoolGetAttribute(
1702 pool,
1703 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1704 &mut used as *mut u64 as *mut core::ffi::c_void,
1705 ) != sys::CUresult::CUDA_SUCCESS
1706 {
1707 return (0, 0);
1708 }
1709 (reserved as usize, used as usize)
1710 }
1711 }
1712
1713 pub fn pool_high_water_reset(&self) -> (usize, usize) {
1722 use cudarc::driver::sys;
1723 unsafe {
1724 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1725 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1726 != sys::CUresult::CUDA_SUCCESS
1727 {
1728 return (0, 0);
1729 }
1730 let (mut reserved, mut used) = (0u64, 0u64);
1731 if sys::cuMemPoolGetAttribute(
1732 pool,
1733 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
1734 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1735 ) != sys::CUresult::CUDA_SUCCESS
1736 {
1737 return (0, 0);
1738 }
1739 if sys::cuMemPoolGetAttribute(
1740 pool,
1741 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
1742 &mut used as *mut u64 as *mut core::ffi::c_void,
1743 ) != sys::CUresult::CUDA_SUCCESS
1744 {
1745 return (0, 0);
1746 }
1747 let mut zero: u64 = 0;
1750 let _ = sys::cuMemPoolSetAttribute(
1751 pool,
1752 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
1753 &mut zero as *mut u64 as *mut core::ffi::c_void,
1754 );
1755 let mut zero2: u64 = 0;
1756 let _ = sys::cuMemPoolSetAttribute(
1757 pool,
1758 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
1759 &mut zero2 as *mut u64 as *mut core::ffi::c_void,
1760 );
1761 (reserved as usize, used as usize)
1762 }
1763 }
1764
1765 pub fn stream(&self) -> Arc<CudaStream> {
1768 self.gpu.stream()
1769 }
1770 pub fn gkv_on() -> bool {
1773 memra_kv::gkv_on()
1774 }
1775
1776 pub fn wkv_on() -> bool {
1788 memra_kv::wkv_on()
1789 }
1790
1791 pub fn kv_fp8_on() -> bool {
1797 memra_kv::kv_fp8_on()
1798 }
1799
1800 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1803 if head_dim == 512 && Self::gkv_on() {
1804 self.func_g(name)
1805 } else {
1806 self.func(name)
1807 }
1808 }
1809
1810 fn func_g(&self, name: &str) -> CudaFunction {
1814 let m = self.flash_g.get_or_init(|| {
1815 self.gpu
1816 .ctx
1817 .load_module(cudarc::nvrtc::Ptx::from_binary(
1818 FLASH_FATBIN_KF8VF8.to_vec(),
1819 ))
1820 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1821 });
1822 let key = format!("g:{name}");
1823 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1824 return f.clone();
1825 }
1826 let f = match m.load_function(name) {
1827 Ok(f) => f,
1828 Err(_) => self.func(name),
1829 };
1830 self.fn_cache.lock().unwrap().insert(key, f.clone());
1831 f
1832 }
1833
1834 fn func(&self, name: &str) -> CudaFunction {
1835 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1838 return f.clone();
1839 }
1840 let f = self
1841 .module
1842 .load_function(name)
1843 .or_else(|_| self.hybrid.load_function(name))
1844 .or_else(|_| self.qmatvec.load_function(name))
1845 .or_else(|_| self.flash.load_function(name))
1846 .or_else(|_| self.gemm.load_function(name))
1847 .or_else(|_| self.router.load_function(name))
1848 .or_else(|_| self.sample.load_function(name))
1849 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1850 self.fn_cache
1851 .lock()
1852 .unwrap()
1853 .insert(name.to_string(), f.clone());
1854 f
1855 }
1856
1857 pub fn scatter_trim_logits(
1860 &self,
1861 src: &CudaSlice<f32>,
1862 d2t: &CudaSlice<u32>,
1863 dst: &mut CudaSlice<f32>,
1864 d_vocab: usize,
1865 n_vocab: usize,
1866 ) -> Result<(), Box<dyn std::error::Error>> {
1867 let f1 = self.func("scatter_trim_logits_f32");
1868 let f2 = self.func("scatter_trim_logits_pass2_f32");
1869 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1870 let cfg1 = LaunchConfig {
1871 grid_dim: (256, 1, 1),
1872 block_dim: (256, 1, 1),
1873 shared_mem_bytes: 0,
1874 };
1875 let __s_b1 = self.gpu.stream();
1876 let mut b1 = __s_b1.launch_builder(&f1);
1877 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1878 unsafe {
1879 b1.launch(cfg1)?;
1880 }
1881 let cfg2 = LaunchConfig {
1882 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1883 block_dim: (256, 1, 1),
1884 shared_mem_bytes: 0,
1885 };
1886 let __s_b2 = self.gpu.stream();
1887 let mut b2 = __s_b2.launch_builder(&f2);
1888 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1889 unsafe {
1890 b2.launch(cfg2)?;
1891 }
1892 Ok(())
1893 }
1894
1895 #[allow(clippy::too_many_arguments)]
1901 pub fn filter_stats(
1902 &self,
1903 x: &CudaSlice<f32>,
1904 row_stride: usize,
1905 rows: &CudaSlice<i32>,
1906 out_th: &mut CudaSlice<f32>,
1907 out_z: &mut CudaSlice<f32>,
1908 out_max: &mut CudaSlice<f32>,
1909 n: usize,
1910 nrow: usize,
1911 temp: f32,
1912 top_k: i32,
1913 top_p: f32,
1914 min_p: f32,
1915 ) -> Result<(), Box<dyn std::error::Error>> {
1916 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1943 let coop_on =
1944 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1945 if coop_on && self.sm_count() >= 16 {
1946 let cap = self.sm_count() as usize / 16;
1947 let mut done = 0usize;
1948 while done < nrow {
1949 let chunk = cap.min(nrow - done);
1950 self.filter_stats_coop_chunk(
1951 x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
1952 top_p, min_p,
1953 )?;
1954 done += chunk;
1955 }
1956 return Ok(());
1957 }
1958 self.filter_stats_plain_program(
1959 x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
1960 )
1961 }
1962
1963 #[allow(clippy::too_many_arguments)]
1968 pub fn filter_stats_coop_chunk(
1969 &self,
1970 x: &CudaSlice<f32>,
1971 row_stride: usize,
1972 rows: &CudaSlice<i32>,
1973 row0: usize,
1974 out_th: &mut CudaSlice<f32>,
1975 out_z: &mut CudaSlice<f32>,
1976 out_max: &mut CudaSlice<f32>,
1977 n: usize,
1978 chunk: usize,
1979 temp: f32,
1980 top_k: i32,
1981 top_p: f32,
1982 min_p: f32,
1983 ) -> Result<(), Box<dyn std::error::Error>> {
1984 let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
1985 let f = self.func("filter_stats_coop_f32");
1986 let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
1987 let cfg = LaunchConfig {
1988 grid_dim: (16, chunk as u32, 1),
1989 block_dim: (512, 1, 1),
1990 shared_mem_bytes: 0,
1991 };
1992 let rows_v = rows.slice(row0..row0 + chunk);
1993 let mut th_v = out_th.slice_mut(row0..row0 + chunk);
1994 let mut z_v = out_z.slice_mut(row0..row0 + chunk);
1995 let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
1996 let __s_b = self.gpu.stream();
1997 let mut b = __s_b.launch_builder(&f);
1998 b.arg(x)
1999 .arg(&rs)
2000 .arg(&rows_v)
2001 .arg(&mut th_v)
2002 .arg(&mut z_v)
2003 .arg(&mut mx_v)
2004 .arg(&mut ws)
2005 .arg(&ni)
2006 .arg(&nr)
2007 .arg(&temp)
2008 .arg(&top_k)
2009 .arg(&top_p)
2010 .arg(&min_p);
2011 unsafe {
2012 b.launch_cooperative(cfg)?;
2013 }
2014 Ok(())
2015 }
2016
2017 #[allow(clippy::too_many_arguments)]
2021 pub fn filter_stats_plain_program(
2022 &self,
2023 x: &CudaSlice<f32>,
2024 row_stride: usize,
2025 rows: &CudaSlice<i32>,
2026 out_th: &mut CudaSlice<f32>,
2027 out_z: &mut CudaSlice<f32>,
2028 out_max: &mut CudaSlice<f32>,
2029 n: usize,
2030 nrow: usize,
2031 temp: f32,
2032 top_k: i32,
2033 top_p: f32,
2034 min_p: f32,
2035 ) -> Result<(), Box<dyn std::error::Error>> {
2036 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
2037 let f = self.func("filter_stats_f32");
2038 let cfg = LaunchConfig {
2039 grid_dim: (nrow as u32, 1, 1),
2040 block_dim: (1024, 1, 1),
2041 shared_mem_bytes: 0,
2042 };
2043 let __s_b = self.gpu.stream();
2044 let mut b = __s_b.launch_builder(&f);
2045 b.arg(x)
2046 .arg(&rs)
2047 .arg(rows)
2048 .arg(&mut *out_th)
2049 .arg(&mut *out_z)
2050 .arg(&mut *out_max)
2051 .arg(&ni)
2052 .arg(&nr)
2053 .arg(&temp)
2054 .arg(&top_k)
2055 .arg(&top_p)
2056 .arg(&min_p);
2057 unsafe {
2058 b.launch(cfg)?;
2059 }
2060 Ok(())
2061 }
2062
2063 #[allow(clippy::too_many_arguments)]
2065 pub fn softmax_gather_filtered(
2066 &self,
2067 x: &CudaSlice<f32>,
2068 row_stride: usize,
2069 ids: &CudaSlice<u32>,
2070 rows: &CudaSlice<i32>,
2071 th: &CudaSlice<f32>,
2072 z: &CudaSlice<f32>,
2073 out: &mut CudaSlice<f32>,
2074 n: usize,
2075 npair: usize,
2076 temp: f32,
2077 ) -> Result<(), Box<dyn std::error::Error>> {
2078 let f = self.func("softmax_gather_filtered_f32");
2079 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
2080 let cfg = LaunchConfig {
2081 grid_dim: (npair as u32, 1, 1),
2082 block_dim: (256, 1, 1),
2083 shared_mem_bytes: 0,
2084 };
2085 let __s_b = self.gpu.stream();
2086 let mut b = __s_b.launch_builder(&f);
2087 b.arg(x)
2088 .arg(&rs)
2089 .arg(ids)
2090 .arg(rows)
2091 .arg(th)
2092 .arg(z)
2093 .arg(&mut *out)
2094 .arg(&ni)
2095 .arg(&np)
2096 .arg(&temp);
2097 unsafe {
2098 b.launch(cfg)?;
2099 }
2100 Ok(())
2101 }
2102
2103 #[allow(clippy::too_many_arguments)]
2105 pub fn residual_sample_filtered(
2106 &self,
2107 p: &CudaSlice<f32>,
2108 q: Option<&CudaSlice<f32>>,
2109 n: usize,
2110 temp: f32,
2111 seed: u64,
2112 stream_pos: u32,
2113 p_stats: (f32, f32, f32),
2114 q_stats: (f32, f32, f32),
2115 out_tok: &mut CudaSlice<u32>,
2116 ) -> Result<(), Box<dyn std::error::Error>> {
2117 let f = self.func("residual_sample_filtered_f32");
2118 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2119 let has_q: i32 = q.is_some() as i32;
2120 let qbuf = q.unwrap_or(p);
2121 let (pm, pth, pz) = p_stats;
2122 let (qm, qth, qz) = q_stats;
2123 let cfg = LaunchConfig {
2124 grid_dim: (1, 1, 1),
2125 block_dim: (1024, 1, 1),
2126 shared_mem_bytes: 0,
2127 };
2128 let __s_b = self.gpu.stream();
2129 let mut b = __s_b.launch_builder(&f);
2130 b.arg(p)
2131 .arg(qbuf)
2132 .arg(&has_q)
2133 .arg(&ni)
2134 .arg(&temp)
2135 .arg(&slo)
2136 .arg(&shi)
2137 .arg(&stream_pos)
2138 .arg(&pm)
2139 .arg(&pth)
2140 .arg(&pz)
2141 .arg(&qm)
2142 .arg(&qth)
2143 .arg(&qz)
2144 .arg(&mut *out_tok);
2145 unsafe {
2146 b.launch(cfg)?;
2147 }
2148 Ok(())
2149 }
2150
2151 #[allow(clippy::too_many_arguments)]
2157 pub fn residual_sample_sparse_q(
2158 &self,
2159 p: &CudaSlice<f32>,
2160 cand_ids: &CudaSlice<u32>,
2161 q_probs: &CudaSlice<f32>,
2162 n_cand: usize,
2163 n: usize,
2164 temp: f32,
2165 seed: u64,
2166 stream_pos: u32,
2167 p_stats: (f32, f32, f32),
2168 out_tok: &mut CudaSlice<u32>,
2169 ) -> Result<(), Box<dyn std::error::Error>> {
2170 assert!(
2171 n_cand >= 1 && n_cand <= 32,
2172 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
2173 );
2174 let f = self.func("residual_sample_sparse_q_f32");
2175 let (ni, nc) = (n as i32, n_cand as i32);
2176 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2177 let (pm, pth, pz) = p_stats;
2178 let cfg = LaunchConfig {
2179 grid_dim: (1, 1, 1),
2180 block_dim: (1024, 1, 1),
2181 shared_mem_bytes: 0,
2182 };
2183 let __s_b = self.gpu.stream();
2184 let mut b = __s_b.launch_builder(&f);
2185 b.arg(p)
2186 .arg(cand_ids)
2187 .arg(q_probs)
2188 .arg(&nc)
2189 .arg(&ni)
2190 .arg(&temp)
2191 .arg(&slo)
2192 .arg(&shi)
2193 .arg(&stream_pos)
2194 .arg(&pm)
2195 .arg(&pth)
2196 .arg(&pz)
2197 .arg(&mut *out_tok);
2198 unsafe {
2199 b.launch(cfg)?;
2200 }
2201 Ok(())
2202 }
2203
2204 #[allow(clippy::too_many_arguments)]
2206 pub fn gumbel_perturb_filtered(
2207 &self,
2208 x: &CudaSlice<f32>,
2209 y: &mut CudaSlice<f32>,
2210 n: usize,
2211 seed: u64,
2212 stream_pos: u32,
2213 temp: f32,
2214 row_max: f32,
2215 th: f32,
2216 ) -> Result<(), Box<dyn std::error::Error>> {
2217 let f = self.func("gumbel_perturb_filtered_f32");
2218 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2219 let cfg = LaunchConfig {
2220 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2221 block_dim: (256, 1, 1),
2222 shared_mem_bytes: 0,
2223 };
2224 let __s_b = self.gpu.stream();
2225 let mut b = __s_b.launch_builder(&f);
2226 b.arg(x)
2227 .arg(&mut *y)
2228 .arg(&ni)
2229 .arg(&slo)
2230 .arg(&shi)
2231 .arg(&stream_pos)
2232 .arg(&temp)
2233 .arg(&row_max)
2234 .arg(&th);
2235 unsafe {
2236 b.launch(cfg)?;
2237 }
2238 Ok(())
2239 }
2240
2241 #[allow(clippy::too_many_arguments)]
2245 pub fn penalize_logits(
2246 &self,
2247 x: &mut CudaSlice<f32>,
2248 hist: &CudaSlice<u32>,
2249 n_hist: usize,
2250 rep: f32,
2251 freq: f32,
2252 present: f32,
2253 n: usize,
2254 ) -> Result<(), Box<dyn std::error::Error>> {
2255 if n_hist == 0 {
2256 return Ok(());
2257 }
2258 let f = self.func("penalize_logits_f32");
2259 let (nh, ni) = (n_hist as i32, n as i32);
2260 let cfg = LaunchConfig {
2261 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
2262 block_dim: (128, 1, 1),
2263 shared_mem_bytes: 0,
2264 };
2265 let __s_b = self.gpu.stream();
2266 let mut b = __s_b.launch_builder(&f);
2267 b.arg(&mut *x)
2268 .arg(hist)
2269 .arg(&nh)
2270 .arg(&rep)
2271 .arg(&freq)
2272 .arg(&present)
2273 .arg(&ni);
2274 unsafe {
2275 b.launch(cfg)?;
2276 }
2277 Ok(())
2278 }
2279
2280 #[allow(clippy::too_many_arguments)]
2282 pub fn penalize_logits_rows(
2283 &self,
2284 x: &mut CudaSlice<f32>,
2285 hist: &CudaSlice<u32>,
2286 n_hist: usize,
2287 rep: f32,
2288 freq: f32,
2289 present: f32,
2290 n: usize,
2291 nrow: usize,
2292 ) -> Result<(), Box<dyn std::error::Error>> {
2293 if n_hist == 0 || nrow == 0 {
2294 return Ok(());
2295 }
2296 let f = self.func("penalize_logits_rows_f32");
2297 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
2298 let cfg = LaunchConfig {
2299 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
2300 block_dim: (128, 1, 1),
2301 shared_mem_bytes: 0,
2302 };
2303 let __s_b = self.gpu.stream();
2304 let mut b = __s_b.launch_builder(&f);
2305 b.arg(&mut *x)
2306 .arg(hist)
2307 .arg(&nh)
2308 .arg(&rep)
2309 .arg(&freq)
2310 .arg(&present)
2311 .arg(&ni)
2312 .arg(&nr);
2313 unsafe {
2314 b.launch(cfg)?;
2315 }
2316 Ok(())
2317 }
2318
2319 #[allow(clippy::too_many_arguments)]
2325 pub fn penalize_logits_sparse_rows(
2326 &self,
2327 x: &mut CudaSlice<f32>,
2328 ids: &[u32],
2329 counts: &[u32],
2330 offsets: &[i32],
2331 rows: &[i32],
2332 reps: &[f32],
2333 freqs: &[f32],
2334 presents: &[f32],
2335 n: usize,
2336 ) -> Result<(), Box<dyn std::error::Error>> {
2337 let nrow = rows.len();
2338 if nrow == 0 {
2339 return Ok(());
2340 }
2341 let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2342 let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2343 let entry_count =
2344 i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
2345 if ids.len() != counts.len()
2346 || offsets.len() != nrow + 1
2347 || reps.len() != nrow
2348 || freqs.len() != nrow
2349 || presents.len() != nrow
2350 || offsets.first().copied() != Some(0)
2351 || offsets.last().copied() != Some(entry_count)
2352 {
2353 return Err("sparse penalty row metadata shape mismatch".into());
2354 }
2355 if counts.contains(&0) {
2356 return Err("sparse penalty counts must be positive".into());
2357 }
2358 let mut max_len = 0usize;
2359 for pair in offsets.windows(2) {
2360 if pair[0] < 0 || pair[1] < pair[0] {
2361 return Err("sparse penalty offsets must be monotonic".into());
2362 }
2363 max_len = max_len.max((pair[1] - pair[0]) as usize);
2364 }
2365 if max_len == 0 {
2366 return Ok(());
2367 }
2368
2369 let mut seen = std::collections::HashSet::with_capacity(ids.len());
2370 for (r, &row) in rows.iter().enumerate() {
2371 if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
2372 return Err("sparse penalty row index exceeds logits shape".into());
2373 }
2374 let begin = offsets[r] as usize;
2375 let end = offsets[r + 1] as usize;
2376 for &id in &ids[begin..end] {
2377 if id as usize >= n {
2378 return Err("sparse penalty token id exceeds logits row".into());
2379 }
2380 if !seen.insert((row, id)) {
2381 return Err("sparse penalty entries must be unique per logits row".into());
2382 }
2383 }
2384 }
2385
2386 unsafe {
2388 self.penalize_logits_sparse_rows_unchecked(
2389 x, ids, counts, offsets, rows, reps, freqs, presents, n,
2390 )
2391 }
2392 }
2393
2394 #[allow(clippy::too_many_arguments)]
2403 pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
2404 &self,
2405 x: &mut CudaSlice<f32>,
2406 ids: &[u32],
2407 counts: &[u32],
2408 offsets: &[i32],
2409 rows: &[i32],
2410 reps: &[f32],
2411 freqs: &[f32],
2412 presents: &[f32],
2413 n: usize,
2414 ) -> Result<(), Box<dyn std::error::Error>> {
2415 let nrow = rows.len();
2416 if nrow == 0 {
2417 return Ok(());
2418 }
2419 let max_len = offsets
2420 .windows(2)
2421 .map(|pair| (pair[1] - pair[0]) as usize)
2422 .max()
2423 .unwrap_or(0);
2424 if max_len == 0 {
2425 return Ok(());
2426 }
2427 let ids_d = self.htod_u32_v(ids)?;
2428 let counts_d = self.htod_u32_v(counts)?;
2429 let offsets_d = self.htod_i32(offsets)?;
2430 let rows_d = self.htod_i32(rows)?;
2431 let reps_d = self.htod(reps)?;
2432 let freqs_d = self.htod(freqs)?;
2433 let presents_d = self.htod(presents)?;
2434 let f = self.func("penalize_logits_sparse_rows_f32");
2435 let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2436 let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2437 let cfg = LaunchConfig {
2438 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2439 block_dim: (128, 1, 1),
2440 shared_mem_bytes: 0,
2441 };
2442 let __s_b = self.gpu.stream();
2443 let mut b = __s_b.launch_builder(&f);
2444 b.arg(&mut *x)
2445 .arg(&ids_d)
2446 .arg(&counts_d)
2447 .arg(&offsets_d)
2448 .arg(&rows_d)
2449 .arg(&reps_d)
2450 .arg(&freqs_d)
2451 .arg(&presents_d)
2452 .arg(&ni)
2453 .arg(&nr);
2454 unsafe {
2455 b.launch(cfg)?;
2456 }
2457 Ok(())
2458 }
2459
2460 #[allow(clippy::too_many_arguments)]
2468 pub fn penalize_logits_rows_inc(
2469 &self,
2470 x: &mut CudaSlice<f32>,
2471 hist: &CudaSlice<u32>,
2472 n_hist0: usize,
2473 rep: f32,
2474 freq: f32,
2475 present: f32,
2476 n: usize,
2477 nrow: usize,
2478 win: usize,
2479 ) -> Result<(), Box<dyn std::error::Error>> {
2480 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2481 return Ok(());
2482 }
2483 debug_assert!(
2484 hist.len() >= n_hist0 + nrow - 1,
2485 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2486 );
2487 let f = self.func("penalize_logits_rows_inc_f32");
2488 let max_len = win.min(n_hist0 + nrow - 1).max(1);
2489 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2490 let cfg = LaunchConfig {
2491 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2492 block_dim: (128, 1, 1),
2493 shared_mem_bytes: 0,
2494 };
2495 let __s_b = self.gpu.stream();
2496 let mut b = __s_b.launch_builder(&f);
2497 b.arg(&mut *x)
2498 .arg(hist)
2499 .arg(&nh)
2500 .arg(&rep)
2501 .arg(&freq)
2502 .arg(&present)
2503 .arg(&ni)
2504 .arg(&nr)
2505 .arg(&wi);
2506 unsafe {
2507 b.launch(cfg)?;
2508 }
2509 Ok(())
2510 }
2511
2512 pub fn wpf_level() -> u32 {
2520 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2521 *ON.get_or_init(|| {
2522 std::env::var("MEMRA_WPF")
2523 .ok()
2524 .and_then(|v| v.parse().ok())
2525 .unwrap_or(1)
2526 })
2527 }
2528
2529 pub fn set_verify_exact(&self, on: bool) {
2545 self.verify_exact
2546 .store(on, std::sync::atomic::Ordering::Relaxed);
2547 }
2548 pub(crate) fn verify_exact_on(&self) -> bool {
2549 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2550 }
2551
2552 pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
2558 ExactScope::set(&self.verify_exact, on)
2559 }
2560
2561 pub fn qkv_append_on() -> bool {
2564 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2565 *ON.get_or_init(|| {
2566 std::env::var("MEMRA_QKV_APPEND")
2567 .map(|v| v != "0")
2568 .unwrap_or(true)
2569 })
2570 }
2571
2572 pub fn pdl_wb_on() -> bool {
2575 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2576 *ON.get_or_init(|| {
2577 std::env::var("MEMRA_PDL_WB")
2578 .map(|v| v != "0")
2579 .unwrap_or(true)
2580 })
2581 }
2582
2583 pub fn norm_ilp_on() -> bool {
2591 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2592 *ON.get_or_init(|| {
2593 std::env::var("MEMRA_NORM_ILP")
2594 .map(|v| v != "0")
2595 .unwrap_or(true)
2596 })
2597 }
2598
2599 pub fn tk_ffn_dual_on() -> bool {
2607 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2608 *ON.get_or_init(|| {
2609 std::env::var("MEMRA_TK_FFN_DUAL")
2610 .map(|v| v != "0")
2611 .unwrap_or(true)
2612 })
2613 }
2614
2615 pub fn pdl_mmvq_on() -> bool {
2619 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2620 *ON.get_or_init(|| {
2621 std::env::var("MEMRA_PDL_MMVQ")
2622 .map(|v| v != "0")
2623 .unwrap_or(true)
2624 })
2625 }
2626
2627 pub fn pdl_on() -> bool {
2628 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2629 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2630 }
2631
2632 pub fn pdl_nvfp4q8_on() -> bool {
2638 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2639 *ON.get_or_init(|| {
2640 std::env::var("MEMRA_PDL_NVFP4")
2641 .map(|v| v != "0")
2642 .unwrap_or(true)
2643 })
2644 }
2645
2646 fn q40_mr1_on() -> bool {
2652 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2653 match *Q40MR.get_or_init(|| {
2654 std::env::var("MEMRA_Q40_MR")
2655 .ok()
2656 .and_then(|v| v.parse().ok())
2657 }) {
2658 Some(v) => v == 1,
2659 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2660 }
2661 }
2662
2663 fn pdl_func_flash(
2668 &self,
2669 g: bool,
2670 name: &'static str,
2671 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2672 use cudarc::driver::sys as cu;
2673 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2680 std::sync::Mutex::new(None);
2681 static FNS: std::sync::Mutex<
2682 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2683 > = std::sync::Mutex::new(None);
2684 let ctx_key = self.ctx().cu_ctx() as usize;
2685 if let Some(&f) = FNS
2686 .lock()
2687 .unwrap()
2688 .get_or_insert_with(Default::default)
2689 .get(&(ctx_key, g, name))
2690 {
2691 return Ok(f as cu::CUfunction);
2692 }
2693 let module = {
2694 let mut mods = MODS.lock().unwrap();
2695 let map = mods.get_or_insert_with(Default::default);
2696 match map.get(&(ctx_key, g)) {
2697 Some(&m) => m,
2698 None => {
2699 let m = self.pdl_load_module_in_ctx(if g {
2700 FLASH_FATBIN_KF8VF8
2701 } else {
2702 FLASH_FATBIN
2703 })?;
2704 map.insert((ctx_key, g), m);
2705 m
2706 }
2707 }
2708 };
2709 let cname = std::ffi::CString::new(name)?;
2710 let mut f: cu::CUfunction = std::ptr::null_mut();
2711 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2712 if r != cu::CUresult::CUDA_SUCCESS {
2713 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2714 }
2715 FNS.lock()
2716 .unwrap()
2717 .get_or_insert_with(Default::default)
2718 .insert((ctx_key, g, name), f as usize);
2719 Ok(f)
2720 }
2721
2722 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2727 use cudarc::driver::sys as cu;
2728 let mut prev: cu::CUcontext = std::ptr::null_mut();
2729 unsafe {
2730 cu::cuCtxGetCurrent(&mut prev).result()?;
2731 }
2732 self.ctx().bind_to_thread()?;
2733 let mut m: cu::CUmodule = std::ptr::null_mut();
2734 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2735 let restore = if prev.is_null() {
2736 cu::CUresult::CUDA_SUCCESS
2737 } else {
2738 unsafe { cu::cuCtxSetCurrent(prev) }
2739 };
2740 if r != cu::CUresult::CUDA_SUCCESS {
2741 return Err(format!("pdl module load: {r:?}").into());
2742 }
2743 if restore != cu::CUresult::CUDA_SUCCESS {
2744 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2745 }
2746 Ok(m as usize)
2747 }
2748
2749 pub fn raw_kernel_function(
2752 &self,
2753 name: &'static str,
2754 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2755 self.pdl_func(name)
2756 }
2757
2758 fn pdl_func(
2759 &self,
2760 name: &'static str,
2761 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2762 use cudarc::driver::sys as cu;
2763 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2766 std::sync::Mutex::new(None);
2767 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2770 std::sync::Mutex::new(None);
2771 static FNS: std::sync::Mutex<
2772 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2773 > = std::sync::Mutex::new(None);
2774 let ctx_key = self.ctx().cu_ctx() as usize;
2775 if let Some(&f) = FNS
2776 .lock()
2777 .unwrap()
2778 .get_or_insert_with(Default::default)
2779 .get(&(ctx_key, name))
2780 {
2781 return Ok(f as cu::CUfunction);
2782 }
2783 let module = {
2784 let mut mods = MODULES.lock().unwrap();
2785 let map = mods.get_or_insert_with(Default::default);
2786 match map.get(&ctx_key) {
2787 Some(&m) => m,
2788 None => {
2789 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2790 map.insert(ctx_key, m);
2791 m
2792 }
2793 }
2794 };
2795 let cname = std::ffi::CString::new(name)?;
2796 let mut f: cu::CUfunction = std::ptr::null_mut();
2797 let mut r =
2798 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2799 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2800 let qmodule = {
2801 let mut mods = QMODULES.lock().unwrap();
2802 let map = mods.get_or_insert_with(Default::default);
2803 match map.get(&ctx_key) {
2804 Some(&m) => m,
2805 None => {
2806 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2807 map.insert(ctx_key, m);
2808 m
2809 }
2810 }
2811 };
2812 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2813 }
2814 if r != cu::CUresult::CUDA_SUCCESS {
2815 return Err(format!("pdl_func {name}: {r:?}").into());
2816 }
2817 FNS.lock()
2818 .unwrap()
2819 .get_or_insert_with(Default::default)
2820 .insert((ctx_key, name), f as usize);
2821 Ok(f)
2822 }
2823
2824 unsafe fn launch_pdl_flash(
2836 &self,
2837 g: bool,
2838 name: &'static str,
2839 grid: (u32, u32, u32),
2840 block: (u32, u32, u32),
2841 smem: u32,
2842 params: &mut [*mut std::ffi::c_void],
2843 ) -> Result<(), Box<dyn std::error::Error>> {
2844 use cudarc::driver::sys as cu;
2845 let f = self.pdl_func_flash(g, name)?;
2846 if smem > 0 {
2847 let r =
2849 unsafe {
2850 cu::cuFuncSetAttribute(f,
2851 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2852 smem as i32)
2853 };
2854 if r != cu::CUresult::CUDA_SUCCESS {
2855 return Err(format!("pdl smem attr {name}: {r:?}").into());
2856 }
2857 }
2858 let mut attr = cu::CUlaunchAttribute {
2859 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2860 pad: [0; 4],
2861 value: cu::CUlaunchAttributeValue {
2862 programmaticStreamSerializationAllowed: 1,
2863 },
2864 };
2865 let cfg = cu::CUlaunchConfig {
2866 gridDimX: grid.0,
2867 gridDimY: grid.1,
2868 gridDimZ: grid.2,
2869 blockDimX: block.0,
2870 blockDimY: block.1,
2871 blockDimZ: block.2,
2872 sharedMemBytes: smem,
2873 hStream: self.gpu.stream().cu_stream(),
2874 attrs: &mut attr,
2875 numAttrs: 1,
2876 };
2877 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2878 if r != cu::CUresult::CUDA_SUCCESS {
2879 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2880 }
2881 Ok(())
2882 }
2883
2884 unsafe fn launch_pdl(
2885 &self,
2886 name: &'static str,
2887 grid: (u32, u32, u32),
2888 block: (u32, u32, u32),
2889 params: &mut [*mut std::ffi::c_void],
2890 ) -> Result<(), Box<dyn std::error::Error>> {
2891 use cudarc::driver::sys as cu;
2892 let f = self.pdl_func(name)?;
2893 let mut attr = cu::CUlaunchAttribute {
2894 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2895 pad: [0; 4],
2896 value: cu::CUlaunchAttributeValue {
2897 programmaticStreamSerializationAllowed: 1,
2898 },
2899 };
2900 let cfg = cu::CUlaunchConfig {
2901 gridDimX: grid.0,
2902 gridDimY: grid.1,
2903 gridDimZ: grid.2,
2904 blockDimX: block.0,
2905 blockDimY: block.1,
2906 blockDimZ: block.2,
2907 sharedMemBytes: 0,
2908 hStream: self.gpu.stream().cu_stream(),
2909 attrs: &mut attr,
2910 numAttrs: 1,
2911 };
2912 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2913 if r != cu::CUresult::CUDA_SUCCESS {
2914 return Err(format!("launch_pdl {name}: {r:?}").into());
2915 }
2916 Ok(())
2917 }
2918
2919 pub fn prefetch_weight_l2(
2922 &self,
2923 w: &crate::model::GpuTensor,
2924 ) -> Result<(), Box<dyn std::error::Error>> {
2925 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2926 let p = rp4.as_ref().unwrap_or(bytes);
2927 self.prefetch_l2(p, p.len())?;
2928 }
2929 Ok(())
2930 }
2931
2932 pub fn gather_row_bf16(
2935 &self,
2936 table: &CudaSlice<u8>,
2937 tok: &CudaSlice<u32>,
2938 idx: usize,
2939 dst: &mut CudaSlice<f32>,
2940 ncols: usize,
2941 ) -> Result<(), Box<dyn std::error::Error>> {
2942 let f = self.func("gather_row_bf16_f32");
2943 let cfg = LaunchConfig {
2944 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2945 block_dim: (256, 1, 1),
2946 shared_mem_bytes: 0,
2947 };
2948 let (nc, ix) = (ncols as i32, idx as i32);
2949 let __s_b = self.gpu.stream();
2950 let mut b = __s_b.launch_builder(&f);
2951 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2952 unsafe {
2953 b.launch(cfg)?;
2954 }
2955 Ok(())
2956 }
2957
2958 #[allow(clippy::too_many_arguments)]
2964 pub fn dflash2_dynconv(
2965 &self,
2966 x: &CudaSlice<f32>,
2967 dyn_: &CudaSlice<f32>,
2968 base: &CudaSlice<f32>,
2969 out: &mut CudaSlice<f32>,
2970 rows: usize,
2971 hidden: usize,
2972 group_size: usize,
2973 ksize: usize,
2974 half: usize,
2975 ) -> Result<(), Box<dyn std::error::Error>> {
2976 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2977 let f = self.func("dflash2_dynconv_f32");
2978 let n = rows * hidden;
2979 let cfg = LaunchConfig {
2980 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2981 block_dim: (256, 1, 1),
2982 shared_mem_bytes: 0,
2983 };
2984 let (ri, hi, gi, ki, hf) = (
2985 rows as i32,
2986 hidden as i32,
2987 group_size as i32,
2988 ksize as i32,
2989 half as i32,
2990 );
2991 let __s_b = self.gpu.stream();
2992 let mut b = __s_b.launch_builder(&f);
2993 b.arg(x)
2994 .arg(dyn_)
2995 .arg(base)
2996 .arg(out)
2997 .arg(&ri)
2998 .arg(&hi)
2999 .arg(&gi)
3000 .arg(&ki)
3001 .arg(&hf);
3002 unsafe {
3003 b.launch(cfg)?;
3004 }
3005 Ok(())
3006 }
3007
3008 pub fn topk_rows(
3012 &self,
3013 logits: &CudaSlice<f32>,
3014 n_rows: usize,
3015 n_cols: usize,
3016 k: usize,
3017 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
3018 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
3019 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
3020 let f = self.func("topk_rows_f32");
3021 let nth = 256usize;
3022 let mut vals = self.uninit(n_rows * k)?;
3023 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
3024 let cfg = LaunchConfig {
3025 grid_dim: (n_rows as u32, 1, 1),
3026 block_dim: (nth as u32, 1, 1),
3027 shared_mem_bytes: (nth * k * 8) as u32,
3028 };
3029 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
3030 let __s_b = self.gpu.stream();
3031 let mut b = __s_b.launch_builder(&f);
3032 b.arg(logits)
3033 .arg(&nr)
3034 .arg(&nc)
3035 .arg(&ki)
3036 .arg(&mut vals)
3037 .arg(&mut idxs);
3038 unsafe {
3039 b.launch(cfg)?;
3040 }
3041 Ok((vals, idxs))
3042 }
3043
3044 pub fn add_row_inplace(
3046 &self,
3047 logits: &mut CudaSlice<f32>,
3048 bias: &CudaSlice<f32>,
3049 n: usize,
3050 row_off: usize,
3051 ) -> Result<(), Box<dyn std::error::Error>> {
3052 let f = self.func("add_row_inplace_f32");
3053 let cfg = LaunchConfig {
3054 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3055 block_dim: (256, 1, 1),
3056 shared_mem_bytes: 0,
3057 };
3058 let (ni, off) = (n as i32, row_off as i64);
3059 let __s_b = self.gpu.stream();
3060 let mut b = __s_b.launch_builder(&f);
3061 b.arg(logits).arg(bias).arg(&ni).arg(&off);
3062 unsafe {
3063 b.launch(cfg)?;
3064 }
3065 Ok(())
3066 }
3067
3068 pub fn prefetch_l2(
3070 &self,
3071 p: &CudaSlice<u8>,
3072 n: usize,
3073 ) -> Result<(), Box<dyn std::error::Error>> {
3074 let f = self.func("prefetch_l2_bytes");
3075 let lines = n.div_ceil(128);
3076 let ni = n as i64;
3077 let cfg = LaunchConfig {
3078 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
3079 block_dim: (256, 1, 1),
3080 shared_mem_bytes: 0,
3081 };
3082 let __s_b = self.gpu.stream();
3083 let mut b = __s_b.launch_builder(&f);
3084 b.arg(p).arg(&ni);
3085 unsafe {
3086 b.launch(cfg)?;
3087 }
3088 Ok(())
3089 }
3090
3091 pub fn router_gemv(
3094 &self,
3095 w: &CudaSlice<f32>,
3096 x: &CudaSlice<f32>,
3097 n_embd: usize,
3098 n_experts: usize,
3099 t: usize,
3100 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3101 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
3107 Ok("0") => false,
3108 Ok(_) => true,
3109 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
3110 };
3111 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
3121 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
3122 }
3123
3124 pub fn router_gemv_form(
3127 &self,
3128 w: &CudaSlice<f32>,
3129 x: &CudaSlice<f32>,
3130 n_embd: usize,
3131 n_experts: usize,
3132 t: usize,
3133 w8: bool,
3134 batch: bool,
3135 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3136 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
3137 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
3138 let f = if batch {
3139 self.func("router_gemv_f32_w8_batch")
3140 } else if w8 {
3141 self.func("router_gemv_f32_w8")
3142 } else {
3143 self.func("router_gemv_f32")
3144 };
3145 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
3146 let cfg = if batch {
3147 LaunchConfig {
3148 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
3149 block_dim: (32, 8, 1),
3150 shared_mem_bytes: 0,
3151 }
3152 } else {
3153 LaunchConfig {
3154 grid_dim: (n_experts as u32, t as u32, 1),
3155 block_dim: (32, if w8 { 8 } else { 1 }, 1),
3156 shared_mem_bytes: 0,
3157 }
3158 };
3159 let __s_b = self.gpu.stream();
3160 let mut b = __s_b.launch_builder(&f);
3161 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
3162 unsafe {
3163 b.launch(cfg)?;
3164 }
3165 Ok(y)
3166 }
3167
3168 pub fn router_gemv_into(
3171 &self,
3172 w: &CudaSlice<f32>,
3173 x: &CudaSlice<f32>,
3174 y: &mut CudaSlice<f32>,
3175 n_embd: usize,
3176 n_experts: usize,
3177 t: usize,
3178 ) -> Result<(), Box<dyn std::error::Error>> {
3179 if y.len() < t * n_experts {
3180 return Err("router_gemv_into output too small".into());
3181 }
3182 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
3183 Ok("0") => false,
3184 Ok(_) => true,
3185 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
3186 };
3187 let f = if w8 {
3188 self.func("router_gemv_f32_w8")
3189 } else {
3190 self.func("router_gemv_f32")
3191 };
3192 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
3193 let cfg = LaunchConfig {
3194 grid_dim: (n_experts as u32, t as u32, 1),
3195 block_dim: (32, if w8 { 8 } else { 1 }, 1),
3196 shared_mem_bytes: 0,
3197 };
3198 let __s_b = self.gpu.stream();
3199 let mut b = __s_b.launch_builder(&f);
3200 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
3201 unsafe {
3202 b.launch(cfg)?;
3203 }
3204 Ok(())
3205 }
3206
3207 pub fn rows_permute(
3209 &self,
3210 src: &CudaSlice<f32>,
3211 idx: &CudaSlice<i32>,
3212 nrows: usize,
3213 ncols: usize,
3214 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3215 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
3216 let f = self.func("rows_permute_f32");
3217 let (nc, nr) = (ncols as i32, nrows as i32);
3218 let cfg = LaunchConfig {
3219 grid_dim: (nrows as u32, 1, 1),
3220 block_dim: (256, 1, 1),
3221 shared_mem_bytes: 0,
3222 };
3223 let __s_b = self.gpu.stream();
3224 let mut b = __s_b.launch_builder(&f);
3225 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
3226 unsafe {
3227 b.launch(cfg)?;
3228 }
3229 Ok(dst)
3230 }
3231
3232 pub fn sigmoid_dot_rows(
3237 &self,
3238 x: &CudaSlice<f32>,
3239 w: &CudaSlice<f32>,
3240 n_embd: usize,
3241 t: usize,
3242 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3243 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3246 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
3247 let gs = self.linear(x, w, t, n_embd, 1)?;
3248 let mut g = self.uninit(t)?;
3249 self.sigmoid(&gs, &mut g, t)?;
3250 return Ok(g);
3251 }
3252 let mut g = self.alloc_uninit::<f32>(t)?;
3258 let f = self.func("sigmoid_dot_rows_f32");
3259 let (ne, ti) = (n_embd as i32, t as i32);
3260 let cfg = LaunchConfig {
3261 grid_dim: (t as u32, 1, 1),
3262 block_dim: (32, 8, 1),
3263 shared_mem_bytes: 0,
3264 };
3265 let __s_b = self.gpu.stream();
3266 let mut b = __s_b.launch_builder(&f);
3267 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
3268 unsafe {
3269 b.launch(cfg)?;
3270 }
3271 Ok(g)
3272 }
3273
3274 pub fn sigmoid_dot_rows_into(
3276 &self,
3277 x: &CudaSlice<f32>,
3278 w: &CudaSlice<f32>,
3279 g: &mut CudaSlice<f32>,
3280 n_embd: usize,
3281 t: usize,
3282 ) -> Result<(), Box<dyn std::error::Error>> {
3283 if g.len() < t {
3284 return Err("sigmoid_dot_rows_into output too small".into());
3285 }
3286 let f = self.func("sigmoid_dot_rows_f32");
3287 let (ne, ti) = (n_embd as i32, t as i32);
3288 let cfg = LaunchConfig {
3289 grid_dim: (t as u32, 1, 1),
3290 block_dim: (32, 8, 1),
3291 shared_mem_bytes: 0,
3292 };
3293 let __s_b = self.gpu.stream();
3294 let mut b = __s_b.launch_builder(&f);
3295 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
3296 unsafe {
3297 b.launch(cfg)?;
3298 }
3299 Ok(())
3300 }
3301
3302 pub fn spec_rollback_stream(
3304 &self,
3305 len_ptrs: &CudaSlice<u64>,
3306 pos_start: &CudaSlice<i32>,
3307 acc: &CudaSlice<u32>,
3308 base: usize,
3309 n_rows: usize,
3310 ) -> Result<(), Box<dyn std::error::Error>> {
3311 let f = self.func("spec_rollback_stream");
3312 let (b, nr) = (base as i32, n_rows as i32);
3313 let cfg = LaunchConfig {
3314 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
3315 block_dim: (64, 1, 1),
3316 shared_mem_bytes: 0,
3317 };
3318 let __s_bl = self.gpu.stream();
3319 let mut bl = __s_bl.launch_builder(&f);
3320 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
3321 unsafe {
3322 bl.launch(cfg)?;
3323 }
3324 Ok(())
3325 }
3326
3327 pub fn plain_tok_ring(
3329 &self,
3330 vam: &CudaSlice<u32>,
3331 pos_start: &CudaSlice<i32>,
3332 base: usize,
3333 ring: &mut CudaSlice<u32>,
3334 ) -> Result<(), Box<dyn std::error::Error>> {
3335 let f = self.func("plain_tok_ring");
3336 let (b, cap) = (base as i32, ring.len() as i32);
3337 let cfg = LaunchConfig {
3338 grid_dim: (1, 1, 1),
3339 block_dim: (32, 1, 1),
3340 shared_mem_bytes: 0,
3341 };
3342 let __s_bl = self.gpu.stream();
3343 let mut bl = __s_bl.launch_builder(&f);
3344 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
3345 unsafe {
3346 bl.launch(cfg)?;
3347 }
3348 Ok(())
3349 }
3350
3351 pub fn spec_ring_commit(
3353 &self,
3354 vtok: &CudaSlice<u32>,
3355 acc: &CudaSlice<u32>,
3356 brk: &CudaSlice<u32>,
3357 ring: &mut CudaSlice<u32>,
3358 pend: &mut CudaSlice<u32>,
3359 ) -> Result<(), Box<dyn std::error::Error>> {
3360 let f = self.func("spec_ring_commit");
3361 let cfg = LaunchConfig {
3362 grid_dim: (1, 1, 1),
3363 block_dim: (32, 1, 1),
3364 shared_mem_bytes: 0,
3365 };
3366 let __s_b = self.gpu.stream();
3367 let mut b = __s_b.launch_builder(&f);
3368 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
3369 unsafe {
3370 b.launch(cfg)?;
3371 }
3372 Ok(())
3373 }
3374 pub fn i32_copy_add(
3375 &self,
3376 src: &CudaSlice<i32>,
3377 dst: &mut CudaSlice<i32>,
3378 delta: i32,
3379 ) -> Result<(), Box<dyn std::error::Error>> {
3380 let f = self.func("i32_copy_add");
3381 let cfg = LaunchConfig {
3382 grid_dim: (1, 1, 1),
3383 block_dim: (32, 1, 1),
3384 shared_mem_bytes: 0,
3385 };
3386 let __s_b = self.gpu.stream();
3387 let mut b = __s_b.launch_builder(&f);
3388 b.arg(src).arg(dst).arg(&delta);
3389 unsafe {
3390 b.launch(cfg)?;
3391 }
3392 Ok(())
3393 }
3394 pub fn u32_copy(
3395 &self,
3396 src: &CudaSlice<u32>,
3397 dst: &mut CudaSlice<u32>,
3398 ) -> Result<(), Box<dyn std::error::Error>> {
3399 let f = self.func("u32_copy");
3400 let cfg = LaunchConfig {
3401 grid_dim: (1, 1, 1),
3402 block_dim: (32, 1, 1),
3403 shared_mem_bytes: 0,
3404 };
3405 let __s_b = self.gpu.stream();
3406 let mut b = __s_b.launch_builder(&f);
3407 b.arg(src).arg(dst);
3408 unsafe {
3409 b.launch(cfg)?;
3410 }
3411 Ok(())
3412 }
3413
3414 pub fn spec_adapt_k(
3418 &self,
3419 acc: &CudaSlice<u32>,
3420 brk: &mut CudaSlice<u32>,
3421 floor: usize,
3422 cap: usize,
3423 ) -> Result<(), Box<dyn std::error::Error>> {
3424 let f = self.func("spec_adapt_k");
3425 let (fl, cp) = (floor as i32, cap as i32);
3426 let cfg = LaunchConfig {
3427 grid_dim: (1, 1, 1),
3428 block_dim: (32, 1, 1),
3429 shared_mem_bytes: 0,
3430 };
3431 let __s_b = self.gpu.stream();
3432 let mut b = __s_b.launch_builder(&f);
3433 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
3434 unsafe {
3435 b.launch(cfg)?;
3436 }
3437 Ok(())
3438 }
3439
3440 pub fn spec_accept_greedy_dc(
3442 &self,
3443 preds: &CudaSlice<u32>,
3444 vtok: &CudaSlice<u32>,
3445 last_pred: &CudaSlice<u32>,
3446 brk: &CudaSlice<u32>,
3447 out: &mut CudaSlice<u32>,
3448 ) -> Result<(), Box<dyn std::error::Error>> {
3449 let f = self.func("spec_accept_greedy_dc");
3450 let cfg = LaunchConfig {
3451 grid_dim: (1, 1, 1),
3452 block_dim: (32, 1, 1),
3453 shared_mem_bytes: 0,
3454 };
3455 let __s_b = self.gpu.stream();
3456 let mut b = __s_b.launch_builder(&f);
3457 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3458 unsafe {
3459 b.launch(cfg)?;
3460 }
3461 Ok(())
3462 }
3463
3464 pub fn pos_iota(
3466 &self,
3467 pos0: &CudaSlice<i32>,
3468 out: &mut CudaSlice<i32>,
3469 t: usize,
3470 ) -> Result<(), Box<dyn std::error::Error>> {
3471 let f = self.func("pos_iota_i32");
3472 let ti = t as i32;
3473 let cfg = LaunchConfig {
3474 grid_dim: (1, 1, 1),
3475 block_dim: (t.max(1) as u32, 1, 1),
3476 shared_mem_bytes: 0,
3477 };
3478 let __s_b = self.gpu.stream();
3479 let mut b = __s_b.launch_builder(&f);
3480 b.arg(pos0).arg(out).arg(&ti);
3481 unsafe {
3482 b.launch(cfg)?;
3483 }
3484 Ok(())
3485 }
3486 #[allow(clippy::too_many_arguments)]
3487 pub fn append_kv_quantized_rows_dc(
3488 &self,
3489 k_rows: &CudaSlice<f32>,
3490 v_rows: &CudaSlice<f32>,
3491 kc: &mut CudaSlice<u8>,
3492 vc: &mut CudaSlice<u8>,
3493 t0_dev: &CudaSlice<i32>,
3494 t: usize,
3495 kv_dim_k: usize,
3496 kv_dim_v: usize,
3497 k_tok_bytes: usize,
3498 v_tok_bytes: usize,
3499 g: bool,
3500 ) -> Result<(), Box<dyn std::error::Error>> {
3501 let f = if g {
3502 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3503 } else {
3504 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3505 };
3506 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3507 let cfg = LaunchConfig {
3508 grid_dim: (nblk, t as u32, 1),
3509 block_dim: (32, 1, 1),
3510 shared_mem_bytes: 0,
3511 };
3512 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3513 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3514 let __s_b = self.gpu.stream();
3515 let mut b = __s_b.launch_builder(&f);
3516 b.arg(k_rows)
3517 .arg(v_rows)
3518 .arg(kc)
3519 .arg(vc)
3520 .arg(t0_dev)
3521 .arg(&kdk)
3522 .arg(&kdv)
3523 .arg(&ktb)
3524 .arg(&vtb);
3525 unsafe {
3526 b.launch(cfg)?;
3527 }
3528 Ok(())
3529 }
3530
3531 #[allow(clippy::too_many_arguments)]
3534 pub fn append_kv_quantized_row_dc_inc(
3535 &self,
3536 k_row: &CudaSlice<f32>,
3537 v_row: &CudaSlice<f32>,
3538 kc: &mut CudaSlice<u8>,
3539 vc: &mut CudaSlice<u8>,
3540 t0_dev: &mut CudaSlice<i32>,
3541 kv_dim_k: usize,
3542 kv_dim_v: usize,
3543 k_tok_bytes: usize,
3544 v_tok_bytes: usize,
3545 g: bool,
3546 ) -> Result<(), Box<dyn std::error::Error>> {
3547 let f = if g {
3548 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3549 } else {
3550 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3551 };
3552 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3553 let cfg = LaunchConfig {
3554 grid_dim: (1, 1, 1),
3555 block_dim: (nthreads, 1, 1),
3556 shared_mem_bytes: 0,
3557 };
3558 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3559 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3560 let __s_b = self.gpu.stream();
3561 let mut b = __s_b.launch_builder(&f);
3562 b.arg(k_row)
3563 .arg(v_row)
3564 .arg(kc)
3565 .arg(vc)
3566 .arg(t0_dev)
3567 .arg(&kdk)
3568 .arg(&kdv)
3569 .arg(&ktb)
3570 .arg(&vtb);
3571 unsafe {
3572 b.launch(cfg)?;
3573 }
3574 Ok(())
3575 }
3576
3577 pub fn pack_tok_p(
3579 &self,
3580 tok: &CudaSlice<u32>,
3581 p: &CudaSlice<f32>,
3582 out: &mut CudaSlice<u32>,
3583 slot: usize,
3584 ) -> Result<(), Box<dyn std::error::Error>> {
3585 let f = self.func("pack_tok_p");
3586 let sl = slot as i32;
3587 let cfg = LaunchConfig {
3588 grid_dim: (1, 1, 1),
3589 block_dim: (32, 1, 1),
3590 shared_mem_bytes: 0,
3591 };
3592 let __s_b = self.gpu.stream();
3593 let mut b = __s_b.launch_builder(&f);
3594 b.arg(tok).arg(p).arg(out).arg(&sl);
3595 unsafe {
3596 b.launch(cfg)?;
3597 }
3598 Ok(())
3599 }
3600 pub fn tok_map_u32(
3601 &self,
3602 tok: &mut CudaSlice<u32>,
3603 map: &CudaSlice<u32>,
3604 ) -> Result<(), Box<dyn std::error::Error>> {
3605 let f = self.func("tok_map_u32");
3606 let cfg = LaunchConfig {
3607 grid_dim: (1, 1, 1),
3608 block_dim: (32, 1, 1),
3609 shared_mem_bytes: 0,
3610 };
3611 let __s_b = self.gpu.stream();
3612 let mut b = __s_b.launch_builder(&f);
3613 b.arg(tok).arg(map);
3614 unsafe {
3615 b.launch(cfg)?;
3616 }
3617 Ok(())
3618 }
3619
3620 #[allow(clippy::too_many_arguments)]
3622 pub fn spec_assemble_verify(
3623 &self,
3624 tokp: &CudaSlice<u32>,
3625 pend: &CudaSlice<u32>,
3626 d2t: Option<&CudaSlice<u32>>,
3627 vtok: &mut CudaSlice<u32>,
3628 brk: &mut CudaSlice<u32>,
3629 p_min: f32,
3630 k: usize,
3631 pmin0: bool,
3632 ) -> Result<(), Box<dyn std::error::Error>> {
3633 let f = self.func("spec_assemble_verify");
3634 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3635 let cfg = LaunchConfig {
3636 grid_dim: (1, 1, 1),
3637 block_dim: (32, 1, 1),
3638 shared_mem_bytes: 0,
3639 };
3640 let __s_b = self.gpu.stream();
3641 let mut b = __s_b.launch_builder(&f);
3642 match d2t {
3643 Some(m) => {
3644 b.arg(tokp)
3645 .arg(pend)
3646 .arg(m)
3647 .arg(vtok)
3648 .arg(brk)
3649 .arg(&p_min)
3650 .arg(&ki)
3651 .arg(&pm);
3652 unsafe {
3653 b.launch(cfg)?;
3654 }
3655 }
3656 None => {
3657 let null: u64 = 0;
3658 b.arg(tokp)
3659 .arg(pend)
3660 .arg(&null)
3661 .arg(vtok)
3662 .arg(brk)
3663 .arg(&p_min)
3664 .arg(&ki)
3665 .arg(&pm);
3666 unsafe {
3667 b.launch(cfg)?;
3668 }
3669 }
3670 }
3671 Ok(())
3672 }
3673
3674 #[allow(clippy::too_many_arguments)]
3676 pub fn ssm_conv_ring_rebuild_dc(
3677 &self,
3678 qkv_tm: &CudaSlice<f32>,
3679 ring_old: &CudaSlice<f32>,
3680 conv_state: &mut CudaSlice<f32>,
3681 conv_dim: usize,
3682 acc: &CudaSlice<u32>,
3683 base: usize,
3684 t_v: usize,
3685 d_conv: usize,
3686 ) -> Result<(), Box<dyn std::error::Error>> {
3687 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3688 let n = conv_dim * (d_conv - 1);
3689 let cfg = LaunchConfig::for_num_elems(n as u32);
3690 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3691 let __s_b = self.gpu.stream();
3692 let mut b = __s_b.launch_builder(&f);
3693 b.arg(qkv_tm)
3694 .arg(ring_old)
3695 .arg(conv_state)
3696 .arg(&cd)
3697 .arg(acc)
3698 .arg(&b0)
3699 .arg(&tv)
3700 .arg(&dc);
3701 unsafe {
3702 b.launch(cfg)?;
3703 }
3704 Ok(())
3705 }
3706 #[allow(clippy::too_many_arguments)]
3707 pub fn gdn_scan_s128_dc(
3708 &self,
3709 q: &CudaSlice<f32>,
3710 k: &CudaSlice<f32>,
3711 v: &CudaSlice<f32>,
3712 g: &CudaSlice<f32>,
3713 beta: &CudaSlice<f32>,
3714 state_in: &CudaSlice<f32>,
3715 state_out: &mut CudaSlice<f32>,
3716 o: &mut CudaSlice<f32>,
3717 n_head: usize,
3718 acc: &CudaSlice<u32>,
3719 base: usize,
3720 t_v: usize,
3721 scale: f32,
3722 ) -> Result<(), Box<dyn std::error::Error>> {
3723 let f = self.func("gdn_scan_s128_dc");
3724 const S_V: u32 = 128;
3725 const WARP: u32 = 32;
3726 const COLS_PER_BLOCK: u32 = 4;
3727 let cfg = LaunchConfig {
3728 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3729 block_dim: (WARP, COLS_PER_BLOCK, 1),
3730 shared_mem_bytes: 0,
3731 };
3732 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3733 let __s_b = self.gpu.stream();
3734 let mut b = __s_b.launch_builder(&f);
3735 b.arg(q)
3736 .arg(k)
3737 .arg(v)
3738 .arg(g)
3739 .arg(beta)
3740 .arg(state_in)
3741 .arg(state_out)
3742 .arg(o)
3743 .arg(&h)
3744 .arg(acc)
3745 .arg(&b0)
3746 .arg(&tv)
3747 .arg(&scale);
3748 unsafe {
3749 b.launch(cfg)?;
3750 }
3751 Ok(())
3752 }
3753
3754 pub fn spec_rollback_kv(
3756 &self,
3757 len_ptrs: &CudaSlice<u64>,
3758 saved: &CudaSlice<i32>,
3759 acc: &CudaSlice<u32>,
3760 base: usize,
3761 n_layer: usize,
3762 ) -> Result<(), Box<dyn std::error::Error>> {
3763 let f = self.func("spec_rollback_kv");
3764 let (b, nl) = (base as i32, n_layer as i32);
3765 let cfg = LaunchConfig {
3766 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3767 block_dim: (64, 1, 1),
3768 shared_mem_bytes: 0,
3769 };
3770 let __s_bl = self.gpu.stream();
3771 let mut bl = __s_bl.launch_builder(&f);
3772 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3773 unsafe {
3774 bl.launch(cfg)?;
3775 }
3776 Ok(())
3777 }
3778
3779 pub fn spec_fork_valid(
3781 &self,
3782 acc: &CudaSlice<u32>,
3783 optimistic_pending: u32,
3784 valid: &mut CudaSlice<u32>,
3785 ) -> Result<(), Box<dyn std::error::Error>> {
3786 let f = self.func("spec_fork_valid");
3787 let cfg = LaunchConfig {
3788 grid_dim: (1, 1, 1),
3789 block_dim: (1, 1, 1),
3790 shared_mem_bytes: 0,
3791 };
3792 let __s_bl = self.gpu.stream();
3793 let mut bl = __s_bl.launch_builder(&f);
3794 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3795 unsafe {
3796 bl.launch(cfg)?;
3797 }
3798 Ok(())
3799 }
3800
3801 pub fn spec_fork_reconcile_kv(
3803 &self,
3804 len_ptrs: &CudaSlice<u64>,
3805 saved: &CudaSlice<i32>,
3806 acc: &CudaSlice<u32>,
3807 valid: &CudaSlice<u32>,
3808 base: usize,
3809 n_layer: usize,
3810 ) -> Result<(), Box<dyn std::error::Error>> {
3811 let f = self.func("spec_fork_reconcile_kv");
3812 let (b, nl) = (base as i32, n_layer as i32);
3813 let cfg = LaunchConfig {
3814 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3815 block_dim: (64, 1, 1),
3816 shared_mem_bytes: 0,
3817 };
3818 let __s_bl = self.gpu.stream();
3819 let mut bl = __s_bl.launch_builder(&f);
3820 bl.arg(len_ptrs)
3821 .arg(saved)
3822 .arg(acc)
3823 .arg(valid)
3824 .arg(&b)
3825 .arg(&nl);
3826 unsafe {
3827 bl.launch(cfg)?;
3828 }
3829 Ok(())
3830 }
3831
3832 pub fn spec_fork_restore_f32(
3834 &self,
3835 snapshot: &CudaSlice<f32>,
3836 state: &mut CudaSlice<f32>,
3837 valid: &CudaSlice<u32>,
3838 ) -> Result<(), Box<dyn std::error::Error>> {
3839 assert_eq!(
3840 snapshot.len(),
3841 state.len(),
3842 "fork recurrent snapshot shape mismatch"
3843 );
3844 let f = self.func("spec_fork_restore_f32");
3845 let n = state.len() as i32;
3846 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3847 let cfg = LaunchConfig {
3848 grid_dim: (blocks, 1, 1),
3849 block_dim: (256, 1, 1),
3850 shared_mem_bytes: 0,
3851 };
3852 let __s_bl = self.gpu.stream();
3853 let mut bl = __s_bl.launch_builder(&f);
3854 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3855 unsafe {
3856 bl.launch(cfg)?;
3857 }
3858 Ok(())
3859 }
3860
3861 pub fn spec_seed_gather(
3864 &self,
3865 vx: &CudaSlice<f32>,
3866 fill_prev: &CudaSlice<f32>,
3867 acc: &CudaSlice<u32>,
3868 h_seed: &mut CudaSlice<f32>,
3869 base: usize,
3870 n_embd: usize,
3871 ) -> Result<(), Box<dyn std::error::Error>> {
3872 let f = self.func("spec_seed_gather");
3873 let (b, ne) = (base as i32, n_embd as i32);
3874 let cfg = LaunchConfig {
3875 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3876 block_dim: (256, 1, 1),
3877 shared_mem_bytes: 0,
3878 };
3879 let __s_bl = self.gpu.stream();
3880 let mut bl = __s_bl.launch_builder(&f);
3881 bl.arg(vx)
3882 .arg(fill_prev)
3883 .arg(acc)
3884 .arg(h_seed)
3885 .arg(&b)
3886 .arg(&ne);
3887 unsafe {
3888 bl.launch(cfg)?;
3889 }
3890 Ok(())
3891 }
3892
3893 pub fn spec_accept_greedy(
3895 &self,
3896 preds: &CudaSlice<u32>,
3897 draft: &CudaSlice<u32>,
3898 last_pred: u32,
3899 base: usize,
3900 k_round: usize,
3901 out: &mut CudaSlice<u32>,
3902 ) -> Result<(), Box<dyn std::error::Error>> {
3903 let f = self.func("spec_accept_greedy");
3904 let (b, k) = (base as i32, k_round as i32);
3905 let cfg = LaunchConfig {
3906 grid_dim: (1, 1, 1),
3907 block_dim: (32, 1, 1),
3908 shared_mem_bytes: 0,
3909 };
3910 let __s_bl = self.gpu.stream();
3911 let mut bl = __s_bl.launch_builder(&f);
3912 bl.arg(preds)
3913 .arg(draft)
3914 .arg(&last_pred)
3915 .arg(&b)
3916 .arg(&k)
3917 .arg(out);
3918 unsafe {
3919 bl.launch(cfg)?;
3920 }
3921 Ok(())
3922 }
3923
3924 pub fn gumbel_perturb(
3931 &self,
3932 x: &CudaSlice<f32>,
3933 y: &mut CudaSlice<f32>,
3934 n: usize,
3935 seed: u64,
3936 stream_pos: u32,
3937 temp: f32,
3938 ) -> Result<(), Box<dyn std::error::Error>> {
3939 let f = self.func("gumbel_perturb_f32");
3940 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3941 let cfg = LaunchConfig {
3942 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3943 block_dim: (256, 1, 1),
3944 shared_mem_bytes: 0,
3945 };
3946 let __s_b = self.gpu.stream();
3947 let mut b = __s_b.launch_builder(&f);
3948 b.arg(x)
3949 .arg(&mut *y)
3950 .arg(&ni)
3951 .arg(&slo)
3952 .arg(&shi)
3953 .arg(&stream_pos)
3954 .arg(&temp);
3955 unsafe {
3956 b.launch(cfg)?;
3957 }
3958 Ok(())
3959 }
3960
3961 pub fn mask_logits_col(
3969 &self,
3970 logits: &mut CudaSlice<f32>,
3971 mask: &CudaSlice<u32>,
3972 col: usize,
3973 n: usize,
3974 mask_words: usize,
3975 ) -> Result<(), Box<dyn std::error::Error>> {
3976 let f = self.func("mask_logits_f32");
3977 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3978 let cfg = LaunchConfig {
3979 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3980 block_dim: (256, 1, 1),
3981 shared_mem_bytes: 0,
3982 };
3983 let __s_b = self.gpu.stream();
3984 let mut b = __s_b.launch_builder(&f);
3985 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3986 unsafe {
3987 b.launch(cfg)?;
3988 }
3989 Ok(())
3990 }
3991
3992 pub fn gumbel_perturb_col(
3999 &self,
4000 x: &CudaSlice<f32>,
4001 col: usize,
4002 y: &mut CudaSlice<f32>,
4003 n: usize,
4004 seed: u64,
4005 stream_pos: u32,
4006 temp: f32,
4007 ) -> Result<(), Box<dyn std::error::Error>> {
4008 let f = self.func("gumbel_perturb_f32");
4009 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4010 let col_view = x.slice(col * n..(col + 1) * n);
4011 let cfg = LaunchConfig {
4012 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4013 block_dim: (256, 1, 1),
4014 shared_mem_bytes: 0,
4015 };
4016 let __s_b = self.gpu.stream();
4017 let mut b = __s_b.launch_builder(&f);
4018 b.arg(&col_view)
4019 .arg(&mut *y)
4020 .arg(&ni)
4021 .arg(&slo)
4022 .arg(&shi)
4023 .arg(&stream_pos)
4024 .arg(&temp);
4025 unsafe {
4026 b.launch(cfg)?;
4027 }
4028 Ok(())
4029 }
4030
4031 #[allow(clippy::too_many_arguments)]
4037 pub fn gumbel_perturb_filtered_col(
4038 &self,
4039 x: &CudaSlice<f32>,
4040 col: usize,
4041 y: &mut CudaSlice<f32>,
4042 n: usize,
4043 seed: u64,
4044 stream_pos: u32,
4045 temp: f32,
4046 stat_max: &CudaSlice<f32>,
4047 stat_th: &CudaSlice<f32>,
4048 stat_idx: usize,
4049 ) -> Result<(), Box<dyn std::error::Error>> {
4050 let f = self.func("gumbel_perturb_filtered_col_f32");
4051 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4052 let (ci, si) = (col as i32, stat_idx as i32);
4053 let cfg = LaunchConfig {
4054 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4055 block_dim: (256, 1, 1),
4056 shared_mem_bytes: 0,
4057 };
4058 let __s_b = self.gpu.stream();
4059 let mut b = __s_b.launch_builder(&f);
4060 b.arg(x)
4061 .arg(&ci)
4062 .arg(&mut *y)
4063 .arg(&ni)
4064 .arg(&slo)
4065 .arg(&shi)
4066 .arg(&stream_pos)
4067 .arg(&temp)
4068 .arg(stat_max)
4069 .arg(stat_th)
4070 .arg(&si);
4071 unsafe {
4072 b.launch(cfg)?;
4073 }
4074 Ok(())
4075 }
4076
4077 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
4082 let f = self.func("memra_sctr_inc");
4083 let cfg = LaunchConfig {
4084 grid_dim: (1, 1, 1),
4085 block_dim: (1, 1, 1),
4086 shared_mem_bytes: 0,
4087 };
4088 let __s_b = self.gpu.stream();
4089 let mut b = __s_b.launch_builder(&f);
4090 b.arg(&mut *ctr);
4091 unsafe {
4092 b.launch(cfg)?;
4093 }
4094 Ok(())
4095 }
4096
4097 pub fn gumbel_perturb_ctr(
4102 &self,
4103 x: &CudaSlice<f32>,
4104 y: &mut CudaSlice<f32>,
4105 n: usize,
4106 seed: u64,
4107 ctr: &CudaSlice<u32>,
4108 temp: f32,
4109 ) -> Result<(), Box<dyn std::error::Error>> {
4110 let f = self.func("gumbel_perturb_ctr_f32");
4111 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4112 let cfg = LaunchConfig {
4113 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4114 block_dim: (256, 1, 1),
4115 shared_mem_bytes: 0,
4116 };
4117 let __s_b = self.gpu.stream();
4118 let mut b = __s_b.launch_builder(&f);
4119 b.arg(x)
4120 .arg(&mut *y)
4121 .arg(&ni)
4122 .arg(&slo)
4123 .arg(&shi)
4124 .arg(ctr)
4125 .arg(&temp);
4126 unsafe {
4127 b.launch(cfg)?;
4128 }
4129 Ok(())
4130 }
4131
4132 #[allow(clippy::too_many_arguments)]
4140 pub fn gumbel_perturb_filtered_ctr(
4141 &self,
4142 x: &CudaSlice<f32>,
4143 y: &mut CudaSlice<f32>,
4144 n: usize,
4145 seed: u64,
4146 ctr: &CudaSlice<u32>,
4147 temp: f32,
4148 stat_max: &CudaSlice<f32>,
4149 stat_th: &CudaSlice<f32>,
4150 ) -> Result<(), Box<dyn std::error::Error>> {
4151 let f = self.func("gumbel_perturb_filtered_ctr_f32");
4152 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4153 let cfg = LaunchConfig {
4154 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4155 block_dim: (256, 1, 1),
4156 shared_mem_bytes: 0,
4157 };
4158 let __s_b = self.gpu.stream();
4159 let mut b = __s_b.launch_builder(&f);
4160 b.arg(x)
4161 .arg(&mut *y)
4162 .arg(&ni)
4163 .arg(&slo)
4164 .arg(&shi)
4165 .arg(ctr)
4166 .arg(&temp)
4167 .arg(stat_max)
4168 .arg(stat_th);
4169 unsafe {
4170 b.launch(cfg)?;
4171 }
4172 Ok(())
4173 }
4174
4175 pub fn softmax_gather(
4179 &self,
4180 x: &CudaSlice<f32>,
4181 row_stride: usize,
4182 ids: &CudaSlice<u32>,
4183 rows: &CudaSlice<i32>,
4184 out: &mut CudaSlice<f32>,
4185 n: usize,
4186 npair: usize,
4187 temp: f32,
4188 ) -> Result<(), Box<dyn std::error::Error>> {
4189 let f = self.func("softmax_gather_f32");
4190 let (ni, rs) = (n as i32, row_stride as i64);
4191 let np = npair as i32;
4192 let cfg = LaunchConfig {
4193 grid_dim: (npair as u32, 1, 1),
4194 block_dim: (256, 1, 1),
4195 shared_mem_bytes: 0,
4196 };
4197 let __s_b = self.gpu.stream();
4198 let mut b = __s_b.launch_builder(&f);
4199 b.arg(x)
4200 .arg(&rs)
4201 .arg(ids)
4202 .arg(rows)
4203 .arg(&mut *out)
4204 .arg(&ni)
4205 .arg(&np)
4206 .arg(&temp);
4207 unsafe {
4208 b.launch(cfg)?;
4209 }
4210 Ok(())
4211 }
4212
4213 pub fn residual_sample(
4217 &self,
4218 p: &CudaSlice<f32>,
4219 q: Option<&CudaSlice<f32>>,
4220 n: usize,
4221 temp: f32,
4222 seed: u64,
4223 stream_pos: u32,
4224 out_tok: &mut CudaSlice<u32>,
4225 ) -> Result<(), Box<dyn std::error::Error>> {
4226 let f = self.func("residual_sample_f32");
4227 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4228 let nth = 1024u32;
4229 let cfg = LaunchConfig {
4230 grid_dim: (1, 1, 1),
4231 block_dim: (nth, 1, 1),
4232 shared_mem_bytes: 0,
4233 };
4234 let has_q: i32 = q.is_some() as i32;
4235 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
4237 let mut b = __s_b.launch_builder(&f);
4238 b.arg(p)
4239 .arg(qbuf)
4240 .arg(&has_q)
4241 .arg(&ni)
4242 .arg(&temp)
4243 .arg(&slo)
4244 .arg(&shi)
4245 .arg(&stream_pos)
4246 .arg(&mut *out_tok);
4247 unsafe {
4248 b.launch(cfg)?;
4249 }
4250 Ok(())
4251 }
4252
4253 pub fn with_moe_cache<R>(
4258 &self,
4259 max_block_bytes: usize,
4260 f: impl FnOnce(
4261 &mut crate::moe_cache::MoeSlotCache,
4262 &Engine,
4263 ) -> Result<R, Box<dyn std::error::Error>>,
4264 ) -> Result<R, Box<dyn std::error::Error>> {
4265 let mut guard = self.moe_cache.lock().unwrap();
4266 if guard.is_none() {
4267 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
4268 }
4269 let cache = guard.as_mut().unwrap();
4270 f(cache, self)
4271 }
4272
4273 pub fn freeze_moe_cache(&self) {
4276 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
4277 cache.freeze();
4278 }
4279 }
4280
4281 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
4284 self.moe_cache
4285 .lock()
4286 .unwrap()
4287 .as_ref()
4288 .map(crate::moe_cache::MoeSlotCache::export_residency)
4289 }
4290
4291 pub(crate) fn moe_cache_frozen(&self) -> bool {
4292 self.moe_cache
4293 .lock()
4294 .unwrap()
4295 .as_ref()
4296 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
4297 }
4298
4299 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
4306 crate::cpu_experts::configured()
4307 && self.moe_cache_frozen()
4308 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
4309 }
4310
4311 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
4313 assert!(
4314 self.moe_cache.lock().unwrap().is_none(),
4315 "MoE cache layout configured after cache construction"
4316 );
4317 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
4318 }
4319
4320 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
4321 self.moe_cache_layout.lock().unwrap().clone()
4322 }
4323
4324 pub fn moe_cache_enabled() -> bool {
4326 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
4327 }
4328
4329 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
4332 let guard = self.moe_cache.lock().unwrap();
4333 guard
4334 .as_ref()
4335 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
4336 }
4337
4338 pub fn cpu_expert_stats(
4342 &self,
4343 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
4344 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
4345 }
4346
4347 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
4350 crate::cpu_experts::predictor_stats()
4351 }
4352
4353 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
4354 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
4355 }
4356
4357 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
4360 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
4361 }
4362
4363 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
4366 let guard = self.moe_cache.lock().unwrap();
4367 guard
4368 .as_ref()
4369 .and_then(|cache| cache.pread_stats())
4370 .map(|stats| {
4371 (
4372 stats.reads,
4373 stats.bytes,
4374 stats.read_errors,
4375 stats.short_reads,
4376 stats.fallbacks,
4377 stats.buffer_waits,
4378 stats.ring_full,
4379 )
4380 })
4381 }
4382
4383 pub fn spill_config_fallbacks(&self) -> u64 {
4385 crate::spill_pread::config_fallbacks()
4386 }
4387
4388 pub fn moe_cache_reset_counters(&self) {
4390 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
4391 c.reset_counters();
4392 }
4393 }
4394
4395 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4396 Ok(self.gpu.stream().clone_htod(v)?)
4397 }
4398
4399 pub fn htod_bytes_padded(
4403 &self,
4404 v: &[u8],
4405 pad: usize,
4406 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4407 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
4408 {
4409 let mut view = d.slice_mut(0..v.len());
4410 self.gpu.stream().memcpy_htod(v, &mut view)?;
4411 }
4412 Ok(d)
4413 }
4414
4415 pub fn copy_into(
4417 &self,
4418 dst: &mut CudaSlice<f32>,
4419 off: usize,
4420 src: &CudaSlice<f32>,
4421 len: usize,
4422 ) -> Result<(), Box<dyn std::error::Error>> {
4423 let mut view = dst.slice_mut(off..off + len);
4424 self.gpu
4425 .stream()
4426 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4427 Ok(())
4428 }
4429
4430 pub fn copy_range_into(
4434 &self,
4435 dst: &mut CudaSlice<f32>,
4436 dst_off: usize,
4437 src: &CudaSlice<f32>,
4438 src_off: usize,
4439 len: usize,
4440 ) -> Result<(), Box<dyn std::error::Error>> {
4441 let mut view = dst.slice_mut(dst_off..dst_off + len);
4442 self.gpu
4443 .stream()
4444 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut view)?;
4445 Ok(())
4446 }
4447
4448 pub fn copy_u8_into(
4451 &self,
4452 dst: &mut CudaSlice<u8>,
4453 off: usize,
4454 src: &CudaSlice<u8>,
4455 len: usize,
4456 ) -> Result<(), Box<dyn std::error::Error>> {
4457 let cap = dst.len();
4461 let mut view = dst.try_slice_mut(off..off + len).ok_or_else(|| {
4462 format!(
4463 "copy_u8_into dst range [{off},{}) exceeds capacity {cap}",
4464 off + len,
4465 )
4466 })?;
4467 self.gpu
4468 .stream()
4469 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4470 Ok(())
4471 }
4472
4473 pub fn copy_u8_range_into(
4475 &self,
4476 dst: &mut CudaSlice<u8>,
4477 dst_off: usize,
4478 src: &CudaSlice<u8>,
4479 src_off: usize,
4480 len: usize,
4481 ) -> Result<(), Box<dyn std::error::Error>> {
4482 let cap = dst.len();
4485 let mut dst_view = dst.try_slice_mut(dst_off..dst_off + len).ok_or_else(|| {
4486 format!(
4487 "copy_u8_range_into dst range [{dst_off},{}) exceeds capacity {cap}",
4488 dst_off + len,
4489 )
4490 })?;
4491 self.gpu
4492 .stream()
4493 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
4494 Ok(())
4495 }
4496
4497 #[track_caller]
4507 pub fn prepare_kv_append(
4508 &self,
4509 kv: &mut crate::cache::KvLayer,
4510 retain_from: usize,
4511 append_rows: usize,
4512 ) -> Result<usize, Box<dyn std::error::Error>> {
4513 let caller = std::panic::Location::caller();
4514 let base_before = kv.ring.as_ref().map(|r| r.base());
4515 let Some(plan) = kv
4516 .ring
4517 .as_ref()
4518 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4519 .transpose()
4520 .map_err(|err| -> Box<dyn std::error::Error> {
4521 format!(
4522 "{err} [append len={} retain_from={retain_from} append_rows={append_rows} base={base_before:?} called from {caller}]",
4523 kv.len
4524 )
4525 .into()
4526 })?
4527 else {
4528 return Ok(kv.len);
4529 };
4530 match plan {
4531 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4532 crate::cache::KvRingAppend::Rebase {
4533 src_row,
4534 keep_rows,
4535 new_base,
4536 write_row,
4537 } => {
4538 if keep_rows > 0 {
4539 let k_len = keep_rows * kv.k_tok_bytes;
4540 let v_len = keep_rows * kv.v_tok_bytes;
4541 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4542 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4543 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4544 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4545 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4546 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4547 }
4548 if std::env::var("MEMRA_KV_REBASE_TRACE").as_deref() == Ok("1") {
4551 eprintln!(
4552 "[kv-rebase] new_base={new_base} keep_rows={keep_rows} len={} \
4553 retain_from={retain_from} called from {caller}",
4554 kv.len
4555 );
4556 }
4557 kv.ring.as_mut().unwrap().apply_rebase(new_base);
4558 if let Some(base_d) = kv.base_d.as_mut() {
4562 self.set_i32_one(base_d, new_base as i32)?;
4563 }
4564 Ok(write_row)
4565 }
4566 }
4567 }
4568
4569 pub fn htod_u8_into(
4572 &self,
4573 dst: &mut CudaSlice<u8>,
4574 off: usize,
4575 src: &[u8],
4576 ) -> Result<(), Box<dyn std::error::Error>> {
4577 let mut view = dst.slice_mut(off..off + src.len());
4578 self.gpu.stream().memcpy_htod(src, &mut view)?;
4579 Ok(())
4580 }
4581
4582 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4583 b.slice(0..len)
4584 }
4585
4586 pub fn view_u8_range<'a>(
4589 &self,
4590 b: &'a CudaSlice<u8>,
4591 start: usize,
4592 end: usize,
4593 ) -> cudarc::driver::CudaView<'a, u8> {
4594 b.slice(start..end)
4595 }
4596 pub fn view_u8<'a>(
4597 &self,
4598 b: &'a CudaSlice<u8>,
4599 len: usize,
4600 ) -> cudarc::driver::CudaView<'a, u8> {
4601 b.slice(0..len)
4602 }
4603
4604 pub fn append_kv_quantized(
4608 &self,
4609 k_row: &CudaSlice<f32>,
4610 v_row: &CudaSlice<f32>,
4611 kc: &mut CudaSlice<u8>,
4612 vc: &mut CudaSlice<u8>,
4613 t: usize,
4614 kv_dim_k: usize,
4615 kv_dim_v: usize,
4616 k_tok_bytes: usize,
4617 v_tok_bytes: usize,
4618 g: bool,
4619 ) -> Result<(), Box<dyn std::error::Error>> {
4620 let f = if g {
4621 self.func_g("append_quantize_kv_q8_0_q5_1")
4622 } else {
4623 self.func("append_quantize_kv_q8_0_q5_1")
4624 };
4625 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4626 let cfg = LaunchConfig {
4627 grid_dim: (nblk, 1, 1),
4628 block_dim: (32, 1, 1),
4629 shared_mem_bytes: 0,
4630 };
4631 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4632 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4633 let __s_b = self.gpu.stream();
4634 let mut b = __s_b.launch_builder(&f);
4635 b.arg(k_row)
4636 .arg(v_row)
4637 .arg(kc)
4638 .arg(vc)
4639 .arg(&ti)
4640 .arg(&kdk)
4641 .arg(&kdv)
4642 .arg(&ktb)
4643 .arg(&vtb);
4644 unsafe {
4645 b.launch(cfg)?;
4646 }
4647 Ok(())
4648 }
4649
4650 pub fn append_kv_quantized_dc(
4654 &self,
4655 k_row: &CudaSlice<f32>,
4656 v_row: &CudaSlice<f32>,
4657 kc: &mut CudaSlice<u8>,
4658 vc: &mut CudaSlice<u8>,
4659 t_dev: &CudaSlice<i32>,
4660 kv_dim_k: usize,
4661 kv_dim_v: usize,
4662 k_tok_bytes: usize,
4663 v_tok_bytes: usize,
4664 g: bool,
4665 ) -> Result<(), Box<dyn std::error::Error>> {
4666 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4667 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4668 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4669 if Self::pdl_on() && Self::pdl_wb_on() {
4671 use cudarc::driver::{DevicePtr, DevicePtrMut};
4672 let s = &self.gpu.stream();
4673 let (pk, _g0) = k_row.device_ptr(s);
4674 let (pv, _g1) = v_row.device_ptr(s);
4675 let (pkc, _g2) = kc.device_ptr_mut(s);
4676 let (pvc, _g3) = vc.device_ptr_mut(s);
4677 let (pt, _g4) = t_dev.device_ptr(s);
4678 let mut ps = [
4679 &pk as *const _ as *mut std::ffi::c_void,
4680 &pv as *const _ as *mut _,
4681 &pkc as *const _ as *mut _,
4682 &pvc as *const _ as *mut _,
4683 &pt as *const _ as *mut _,
4684 &kdk as *const _ as *mut _,
4685 &kdv as *const _ as *mut _,
4686 &ktb as *const _ as *mut _,
4687 &vtb as *const _ as *mut _,
4688 ];
4689 unsafe {
4690 self.launch_pdl_flash(
4691 g,
4692 "append_quantize_kv_q8_0_q5_1_dc",
4693 (nblk, 1, 1),
4694 (32, 1, 1),
4695 0,
4696 &mut ps,
4697 )?;
4698 }
4699 return Ok(());
4700 }
4701 let f = if g {
4702 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4703 } else {
4704 self.func("append_quantize_kv_q8_0_q5_1_dc")
4705 };
4706 let cfg = LaunchConfig {
4707 grid_dim: (nblk, 1, 1),
4708 block_dim: (32, 1, 1),
4709 shared_mem_bytes: 0,
4710 };
4711 let __s_b = self.gpu.stream();
4712 let mut b = __s_b.launch_builder(&f);
4713 b.arg(k_row)
4714 .arg(v_row)
4715 .arg(kc)
4716 .arg(vc)
4717 .arg(t_dev)
4718 .arg(&kdk)
4719 .arg(&kdv)
4720 .arg(&ktb)
4721 .arg(&vtb);
4722 unsafe {
4723 b.launch(cfg)?;
4724 }
4725 Ok(())
4726 }
4727
4728 #[allow(clippy::too_many_arguments)]
4735 pub fn append_kv_quantized_rows(
4736 &self,
4737 k_rows: &CudaSlice<f32>,
4738 v_rows: &CudaSlice<f32>,
4739 kc: &mut CudaSlice<u8>,
4740 vc: &mut CudaSlice<u8>,
4741 t0: usize,
4742 t: usize,
4743 kv_dim_k: usize,
4744 kv_dim_v: usize,
4745 k_tok_bytes: usize,
4746 v_tok_bytes: usize,
4747 g: bool,
4748 ) -> Result<(), Box<dyn std::error::Error>> {
4749 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4750 for i in 0..t {
4751 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4752 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4753 self.append_kv_quantized_view(
4754 &k_row,
4755 &v_row,
4756 kc,
4757 vc,
4758 t0 + i,
4759 kv_dim_k,
4760 kv_dim_v,
4761 k_tok_bytes,
4762 v_tok_bytes,
4763 g,
4764 )?;
4765 }
4766 return Ok(());
4767 }
4768 let f = if g {
4769 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4770 } else {
4771 self.func("append_quantize_kv_q8_0_q5_1_rows")
4772 };
4773 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4774 let cfg = LaunchConfig {
4775 grid_dim: (nblk, t as u32, 1),
4776 block_dim: (32, 1, 1),
4777 shared_mem_bytes: 0,
4778 };
4779 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4780 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4781 let __s_b = self.gpu.stream();
4782 let mut b = __s_b.launch_builder(&f);
4783 b.arg(k_rows)
4784 .arg(v_rows)
4785 .arg(kc)
4786 .arg(vc)
4787 .arg(&t0i)
4788 .arg(&kdk)
4789 .arg(&kdv)
4790 .arg(&ktb)
4791 .arg(&vtb);
4792 unsafe {
4793 b.launch(cfg)?;
4794 }
4795 Ok(())
4796 }
4797
4798 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4802 let f = self.func("inc_i32");
4803 let cfg = LaunchConfig {
4804 grid_dim: (1, 1, 1),
4805 block_dim: (1, 1, 1),
4806 shared_mem_bytes: 0,
4807 };
4808 let __s_b = self.gpu.stream();
4809 let mut b = __s_b.launch_builder(&f);
4810 b.arg(p);
4811 unsafe {
4812 b.launch(cfg)?;
4813 }
4814 Ok(())
4815 }
4816
4817 pub fn append_kv_quantized_view(
4820 &self,
4821 k_row: &cudarc::driver::CudaView<f32>,
4822 v_row: &cudarc::driver::CudaView<f32>,
4823 kc: &mut CudaSlice<u8>,
4824 vc: &mut CudaSlice<u8>,
4825 t: usize,
4826 kv_dim_k: usize,
4827 kv_dim_v: usize,
4828 k_tok_bytes: usize,
4829 v_tok_bytes: usize,
4830 g: bool,
4831 ) -> Result<(), Box<dyn std::error::Error>> {
4832 let stream = self.gpu.stream();
4833 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4834 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4835 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4836 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4837 let f = if g {
4838 self.func_g("append_quantize_kv_q8_0_q5_1")
4839 } else {
4840 self.func("append_quantize_kv_q8_0_q5_1")
4841 };
4842 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4843 let cfg = LaunchConfig {
4844 grid_dim: (nblk, 1, 1),
4845 block_dim: (32, 1, 1),
4846 shared_mem_bytes: 0,
4847 };
4848 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4849 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4850 let mut b = stream.launch_builder(&f);
4851 b.arg(k_row)
4852 .arg(v_row)
4853 .arg(kc)
4854 .arg(vc)
4855 .arg(&ti)
4856 .arg(&kdk)
4857 .arg(&kdv)
4858 .arg(&ktb)
4859 .arg(&vtb);
4860 unsafe {
4861 b.launch(cfg)?;
4862 }
4863 Ok(())
4864 }
4865
4866 pub fn copy_view_into(
4869 &self,
4870 dst: &mut CudaSlice<f32>,
4871 off: usize,
4872 src: &cudarc::driver::CudaView<f32>,
4873 len: usize,
4874 ) -> Result<(), Box<dyn std::error::Error>> {
4875 let mut view = dst.slice_mut(off..off + len);
4876 self.gpu
4877 .stream()
4878 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4879 Ok(())
4880 }
4881
4882 pub fn clone_dtod(
4886 &self,
4887 src: &CudaSlice<f32>,
4888 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4889 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4890 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4891 Ok(dst)
4892 }
4893
4894 pub fn dtod_copy_view(
4897 &self,
4898 src: &cudarc::driver::CudaView<f32>,
4899 dst: &mut CudaSlice<f32>,
4900 ) -> Result<(), Box<dyn std::error::Error>> {
4901 self.gpu.stream().memcpy_dtod(src, dst)?;
4902 Ok(())
4903 }
4904
4905 pub fn dtod_copy_view_i8(
4907 &self,
4908 src: &cudarc::driver::CudaView<i8>,
4909 dst: &mut CudaSlice<i8>,
4910 ) -> Result<(), Box<dyn std::error::Error>> {
4911 self.gpu.stream().memcpy_dtod(src, dst)?;
4912 Ok(())
4913 }
4914
4915 pub fn dtod_copy_into(
4917 &self,
4918 src: &CudaSlice<f32>,
4919 dst: &mut CudaSlice<f32>,
4920 offset: usize,
4921 ) -> Result<(), Box<dyn std::error::Error>> {
4922 let n = src.len();
4923 let mut dv = dst.slice_mut(offset..offset + n);
4924 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4925 Ok(())
4926 }
4927
4928 pub fn copy_batch_uniform_f32(
4934 &self,
4935 table: &CudaSlice<u64>,
4936 n: usize,
4937 words: usize,
4938 ) -> Result<(), Box<dyn std::error::Error>> {
4939 if n == 0 || words == 0 {
4940 return Ok(());
4941 }
4942 debug_assert!(
4943 table.len() >= 2 * n,
4944 "pointer table must hold n srcs + n dsts"
4945 );
4946 let f = self.func("copy_batch_uniform_f32");
4947 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4950 let (ni, wi) = (n as i32, words as i32);
4951 let cfg = LaunchConfig {
4952 grid_dim: (chunks, n as u32, 1),
4953 block_dim: (256, 1, 1),
4954 shared_mem_bytes: 0,
4955 };
4956 let __s = self.gpu.stream();
4957 let mut b = __s.launch_builder(&f);
4958 b.arg(table).arg(&ni).arg(&wi);
4959 unsafe {
4960 b.launch(cfg)?;
4961 }
4962 Ok(())
4963 }
4964
4965 pub fn htod_u64_into(
4968 &self,
4969 v: &[u64],
4970 dst: &mut CudaSlice<u64>,
4971 ) -> Result<(), Box<dyn std::error::Error>> {
4972 let mut view = dst.slice_mut(0..v.len());
4973 self.gpu.stream().memcpy_htod(v, &mut view)?;
4974 Ok(())
4975 }
4976
4977 pub fn copy_indirect_src_f32(
4982 &self,
4983 src_entry: &cudarc::driver::CudaView<u64>,
4984 dst: &mut CudaSlice<f32>,
4985 dst_off: usize,
4986 words: usize,
4987 ) -> Result<(), Box<dyn std::error::Error>> {
4988 let f = self.func("copy_indirect_src_f32");
4989 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4990 let wi = words as i32;
4991 let cfg = LaunchConfig {
4992 grid_dim: (chunks, 1, 1),
4993 block_dim: (256, 1, 1),
4994 shared_mem_bytes: 0,
4995 };
4996 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4997 let __s = self.gpu.stream();
4998 let mut b = __s.launch_builder(&f);
4999 b.arg(src_entry).arg(&mut dv).arg(&wi);
5000 unsafe {
5001 b.launch(cfg)?;
5002 }
5003 Ok(())
5004 }
5005
5006 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
5008 self.alloc_uninit::<i8>(n)
5009 }
5010
5011 pub fn qmatvec(
5013 &self,
5014 w: &CudaSlice<u8>,
5015 x: &CudaSlice<f32>,
5016 m: usize,
5017 in_f: usize,
5018 out_f: usize,
5019 qtype: i32,
5020 row_bytes: usize,
5021 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5022 let f = self.func("qmatvec_f32");
5023 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
5025 grid_dim: (out_f as u32, m as u32, 1),
5026 block_dim: (256, 1, 1),
5027 shared_mem_bytes: 0,
5028 };
5029 let (inf, outf, mi, qt, rb) =
5030 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
5031 let __s_b = self.gpu.stream();
5032 let mut b = __s_b.launch_builder(&f);
5033 b.arg(w)
5034 .arg(x)
5035 .arg(&mut y)
5036 .arg(&inf)
5037 .arg(&outf)
5038 .arg(&mi)
5039 .arg(&qt)
5040 .arg(&rb);
5041 unsafe {
5042 b.launch(cfg)?;
5043 }
5044 Ok(y)
5045 }
5046
5047 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5049 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
5050 self.keep_if_capturing(&s);
5051 Ok(s)
5052 }
5053
5054 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5058 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
5059 self.keep_if_capturing(&s);
5060 Ok(s)
5061 }
5062
5063 pub fn memset_zeros_view(
5066 &self,
5067 dst: &mut cudarc::driver::CudaViewMut<f32>,
5068 ) -> Result<(), Box<dyn std::error::Error>> {
5069 self.gpu.stream().memset_zeros(dst)?;
5070 Ok(())
5071 }
5072
5073 pub fn stage_expert(
5079 &self,
5080 host_bytes: &[u8],
5081 scratch: &mut CudaSlice<u8>,
5082 off: usize,
5083 ) -> Result<(), Box<dyn std::error::Error>> {
5084 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
5087 }
5088
5089 pub fn moe_router_topk(
5095 &self,
5096 logits: &CudaSlice<f32>,
5097 t: usize,
5098 n_expert: usize,
5099 n_used: usize,
5100 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5101 let f = self.func("moe_router_topk_f32");
5102 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 {
5105 grid_dim: (t as u32, 1, 1),
5106 block_dim: (n_expert as u32, 1, 1),
5107 shared_mem_bytes: 0,
5108 };
5109 let (ne, nu) = (n_expert as i32, n_used as i32);
5110 let __s_b = self.gpu.stream();
5111 let mut b = __s_b.launch_builder(&f);
5112 b.arg(logits)
5113 .arg(&mut sel_idx)
5114 .arg(&mut sel_w)
5115 .arg(&ne)
5116 .arg(&nu);
5117 unsafe {
5118 b.launch(cfg)?;
5119 }
5120 Ok((sel_idx, sel_w))
5121 }
5122
5123 pub fn moe_router_topk_scaled(
5126 &self,
5127 logits: &CudaSlice<f32>,
5128 t: usize,
5129 n_expert: usize,
5130 n_used: usize,
5131 ex_scale: &CudaSlice<f32>,
5132 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5133 let f = self.func("moe_router_topk_scaled_f32");
5138 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
5139 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
5140 let cfg = LaunchConfig {
5141 grid_dim: (t as u32, 1, 1),
5142 block_dim: (n_expert as u32, 1, 1),
5143 shared_mem_bytes: 0,
5144 };
5145 let (ne, nu) = (n_expert as i32, n_used as i32);
5146 let __s_b = self.gpu.stream();
5147 let mut b = __s_b.launch_builder(&f);
5148 b.arg(logits)
5149 .arg(&mut sel_idx)
5150 .arg(&mut sel_w)
5151 .arg(&ne)
5152 .arg(&nu)
5153 .arg(ex_scale);
5154 unsafe {
5155 b.launch(cfg)?;
5156 }
5157 Ok((sel_idx, sel_w))
5158 }
5159
5160 pub fn moe_router_topk_host(
5168 &self,
5169 logits: &CudaSlice<f32>,
5170 t: usize,
5171 n_expert: usize,
5172 n_used: usize,
5173 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5174 let f = self.func("moe_router_topk_f32");
5175 let n = t * n_used;
5176 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
5177 let mut sel_w = self.alloc_uninit::<f32>(n)?;
5178 let cfg = LaunchConfig {
5179 grid_dim: (t as u32, 1, 1),
5180 block_dim: (n_expert as u32, 1, 1),
5181 shared_mem_bytes: 0,
5182 };
5183 let (ne, nu) = (n_expert as i32, n_used as i32);
5184 let __s_b = self.gpu.stream();
5185 let mut b = __s_b.launch_builder(&f);
5186 b.arg(logits)
5187 .arg(&mut sel_idx)
5188 .arg(&mut sel_w)
5189 .arg(&ne)
5190 .arg(&nu);
5191 unsafe {
5192 b.launch(cfg)?;
5193 }
5194 let bytes = n * 8;
5196 let mut guard = self.router_stage.lock().unwrap();
5197 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
5198 *guard = Some(PinnedStage::new(bytes.max(4096))?);
5199 }
5200 let stage = guard.as_mut().unwrap();
5201 let (si, sw) = unsafe {
5202 (
5203 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
5204 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
5205 )
5206 };
5207 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()))
5211 }
5212
5213 #[allow(clippy::too_many_arguments)]
5217 pub fn moe_router_sigmoid_topk(
5218 &self,
5219 logits: &CudaSlice<f32>,
5220 t: usize,
5221 n_expert: usize,
5222 n_used: usize,
5223 active_count: usize,
5224 correction_bias: &CudaSlice<f32>,
5225 active: &CudaSlice<u8>,
5226 scaling_factor: f32,
5227 route_norm: bool,
5228 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5229 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5230 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
5231 return Err(format!(
5232 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
5233 )
5234 .into());
5235 }
5236 if logits.len() < t * n_expert
5237 || correction_bias.len() != n_expert
5238 || active.len() != n_expert
5239 {
5240 return Err(format!(
5241 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
5242 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
5243 ).into());
5244 }
5245 let f = self.func(crate::sigmoid_topk_kernel(
5246 crate::sig_expf_dev_on(),
5247 crate::topk_fast_on(),
5248 n_used,
5249 ));
5250 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
5251 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
5252 let threads = n_expert.div_ceil(32) * 32;
5253 let cfg = LaunchConfig {
5254 grid_dim: (t as u32, 1, 1),
5255 block_dim: (threads as u32, 1, 1),
5256 shared_mem_bytes: 0,
5257 };
5258 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
5259 let __s_b = self.gpu.stream();
5260 let mut b = __s_b.launch_builder(&f);
5261 b.arg(logits)
5262 .arg(correction_bias)
5263 .arg(active)
5264 .arg(&mut sel_idx)
5265 .arg(&mut sel_w)
5266 .arg(&ne)
5267 .arg(&nu)
5268 .arg(&scaling_factor)
5269 .arg(&rn);
5270 unsafe {
5271 b.launch(cfg)?;
5272 }
5273 Ok((sel_idx, sel_w))
5274 }
5275
5276 #[allow(clippy::too_many_arguments)]
5279 pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
5283 if ptr == 0 {
5284 return Err("ring_flag_raw: unarmed flag".into());
5285 }
5286 let f = self.func("memra_ring_flag");
5287 let cfg = LaunchConfig {
5288 grid_dim: (1, 1, 1),
5289 block_dim: (32, 1, 1),
5290 shared_mem_bytes: 0,
5291 };
5292 let __s_b = self.gpu.stream();
5293 let mut b = __s_b.launch_builder(&f);
5294 b.arg(&ptr).arg(&value);
5295 unsafe {
5296 b.launch(cfg)?;
5297 }
5298 Ok(())
5299 }
5300
5301 pub fn moe_sel_w_mirror(
5304 &self,
5305 sel_src: &CudaSlice<i32>,
5306 w_src: &CudaSlice<f32>,
5307 sel_dst: &mut CudaSlice<i32>,
5308 w_dst: &mut CudaSlice<f32>,
5309 n: usize,
5310 ) -> Result<(), Box<dyn std::error::Error>> {
5311 if n == 0
5312 || n > 32
5313 || sel_src.len() < n
5314 || w_src.len() < n
5315 || sel_dst.len() < n
5316 || w_dst.len() < n
5317 {
5318 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
5319 }
5320 let f = self.func("moe_sel_w_mirror");
5321 let cfg = LaunchConfig {
5322 grid_dim: (1, 1, 1),
5323 block_dim: (32, 1, 1),
5324 shared_mem_bytes: 0,
5325 };
5326 let ni = n as i32;
5327 let __s_b = self.gpu.stream();
5328 let mut b = __s_b.launch_builder(&f);
5329 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
5330 unsafe {
5331 b.launch(cfg)?;
5332 }
5333 Ok(())
5334 }
5335
5336 pub fn moe_router_sigmoid_topk_into(
5337 &self,
5338 logits: &CudaSlice<f32>,
5339 t: usize,
5340 n_expert: usize,
5341 n_used: usize,
5342 active_count: usize,
5343 correction_bias: &CudaSlice<f32>,
5344 active: &CudaSlice<u8>,
5345 scaling_factor: f32,
5346 route_norm: bool,
5347 sel_idx: &mut CudaSlice<i32>,
5348 sel_w: &mut CudaSlice<f32>,
5349 ) -> Result<(), Box<dyn std::error::Error>> {
5350 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5351 if n_expert == 0
5352 || n_expert > 1024
5353 || n_used == 0
5354 || n_used > 32 || n_used > n_expert
5356 || logits.len() < t * n_expert
5357 || correction_bias.len() != n_expert
5358 || active.len() != n_expert
5359 || sel_idx.len() < t * n_used
5360 || sel_w.len() < t * n_used
5361 {
5362 return Err("sigmoid router _into geometry mismatch".into());
5363 }
5364 let f = self.func(crate::sigmoid_topk_kernel(
5365 crate::sig_expf_dev_on(),
5366 crate::topk_fast_on(),
5367 n_used,
5368 ));
5369 let threads = n_expert.div_ceil(32) * 32;
5370 let cfg = LaunchConfig {
5371 grid_dim: (t as u32, 1, 1),
5372 block_dim: (threads as u32, 1, 1),
5373 shared_mem_bytes: 0,
5374 };
5375 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
5376 let __s_b = self.gpu.stream();
5377 let mut b = __s_b.launch_builder(&f);
5378 b.arg(logits)
5379 .arg(correction_bias)
5380 .arg(active)
5381 .arg(&mut *sel_idx)
5382 .arg(&mut *sel_w)
5383 .arg(&ne)
5384 .arg(&nu)
5385 .arg(&scaling_factor)
5386 .arg(&rn);
5387 unsafe {
5388 b.launch(cfg)?;
5389 }
5390 Ok(())
5391 }
5392
5393 #[allow(clippy::too_many_arguments)]
5396 pub fn moe_router_sigmoid_topk_host(
5397 &self,
5398 logits: &CudaSlice<f32>,
5399 t: usize,
5400 n_expert: usize,
5401 n_used: usize,
5402 active_count: usize,
5403 correction_bias: &CudaSlice<f32>,
5404 active: &CudaSlice<u8>,
5405 scaling_factor: f32,
5406 route_norm: bool,
5407 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5408 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
5409 logits,
5410 t,
5411 n_expert,
5412 n_used,
5413 active_count,
5414 correction_bias,
5415 active,
5416 scaling_factor,
5417 route_norm,
5418 )?;
5419 let n = t * n_used;
5420 let bytes = n * 8;
5421 let mut guard = self.router_stage.lock().unwrap();
5422 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
5423 *guard = Some(PinnedStage::new(bytes.max(4096))?);
5424 }
5425 let stage = guard.as_mut().unwrap();
5426 let (si, sw) = unsafe {
5427 (
5428 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
5429 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
5430 )
5431 };
5432 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
5433 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
5434 self.gpu.stream().synchronize()?;
5435 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
5436 }
5437
5438 pub fn stage_expert_async(
5442 &self,
5443 host_bytes: &[u8],
5444 scratch: &mut CudaSlice<u8>,
5445 off: usize,
5446 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
5447 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
5448 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
5449 Ok(self.copy_stream.record_event(None)?)
5450 }
5451
5452 pub fn compute_wait(
5454 &self,
5455 ev: &cudarc::driver::CudaEvent,
5456 ) -> Result<(), Box<dyn std::error::Error>> {
5457 self.gpu.stream().wait(ev)?;
5458 Ok(())
5459 }
5460
5461 pub fn qmatvec_view(
5466 &self,
5467 w: &CudaSlice<u8>,
5468 range: std::ops::Range<usize>,
5469 x: &cudarc::driver::CudaView<f32>,
5470 m: usize,
5471 in_f: usize,
5472 out_f: usize,
5473 qtype: i32,
5474 row_bytes: usize,
5475 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5476 let f = self.func("qmatvec_f32");
5477 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
5480 grid_dim: (out_f as u32, m as u32, 1),
5481 block_dim: (256, 1, 1),
5482 shared_mem_bytes: 0,
5483 };
5484 let (inf, outf, mi, qt, rb) =
5485 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
5486 let __s_b = self.gpu.stream();
5487 let mut b = __s_b.launch_builder(&f);
5488 b.arg(&wv)
5489 .arg(x)
5490 .arg(&mut y)
5491 .arg(&inf)
5492 .arg(&outf)
5493 .arg(&mi)
5494 .arg(&qt)
5495 .arg(&rb);
5496 unsafe {
5497 b.launch(cfg)?;
5498 }
5499 Ok(y)
5500 }
5501
5502 #[allow(clippy::too_many_arguments)]
5509 pub fn moe_gate_up_silu8_q8(
5513 &self,
5514 gp: WPtr8,
5515 up: WPtr8,
5516 aq: &CudaSlice<i8>,
5517 ad: &CudaSlice<f32>,
5518 in_f: usize,
5519 n_ff: usize,
5520 n_used: usize,
5521 qt_g: i32,
5522 qt_u: i32,
5523 rb_g: usize,
5524 rb_u: usize,
5525 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5526 let f = self.func("moe_gate_up_silu8_q8");
5527 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5528 let cfg = LaunchConfig {
5529 grid_dim: (n_ff as u32, n_used as u32, 1),
5530 block_dim: (32, 1, 1),
5531 shared_mem_bytes: 0,
5532 };
5533 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5534 let __s_b = self.gpu.stream();
5535 let mut b = __s_b.launch_builder(&f);
5536 b.arg(&gp)
5537 .arg(&up)
5538 .arg(aq)
5539 .arg(ad)
5540 .arg(&mut act)
5541 .arg(&inf)
5542 .arg(&nff)
5543 .arg(&qt_g)
5544 .arg(&qt_u)
5545 .arg(&rbg)
5546 .arg(&rbu);
5547 unsafe {
5548 b.launch(cfg)?;
5549 }
5550 Ok(act)
5551 }
5552
5553 #[allow(clippy::too_many_arguments)]
5554 pub fn moe_down8_fma_q8(
5555 &self,
5556 dp: WPtr8,
5557 w: F32x8,
5558 aq2: &CudaSlice<i8>,
5559 ad2: &CudaSlice<f32>,
5560 dst: &mut cudarc::driver::CudaViewMut<f32>,
5561 in_f: usize,
5562 out_f: usize,
5563 n_used: usize,
5564 qt: i32,
5565 rb: usize,
5566 ) -> Result<(), Box<dyn std::error::Error>> {
5567 let f = self.func("moe_down8_fma_q8");
5568 let cfg = LaunchConfig {
5569 grid_dim: (out_f as u32, 1, 1),
5570 block_dim: (32, 1, 1),
5571 shared_mem_bytes: 0,
5572 };
5573 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5574 let __s_b = self.gpu.stream();
5575 let mut b = __s_b.launch_builder(&f);
5576 b.arg(&dp)
5577 .arg(&w)
5578 .arg(aq2)
5579 .arg(ad2)
5580 .arg(dst)
5581 .arg(&inf)
5582 .arg(&outf)
5583 .arg(&nu)
5584 .arg(&qt)
5585 .arg(&rbi);
5586 unsafe {
5587 b.launch(cfg)?;
5588 }
5589 Ok(())
5590 }
5591
5592 pub fn qmatvec_expert_q8(
5594 &self,
5595 w: &CudaSlice<u8>,
5596 range: std::ops::Range<usize>,
5597 aq: &CudaSlice<i8>,
5598 ad: &CudaSlice<f32>,
5599 m: usize,
5600 in_f: usize,
5601 out_f: usize,
5602 qtype: i32,
5603 row_bytes: usize,
5604 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5605 let f = self.func("qmatvec_expert_q8");
5606 let wv = w.slice(range);
5607 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5608 const ROWS: u32 = 4; let cfg = LaunchConfig {
5610 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5611 block_dim: (32, ROWS, 1),
5612 shared_mem_bytes: 0,
5613 };
5614 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5615 let __s_b = self.gpu.stream();
5616 let mut b = __s_b.launch_builder(&f);
5617 b.arg(&wv)
5618 .arg(aq)
5619 .arg(ad)
5620 .arg(&mut y)
5621 .arg(&inf)
5622 .arg(&outf)
5623 .arg(&mi)
5624 .arg(&qtype)
5625 .arg(&rbi);
5626 unsafe {
5627 b.launch(cfg)?;
5628 }
5629 Ok(y)
5630 }
5631
5632 pub fn moe_gate_up_silu8(
5633 &self,
5634 gp: WPtr8,
5635 up: WPtr8,
5636 x: &cudarc::driver::CudaView<f32>,
5637 in_f: usize,
5638 n_ff: usize,
5639 n_used: usize,
5640 qt_g: i32,
5641 qt_u: i32,
5642 rb_g: usize,
5643 rb_u: usize,
5644 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5645 let f = self.func("moe_gate_up_silu8_f32");
5646 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
5648 grid_dim: (n_ff as u32, n_used as u32, 1),
5649 block_dim: (256, 1, 1),
5650 shared_mem_bytes: 0,
5651 };
5652 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5653 let __s_b = self.gpu.stream();
5654 let mut b = __s_b.launch_builder(&f);
5655 b.arg(&gp)
5656 .arg(&up)
5657 .arg(x)
5658 .arg(&mut act)
5659 .arg(&inf)
5660 .arg(&nff)
5661 .arg(&qt_g)
5662 .arg(&qt_u)
5663 .arg(&rbg)
5664 .arg(&rbu);
5665 unsafe {
5666 b.launch(cfg)?;
5667 }
5668 Ok(act)
5669 }
5670
5671 #[allow(clippy::too_many_arguments)]
5677 pub fn moe_down8_fma_into(
5678 &self,
5679 dp: WPtr8,
5680 w: F32x8,
5681 act: &CudaSlice<f32>,
5682 dst: &mut cudarc::driver::CudaViewMut<f32>,
5683 in_f: usize,
5684 out_f: usize,
5685 n_used: usize,
5686 qt: i32,
5687 rb: usize,
5688 ) -> Result<(), Box<dyn std::error::Error>> {
5689 let f = self.func("moe_down8_fma_f32");
5690 let cfg = LaunchConfig {
5691 grid_dim: (out_f as u32, 1, 1),
5692 block_dim: (256, 1, 1),
5693 shared_mem_bytes: 0,
5694 };
5695 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5696 let __s_b = self.gpu.stream();
5697 let mut b = __s_b.launch_builder(&f);
5698 b.arg(&dp)
5699 .arg(&w)
5700 .arg(act)
5701 .arg(dst)
5702 .arg(&inf)
5703 .arg(&outf)
5704 .arg(&nu)
5705 .arg(&qt)
5706 .arg(&rbv);
5707 unsafe {
5708 b.launch(cfg)?;
5709 }
5710 Ok(())
5711 }
5712
5713 #[allow(clippy::too_many_arguments)]
5718 #[allow(clippy::too_many_arguments)]
5733 #[allow(clippy::too_many_arguments)]
5735 pub fn moe_pairs_matvec_q8(
5736 &self,
5737 table: &CudaSlice<u64>,
5738 proj: i32,
5739 pair_tok: &CudaSlice<i32>,
5740 pair_ex: &CudaSlice<i32>,
5741 aq: &CudaSlice<i8>,
5742 ad: &CudaSlice<f32>,
5743 in_f: usize,
5744 out_f: usize,
5745 n_expert: usize,
5746 n_pairs: usize,
5747 qtype: i32,
5748 row_bytes: usize,
5749 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5750 let f = self.func("moe_pairs_matvec_q8");
5751 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5752 const ROWS: u32 = 4;
5753 let cfg = LaunchConfig {
5754 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5755 block_dim: (32, ROWS, 1),
5756 shared_mem_bytes: 0,
5757 };
5758 let (inf, outf, ne, np, rbi) = (
5759 in_f as i32,
5760 out_f as i32,
5761 n_expert as i32,
5762 n_pairs as i32,
5763 row_bytes as i64,
5764 );
5765 let __s_b = self.gpu.stream();
5766 let mut b = __s_b.launch_builder(&f);
5767 b.arg(table)
5768 .arg(&proj)
5769 .arg(pair_tok)
5770 .arg(pair_ex)
5771 .arg(aq)
5772 .arg(ad)
5773 .arg(&mut y)
5774 .arg(&inf)
5775 .arg(&outf)
5776 .arg(&ne)
5777 .arg(&np)
5778 .arg(&qtype)
5779 .arg(&rbi);
5780 unsafe {
5781 b.launch(cfg)?;
5782 }
5783 Ok(y)
5784 }
5785
5786 #[allow(clippy::too_many_arguments)]
5788 pub fn moe_pairs_matvec_q8_em(
5789 &self,
5790 table: &CudaSlice<u64>,
5791 proj: i32,
5792 ex_ids: &CudaSlice<i32>,
5793 ex_off: &CudaSlice<i32>,
5794 ex_pairs: &CudaSlice<i32>,
5795 pair_tok: &CudaSlice<i32>,
5796 aq: &CudaSlice<i8>,
5797 ad: &CudaSlice<f32>,
5798 in_f: usize,
5799 out_f: usize,
5800 n_expert: usize,
5801 n_active: usize,
5802 n_pairs: usize,
5803 qtype: i32,
5804 row_bytes: usize,
5805 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5806 let f = self.func("moe_pairs_matvec_q8_em");
5807 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5808 const ROWS: u32 = 4;
5809 let cfg = LaunchConfig {
5810 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5811 block_dim: (32, ROWS, 1),
5812 shared_mem_bytes: 0,
5813 };
5814 let (inf, outf, ne, na, rbi) = (
5815 in_f as i32,
5816 out_f as i32,
5817 n_expert as i32,
5818 n_active as i32,
5819 row_bytes as i64,
5820 );
5821 let __s_b = self.gpu.stream();
5822 let mut b = __s_b.launch_builder(&f);
5823 b.arg(table)
5824 .arg(&proj)
5825 .arg(ex_ids)
5826 .arg(ex_off)
5827 .arg(ex_pairs)
5828 .arg(pair_tok)
5829 .arg(aq)
5830 .arg(ad)
5831 .arg(&mut y)
5832 .arg(&inf)
5833 .arg(&outf)
5834 .arg(&ne)
5835 .arg(&na)
5836 .arg(&qtype)
5837 .arg(&rbi);
5838 unsafe {
5839 b.launch(cfg)?;
5840 }
5841 Ok(y)
5842 }
5843
5844 #[allow(clippy::too_many_arguments)]
5847 pub fn moe_pairs_matvec_q8_dec(
5848 &self,
5849 table: &CudaSlice<u64>,
5850 proj: i32,
5851 ex_ids: &CudaSlice<i32>,
5852 ex_off: &CudaSlice<i32>,
5853 ex_pairs: &CudaSlice<i32>,
5854 pair_tok: &CudaSlice<i32>,
5855 aq: &CudaSlice<i8>,
5856 ad: &CudaSlice<f32>,
5857 in_f: usize,
5858 out_f: usize,
5859 n_expert: usize,
5860 n_active: usize,
5861 n_pairs: usize,
5862 qtype: i32,
5863 row_bytes: usize,
5864 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5865 let f = self.func("moe_pairs_matvec_q8_dec");
5866 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5867 const ROWS: u32 = 4;
5868 let cfg = LaunchConfig {
5869 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5870 block_dim: (32, ROWS, 1),
5871 shared_mem_bytes: 0,
5872 };
5873 let (inf, outf, ne, na, rbi) = (
5874 in_f as i32,
5875 out_f as i32,
5876 n_expert as i32,
5877 n_active as i32,
5878 row_bytes as i64,
5879 );
5880 let __s_b = self.gpu.stream();
5881 let mut b = __s_b.launch_builder(&f);
5882 b.arg(table)
5883 .arg(&proj)
5884 .arg(ex_ids)
5885 .arg(ex_off)
5886 .arg(ex_pairs)
5887 .arg(pair_tok)
5888 .arg(aq)
5889 .arg(ad)
5890 .arg(&mut y)
5891 .arg(&inf)
5892 .arg(&outf)
5893 .arg(&ne)
5894 .arg(&na)
5895 .arg(&qtype)
5896 .arg(&rbi);
5897 unsafe {
5898 b.launch(cfg)?;
5899 }
5900 Ok(y)
5901 }
5902
5903 pub fn moe_pairs_gelu_mul(
5904 &self,
5905 gate: &CudaSlice<f32>,
5906 up: &CudaSlice<f32>,
5907 n: usize,
5908 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5909 let f = self.func("moe_pairs_gelu_mul");
5910 let mut act = self.alloc_uninit::<f32>(n)?;
5911 let cfg = LaunchConfig::for_num_elems(n as u32);
5912 let nl = n as i64;
5913 let __s_b = self.gpu.stream();
5914 let mut b = __s_b.launch_builder(&f);
5915 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5916 unsafe {
5917 b.launch(cfg)?;
5918 }
5919 Ok(act)
5920 }
5921
5922 pub fn moe_pairs_silu_mul(
5923 &self,
5924 gate: &CudaSlice<f32>,
5925 up: &CudaSlice<f32>,
5926 n: usize,
5927 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5928 let f = self.func("moe_pairs_silu_mul");
5929 let mut act = self.alloc_uninit::<f32>(n)?;
5930 let cfg = LaunchConfig::for_num_elems(n as u32);
5931 let nl = n as i64;
5932 let __s_b = self.gpu.stream();
5933 let mut b = __s_b.launch_builder(&f);
5934 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5935 unsafe {
5936 b.launch(cfg)?;
5937 }
5938 Ok(act)
5939 }
5940
5941 #[allow(clippy::too_many_arguments)]
5942 pub fn moe_pairs_scatter(
5943 &self,
5944 y_down: &CudaSlice<f32>,
5945 pair_w: &CudaSlice<f32>,
5946 tok_pair_off: &CudaSlice<i32>,
5947 tok_pair_ids: &CudaSlice<i32>,
5948 moe_out: &mut CudaSlice<f32>,
5949 t: usize,
5950 n_embd: usize,
5951 ) -> Result<(), Box<dyn std::error::Error>> {
5952 let f = self.func("moe_pairs_scatter");
5953 let cfg = LaunchConfig {
5954 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5955 block_dim: (256, 1, 1),
5956 shared_mem_bytes: 0,
5957 };
5958 let ne = n_embd as i32;
5959 let __s_b = self.gpu.stream();
5960 let mut b = __s_b.launch_builder(&f);
5961 b.arg(y_down)
5962 .arg(pair_w)
5963 .arg(tok_pair_off)
5964 .arg(tok_pair_ids)
5965 .arg(moe_out)
5966 .arg(&ne);
5967 unsafe {
5968 b.launch(cfg)?;
5969 }
5970 Ok(())
5971 }
5972
5973 #[allow(clippy::too_many_arguments)]
5977 pub fn moe_gate_up_gelu8_dev_q8(
5978 &self,
5979 table: &CudaSlice<u64>,
5980 sel: &cudarc::driver::CudaView<i32>,
5981 aq: &CudaSlice<i8>,
5982 ad: &CudaSlice<f32>,
5983 in_f: usize,
5984 n_ff: usize,
5985 n_used: usize,
5986 n_expert: usize,
5987 qt_g: i32,
5988 qt_u: i32,
5989 rb_g: usize,
5990 rb_u: usize,
5991 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5992 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5993 let (inf, nff, ne, rbg, rbu) = (
5994 in_f as i32,
5995 n_ff as i32,
5996 n_expert as i32,
5997 rb_g as i64,
5998 rb_u as i64,
5999 );
6000 let f = self.func("moe_gate_up_gelu8_dev_q8");
6001 let cfg = LaunchConfig {
6002 grid_dim: (n_ff as u32, n_used as u32, 1),
6003 block_dim: (32, 1, 1),
6004 shared_mem_bytes: 0,
6005 };
6006 let __s_b = self.gpu.stream();
6007 let mut b = __s_b.launch_builder(&f);
6008 b.arg(table)
6009 .arg(sel)
6010 .arg(aq)
6011 .arg(ad)
6012 .arg(&mut act)
6013 .arg(&inf)
6014 .arg(&nff)
6015 .arg(&ne)
6016 .arg(&qt_g)
6017 .arg(&qt_u)
6018 .arg(&rbg)
6019 .arg(&rbu);
6020 unsafe {
6021 b.launch(cfg)?;
6022 }
6023 Ok(act)
6024 }
6025
6026 #[allow(clippy::too_many_arguments)]
6028 pub fn moe_gate_up_gelu8_dev_q8_rows(
6029 &self,
6030 table: &CudaSlice<u64>,
6031 sel: &CudaSlice<i32>,
6032 aq: &CudaSlice<i8>,
6033 ad: &CudaSlice<f32>,
6034 t: usize,
6035 in_f: usize,
6036 n_ff: usize,
6037 n_used: usize,
6038 n_expert: usize,
6039 qt_g: i32,
6040 qt_u: i32,
6041 rb_g: usize,
6042 rb_u: usize,
6043 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6044 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6045 let (inf, nff, ne, rbg, rbu, nu) = (
6046 in_f as i32,
6047 n_ff as i32,
6048 n_expert as i32,
6049 rb_g as i64,
6050 rb_u as i64,
6051 n_used as i32,
6052 );
6053 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
6054 let cfg = LaunchConfig {
6055 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6056 block_dim: (32, 1, 1),
6057 shared_mem_bytes: 0,
6058 };
6059 let __s_b = self.gpu.stream();
6060 let mut b = __s_b.launch_builder(&f);
6061 b.arg(table)
6062 .arg(sel)
6063 .arg(aq)
6064 .arg(ad)
6065 .arg(&mut act)
6066 .arg(&inf)
6067 .arg(&nff)
6068 .arg(&ne)
6069 .arg(&qt_g)
6070 .arg(&qt_u)
6071 .arg(&rbg)
6072 .arg(&rbu)
6073 .arg(&nu);
6074 unsafe {
6075 b.launch(cfg)?;
6076 }
6077 Ok(act)
6078 }
6079
6080 #[allow(clippy::too_many_arguments)]
6082 pub fn moe_gate_up_gelu8_dev_q8_csr(
6083 &self,
6084 table: &CudaSlice<u64>,
6085 sel: &CudaSlice<i32>,
6086 aq: &CudaSlice<i8>,
6087 ad: &CudaSlice<f32>,
6088 n_pairs: usize,
6089 in_f: usize,
6090 n_ff: usize,
6091 n_used: usize,
6092 n_expert: usize,
6093 qt_g: i32,
6094 qt_u: i32,
6095 rb_g: usize,
6096 rb_u: usize,
6097 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6098 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6099 let (inf, nff, ne, rbg, rbu, nu, npi) = (
6100 in_f as i32,
6101 n_ff as i32,
6102 n_expert as i32,
6103 rb_g as i64,
6104 rb_u as i64,
6105 n_used as i32,
6106 n_pairs as i32,
6107 );
6108 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
6109 let cfg = LaunchConfig {
6110 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6111 block_dim: (32, 1, 1),
6112 shared_mem_bytes: 0,
6113 };
6114 let __s_b = self.gpu.stream();
6115 let mut b = __s_b.launch_builder(&f);
6116 b.arg(table)
6117 .arg(sel)
6118 .arg(aq)
6119 .arg(ad)
6120 .arg(&mut act)
6121 .arg(&inf)
6122 .arg(&nff)
6123 .arg(&ne)
6124 .arg(&qt_g)
6125 .arg(&qt_u)
6126 .arg(&rbg)
6127 .arg(&rbu)
6128 .arg(&nu)
6129 .arg(&npi);
6130 unsafe {
6131 b.launch(cfg)?;
6132 }
6133 Ok(act)
6134 }
6135
6136 #[allow(clippy::too_many_arguments)]
6138 pub fn moe_down8_fma_dev_q8_rows_g(
6139 &self,
6140 table: &CudaSlice<u64>,
6141 sel: &CudaSlice<i32>,
6142 w: &CudaSlice<f32>,
6143 aq2: &CudaSlice<i8>,
6144 ad2: &CudaSlice<f32>,
6145 dst: &mut CudaSlice<f32>,
6146 t: usize,
6147 in_f: usize,
6148 out_f: usize,
6149 n_used: usize,
6150 n_expert: usize,
6151 qt: i32,
6152 rb: usize,
6153 ) -> Result<(), Box<dyn std::error::Error>> {
6154 let (inf, outf, nu, ne, rbi) = (
6155 in_f as i32,
6156 out_f as i32,
6157 n_used as i32,
6158 n_expert as i32,
6159 rb as i64,
6160 );
6161 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
6165 let f = self.func(if step_b1_w8 {
6166 "moe_down8_fma_dev_q8_rows_w8"
6167 } else {
6168 "moe_down8_fma_dev_q8_rows_g"
6169 });
6170 let cfg = LaunchConfig {
6171 grid_dim: (out_f as u32, 1, t as u32),
6172 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
6173 shared_mem_bytes: 0,
6174 };
6175 let __s_b = self.gpu.stream();
6176 let mut b = __s_b.launch_builder(&f);
6177 b.arg(table)
6178 .arg(sel)
6179 .arg(w)
6180 .arg(aq2)
6181 .arg(ad2)
6182 .arg(dst)
6183 .arg(&inf)
6184 .arg(&outf)
6185 .arg(&nu)
6186 .arg(&ne)
6187 .arg(&qt)
6188 .arg(&rbi);
6189 unsafe {
6190 b.launch(cfg)?;
6191 }
6192 Ok(())
6193 }
6194
6195 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
6199 let (out_f, in_f) = (2048usize, 2816usize);
6200 let nblk = in_f / 32;
6201 let mut seed = 0x9E3779B97F4A7C15u64;
6202 let mut rng = move || {
6203 seed = seed
6204 .wrapping_mul(6364136223846793005)
6205 .wrapping_add(1442695040888963407);
6206 (seed >> 33) as u8
6207 };
6208 let mut w = vec![0u8; out_f * nblk * 18];
6209 for b in w.iter_mut() {
6210 *b = rng();
6211 }
6212 for r in 0..out_f {
6213 for g in 0..nblk {
6214 let off = (r * nblk + g) * 18;
6215 w[off] = 0x00;
6216 w[off + 1] = 0x2C; }
6218 }
6219 let qplane = out_f * nblk * 16;
6220 let mut wrp = vec![0u8; w.len()];
6221 for r in 0..out_f {
6222 for g in 0..nblk {
6223 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
6224 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
6225 .copy_from_slice(&src[0..2]);
6226 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
6227 }
6228 }
6229 let w_d = self.htod_bytes(&w)?;
6230 let wrp_d = self.htod_bytes(&wrp)?;
6231 let mut aq = vec![0i8; m * in_f];
6232 for v in aq.iter_mut() {
6233 *v = rng() as i8;
6234 }
6235 let aq_d = self.htod_i8(&aq)?;
6236 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
6237 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
6238 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
6239 const RPB: u32 = 4;
6240 let cfg = LaunchConfig {
6241 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
6242 block_dim: (32, RPB, 1),
6243 shared_mem_bytes: 0,
6244 };
6245 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
6246 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
6247 let fb = self.func("qmatvec_q4_0_mmvq_b4");
6248 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
6249 {
6250 let __s_b = self.gpu.stream();
6251 let mut b = __s_b.launch_builder(&fb);
6252 b.arg(&w_d)
6253 .arg(&aq_d)
6254 .arg(&ad_d)
6255 .arg(&mut y0)
6256 .arg(&inf)
6257 .arg(&outf)
6258 .arg(&mi)
6259 .arg(&rb);
6260 unsafe {
6261 b.launch(cfg)?;
6262 }
6263 let __s_b = self.gpu.stream();
6264 let mut b = __s_b.launch_builder(&fr);
6265 b.arg(&wrp_d)
6266 .arg(&aq_d)
6267 .arg(&ad_d)
6268 .arg(&mut y1)
6269 .arg(&inf)
6270 .arg(&outf)
6271 .arg(&mi)
6272 .arg(&qp);
6273 unsafe {
6274 b.launch(cfg)?;
6275 }
6276 }
6277 self.gpu.stream().synchronize()?;
6278 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
6279 let nd = h0
6280 .iter()
6281 .zip(&h1)
6282 .filter(|(a, b)| a.to_bits() != b.to_bits())
6283 .count();
6284 if nd != 0 {
6285 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
6286 }
6287 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
6288 self.gpu.stream().synchronize()?;
6289 let t0 = std::time::Instant::now();
6290 for _ in 0..500 {
6291 if rp {
6292 let __s_b = self.gpu.stream();
6293 let mut b = __s_b.launch_builder(&fr);
6294 b.arg(&wrp_d)
6295 .arg(&aq_d)
6296 .arg(&ad_d)
6297 .arg(&mut y1)
6298 .arg(&inf)
6299 .arg(&outf)
6300 .arg(&mi)
6301 .arg(&qp);
6302 unsafe {
6303 b.launch(cfg)?;
6304 }
6305 } else {
6306 let __s_b = self.gpu.stream();
6307 let mut b = __s_b.launch_builder(&fb);
6308 b.arg(&w_d)
6309 .arg(&aq_d)
6310 .arg(&ad_d)
6311 .arg(&mut y0)
6312 .arg(&inf)
6313 .arg(&outf)
6314 .arg(&mi)
6315 .arg(&rb);
6316 unsafe {
6317 b.launch(cfg)?;
6318 }
6319 }
6320 }
6321 self.gpu.stream().synchronize()?;
6322 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
6323 };
6324 let _ = time(false)?;
6325 let _ = time(true)?; Ok((time(false)?, time(true)?))
6327 }
6328
6329 pub fn build_q4_rp4(
6334 &self,
6335 t: &mut crate::model::GpuTensor,
6336 ) -> Result<(), Box<dyn std::error::Error>> {
6337 use crate::model::GpuTensor;
6338 let GpuTensor::Quant {
6339 bytes,
6340 qtype,
6341 row_bytes,
6342 ne,
6343 rp4,
6344 ..
6345 } = t
6346 else {
6347 return Ok(());
6348 };
6349 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
6350 return Ok(());
6351 }
6352 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6353 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
6354 return Ok(());
6355 }
6356 let nblk = in_f / 32;
6357 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
6358 let f = self.func("q4_0_split_rp_build");
6359 let n = (out_f * nblk) as i32;
6360 let cfg = LaunchConfig {
6361 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6362 block_dim: (256, 1, 1),
6363 shared_mem_bytes: 0,
6364 };
6365 let (of, nb) = (out_f as i32, nblk as i32);
6366 let _ = n;
6367 let __s_b = self.gpu.stream();
6368 let mut b = __s_b.launch_builder(&f);
6369 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6370 unsafe {
6371 b.launch(cfg)?;
6372 }
6373 *rp4 = Some(dst);
6374 Ok(())
6375 }
6376
6377 pub fn build_q8_rp4(
6382 &self,
6383 t: &mut crate::model::GpuTensor,
6384 ) -> Result<(), Box<dyn std::error::Error>> {
6385 use crate::model::GpuTensor;
6386 let GpuTensor::Quant {
6387 bytes,
6388 qtype,
6389 row_bytes,
6390 ne,
6391 rp4,
6392 ..
6393 } = t
6394 else {
6395 return Ok(());
6396 };
6397 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
6398 return Ok(());
6399 }
6400 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6401 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
6402 return Ok(());
6403 }
6404 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
6405 Ok(())
6406 }
6407
6408 pub fn build_q8_rp4_raw(
6411 &self,
6412 bytes: &CudaSlice<u8>,
6413 in_f: usize,
6414 out_f: usize,
6415 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6416 assert!(in_f % 32 == 0);
6417 let nblk = in_f / 32;
6418 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
6419 let f = self.func("q8_0_split_rp_build");
6420 let cfg = LaunchConfig {
6421 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6422 block_dim: (256, 1, 1),
6423 shared_mem_bytes: 0,
6424 };
6425 let (of, nb) = (out_f as i32, nblk as i32);
6426 let __s_b = self.gpu.stream();
6427 let mut b = __s_b.launch_builder(&f);
6428 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6429 unsafe {
6430 b.launch(cfg)?;
6431 }
6432 Ok(dst)
6433 }
6434
6435 pub fn build_q4k_rp4(
6443 &self,
6444 t: &mut crate::model::GpuTensor,
6445 ) -> Result<(), Box<dyn std::error::Error>> {
6446 use crate::model::GpuTensor;
6447 let GpuTensor::Quant {
6448 bytes,
6449 qtype,
6450 row_bytes,
6451 ne,
6452 rp4,
6453 ..
6454 } = t
6455 else {
6456 return Ok(());
6457 };
6458 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
6459 return Ok(());
6460 }
6461 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6462 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
6463 return Ok(());
6464 }
6465 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
6466 Ok(())
6467 }
6468
6469 pub fn build_q6k_rp4(
6470 &self,
6471 t: &mut crate::model::GpuTensor,
6472 ) -> Result<(), Box<dyn std::error::Error>> {
6473 use crate::model::GpuTensor;
6474 let GpuTensor::Quant {
6475 bytes,
6476 qtype,
6477 row_bytes,
6478 ne,
6479 rp4,
6480 ..
6481 } = t
6482 else {
6483 return Ok(());
6484 };
6485 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
6486 return Ok(());
6487 }
6488 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6489 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
6490 return Ok(());
6491 }
6492 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
6493 Ok(())
6494 }
6495
6496 pub fn build_kq_rp4_raw(
6498 &self,
6499 bytes: &CudaSlice<u8>,
6500 in_f: usize,
6501 out_f: usize,
6502 qtype: i32,
6503 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6504 assert!(in_f % 256 == 0);
6505 let nsbk = in_f / 256;
6506 let (sb_bytes, kname) = match qtype {
6507 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
6508 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
6509 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
6510 };
6511 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
6512 let f = self.func(kname);
6513 let cfg = LaunchConfig {
6514 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
6515 block_dim: (256, 1, 1),
6516 shared_mem_bytes: 0,
6517 };
6518 let (of, nb) = (out_f as i32, nsbk as i32);
6519 let __s_b = self.gpu.stream();
6520 let mut b = __s_b.launch_builder(&f);
6521 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6522 unsafe {
6523 b.launch(cfg)?;
6524 }
6525 Ok(dst)
6526 }
6527
6528 pub fn kqrp_enabled() -> bool {
6532 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6533 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6534 Ok("0") => false,
6535 Ok(_) => true,
6536 Err(_) => cfg!(memra_hopper_mma),
6537 })
6538 }
6539
6540 pub fn build_q4_rp_swap(
6546 &self,
6547 t: &mut crate::model::GpuTensor,
6548 ) -> Result<bool, Box<dyn std::error::Error>> {
6549 use crate::model::GpuTensor;
6550 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6560 return Ok(false);
6561 }
6562 self.build_q4_rp4(t)?;
6563 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6565 return Ok(false);
6566 };
6567 match rp4.take() {
6568 Some(split) => {
6569 *bytes = split; *rp = true;
6571 Ok(true)
6572 }
6573 None => Ok(false),
6574 }
6575 }
6576
6577 pub fn q4rp_enabled() -> bool {
6579 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6580 *ON.get_or_init(|| {
6581 std::env::var("MEMRA_Q4RP")
6582 .map(|v| v != "0")
6583 .unwrap_or(true)
6584 })
6585 }
6586
6587 pub fn copy_rows_strided(
6590 &self,
6591 src: &CudaSlice<f32>,
6592 dst: &mut CudaSlice<f32>,
6593 row_elems: usize,
6594 n_rows: usize,
6595 src_stride: usize,
6596 src_off: usize,
6597 ) -> Result<(), Box<dyn std::error::Error>> {
6598 let f = self.func("copy_rows_strided_f32");
6599 let cfg = LaunchConfig {
6600 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6601 block_dim: (256, 1, 1),
6602 shared_mem_bytes: 0,
6603 };
6604 let (re, nr) = (row_elems as i32, n_rows as i32);
6605 let (st, off) = (src_stride as i64, src_off as i64);
6606 let __s_b = self.gpu.stream();
6607 let mut b = __s_b.launch_builder(&f);
6608 b.arg(src)
6609 .arg(&mut *dst)
6610 .arg(&re)
6611 .arg(&nr)
6612 .arg(&st)
6613 .arg(&off);
6614 unsafe {
6615 b.launch(cfg)?;
6616 }
6617 Ok(())
6618 }
6619
6620 pub fn place_rows_strided(
6626 &self,
6627 src: &CudaSlice<f32>,
6628 dst: &mut CudaSlice<f32>,
6629 row_elems: usize,
6630 n_rows: usize,
6631 dst_stride: usize,
6632 dst_off: usize,
6633 ) -> Result<(), Box<dyn std::error::Error>> {
6634 if row_elems == 0 || n_rows == 0 {
6635 return Err("strided row placement requires nonzero rows and row width".into());
6636 }
6637 let src_len = n_rows
6638 .checked_mul(row_elems)
6639 .ok_or("strided row placement source size overflow")?;
6640 let dst_len = n_rows
6641 .checked_sub(1)
6642 .and_then(|rows| rows.checked_mul(dst_stride))
6643 .and_then(|base| base.checked_add(dst_off))
6644 .and_then(|base| base.checked_add(row_elems))
6645 .ok_or("strided row placement destination size overflow")?;
6646 let row_end = dst_off
6647 .checked_add(row_elems)
6648 .ok_or("strided row placement row size overflow")?;
6649 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6650 return Err(format!(
6651 "strided row placement geometry mismatch: src={} need_src={src_len} \
6652 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6653 dst_stride={dst_stride} dst_off={dst_off}",
6654 src.len(),
6655 dst.len(),
6656 )
6657 .into());
6658 }
6659 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6660 return Err("strided row placement exceeds CUDA kernel geometry".into());
6661 }
6662 let f = self.func("place_rows_strided_f32");
6663 let cfg = LaunchConfig {
6664 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6665 block_dim: (256, 1, 1),
6666 shared_mem_bytes: 0,
6667 };
6668 let (re, nr) = (row_elems as i32, n_rows as i32);
6669 let (st, off) = (dst_stride as i64, dst_off as i64);
6670 let __s_b = self.gpu.stream();
6671 let mut b = __s_b.launch_builder(&f);
6672 b.arg(src)
6673 .arg(&mut *dst)
6674 .arg(&re)
6675 .arg(&nr)
6676 .arg(&st)
6677 .arg(&off);
6678 unsafe {
6679 b.launch(cfg)?;
6680 }
6681 Ok(())
6682 }
6683
6684 pub fn u32_set_k(
6686 &self,
6687 dst: &mut CudaSlice<u32>,
6688 v: u32,
6689 idx: usize,
6690 ) -> Result<(), Box<dyn std::error::Error>> {
6691 let f = self.func("u32_set_k");
6692 let cfg = LaunchConfig {
6693 grid_dim: (1, 1, 1),
6694 block_dim: (1, 1, 1),
6695 shared_mem_bytes: 0,
6696 };
6697 let ii = idx as i32;
6698 let __s_b = self.gpu.stream();
6699 let mut b = __s_b.launch_builder(&f);
6700 b.arg(dst).arg(&v).arg(&ii);
6701 unsafe {
6702 b.launch(cfg)?;
6703 }
6704 Ok(())
6705 }
6706
6707 pub fn i32_add_k(
6709 &self,
6710 d: &mut CudaSlice<i32>,
6711 v: i32,
6712 ) -> Result<(), Box<dyn std::error::Error>> {
6713 let f = self.func("i32_add_k");
6714 let cfg = LaunchConfig {
6715 grid_dim: (1, 1, 1),
6716 block_dim: (32, 1, 1),
6717 shared_mem_bytes: 0,
6718 };
6719 let __s_b = self.gpu.stream();
6720 let mut b = __s_b.launch_builder(&f);
6721 b.arg(d).arg(&v);
6722 unsafe {
6723 b.launch(cfg)?;
6724 }
6725 Ok(())
6726 }
6727
6728 pub fn i32_iota_from(
6730 &self,
6731 ctr: &CudaSlice<i32>,
6732 dst: &mut CudaSlice<i32>,
6733 n: usize,
6734 ) -> Result<(), Box<dyn std::error::Error>> {
6735 let f = self.func("i32_iota_from");
6736 let cfg = LaunchConfig::for_num_elems(n as u32);
6737 let ni = n as i32;
6738 let __s_b = self.gpu.stream();
6739 let mut b = __s_b.launch_builder(&f);
6740 b.arg(ctr).arg(dst).arg(&ni);
6741 unsafe {
6742 b.launch(cfg)?;
6743 }
6744 Ok(())
6745 }
6746
6747 pub fn u32_map_k(
6749 &self,
6750 buf: &mut CudaSlice<u32>,
6751 map: &CudaSlice<u32>,
6752 idx: usize,
6753 ) -> Result<(), Box<dyn std::error::Error>> {
6754 let f = self.func("u32_map_k");
6755 let cfg = LaunchConfig {
6756 grid_dim: (1, 1, 1),
6757 block_dim: (1, 1, 1),
6758 shared_mem_bytes: 0,
6759 };
6760 let ii = idx as i32;
6761 let __s_b = self.gpu.stream();
6762 let mut b = __s_b.launch_builder(&f);
6763 b.arg(buf).arg(map).arg(&ii);
6764 unsafe {
6765 b.launch(cfg)?;
6766 }
6767 Ok(())
6768 }
6769
6770 #[allow(clippy::too_many_arguments)]
6772 pub fn u32_pack2(
6773 &self,
6774 a: &CudaSlice<u32>,
6775 off_a: usize,
6776 n1: usize,
6777 b_in: &CudaSlice<u32>,
6778 n2: usize,
6779 out: &mut CudaSlice<u32>,
6780 ) -> Result<(), Box<dyn std::error::Error>> {
6781 let f = self.func("u32_pack2");
6782 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6783 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6784 let __s_b = self.gpu.stream();
6785 let mut b = __s_b.launch_builder(&f);
6786 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6787 unsafe {
6788 b.launch(cfg)?;
6789 }
6790 Ok(())
6791 }
6792
6793 pub fn moe_w_exscale(
6795 &self,
6796 w: &mut CudaSlice<f32>,
6797 sel: &CudaSlice<i32>,
6798 s: &CudaSlice<f32>,
6799 n: usize,
6800 ) -> Result<(), Box<dyn std::error::Error>> {
6801 let f = self.func("moe_w_exscale");
6802 let cfg = LaunchConfig::for_num_elems(n as u32);
6803 let ni = n as i32;
6804 let __s_b = self.gpu.stream();
6805 let mut b = __s_b.launch_builder(&f);
6806 b.arg(w).arg(sel).arg(s).arg(&ni);
6807 unsafe {
6808 b.launch(cfg)?;
6809 }
6810 Ok(())
6811 }
6812
6813 pub fn moe_w_scale_by_expert(
6816 &self,
6817 w: &mut CudaSlice<f32>,
6818 sel: &CudaSlice<i32>,
6819 macros: &CudaSlice<f32>,
6820 n_expert: usize,
6821 n: usize,
6822 ) -> Result<(), Box<dyn std::error::Error>> {
6823 let f = self.func("moe_w_scale_by_expert");
6824 let cfg = LaunchConfig {
6825 grid_dim: (n.div_ceil(64) as u32, 1, 1),
6826 block_dim: (64, 1, 1),
6827 shared_mem_bytes: 0,
6828 };
6829 let (ne, nn) = (n_expert as i32, n as i32);
6830 let __s_b = self.gpu.stream();
6831 let mut b = __s_b.launch_builder(&f);
6832 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6833 unsafe {
6834 b.launch(cfg)?;
6835 }
6836 Ok(())
6837 }
6838
6839 pub fn moe_gate_up_silu8_dev_q8(
6840 &self,
6841 table: &CudaSlice<u64>,
6842 sel: &cudarc::driver::CudaView<i32>,
6843 aq: &CudaSlice<i8>,
6844 ad: &CudaSlice<f32>,
6845 in_f: usize,
6846 n_ff: usize,
6847 n_used: usize,
6848 n_expert: usize,
6849 qt_g: i32,
6850 qt_u: i32,
6851 rb_g: usize,
6852 rb_u: usize,
6853 macros: &CudaSlice<f32>,
6854 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6855 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6856 let (mode, wpb) = GU.get_or_init(|| {
6857 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6858 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6859 .ok()
6860 .and_then(|v| v.parse().ok())
6861 .unwrap_or(4u32)
6862 .clamp(1, 16);
6863 (mode, wpb)
6864 });
6865 let (mode, wpb) = (mode.as_str(), *wpb);
6866 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6867 let (inf, nff, ne, rbg, rbu) = (
6868 in_f as i32,
6869 n_ff as i32,
6870 n_expert as i32,
6871 rb_g as i64,
6872 rb_u as i64,
6873 );
6874 let (f, cfg) = match mode {
6875 "1" | "2" | "4" => {
6876 let rpw: u32 = mode.parse().unwrap();
6877 let f = self.func(match rpw {
6878 1 => "moe_gate_up_silu8_dev_q8_r1",
6879 2 => "moe_gate_up_silu8_dev_q8_r2",
6880 _ => "moe_gate_up_silu8_dev_q8_r4",
6881 });
6882 let rows_per_block = (rpw * wpb) as usize;
6883 let gx = n_ff.div_ceil(rows_per_block) as u32;
6884 (
6885 f,
6886 LaunchConfig {
6887 grid_dim: (gx, n_used as u32, 1),
6888 block_dim: (32, wpb, 1),
6889 shared_mem_bytes: 0,
6890 },
6891 )
6892 }
6893 "j8" if n_used <= 32 => (
6894 self.func("moe_gate_up_silu8_dev_q8_j8"),
6895 LaunchConfig {
6896 grid_dim: (n_ff as u32, 1, 1),
6897 block_dim: (32, n_used as u32, 1),
6898 shared_mem_bytes: 0,
6899 },
6900 ),
6901 "vsm2" => {
6903 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6904 let sh = (rb_g + rb_u) as u32;
6905 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6906 f.set_attribute(
6907 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6908 sh as i32,
6909 )?;
6910 (
6911 f,
6912 LaunchConfig {
6913 grid_dim: (n_ff as u32, n_used as u32, 1),
6914 block_dim: (32, 1, 1),
6915 shared_mem_bytes: sh,
6916 },
6917 )
6918 }
6919 "vsm" => {
6920 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6921 let sh = (rb_g + rb_u) as u32;
6922 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6923 f.set_attribute(
6924 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6925 sh as i32,
6926 )?;
6927 (
6928 f,
6929 LaunchConfig {
6930 grid_dim: (n_ff as u32, n_used as u32, 1),
6931 block_dim: (32, 1, 1),
6932 shared_mem_bytes: sh,
6933 },
6934 )
6935 }
6936 "sg" => (
6937 self.func("moe_gate_up_silu8_dev_q8_sg"),
6938 LaunchConfig {
6939 grid_dim: (n_ff as u32, n_used as u32, 1),
6940 block_dim: (32, 1, 1),
6941 shared_mem_bytes: 0,
6942 },
6943 ),
6944 "j8sg" if n_used <= 32 => (
6945 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6946 LaunchConfig {
6947 grid_dim: (n_ff as u32, 1, 1),
6948 block_dim: (32, n_used as u32, 1),
6949 shared_mem_bytes: 0,
6950 },
6951 ),
6952 "u64" if in_f == 2048 => (
6953 self.func("moe_gate_up_silu8_dev_q8_u64"),
6954 LaunchConfig {
6955 grid_dim: (n_ff as u32, n_used as u32, 1),
6956 block_dim: (32, 1, 1),
6957 shared_mem_bytes: 0,
6958 },
6959 ),
6960 "gs4" if in_f == 2048 => (
6961 self.func("moe_gate_up_silu8_dev_q8_gs4"),
6962 LaunchConfig {
6963 grid_dim: (n_ff as u32, n_used as u32, 1),
6964 block_dim: (32, 4, 1),
6965 shared_mem_bytes: 0,
6966 },
6967 ),
6968 "v" | "" => (
6970 self.func("moe_gate_up_silu8_dev_q8_v"),
6971 LaunchConfig {
6972 grid_dim: (n_ff as u32, n_used as u32, 1),
6973 block_dim: (32, 1, 1),
6974 shared_mem_bytes: 0,
6975 },
6976 ),
6977 "s2" => (
6978 self.func("moe_gate_up_silu8_dev_q8_s2"),
6979 LaunchConfig {
6980 grid_dim: (n_ff as u32, n_used as u32, 1),
6981 block_dim: (32, 2, 1),
6982 shared_mem_bytes: 0,
6983 },
6984 ),
6985 "s2z" => {
6986 let rz = wpb.min(16); (
6988 self.func("moe_gate_up_silu8_dev_q8_s2z"),
6989 LaunchConfig {
6990 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6991 block_dim: (32, 2, rz),
6992 shared_mem_bytes: 0,
6993 },
6994 )
6995 }
6996 _ => (
6997 self.func("moe_gate_up_silu8_dev_q8"),
6998 LaunchConfig {
6999 grid_dim: (n_ff as u32, n_used as u32, 1),
7000 block_dim: (32, 1, 1),
7001 shared_mem_bytes: 0,
7002 },
7003 ),
7004 };
7005 let __s_b = self.gpu.stream();
7006 let mut b = __s_b.launch_builder(&f);
7007 b.arg(table)
7008 .arg(sel)
7009 .arg(aq)
7010 .arg(ad)
7011 .arg(&mut act)
7012 .arg(&inf)
7013 .arg(&nff)
7014 .arg(&ne)
7015 .arg(&qt_g)
7016 .arg(&qt_u)
7017 .arg(&rbg)
7018 .arg(&rbu)
7019 .arg(macros);
7020 unsafe {
7021 b.launch(cfg)?;
7022 }
7023 Ok(act)
7024 }
7025
7026 #[allow(clippy::too_many_arguments)]
7027 pub fn moe_down8_fma_dev_q8(
7028 &self,
7029 table: &CudaSlice<u64>,
7030 sel: &cudarc::driver::CudaView<i32>,
7031 w: &cudarc::driver::CudaView<f32>,
7032 aq2: &CudaSlice<i8>,
7033 ad2: &CudaSlice<f32>,
7034 dst: &mut cudarc::driver::CudaViewMut<f32>,
7035 in_f: usize,
7036 out_f: usize,
7037 n_used: usize,
7038 n_expert: usize,
7039 qt: i32,
7040 rb: usize,
7041 ) -> Result<(), Box<dyn std::error::Error>> {
7042 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
7043 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
7044 let (inf, outf, nu, ne, rbi) = (
7045 in_f as i32,
7046 out_f as i32,
7047 n_used as i32,
7048 n_expert as i32,
7049 rb as i64,
7050 );
7051 let (f, cfg) = match mode.as_str() {
7054 m @ ("1" | "2" | "4") if n_used <= 8 => {
7055 let rpw: usize = m.parse().unwrap();
7056 let f = self.func(match rpw {
7057 1 => "moe_down8_fma_dev_q8_w8r1",
7058 2 => "moe_down8_fma_dev_q8_w8r2",
7059 _ => "moe_down8_fma_dev_q8_w8r4",
7060 });
7061 (
7062 f,
7063 LaunchConfig {
7064 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
7065 block_dim: (32, n_used as u32, 1),
7066 shared_mem_bytes: 0,
7067 },
7068 )
7069 }
7070 "h2" if in_f == 512 => (
7071 self.func("moe_down8_fma_dev_q8_h2"),
7072 LaunchConfig {
7073 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7074 block_dim: (32, 1, 1),
7075 shared_mem_bytes: 0,
7076 },
7077 ),
7078 "" if in_f == 704 && n_used <= 8 => (
7081 self.func("moe_down8_fma_dev_q8_w8r2"),
7082 LaunchConfig {
7083 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7084 block_dim: (32, n_used as u32, 1),
7085 shared_mem_bytes: 0,
7086 },
7087 ),
7088 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
7092 self.func("moe_down8_fma_dev_q8_w8h2v"),
7093 LaunchConfig {
7094 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7095 block_dim: (32, n_used as u32, 1),
7096 shared_mem_bytes: 0,
7097 },
7098 ),
7099 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
7100 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
7101 LaunchConfig {
7102 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
7103 block_dim: (32, n_used as u32, 1),
7104 shared_mem_bytes: 0,
7105 },
7106 ),
7107 "w8h2r2" if in_f == 512 && n_used <= 8 => (
7108 self.func("moe_down8_fma_dev_q8_w8h2r2"),
7109 LaunchConfig {
7110 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
7111 block_dim: (32, n_used as u32, 1),
7112 shared_mem_bytes: 0,
7113 },
7114 ),
7115 "w8h2" if in_f == 512 && n_used <= 8 => (
7116 self.func("moe_down8_fma_dev_q8_w8h2"),
7117 LaunchConfig {
7118 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7119 block_dim: (32, n_used as u32, 1),
7120 shared_mem_bytes: 0,
7121 },
7122 ),
7123 _ => (
7124 self.func("moe_down8_fma_dev_q8"),
7125 LaunchConfig {
7126 grid_dim: (out_f as u32, 1, 1),
7127 block_dim: (32, 1, 1),
7128 shared_mem_bytes: 0,
7129 },
7130 ),
7131 };
7132 let __s_b = self.gpu.stream();
7133 let mut b = __s_b.launch_builder(&f);
7134 b.arg(table)
7135 .arg(sel)
7136 .arg(w)
7137 .arg(aq2)
7138 .arg(ad2)
7139 .arg(dst)
7140 .arg(&inf)
7141 .arg(&outf)
7142 .arg(&nu)
7143 .arg(&ne)
7144 .arg(&qt)
7145 .arg(&rbi);
7146 unsafe {
7147 b.launch(cfg)?;
7148 }
7149 Ok(())
7150 }
7151
7152 #[allow(clippy::too_many_arguments)]
7159 pub fn moe_gate_up_silu8_dev_q8_rows(
7160 &self,
7161 table: &CudaSlice<u64>,
7162 sel: &CudaSlice<i32>,
7163 aq: &CudaSlice<i8>,
7164 ad: &CudaSlice<f32>,
7165 t: usize,
7166 in_f: usize,
7167 n_ff: usize,
7168 n_used: usize,
7169 n_expert: usize,
7170 qt_g: i32,
7171 qt_u: i32,
7172 rb_g: usize,
7173 rb_u: usize,
7174 macros: &CudaSlice<f32>,
7175 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7176 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
7177 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
7178 let cfg = LaunchConfig {
7179 grid_dim: (n_ff as u32, n_used as u32, t as u32),
7180 block_dim: (32, 1, 1),
7181 shared_mem_bytes: 0,
7182 };
7183 let (inf, nff, ne, nu, rbg, rbu) = (
7184 in_f as i32,
7185 n_ff as i32,
7186 n_expert as i32,
7187 n_used as i32,
7188 rb_g as i64,
7189 rb_u as i64,
7190 );
7191 let __s_b = self.gpu.stream();
7192 let mut b = __s_b.launch_builder(&f);
7193 b.arg(table)
7194 .arg(sel)
7195 .arg(aq)
7196 .arg(ad)
7197 .arg(&mut act)
7198 .arg(&inf)
7199 .arg(&nff)
7200 .arg(&ne)
7201 .arg(&qt_g)
7202 .arg(&qt_u)
7203 .arg(&rbg)
7204 .arg(&rbu)
7205 .arg(&nu)
7206 .arg(macros);
7207 unsafe {
7208 b.launch(cfg)?;
7209 }
7210 Ok(act)
7211 }
7212
7213 #[allow(clippy::too_many_arguments)]
7218 pub fn moe_down8_fma_dev_q8_rows(
7219 &self,
7220 table: &CudaSlice<u64>,
7221 sel: &CudaSlice<i32>,
7222 w: &CudaSlice<f32>,
7223 aq2: &CudaSlice<i8>,
7224 ad2: &CudaSlice<f32>,
7225 dst: &mut CudaSlice<f32>,
7226 t: usize,
7227 in_f: usize,
7228 out_f: usize,
7229 n_used: usize,
7230 n_expert: usize,
7231 qt: i32,
7232 rb: usize,
7233 ) -> Result<(), Box<dyn std::error::Error>> {
7234 assert!(
7235 in_f == 512 && n_used <= 8,
7236 "down rows twin is w8h2v shape-gated"
7237 );
7238 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
7239 let cfg = LaunchConfig {
7240 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
7241 block_dim: (32, n_used as u32, 1),
7242 shared_mem_bytes: 0,
7243 };
7244 let (inf, outf, nu, ne, rbi) = (
7245 in_f as i32,
7246 out_f as i32,
7247 n_used as i32,
7248 n_expert as i32,
7249 rb as i64,
7250 );
7251 let __s_b = self.gpu.stream();
7252 let mut b = __s_b.launch_builder(&f);
7253 b.arg(table)
7254 .arg(sel)
7255 .arg(w)
7256 .arg(aq2)
7257 .arg(ad2)
7258 .arg(dst)
7259 .arg(&inf)
7260 .arg(&outf)
7261 .arg(&nu)
7262 .arg(&ne)
7263 .arg(&qt)
7264 .arg(&rbi);
7265 unsafe {
7266 b.launch(cfg)?;
7267 }
7268 Ok(())
7269 }
7270
7271 #[allow(clippy::too_many_arguments)]
7275 pub fn moe_gate_up_silu8_dev_q8_csr(
7276 &self,
7277 table: &CudaSlice<u64>,
7278 sel: &CudaSlice<i32>,
7279 aq: &CudaSlice<i8>,
7280 ad: &CudaSlice<f32>,
7281 n_pairs: usize,
7282 in_f: usize,
7283 n_ff: usize,
7284 n_used: usize,
7285 n_expert: usize,
7286 qt_g: i32,
7287 qt_u: i32,
7288 rb_g: usize,
7289 rb_u: usize,
7290 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7291 let f = if qt_g == crate::QT_NVFP4 {
7294 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
7295 } else {
7296 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
7297 };
7298 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
7299 let cfg = LaunchConfig {
7300 grid_dim: (n_ff as u32, n_pairs as u32, 1),
7301 block_dim: (32, 1, 1),
7302 shared_mem_bytes: 0,
7303 };
7304 let (inf, nff, ne, nu, npi, rbg, rbu) = (
7305 in_f as i32,
7306 n_ff as i32,
7307 n_expert as i32,
7308 n_used as i32,
7309 n_pairs as i32,
7310 rb_g as i64,
7311 rb_u as i64,
7312 );
7313 let __s_b = self.gpu.stream();
7314 let mut b = __s_b.launch_builder(&f);
7315 b.arg(table)
7316 .arg(sel)
7317 .arg(aq)
7318 .arg(ad)
7319 .arg(&mut act)
7320 .arg(&inf)
7321 .arg(&nff)
7322 .arg(&ne)
7323 .arg(&qt_g)
7324 .arg(&qt_u)
7325 .arg(&rbg)
7326 .arg(&rbu)
7327 .arg(&nu)
7328 .arg(&npi);
7329 unsafe {
7330 b.launch(cfg)?;
7331 }
7332 Ok(act)
7333 }
7334
7335 #[allow(clippy::too_many_arguments)]
7339 pub fn moe_down8_fma_dev_q8_variant(
7340 &self,
7341 variant: &str,
7342 table: &CudaSlice<u64>,
7343 sel: &cudarc::driver::CudaView<i32>,
7344 w: &cudarc::driver::CudaView<f32>,
7345 aq2: &CudaSlice<i8>,
7346 ad2: &CudaSlice<f32>,
7347 dst: &mut cudarc::driver::CudaViewMut<f32>,
7348 in_f: usize,
7349 out_f: usize,
7350 n_used: usize,
7351 n_expert: usize,
7352 qt: i32,
7353 rb: usize,
7354 ) -> Result<(), Box<dyn std::error::Error>> {
7355 let (inf, outf, nu, ne, rbi) = (
7356 in_f as i32,
7357 out_f as i32,
7358 n_used as i32,
7359 n_expert as i32,
7360 rb as i64,
7361 );
7362 let (f, cfg) = match variant {
7363 "w8h2" | "w8h2v" => (
7364 self.func(if variant == "w8h2" {
7365 "moe_down8_fma_dev_q8_w8h2"
7366 } else {
7367 "moe_down8_fma_dev_q8_w8h2v"
7368 }),
7369 LaunchConfig {
7370 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7371 block_dim: (32, n_used as u32, 1),
7372 shared_mem_bytes: 0,
7373 },
7374 ),
7375 "w8h2r2" | "w8h2r2v" => (
7376 self.func(if variant == "w8h2r2" {
7377 "moe_down8_fma_dev_q8_w8h2r2"
7378 } else {
7379 "moe_down8_fma_dev_q8_w8h2r2v"
7380 }),
7381 LaunchConfig {
7382 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
7383 block_dim: (32, n_used as u32, 1),
7384 shared_mem_bytes: 0,
7385 },
7386 ),
7387 _ => (
7388 self.func("moe_down8_fma_dev_q8"),
7389 LaunchConfig {
7390 grid_dim: (out_f as u32, 1, 1),
7391 block_dim: (32, 1, 1),
7392 shared_mem_bytes: 0,
7393 },
7394 ),
7395 };
7396 let __s_b = self.gpu.stream();
7397 let mut b = __s_b.launch_builder(&f);
7398 b.arg(table)
7399 .arg(sel)
7400 .arg(w)
7401 .arg(aq2)
7402 .arg(ad2)
7403 .arg(dst)
7404 .arg(&inf)
7405 .arg(&outf)
7406 .arg(&nu)
7407 .arg(&ne)
7408 .arg(&qt)
7409 .arg(&rbi);
7410 unsafe {
7411 b.launch(cfg)?;
7412 }
7413 Ok(())
7414 }
7415
7416 #[allow(clippy::too_many_arguments)]
7418 pub fn moe_gate_up_silu8_dev_q8_variant(
7419 &self,
7420 variant: &str,
7421 table: &CudaSlice<u64>,
7422 sel: &cudarc::driver::CudaView<i32>,
7423 aq: &CudaSlice<i8>,
7424 ad: &CudaSlice<f32>,
7425 in_f: usize,
7426 n_ff: usize,
7427 n_used: usize,
7428 n_expert: usize,
7429 qt_g: i32,
7430 qt_u: i32,
7431 rb_g: usize,
7432 rb_u: usize,
7433 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7434 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7435 let (inf, nff, ne, rbg, rbu) = (
7436 in_f as i32,
7437 n_ff as i32,
7438 n_expert as i32,
7439 rb_g as i64,
7440 rb_u as i64,
7441 );
7442 let f = self.func(if variant == "v" {
7443 "moe_gate_up_silu8_dev_q8_v"
7444 } else {
7445 "moe_gate_up_silu8_dev_q8"
7446 });
7447 let cfg = LaunchConfig {
7448 grid_dim: (n_ff as u32, n_used as u32, 1),
7449 block_dim: (32, 1, 1),
7450 shared_mem_bytes: 0,
7451 };
7452 let __s_b = self.gpu.stream();
7453 let mut b = __s_b.launch_builder(&f);
7454 b.arg(table)
7455 .arg(sel)
7456 .arg(aq)
7457 .arg(ad)
7458 .arg(&mut act)
7459 .arg(&inf)
7460 .arg(&nff)
7461 .arg(&ne)
7462 .arg(&qt_g)
7463 .arg(&qt_u)
7464 .arg(&rbg)
7465 .arg(&rbu);
7466 unsafe {
7467 b.launch(cfg)?;
7468 }
7469 Ok(act)
7470 }
7471
7472 pub fn moe_gate_up_silu8_dev(
7473 &self,
7474 table: &CudaSlice<u64>,
7475 sel: &cudarc::driver::CudaView<i32>,
7476 x: &cudarc::driver::CudaView<f32>,
7477 in_f: usize,
7478 n_ff: usize,
7479 n_used: usize,
7480 n_expert: usize,
7481 qt_g: i32,
7482 qt_u: i32,
7483 rb_g: usize,
7484 rb_u: usize,
7485 macros: &CudaSlice<f32>,
7486 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7487 let f = self.func("moe_gate_up_silu8_dev");
7488 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
7490 grid_dim: (n_ff as u32, n_used as u32, 1),
7491 block_dim: (256, 1, 1),
7492 shared_mem_bytes: 0,
7493 };
7494 let (inf, nff, ne, rbg, rbu) = (
7495 in_f as i32,
7496 n_ff as i32,
7497 n_expert as i32,
7498 rb_g as i64,
7499 rb_u as i64,
7500 );
7501 let __s_b = self.gpu.stream();
7502 let mut b = __s_b.launch_builder(&f);
7503 b.arg(table)
7504 .arg(sel)
7505 .arg(x)
7506 .arg(&mut act)
7507 .arg(&inf)
7508 .arg(&nff)
7509 .arg(&ne)
7510 .arg(&qt_g)
7511 .arg(&qt_u)
7512 .arg(&rbg)
7513 .arg(&rbu)
7514 .arg(macros);
7515 unsafe {
7516 b.launch(cfg)?;
7517 }
7518 Ok(act)
7519 }
7520
7521 #[allow(clippy::too_many_arguments)]
7524 pub fn moe_down8_fma_dev(
7525 &self,
7526 table: &CudaSlice<u64>,
7527 sel: &cudarc::driver::CudaView<i32>,
7528 w: &cudarc::driver::CudaView<f32>,
7529 act: &CudaSlice<f32>,
7530 dst: &mut cudarc::driver::CudaViewMut<f32>,
7531 in_f: usize,
7532 out_f: usize,
7533 n_used: usize,
7534 n_expert: usize,
7535 qt: i32,
7536 rb: usize,
7537 ) -> Result<(), Box<dyn std::error::Error>> {
7538 let f = self.func("moe_down8_fma_dev");
7539 let cfg = LaunchConfig {
7540 grid_dim: (out_f as u32, 1, 1),
7541 block_dim: (256, 1, 1),
7542 shared_mem_bytes: 0,
7543 };
7544 let (inf, outf, nu, ne, rbv) = (
7545 in_f as i32,
7546 out_f as i32,
7547 n_used as i32,
7548 n_expert as i32,
7549 rb as i64,
7550 );
7551 let __s_b = self.gpu.stream();
7552 let mut b = __s_b.launch_builder(&f);
7553 b.arg(table)
7554 .arg(sel)
7555 .arg(w)
7556 .arg(act)
7557 .arg(dst)
7558 .arg(&inf)
7559 .arg(&outf)
7560 .arg(&nu)
7561 .arg(&ne)
7562 .arg(&qt)
7563 .arg(&rbv);
7564 unsafe {
7565 b.launch(cfg)?;
7566 }
7567 Ok(())
7568 }
7569
7570 pub fn axpy_into(
7572 &self,
7573 src: &CudaSlice<f32>,
7574 alpha: f32,
7575 dst: &mut cudarc::driver::CudaViewMut<f32>,
7576 n: usize,
7577 ) -> Result<(), Box<dyn std::error::Error>> {
7578 let f = self.func("axpy_f32");
7579 let cfg = LaunchConfig::for_num_elems(n as u32);
7580 let (a, ni) = (alpha, n as i32);
7581 let __s_b = self.gpu.stream();
7582 let mut b = __s_b.launch_builder(&f);
7583 b.arg(src).arg(dst).arg(&a).arg(&ni);
7584 unsafe {
7585 b.launch(cfg)?;
7586 }
7587 Ok(())
7588 }
7589
7590 pub fn axpy_host_into(
7592 &self,
7593 src: &cudarc::driver::CudaView<'_, f32>,
7594 alpha: f32,
7595 dst: &mut cudarc::driver::CudaViewMut<f32>,
7596 n: usize,
7597 ) -> Result<(), Box<dyn std::error::Error>> {
7598 let f = self.func("axpy_host_f32");
7599 let cfg = LaunchConfig::for_num_elems(n as u32);
7600 let (a, ni) = (alpha, n as i32);
7601 let __s_b = self.gpu.stream();
7602 let mut b = __s_b.launch_builder(&f);
7603 b.arg(src).arg(dst).arg(&a).arg(&ni);
7604 unsafe {
7605 b.launch(cfg)?;
7606 }
7607 Ok(())
7608 }
7609
7610 pub fn add_scaled_rows(
7612 &self,
7613 src: &CudaSlice<f32>,
7614 scale: &CudaSlice<f32>,
7615 dst: &mut CudaSlice<f32>,
7616 ncols: usize,
7617 nrows: usize,
7618 ) -> Result<(), Box<dyn std::error::Error>> {
7619 let f = self.func("add_scaled_rows_f32");
7620 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7621 let (nc, nr) = (ncols as i32, nrows as i32);
7622 let __s_b = self.gpu.stream();
7623 let mut b = __s_b.launch_builder(&f);
7624 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7625 unsafe {
7626 b.launch(cfg)?;
7627 }
7628 Ok(())
7629 }
7630
7631 pub fn scale_rows(
7634 &self,
7635 y: &mut CudaSlice<f32>,
7636 s: &CudaSlice<f32>,
7637 ncols: usize,
7638 nrows: usize,
7639 ) -> Result<(), Box<dyn std::error::Error>> {
7640 let f = self.func("scale_rows_f32");
7641 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7642 let (nc, nr) = (ncols as i32, nrows as i32);
7643 let __s_b = self.gpu.stream();
7644 let mut b = __s_b.launch_builder(&f);
7645 b.arg(&mut *y).arg(s).arg(&nc).arg(&nr);
7646 unsafe {
7647 b.launch(cfg)?;
7648 }
7649 Ok(())
7650 }
7651
7652 #[allow(clippy::too_many_arguments)]
7656 pub fn moe_prime_join_scatter(
7657 &self,
7658 y0: &CudaSlice<f32>,
7659 y1: &CudaSlice<f32>,
7660 inv: &CudaSlice<i32>,
7661 w: &CudaSlice<f32>,
7662 out: &mut CudaSlice<f32>,
7663 ncols: usize,
7664 n_used: usize,
7665 t: usize,
7666 ) -> Result<(), Box<dyn std::error::Error>> {
7667 let f = self.func("moe_prime_join_scatter_f32");
7668 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7669 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7670 let __s_b = self.gpu.stream();
7671 let mut b = __s_b.launch_builder(&f);
7672 b.arg(y0)
7673 .arg(y1)
7674 .arg(inv)
7675 .arg(w)
7676 .arg(&mut *out)
7677 .arg(&nc)
7678 .arg(&nu)
7679 .arg(&ti);
7680 unsafe {
7681 b.launch(cfg)?;
7682 }
7683 Ok(())
7684 }
7685
7686 pub fn moe_pairs_weighted_scatter(
7689 &self,
7690 y: &CudaSlice<f32>,
7691 w: &CudaSlice<f32>,
7692 out: &mut CudaSlice<f32>,
7693 ncols: usize,
7694 n_used: usize,
7695 t: usize,
7696 ) -> Result<(), Box<dyn std::error::Error>> {
7697 let f = self.func("moe_pairs_weighted_scatter_f32");
7698 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7699 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7700 let __s_b = self.gpu.stream();
7701 let mut b = __s_b.launch_builder(&f);
7702 b.arg(y).arg(w).arg(&mut *out).arg(&nc).arg(&nu).arg(&ti);
7703 unsafe {
7704 b.launch(cfg)?;
7705 }
7706 Ok(())
7707 }
7708
7709 pub fn gather_rows(
7713 &self,
7714 src: &CudaSlice<f32>,
7715 idx: &CudaSlice<i32>,
7716 dst: &mut CudaSlice<f32>,
7717 ncols: usize,
7718 m_e: usize,
7719 ) -> Result<(), Box<dyn std::error::Error>> {
7720 let f = self.func("gather_rows_f32");
7721 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7722 let (nc, me) = (ncols as i32, m_e as i32);
7723 let __s_b = self.gpu.stream();
7724 let mut b = __s_b.launch_builder(&f);
7725 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7726 unsafe {
7727 b.launch(cfg)?;
7728 }
7729 Ok(())
7730 }
7731
7732 pub fn scatter_slot(
7737 &self,
7738 src: &CudaSlice<f32>,
7739 tok_idx: &CudaSlice<i32>,
7740 slot_idx: &CudaSlice<i32>,
7741 weight: &CudaSlice<f32>,
7742 dst: &mut CudaSlice<f32>,
7743 wbuf: &mut CudaSlice<f32>,
7744 ncols: usize,
7745 n_used: usize,
7746 m_e: usize,
7747 ) -> Result<(), Box<dyn std::error::Error>> {
7748 let f = self.func("scatter_add_slot_f32");
7749 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7750 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7751 let __s_b = self.gpu.stream();
7752 let mut b = __s_b.launch_builder(&f);
7753 b.arg(src)
7754 .arg(tok_idx)
7755 .arg(slot_idx)
7756 .arg(weight)
7757 .arg(dst)
7758 .arg(wbuf)
7759 .arg(&nc)
7760 .arg(&nu)
7761 .arg(&me);
7762 unsafe {
7763 b.launch(cfg)?;
7764 }
7765 Ok(())
7766 }
7767
7768 pub fn reduce_slots(
7772 &self,
7773 slots: &CudaSlice<f32>,
7774 wbuf: &CudaSlice<f32>,
7775 dst: &mut CudaSlice<f32>,
7776 ncols: usize,
7777 n_used: usize,
7778 t: usize,
7779 ) -> Result<(), Box<dyn std::error::Error>> {
7780 let f = self.func("reduce_slots_f32");
7781 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7782 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7783 let __s_b = self.gpu.stream();
7784 let mut b = __s_b.launch_builder(&f);
7785 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7786 unsafe {
7787 b.launch(cfg)?;
7788 }
7789 Ok(())
7790 }
7791
7792 pub fn reduce_slots_host(
7797 &self,
7798 slots: &CudaSlice<f32>,
7799 wbuf: &CudaSlice<f32>,
7800 dst: &mut CudaSlice<f32>,
7801 ncols: usize,
7802 n_used: usize,
7803 t: usize,
7804 ) -> Result<(), Box<dyn std::error::Error>> {
7805 let f = self.func("reduce_slots_host_f32");
7806 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7807 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7808 let __s_b = self.gpu.stream();
7809 let mut b = __s_b.launch_builder(&f);
7810 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7811 unsafe {
7812 b.launch(cfg)?;
7813 }
7814 Ok(())
7815 }
7816
7817 pub fn quantize_q8_1_view(
7824 &self,
7825 x: &cudarc::driver::CudaView<f32>,
7826 m: usize,
7827 in_f: usize,
7828 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7829 let f = self.func("quantize_q8_1");
7830 let nblk = in_f / 32;
7831 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7832 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7833 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7834 let (inf, mi) = (in_f as i32, m as i32);
7835 let __s_b = self.gpu.stream();
7836 let mut b = __s_b.launch_builder(&f);
7837 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7838 unsafe {
7839 b.launch(cfg)?;
7840 }
7841 Ok((q, d))
7842 }
7843
7844 pub fn quantize_q8_1(
7845 &self,
7846 x: &CudaSlice<f32>,
7847 m: usize,
7848 in_f: usize,
7849 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7850 let nblk = in_f / 32;
7851 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);
7855 let (inf, mi) = (in_f as i32, m as i32);
7856 if Self::pdl_on() && Self::pdl_wb_on() {
7857 {
7858 use cudarc::driver::{DevicePtr, DevicePtrMut};
7859 let s = &self.gpu.stream();
7860 let (px, _g0) = x.device_ptr(s);
7861 let (pq, _g1) = q.device_ptr_mut(s);
7862 let (pd, _g2) = d.device_ptr_mut(s);
7863 let mut ps = [
7864 &px as *const _ as *mut std::ffi::c_void,
7865 &pq as *const _ as *mut _,
7866 &pd as *const _ as *mut _,
7867 &inf as *const _ as *mut _,
7868 &mi as *const _ as *mut _,
7869 ];
7870 unsafe {
7871 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7872 }
7873 }
7874 return Ok((q, d));
7875 }
7876 let f = self.func("quantize_q8_1");
7877 let __s_b = self.gpu.stream();
7878 let mut b = __s_b.launch_builder(&f);
7879 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7880 unsafe {
7881 b.launch(cfg)?;
7882 }
7883 Ok((q, d))
7884 }
7885
7886 pub fn quantize_fp4_act(
7890 &self,
7891 x: &CudaSlice<f32>,
7892 m: usize,
7893 in_f: usize,
7894 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7895 let f = self.func("quantize_fp4_act");
7896 let nb16 = in_f / 16;
7897 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);
7900 let (inf, mi) = (in_f as i32, m as i32);
7901 let __s_b = self.gpu.stream();
7902 let mut b = __s_b.launch_builder(&f);
7903 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7904 unsafe {
7905 b.launch(cfg)?;
7906 }
7907 Ok((aq4, ad4))
7908 }
7909
7910 pub fn qmatvec_gemm_nvfp4_fp4(
7915 &self,
7916 bytes: &CudaSlice<u8>,
7917 x: &CudaSlice<f32>,
7918 m: usize,
7919 in_f: usize,
7920 out_f: usize,
7921 row_bytes: usize,
7922 scale: f32,
7923 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7924 assert!(
7925 in_f % 64 == 0,
7926 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7927 );
7928 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7929 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7930 if scale != 1.0 {
7931 self.scale_inplace(&mut y, scale, m * out_f)?;
7932 }
7933 Ok(y)
7934 }
7935
7936 fn fp4_gemm_launch(
7939 &self,
7940 bytes: &CudaSlice<u8>,
7941 aq4: &CudaSlice<u32>,
7942 ad4: &CudaSlice<u8>,
7943 m: usize,
7944 in_f: usize,
7945 out_f: usize,
7946 row_bytes: usize,
7947 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7948 let f = self.func("qmatvec_gemm_nvfp4_fp4");
7949 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
7951 const BN: u32 = 256;
7952 let cfg = LaunchConfig {
7953 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7954 block_dim: (32, 4, 1),
7955 shared_mem_bytes: 0,
7956 };
7957 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7958 let __s_b = self.gpu.stream();
7959 let mut b = __s_b.launch_builder(&f);
7960 b.arg(bytes)
7961 .arg(aq4)
7962 .arg(ad4)
7963 .arg(&mut y)
7964 .arg(&inf)
7965 .arg(&outf)
7966 .arg(&mi)
7967 .arg(&rb);
7968 unsafe {
7969 b.launch(cfg)?;
7970 }
7971 Ok(y)
7972 }
7973
7974 pub fn qmatvec_gemm_nvfp4_fp4_raw(
7976 &self,
7977 bytes: &CudaSlice<u8>,
7978 x: &CudaSlice<f32>,
7979 m: usize,
7980 in_f: usize,
7981 out_f: usize,
7982 row_bytes: usize,
7983 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7984 assert!(
7985 in_f % 64 == 0,
7986 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7987 );
7988 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7989 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7990 }
7991
7992 pub fn qmatvec_q8_0_fast(
7994 &self,
7995 w: &CudaSlice<u8>,
7996 x: &CudaSlice<f32>,
7997 m: usize,
7998 in_f: usize,
7999 out_f: usize,
8000 row_bytes: usize,
8001 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8002 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
8003 let f = self.func("qmatvec_q8_0_dp4a");
8004 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
8006 grid_dim: (out_f as u32, m as u32, 1),
8007 block_dim: (128, 1, 1),
8008 shared_mem_bytes: 0,
8009 };
8010 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8011 let __s_b = self.gpu.stream();
8012 let mut b = __s_b.launch_builder(&f);
8013 b.arg(w)
8014 .arg(&aq)
8015 .arg(&ad)
8016 .arg(&mut y)
8017 .arg(&inf)
8018 .arg(&outf)
8019 .arg(&mi)
8020 .arg(&rb);
8021 unsafe {
8022 b.launch(cfg)?;
8023 }
8024 Ok(y)
8025 }
8026
8027 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
8030 &self,
8031 w: &CudaSlice<u8>,
8032 x: &CudaSlice<f32>,
8033 m: usize,
8034 in_f: usize,
8035 out_f: usize,
8036 row_bytes: usize,
8037 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8038 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
8039 let f = self.func("qmatvec_q4_K_dp4a");
8040 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
8042 grid_dim: (out_f as u32, m as u32, 1),
8043 block_dim: (128, 1, 1),
8044 shared_mem_bytes: 0,
8045 };
8046 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8047 let __s_b = self.gpu.stream();
8048 let mut b = __s_b.launch_builder(&f);
8049 b.arg(w)
8050 .arg(&aq)
8051 .arg(&ad)
8052 .arg(&mut y)
8053 .arg(&inf)
8054 .arg(&outf)
8055 .arg(&mi)
8056 .arg(&rb);
8057 unsafe {
8058 b.launch(cfg)?;
8059 }
8060 Ok(y)
8061 }
8062
8063 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
8066 &self,
8067 w: &CudaSlice<u8>,
8068 x: &CudaSlice<f32>,
8069 m: usize,
8070 in_f: usize,
8071 out_f: usize,
8072 row_bytes: usize,
8073 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8074 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
8075 let f = self.func("qmatvec_q6_K_dp4a");
8076 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
8078 grid_dim: (out_f as u32, m as u32, 1),
8079 block_dim: (128, 1, 1),
8080 shared_mem_bytes: 0,
8081 };
8082 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8083 let __s_b = self.gpu.stream();
8084 let mut b = __s_b.launch_builder(&f);
8085 b.arg(w)
8086 .arg(&aq)
8087 .arg(&ad)
8088 .arg(&mut y)
8089 .arg(&inf)
8090 .arg(&outf)
8091 .arg(&mi)
8092 .arg(&rb);
8093 unsafe {
8094 b.launch(cfg)?;
8095 }
8096 Ok(y)
8097 }
8098
8099 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
8102 &self,
8103 w: &CudaSlice<u8>,
8104 x: &CudaSlice<f32>,
8105 m: usize,
8106 in_f: usize,
8107 out_f: usize,
8108 row_bytes: usize,
8109 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8110 self.qmatvec_dp4a_named(
8111 "qmatvec_q5_K_dp4a",
8112 &w.slice(0..w.len()),
8113 x,
8114 m,
8115 in_f,
8116 out_f,
8117 row_bytes,
8118 )
8119 }
8120 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
8123 &self,
8124 w: &CudaSlice<u8>,
8125 x: &CudaSlice<f32>,
8126 m: usize,
8127 in_f: usize,
8128 out_f: usize,
8129 row_bytes: usize,
8130 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8131 self.qmatvec_dp4a_named(
8132 "qmatvec_q3_K_dp4a",
8133 &w.slice(0..w.len()),
8134 x,
8135 m,
8136 in_f,
8137 out_f,
8138 row_bytes,
8139 )
8140 }
8141 pub fn qmatvec_nvfp4_fast_rp(
8143 &self,
8144 w: &CudaSlice<u8>,
8145 x: &CudaSlice<f32>,
8146 m: usize,
8147 in_f: usize,
8148 out_f: usize,
8149 row_bytes: usize,
8150 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8151 assert!(
8152 in_f % 64 == 0,
8153 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8154 );
8155 self.qmatvec_dp4a_named(
8156 "qmatvec_nvfp4_dp4a_rp",
8157 &w.slice(0..w.len()),
8158 x,
8159 m,
8160 in_f,
8161 out_f,
8162 row_bytes,
8163 )
8164 }
8165 pub fn qmatvec_nvfp4_fast(
8167 &self,
8168 w: &cudarc::driver::CudaView<'_, u8>,
8169 x: &CudaSlice<f32>,
8170 m: usize,
8171 in_f: usize,
8172 out_f: usize,
8173 row_bytes: usize,
8174 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8175 assert!(
8178 in_f % 64 == 0,
8179 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8180 );
8181 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
8182 }
8183 pub fn qmatvec_nvfp4_fast_v2(
8188 &self,
8189 w: &cudarc::driver::CudaView<'_, u8>,
8190 x: &CudaSlice<f32>,
8191 m: usize,
8192 in_f: usize,
8193 out_f: usize,
8194 row_bytes: usize,
8195 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8196 assert!(
8197 in_f % 64 == 0,
8198 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8199 );
8200 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
8201 }
8202 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
8205 &self,
8206 w: &CudaSlice<u8>,
8207 x: &CudaSlice<f32>,
8208 m: usize,
8209 in_f: usize,
8210 out_f: usize,
8211 row_bytes: usize,
8212 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8213 self.qmatvec_dp4a_named(
8214 "qmatvec_iq4_XS_dp4a",
8215 &w.slice(0..w.len()),
8216 x,
8217 m,
8218 in_f,
8219 out_f,
8220 row_bytes,
8221 )
8222 }
8223
8224 fn qmatvec_dp4a_named(
8226 &self,
8227 name: &str,
8228 w: &cudarc::driver::CudaView<'_, u8>,
8229 x: &CudaSlice<f32>,
8230 m: usize,
8231 in_f: usize,
8232 out_f: usize,
8233 row_bytes: usize,
8234 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8235 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
8236 let f = self.func(name);
8237 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
8239 grid_dim: (out_f as u32, m as u32, 1),
8240 block_dim: (128, 1, 1),
8241 shared_mem_bytes: 0,
8242 };
8243 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8244 let __s_b = self.gpu.stream();
8245 let mut b = __s_b.launch_builder(&f);
8246 b.arg(w)
8247 .arg(&aq)
8248 .arg(&ad)
8249 .arg(&mut y)
8250 .arg(&inf)
8251 .arg(&outf)
8252 .arg(&mi)
8253 .arg(&rb);
8254 unsafe {
8255 b.launch(cfg)?;
8256 }
8257 Ok(y)
8258 }
8259
8260 #[allow(clippy::too_many_arguments)]
8266 pub fn qmatvec_nvfp4_fast_prequant_into(
8267 &self,
8268 w: &CudaSlice<u8>,
8269 aq: &CudaSlice<i8>,
8270 ad: &CudaSlice<f32>,
8271 y: &mut CudaSlice<f32>,
8272 m: usize,
8273 in_f: usize,
8274 out_f: usize,
8275 row_bytes: usize,
8276 ) -> Result<(), Box<dyn std::error::Error>> {
8277 assert!(
8278 in_f % 64 == 0,
8279 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8280 );
8281 if y.len() < m * out_f {
8282 return Err(format!(
8283 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
8284 y.len()
8285 )
8286 .into());
8287 }
8288 let f = self.func("qmatvec_nvfp4_dp4a");
8289 let cfg = LaunchConfig {
8290 grid_dim: (out_f as u32, m as u32, 1),
8291 block_dim: (128, 1, 1),
8292 shared_mem_bytes: 0,
8293 };
8294 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8295 let __s_b = self.gpu.stream();
8296 let mut b = __s_b.launch_builder(&f);
8297 b.arg(w)
8298 .arg(aq)
8299 .arg(ad)
8300 .arg(y)
8301 .arg(&inf)
8302 .arg(&outf)
8303 .arg(&mi)
8304 .arg(&rb);
8305 unsafe {
8306 b.launch(cfg)?;
8307 }
8308 Ok(())
8309 }
8310
8311 #[allow(clippy::too_many_arguments)]
8314 pub fn matvec_f32_qkv_into(
8315 &self,
8316 wq: &CudaSlice<f32>,
8317 wk: &CudaSlice<f32>,
8318 wv: &CudaSlice<f32>,
8319 wg: &CudaSlice<f32>,
8320 x: &CudaSlice<f32>,
8321 yq: &mut CudaSlice<f32>,
8322 yk: &mut CudaSlice<f32>,
8323 yv: &mut CudaSlice<f32>,
8324 yg: &mut CudaSlice<f32>,
8325 in_f: usize,
8326 out_q: usize,
8327 out_kv: usize,
8328 out_g: usize,
8329 ) -> Result<(), Box<dyn std::error::Error>> {
8330 if in_f % 4 != 0
8331 || wq.len() != out_q * in_f
8332 || wk.len() != out_kv * in_f
8333 || wv.len() != out_kv * in_f
8334 || wg.len() < out_g * in_f
8335 || x.len() < in_f
8336 || yq.len() < out_q
8337 || yk.len() < out_kv
8338 || yv.len() < out_kv
8339 || (out_g > 0 && yg.len() < out_g)
8340 {
8341 return Err(format!(
8342 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
8343 wq={} wk={} wv={} wg={}",
8344 wq.len(),
8345 wk.len(),
8346 wv.len(),
8347 wg.len()
8348 )
8349 .into());
8350 }
8351 let f = self.func("matvec_f32_qkv");
8352 let cfg = LaunchConfig {
8353 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
8354 block_dim: (128, 1, 1),
8355 shared_mem_bytes: 0,
8356 };
8357 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
8358 let __s_b = self.gpu.stream();
8359 let mut b = __s_b.launch_builder(&f);
8360 b.arg(wq)
8361 .arg(wk)
8362 .arg(wv)
8363 .arg(wg)
8364 .arg(x)
8365 .arg(yq)
8366 .arg(yk)
8367 .arg(yv)
8368 .arg(yg)
8369 .arg(&inf)
8370 .arg(&oq)
8371 .arg(&okv)
8372 .arg(&og);
8373 unsafe {
8374 b.launch(cfg)?;
8375 }
8376 Ok(())
8377 }
8378
8379 #[allow(clippy::too_many_arguments)]
8382 pub fn qmatvec_nvfp4_sel_gu_ep_into(
8383 &self,
8384 gate_bank: &CudaSlice<u8>,
8385 up_bank: &CudaSlice<u8>,
8386 sel: &CudaSlice<i32>,
8387 aq: &CudaSlice<i8>,
8388 ad: &CudaSlice<f32>,
8389 yg: &mut CudaSlice<f32>,
8390 yu: &mut CudaSlice<f32>,
8391 n_sel: usize,
8392 in_f: usize,
8393 out_f: usize,
8394 row_bytes: usize,
8395 expert_stride: usize,
8396 owner: usize,
8397 ) -> Result<(), Box<dyn std::error::Error>> {
8398 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8399 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8400 return Err("NVFP4 gu ep geometry".into());
8401 }
8402 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8403 let cfg = LaunchConfig {
8404 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8405 block_dim: (128, 1, 1),
8406 shared_mem_bytes: 0,
8407 };
8408 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8409 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8410 let (ars, adrs) = (0i64, 0i64);
8411 let __s_b = self.gpu.stream();
8412 let mut b = __s_b.launch_builder(&f);
8413 b.arg(gate_bank)
8414 .arg(up_bank)
8415 .arg(sel)
8416 .arg(aq)
8417 .arg(ad)
8418 .arg(yg)
8419 .arg(yu)
8420 .arg(&inf)
8421 .arg(&outf)
8422 .arg(&ns)
8423 .arg(&rb)
8424 .arg(&es)
8425 .arg(&ars)
8426 .arg(&adrs)
8427 .arg(&own);
8428 unsafe {
8429 b.launch(cfg)?;
8430 }
8431 Ok(())
8432 }
8433
8434 #[allow(clippy::too_many_arguments)]
8436 pub fn silu_mul_scaled_q8_1_sel_ep_into(
8437 &self,
8438 gate: &CudaSlice<f32>,
8439 up: &CudaSlice<f32>,
8440 gmac: &CudaSlice<f32>,
8441 umac: &CudaSlice<f32>,
8442 sel: &CudaSlice<i32>,
8443 limit: Option<f32>,
8444 out_q: &mut CudaSlice<i8>,
8445 out_d: &mut CudaSlice<f32>,
8446 n_per: usize,
8447 n_sel: usize,
8448 owner: usize,
8449 ) -> Result<(), Box<dyn std::error::Error>> {
8450 if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8451 return Err("NVFP4 silu ep geometry".into());
8452 }
8453 let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8454 let warps = n_sel * n_per / 32;
8455 let cfg = LaunchConfig {
8456 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8457 block_dim: (128, 1, 1),
8458 shared_mem_bytes: 0,
8459 };
8460 let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8461 let (lim, has) = match limit {
8462 Some(l) => (l, 1i32),
8463 None => (0.0f32, 0i32),
8464 };
8465 let __s_b = self.gpu.stream();
8466 let mut b = __s_b.launch_builder(&f);
8467 b.arg(gate)
8468 .arg(up)
8469 .arg(gmac)
8470 .arg(umac)
8471 .arg(sel)
8472 .arg(&lim)
8473 .arg(&has)
8474 .arg(out_q)
8475 .arg(out_d)
8476 .arg(&np)
8477 .arg(&ns)
8478 .arg(&own);
8479 unsafe {
8480 b.launch(cfg)?;
8481 }
8482 Ok(())
8483 }
8484
8485 #[allow(clippy::too_many_arguments)]
8487 pub fn qmatvec_nvfp4_sel_down8_ep_into(
8488 &self,
8489 bank: &CudaSlice<u8>,
8490 sel: &CudaSlice<i32>,
8491 aq: &CudaSlice<i8>,
8492 ad: &CudaSlice<f32>,
8493 route_w: &CudaSlice<f32>,
8494 md: &CudaSlice<f32>,
8495 dst: &mut CudaSlice<f32>,
8496 n_sel: usize,
8497 in_f: usize,
8498 out_f: usize,
8499 row_bytes: usize,
8500 expert_stride: usize,
8501 act_row_stride: usize,
8502 ad_row_stride: usize,
8503 owner: usize,
8504 ) -> Result<(), Box<dyn std::error::Error>> {
8505 if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8506 return Err("NVFP4 down8 ep geometry".into());
8507 }
8508 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8509 let cfg = LaunchConfig {
8510 grid_dim: (out_f as u32, 1, 1),
8511 block_dim: (32, n_sel as u32, 1),
8512 shared_mem_bytes: 0,
8513 };
8514 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8515 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8516 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8517 let __s_b = self.gpu.stream();
8518 let mut b = __s_b.launch_builder(&f);
8519 b.arg(bank)
8520 .arg(sel)
8521 .arg(aq)
8522 .arg(ad)
8523 .arg(route_w)
8524 .arg(md)
8525 .arg(dst)
8526 .arg(&inf)
8527 .arg(&outf)
8528 .arg(&ns)
8529 .arg(&rb)
8530 .arg(&es)
8531 .arg(&ars)
8532 .arg(&adrs)
8533 .arg(&own);
8534 unsafe {
8535 b.launch(cfg)?;
8536 }
8537 Ok(())
8538 }
8539
8540 #[allow(clippy::too_many_arguments)]
8546 pub fn qmatvec_nvfp4_sel_into(
8547 &self,
8548 bank: &CudaSlice<u8>,
8549 sel: &CudaSlice<i32>,
8550 aq: &CudaSlice<i8>,
8551 ad: &CudaSlice<f32>,
8552 y: &mut CudaSlice<f32>,
8553 n_sel: usize,
8554 in_f: usize,
8555 out_f: usize,
8556 row_bytes: usize,
8557 expert_stride: usize,
8558 act_row_stride: usize,
8559 ad_row_stride: usize,
8560 ) -> Result<(), Box<dyn std::error::Error>> {
8561 assert!(
8562 in_f % 64 == 0,
8563 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8564 );
8565 if y.len() < n_sel * out_f || sel.len() < n_sel {
8566 return Err(format!(
8567 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8568 y.len(),
8569 sel.len()
8570 )
8571 .into());
8572 }
8573 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8580 let mode = *MR.get_or_init(|| {
8581 if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8582 2
8583 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8584 1
8585 } else {
8586 0
8587 }
8588 });
8589 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8590 let f = match mode {
8591 2 => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8592 1 => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8593 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8594 };
8595 let nsb = in_f >> 5;
8600 let fit_block: u32 = if mode == 0 && nsb <= 32 {
8601 32
8602 } else if mode == 1 {
8603 512
8604 } else {
8605 128
8606 };
8607 let cfg = LaunchConfig {
8608 grid_dim: (
8609 match mode {
8610 2 => (out_f as u32).div_ceil(16),
8611 1 => (out_f as u32).div_ceil(4),
8612 _ => out_f as u32,
8613 },
8614 n_sel as u32,
8615 1,
8616 ),
8617 block_dim: (fit_block, 1, 1),
8618 shared_mem_bytes: 0,
8619 };
8620 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8621 let (rb, es, ars, adrs) = (
8622 row_bytes as i64,
8623 expert_stride as i64,
8624 act_row_stride as i64,
8625 ad_row_stride as i64,
8626 );
8627 let __s_b = self.gpu.stream();
8628 let mut b = __s_b.launch_builder(&f);
8629 b.arg(bank)
8630 .arg(sel)
8631 .arg(aq)
8632 .arg(ad)
8633 .arg(y)
8634 .arg(&inf)
8635 .arg(&outf)
8636 .arg(&ns)
8637 .arg(&rb)
8638 .arg(&es)
8639 .arg(&ars)
8640 .arg(&adrs);
8641 unsafe {
8642 b.launch(cfg)?;
8643 }
8644 Ok(())
8645 }
8646
8647 #[allow(clippy::too_many_arguments)]
8652 pub fn silu_mul_scaled_q8_1_sel_into(
8653 &self,
8654 gate: &CudaSlice<f32>,
8655 up: &CudaSlice<f32>,
8656 gmac: &CudaSlice<f32>,
8657 umac: &CudaSlice<f32>,
8658 sel: &CudaSlice<i32>,
8659 limit: Option<f32>,
8660 out_q: &mut CudaSlice<i8>,
8661 out_d: &mut CudaSlice<f32>,
8662 n_per: usize,
8663 n_sel: usize,
8664 ) -> Result<(), Box<dyn std::error::Error>> {
8665 let n = n_per * n_sel;
8666 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8667 return Err(format!(
8668 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8669 out_q.len(),
8670 out_d.len()
8671 )
8672 .into());
8673 }
8674 if let Some(limit) = limit {
8675 if limit <= 1e-6 {
8676 return Err(format!(
8677 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8678 )
8679 .into());
8680 }
8681 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8682 let cfg = LaunchConfig::for_num_elems(n as u32);
8683 let (np, ns) = (n_per as i32, n_sel as i32);
8684 let __s_b = self.gpu.stream();
8685 let mut b = __s_b.launch_builder(&f);
8686 b.arg(gate)
8687 .arg(up)
8688 .arg(gmac)
8689 .arg(umac)
8690 .arg(sel)
8691 .arg(&limit)
8692 .arg(out_q)
8693 .arg(out_d)
8694 .arg(&np)
8695 .arg(&ns);
8696 unsafe {
8697 b.launch(cfg)?;
8698 }
8699 return Ok(());
8700 }
8701 let f = self.func("silu_mul_scaled_q8_1_sel");
8702 let cfg = LaunchConfig::for_num_elems(n as u32);
8703 let (np, ns) = (n_per as i32, n_sel as i32);
8704 let __s_b = self.gpu.stream();
8705 let mut b = __s_b.launch_builder(&f);
8706 b.arg(gate)
8707 .arg(up)
8708 .arg(gmac)
8709 .arg(umac)
8710 .arg(sel)
8711 .arg(out_q)
8712 .arg(out_d)
8713 .arg(&np)
8714 .arg(&ns);
8715 unsafe {
8716 b.launch(cfg)?;
8717 }
8718 Ok(())
8719 }
8720
8721 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8722 Ok(self.gpu.stream().clone_htod(v)?)
8723 }
8724 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8725 Ok(self.gpu.stream().clone_htod(v)?)
8726 }
8727 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8729 Ok(self.gpu.stream().clone_htod(v)?)
8730 }
8731 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8732 Ok(self.gpu.stream().clone_htod(v)?)
8733 }
8734 pub fn dtoh_view(
8736 &self,
8737 d: &cudarc::driver::CudaView<f32>,
8738 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8739 let v = self.gpu.stream().clone_dtoh(d)?;
8740 self.gpu.stream().synchronize()?;
8741 Ok(v)
8742 }
8743 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8744 let v = self.gpu.stream().clone_dtoh(d)?;
8745 self.gpu.stream().synchronize()?;
8746 Ok(v)
8747 }
8748 pub fn dtoh_pair(
8752 &self,
8753 a: &CudaSlice<f32>,
8754 b: &CudaSlice<f32>,
8755 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8756 let av = self.gpu.stream().clone_dtoh(a)?;
8757 let bv = self.gpu.stream().clone_dtoh(b)?;
8758 self.gpu.stream().synchronize()?;
8759 Ok((av, bv))
8760 }
8761 pub fn dtoh_pair_views(
8764 &self,
8765 a: &cudarc::driver::CudaView<f32>,
8766 b: &cudarc::driver::CudaView<f32>,
8767 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8768 let av = self.gpu.stream().clone_dtoh(a)?;
8769 let bv = self.gpu.stream().clone_dtoh(b)?;
8770 self.gpu.stream().synchronize()?;
8771 Ok((av, bv))
8772 }
8773 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8775 let v = self.gpu.stream().clone_dtoh(d)?;
8776 self.gpu.stream().synchronize()?;
8777 Ok(v)
8778 }
8779 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8781 let v = self.gpu.stream().clone_dtoh(d)?;
8782 self.gpu.stream().synchronize()?;
8783 Ok(v)
8784 }
8785 pub fn dtoh_u8_view(
8786 &self,
8787 d: &cudarc::driver::CudaView<u8>,
8788 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8789 let v = self.gpu.stream().clone_dtoh(d)?;
8790 self.gpu.stream().synchronize()?;
8791 Ok(v)
8792 }
8793 pub fn dtoh_u8_into_pinned(
8801 &self,
8802 d: &CudaSlice<u8>,
8803 dst: &mut PinnedHostBuf,
8804 n: usize,
8805 ) -> Result<(), Box<dyn std::error::Error>> {
8806 if n > d.len() || n > dst.len() {
8807 return Err(format!(
8808 "dtoh_u8_into_pinned range {n} exceeds src {} or pinned dst {}",
8809 d.len(),
8810 dst.len(),
8811 )
8812 .into());
8813 }
8814 if n == 0 {
8815 return Ok(());
8816 }
8817 let host = &mut dst.as_mut_slice()[..n];
8818 self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
8819 self.gpu.stream().synchronize()?;
8820 Ok(())
8821 }
8822 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8823 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8824 self.keep_if_capturing(&s);
8825 Ok(s)
8826 }
8827
8828 pub fn prob_of_token_device(
8837 &self,
8838 logits: &CudaSlice<f32>,
8839 tok: &CudaSlice<u32>,
8840 n_vocab: usize,
8841 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8842 let nb = ARGMAX_NB;
8843 let mut part = self.alloc_uninit::<f32>(nb)?;
8844 let mut p = self.alloc_uninit::<f32>(1)?;
8845 let f1 = self.func("prob_of_token_partial_f32");
8846 let cfg1 = LaunchConfig {
8847 grid_dim: (nb as u32, 1, 1),
8848 block_dim: (256, 1, 1),
8849 shared_mem_bytes: 0,
8850 };
8851 let nv = n_vocab as i32;
8852 let __s_b1 = self.gpu.stream();
8853 let mut b1 = __s_b1.launch_builder(&f1);
8854 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8855 unsafe {
8856 b1.launch(cfg1)?;
8857 }
8858 let f2 = self.func("prob_of_token_final_f32");
8859 let cfg2 = LaunchConfig {
8860 grid_dim: (1, 1, 1),
8861 block_dim: (256, 1, 1),
8862 shared_mem_bytes: 0,
8863 };
8864 let nbi = nb as i32;
8865 let __s_b2 = self.gpu.stream();
8866 let mut b2 = __s_b2.launch_builder(&f2);
8867 b2.arg(&part).arg(&mut p).arg(&nbi);
8868 unsafe {
8869 b2.launch(cfg2)?;
8870 }
8871 Ok(p)
8872 }
8873
8874 pub fn prob_of_token_device_col(
8881 &self,
8882 logits: &CudaSlice<f32>,
8883 tok_all: &CudaSlice<u32>,
8884 tok_idx: usize,
8885 p_out: &mut CudaSlice<f32>,
8886 p_idx: usize,
8887 n_vocab: usize,
8888 ) -> Result<(), Box<dyn std::error::Error>> {
8889 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8890 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8891 let nb = ARGMAX_NB;
8892 let mut part = self.alloc_uninit::<f32>(nb)?;
8893 let f1 = self.func("prob_of_token_partial_f32");
8894 let cfg1 = LaunchConfig {
8895 grid_dim: (nb as u32, 1, 1),
8896 block_dim: (256, 1, 1),
8897 shared_mem_bytes: 0,
8898 };
8899 let nv = n_vocab as i32;
8900 let __s_b1 = self.gpu.stream();
8901 let mut b1 = __s_b1.launch_builder(&f1);
8902 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8903 unsafe {
8904 b1.launch(cfg1)?;
8905 }
8906 let f2 = self.func("prob_of_token_final_f32");
8907 let cfg2 = LaunchConfig {
8908 grid_dim: (1, 1, 1),
8909 block_dim: (256, 1, 1),
8910 shared_mem_bytes: 0,
8911 };
8912 let nbi = nb as i32;
8913 let __s_b2 = self.gpu.stream();
8914 let mut b2 = __s_b2.launch_builder(&f2);
8915 b2.arg(&part).arg(&mut p_v).arg(&nbi);
8916 unsafe {
8917 b2.launch(cfg2)?;
8918 }
8919 Ok(())
8920 }
8921
8922 pub fn prob_of_token_device_into(
8923 &self,
8924 logits: &CudaSlice<f32>,
8925 tok: &CudaSlice<u32>,
8926 p_out: &mut CudaSlice<f32>,
8927 n_vocab: usize,
8928 ) -> Result<(), Box<dyn std::error::Error>> {
8929 let nb = ARGMAX_NB;
8930 let mut part = self.alloc_uninit::<f32>(nb)?;
8931 let f1 = self.func("prob_of_token_partial_f32");
8932 let cfg1 = LaunchConfig {
8933 grid_dim: (nb as u32, 1, 1),
8934 block_dim: (256, 1, 1),
8935 shared_mem_bytes: 0,
8936 };
8937 let nv = n_vocab as i32;
8938 let __s_b1 = self.gpu.stream();
8939 let mut b1 = __s_b1.launch_builder(&f1);
8940 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8941 unsafe {
8942 b1.launch(cfg1)?;
8943 }
8944 let f2 = self.func("prob_of_token_final_f32");
8945 let cfg2 = LaunchConfig {
8946 grid_dim: (1, 1, 1),
8947 block_dim: (256, 1, 1),
8948 shared_mem_bytes: 0,
8949 };
8950 let nbi = nb as i32;
8951 let __s_b2 = self.gpu.stream();
8952 let mut b2 = __s_b2.launch_builder(&f2);
8953 b2.arg(&part).arg(p_out).arg(&nbi);
8954 unsafe {
8955 b2.launch(cfg2)?;
8956 }
8957 Ok(())
8958 }
8959
8960 pub fn u32_hist_append(
8963 &self,
8964 tok: &CudaSlice<u32>,
8965 hist: &mut CudaSlice<u32>,
8966 idx: &mut CudaSlice<i32>,
8967 ) -> Result<(), Box<dyn std::error::Error>> {
8968 let f = self.func("u32_hist_append");
8969 let cfg = LaunchConfig {
8970 grid_dim: (1, 1, 1),
8971 block_dim: (32, 1, 1),
8972 shared_mem_bytes: 0,
8973 };
8974 let __s_b = self.gpu.stream();
8975 let mut b = __s_b.launch_builder(&f);
8976 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8977 unsafe {
8978 b.launch(cfg)?;
8979 }
8980 Ok(())
8981 }
8982
8983 pub fn argmax_token_device(
8984 &self,
8985 logits: &CudaSlice<f32>,
8986 n_vocab: usize,
8987 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8988 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8989 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8990 Ok(tok)
8991 }
8992 pub fn argmax_token_device_into(
8999 &self,
9000 logits: &CudaSlice<f32>,
9001 tok: &mut CudaSlice<u32>,
9002 n_vocab: usize,
9003 ) -> Result<(), Box<dyn std::error::Error>> {
9004 let nb = ARGMAX_NB;
9005 let f1 = self.func("argmax_partial_f32");
9006 let f2 = self.func("argmax_final_f32");
9007 let mut guard = self.argmax_partials.lock().unwrap();
9008 if guard.is_none() {
9009 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
9012 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
9013 *guard = Some((pv, pi));
9014 }
9015 let (part_v, part_i) = guard.as_mut().unwrap();
9016 let nv = n_vocab as i32;
9017 let nbi = nb as i32;
9018 let cfg1 = LaunchConfig {
9020 grid_dim: (nb as u32, 1, 1),
9021 block_dim: (256, 1, 1),
9022 shared_mem_bytes: 0,
9023 };
9024 let __s_b1 = self.gpu.stream();
9025 let mut b1 = __s_b1.launch_builder(&f1);
9026 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
9027 unsafe {
9028 b1.launch(cfg1)?;
9029 }
9030 let cfg2 = LaunchConfig {
9032 grid_dim: (1, 1, 1),
9033 block_dim: (256, 1, 1),
9034 shared_mem_bytes: 0,
9035 };
9036 let __s_b2 = self.gpu.stream();
9037 let mut b2 = __s_b2.launch_builder(&f2);
9038 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
9039 unsafe {
9040 b2.launch(cfg2)?;
9041 }
9042 Ok(())
9043 }
9044 pub fn argmax_token_device_col(
9050 &self,
9051 logits: &CudaSlice<f32>,
9052 col: usize,
9053 n_vocab: usize,
9054 toks: &mut CudaSlice<u32>,
9055 out_idx: usize,
9056 ) -> Result<(), Box<dyn std::error::Error>> {
9057 let nb = ARGMAX_NB;
9058 let f1 = self.func("argmax_partial_f32");
9059 let f2 = self.func("argmax_final_f32");
9060 let mut guard = self.argmax_partials.lock().unwrap();
9061 if guard.is_none() {
9062 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
9063 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
9064 *guard = Some((pv, pi));
9065 }
9066 let (part_v, part_i) = guard.as_mut().unwrap();
9067 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
9068 let nv = n_vocab as i32;
9069 let nbi = nb as i32;
9070 let cfg1 = LaunchConfig {
9071 grid_dim: (nb as u32, 1, 1),
9072 block_dim: (256, 1, 1),
9073 shared_mem_bytes: 0,
9074 };
9075 let __s_b1 = self.gpu.stream();
9076 let mut b1 = __s_b1.launch_builder(&f1);
9077 b1.arg(&col_view)
9078 .arg(&mut *part_v)
9079 .arg(&mut *part_i)
9080 .arg(&nv);
9081 unsafe {
9082 b1.launch(cfg1)?;
9083 }
9084 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
9085 let cfg2 = LaunchConfig {
9086 grid_dim: (1, 1, 1),
9087 block_dim: (256, 1, 1),
9088 shared_mem_bytes: 0,
9089 };
9090 let __s_b2 = self.gpu.stream();
9091 let mut b2 = __s_b2.launch_builder(&f2);
9092 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
9093 unsafe {
9094 b2.launch(cfg2)?;
9095 }
9096 Ok(())
9097 }
9098 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
9100 Ok(self.gpu.stream().clone_htod(v)?)
9101 }
9102 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
9103 let v = self.gpu.stream().clone_dtoh(d)?;
9104 self.gpu.stream().synchronize()?;
9105 Ok(v)
9106 }
9107
9108 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
9109 let v = self.gpu.stream().clone_dtoh(d)?;
9110 self.gpu.stream().synchronize()?;
9111 Ok(v)
9112 }
9113 pub fn htod_u32_into(
9117 &self,
9118 dst: &mut CudaSlice<u32>,
9119 src: &[u32],
9120 ) -> Result<(), Box<dyn std::error::Error>> {
9121 let mut view = dst.slice_mut(0..src.len());
9122 self.gpu.stream().memcpy_htod(src, &mut view)?;
9123 Ok(())
9124 }
9125
9126 pub fn htod_i32_into(
9129 &self,
9130 dst: &mut CudaSlice<i32>,
9131 src: &[i32],
9132 ) -> Result<(), Box<dyn std::error::Error>> {
9133 let mut view = dst.slice_mut(0..src.len());
9134 self.gpu.stream().memcpy_htod(src, &mut view)?;
9135 Ok(())
9136 }
9137
9138 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
9139 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
9140 self.keep_if_capturing(&s);
9141 Ok(s)
9142 }
9143 pub fn embed_gather_device_into(
9146 &self,
9147 embd: &CudaSlice<u8>,
9148 token_d: &CudaSlice<u32>,
9149 x_out: &mut CudaSlice<f32>,
9150 n_embd: usize,
9151 qtype: i32,
9152 row_bytes: usize,
9153 ) -> Result<(), Box<dyn std::error::Error>> {
9154 let f = self.func("embed_gather_u32");
9155 let cfg = LaunchConfig {
9156 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9157 block_dim: (256, 1, 1),
9158 shared_mem_bytes: 0,
9159 };
9160 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9161 let __s_b = self.gpu.stream();
9162 let mut b = __s_b.launch_builder(&f);
9163 b.arg(embd)
9164 .arg(token_d)
9165 .arg(x_out)
9166 .arg(&ne)
9167 .arg(&qt)
9168 .arg(&rb);
9169 unsafe {
9170 b.launch(cfg)?;
9171 }
9172 Ok(())
9173 }
9174 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
9176 let v = self.gpu.stream().clone_dtoh(d)?;
9177 self.gpu.stream().synchronize()?;
9178 Ok(v[0])
9179 }
9180 pub fn i32_set_k(
9187 &self,
9188 dst: &mut CudaSlice<i32>,
9189 v: i32,
9190 ) -> Result<(), Box<dyn std::error::Error>> {
9191 let f = self.func("i32_set_k");
9192 let cfg = LaunchConfig {
9193 grid_dim: (1, 1, 1),
9194 block_dim: (1, 1, 1),
9195 shared_mem_bytes: 0,
9196 };
9197 let idx = 0i32;
9198 let __s_b = self.gpu.stream();
9199 let mut b = __s_b.launch_builder(&f);
9200 b.arg(dst).arg(&v).arg(&idx);
9201 unsafe {
9202 b.launch(cfg)?;
9203 }
9204 Ok(())
9205 }
9206
9207 pub fn set_i32_one(
9208 &self,
9209 d: &mut CudaSlice<i32>,
9210 v: i32,
9211 ) -> Result<(), Box<dyn std::error::Error>> {
9212 self.gpu.stream().memcpy_htod(&[v], d)?;
9213 Ok(())
9214 }
9215 pub fn set_u32_one(
9218 &self,
9219 d: &mut CudaSlice<u32>,
9220 v: u32,
9221 ) -> Result<(), Box<dyn std::error::Error>> {
9222 self.gpu.stream().memcpy_htod(&[v], d)?;
9223 Ok(())
9224 }
9225 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
9227 let v = self.gpu.stream().clone_dtoh(d)?;
9228 self.gpu.stream().synchronize()?;
9229 Ok(v[0])
9230 }
9231 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9233 Ok(self.gpu.stream().clone_htod(bytes)?)
9234 }
9235 pub fn embed_gather_device(
9239 &self,
9240 embd: &CudaSlice<u8>,
9241 token_d: &CudaSlice<u32>,
9242 n_embd: usize,
9243 qtype: i32,
9244 row_bytes: usize,
9245 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9246 let f = self.func("embed_gather_u32");
9247 let mut x = self.alloc_uninit::<f32>(n_embd)?;
9248 let cfg = LaunchConfig {
9249 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9250 block_dim: (256, 1, 1),
9251 shared_mem_bytes: 0,
9252 };
9253 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9254 let __s_b = self.gpu.stream();
9255 let mut b = __s_b.launch_builder(&f);
9256 b.arg(embd)
9257 .arg(token_d)
9258 .arg(&mut x)
9259 .arg(&ne)
9260 .arg(&qt)
9261 .arg(&rb);
9262 unsafe {
9263 b.launch(cfg)?;
9264 }
9265 Ok(x)
9266 }
9267
9268 pub fn embed_gather_device_t(
9272 &self,
9273 embd: &CudaSlice<u8>,
9274 tokens: &[u32],
9275 n_embd: usize,
9276 qtype: i32,
9277 row_bytes: usize,
9278 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9279 let t = tokens.len();
9280 let tok_d = self.gpu.stream().clone_htod(tokens)?;
9281 let f = self.func("embed_gather_u32_t");
9282 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9283 let cfg = LaunchConfig {
9284 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9285 block_dim: (256, 1, 1),
9286 shared_mem_bytes: 0,
9287 };
9288 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9289 let __s_b = self.gpu.stream();
9290 let mut b = __s_b.launch_builder(&f);
9291 b.arg(embd)
9292 .arg(&tok_d)
9293 .arg(&mut x)
9294 .arg(&ne)
9295 .arg(&qt)
9296 .arg(&rb)
9297 .arg(&ti);
9298 unsafe {
9299 b.launch(cfg)?;
9300 }
9301 Ok(x)
9302 }
9303
9304 pub fn embed_gather_device_tv(
9309 &self,
9310 embd: &CudaSlice<u8>,
9311 tok_v: &cudarc::driver::CudaView<u32>,
9312 t: usize,
9313 n_embd: usize,
9314 qtype: i32,
9315 row_bytes: usize,
9316 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9317 let f = self.func("embed_gather_u32_t");
9318 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9319 let cfg = LaunchConfig {
9320 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9321 block_dim: (256, 1, 1),
9322 shared_mem_bytes: 0,
9323 };
9324 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9325 let __s_b = self.gpu.stream();
9326 let mut b = __s_b.launch_builder(&f);
9327 b.arg(embd)
9328 .arg(tok_v)
9329 .arg(&mut x)
9330 .arg(&ne)
9331 .arg(&qt)
9332 .arg(&rb)
9333 .arg(&ti);
9334 unsafe {
9335 b.launch(cfg)?;
9336 }
9337 Ok(x)
9338 }
9339
9340 pub fn embed_gather_device_td(
9341 &self,
9342 embd: &CudaSlice<u8>,
9343 tok_d: &CudaSlice<u32>,
9344 t: usize,
9345 n_embd: usize,
9346 qtype: i32,
9347 row_bytes: usize,
9348 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9349 let f = self.func("embed_gather_u32_t");
9350 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9351 let cfg = LaunchConfig {
9352 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9353 block_dim: (256, 1, 1),
9354 shared_mem_bytes: 0,
9355 };
9356 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9357 let __s_b = self.gpu.stream();
9358 let mut b = __s_b.launch_builder(&f);
9359 b.arg(embd)
9360 .arg(tok_d)
9361 .arg(&mut x)
9362 .arg(&ne)
9363 .arg(&qt)
9364 .arg(&rb)
9365 .arg(&ti);
9366 unsafe {
9367 b.launch(cfg)?;
9368 }
9369 Ok(x)
9370 }
9371
9372 #[inline]
9378 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
9380 if self
9381 .capture_keep_on
9382 .load(std::sync::atomic::Ordering::Relaxed)
9383 {
9384 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
9385 }
9386 }
9387
9388 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9389 &self,
9390 n: usize,
9391 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9392 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9393 {
9397 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9398 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9399 use cudarc::driver::DevicePtrMut;
9401 let n_bytes = s.len() * std::mem::size_of::<T>();
9402 let stream = self.gpu.stream();
9403 let (p_, _g) = s.device_ptr_mut(&stream);
9404 unsafe {
9405 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9406 .result()?;
9407 }
9408 }
9409 }
9410 self.keep_if_capturing(&s);
9411 Ok(s)
9412 }
9413
9414 pub fn uninit_q8_pair(
9419 &self,
9420 n: usize,
9421 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9422 Ok((
9423 self.alloc_uninit::<i8>(n)?,
9424 self.alloc_uninit::<f32>(n / 32)?,
9425 ))
9426 }
9427
9428 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9429 self.alloc_uninit::<f32>(n)
9430 }
9431
9432 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9434 self.alloc_uninit::<i8>(n)
9435 }
9436
9437 #[allow(clippy::too_many_arguments)]
9441 pub fn rms_norm3(
9442 &self,
9443 x: &CudaSlice<f32>,
9444 w0: &CudaSlice<f32>,
9445 w1: &CudaSlice<f32>,
9446 w2: &CudaSlice<f32>,
9447 d0: &mut CudaSlice<f32>,
9448 d1: &mut CudaSlice<f32>,
9449 d2: &mut CudaSlice<f32>,
9450 ncols: usize,
9451 nrows: usize,
9452 eps: f32,
9453 ) -> Result<(), Box<dyn std::error::Error>> {
9454 let f = self.func("rms_norm3_f32");
9455 let cfg = LaunchConfig {
9456 grid_dim: (nrows as u32, 1, 1),
9457 block_dim: (rms_block(), 1, 1),
9458 shared_mem_bytes: 0,
9459 };
9460 let (nc, e) = (ncols as i32, eps);
9461 let __s_b = self.gpu.stream();
9462 let mut b = __s_b.launch_builder(&f);
9463 b.arg(x)
9464 .arg(w0)
9465 .arg(w1)
9466 .arg(w2)
9467 .arg(d0)
9468 .arg(d1)
9469 .arg(d2)
9470 .arg(&nc)
9471 .arg(&e);
9472 unsafe {
9473 b.launch(cfg)?;
9474 }
9475 Ok(())
9476 }
9477
9478 #[allow(clippy::too_many_arguments)]
9480 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9483 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9484 *WARP_ON.get_or_init(|| {
9485 std::env::var("MEMRA_QKVNORM_W")
9486 .map(|v| v != "0")
9487 .unwrap_or(true)
9488 }) && ncols % 4 == 0
9489 && rows >= 64
9490 }
9491
9492 #[allow(clippy::too_many_arguments)]
9495 pub fn rms_norm_qkv_w4b(
9496 &self,
9497 q: &CudaSlice<f32>,
9498 k: &CudaSlice<f32>,
9499 v: &CudaSlice<f32>,
9500 wq: &CudaSlice<f32>,
9501 wk: &CudaSlice<f32>,
9502 wv: &CudaSlice<f32>,
9503 dq: &mut CudaSlice<f32>,
9504 dk: &mut CudaSlice<f32>,
9505 dv: &mut CudaSlice<f32>,
9506 dvb: &mut CudaSlice<u8>,
9507 ncols: usize,
9508 rq: usize,
9509 rk: usize,
9510 eps: f32,
9511 vf16: bool,
9512 ) -> Result<(), Box<dyn std::error::Error>> {
9513 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9514 let f = self.func("rms_norm_qkv_w4b_f32");
9515 let rows = (rq + 2 * rk) as u32;
9516 let cfg = LaunchConfig {
9517 grid_dim: (rows.div_ceil(8), 1, 1),
9518 block_dim: (256, 1, 1),
9519 shared_mem_bytes: 0,
9520 };
9521 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9522 let vf = vf16 as i32;
9523 let __s_b = self.gpu.stream();
9524 let mut b = __s_b.launch_builder(&f);
9525 b.arg(q)
9526 .arg(k)
9527 .arg(v)
9528 .arg(wq)
9529 .arg(wk)
9530 .arg(wv)
9531 .arg(dq)
9532 .arg(dk)
9533 .arg(dv)
9534 .arg(&mut *dvb)
9535 .arg(&nc)
9536 .arg(&rqi)
9537 .arg(&rki)
9538 .arg(&rvi)
9539 .arg(&e)
9540 .arg(&vf);
9541 unsafe {
9542 b.launch(cfg)?;
9543 }
9544 Ok(())
9545 }
9546
9547 pub fn rms_norm_qkv(
9548 &self,
9549 q: &CudaSlice<f32>,
9550 k: &CudaSlice<f32>,
9551 v: &CudaSlice<f32>,
9552 wq: &CudaSlice<f32>,
9553 wk: &CudaSlice<f32>,
9554 wv: &CudaSlice<f32>,
9555 dq: &mut CudaSlice<f32>,
9556 dk: &mut CudaSlice<f32>,
9557 dv: &mut CudaSlice<f32>,
9558 ncols: usize,
9559 rq: usize,
9560 rk: usize,
9561 eps: f32,
9562 ) -> Result<(), Box<dyn std::error::Error>> {
9563 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9567 let warp_on = *WARP_ON.get_or_init(|| {
9568 std::env::var("MEMRA_QKVNORM_W")
9569 .map(|v| v != "0")
9570 .unwrap_or(true)
9571 });
9572 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9575 let f = self.func("rms_norm_qkv_w4_f32");
9576 let rows = (rq + 2 * rk) as u32;
9577 let cfg = LaunchConfig {
9578 grid_dim: (rows.div_ceil(8), 1, 1),
9579 block_dim: (256, 1, 1),
9580 shared_mem_bytes: 0,
9581 };
9582 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9583 let __s_b = self.gpu.stream();
9584 let mut b = __s_b.launch_builder(&f);
9585 b.arg(q)
9586 .arg(k)
9587 .arg(v)
9588 .arg(wq)
9589 .arg(wk)
9590 .arg(wv)
9591 .arg(dq)
9592 .arg(dk)
9593 .arg(dv)
9594 .arg(&nc)
9595 .arg(&rqi)
9596 .arg(&rki)
9597 .arg(&rvi)
9598 .arg(&e);
9599 unsafe {
9600 b.launch(cfg)?;
9601 }
9602 return Ok(());
9603 }
9604 let f = self.func("rms_norm_qkv_f32");
9605 let grid = (rq + 2 * rk) as u32;
9606 let cfg = LaunchConfig {
9607 grid_dim: (grid, 1, 1),
9608 block_dim: (rms_block(), 1, 1),
9609 shared_mem_bytes: 0,
9610 };
9611 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9612 let __s_b = self.gpu.stream();
9613 let mut b = __s_b.launch_builder(&f);
9614 b.arg(q)
9615 .arg(k)
9616 .arg(v)
9617 .arg(wq)
9618 .arg(wk)
9619 .arg(wv)
9620 .arg(dq)
9621 .arg(dk)
9622 .arg(dv)
9623 .arg(&nc)
9624 .arg(&rqi)
9625 .arg(&rki)
9626 .arg(&e);
9627 unsafe {
9628 b.launch(cfg)?;
9629 }
9630 Ok(())
9631 }
9632
9633 #[allow(clippy::too_many_arguments)]
9635 pub fn rms_norm2x(
9636 &self,
9637 a: &CudaSlice<f32>,
9638 bb: &CudaSlice<f32>,
9639 wa: &CudaSlice<f32>,
9640 wb: &CudaSlice<f32>,
9641 da: &mut CudaSlice<f32>,
9642 db: &mut CudaSlice<f32>,
9643 ncols: usize,
9644 nrows: usize,
9645 eps: f32,
9646 ) -> Result<(), Box<dyn std::error::Error>> {
9647 let f = self.func("rms_norm2x_f32");
9648 let cfg = LaunchConfig {
9649 grid_dim: (2 * nrows as u32, 1, 1),
9650 block_dim: (rms_block(), 1, 1),
9651 shared_mem_bytes: 0,
9652 };
9653 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9654 let __s_b = self.gpu.stream();
9655 let mut b = __s_b.launch_builder(&f);
9656 b.arg(a)
9657 .arg(bb)
9658 .arg(wa)
9659 .arg(wb)
9660 .arg(da)
9661 .arg(db)
9662 .arg(&nc)
9663 .arg(&nr)
9664 .arg(&e);
9665 unsafe {
9666 b.launch(cfg)?;
9667 }
9668 Ok(())
9669 }
9670
9671 pub fn softcap(
9673 &self,
9674 y: &mut CudaSlice<f32>,
9675 cap: f32,
9676 n: usize,
9677 ) -> Result<(), Box<dyn std::error::Error>> {
9678 let f = self.func("softcap_f32");
9679 let cfg = LaunchConfig::for_num_elems(n as u32);
9680 let ni = n as i32;
9681 let __s_b = self.gpu.stream();
9682 let mut b = __s_b.launch_builder(&f);
9683 b.arg(y).arg(&cap).arg(&ni);
9684 unsafe {
9685 b.launch(cfg)?;
9686 }
9687 Ok(())
9688 }
9689
9690 pub fn mask_ids_rows(
9693 &self,
9694 y: &mut CudaSlice<f32>,
9695 ids: &CudaSlice<i32>,
9696 n_ids: usize,
9697 n_vocab: usize,
9698 t: usize,
9699 ) -> Result<(), Box<dyn std::error::Error>> {
9700 let f = self.func("mask_ids_rows_f32");
9701 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9702 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9703 let __s_b = self.gpu.stream();
9704 let mut b = __s_b.launch_builder(&f);
9705 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9706 unsafe {
9707 b.launch(cfg)?;
9708 }
9709 Ok(())
9710 }
9711
9712 #[allow(clippy::too_many_arguments)]
9714 pub fn add_scale_rms_norm(
9715 &self,
9716 a: &CudaSlice<f32>,
9717 b_in: &CudaSlice<f32>,
9718 c: f32,
9719 w: &CudaSlice<f32>,
9720 res: &mut CudaSlice<f32>,
9721 dst: &mut CudaSlice<f32>,
9722 ncols: usize,
9723 nrows: usize,
9724 eps: f32,
9725 ) -> Result<(), Box<dyn std::error::Error>> {
9726 let f = self.func("add_scale_rms_norm_f32");
9727 let cfg = LaunchConfig {
9728 grid_dim: (nrows as u32, 1, 1),
9729 block_dim: (rms_block(), 1, 1),
9730 shared_mem_bytes: 0,
9731 };
9732 let (nc, e2) = (ncols as i32, eps);
9733 let __s_b = self.gpu.stream();
9734 let mut b = __s_b.launch_builder(&f);
9735 b.arg(a)
9736 .arg(b_in)
9737 .arg(&c)
9738 .arg(w)
9739 .arg(res)
9740 .arg(dst)
9741 .arg(&nc)
9742 .arg(&e2);
9743 unsafe {
9744 b.launch(cfg)?;
9745 }
9746 Ok(())
9747 }
9748
9749 #[allow(clippy::too_many_arguments)]
9752 pub fn add_scale_rms_norm_q8_1(
9753 &self,
9754 a: &CudaSlice<f32>,
9755 b_in: &CudaSlice<f32>,
9756 c: f32,
9757 w: &CudaSlice<f32>,
9758 res: &mut CudaSlice<f32>,
9759 ncols: usize,
9760 nrows: usize,
9761 eps: f32,
9762 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9763 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9764 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9765 let (nc, e2) = (ncols as i32, eps);
9766 if Self::pdl_on() && Self::pdl_wb_on() {
9767 {
9768 use cudarc::driver::{DevicePtr, DevicePtrMut};
9769 let s = &self.gpu.stream();
9770 let (pa, _g0) = a.device_ptr(s);
9771 let (pb, _g1) = b_in.device_ptr(s);
9772 let (pw, _g2) = w.device_ptr(s);
9773 let (pr, _g3) = res.device_ptr_mut(s);
9774 let (pq, _g4) = out_q.device_ptr_mut(s);
9775 let (pd, _g5) = out_d.device_ptr_mut(s);
9776 let mut ps = [
9777 &pa as *const _ as *mut std::ffi::c_void,
9778 &pb as *const _ as *mut _,
9779 &c as *const _ as *mut _,
9780 &pw as *const _ as *mut _,
9781 &pr as *const _ as *mut _,
9782 &pq as *const _ as *mut _,
9783 &pd as *const _ as *mut _,
9784 &nc as *const _ as *mut _,
9785 &e2 as *const _ as *mut _,
9786 ];
9787 unsafe {
9788 self.launch_pdl(
9789 "add_scale_rms_norm_q8_1",
9790 (nrows as u32, 1, 1),
9791 (rms_block(), 1, 1),
9792 &mut ps,
9793 )?;
9794 }
9795 }
9796 return Ok((out_q, out_d));
9797 }
9798 let f = self.func("add_scale_rms_norm_q8_1");
9799 let cfg = LaunchConfig {
9800 grid_dim: (nrows as u32, 1, 1),
9801 block_dim: (rms_block(), 1, 1),
9802 shared_mem_bytes: 0,
9803 };
9804 let __s_b = self.gpu.stream();
9805 let mut b = __s_b.launch_builder(&f);
9806 b.arg(a)
9807 .arg(b_in)
9808 .arg(&c)
9809 .arg(w)
9810 .arg(res)
9811 .arg(&mut out_q)
9812 .arg(&mut out_d)
9813 .arg(&nc)
9814 .arg(&e2);
9815 unsafe {
9816 b.launch(cfg)?;
9817 }
9818 Ok((out_q, out_d))
9819 }
9820
9821 #[allow(clippy::too_many_arguments)]
9823 pub fn add_scale_rms_norm_q8_1_into(
9824 &self,
9825 a: &CudaSlice<f32>,
9826 b_in: &CudaSlice<f32>,
9827 c: f32,
9828 w: &CudaSlice<f32>,
9829 res: &mut CudaSlice<f32>,
9830 ncols: usize,
9831 nrows: usize,
9832 eps: f32,
9833 out_q: &mut CudaSlice<i8>,
9834 out_d: &mut CudaSlice<f32>,
9835 ) -> Result<(), Box<dyn std::error::Error>> {
9836 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9837 let (nc, e2) = (ncols as i32, eps);
9838 if Self::pdl_on() && Self::pdl_wb_on() {
9839 use cudarc::driver::{DevicePtr, DevicePtrMut};
9840 let s = &self.gpu.stream();
9841 let (pa, _g0) = a.device_ptr(s);
9842 let (pb, _g1) = b_in.device_ptr(s);
9843 let (pw, _g2) = w.device_ptr(s);
9844 let (pr, _g3) = res.device_ptr_mut(s);
9845 let (pq, _g4) = out_q.device_ptr_mut(s);
9846 let (pd, _g5) = out_d.device_ptr_mut(s);
9847 let mut ps = [
9848 &pa as *const _ as *mut std::ffi::c_void,
9849 &pb as *const _ as *mut _,
9850 &c as *const _ as *mut _,
9851 &pw as *const _ as *mut _,
9852 &pr as *const _ as *mut _,
9853 &pq as *const _ as *mut _,
9854 &pd as *const _ as *mut _,
9855 &nc as *const _ as *mut _,
9856 &e2 as *const _ as *mut _,
9857 ];
9858 unsafe {
9859 self.launch_pdl(
9860 "add_scale_rms_norm_q8_1",
9861 (nrows as u32, 1, 1),
9862 (rms_block(), 1, 1),
9863 &mut ps,
9864 )?;
9865 }
9866 return Ok(());
9867 }
9868 let f = self.func("add_scale_rms_norm_q8_1");
9869 let cfg = LaunchConfig {
9870 grid_dim: (nrows as u32, 1, 1),
9871 block_dim: (rms_block(), 1, 1),
9872 shared_mem_bytes: 0,
9873 };
9874 let __s_b = self.gpu.stream();
9875 let mut b = __s_b.launch_builder(&f);
9876 b.arg(a)
9877 .arg(b_in)
9878 .arg(&c)
9879 .arg(w)
9880 .arg(res)
9881 .arg(&mut *out_q)
9882 .arg(&mut *out_d)
9883 .arg(&nc)
9884 .arg(&e2);
9885 unsafe {
9886 b.launch(cfg)?;
9887 }
9888 Ok(())
9889 }
9890
9891 #[allow(clippy::too_many_arguments)]
9894 pub fn rms_pre_add_scale_rms_norm_q8_1(
9895 &self,
9896 a: &CudaSlice<f32>,
9897 wa: &CudaSlice<f32>,
9898 b_in: &CudaSlice<f32>,
9899 c: f32,
9900 w: &CudaSlice<f32>,
9901 res: &mut CudaSlice<f32>,
9902 ncols: usize,
9903 nrows: usize,
9904 eps: f32,
9905 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9906 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9907 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9908 let (nc, e2) = (ncols as i32, eps);
9909 if Self::pdl_on() {
9910 {
9911 use cudarc::driver::{DevicePtr, DevicePtrMut};
9912 let s = &self.gpu.stream();
9913 let (pa, _g0) = a.device_ptr(s);
9914 let (pwa, _g1) = wa.device_ptr(s);
9915 let (pb, _g2) = b_in.device_ptr(s);
9916 let (pw, _g3) = w.device_ptr(s);
9917 let (pr, _g4) = res.device_ptr_mut(s);
9918 let (pq, _g5) = out_q.device_ptr_mut(s);
9919 let (pd, _g6) = out_d.device_ptr_mut(s);
9920 let mut ps = [
9921 &pa as *const _ as *mut std::ffi::c_void,
9922 &pwa as *const _ as *mut _,
9923 &pb as *const _ as *mut _,
9924 &c as *const _ as *mut _,
9925 &pw as *const _ as *mut _,
9926 &pr as *const _ as *mut _,
9927 &pq as *const _ as *mut _,
9928 &pd as *const _ as *mut _,
9929 &nc as *const _ as *mut _,
9930 &e2 as *const _ as *mut _,
9931 ];
9932 unsafe {
9933 self.launch_pdl(
9934 "rms_pre_add_scale_rms_norm_q8_1",
9935 (nrows as u32, 1, 1),
9936 (rms_block(), 1, 1),
9937 &mut ps,
9938 )?;
9939 }
9940 }
9941 return Ok((out_q, out_d));
9942 }
9943 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9944 let cfg = LaunchConfig {
9945 grid_dim: (nrows as u32, 1, 1),
9946 block_dim: (rms_block(), 1, 1),
9947 shared_mem_bytes: 0,
9948 };
9949 let __s_b = self.gpu.stream();
9950 let mut b = __s_b.launch_builder(&f);
9951 b.arg(a)
9952 .arg(wa)
9953 .arg(b_in)
9954 .arg(&c)
9955 .arg(w)
9956 .arg(res)
9957 .arg(&mut out_q)
9958 .arg(&mut out_d)
9959 .arg(&nc)
9960 .arg(&e2);
9961 unsafe {
9962 b.launch(cfg)?;
9963 }
9964 Ok((out_q, out_d))
9965 }
9966
9967 pub fn gelu_tanh_mul_q8_1(
9970 &self,
9971 gate: &CudaSlice<f32>,
9972 up: &cudarc::driver::CudaView<f32>,
9973 act: &mut CudaSlice<f32>,
9974 ncols: usize,
9975 nrows: usize,
9976 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9977 debug_assert!(ncols % 128 == 0);
9978 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9979 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9980 let nc = ncols as i32;
9981 if Self::pdl_on() {
9982 {
9983 use cudarc::driver::{DevicePtr, DevicePtrMut};
9984 let s = &self.gpu.stream();
9985 let (pg, _g0) = gate.device_ptr(s);
9986 let (pu, _g1) = up.device_ptr(s);
9987 let (pact, _g2) = act.device_ptr_mut(s);
9988 let (pq, _g3) = out_q.device_ptr_mut(s);
9989 let (pd, _g4) = out_d.device_ptr_mut(s);
9990 let mut ps = [
9991 &pg as *const _ as *mut std::ffi::c_void,
9992 &pu as *const _ as *mut _,
9993 &pact as *const _ as *mut _,
9994 &pq as *const _ as *mut _,
9995 &pd as *const _ as *mut _,
9996 &nc as *const _ as *mut _,
9997 ];
9998 unsafe {
9999 self.launch_pdl(
10000 "gelu_tanh_mul_q8_1",
10001 (nrows as u32, 1, 1),
10002 (rms_block(), 1, 1),
10003 &mut ps,
10004 )?;
10005 }
10006 }
10007 return Ok((out_q, out_d));
10008 }
10009 let f = self.func("gelu_tanh_mul_q8_1");
10010 let cfg = LaunchConfig {
10011 grid_dim: (nrows as u32, 1, 1),
10012 block_dim: (rms_block(), 1, 1),
10013 shared_mem_bytes: 0,
10014 };
10015 let __s_b = self.gpu.stream();
10016 let mut b = __s_b.launch_builder(&f);
10017 b.arg(gate)
10018 .arg(up)
10019 .arg(act)
10020 .arg(&mut out_q)
10021 .arg(&mut out_d)
10022 .arg(&nc);
10023 unsafe {
10024 b.launch(cfg)?;
10025 }
10026 Ok((out_q, out_d))
10027 }
10028
10029 #[allow(clippy::too_many_arguments)]
10031 pub fn gelu_tanh_mul_q8_1_into(
10032 &self,
10033 gate: &CudaSlice<f32>,
10034 up: &cudarc::driver::CudaView<f32>,
10035 act: &mut CudaSlice<f32>,
10036 ncols: usize,
10037 nrows: usize,
10038 out_q: &mut CudaSlice<i8>,
10039 out_d: &mut CudaSlice<f32>,
10040 ) -> Result<(), Box<dyn std::error::Error>> {
10041 debug_assert!(ncols % 128 == 0);
10042 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
10043 let nc = ncols as i32;
10044 if Self::pdl_on() {
10045 use cudarc::driver::{DevicePtr, DevicePtrMut};
10046 let s = &self.gpu.stream();
10047 let (pg, _g0) = gate.device_ptr(s);
10048 let (pu, _g1) = up.device_ptr(s);
10049 let (pact, _g2) = act.device_ptr_mut(s);
10050 let (pq, _g3) = out_q.device_ptr_mut(s);
10051 let (pd, _g4) = out_d.device_ptr_mut(s);
10052 let mut ps = [
10053 &pg as *const _ as *mut std::ffi::c_void,
10054 &pu as *const _ as *mut _,
10055 &pact as *const _ as *mut _,
10056 &pq as *const _ as *mut _,
10057 &pd as *const _ as *mut _,
10058 &nc as *const _ as *mut _,
10059 ];
10060 unsafe {
10061 self.launch_pdl(
10062 "gelu_tanh_mul_q8_1",
10063 (nrows as u32, 1, 1),
10064 (rms_block(), 1, 1),
10065 &mut ps,
10066 )?;
10067 }
10068 return Ok(());
10069 }
10070 let f = self.func("gelu_tanh_mul_q8_1");
10071 let cfg = LaunchConfig {
10072 grid_dim: (nrows as u32, 1, 1),
10073 block_dim: (rms_block(), 1, 1),
10074 shared_mem_bytes: 0,
10075 };
10076 let __s_b = self.gpu.stream();
10077 let mut b = __s_b.launch_builder(&f);
10078 b.arg(gate)
10079 .arg(up)
10080 .arg(&mut *act)
10081 .arg(&mut *out_q)
10082 .arg(&mut *out_d)
10083 .arg(&nc);
10084 unsafe {
10085 b.launch(cfg)?;
10086 }
10087 Ok(())
10088 }
10089
10090 #[allow(clippy::too_many_arguments)]
10092 pub fn add_rms_norm3_q8z(
10093 &self,
10094 a: &CudaSlice<f32>,
10095 b_in: &CudaSlice<f32>,
10096 w0: &CudaSlice<f32>,
10097 w1: &CudaSlice<f32>,
10098 w2: &CudaSlice<f32>,
10099 res: &mut CudaSlice<f32>,
10100 out1: &mut CudaSlice<f32>,
10101 ncols: usize,
10102 nrows: usize,
10103 eps: f32,
10104 ) -> Result<
10105 (
10106 (CudaSlice<i8>, CudaSlice<f32>),
10107 (CudaSlice<i8>, CudaSlice<f32>),
10108 ),
10109 Box<dyn std::error::Error>,
10110 > {
10111 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
10112 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10113 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
10114 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10115 let f = self.func("add_rms_norm3_q8z_f32");
10116 let cfg = LaunchConfig {
10117 grid_dim: (nrows as u32, 1, 1),
10118 block_dim: (rms_block(), 1, 1),
10119 shared_mem_bytes: 0,
10120 };
10121 let (nc, e2) = (ncols as i32, eps);
10122 let __s_b = self.gpu.stream();
10123 let mut b = __s_b.launch_builder(&f);
10124 b.arg(a)
10125 .arg(b_in)
10126 .arg(w0)
10127 .arg(w1)
10128 .arg(w2)
10129 .arg(res)
10130 .arg(&mut q0)
10131 .arg(&mut d0)
10132 .arg(out1)
10133 .arg(&mut q2)
10134 .arg(&mut d2)
10135 .arg(&nc)
10136 .arg(&e2);
10137 unsafe {
10138 b.launch(cfg)?;
10139 }
10140 Ok(((q0, d0), (q2, d2)))
10141 }
10142
10143 #[allow(clippy::too_many_arguments)]
10145 pub fn add_rms_norm3(
10146 &self,
10147 a: &CudaSlice<f32>,
10148 b_in: &CudaSlice<f32>,
10149 w0: &CudaSlice<f32>,
10150 w1: &CudaSlice<f32>,
10151 w2: &CudaSlice<f32>,
10152 res: &mut CudaSlice<f32>,
10153 d0: &mut CudaSlice<f32>,
10154 d1: &mut CudaSlice<f32>,
10155 d2: &mut CudaSlice<f32>,
10156 ncols: usize,
10157 nrows: usize,
10158 eps: f32,
10159 ) -> Result<(), Box<dyn std::error::Error>> {
10160 let f = self.func("add_rms_norm3_f32");
10161 let cfg = LaunchConfig {
10162 grid_dim: (nrows as u32, 1, 1),
10163 block_dim: (rms_block(), 1, 1),
10164 shared_mem_bytes: 0,
10165 };
10166 let (nc, e2) = (ncols as i32, eps);
10167 let __s_b = self.gpu.stream();
10168 let mut b = __s_b.launch_builder(&f);
10169 b.arg(a)
10170 .arg(b_in)
10171 .arg(w0)
10172 .arg(w1)
10173 .arg(w2)
10174 .arg(res)
10175 .arg(d0)
10176 .arg(d1)
10177 .arg(d2)
10178 .arg(&nc)
10179 .arg(&e2);
10180 unsafe {
10181 b.launch(cfg)?;
10182 }
10183 Ok(())
10184 }
10185
10186 pub fn add_scale(
10188 &self,
10189 a: &CudaSlice<f32>,
10190 b_in: &CudaSlice<f32>,
10191 c: f32,
10192 dst: &mut CudaSlice<f32>,
10193 n: usize,
10194 ) -> Result<(), Box<dyn std::error::Error>> {
10195 let f = self.func("add_scale_f32");
10196 let cfg = LaunchConfig::for_num_elems(n as u32);
10197 let ni = n as i32;
10198 let __s_b = self.gpu.stream();
10199 let mut b = __s_b.launch_builder(&f);
10200 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
10201 unsafe {
10202 b.launch(cfg)?;
10203 }
10204 Ok(())
10205 }
10206
10207 pub fn layer_norm_bias(
10209 &self,
10210 x: &CudaSlice<f32>,
10211 w: &CudaSlice<f32>,
10212 b: &CudaSlice<f32>,
10213 dst: &mut CudaSlice<f32>,
10214 ncols: usize,
10215 nrows: usize,
10216 eps: f32,
10217 ) -> Result<(), Box<dyn std::error::Error>> {
10218 let f = self.func("layer_norm_bias_f32");
10219 let (nc, e) = (ncols as i32, eps);
10220 let cfg = LaunchConfig {
10221 grid_dim: (nrows as u32, 1, 1),
10222 block_dim: (256, 1, 1),
10223 shared_mem_bytes: 0,
10224 };
10225 let __s_b = self.gpu.stream();
10226 let mut lb = __s_b.launch_builder(&f);
10227 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
10228 unsafe {
10229 lb.launch(cfg)?;
10230 }
10231 Ok(())
10232 }
10233
10234 pub fn gelu_tanh(
10236 &self,
10237 x: &CudaSlice<f32>,
10238 dst: &mut CudaSlice<f32>,
10239 n: usize,
10240 ) -> Result<(), Box<dyn std::error::Error>> {
10241 let f = self.func("gelu_tanh_f32");
10242 let ni = n as i64;
10243 let cfg = LaunchConfig {
10244 grid_dim: (n.div_ceil(256) as u32, 1, 1),
10245 block_dim: (256, 1, 1),
10246 shared_mem_bytes: 0,
10247 };
10248 let __s_b = self.gpu.stream();
10249 let mut lb = __s_b.launch_builder(&f);
10250 lb.arg(x).arg(&mut *dst).arg(&ni);
10251 unsafe {
10252 lb.launch(cfg)?;
10253 }
10254 Ok(())
10255 }
10256
10257 pub fn row_softmax(
10259 &self,
10260 x: &mut CudaSlice<f32>,
10261 ncols: usize,
10262 nrows: usize,
10263 ) -> Result<(), Box<dyn std::error::Error>> {
10264 let f = self.func("row_softmax_f32");
10265 let nc = ncols as i32;
10266 let cfg = LaunchConfig {
10267 grid_dim: (nrows as u32, 1, 1),
10268 block_dim: (256, 1, 1),
10269 shared_mem_bytes: 0,
10270 };
10271 let __s_b = self.gpu.stream();
10272 let mut lb = __s_b.launch_builder(&f);
10273 lb.arg(&mut *x).arg(&nc);
10274 unsafe {
10275 lb.launch(cfg)?;
10276 }
10277 Ok(())
10278 }
10279
10280 pub fn rms_norm(
10281 &self,
10282 x: &CudaSlice<f32>,
10283 w: &CudaSlice<f32>,
10284 dst: &mut CudaSlice<f32>,
10285 ncols: usize,
10286 nrows: usize,
10287 eps: f32,
10288 ) -> Result<(), Box<dyn std::error::Error>> {
10289 let (nc, e) = (ncols as i32, eps);
10290 let kname = if Self::norm_ilp_on() {
10291 "rms_norm_f32_v2"
10292 } else {
10293 "rms_norm_f32"
10294 };
10295 if Self::pdl_on() && Self::pdl_wb_on() {
10296 use cudarc::driver::{DevicePtr, DevicePtrMut};
10297 let s = &self.gpu.stream();
10298 let (px, _g0) = x.device_ptr(s);
10299 let (pw, _g1) = w.device_ptr(s);
10300 let (pd, _g2) = dst.device_ptr_mut(s);
10301 let mut ps = [
10302 &px as *const _ as *mut std::ffi::c_void,
10303 &pw as *const _ as *mut _,
10304 &pd as *const _ as *mut _,
10305 &nc as *const _ as *mut _,
10306 &e as *const _ as *mut _,
10307 ];
10308 unsafe {
10309 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10310 }
10311 return Ok(());
10312 }
10313 let f = self.func(kname);
10314 let cfg = LaunchConfig {
10315 grid_dim: (nrows as u32, 1, 1),
10316 block_dim: (rms_block(), 1, 1),
10317 shared_mem_bytes: 0,
10318 };
10319 let __s_b = self.gpu.stream();
10320 let mut b = __s_b.launch_builder(&f);
10321 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10322 unsafe {
10323 b.launch(cfg)?;
10324 }
10325 Ok(())
10326 }
10327
10328 pub fn rms_norm_decode(
10336 &self,
10337 x: &CudaSlice<f32>,
10338 w: &CudaSlice<f32>,
10339 dst: &mut CudaSlice<f32>,
10340 ncols: usize,
10341 nrows: usize,
10342 eps: f32,
10343 ) -> Result<(), Box<dyn std::error::Error>> {
10344 let f = self.func(if Self::norm_ilp_on() {
10345 "rms_norm_f32_v2"
10346 } else {
10347 "rms_norm_f32"
10348 });
10349 let cfg = LaunchConfig {
10350 grid_dim: (nrows as u32, 1, 1),
10351 block_dim: (1024, 1, 1),
10352 shared_mem_bytes: 0,
10353 };
10354 let (nc, e) = (ncols as i32, eps);
10355 let __s_b = self.gpu.stream();
10356 let mut b = __s_b.launch_builder(&f);
10357 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10358 unsafe {
10359 b.launch(cfg)?;
10360 }
10361 Ok(())
10362 }
10363
10364 pub fn rms_norm_q8_1(
10368 &self,
10369 x: &CudaSlice<f32>,
10370 w: &CudaSlice<f32>,
10371 ncols: usize,
10372 nrows: usize,
10373 eps: f32,
10374 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10375 let nblk = ncols / 32;
10376 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10377 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10378 let (nc, e) = (ncols as i32, eps);
10379 if Self::pdl_on() {
10380 {
10381 use cudarc::driver::{DevicePtr, DevicePtrMut};
10382 let s = &self.gpu.stream();
10383 let (px, _g0) = x.device_ptr(s);
10384 let (pw, _g1) = w.device_ptr(s);
10385 let (pq, _g2) = q.device_ptr_mut(s);
10386 let (pd, _g3) = d.device_ptr_mut(s);
10387 let mut ps = [
10388 &px as *const _ as *mut std::ffi::c_void,
10389 &pw as *const _ as *mut _,
10390 &pq as *const _ as *mut _,
10391 &pd as *const _ as *mut _,
10392 &nc as *const _ as *mut _,
10393 &e as *const _ as *mut _,
10394 ];
10395 unsafe {
10396 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10397 }
10398 }
10399 return Ok((q, d));
10400 }
10401 let f = self.func("rms_norm_q8_1");
10402 let cfg = LaunchConfig {
10405 grid_dim: (nrows as u32, 1, 1),
10406 block_dim: (1024, 1, 1),
10407 shared_mem_bytes: 0,
10408 };
10409 let __s_b = self.gpu.stream();
10410 let mut b = __s_b.launch_builder(&f);
10411 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10412 unsafe {
10413 b.launch(cfg)?;
10414 }
10415 Ok((q, d))
10416 }
10417
10418 pub fn rms_norm_q8_1_into(
10421 &self,
10422 x: &CudaSlice<f32>,
10423 w: &CudaSlice<f32>,
10424 ncols: usize,
10425 nrows: usize,
10426 eps: f32,
10427 q: &mut CudaSlice<i8>,
10428 d: &mut CudaSlice<f32>,
10429 ) -> Result<(), Box<dyn std::error::Error>> {
10430 let nblk = ncols / 32;
10431 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10432 let (nc, e) = (ncols as i32, eps);
10433 if Self::pdl_on() {
10434 use cudarc::driver::{DevicePtr, DevicePtrMut};
10435 let s = &self.gpu.stream();
10436 let (px, _g0) = x.device_ptr(s);
10437 let (pw, _g1) = w.device_ptr(s);
10438 let (pq, _g2) = q.device_ptr_mut(s);
10439 let (pd, _g3) = d.device_ptr_mut(s);
10440 let mut ps = [
10441 &px as *const _ as *mut std::ffi::c_void,
10442 &pw as *const _ as *mut _,
10443 &pq as *const _ as *mut _,
10444 &pd as *const _ as *mut _,
10445 &nc as *const _ as *mut _,
10446 &e as *const _ as *mut _,
10447 ];
10448 unsafe {
10449 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10450 }
10451 return Ok(());
10452 }
10453 let f = self.func("rms_norm_q8_1");
10454 let cfg = LaunchConfig {
10455 grid_dim: (nrows as u32, 1, 1),
10456 block_dim: (1024, 1, 1),
10457 shared_mem_bytes: 0,
10458 };
10459 let __s_b = self.gpu.stream();
10460 let mut b = __s_b.launch_builder(&f);
10461 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10462 unsafe {
10463 b.launch(cfg)?;
10464 }
10465 Ok(())
10466 }
10467
10468 pub fn quantize_q8_1_into(
10470 &self,
10471 x: &CudaSlice<f32>,
10472 m: usize,
10473 in_f: usize,
10474 q: &mut CudaSlice<i8>,
10475 d: &mut CudaSlice<f32>,
10476 ) -> Result<(), Box<dyn std::error::Error>> {
10477 let nblk = in_f / 32;
10478 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10479 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10480 let (inf, mi) = (in_f as i32, m as i32);
10481 if Self::pdl_on() && Self::pdl_wb_on() {
10482 use cudarc::driver::{DevicePtr, DevicePtrMut};
10483 let s = &self.gpu.stream();
10484 let (px, _g0) = x.device_ptr(s);
10485 let (pq, _g1) = q.device_ptr_mut(s);
10486 let (pd, _g2) = d.device_ptr_mut(s);
10487 let mut ps = [
10488 &px as *const _ as *mut std::ffi::c_void,
10489 &pq as *const _ as *mut _,
10490 &pd as *const _ as *mut _,
10491 &inf as *const _ as *mut _,
10492 &mi as *const _ as *mut _,
10493 ];
10494 unsafe {
10495 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10496 }
10497 return Ok(());
10498 }
10499 let f = self.func("quantize_q8_1");
10500 let __s_b = self.gpu.stream();
10501 let mut b = __s_b.launch_builder(&f);
10502 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10503 unsafe {
10504 b.launch(cfg)?;
10505 }
10506 Ok(())
10507 }
10508
10509 pub fn add_rms_norm_q8_1(
10513 &self,
10514 a: &CudaSlice<f32>,
10515 b_in: &CudaSlice<f32>,
10516 w: &CudaSlice<f32>,
10517 res: &mut CudaSlice<f32>,
10518 ncols: usize,
10519 nrows: usize,
10520 eps: f32,
10521 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10522 let nblk = ncols / 32;
10523 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10524 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10525 let f = self.func("add_rms_norm_q8_1");
10526 let cfg = LaunchConfig {
10528 grid_dim: (nrows as u32, 1, 1),
10529 block_dim: (1024, 1, 1),
10530 shared_mem_bytes: 0,
10531 };
10532 let (nc, e) = (ncols as i32, eps);
10533 let __s_bld = self.gpu.stream();
10534 let mut bld = __s_bld.launch_builder(&f);
10535 bld.arg(a)
10536 .arg(b_in)
10537 .arg(w)
10538 .arg(res)
10539 .arg(&mut q)
10540 .arg(&mut d)
10541 .arg(&nc)
10542 .arg(&e);
10543 unsafe {
10544 bld.launch(cfg)?;
10545 }
10546 Ok((q, d))
10547 }
10548
10549 #[allow(clippy::too_many_arguments)]
10555 pub fn join_add_rms_norm_raw(
10556 &self,
10557 a0_raw: u64,
10558 a1_raw: u64,
10559 x: &CudaSlice<f32>,
10560 w: &CudaSlice<f32>,
10561 res: &mut CudaSlice<f32>,
10562 dst: &mut CudaSlice<f32>,
10563 ncols: usize,
10564 eps: f32,
10565 ) -> Result<(), Box<dyn std::error::Error>> {
10566 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10567 return Err("join_add_rms_norm geometry".into());
10568 }
10569 let f = self.func("join_add_rms_norm_f32");
10570 let cfg = LaunchConfig {
10571 grid_dim: (1, 1, 1),
10572 block_dim: (rms_block(), 1, 1),
10573 shared_mem_bytes: 0,
10574 };
10575 let (nc, e) = (ncols as i32, eps);
10576 let __s_b = self.gpu.stream();
10577 let mut b = __s_b.launch_builder(&f);
10578 b.arg(&a0_raw)
10579 .arg(&a1_raw)
10580 .arg(x)
10581 .arg(w)
10582 .arg(&mut *res)
10583 .arg(&mut *dst)
10584 .arg(&nc)
10585 .arg(&e);
10586 unsafe {
10587 b.launch(cfg)?;
10588 }
10589 Ok(())
10590 }
10591
10592 pub fn add_rms_norm(
10593 &self,
10594 a: &CudaSlice<f32>,
10595 b: &CudaSlice<f32>,
10596 w: &CudaSlice<f32>,
10597 res: &mut CudaSlice<f32>,
10598 dst: &mut CudaSlice<f32>,
10599 ncols: usize,
10600 nrows: usize,
10601 eps: f32,
10602 ) -> Result<(), Box<dyn std::error::Error>> {
10603 let (nc, e) = (ncols as i32, eps);
10604 let kname = if Self::norm_ilp_on() {
10605 "add_rms_norm_f32_v2"
10606 } else {
10607 "add_rms_norm_f32"
10608 };
10609 if Self::pdl_on() && Self::pdl_wb_on() {
10610 use cudarc::driver::{DevicePtr, DevicePtrMut};
10611 let s = &self.gpu.stream();
10612 let (pa, _g0) = a.device_ptr(s);
10613 let (pb, _g1) = b.device_ptr(s);
10614 let (pw, _g2) = w.device_ptr(s);
10615 let (pr, _g3) = res.device_ptr_mut(s);
10616 let (pd, _g4) = dst.device_ptr_mut(s);
10617 let mut ps = [
10618 &pa as *const _ as *mut std::ffi::c_void,
10619 &pb as *const _ as *mut _,
10620 &pw as *const _ as *mut _,
10621 &pr as *const _ as *mut _,
10622 &pd as *const _ as *mut _,
10623 &nc as *const _ as *mut _,
10624 &e as *const _ as *mut _,
10625 ];
10626 unsafe {
10627 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10628 }
10629 return Ok(());
10630 }
10631 let f = self.func(kname);
10632 let cfg = LaunchConfig {
10633 grid_dim: (nrows as u32, 1, 1),
10634 block_dim: (rms_block(), 1, 1),
10635 shared_mem_bytes: 0,
10636 };
10637 let __s_b2 = self.gpu.stream();
10638 let mut b2 = __s_b2.launch_builder(&f);
10639 b2.arg(a)
10640 .arg(b)
10641 .arg(w)
10642 .arg(&mut *res)
10643 .arg(&mut *dst)
10644 .arg(&nc)
10645 .arg(&e);
10646 unsafe {
10647 b2.launch(cfg)?;
10648 }
10649 Ok(())
10650 }
10651
10652 #[allow(clippy::too_many_arguments)]
10655 pub fn rms_pre_add_rms_norm(
10656 &self,
10657 a: &CudaSlice<f32>,
10658 wa: &CudaSlice<f32>,
10659 b: &CudaSlice<f32>,
10660 w: &CudaSlice<f32>,
10661 res: &mut CudaSlice<f32>,
10662 dst: &mut CudaSlice<f32>,
10663 ncols: usize,
10664 nrows: usize,
10665 eps: f32,
10666 ) -> Result<(), Box<dyn std::error::Error>> {
10667 let f = self.func("rms_pre_add_rms_norm_f32");
10668 let cfg = LaunchConfig {
10669 grid_dim: (nrows as u32, 1, 1),
10670 block_dim: (rms_block(), 1, 1),
10671 shared_mem_bytes: 0,
10672 };
10673 let (nc, e) = (ncols as i32, eps);
10674 let __s_b2 = self.gpu.stream();
10675 let mut b2 = __s_b2.launch_builder(&f);
10676 b2.arg(a)
10677 .arg(wa)
10678 .arg(b)
10679 .arg(w)
10680 .arg(&mut *res)
10681 .arg(&mut *dst)
10682 .arg(&nc)
10683 .arg(&e);
10684 unsafe {
10685 b2.launch(cfg)?;
10686 }
10687 Ok(())
10688 }
10689
10690 #[allow(clippy::too_many_arguments)]
10692 pub fn rms_pre_add_rms_norm_q8z(
10693 &self,
10694 a: &CudaSlice<f32>,
10695 wa: &CudaSlice<f32>,
10696 b: &CudaSlice<f32>,
10697 w: &CudaSlice<f32>,
10698 res: &mut CudaSlice<f32>,
10699 dst: &mut CudaSlice<f32>,
10700 ncols: usize,
10701 nrows: usize,
10702 eps: f32,
10703 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10704 debug_assert!(ncols % 128 == 0);
10705 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10706 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10707 let (nc, e) = (ncols as i32, eps);
10708 if Self::pdl_on() {
10709 {
10710 use cudarc::driver::{DevicePtr, DevicePtrMut};
10711 let s = &self.gpu.stream();
10712 let (pa, _g0) = a.device_ptr(s);
10713 let (pwa, _g1) = wa.device_ptr(s);
10714 let (pb, _g2) = b.device_ptr(s);
10715 let (pw, _g3) = w.device_ptr(s);
10716 let (pr, _g4) = res.device_ptr_mut(s);
10717 let (pdst, _g5) = dst.device_ptr_mut(s);
10718 let (pq, _g6) = out_q.device_ptr_mut(s);
10719 let (pd, _g7) = out_d.device_ptr_mut(s);
10720 let mut ps = [
10721 &pa as *const _ as *mut std::ffi::c_void,
10722 &pwa as *const _ as *mut _,
10723 &pb as *const _ as *mut _,
10724 &pw as *const _ as *mut _,
10725 &pr as *const _ as *mut _,
10726 &pdst as *const _ as *mut _,
10727 &pq as *const _ as *mut _,
10728 &pd as *const _ as *mut _,
10729 &nc as *const _ as *mut _,
10730 &e as *const _ as *mut _,
10731 ];
10732 unsafe {
10733 self.launch_pdl(
10734 "rms_pre_add_rms_norm_q8z_f32",
10735 (nrows as u32, 1, 1),
10736 (rms_block(), 1, 1),
10737 &mut ps,
10738 )?;
10739 }
10740 }
10741 return Ok((out_q, out_d));
10742 }
10743 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10744 let cfg = LaunchConfig {
10745 grid_dim: (nrows as u32, 1, 1),
10746 block_dim: (rms_block(), 1, 1),
10747 shared_mem_bytes: 0,
10748 };
10749 let __s_b2 = self.gpu.stream();
10750 let mut b2 = __s_b2.launch_builder(&f);
10751 b2.arg(a)
10752 .arg(wa)
10753 .arg(b)
10754 .arg(w)
10755 .arg(&mut *res)
10756 .arg(&mut *dst)
10757 .arg(&mut out_q)
10758 .arg(&mut out_d)
10759 .arg(&nc)
10760 .arg(&e);
10761 unsafe {
10762 b2.launch(cfg)?;
10763 }
10764 Ok((out_q, out_d))
10765 }
10766
10767 #[allow(clippy::too_many_arguments)]
10771 pub fn rms_pre_add_rms_norm_q8z_into(
10772 &self,
10773 a: &CudaSlice<f32>,
10774 wa: &CudaSlice<f32>,
10775 b: &CudaSlice<f32>,
10776 w: &CudaSlice<f32>,
10777 res: &mut CudaSlice<f32>,
10778 dst: &mut CudaSlice<f32>,
10779 ncols: usize,
10780 nrows: usize,
10781 eps: f32,
10782 out_q: &mut CudaSlice<i8>,
10783 out_d: &mut CudaSlice<f32>,
10784 ) -> Result<(), Box<dyn std::error::Error>> {
10785 debug_assert!(ncols % 128 == 0);
10786 let (nc, e) = (ncols as i32, eps);
10787 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10788 let cfg = LaunchConfig {
10789 grid_dim: (nrows as u32, 1, 1),
10790 block_dim: (rms_block(), 1, 1),
10791 shared_mem_bytes: 0,
10792 };
10793 let __s_b = self.gpu.stream();
10794 let mut b2 = __s_b.launch_builder(&f);
10795 b2.arg(a)
10796 .arg(wa)
10797 .arg(b)
10798 .arg(w)
10799 .arg(&mut *res)
10800 .arg(&mut *dst)
10801 .arg(&mut *out_q)
10802 .arg(&mut *out_d)
10803 .arg(&nc)
10804 .arg(&e);
10805 unsafe {
10806 b2.launch(cfg)?;
10807 }
10808 Ok(())
10809 }
10810
10811 #[allow(clippy::too_many_arguments)]
10814 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10815 &self,
10816 a: &CudaSlice<f32>,
10817 wa: &CudaSlice<f32>,
10818 b_in: &CudaSlice<f32>,
10819 c: f32,
10820 w: &CudaSlice<f32>,
10821 res: &mut CudaSlice<f32>,
10822 ncols: usize,
10823 nrows: usize,
10824 eps: f32,
10825 out_q: &mut CudaSlice<i8>,
10826 out_d: &mut CudaSlice<f32>,
10827 ) -> Result<(), Box<dyn std::error::Error>> {
10828 debug_assert!(ncols % 128 == 0);
10829 let (nc, e2) = (ncols as i32, eps);
10830 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10831 let cfg = LaunchConfig {
10832 grid_dim: (nrows as u32, 1, 1),
10833 block_dim: (rms_block(), 1, 1),
10834 shared_mem_bytes: 0,
10835 };
10836 let __s_b = self.gpu.stream();
10837 let mut b2 = __s_b.launch_builder(&f);
10838 b2.arg(a)
10839 .arg(wa)
10840 .arg(b_in)
10841 .arg(&c)
10842 .arg(w)
10843 .arg(&mut *res)
10844 .arg(&mut *out_q)
10845 .arg(&mut *out_d)
10846 .arg(&nc)
10847 .arg(&e2);
10848 unsafe {
10849 b2.launch(cfg)?;
10850 }
10851 Ok(())
10852 }
10853
10854 pub fn g4_pnfold_on() -> bool {
10862 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10863 *ON.get_or_init(|| {
10864 std::env::var("MEMRA_G4_PNFOLD")
10865 .map(|v| v != "0")
10866 .unwrap_or(true)
10867 })
10868 }
10869
10870 pub fn build_q4_out_concat3(
10874 &self,
10875 w0: &crate::model::GpuTensor,
10876 w1: &crate::model::GpuTensor,
10877 w2: &crate::model::GpuTensor,
10878 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10879 use crate::model::GpuTensor;
10880 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10881 match w {
10882 GpuTensor::Quant {
10883 qtype,
10884 row_bytes,
10885 rp,
10886 ..
10887 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10888 _ => None,
10889 }
10890 };
10891 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10892 else {
10893 return Ok(None);
10894 };
10895 if rb0 != rb1
10896 || rb0 != rb2
10897 || w0.in_features() != w1.in_features()
10898 || w0.in_features() != w2.in_features()
10899 {
10900 return Ok(None);
10901 }
10902 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10903 match w {
10904 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10905 _ => unreachable!(),
10906 }
10907 }
10908 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10909 let total = rb0 * (o0 + o1 + o2);
10910 let mut cat = self.alloc_u8(total)?;
10911 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10912 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10913 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10914 Ok(Some(GpuTensor::Quant {
10915 bytes: cat,
10916 qtype: QT_Q4_0,
10917 row_bytes: rb0,
10918 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10919 scale: 1.0,
10920 rp: false,
10921 #[cfg(memra_cutlass)]
10922 cutlass: None,
10923 fp8: None,
10924 blk: None,
10925 rp4: None,
10926 f16: None,
10927 }))
10928 }
10929
10930 fn full_width_rope_only(
10948 kernel: &str,
10949 n_rot: usize,
10950 head_dim: usize,
10951 ) -> Result<(), Box<dyn std::error::Error>> {
10952 if n_rot == head_dim {
10953 return Ok(());
10954 }
10955 Err(format!(
10956 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10957 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10958 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10959 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10960 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10961 )
10962 .into())
10963 }
10964
10965 #[allow(clippy::too_many_arguments)]
10969 pub fn rms_norm_qkv_rope_cat(
10970 &self,
10971 qkv: &CudaSlice<f32>,
10972 wq: &CudaSlice<f32>,
10973 wk: &CudaSlice<f32>,
10974 wv: &CudaSlice<f32>,
10975 q: &mut CudaSlice<f32>,
10976 k: &mut CudaSlice<f32>,
10977 v: &mut CudaSlice<f32>,
10978 head_dim: usize,
10979 n_rot: usize,
10980 rq: usize,
10981 rk: usize,
10982 pos: &CudaSlice<i32>,
10983 nh_q: usize,
10984 nh_k: usize,
10985 base: f32,
10986 freq_scale: f32,
10987 ff: Option<&CudaSlice<f32>>,
10988 eps: f32,
10989 ) -> Result<(), Box<dyn std::error::Error>> {
10990 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10991 let rows = rq + rk + rk;
10992 let theta_scale = base.powf(-2.0 / head_dim as f32);
10993 let (nc, rqi, rki, nhq, nhk) = (
10994 head_dim as i32,
10995 rq as i32,
10996 rk as i32,
10997 nh_q as i32,
10998 nh_k as i32,
10999 );
11000 if Self::pdl_on() {
11001 use cudarc::driver::{DevicePtr, DevicePtrMut};
11002 let s = &self.gpu.stream();
11003 let (pqkv, _g0) = qkv.device_ptr(s);
11004 let (pwq, _g1) = wq.device_ptr(s);
11005 let (pwk, _g2) = wk.device_ptr(s);
11006 let (pwv, _g3) = wv.device_ptr(s);
11007 let (pq, _g4) = q.device_ptr_mut(s);
11008 let (pk, _g5) = k.device_ptr_mut(s);
11009 let (pv, _g6) = v.device_ptr_mut(s);
11010 let (ppos, _g7) = pos.device_ptr(s);
11011 let (pff, _g8) = match ff {
11012 Some(t) => {
11013 let (p, g) = t.device_ptr(s);
11014 (p, Some(g))
11015 }
11016 None => (0, None),
11017 };
11018 let mut ps = [
11019 &pqkv as *const _ as *mut std::ffi::c_void,
11020 &pwq as *const _ as *mut _,
11021 &pwk as *const _ as *mut _,
11022 &pwv as *const _ as *mut _,
11023 &pq as *const _ as *mut _,
11024 &pk as *const _ as *mut _,
11025 &pv as *const _ as *mut _,
11026 &nc as *const _ as *mut _,
11027 &rqi as *const _ as *mut _,
11028 &rki as *const _ as *mut _,
11029 &ppos as *const _ as *mut _,
11030 &nhq as *const _ as *mut _,
11031 &nhk as *const _ as *mut _,
11032 &theta_scale as *const _ as *mut _,
11033 &freq_scale as *const _ as *mut _,
11034 &pff as *const _ as *mut _,
11035 &eps as *const _ as *mut _,
11036 ];
11037 unsafe {
11038 self.launch_pdl(
11039 "rms_norm_qkv_rope_cat_f32",
11040 (rows as u32, 1, 1),
11041 (rms_block(), 1, 1),
11042 &mut ps,
11043 )?;
11044 }
11045 return Ok(());
11046 }
11047 let f = self.func("rms_norm_qkv_rope_cat_f32");
11048 let cfg = LaunchConfig {
11049 grid_dim: (rows as u32, 1, 1),
11050 block_dim: (rms_block(), 1, 1),
11051 shared_mem_bytes: 0,
11052 };
11053 let __s_b = self.gpu.stream();
11054 let mut b = __s_b.launch_builder(&f);
11055 match ff {
11056 Some(t) => {
11057 b.arg(qkv)
11058 .arg(wq)
11059 .arg(wk)
11060 .arg(wv)
11061 .arg(&mut *q)
11062 .arg(&mut *k)
11063 .arg(&mut *v)
11064 .arg(&nc)
11065 .arg(&rqi)
11066 .arg(&rki)
11067 .arg(pos)
11068 .arg(&nhq)
11069 .arg(&nhk)
11070 .arg(&theta_scale)
11071 .arg(&freq_scale)
11072 .arg(t)
11073 .arg(&eps);
11074 unsafe {
11075 b.launch(cfg)?;
11076 }
11077 }
11078 None => {
11079 let null: u64 = 0;
11080 b.arg(qkv)
11081 .arg(wq)
11082 .arg(wk)
11083 .arg(wv)
11084 .arg(&mut *q)
11085 .arg(&mut *k)
11086 .arg(&mut *v)
11087 .arg(&nc)
11088 .arg(&rqi)
11089 .arg(&rki)
11090 .arg(pos)
11091 .arg(&nhq)
11092 .arg(&nhk)
11093 .arg(&theta_scale)
11094 .arg(&freq_scale)
11095 .arg(&null)
11096 .arg(&eps);
11097 unsafe {
11098 b.launch(cfg)?;
11099 }
11100 }
11101 }
11102 Ok(())
11103 }
11104
11105 #[allow(clippy::too_many_arguments)]
11109 pub fn rms_norm_qkv_rope(
11110 &self,
11111 q0: &CudaSlice<f32>,
11112 k0: &CudaSlice<f32>,
11113 v0: &CudaSlice<f32>,
11114 wq: &CudaSlice<f32>,
11115 wk: &CudaSlice<f32>,
11116 wv: &CudaSlice<f32>,
11117 q: &mut CudaSlice<f32>,
11118 k: &mut CudaSlice<f32>,
11119 v: &mut CudaSlice<f32>,
11120 head_dim: usize,
11121 n_rot: usize,
11122 rq: usize,
11123 rk: usize,
11124 pos: &CudaSlice<i32>,
11125 nh_q: usize,
11126 nh_k: usize,
11127 base: f32,
11128 freq_scale: f32,
11129 ff: Option<&CudaSlice<f32>>,
11130 eps: f32,
11131 ) -> Result<(), Box<dyn std::error::Error>> {
11132 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
11133 let f = self.func("rms_norm_qkv_rope_f32");
11134 let rows = rq + rk + rk; let cfg = LaunchConfig {
11136 grid_dim: (rows as u32, 1, 1),
11137 block_dim: (rms_block(), 1, 1),
11138 shared_mem_bytes: 0,
11139 };
11140 let theta_scale = base.powf(-2.0 / head_dim as f32);
11141 let (nc, rqi, rki, nhq, nhk) = (
11142 head_dim as i32,
11143 rq as i32,
11144 rk as i32,
11145 nh_q as i32,
11146 nh_k as i32,
11147 );
11148 let __s_b = self.gpu.stream();
11149 let mut b = __s_b.launch_builder(&f);
11150 match ff {
11151 Some(t) => {
11152 b.arg(q0)
11153 .arg(k0)
11154 .arg(v0)
11155 .arg(wq)
11156 .arg(wk)
11157 .arg(wv)
11158 .arg(&mut *q)
11159 .arg(&mut *k)
11160 .arg(&mut *v)
11161 .arg(&nc)
11162 .arg(&rqi)
11163 .arg(&rki)
11164 .arg(pos)
11165 .arg(&nhq)
11166 .arg(&nhk)
11167 .arg(&theta_scale)
11168 .arg(&freq_scale)
11169 .arg(t)
11170 .arg(&eps);
11171 unsafe {
11172 b.launch(cfg)?;
11173 }
11174 }
11175 None => {
11176 let null: u64 = 0;
11177 b.arg(q0)
11178 .arg(k0)
11179 .arg(v0)
11180 .arg(wq)
11181 .arg(wk)
11182 .arg(wv)
11183 .arg(&mut *q)
11184 .arg(&mut *k)
11185 .arg(&mut *v)
11186 .arg(&nc)
11187 .arg(&rqi)
11188 .arg(&rki)
11189 .arg(pos)
11190 .arg(&nhq)
11191 .arg(&nhk)
11192 .arg(&theta_scale)
11193 .arg(&freq_scale)
11194 .arg(&null)
11195 .arg(&eps);
11196 unsafe {
11197 b.launch(cfg)?;
11198 }
11199 }
11200 }
11201 Ok(())
11202 }
11203
11204 #[allow(clippy::too_many_arguments)]
11210 pub fn rms_norm_qkv_rope_append_dc(
11211 &self,
11212 q0: &CudaSlice<f32>,
11213 k0: &CudaSlice<f32>,
11214 v0: &CudaSlice<f32>,
11215 wq: &CudaSlice<f32>,
11216 wk: &CudaSlice<f32>,
11217 wv: &CudaSlice<f32>,
11218 q: &mut CudaSlice<f32>,
11219 k: &mut CudaSlice<f32>,
11220 v: &mut CudaSlice<f32>,
11221 head_dim: usize,
11222 n_rot: usize,
11223 rq: usize,
11224 rk: usize,
11225 pos: &CudaSlice<i32>,
11226 nh_q: usize,
11227 nh_k: usize,
11228 base: f32,
11229 freq_scale: f32,
11230 ff: Option<&CudaSlice<f32>>,
11231 eps: f32,
11232 kc: &mut CudaSlice<u8>,
11233 vc: &mut CudaSlice<u8>,
11234 t_dev: &CudaSlice<i32>,
11235 k_tok_bytes: usize,
11236 v_tok_bytes: usize,
11237 g: bool,
11238 ) -> Result<(), Box<dyn std::error::Error>> {
11239 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
11240 let rows = rq + rk + rk;
11241 let theta_scale = base.powf(-2.0 / head_dim as f32);
11242 let (nc, rqi, rki, nhq, nhk) = (
11243 head_dim as i32,
11244 rq as i32,
11245 rk as i32,
11246 nh_q as i32,
11247 nh_k as i32,
11248 );
11249 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11250 if Self::pdl_on() && Self::pdl_wb_on() {
11251 use cudarc::driver::{DevicePtr, DevicePtrMut};
11252 let s = &self.gpu.stream();
11253 let (p0, _a0) = q0.device_ptr(s);
11254 let (p1, _a1) = k0.device_ptr(s);
11255 let (p2, _a2) = v0.device_ptr(s);
11256 let (pwq, _a3) = wq.device_ptr(s);
11257 let (pwk, _a4) = wk.device_ptr(s);
11258 let (pwv, _a5) = wv.device_ptr(s);
11259 let (pq, _a6) = q.device_ptr_mut(s);
11260 let (pk, _a7) = k.device_ptr_mut(s);
11261 let (pv, _a8) = v.device_ptr_mut(s);
11262 let (pp, _a9) = pos.device_ptr(s);
11263 let pff: u64 = match ff {
11264 Some(t) => {
11265 let (p, _gg) = t.device_ptr(s);
11266 p as u64
11267 }
11268 None => 0,
11269 };
11270 let (pkc, _a10) = kc.device_ptr_mut(s);
11271 let (pvc, _a11) = vc.device_ptr_mut(s);
11272 let (pt, _a12) = t_dev.device_ptr(s);
11273 let mut ps = [
11274 &p0 as *const _ as *mut std::ffi::c_void,
11275 &p1 as *const _ as *mut _,
11276 &p2 as *const _ as *mut _,
11277 &pwq as *const _ as *mut _,
11278 &pwk as *const _ as *mut _,
11279 &pwv as *const _ as *mut _,
11280 &pq as *const _ as *mut _,
11281 &pk as *const _ as *mut _,
11282 &pv as *const _ as *mut _,
11283 &nc as *const _ as *mut _,
11284 &rqi as *const _ as *mut _,
11285 &rki as *const _ as *mut _,
11286 &pp as *const _ as *mut _,
11287 &nhq as *const _ as *mut _,
11288 &nhk as *const _ as *mut _,
11289 &theta_scale as *const _ as *mut _,
11290 &freq_scale as *const _ as *mut _,
11291 &pff as *const _ as *mut _,
11292 &eps as *const _ as *mut _,
11293 &pkc as *const _ as *mut _,
11294 &pvc as *const _ as *mut _,
11295 &pt as *const _ as *mut _,
11296 &ktb as *const _ as *mut _,
11297 &vtb as *const _ as *mut _,
11298 ];
11299 unsafe {
11300 self.launch_pdl_flash(
11301 g,
11302 "rms_norm_qkv_rope_append_dc_f32",
11303 (rows as u32, 1, 1),
11304 (rms_block(), 1, 1),
11305 0,
11306 &mut ps,
11307 )?;
11308 }
11309 return Ok(());
11310 }
11311 let f = if g {
11312 self.func_g("rms_norm_qkv_rope_append_dc_f32")
11313 } else {
11314 self.func("rms_norm_qkv_rope_append_dc_f32")
11315 };
11316 let cfg = LaunchConfig {
11317 grid_dim: (rows as u32, 1, 1),
11318 block_dim: (rms_block(), 1, 1),
11319 shared_mem_bytes: 0,
11320 };
11321 let __s_b = self.gpu.stream();
11322 let mut b = __s_b.launch_builder(&f);
11323 match ff {
11324 Some(t) => {
11325 b.arg(q0)
11326 .arg(k0)
11327 .arg(v0)
11328 .arg(wq)
11329 .arg(wk)
11330 .arg(wv)
11331 .arg(&mut *q)
11332 .arg(&mut *k)
11333 .arg(&mut *v)
11334 .arg(&nc)
11335 .arg(&rqi)
11336 .arg(&rki)
11337 .arg(pos)
11338 .arg(&nhq)
11339 .arg(&nhk)
11340 .arg(&theta_scale)
11341 .arg(&freq_scale)
11342 .arg(t)
11343 .arg(&eps)
11344 .arg(&mut *kc)
11345 .arg(&mut *vc)
11346 .arg(t_dev)
11347 .arg(&ktb)
11348 .arg(&vtb);
11349 unsafe {
11350 b.launch(cfg)?;
11351 }
11352 }
11353 None => {
11354 let null: u64 = 0;
11355 b.arg(q0)
11356 .arg(k0)
11357 .arg(v0)
11358 .arg(wq)
11359 .arg(wk)
11360 .arg(wv)
11361 .arg(&mut *q)
11362 .arg(&mut *k)
11363 .arg(&mut *v)
11364 .arg(&nc)
11365 .arg(&rqi)
11366 .arg(&rki)
11367 .arg(pos)
11368 .arg(&nhq)
11369 .arg(&nhk)
11370 .arg(&theta_scale)
11371 .arg(&freq_scale)
11372 .arg(&null)
11373 .arg(&eps)
11374 .arg(&mut *kc)
11375 .arg(&mut *vc)
11376 .arg(t_dev)
11377 .arg(&ktb)
11378 .arg(&vtb);
11379 unsafe {
11380 b.launch(cfg)?;
11381 }
11382 }
11383 }
11384 Ok(())
11385 }
11386
11387 #[allow(clippy::too_many_arguments)]
11395 pub fn rms_norm_qkv_rope_append(
11396 &self,
11397 q0: &CudaSlice<f32>,
11398 k0: &CudaSlice<f32>,
11399 v0: &CudaSlice<f32>,
11400 wq: &CudaSlice<f32>,
11401 wk: &CudaSlice<f32>,
11402 wv: &CudaSlice<f32>,
11403 q: &mut CudaSlice<f32>,
11404 k: &mut CudaSlice<f32>,
11405 v: &mut CudaSlice<f32>,
11406 head_dim: usize,
11407 n_rot: usize,
11408 rq: usize,
11409 rk: usize,
11410 pos: &CudaSlice<i32>,
11411 nh_q: usize,
11412 nh_k: usize,
11413 base: f32,
11414 freq_scale: f32,
11415 ff: Option<&CudaSlice<f32>>,
11416 eps: f32,
11417 kc: &mut CudaSlice<u8>,
11418 vc: &mut CudaSlice<u8>,
11419 t: usize,
11420 k_tok_bytes: usize,
11421 v_tok_bytes: usize,
11422 g: bool,
11423 ) -> Result<(), Box<dyn std::error::Error>> {
11424 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11425 let rows = rq + rk + rk;
11426 let theta_scale = base.powf(-2.0 / head_dim as f32);
11427 let (nc, rqi, rki, nhq, nhk) = (
11428 head_dim as i32,
11429 rq as i32,
11430 rk as i32,
11431 nh_q as i32,
11432 nh_k as i32,
11433 );
11434 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11435 let ti = t as i32;
11436 if Self::pdl_on() && Self::pdl_wb_on() {
11437 use cudarc::driver::{DevicePtr, DevicePtrMut};
11438 let s = &self.gpu.stream();
11439 let (p0, _a0) = q0.device_ptr(s);
11440 let (p1, _a1) = k0.device_ptr(s);
11441 let (p2, _a2) = v0.device_ptr(s);
11442 let (pwq, _a3) = wq.device_ptr(s);
11443 let (pwk, _a4) = wk.device_ptr(s);
11444 let (pwv, _a5) = wv.device_ptr(s);
11445 let (pq, _a6) = q.device_ptr_mut(s);
11446 let (pk, _a7) = k.device_ptr_mut(s);
11447 let (pv, _a8) = v.device_ptr_mut(s);
11448 let (pp, _a9) = pos.device_ptr(s);
11449 let pff: u64 = match ff {
11450 Some(t) => {
11451 let (p, _gg) = t.device_ptr(s);
11452 p as u64
11453 }
11454 None => 0,
11455 };
11456 let (pkc, _a10) = kc.device_ptr_mut(s);
11457 let (pvc, _a11) = vc.device_ptr_mut(s);
11458 let mut ps = [
11459 &p0 as *const _ as *mut std::ffi::c_void,
11460 &p1 as *const _ as *mut _,
11461 &p2 as *const _ as *mut _,
11462 &pwq as *const _ as *mut _,
11463 &pwk as *const _ as *mut _,
11464 &pwv as *const _ as *mut _,
11465 &pq as *const _ as *mut _,
11466 &pk as *const _ as *mut _,
11467 &pv as *const _ as *mut _,
11468 &nc as *const _ as *mut _,
11469 &rqi as *const _ as *mut _,
11470 &rki as *const _ as *mut _,
11471 &pp as *const _ as *mut _,
11472 &nhq as *const _ as *mut _,
11473 &nhk as *const _ as *mut _,
11474 &theta_scale as *const _ as *mut _,
11475 &freq_scale as *const _ as *mut _,
11476 &pff as *const _ as *mut _,
11477 &eps as *const _ as *mut _,
11478 &pkc as *const _ as *mut _,
11479 &pvc as *const _ as *mut _,
11480 &ti as *const _ as *mut _,
11481 &ktb as *const _ as *mut _,
11482 &vtb as *const _ as *mut _,
11483 ];
11484 unsafe {
11485 self.launch_pdl_flash(
11486 g,
11487 "rms_norm_qkv_rope_append_f32",
11488 (rows as u32, 1, 1),
11489 (rms_block(), 1, 1),
11490 0,
11491 &mut ps,
11492 )?;
11493 }
11494 return Ok(());
11495 }
11496 let f = if g {
11497 self.func_g("rms_norm_qkv_rope_append_f32")
11498 } else {
11499 self.func("rms_norm_qkv_rope_append_f32")
11500 };
11501 let cfg = LaunchConfig {
11502 grid_dim: (rows as u32, 1, 1),
11503 block_dim: (rms_block(), 1, 1),
11504 shared_mem_bytes: 0,
11505 };
11506 let __s_b = self.gpu.stream();
11507 let mut b = __s_b.launch_builder(&f);
11508 let null: u64 = 0;
11509 b.arg(q0)
11510 .arg(k0)
11511 .arg(v0)
11512 .arg(wq)
11513 .arg(wk)
11514 .arg(wv)
11515 .arg(&mut *q)
11516 .arg(&mut *k)
11517 .arg(&mut *v)
11518 .arg(&nc)
11519 .arg(&rqi)
11520 .arg(&rki)
11521 .arg(pos)
11522 .arg(&nhq)
11523 .arg(&nhk)
11524 .arg(&theta_scale)
11525 .arg(&freq_scale);
11526 match ff {
11527 Some(t) => {
11528 b.arg(t);
11529 }
11530 None => {
11531 b.arg(&null);
11532 }
11533 }
11534 b.arg(&eps)
11535 .arg(&mut *kc)
11536 .arg(&mut *vc)
11537 .arg(&ti)
11538 .arg(&ktb)
11539 .arg(&vtb);
11540 unsafe {
11541 b.launch(cfg)?;
11542 }
11543 Ok(())
11544 }
11545
11546 pub fn add_q8_1(
11547 &self,
11548 a: &CudaSlice<f32>,
11549 b: &CudaSlice<f32>,
11550 res: &mut CudaSlice<f32>,
11551 ncols: usize,
11552 nrows: usize,
11553 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11554 debug_assert!(ncols % 128 == 0);
11555 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11556 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11557 let f = self.func("add_q8_1_f32");
11558 let cfg = LaunchConfig {
11559 grid_dim: (nrows as u32, 1, 1),
11560 block_dim: (rms_block(), 1, 1),
11561 shared_mem_bytes: 0,
11562 };
11563 let nc = ncols as i32;
11564 let __s_b2 = self.gpu.stream();
11565 let mut b2 = __s_b2.launch_builder(&f);
11566 b2.arg(a)
11567 .arg(b)
11568 .arg(&mut *res)
11569 .arg(&mut out_q)
11570 .arg(&mut out_d)
11571 .arg(&nc);
11572 unsafe {
11573 b2.launch(cfg)?;
11574 }
11575 Ok((out_q, out_d))
11576 }
11577
11578 pub fn rms_pre_add_q8_1(
11582 &self,
11583 a: &CudaSlice<f32>,
11584 wa: &CudaSlice<f32>,
11585 b: &CudaSlice<f32>,
11586 res: &mut CudaSlice<f32>,
11587 ncols: usize,
11588 nrows: usize,
11589 eps: f32,
11590 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11591 debug_assert!(ncols % 128 == 0);
11592 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11593 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11594 let f = self.func("rms_pre_add_q8_1_f32");
11595 let cfg = LaunchConfig {
11596 grid_dim: (nrows as u32, 1, 1),
11597 block_dim: (rms_block(), 1, 1),
11598 shared_mem_bytes: 0,
11599 };
11600 let (nc, ep) = (ncols as i32, eps);
11601 let __s_b2 = self.gpu.stream();
11602 let mut b2 = __s_b2.launch_builder(&f);
11603 b2.arg(a)
11604 .arg(wa)
11605 .arg(b)
11606 .arg(&mut *res)
11607 .arg(&mut out_q)
11608 .arg(&mut out_d)
11609 .arg(&nc)
11610 .arg(&ep);
11611 unsafe {
11612 b2.launch(cfg)?;
11613 }
11614 Ok((out_q, out_d))
11615 }
11616
11617 pub fn l2_v2_on(ncols: usize) -> bool {
11621 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11622 }
11623
11624 pub fn l2_norm_pp(
11625 &self,
11626 x: &CudaSlice<f32>,
11627 dst: &mut CudaSlice<f32>,
11628 dst16: Option<&mut CudaSlice<u8>>,
11629 ncols: usize,
11630 nrows: usize,
11631 eps: f32,
11632 ) -> Result<(), Box<dyn std::error::Error>> {
11633 if Self::l2_v2_on(ncols) {
11634 let f = self.func("l2_norm_pp_v2_f32");
11635 let rows_per_block = 8u32; let cfg = LaunchConfig {
11637 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11638 block_dim: (256, 1, 1),
11639 shared_mem_bytes: 0,
11640 };
11641 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11642 let d16: u64 = match dst16 {
11644 Some(d) => self.addr_u8(d),
11645 None => 0,
11646 };
11647 let __s_b = self.gpu.stream();
11648 let mut b = __s_b.launch_builder(&f);
11649 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11650 unsafe {
11651 b.launch(cfg)?;
11652 }
11653 return Ok(());
11654 }
11655 self.l2_norm(x, dst, ncols, nrows, eps)
11656 }
11657
11658 pub fn l2_norm(
11659 &self,
11660 x: &CudaSlice<f32>,
11661 dst: &mut CudaSlice<f32>,
11662 ncols: usize,
11663 nrows: usize,
11664 eps: f32,
11665 ) -> Result<(), Box<dyn std::error::Error>> {
11666 let f = self.func("l2_norm_f32");
11667 let cfg = LaunchConfig {
11668 grid_dim: (nrows as u32, 1, 1),
11669 block_dim: (256, 1, 1),
11670 shared_mem_bytes: 0,
11671 };
11672 let (nc, e) = (ncols as i32, eps);
11673 let __s_b = self.gpu.stream();
11674 let mut b = __s_b.launch_builder(&f);
11675 b.arg(x).arg(dst).arg(&nc).arg(&e);
11676 unsafe {
11677 b.launch(cfg)?;
11678 }
11679 Ok(())
11680 }
11681
11682 pub fn l2_norm_decode(
11688 &self,
11689 x: &CudaSlice<f32>,
11690 dst: &mut CudaSlice<f32>,
11691 ncols: usize,
11692 nrows: usize,
11693 eps: f32,
11694 ) -> Result<(), Box<dyn std::error::Error>> {
11695 let f = self.func("l2_norm_f32");
11696 let cfg = LaunchConfig {
11697 grid_dim: (nrows as u32, 1, 1),
11698 block_dim: (32, 1, 1),
11699 shared_mem_bytes: 0,
11700 };
11701 let (nc, e) = (ncols as i32, eps);
11702 let __s_b = self.gpu.stream();
11703 let mut b = __s_b.launch_builder(&f);
11704 b.arg(x).arg(dst).arg(&nc).arg(&e);
11705 unsafe {
11706 b.launch(cfg)?;
11707 }
11708 Ok(())
11709 }
11710
11711 pub fn rope_neox(
11713 &self,
11714 x: &mut CudaSlice<f32>,
11715 pos: &CudaSlice<i32>,
11716 head_dim: usize,
11717 n_dims: usize,
11718 n_heads: usize,
11719 n_tokens: usize,
11720 freq_base: f32,
11721 freq_scale: f32,
11722 ) -> Result<(), Box<dyn std::error::Error>> {
11723 let f = self.func("rope_neox_f32");
11724 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11725 let grid = (n_heads * n_tokens) as u32;
11726 let cfg = LaunchConfig {
11727 grid_dim: (grid, 1, 1),
11728 block_dim: ((head_dim / 2) as u32, 1, 1),
11729 shared_mem_bytes: 0,
11730 };
11731 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11732 let __s_b = self.gpu.stream();
11733 let mut b = __s_b.launch_builder(&f);
11734 b.arg(x)
11735 .arg(pos)
11736 .arg(&hd)
11737 .arg(&nd)
11738 .arg(&nh)
11739 .arg(&theta_scale)
11740 .arg(&freq_scale);
11741 unsafe {
11742 b.launch(cfg)?;
11743 }
11744 Ok(())
11745 }
11746
11747 pub fn rope_neox_ff(
11749 &self,
11750 x: &mut CudaSlice<f32>,
11751 pos: &CudaSlice<i32>,
11752 head_dim: usize,
11753 n_dims: usize,
11754 n_heads: usize,
11755 n_tokens: usize,
11756 freq_base: f32,
11757 freq_scale: f32,
11758 ff: &CudaSlice<f32>,
11759 ) -> Result<(), Box<dyn std::error::Error>> {
11760 let f = self.func("rope_neox_ff_f32");
11761 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11762 let grid = (n_heads * n_tokens) as u32;
11763 let cfg = LaunchConfig {
11764 grid_dim: (grid, 1, 1),
11765 block_dim: ((head_dim / 2) as u32, 1, 1),
11766 shared_mem_bytes: 0,
11767 };
11768 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11769 let __s_b = self.gpu.stream();
11770 let mut b = __s_b.launch_builder(&f);
11771 b.arg(x)
11772 .arg(pos)
11773 .arg(&hd)
11774 .arg(&nd)
11775 .arg(&nh)
11776 .arg(&theta_scale)
11777 .arg(&freq_scale)
11778 .arg(ff);
11779 unsafe {
11780 b.launch(cfg)?;
11781 }
11782 Ok(())
11783 }
11784
11785 #[allow(clippy::too_many_arguments)]
11787 pub fn rope_neox2(
11788 &self,
11789 q: &mut CudaSlice<f32>,
11790 k: &mut CudaSlice<f32>,
11791 pos: &CudaSlice<i32>,
11792 head_dim: usize,
11793 n_dims: usize,
11794 nh_q: usize,
11795 nh_k: usize,
11796 n_tokens: usize,
11797 freq_base: f32,
11798 freq_scale: f32,
11799 ff: Option<&CudaSlice<f32>>,
11800 ) -> Result<(), Box<dyn std::error::Error>> {
11801 let f = self.func("rope_neox2_f32");
11802 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11803 let grid = ((nh_q + nh_k) * n_tokens) as u32;
11804 let cfg = LaunchConfig {
11805 grid_dim: (grid, 1, 1),
11806 block_dim: ((head_dim / 2) as u32, 1, 1),
11807 shared_mem_bytes: 0,
11808 };
11809 let (hd, nd, nq, nk, nt) = (
11810 head_dim as i32,
11811 n_dims as i32,
11812 nh_q as i32,
11813 nh_k as i32,
11814 n_tokens as i32,
11815 );
11816 let __s_b = self.gpu.stream();
11817 let mut b = __s_b.launch_builder(&f);
11818 b.arg(q)
11819 .arg(k)
11820 .arg(pos)
11821 .arg(&hd)
11822 .arg(&nd)
11823 .arg(&nq)
11824 .arg(&nk)
11825 .arg(&nt)
11826 .arg(&theta_scale)
11827 .arg(&freq_scale);
11828 match ff {
11829 Some(ffv) => {
11830 b.arg(ffv);
11831 unsafe {
11832 b.launch(cfg)?;
11833 }
11834 }
11835 None => {
11836 let null: u64 = 0;
11837 b.arg(&null);
11838 unsafe {
11839 b.launch(cfg)?;
11840 }
11841 }
11842 }
11843 Ok(())
11844 }
11845
11846 pub fn gelu_tanh_mul(
11848 &self,
11849 gate: &CudaSlice<f32>,
11850 up: &CudaSlice<f32>,
11851 dst: &mut CudaSlice<f32>,
11852 n: usize,
11853 ) -> Result<(), Box<dyn std::error::Error>> {
11854 let f = self.func("gelu_tanh_mul_f32");
11855 let cfg = LaunchConfig::for_num_elems(n as u32);
11856 let ni = n as i32;
11857 let __s_b = self.gpu.stream();
11858 let mut b = __s_b.launch_builder(&f);
11859 b.arg(gate).arg(up).arg(dst).arg(&ni);
11860 unsafe {
11861 b.launch(cfg)?;
11862 }
11863 Ok(())
11864 }
11865
11866 pub fn silu_mul(
11867 &self,
11868 gate: &CudaSlice<f32>,
11869 up: &CudaSlice<f32>,
11870 dst: &mut CudaSlice<f32>,
11871 n: usize,
11872 ) -> Result<(), Box<dyn std::error::Error>> {
11873 let f = self.func("silu_mul_f32");
11874 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11876 let ni = n as i32;
11877 let __s_b = self.gpu.stream();
11878 let mut b = __s_b.launch_builder(&f);
11879 b.arg(gate).arg(up).arg(dst).arg(&ni);
11880 unsafe {
11881 b.launch(cfg)?;
11882 }
11883 Ok(())
11884 }
11885
11886 pub fn silu_mul_host_expf(
11888 &self,
11889 gate: &CudaSlice<f32>,
11890 up: &CudaSlice<f32>,
11891 dst: &mut CudaSlice<f32>,
11892 n: usize,
11893 ) -> Result<(), Box<dyn std::error::Error>> {
11894 let f = self.func("silu_mul_host_expf_f32");
11895 let cfg = LaunchConfig::for_num_elems(n as u32);
11896 let ni = n as i32;
11897 let __s_b = self.gpu.stream();
11898 let mut b = __s_b.launch_builder(&f);
11899 b.arg(gate).arg(up).arg(dst).arg(&ni);
11900 unsafe {
11901 b.launch(cfg)?;
11902 }
11903 Ok(())
11904 }
11905
11906 pub fn silu_clamped_mul_host_expf(
11908 &self,
11909 gate: &CudaSlice<f32>,
11910 up: &CudaSlice<f32>,
11911 limit: f32,
11912 dst: &mut CudaSlice<f32>,
11913 n: usize,
11914 ) -> Result<(), Box<dyn std::error::Error>> {
11915 if !limit.is_finite() || limit <= 0.0 {
11916 return Err(
11917 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11918 );
11919 }
11920 let f = self.func("silu_clamped_mul_host_expf_f32");
11921 let cfg = LaunchConfig::for_num_elems(n as u32);
11922 let ni = n as i32;
11923 let __s_b = self.gpu.stream();
11924 let mut b = __s_b.launch_builder(&f);
11925 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11926 unsafe {
11927 b.launch(cfg)?;
11928 }
11929 Ok(())
11930 }
11931
11932 pub fn silu_mul_f16out(
11935 &self,
11936 gate: &CudaSlice<f32>,
11937 up: &CudaSlice<f32>,
11938 dst: &mut CudaSlice<f32>,
11939 dst16: &mut CudaSlice<u8>,
11940 n: usize,
11941 ) -> Result<(), Box<dyn std::error::Error>> {
11942 let f = self.func("silu_mul_f16out_f32");
11943 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11944 let ni = n as i32;
11945 let __s_b = self.gpu.stream();
11946 let mut b = __s_b.launch_builder(&f);
11947 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11948 unsafe {
11949 b.launch(cfg)?;
11950 }
11951 Ok(())
11952 }
11953
11954 pub fn silu_mul_scaled(
11961 &self,
11962 gate: &CudaSlice<f32>,
11963 up: &CudaSlice<f32>,
11964 gs: f32,
11965 us: f32,
11966 dst: &mut CudaSlice<f32>,
11967 n: usize,
11968 ) -> Result<(), Box<dyn std::error::Error>> {
11969 let f = self.func("silu_mul_scaled_f32");
11970 let cfg = LaunchConfig::for_num_elems(n as u32);
11971 let ni = n as i32;
11972 let (gsf, usf) = (gs, us);
11973 let __s_b = self.gpu.stream();
11974 let mut b = __s_b.launch_builder(&f);
11975 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11976 unsafe {
11977 b.launch(cfg)?;
11978 }
11979 Ok(())
11980 }
11981
11982 #[allow(clippy::too_many_arguments)]
11986 pub fn swigluoai_mul_scaled(
11987 &self,
11988 gate: &CudaSlice<f32>,
11989 up: &CudaSlice<f32>,
11990 gs: f32,
11991 us: f32,
11992 alpha: f32,
11993 limit: f32,
11994 dst: &mut CudaSlice<f32>,
11995 n: usize,
11996 ) -> Result<(), Box<dyn std::error::Error>> {
11997 let f = self.func("swigluoai_mul_scaled_f32");
11998 let cfg = LaunchConfig::for_num_elems(n as u32);
11999 let ni = n as i32;
12000 let __s_b = self.gpu.stream();
12001 let mut b = __s_b.launch_builder(&f);
12002 b.arg(gate)
12003 .arg(up)
12004 .arg(&gs)
12005 .arg(&us)
12006 .arg(&alpha)
12007 .arg(&limit)
12008 .arg(dst)
12009 .arg(&ni);
12010 unsafe {
12011 b.launch(cfg)?;
12012 }
12013 Ok(())
12014 }
12015
12016 pub fn silu_mul_scaled_q8_1(
12024 &self,
12025 gate: &CudaSlice<f32>,
12026 up: &CudaSlice<f32>,
12027 gs: f32,
12028 us: f32,
12029 n: usize,
12030 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12031 let f = self.func("silu_mul_scaled_q8_1");
12032 let nblk = n / 32;
12033 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);
12037 let (gsf, usf, ni) = (gs, us, n as i32);
12038 let __s_b = self.gpu.stream();
12039 let mut b = __s_b.launch_builder(&f);
12040 b.arg(gate)
12041 .arg(up)
12042 .arg(&gsf)
12043 .arg(&usf)
12044 .arg(&mut aq)
12045 .arg(&mut ad)
12046 .arg(&ni);
12047 unsafe {
12048 b.launch(cfg)?;
12049 }
12050 Ok((aq, ad))
12051 }
12052
12053 pub fn add(
12054 &self,
12055 a: &CudaSlice<f32>,
12056 b_in: &CudaSlice<f32>,
12057 dst: &mut CudaSlice<f32>,
12058 n: usize,
12059 ) -> Result<(), Box<dyn std::error::Error>> {
12060 let f = self.func("add_f32");
12061 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
12063 let ni = n as i32;
12064 let __s_bld = self.gpu.stream();
12065 let mut bld = __s_bld.launch_builder(&f);
12066 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
12067 unsafe {
12068 bld.launch(cfg)?;
12069 }
12070 Ok(())
12071 }
12072
12073 pub fn mul(
12074 &self,
12075 a: &CudaSlice<f32>,
12076 b_in: &CudaSlice<f32>,
12077 dst: &mut CudaSlice<f32>,
12078 n: usize,
12079 ) -> Result<(), Box<dyn std::error::Error>> {
12080 let f = self.func("mul_f32");
12081 let cfg = LaunchConfig::for_num_elems(n as u32);
12082 let ni = n as i32;
12083 let __s_bld = self.gpu.stream();
12084 let mut bld = __s_bld.launch_builder(&f);
12085 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
12086 unsafe {
12087 bld.launch(cfg)?;
12088 }
12089 Ok(())
12090 }
12091
12092 pub fn matmul(
12095 &self,
12096 w: &crate::model::GpuTensor,
12097 x: &CudaSlice<f32>,
12098 m: usize,
12099 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12100 use crate::model::GpuTensor;
12101 let in_f = w.in_features();
12102 let out_f = w.out_features();
12103 #[allow(non_snake_case)]
12111 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
12114 usize::MAX
12115 } else {
12116 16usize
12117 };
12118
12119 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
12144 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
12145 return Ok(y);
12146 }
12147 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
12154 return Ok(y);
12155 }
12156 if let Some(y) = self.try_f16_gemm(w, x, m)? {
12159 return Ok(y);
12160 }
12161 }
12162 if let GpuTensor::Quant { qtype, .. } = w {
12177 if *qtype == QT_F8_E4M3_BLK {
12178 if m >= GEMM_M_THRESHOLD {
12179 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
12180 return Ok(y);
12181 }
12182 }
12183 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12184 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12185 return Ok(y);
12186 }
12187 }
12188 }
12189 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
12190 return self.qmatvec_mmq(w, x, m);
12191 }
12192 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
12193 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12194 return self.qmatvec_gemm(w, &aq, &ad, m);
12195 }
12196 if m >= GEMM_M_THRESHOLD {
12199 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
12200 return Ok(y);
12201 }
12202 }
12203 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
12207 if m == 1 && fast {
12212 if let GpuTensor::Quant {
12213 bytes,
12214 qtype,
12215 row_bytes,
12216 rp,
12217 rp4,
12218 scale,
12219 ..
12220 } = w
12221 {
12222 if self.mmvq_supports(*qtype) {
12223 let (bytes, rp) = match rp4 {
12227 Some(m4) => (m4, true),
12228 None => (bytes, *rp),
12229 };
12230 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12231 return self.qmatvec_mmvq(
12232 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
12233 );
12234 }
12235 }
12236 }
12237 if (2..=16).contains(&m)
12253 && fast
12254 && std::env::var("MEMRA_NO_BATCHED").is_err()
12255 && (m <= 4 || Self::b8_enabled())
12256 {
12257 let m_ok = m <= 8
12267 || matches!(w, GpuTensor::Quant { qtype, .. }
12268 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
12269 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
12270 if m_ok {
12271 if let GpuTensor::Quant {
12272 bytes,
12273 qtype,
12274 row_bytes,
12275 rp,
12276 rp4,
12277 ..
12278 } = w
12279 {
12280 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
12281 let (bytes, rp) = match rp4 {
12282 Some(m4) => (m4, true),
12283 None => (bytes, *rp),
12284 };
12285 let mcols = Self::batched_mcols(m);
12286 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12287 let mut y = self.qmatvec_mmvq_batched(
12288 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
12289 )?;
12290 if let GpuTensor::Quant { scale, .. } = w {
12291 if *scale != 1.0 {
12292 self.scale_inplace(&mut y, *scale, m * out_f)?;
12293 }
12294 }
12295 return Ok(y);
12296 }
12297 }
12298 }
12299 }
12300 if fast {
12306 if let GpuTensor::Quant {
12307 bytes,
12308 qtype,
12309 row_bytes,
12310 scale,
12311 ..
12312 } = w
12313 {
12314 if *qtype == QT_F8_E4M3 {
12315 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12316 return self.qmatvec_mmvq(
12317 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
12318 );
12319 }
12320 }
12321 }
12322 let mut y = match w {
12323 GpuTensor::Quant {
12324 bytes,
12325 qtype,
12326 row_bytes,
12327 ..
12328 } if fast && *qtype == QT_Q8_0 => {
12329 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12330 }
12331 GpuTensor::Quant {
12332 bytes,
12333 qtype,
12334 row_bytes,
12335 ..
12336 } if fast && *qtype == QT_Q4_K => {
12337 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12338 }
12339 GpuTensor::Quant {
12340 bytes,
12341 qtype,
12342 row_bytes,
12343 ..
12344 } if fast && *qtype == QT_Q6_K => {
12345 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12346 }
12347 GpuTensor::Quant {
12348 bytes,
12349 qtype,
12350 row_bytes,
12351 ..
12352 } if fast && *qtype == QT_Q5_K => {
12353 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12354 }
12355 GpuTensor::Quant {
12356 bytes,
12357 qtype,
12358 row_bytes,
12359 ..
12360 } if fast && *qtype == QT_Q3_K => {
12361 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12362 }
12363 GpuTensor::Quant {
12364 bytes,
12365 qtype,
12366 row_bytes,
12367 rp,
12368 ..
12369 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
12370 if *rp {
12371 "qmatvec_nvfp4_dp4a_rp"
12372 } else {
12373 "qmatvec_nvfp4_dp4a"
12374 },
12375 &bytes.slice(0..bytes.len()),
12376 x,
12377 m,
12378 in_f,
12379 out_f,
12380 *row_bytes,
12381 )?,
12382 GpuTensor::Quant {
12386 bytes,
12387 qtype,
12388 row_bytes,
12389 ..
12390 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12391 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12392 }
12393 GpuTensor::Quant {
12398 bytes,
12399 qtype,
12400 row_bytes,
12401 rp,
12402 ..
12403 } =>
12404 {
12407 self.qmatvec(
12408 bytes,
12409 x,
12410 m,
12411 in_f,
12412 out_f,
12413 if *rp && *qtype == QT_NVFP4 {
12414 QT_NVFP4_RP
12415 } else {
12416 *qtype
12417 },
12418 *row_bytes,
12419 )?
12420 }
12421 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12422 GpuTensor::FloatBf16 { data, .. } => {
12425 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12432 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12433 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12434 y
12435 } else {
12436 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12437 }
12438 }
12439 };
12440 if let GpuTensor::Quant { scale, .. } = w {
12442 if *scale != 1.0 {
12443 self.scale_inplace(&mut y, *scale, m * out_f)?;
12444 }
12445 }
12446 Ok(y)
12447 }
12448
12449 pub fn stage_a_raw_needed() -> bool {
12459 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12460 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12461 }
12462
12463 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12466 use crate::model::GpuTensor;
12467 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12468 return false;
12469 }
12470 match w {
12471 GpuTensor::Quant { qtype, .. } => {
12478 matches!(
12479 *qtype,
12480 QT_Q8_0
12481 | QT_Q4_K
12482 | QT_Q6_K
12483 | QT_Q5_K
12484 | QT_Q3_K
12485 | QT_NVFP4
12486 | QT_F8_E4M3
12487 | QT_F8_E4M3_BLK
12488 | QT_Q4_0
12489 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12490 }
12491 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12492 }
12493 }
12494
12495 pub fn matmul_pre(
12500 &self,
12501 w: &crate::model::GpuTensor,
12502 aq: &CudaSlice<i8>,
12503 ad: &CudaSlice<f32>,
12504 x_fallback: &CudaSlice<f32>,
12505 m: usize,
12506 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12507 use crate::model::GpuTensor;
12508 let x_raw_ok = x_fallback.len() >= m * w.in_features();
12514 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12517 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12518 return Ok(y);
12519 }
12520 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12523 return Ok(y);
12524 }
12525 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12527 return Ok(y);
12528 }
12529 }
12530 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12536 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12537 return Ok(y);
12538 }
12539 }
12540 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12541 return Ok(y);
12542 }
12543 if m >= 16
12548 && w.out_features() >= 128
12549 && self.mmq_supports(w)
12550 && !self.verify_exact_on()
12551 && x_raw_ok
12552 {
12553 return self.qmatvec_mmq(w, x_fallback, m);
12554 }
12555 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12558 if let Some(y) =
12559 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12560 {
12561 return Ok(y);
12562 }
12563 }
12564 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12567 return self.qmatvec_gemm(w, aq, ad, m);
12568 }
12569 if !self.uses_q8_1_fast(w) {
12588 if !x_raw_ok {
12589 return Err(format!(
12590 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12591 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12592 dropped the f32, so there is nothing to fall back to — pass the real f32 \
12593 activation (see Engine::rms_norm_decode, which is bit-identical to \
12594 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12595 x_fallback.len(),
12596 m,
12597 w.in_features(),
12598 m * w.in_features()
12599 )
12600 .into());
12601 }
12602 return self.matmul(w, x_fallback, m);
12603 }
12604 let in_f = w.in_features();
12605 let out_f = w.out_features();
12606 let (bytes, qtype, row_bytes, scale, rp) = match w {
12607 GpuTensor::Quant {
12608 bytes,
12609 qtype,
12610 row_bytes,
12611 scale,
12612 rp,
12613 ..
12614 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12615 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12616 };
12617 let (mbytes, mrp) = match w {
12620 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12621 _ => (bytes, rp),
12622 };
12623 if m == 1 && self.mmvq_supports(qtype) {
12627 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12628 }
12629 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12642 && std::env::var("MEMRA_NO_BATCHED").is_err()
12643 && (m <= 4 || Self::b8_enabled())
12644 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12648 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12649 {
12650 let mcols = Self::batched_mcols(m);
12651 return self.qmatvec_mmvq_batched(
12652 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12653 );
12654 }
12655 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12661 let (b2, r2) = if qtype == QT_Q4_0 {
12662 (mbytes, mrp)
12663 } else {
12664 (bytes, rp)
12665 };
12666 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12667 }
12668 let name = match qtype {
12669 QT_Q8_0 => "qmatvec_q8_0_dp4a",
12670 QT_Q4_K => "qmatvec_q4_K_dp4a",
12671 QT_Q6_K => "qmatvec_q6_K_dp4a",
12672 QT_Q5_K => "qmatvec_q5_K_dp4a",
12673 QT_Q3_K => "qmatvec_q3_K_dp4a",
12674 QT_NVFP4 => {
12675 if rp {
12676 "qmatvec_nvfp4_dp4a_rp"
12677 } else {
12678 "qmatvec_nvfp4_dp4a"
12679 }
12680 }
12681 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12682 _ => unreachable!(),
12683 };
12684 let f = self.func(name);
12685 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
12687 grid_dim: (out_f as u32, m as u32, 1),
12688 block_dim: (128, 1, 1),
12689 shared_mem_bytes: 0,
12690 };
12691 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12692 let __s_b = self.gpu.stream();
12693 let mut b = __s_b.launch_builder(&f);
12694 b.arg(bytes)
12695 .arg(aq)
12696 .arg(ad)
12697 .arg(&mut y)
12698 .arg(&inf)
12699 .arg(&outf)
12700 .arg(&mi)
12701 .arg(&rb);
12702 unsafe {
12703 b.launch(cfg)?;
12704 }
12705 if scale != 1.0 {
12706 self.scale_inplace(&mut y, scale, m * out_f)?;
12707 }
12708 Ok(y)
12709 }
12710
12711 pub fn matmul_decode_exact(
12719 &self,
12720 w: &crate::model::GpuTensor,
12721 x: &CudaSlice<f32>,
12722 m: usize,
12723 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12724 use crate::model::GpuTensor;
12725 if let GpuTensor::Float { data, .. } = w {
12733 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12734 }
12735 if let GpuTensor::FloatBf16 { data, .. } = w {
12738 let (in_f, out_f) = (w.in_features(), w.out_features());
12739 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12742 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12743 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12744 return Ok(y);
12745 }
12746 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12747 }
12748 if !self.uses_q8_1_fast(w) {
12749 return self.matmul(w, x, m);
12750 }
12751 let in_f = w.in_features();
12752 let out_f = w.out_features();
12753 let (bytes, qtype, row_bytes, scale, rp) = match w {
12754 GpuTensor::Quant {
12755 bytes,
12756 qtype,
12757 row_bytes,
12758 scale,
12759 rp,
12760 ..
12761 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12762 _ => return self.matmul(w, x, m),
12763 };
12764 let (bytes, rp) = match w {
12767 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12768 _ => (bytes, rp),
12769 };
12770 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12771 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12775 return Ok(y);
12776 }
12777 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12786 && std::env::var("MEMRA_NO_BATCHED").is_err()
12787 && (m <= 4 || Self::b8_enabled())
12788 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12791 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12792 {
12793 let mcols = Self::batched_mcols(m);
12794 return self.qmatvec_mmvq_batched(
12795 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12796 );
12797 }
12798 if self.mmvq_supports(qtype) {
12799 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12802 }
12803 self.matmul_pre(w, &aq, &ad, x, m)
12806 }
12807
12808 pub fn matmul_decode_exact_pre(
12818 &self,
12819 w: &crate::model::GpuTensor,
12820 aq: &CudaSlice<i8>,
12821 ad: &CudaSlice<f32>,
12822 m: usize,
12823 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12824 use crate::model::GpuTensor;
12825 debug_assert!(
12826 self.uses_q8_1_fast(w),
12827 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12828 );
12829 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12831 return Ok(y);
12832 }
12833 let in_f = w.in_features();
12834 let out_f = w.out_features();
12835 let (bytes, qtype, row_bytes, scale, rp) = match w {
12836 GpuTensor::Quant {
12837 bytes,
12838 qtype,
12839 row_bytes,
12840 scale,
12841 rp,
12842 ..
12843 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12844 _ => {
12845 return Err(
12846 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12847 );
12848 }
12849 };
12850 let (bytes, rp) = match w {
12852 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12853 _ => (bytes, rp),
12854 };
12855 if (2..=16).contains(&m)
12857 && self.batched_supports(qtype)
12858 && self.mmvq_supports(qtype)
12859 && std::env::var("MEMRA_NO_BATCHED").is_err()
12860 && (m <= 4 || Self::b8_enabled())
12861 && (m <= 8
12862 || qtype == QT_Q4_0
12863 || qtype == QT_Q6_K
12864 || qtype == QT_F8_E4M3
12865 || qtype == QT_NVFP4
12866 || qtype == QT_Q4_K
12867 || qtype == QT_Q5_K
12868 || qtype == QT_Q8_0)
12869 {
12870 let mcols = Self::batched_mcols(m);
12871 return self.qmatvec_mmvq_batched(
12872 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12873 );
12874 }
12875 if self.mmvq_supports(qtype) {
12876 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12877 }
12878 let x0 = self.zeros(0)?;
12881 self.matmul_pre(w, aq, ad, &x0, m)
12882 }
12883
12884 pub fn matmul_decode_exact_dual_pre(
12893 &self,
12894 w0: &crate::model::GpuTensor,
12895 w1: &crate::model::GpuTensor,
12896 aq: &CudaSlice<i8>,
12897 ad: &CudaSlice<f32>,
12898 m: usize,
12899 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12900 {
12901 use crate::model::GpuTensor;
12902 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12903 let on = *ON.get_or_init(|| {
12904 std::env::var("MEMRA_SPEC_DUAL_T")
12905 .map(|v| v != "0")
12906 .unwrap_or(true)
12907 });
12908 if !on
12909 || !(2..=7).contains(&m)
12910 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12911 || !self.uses_q8_1_fast(w0)
12912 || !self.uses_q8_1_fast(w1)
12913 {
12914 return Ok(None);
12915 }
12916 if !self.mmvq_supports(QT_NVFP4) {
12921 return Ok(None);
12922 }
12923 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12924 if w1.in_features() != in_f || w1.out_features() != out_f {
12925 return Ok(None);
12926 }
12927 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12928 (
12929 GpuTensor::Quant {
12930 bytes: b0,
12931 qtype: q0,
12932 row_bytes: rb0,
12933 scale: s0,
12934 rp: rp0,
12935 rp4: None,
12936 ..
12937 },
12938 GpuTensor::Quant {
12939 bytes: b1,
12940 qtype: q1,
12941 row_bytes: rb1,
12942 scale: s1,
12943 rp: rp1,
12944 rp4: None,
12945 ..
12946 },
12947 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12948 (b0, b1, *rb0, *s0, *s1, *rp0)
12949 }
12950 _ => return Ok(None),
12951 };
12952 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12955 {
12956 return Ok(None);
12957 }
12958 let (y0, y1) =
12959 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12960 Ok(Some(((y0, s0), (y1, s1))))
12961 }
12962
12963 pub fn matmul_decode_exact_group4_pre(
12975 &self,
12976 ws: [&crate::model::GpuTensor; 4],
12977 aq: &CudaSlice<i8>,
12978 ad: &CudaSlice<f32>,
12979 m: usize,
12980 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12981 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12982 let on = *ON.get_or_init(|| {
12983 std::env::var("MEMRA_TK_GDN_GROUP")
12984 .map(|v| v != "0")
12985 .unwrap_or(true)
12986 });
12987 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12988 }
12989
12990 pub fn matmul_decode_exact_group3_pre(
12995 &self,
12996 ws: [&crate::model::GpuTensor; 3],
12997 aq: &CudaSlice<i8>,
12998 ad: &CudaSlice<f32>,
12999 m: usize,
13000 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
13001 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13002 let on = *ON.get_or_init(|| {
13003 std::env::var("MEMRA_TK_FA_GROUP")
13004 .map(|v| v != "0")
13005 .unwrap_or(true)
13006 });
13007 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
13008 }
13009
13010 fn matmul_decode_exact_group_pre(
13014 &self,
13015 ws: &[&crate::model::GpuTensor],
13016 aq: &CudaSlice<i8>,
13017 ad: &CudaSlice<f32>,
13018 m: usize,
13019 on: bool,
13020 tag: &'static str,
13021 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
13022 use crate::model::GpuTensor;
13023 if !on
13024 || !(2..=16).contains(&m)
13025 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13026 || (m > 4 && !Self::b8_enabled())
13027 || !self.mmvq_supports(QT_NVFP4)
13028 || !self.batched_supports(QT_NVFP4)
13029 {
13030 return Ok(None);
13031 }
13032 let in_f = ws[0].in_features();
13033 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
13034 for w in ws {
13035 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
13036 return Ok(None);
13037 }
13038 match w {
13039 GpuTensor::Quant {
13040 bytes,
13041 qtype,
13042 scale,
13043 rp: true,
13044 rp4: None,
13045 ..
13046 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
13047 parts.push((bytes, w.out_features(), *scale));
13048 }
13049 _ => return Ok(None),
13050 }
13051 }
13052 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13054 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
13055 let mcols = if (5..=7).contains(&m) && b567 {
13056 m
13057 } else {
13058 Self::batched_mcols(m)
13059 };
13060 let kname: &'static str = match mcols {
13061 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
13062 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
13063 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
13064 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
13065 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
13066 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
13067 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
13068 _ => return Ok(None),
13069 };
13070 if std::env::var("MEMRA_DEBUG").is_ok() {
13073 use std::sync::Mutex;
13074 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
13075 let mut seen = SEEN.lock().unwrap();
13076 if !seen.contains(&tag) {
13077 seen.push(tag);
13078 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
13079 }
13080 }
13081 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
13084 let three = parts.len() == 3;
13085 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
13086 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
13087 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
13088 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
13091 let cfg = LaunchConfig {
13092 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
13093 block_dim: (32, ROWS_PER_BLOCK, 1),
13094 shared_mem_bytes: 0,
13095 };
13096 let (inf, mi) = (in_f as i32, m as i32);
13097 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
13098 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
13099 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
13100 let s3 = if three { 1.0f32 } else { parts[3].2 };
13101 let w3 = if three { parts[0].0 } else { parts[3].0 };
13102 let f = self.func(kname);
13103 let __s_b = self.gpu.stream();
13104 let mut b = __s_b.launch_builder(&f);
13105 b.arg(parts[0].0)
13106 .arg(parts[1].0)
13107 .arg(parts[2].0)
13108 .arg(w3)
13109 .arg(aq)
13110 .arg(ad)
13111 .arg(&mut y0)
13112 .arg(&mut y1)
13113 .arg(&mut y2)
13114 .arg(&mut y3)
13115 .arg(&inf)
13116 .arg(&n0)
13117 .arg(&n1)
13118 .arg(&n2)
13119 .arg(&n3)
13120 .arg(&mi)
13121 .arg(&s0)
13122 .arg(&s1)
13123 .arg(&s2)
13124 .arg(&s3);
13125 unsafe {
13126 b.launch(cfg)?;
13127 }
13128 Ok(Some(if three {
13129 vec![y0, y1, y2]
13130 } else {
13131 vec![y0, y1, y2, y3]
13132 }))
13133 }
13134
13135 pub fn matmul_decode_exact_dual(
13151 &self,
13152 w0: &crate::model::GpuTensor,
13153 w1: &crate::model::GpuTensor,
13154 x: &CudaSlice<f32>,
13155 m: usize,
13156 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13157 use crate::model::GpuTensor;
13158 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13159 let on = *ON.get_or_init(|| {
13160 std::env::var("MEMRA_SPEC_DUAL_T")
13161 .map(|v| v != "0")
13162 .unwrap_or(true)
13163 });
13164 if !on
13165 || !(2..=4).contains(&m)
13166 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13167 || !self.uses_q8_1_fast(w0)
13168 || !self.uses_q8_1_fast(w1)
13169 {
13170 return Ok(None);
13171 }
13172 if !self.mmvq_supports(QT_NVFP4) {
13177 return Ok(None);
13178 }
13179 let (in_f, out_f) = (w0.in_features(), w0.out_features());
13180 if w1.in_features() != in_f || w1.out_features() != out_f {
13181 return Ok(None);
13182 }
13183 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
13184 (
13185 GpuTensor::Quant {
13186 bytes: b0,
13187 qtype: q0,
13188 row_bytes: rb0,
13189 scale: s0,
13190 rp: rp0,
13191 rp4: None,
13192 ..
13193 },
13194 GpuTensor::Quant {
13195 bytes: b1,
13196 qtype: q1,
13197 row_bytes: rb1,
13198 scale: s1,
13199 rp: rp1,
13200 rp4: None,
13201 ..
13202 },
13203 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
13204 (b0, b1, *rb0, *s0, *s1, *rp0)
13205 }
13206 _ => return Ok(None),
13207 };
13208 if std::env::var("MEMRA_DEBUG").is_ok() {
13211 static ONCE: std::sync::Once = std::sync::Once::new();
13212 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
13213 }
13214 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13215 let (y0, y1) =
13216 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
13217 let mut y0 = y0;
13218 let mut y1 = y1;
13219 if s0 != 1.0 {
13220 self.scale_inplace(&mut y0, s0, m * out_f)?;
13221 }
13222 if s1 != 1.0 {
13223 self.scale_inplace(&mut y1, s1, m * out_f)?;
13224 }
13225 Ok(Some((y0, y1)))
13226 }
13227
13228 #[allow(clippy::too_many_arguments)]
13233 pub fn qmatvec_batched_dual_raw(
13234 &self,
13235 b0: &CudaSlice<u8>,
13236 b1: &CudaSlice<u8>,
13237 aq: &CudaSlice<i8>,
13238 ad: &CudaSlice<f32>,
13239 m: usize,
13240 in_f: usize,
13241 out_f: usize,
13242 row_bytes: usize,
13243 rp: bool,
13244 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13245 const ROWS_PER_BLOCK: u32 = 4;
13246 let mcols = Self::batched_mcols(m);
13247 let tiny_rp1 = rp
13250 && mcols == 4
13251 && out_f <= 128
13252 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
13253 let (name, rows_per_block) = if tiny_rp1 {
13254 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
13255 } else {
13256 match (mcols, rp, m) {
13257 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
13258 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
13259 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
13260 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
13261 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
13262 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
13263 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
13264 _ => {
13265 return Err(
13266 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
13267 );
13268 }
13269 }
13270 };
13271 let f = self.func(name);
13272 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
13273 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
13274 let cfg = LaunchConfig {
13275 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13276 block_dim: (32, ROWS_PER_BLOCK, 1),
13277 shared_mem_bytes: 0,
13278 };
13279 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13280 let __s_b = self.gpu.stream();
13281 let mut b = __s_b.launch_builder(&f);
13282 b.arg(b0)
13283 .arg(b1)
13284 .arg(aq)
13285 .arg(ad)
13286 .arg(&mut y0)
13287 .arg(&mut y1)
13288 .arg(&inf)
13289 .arg(&outf)
13290 .arg(&mi)
13291 .arg(&rb);
13292 unsafe {
13293 b.launch(cfg)?;
13294 }
13295 Ok((y0, y1))
13296 }
13297
13298 pub fn matmul_pre_dual_noscale(
13310 &self,
13311 w0: &crate::model::GpuTensor,
13312 w1: &crate::model::GpuTensor,
13313 aq: &CudaSlice<i8>,
13314 ad: &CudaSlice<f32>,
13315 m: usize,
13316 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
13317 {
13318 use crate::model::GpuTensor;
13319 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13320 return Ok(None);
13321 }
13322 if !self.mmvq_supports(QT_NVFP4) {
13332 return Ok(None);
13333 }
13334 let (in_f, out_f) = (w0.in_features(), w0.out_features());
13335 if w1.in_features() != in_f || w1.out_features() != out_f {
13336 return Ok(None);
13337 }
13338 let no_mirror =
13351 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
13352 if self.q8_ffn_fuse2_on()
13353 && no_mirror(w0)
13354 && no_mirror(w1)
13355 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
13356 {
13357 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
13358 return Ok(Some(((y0, 1.0), (y1, 1.0))));
13359 }
13360 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13370 let (y0, y1) =
13371 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
13372 return Ok(Some(((y0, p0.3), (y1, p1.3))));
13373 }
13374 let (b0, q0, rb0, s0, rp0) = match w0 {
13375 GpuTensor::Quant {
13376 bytes,
13377 qtype,
13378 row_bytes,
13379 scale,
13380 rp,
13381 ..
13382 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13383 _ => return Ok(None),
13384 };
13385 let (b1, q1, rb1, s1, rp1) = match w1 {
13386 GpuTensor::Quant {
13387 bytes,
13388 qtype,
13389 row_bytes,
13390 scale,
13391 rp,
13392 ..
13393 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13394 _ => return Ok(None),
13395 };
13396 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13397 return Ok(None);
13398 }
13399 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13401 let rows_per_block = ROWS_PER_BLOCK * RPW;
13402 let f = self.func(if rp0 {
13403 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13404 } else {
13405 "qmatvec_nvfp4_mmvq_dual_mr2"
13406 });
13407 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13408 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13409 let cfg = LaunchConfig {
13410 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13411 block_dim: (32, ROWS_PER_BLOCK, 1),
13412 shared_mem_bytes: 0,
13413 };
13414 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13415 let one = 1.0f32;
13418 let __s_b = self.gpu.stream();
13419 let mut b = __s_b.launch_builder(&f);
13420 b.arg(b0)
13421 .arg(b1)
13422 .arg(aq)
13423 .arg(ad)
13424 .arg(&mut y0)
13425 .arg(&mut y1)
13426 .arg(&inf)
13427 .arg(&outf)
13428 .arg(&mi)
13429 .arg(&rb)
13430 .arg(&one)
13431 .arg(&one);
13432 unsafe {
13433 b.launch(cfg)?;
13434 }
13435 Ok(Some(((y0, s0), (y1, s1))))
13436 }
13437
13438 #[allow(clippy::too_many_arguments)]
13446 pub fn matmul_nvfp4_fused3(
13447 &self,
13448 w0: &crate::model::GpuTensor,
13449 w1: &crate::model::GpuTensor,
13450 w2: &crate::model::GpuTensor,
13451 aq: &CudaSlice<i8>,
13452 ad: &CudaSlice<f32>,
13453 m: usize,
13454 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13455 {
13456 use crate::model::GpuTensor;
13457 if !self.mmvq_supports(QT_NVFP4)
13464 || !self.uses_q8_1_fast(w0)
13465 || !self.uses_q8_1_fast(w1)
13466 || !self.uses_q8_1_fast(w2)
13467 {
13468 return Ok(None);
13469 }
13470 if (9..=16).contains(&m) {
13473 return Ok(
13474 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13475 Some(mut ys) => {
13476 let y2 = ys.pop().unwrap();
13477 let y1 = ys.pop().unwrap();
13478 let y0 = ys.pop().unwrap();
13479 Some((y0, y1, y2))
13480 }
13481 None => None,
13482 },
13483 );
13484 }
13485 if !(1..=8).contains(&m) {
13486 return Ok(None);
13487 }
13488 if m > 1 {
13489 let in_f = w0.in_features();
13490 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13491 || !self.batched_supports(QT_NVFP4)
13492 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13493 || (m > 4 && !Self::b8_enabled())
13494 || in_f % 512 != 0
13495 || in_f / 64 > 272
13496 {
13497 return Ok(None);
13498 }
13499 }
13500 let unpack = |w: &crate::model::GpuTensor| match w {
13501 GpuTensor::Quant {
13502 bytes,
13503 qtype,
13504 scale,
13505 rp,
13506 ..
13507 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13508 _ => None,
13509 };
13510 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13511 return Ok(None);
13512 };
13513 let in_f = w0.in_features();
13514 if w1.in_features() != in_f || w2.in_features() != in_f {
13515 return Ok(None);
13516 }
13517 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13518 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13520 let rows_pb = ROWS_PER_BLOCK * RPW;
13521 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13522 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13523 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13524 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13525 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13526 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13529 if m > 1 {
13530 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13532 return Ok(None);
13533 }
13534 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13535 let cfg = LaunchConfig {
13536 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13537 block_dim: (32, ROWS_PER_BLOCK, 1),
13538 shared_mem_bytes: 0,
13539 };
13540 let __s_b = self.gpu.stream();
13541 let mut b = __s_b.launch_builder(&f);
13542 b.arg(b0)
13543 .arg(b1)
13544 .arg(b2)
13545 .arg(aq)
13546 .arg(ad)
13547 .arg(&mut y0)
13548 .arg(&mut y1)
13549 .arg(&mut y2)
13550 .arg(&inf)
13551 .arg(&oi0)
13552 .arg(&oi1)
13553 .arg(&oi2)
13554 .arg(&mi);
13555 unsafe {
13556 b.launch(cfg)?;
13557 }
13558 return Ok(Some((y0, y1, y2)));
13559 }
13560 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13561 let cfg = LaunchConfig {
13562 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13563 block_dim: (32, ROWS_PER_BLOCK, 1),
13564 shared_mem_bytes: 0,
13565 };
13566 let __s_b = self.gpu.stream();
13567 let mut b = __s_b.launch_builder(&f);
13568 b.arg(b0)
13569 .arg(b1)
13570 .arg(b2)
13571 .arg(aq)
13572 .arg(ad)
13573 .arg(&mut y0)
13574 .arg(&mut y1)
13575 .arg(&mut y2)
13576 .arg(&inf)
13577 .arg(&oi0)
13578 .arg(&oi1)
13579 .arg(&oi2)
13580 .arg(&mi)
13581 .arg(&p0.1)
13582 .arg(&p1.1)
13583 .arg(&p2.1);
13584 unsafe {
13585 b.launch(cfg)?;
13586 }
13587 Ok(Some((y0, y1, y2)))
13588 }
13589
13590 pub fn matmul_nvfp4_fused2(
13599 &self,
13600 w0: &crate::model::GpuTensor,
13601 w1: &crate::model::GpuTensor,
13602 aq: &CudaSlice<i8>,
13603 ad: &CudaSlice<f32>,
13604 m: usize,
13605 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13606 use crate::model::GpuTensor;
13607 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13608 let off =
13609 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13610 if off
13613 || m != 1
13614 || !self.mmvq_supports(QT_NVFP4)
13615 || !self.uses_q8_1_fast(w0)
13616 || !self.uses_q8_1_fast(w1)
13617 {
13618 return Ok(None);
13619 }
13620 let unpack = |w: &crate::model::GpuTensor| match w {
13621 GpuTensor::Quant {
13622 bytes,
13623 qtype,
13624 scale,
13625 rp,
13626 ..
13627 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13628 _ => None,
13629 };
13630 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13631 return Ok(None);
13632 };
13633 let in_f = w0.in_features();
13634 if w1.in_features() != in_f {
13635 return Ok(None);
13636 }
13637 let (o0, o1) = (w0.out_features(), w1.out_features());
13638 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13640 let rows_pb = ROWS_PER_BLOCK * RPW;
13641 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13642 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13643 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13644 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13645 let cfg = LaunchConfig {
13646 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13647 block_dim: (32, ROWS_PER_BLOCK, 1),
13648 shared_mem_bytes: 0,
13649 };
13650 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13651 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13654 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13657 {
13658 use cudarc::driver::{DevicePtr, DevicePtrMut};
13659 let s = &self.gpu.stream();
13660 let (pw0, _g0) = b0.device_ptr(s);
13661 let (pw1, _g1) = b1.device_ptr(s);
13662 let (paq, _g2) = aq.device_ptr(s);
13663 let (pad, _g3) = ad.device_ptr(s);
13664 let (py0, _g4) = y0.device_ptr_mut(s);
13665 let (py1, _g5) = y1.device_ptr_mut(s);
13666 let (s0, s1) = (p0.1, p1.1);
13667 let mut ps = [
13668 &pw0 as *const _ as *mut std::ffi::c_void,
13669 &pw1 as *const _ as *mut _,
13670 &paq as *const _ as *mut _,
13671 &pad as *const _ as *mut _,
13672 &py0 as *const _ as *mut _,
13673 &py1 as *const _ as *mut _,
13674 &inf as *const _ as *mut _,
13675 &oi0 as *const _ as *mut _,
13676 &oi1 as *const _ as *mut _,
13677 &mi as *const _ as *mut _,
13678 &s0 as *const _ as *mut _,
13679 &s1 as *const _ as *mut _,
13680 ];
13681 unsafe {
13682 self.launch_pdl(
13683 "qmatvec_nvfp4_mmvq_fused2_rp",
13684 cfg.grid_dim,
13685 cfg.block_dim,
13686 &mut ps,
13687 )?;
13688 }
13689 }
13690 return Ok(Some((y0, y1)));
13691 }
13692 let __s_b = self.gpu.stream();
13693 let mut b = __s_b.launch_builder(&f);
13694 b.arg(b0)
13695 .arg(b1)
13696 .arg(aq)
13697 .arg(ad)
13698 .arg(&mut y0)
13699 .arg(&mut y1)
13700 .arg(&inf)
13701 .arg(&oi0)
13702 .arg(&oi1)
13703 .arg(&mi)
13704 .arg(&p0.1)
13705 .arg(&p1.1);
13706 unsafe {
13707 b.launch(cfg)?;
13708 }
13709 Ok(Some((y0, y1)))
13710 }
13711
13712 pub fn matmul_nvfp4_fused2_into(
13717 &self,
13718 w0: &crate::model::GpuTensor,
13719 w1: &crate::model::GpuTensor,
13720 aq: &CudaSlice<i8>,
13721 ad: &CudaSlice<f32>,
13722 y0: &mut CudaSlice<f32>,
13723 y1: &mut CudaSlice<f32>,
13724 ) -> Result<bool, Box<dyn std::error::Error>> {
13725 use crate::model::GpuTensor;
13726 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13727 let off =
13728 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13729 if off
13730 || !self.mmvq_supports(QT_NVFP4)
13731 || !self.uses_q8_1_fast(w0)
13732 || !self.uses_q8_1_fast(w1)
13733 {
13734 return Ok(false);
13735 }
13736 let unpack = |w: &crate::model::GpuTensor| match w {
13737 GpuTensor::Quant {
13738 bytes,
13739 qtype,
13740 scale,
13741 rp,
13742 ..
13743 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13744 _ => None,
13745 };
13746 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13747 return Ok(false);
13748 };
13749 let in_f = w0.in_features();
13750 if w1.in_features() != in_f {
13751 return Ok(false);
13752 }
13753 let (o0, o1) = (w0.out_features(), w1.out_features());
13754 if y0.len() < o0 || y1.len() < o1 {
13755 return Ok(false);
13756 }
13757 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13759 let rows_pb = ROWS_PER_BLOCK * RPW;
13760 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13761 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13762 let cfg = LaunchConfig {
13763 grid_dim: (nb(o0) + nb(o1), 1, 1),
13764 block_dim: (32, ROWS_PER_BLOCK, 1),
13765 shared_mem_bytes: 0,
13766 };
13767 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13768 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13771 let __s_b = self.gpu.stream();
13772 let mut b = __s_b.launch_builder(&f);
13773 b.arg(b0)
13774 .arg(b1)
13775 .arg(aq)
13776 .arg(ad)
13777 .arg(&mut *y0)
13778 .arg(&mut *y1)
13779 .arg(&inf)
13780 .arg(&oi0)
13781 .arg(&oi1)
13782 .arg(&mi)
13783 .arg(&p0.1)
13784 .arg(&p1.1);
13785 unsafe {
13786 b.launch(cfg)?;
13787 }
13788 Ok(true)
13789 }
13790
13791 #[allow(clippy::type_complexity)]
13796 pub fn matmul_nvfp4_fused4(
13797 &self,
13798 w0: &crate::model::GpuTensor,
13799 w1: &crate::model::GpuTensor,
13800 w2: &crate::model::GpuTensor,
13801 w3: &crate::model::GpuTensor,
13802 aq: &CudaSlice<i8>,
13803 ad: &CudaSlice<f32>,
13804 m: usize,
13805 ) -> Result<
13806 Option<(
13807 CudaSlice<f32>,
13808 CudaSlice<f32>,
13809 CudaSlice<f32>,
13810 CudaSlice<f32>,
13811 )>,
13812 Box<dyn std::error::Error>,
13813 > {
13814 use crate::model::GpuTensor;
13815 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13822 || !self.mmvq_supports(QT_NVFP4)
13823 || !self.uses_q8_1_fast(w0)
13824 || !self.uses_q8_1_fast(w1)
13825 || !self.uses_q8_1_fast(w2)
13826 || !self.uses_q8_1_fast(w3)
13827 {
13828 return Ok(None);
13829 }
13830 if (9..=16).contains(&m) {
13835 return Ok(
13836 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13837 Some(mut ys) => {
13838 let y3 = ys.pop().unwrap();
13839 let y2 = ys.pop().unwrap();
13840 let y1 = ys.pop().unwrap();
13841 let y0 = ys.pop().unwrap();
13842 Some((y0, y1, y2, y3))
13843 }
13844 None => None,
13845 },
13846 );
13847 }
13848 if !(1..=8).contains(&m) {
13849 return Ok(None);
13850 }
13851 if m > 1 {
13852 let in_f = w0.in_features();
13855 if !self.batched_supports(QT_NVFP4)
13856 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13857 || (m > 4 && !Self::b8_enabled())
13858 || in_f % 512 != 0
13859 || in_f / 64 > 272
13860 {
13861 return Ok(None);
13862 }
13863 }
13864 let unpack = |w: &crate::model::GpuTensor| match w {
13865 GpuTensor::Quant {
13866 bytes,
13867 qtype,
13868 scale,
13869 rp,
13870 ..
13871 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13872 _ => None,
13873 };
13874 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13875 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13876 else {
13877 return Ok(None);
13878 };
13879 let in_f = w0.in_features();
13880 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13881 return Ok(None);
13882 }
13883 let (o0, o1, o2, o3) = (
13884 w0.out_features(),
13885 w1.out_features(),
13886 w2.out_features(),
13887 w3.out_features(),
13888 );
13889 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13891 let rows_pb = ROWS_PER_BLOCK * RPW;
13892 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13893 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13894 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13895 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13896 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13897 let (inf, oi0, oi1, oi2, oi3, mi) = (
13898 in_f as i32,
13899 o0 as i32,
13900 o1 as i32,
13901 o2 as i32,
13902 o3 as i32,
13903 m as i32,
13904 );
13905 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13908 if m > 1 {
13909 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13912 return Ok(None);
13913 }
13914 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13915 let cfg = LaunchConfig {
13916 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13917 block_dim: (32, ROWS_PER_BLOCK, 1),
13918 shared_mem_bytes: 0,
13919 };
13920 let __s_b = self.gpu.stream();
13921 let mut b = __s_b.launch_builder(&f);
13922 b.arg(b0)
13923 .arg(b1)
13924 .arg(b2)
13925 .arg(b3)
13926 .arg(aq)
13927 .arg(ad)
13928 .arg(&mut y0)
13929 .arg(&mut y1)
13930 .arg(&mut y2)
13931 .arg(&mut y3)
13932 .arg(&inf)
13933 .arg(&oi0)
13934 .arg(&oi1)
13935 .arg(&oi2)
13936 .arg(&oi3)
13937 .arg(&mi);
13938 unsafe {
13939 b.launch(cfg)?;
13940 }
13941 return Ok(Some((y0, y1, y2, y3)));
13942 }
13943 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13944 let cfg = LaunchConfig {
13945 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13946 block_dim: (32, ROWS_PER_BLOCK, 1),
13947 shared_mem_bytes: 0,
13948 };
13949 let __s_b = self.gpu.stream();
13950 let mut b = __s_b.launch_builder(&f);
13951 b.arg(b0)
13952 .arg(b1)
13953 .arg(b2)
13954 .arg(b3)
13955 .arg(aq)
13956 .arg(ad)
13957 .arg(&mut y0)
13958 .arg(&mut y1)
13959 .arg(&mut y2)
13960 .arg(&mut y3)
13961 .arg(&inf)
13962 .arg(&oi0)
13963 .arg(&oi1)
13964 .arg(&oi2)
13965 .arg(&oi3)
13966 .arg(&mi)
13967 .arg(&p0.1)
13968 .arg(&p1.1)
13969 .arg(&p2.1)
13970 .arg(&p3.1);
13971 unsafe {
13972 b.launch(cfg)?;
13973 }
13974 Ok(Some((y0, y1, y2, y3)))
13975 }
13976
13977 pub fn matmul_q8_fused2(
13985 &self,
13986 w0: &crate::model::GpuTensor,
13987 w1: &crate::model::GpuTensor,
13988 aq: &CudaSlice<i8>,
13989 ad: &CudaSlice<f32>,
13990 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13991 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13997 return Ok(Some(self.e4m3_fused2_core(
13998 p0.0,
13999 p1.0,
14000 aq,
14001 ad,
14002 w0.in_features(),
14003 p0.1,
14004 p1.1,
14005 p0.2,
14006 p0.3,
14007 p1.3,
14008 )?));
14009 }
14010 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14011 return Ok(None);
14012 };
14013 Ok(Some(self.q8_fused2_core(
14014 p0.0,
14015 p1.0,
14016 aq,
14017 ad,
14018 w0.in_features(),
14019 p0.1,
14020 p1.1,
14021 p0.2,
14022 )?))
14023 }
14024
14025 #[allow(clippy::too_many_arguments)]
14026 fn q8_fused2_core(
14027 &self,
14028 b0: &CudaSlice<u8>,
14029 b1: &CudaSlice<u8>,
14030 aq: &CudaSlice<i8>,
14031 ad: &CudaSlice<f32>,
14032 in_f: usize,
14033 out0: usize,
14034 out1: usize,
14035 row_bytes: usize,
14036 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14037 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14039 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14040 let f = self.func("qmatvec_q8_0_mmvq_fused2");
14041 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14042 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14043 let cfg = LaunchConfig {
14044 grid_dim: (nb0 + nb1, 1, 1),
14045 block_dim: (32, ROWS_PER_BLOCK, 1),
14046 shared_mem_bytes: 0,
14047 };
14048 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14049 let __s_b = self.gpu.stream();
14050 let mut b = __s_b.launch_builder(&f);
14051 b.arg(b0)
14052 .arg(b1)
14053 .arg(aq)
14054 .arg(ad)
14055 .arg(&mut y0)
14056 .arg(&mut y1)
14057 .arg(&inf)
14058 .arg(&o0)
14059 .arg(&o1)
14060 .arg(&rbl);
14061 unsafe {
14062 b.launch(cfg)?;
14063 }
14064 Ok((y0, y1))
14065 }
14066
14067 pub fn matmul_q8_fused2_x(
14073 &self,
14074 w0: &crate::model::GpuTensor,
14075 w1: &crate::model::GpuTensor,
14076 x: &CudaSlice<f32>,
14077 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14078 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
14079 return Ok(None);
14080 }
14081 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14082 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
14083 return Ok(Some(self.e4m3_fused2_core(
14084 p0.0,
14085 p1.0,
14086 &aq,
14087 &ad,
14088 w0.in_features(),
14089 p0.1,
14090 p1.1,
14091 p0.2,
14092 p0.3,
14093 p1.3,
14094 )?));
14095 }
14096 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14097 return Ok(None);
14098 };
14099 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
14100 Ok(Some(self.q8_fused2_core(
14101 p0.0,
14102 p1.0,
14103 &aq,
14104 &ad,
14105 w0.in_features(),
14106 p0.1,
14107 p1.1,
14108 p0.2,
14109 )?))
14110 }
14111
14112 #[allow(clippy::too_many_arguments)]
14115 pub fn qmatvec_q8_fused2_raw(
14116 &self,
14117 b0: &CudaSlice<u8>,
14118 b1: &CudaSlice<u8>,
14119 x: &CudaSlice<f32>,
14120 in_f: usize,
14121 out0: usize,
14122 out1: usize,
14123 row_bytes: usize,
14124 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14125 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14126 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
14127 }
14128
14129 pub fn matmul_q4_fused3(
14135 &self,
14136 w0: &crate::model::GpuTensor,
14137 w1: &crate::model::GpuTensor,
14138 w2: &crate::model::GpuTensor,
14139 aq: &CudaSlice<i8>,
14140 ad: &CudaSlice<f32>,
14141 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14142 {
14143 use crate::model::GpuTensor;
14144 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14145 match w {
14146 GpuTensor::Quant {
14147 qtype, row_bytes, ..
14148 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14149 _ => None,
14150 }
14151 };
14152 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14153 return Ok(None);
14154 };
14155 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14156 return Ok(None);
14157 }
14158 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14162 match w {
14163 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14164 Some(m) => (m, true),
14165 None => (bytes, *rp),
14166 },
14167 _ => unreachable!(),
14168 }
14169 }
14170 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14171 if rp0 != rp1 || rp1 != rp2 {
14172 return Ok(None);
14173 }
14174 let rp = rp0;
14175 let rpb: u32 = 4;
14176 let mr1 = rp && Self::q40_mr1_on();
14180 let nb = |o: usize| {
14181 if mr1 {
14182 (o as u32).div_ceil(rpb)
14183 } else {
14184 (o as u32).div_ceil(2).div_ceil(rpb)
14185 }
14186 };
14187 let grid = nb(o0) + nb(o1) + nb(o2);
14188 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14189 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14190 let mut y2 = self.alloc_uninit::<f32>(o2)?;
14191 let f = self.func(if mr1 {
14192 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14193 } else if rp {
14194 "qmatvec_q4_0_mmvq_fused3_rp"
14195 } else {
14196 "qmatvec_q4_0_mmvq_fused3"
14197 });
14198 let cfg = LaunchConfig {
14199 grid_dim: (grid, 1, 1),
14200 block_dim: (32, rpb, 1),
14201 shared_mem_bytes: 0,
14202 };
14203 let inf = w0.in_features() as i32;
14204 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14205 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14206 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14209 {
14210 use cudarc::driver::{DevicePtr, DevicePtrMut};
14211 let s = &self.gpu.stream();
14212 let (p0, _g0) = b0.device_ptr(s);
14213 let (p1, _g1) = b1.device_ptr(s);
14214 let (p2, _g2) = b2.device_ptr(s);
14215 let (paq, _g3) = aq.device_ptr(s);
14216 let (pad, _g4) = ad.device_ptr(s);
14217 let (py0, _g5) = y0.device_ptr_mut(s);
14218 let (py1, _g6) = y1.device_ptr_mut(s);
14219 let (py2, _g7) = y2.device_ptr_mut(s);
14220 let mut ps = [
14221 &p0 as *const _ as *mut std::ffi::c_void,
14222 &p1 as *const _ as *mut _,
14223 &p2 as *const _ as *mut _,
14224 &paq as *const _ as *mut _,
14225 &pad as *const _ as *mut _,
14226 &py0 as *const _ as *mut _,
14227 &py1 as *const _ as *mut _,
14228 &py2 as *const _ as *mut _,
14229 &inf as *const _ as *mut _,
14230 &oo0 as *const _ as *mut _,
14231 &oo1 as *const _ as *mut _,
14232 &oo2 as *const _ as *mut _,
14233 &r0 as *const _ as *mut _,
14234 &r1 as *const _ as *mut _,
14235 &r2 as *const _ as *mut _,
14236 ];
14237 unsafe {
14238 self.launch_pdl(
14239 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14240 (grid, 1, 1),
14241 (32, rpb, 1),
14242 &mut ps,
14243 )?;
14244 }
14245 }
14246 return Ok(Some((y0, y1, y2)));
14247 }
14248 let __s_b = self.gpu.stream();
14249 let mut b = __s_b.launch_builder(&f);
14250 b.arg(b0)
14251 .arg(b1)
14252 .arg(b2)
14253 .arg(aq)
14254 .arg(ad)
14255 .arg(&mut y0)
14256 .arg(&mut y1)
14257 .arg(&mut y2)
14258 .arg(&inf)
14259 .arg(&oo0)
14260 .arg(&oo1)
14261 .arg(&oo2)
14262 .arg(&r0)
14263 .arg(&r1)
14264 .arg(&r2);
14265 unsafe {
14266 b.launch(cfg)?;
14267 }
14268 Ok(Some((y0, y1, y2)))
14269 }
14270
14271 #[allow(clippy::too_many_arguments)]
14274 pub fn matmul_q4_fused3_into(
14275 &self,
14276 w0: &crate::model::GpuTensor,
14277 w1: &crate::model::GpuTensor,
14278 w2: &crate::model::GpuTensor,
14279 aq: &CudaSlice<i8>,
14280 ad: &CudaSlice<f32>,
14281 y0: &mut CudaSlice<f32>,
14282 y1: &mut CudaSlice<f32>,
14283 y2: &mut CudaSlice<f32>,
14284 ) -> Result<bool, Box<dyn std::error::Error>> {
14285 use crate::model::GpuTensor;
14286 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14287 match w {
14288 GpuTensor::Quant {
14289 qtype, row_bytes, ..
14290 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14291 _ => None,
14292 }
14293 };
14294 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14295 return Ok(false);
14296 };
14297 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14298 return Ok(false);
14299 }
14300 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14301 match w {
14302 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14303 Some(m) => (m, true),
14304 None => (bytes, *rp),
14305 },
14306 _ => unreachable!(),
14307 }
14308 }
14309 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14310 if rp0 != rp1 || rp1 != rp2 {
14311 return Ok(false);
14312 }
14313 let rp = rp0;
14314 let rpb: u32 = 4;
14315 let mr1 = rp && Self::q40_mr1_on();
14316 let nb = |o: usize| {
14317 if mr1 {
14318 (o as u32).div_ceil(rpb)
14319 } else {
14320 (o as u32).div_ceil(2).div_ceil(rpb)
14321 }
14322 };
14323 let grid = nb(o0) + nb(o1) + nb(o2);
14324 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
14325 let f = self.func(if mr1 {
14326 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14327 } else if rp {
14328 "qmatvec_q4_0_mmvq_fused3_rp"
14329 } else {
14330 "qmatvec_q4_0_mmvq_fused3"
14331 });
14332 let cfg = LaunchConfig {
14333 grid_dim: (grid, 1, 1),
14334 block_dim: (32, rpb, 1),
14335 shared_mem_bytes: 0,
14336 };
14337 let inf = w0.in_features() as i32;
14338 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14339 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14340 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14342 use cudarc::driver::{DevicePtr, DevicePtrMut};
14343 let s = &self.gpu.stream();
14344 let (p0, _g0) = b0.device_ptr(s);
14345 let (p1, _g1) = b1.device_ptr(s);
14346 let (p2, _g2) = b2.device_ptr(s);
14347 let (paq, _g3) = aq.device_ptr(s);
14348 let (pad, _g4) = ad.device_ptr(s);
14349 let (py0, _g5) = y0.device_ptr_mut(s);
14350 let (py1, _g6) = y1.device_ptr_mut(s);
14351 let (py2, _g7) = y2.device_ptr_mut(s);
14352 let mut ps = [
14353 &p0 as *const _ as *mut std::ffi::c_void,
14354 &p1 as *const _ as *mut _,
14355 &p2 as *const _ as *mut _,
14356 &paq as *const _ as *mut _,
14357 &pad as *const _ as *mut _,
14358 &py0 as *const _ as *mut _,
14359 &py1 as *const _ as *mut _,
14360 &py2 as *const _ as *mut _,
14361 &inf as *const _ as *mut _,
14362 &oo0 as *const _ as *mut _,
14363 &oo1 as *const _ as *mut _,
14364 &oo2 as *const _ as *mut _,
14365 &r0 as *const _ as *mut _,
14366 &r1 as *const _ as *mut _,
14367 &r2 as *const _ as *mut _,
14368 ];
14369 unsafe {
14370 self.launch_pdl(
14371 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14372 (grid, 1, 1),
14373 (32, rpb, 1),
14374 &mut ps,
14375 )?;
14376 }
14377 return Ok(true);
14378 }
14379 let __s_b = self.gpu.stream();
14380 let mut b = __s_b.launch_builder(&f);
14381 b.arg(b0)
14382 .arg(b1)
14383 .arg(b2)
14384 .arg(aq)
14385 .arg(ad)
14386 .arg(&mut *y0)
14387 .arg(&mut *y1)
14388 .arg(&mut *y2)
14389 .arg(&inf)
14390 .arg(&oo0)
14391 .arg(&oo1)
14392 .arg(&oo2)
14393 .arg(&r0)
14394 .arg(&r1)
14395 .arg(&r2);
14396 unsafe {
14397 b.launch(cfg)?;
14398 }
14399 Ok(true)
14400 }
14401
14402 pub fn matmul_q4_fused2(
14404 &self,
14405 w0: &crate::model::GpuTensor,
14406 w1: &crate::model::GpuTensor,
14407 aq: &CudaSlice<i8>,
14408 ad: &CudaSlice<f32>,
14409 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14410 use crate::model::GpuTensor;
14411 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14412 match w {
14413 GpuTensor::Quant {
14414 qtype, row_bytes, ..
14415 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14416 _ => None,
14417 }
14418 };
14419 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14420 return Ok(None);
14421 };
14422 if w0.in_features() != w1.in_features() {
14423 return Ok(None);
14424 }
14425 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14427 match w {
14428 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14429 Some(m) => (m, true),
14430 None => (bytes, *rp),
14431 },
14432 _ => unreachable!(),
14433 }
14434 }
14435 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14436 if rp0 != rp1 {
14437 return Ok(None);
14438 }
14439 let rp = rp0;
14440 let rpb: u32 = 4;
14441 let mr1 = rp && Self::q40_mr1_on();
14443 let nb = |o: usize| {
14444 if mr1 {
14445 (o as u32).div_ceil(rpb)
14446 } else {
14447 (o as u32).div_ceil(2).div_ceil(rpb)
14448 }
14449 };
14450 let grid = nb(o0) + nb(o1);
14451 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14452 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14453 let f = self.func(if mr1 {
14454 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14455 } else if rp {
14456 "qmatvec_q4_0_mmvq_fused2_rp"
14457 } else {
14458 "qmatvec_q4_0_mmvq_fused2"
14459 });
14460 let cfg = LaunchConfig {
14461 grid_dim: (grid, 1, 1),
14462 block_dim: (32, rpb, 1),
14463 shared_mem_bytes: 0,
14464 };
14465 let inf = w0.in_features() as i32;
14466 let (oo0, oo1) = (o0 as i32, o1 as i32);
14467 let (r0, r1) = (rb0 as i64, rb1 as i64);
14468 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14470 {
14471 use cudarc::driver::{DevicePtr, DevicePtrMut};
14472 let s = &self.gpu.stream();
14473 let (p0, _g0) = b0.device_ptr(s);
14474 let (p1, _g1) = b1.device_ptr(s);
14475 let (paq, _g2) = aq.device_ptr(s);
14476 let (pad, _g3) = ad.device_ptr(s);
14477 let (py0, _g4) = y0.device_ptr_mut(s);
14478 let (py1, _g5) = y1.device_ptr_mut(s);
14479 let mut ps = [
14480 &p0 as *const _ as *mut std::ffi::c_void,
14481 &p1 as *const _ as *mut _,
14482 &paq as *const _ as *mut _,
14483 &pad as *const _ as *mut _,
14484 &py0 as *const _ as *mut _,
14485 &py1 as *const _ as *mut _,
14486 &inf as *const _ as *mut _,
14487 &oo0 as *const _ as *mut _,
14488 &oo1 as *const _ as *mut _,
14489 &r0 as *const _ as *mut _,
14490 &r1 as *const _ as *mut _,
14491 ];
14492 unsafe {
14493 self.launch_pdl(
14494 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14495 (grid, 1, 1),
14496 (32, rpb, 1),
14497 &mut ps,
14498 )?;
14499 }
14500 }
14501 return Ok(Some((y0, y1)));
14502 }
14503 let __s_b = self.gpu.stream();
14504 let mut b = __s_b.launch_builder(&f);
14505 b.arg(b0)
14506 .arg(b1)
14507 .arg(aq)
14508 .arg(ad)
14509 .arg(&mut y0)
14510 .arg(&mut y1)
14511 .arg(&inf)
14512 .arg(&oo0)
14513 .arg(&oo1)
14514 .arg(&r0)
14515 .arg(&r1);
14516 unsafe {
14517 b.launch(cfg)?;
14518 }
14519 Ok(Some((y0, y1)))
14520 }
14521
14522 pub fn matmul_q4_fused2_into(
14524 &self,
14525 w0: &crate::model::GpuTensor,
14526 w1: &crate::model::GpuTensor,
14527 aq: &CudaSlice<i8>,
14528 ad: &CudaSlice<f32>,
14529 y0: &mut CudaSlice<f32>,
14530 y1: &mut CudaSlice<f32>,
14531 ) -> Result<bool, Box<dyn std::error::Error>> {
14532 use crate::model::GpuTensor;
14533 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14534 match w {
14535 GpuTensor::Quant {
14536 qtype, row_bytes, ..
14537 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14538 _ => None,
14539 }
14540 };
14541 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14542 return Ok(false);
14543 };
14544 if w0.in_features() != w1.in_features() {
14545 return Ok(false);
14546 }
14547 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14548 match w {
14549 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14550 Some(m) => (m, true),
14551 None => (bytes, *rp),
14552 },
14553 _ => unreachable!(),
14554 }
14555 }
14556 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14557 if rp0 != rp1 {
14558 return Ok(false);
14559 }
14560 let rp = rp0;
14561 let rpb: u32 = 4;
14562 let mr1 = rp && Self::q40_mr1_on();
14563 let nb = |o: usize| {
14564 if mr1 {
14565 (o as u32).div_ceil(rpb)
14566 } else {
14567 (o as u32).div_ceil(2).div_ceil(rpb)
14568 }
14569 };
14570 let grid = nb(o0) + nb(o1);
14571 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14572 let f = self.func(if mr1 {
14573 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14574 } else if rp {
14575 "qmatvec_q4_0_mmvq_fused2_rp"
14576 } else {
14577 "qmatvec_q4_0_mmvq_fused2"
14578 });
14579 let cfg = LaunchConfig {
14580 grid_dim: (grid, 1, 1),
14581 block_dim: (32, rpb, 1),
14582 shared_mem_bytes: 0,
14583 };
14584 let inf = w0.in_features() as i32;
14585 let (oo0, oo1) = (o0 as i32, o1 as i32);
14586 let (r0, r1) = (rb0 as i64, rb1 as i64);
14587 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14589 use cudarc::driver::{DevicePtr, DevicePtrMut};
14590 let s = &self.gpu.stream();
14591 let (p0, _g0) = b0.device_ptr(s);
14592 let (p1, _g1) = b1.device_ptr(s);
14593 let (paq, _g2) = aq.device_ptr(s);
14594 let (pad, _g3) = ad.device_ptr(s);
14595 let (py0, _g4) = y0.device_ptr_mut(s);
14596 let (py1, _g5) = y1.device_ptr_mut(s);
14597 let mut ps = [
14598 &p0 as *const _ as *mut std::ffi::c_void,
14599 &p1 as *const _ as *mut _,
14600 &paq as *const _ as *mut _,
14601 &pad as *const _ as *mut _,
14602 &py0 as *const _ as *mut _,
14603 &py1 as *const _ as *mut _,
14604 &inf as *const _ as *mut _,
14605 &oo0 as *const _ as *mut _,
14606 &oo1 as *const _ as *mut _,
14607 &r0 as *const _ as *mut _,
14608 &r1 as *const _ as *mut _,
14609 ];
14610 unsafe {
14611 self.launch_pdl(
14612 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14613 (grid, 1, 1),
14614 (32, rpb, 1),
14615 &mut ps,
14616 )?;
14617 }
14618 return Ok(true);
14619 }
14620 let __s_b = self.gpu.stream();
14621 let mut b = __s_b.launch_builder(&f);
14622 b.arg(b0)
14623 .arg(b1)
14624 .arg(aq)
14625 .arg(ad)
14626 .arg(&mut *y0)
14627 .arg(&mut *y1)
14628 .arg(&inf)
14629 .arg(&oo0)
14630 .arg(&oo1)
14631 .arg(&r0)
14632 .arg(&r1);
14633 unsafe {
14634 b.launch(cfg)?;
14635 }
14636 Ok(true)
14637 }
14638
14639 pub fn matmul_q4_fused2_batched(
14644 &self,
14645 w0: &crate::model::GpuTensor,
14646 w1: &crate::model::GpuTensor,
14647 aq: &CudaSlice<i8>,
14648 ad: &CudaSlice<f32>,
14649 m: usize,
14650 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14651 use crate::model::GpuTensor;
14652 if m < 2 || m > 8 {
14653 return Ok(None);
14654 }
14655 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14656 match w {
14657 GpuTensor::Quant {
14658 qtype, row_bytes, ..
14659 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14660 _ => None,
14661 }
14662 };
14663 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14664 return Ok(None);
14665 };
14666 if w0.in_features() != w1.in_features() {
14667 return Ok(None);
14668 }
14669 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14670 match w {
14671 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14672 Some(mr) => (mr, true),
14673 None => (bytes, *rp),
14674 },
14675 _ => unreachable!(),
14676 }
14677 }
14678 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14679 if !rp0 || !rp1 {
14680 return Ok(None);
14681 }
14682 let mcols = Self::batched_mcols(m);
14683 let rpb: u32 = 4;
14684 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14685 let grid = nb(o0) + nb(o1);
14686 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14687 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14688 let f = self.func(match mcols {
14689 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14690 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14691 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14692 });
14693 let cfg = LaunchConfig {
14694 grid_dim: (grid, 1, 1),
14695 block_dim: (32, rpb, 1),
14696 shared_mem_bytes: 0,
14697 };
14698 let inf = w0.in_features() as i32;
14699 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14700 let rb = rb0 as i64;
14701 let __s_b = self.gpu.stream();
14702 let mut b = __s_b.launch_builder(&f);
14703 b.arg(b0)
14704 .arg(b1)
14705 .arg(aq)
14706 .arg(ad)
14707 .arg(&mut y0)
14708 .arg(&mut y1)
14709 .arg(&inf)
14710 .arg(&oo0)
14711 .arg(&oo1)
14712 .arg(&mi)
14713 .arg(&rb);
14714 unsafe {
14715 b.launch(cfg)?;
14716 }
14717 Ok(Some((y0, y1)))
14718 }
14719
14720 #[allow(clippy::too_many_arguments)]
14723 pub fn matmul_q4_fused3_batched(
14724 &self,
14725 w0: &crate::model::GpuTensor,
14726 w1: &crate::model::GpuTensor,
14727 w2: &crate::model::GpuTensor,
14728 aq: &CudaSlice<i8>,
14729 ad: &CudaSlice<f32>,
14730 m: usize,
14731 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14732 {
14733 use crate::model::GpuTensor;
14734 if m < 2 || m > 8 {
14735 return Ok(None);
14736 }
14737 let q4 = |w: &GpuTensor| -> Option<usize> {
14738 match w {
14739 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14740 _ => None,
14741 }
14742 };
14743 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14744 return Ok(None);
14745 };
14746 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14747 return Ok(None);
14748 }
14749 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14750 match w {
14751 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14752 Some(mr) => (mr, true),
14753 None => (bytes, *rp),
14754 },
14755 _ => unreachable!(),
14756 }
14757 }
14758 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14759 if !rp0 || !rp1 || !rp2 {
14760 return Ok(None);
14761 }
14762 let mcols = Self::batched_mcols(m);
14763 let rpb: u32 = 4;
14764 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14765 let grid = nb(o0) + nb(o1) + nb(o2);
14766 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14767 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14768 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14769 let f = self.func(match mcols {
14770 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14771 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14772 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14773 });
14774 let cfg = LaunchConfig {
14775 grid_dim: (grid, 1, 1),
14776 block_dim: (32, rpb, 1),
14777 shared_mem_bytes: 0,
14778 };
14779 let inf = w0.in_features() as i32;
14780 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14781 let rb = 0i64;
14782 let __s_b = self.gpu.stream();
14783 let mut b = __s_b.launch_builder(&f);
14784 b.arg(b0)
14785 .arg(b1)
14786 .arg(b2)
14787 .arg(aq)
14788 .arg(ad)
14789 .arg(&mut y0)
14790 .arg(&mut y1)
14791 .arg(&mut y2)
14792 .arg(&inf)
14793 .arg(&oo0)
14794 .arg(&oo1)
14795 .arg(&oo2)
14796 .arg(&mi)
14797 .arg(&rb);
14798 unsafe {
14799 b.launch(cfg)?;
14800 }
14801 Ok(Some((y0, y1, y2)))
14802 }
14803
14804 pub fn matmul_q8_fused3(
14805 &self,
14806 w0: &crate::model::GpuTensor,
14807 w1: &crate::model::GpuTensor,
14808 w2: &crate::model::GpuTensor,
14809 aq: &CudaSlice<i8>,
14810 ad: &CudaSlice<f32>,
14811 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14812 {
14813 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14816 return Ok(Some(self.e4m3_fused3_core(
14817 p0.0,
14818 p1.0,
14819 p2.0,
14820 aq,
14821 ad,
14822 w0.in_features(),
14823 p0.1,
14824 p1.1,
14825 p2.1,
14826 p0.2,
14827 p0.3,
14828 p1.3,
14829 p2.3,
14830 )?));
14831 }
14832 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14833 return Ok(None);
14834 };
14835 Ok(Some(self.q8_fused3_core(
14836 p0.0,
14837 p1.0,
14838 p2.0,
14839 aq,
14840 ad,
14841 w0.in_features(),
14842 p0.1,
14843 p1.1,
14844 p2.1,
14845 p0.2,
14846 )?))
14847 }
14848
14849 #[allow(clippy::too_many_arguments)]
14850 fn q8_fused3_core(
14851 &self,
14852 b0: &CudaSlice<u8>,
14853 b1: &CudaSlice<u8>,
14854 b2: &CudaSlice<u8>,
14855 aq: &CudaSlice<i8>,
14856 ad: &CudaSlice<f32>,
14857 in_f: usize,
14858 out0: usize,
14859 out1: usize,
14860 out2: usize,
14861 row_bytes: usize,
14862 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14863 const ROWS_PER_BLOCK: u32 = 4;
14864 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14865 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14866 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14867 let f = self.func("qmatvec_q8_0_mmvq_fused3");
14868 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14869 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14870 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14871 let cfg = LaunchConfig {
14872 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14873 block_dim: (32, ROWS_PER_BLOCK, 1),
14874 shared_mem_bytes: 0,
14875 };
14876 let (inf, o0, o1, o2, rbl) = (
14877 in_f as i32,
14878 out0 as i32,
14879 out1 as i32,
14880 out2 as i32,
14881 row_bytes as i64,
14882 );
14883 let __s_b = self.gpu.stream();
14884 let mut b = __s_b.launch_builder(&f);
14885 b.arg(b0)
14886 .arg(b1)
14887 .arg(b2)
14888 .arg(aq)
14889 .arg(ad)
14890 .arg(&mut y0)
14891 .arg(&mut y1)
14892 .arg(&mut y2)
14893 .arg(&inf)
14894 .arg(&o0)
14895 .arg(&o1)
14896 .arg(&o2)
14897 .arg(&rbl);
14898 unsafe {
14899 b.launch(cfg)?;
14900 }
14901 Ok((y0, y1, y2))
14902 }
14903
14904 #[allow(clippy::too_many_arguments)]
14906 pub fn qmatvec_q8_fused3_raw(
14907 &self,
14908 b0: &CudaSlice<u8>,
14909 b1: &CudaSlice<u8>,
14910 b2: &CudaSlice<u8>,
14911 x: &CudaSlice<f32>,
14912 in_f: usize,
14913 out0: usize,
14914 out1: usize,
14915 out2: usize,
14916 row_bytes: usize,
14917 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14918 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14919 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14920 }
14921
14922 pub fn matmul_q8_fused2_t(
14933 &self,
14934 w0: &crate::model::GpuTensor,
14935 w1: &crate::model::GpuTensor,
14936 aq: &CudaSlice<i8>,
14937 ad: &CudaSlice<f32>,
14938 m: usize,
14939 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14940 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14944 return Ok(None);
14945 }
14946 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14949 if m > 4 && !Self::b8_enabled() {
14950 return Ok(None);
14951 }
14952 return Ok(Some(self.e4m3_fused2_t_core(
14953 p0.0,
14954 p1.0,
14955 aq,
14956 ad,
14957 m,
14958 w0.in_features(),
14959 p0.1,
14960 p1.1,
14961 p0.2,
14962 p0.3,
14963 p1.3,
14964 )?));
14965 }
14966 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14967 return Ok(None);
14968 };
14969 Ok(Some(self.q8_fused2_t_core(
14970 p0.0,
14971 p1.0,
14972 aq,
14973 ad,
14974 m,
14975 w0.in_features(),
14976 p0.1,
14977 p1.1,
14978 p0.2,
14979 )?))
14980 }
14981
14982 #[allow(clippy::too_many_arguments)]
14983 fn q8_fused2_t_core(
14984 &self,
14985 b0: &CudaSlice<u8>,
14986 b1: &CudaSlice<u8>,
14987 aq: &CudaSlice<i8>,
14988 ad: &CudaSlice<f32>,
14989 m: usize,
14990 in_f: usize,
14991 out0: usize,
14992 out1: usize,
14993 row_bytes: usize,
14994 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14995 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14997 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14998 let f = self.func(match Self::batched_mcols(m) {
14999 2 => "qmatvec_q8_0_mmvq_fused2_b2",
15000 4 => "qmatvec_q8_0_mmvq_fused2_b4",
15001 _ => "qmatvec_q8_0_mmvq_fused2_b8",
15003 });
15004 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15005 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15006 let cfg = LaunchConfig {
15007 grid_dim: (nb0 + nb1, 1, 1),
15008 block_dim: (32, ROWS_PER_BLOCK, 1),
15009 shared_mem_bytes: 0,
15010 };
15011 let (inf, o0, o1, mi, rbl) = (
15012 in_f as i32,
15013 out0 as i32,
15014 out1 as i32,
15015 m as i32,
15016 row_bytes as i64,
15017 );
15018 let __s_b = self.gpu.stream();
15019 let mut b = __s_b.launch_builder(&f);
15020 b.arg(b0)
15021 .arg(b1)
15022 .arg(aq)
15023 .arg(ad)
15024 .arg(&mut y0)
15025 .arg(&mut y1)
15026 .arg(&inf)
15027 .arg(&o0)
15028 .arg(&o1)
15029 .arg(&mi)
15030 .arg(&rbl);
15031 unsafe {
15032 b.launch(cfg)?;
15033 }
15034 Ok((y0, y1))
15035 }
15036
15037 #[allow(clippy::too_many_arguments)]
15040 pub fn qmatvec_q8_fused2_t_raw(
15041 &self,
15042 b0: &CudaSlice<u8>,
15043 b1: &CudaSlice<u8>,
15044 x: &CudaSlice<f32>,
15045 m: usize,
15046 in_f: usize,
15047 out0: usize,
15048 out1: usize,
15049 row_bytes: usize,
15050 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15051 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15052 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
15053 }
15054
15055 #[allow(clippy::too_many_arguments)]
15058 pub fn matmul_q8_fused3_t(
15059 &self,
15060 w0: &crate::model::GpuTensor,
15061 w1: &crate::model::GpuTensor,
15062 w2: &crate::model::GpuTensor,
15063 aq: &CudaSlice<i8>,
15064 ad: &CudaSlice<f32>,
15065 m: usize,
15066 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
15067 {
15068 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
15069 return Ok(None);
15070 }
15071 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
15072 return Ok(Some(self.e4m3_fused3_t_core(
15073 p0.0,
15074 p1.0,
15075 p2.0,
15076 aq,
15077 ad,
15078 m,
15079 w0.in_features(),
15080 p0.1,
15081 p1.1,
15082 p2.1,
15083 p0.2,
15084 p0.3,
15085 p1.3,
15086 p2.3,
15087 )?));
15088 }
15089 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
15090 return Ok(None);
15091 };
15092 Ok(Some(self.q8_fused3_t_core(
15093 p0.0,
15094 p1.0,
15095 p2.0,
15096 aq,
15097 ad,
15098 m,
15099 w0.in_features(),
15100 p0.1,
15101 p1.1,
15102 p2.1,
15103 p0.2,
15104 )?))
15105 }
15106
15107 #[allow(clippy::too_many_arguments)]
15108 fn q8_fused3_t_core(
15109 &self,
15110 b0: &CudaSlice<u8>,
15111 b1: &CudaSlice<u8>,
15112 b2: &CudaSlice<u8>,
15113 aq: &CudaSlice<i8>,
15114 ad: &CudaSlice<f32>,
15115 m: usize,
15116 in_f: usize,
15117 out0: usize,
15118 out1: usize,
15119 out2: usize,
15120 row_bytes: usize,
15121 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15122 const ROWS_PER_BLOCK: u32 = 4;
15123 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15124 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15125 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15126 let f = self.func(if Self::batched_mcols(m) == 2 {
15127 "qmatvec_q8_0_mmvq_fused3_b2"
15128 } else {
15129 "qmatvec_q8_0_mmvq_fused3_b4"
15130 });
15131 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15132 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15133 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15134 let cfg = LaunchConfig {
15135 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15136 block_dim: (32, ROWS_PER_BLOCK, 1),
15137 shared_mem_bytes: 0,
15138 };
15139 let (inf, o0, o1, o2, mi, rbl) = (
15140 in_f as i32,
15141 out0 as i32,
15142 out1 as i32,
15143 out2 as i32,
15144 m as i32,
15145 row_bytes as i64,
15146 );
15147 let __s_b = self.gpu.stream();
15148 let mut b = __s_b.launch_builder(&f);
15149 b.arg(b0)
15150 .arg(b1)
15151 .arg(b2)
15152 .arg(aq)
15153 .arg(ad)
15154 .arg(&mut y0)
15155 .arg(&mut y1)
15156 .arg(&mut y2)
15157 .arg(&inf)
15158 .arg(&o0)
15159 .arg(&o1)
15160 .arg(&o2)
15161 .arg(&mi)
15162 .arg(&rbl);
15163 unsafe {
15164 b.launch(cfg)?;
15165 }
15166 Ok((y0, y1, y2))
15167 }
15168
15169 #[allow(clippy::too_many_arguments)]
15171 pub fn qmatvec_q8_fused3_t_raw(
15172 &self,
15173 b0: &CudaSlice<u8>,
15174 b1: &CudaSlice<u8>,
15175 b2: &CudaSlice<u8>,
15176 x: &CudaSlice<f32>,
15177 m: usize,
15178 in_f: usize,
15179 out0: usize,
15180 out1: usize,
15181 out2: usize,
15182 row_bytes: usize,
15183 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15184 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15185 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
15186 }
15187
15188 pub fn q8_ffn_fuse2_on(&self) -> bool {
15192 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15193 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
15194 }
15195
15196 #[allow(clippy::type_complexity)]
15202 fn q8_fused_params<'w, const N: usize>(
15203 &self,
15204 ws: &[&'w crate::model::GpuTensor; N],
15205 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
15206 use crate::model::GpuTensor;
15207 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15208 return None;
15209 }
15210 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
15211 return None;
15212 }
15213 let in_f = ws[0].in_features();
15214 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
15215 for (i, w) in ws.iter().enumerate() {
15216 match w {
15217 GpuTensor::Quant {
15218 bytes,
15219 qtype,
15220 row_bytes,
15221 scale,
15222 ..
15223 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
15224 out[i] = Some((bytes, w.out_features(), *row_bytes))
15225 }
15226 _ => return None,
15227 }
15228 }
15229 Some(out.map(|o| o.unwrap()))
15230 }
15231
15232 pub fn e4m3_dual_on(&self) -> bool {
15235 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15236 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
15237 }
15238
15239 #[allow(clippy::type_complexity)]
15251 fn e4m3_fused_params<'w, const N: usize>(
15252 &self,
15253 ws: &[&'w crate::model::GpuTensor; N],
15254 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
15255 use crate::model::GpuTensor;
15256 if !self.e4m3_dual_on() {
15257 return None;
15258 }
15259 let in_f = ws[0].in_features();
15260 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
15261 for (i, w) in ws.iter().enumerate() {
15262 match w {
15263 GpuTensor::Quant {
15264 bytes,
15265 qtype,
15266 row_bytes,
15267 scale,
15268 rp,
15269 rp4,
15270 ..
15271 } if *qtype == QT_F8_E4M3
15272 && w.in_features() == in_f
15273 && *row_bytes == in_f
15274 && !*rp
15275 && rp4.is_none() =>
15276 {
15277 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
15278 }
15279 _ => return None,
15280 }
15281 }
15282 Some(out.map(|o| o.unwrap()))
15283 }
15284
15285 #[allow(clippy::too_many_arguments)]
15289 fn e4m3_fused2_core(
15290 &self,
15291 b0: &CudaSlice<u8>,
15292 b1: &CudaSlice<u8>,
15293 aq: &CudaSlice<i8>,
15294 ad: &CudaSlice<f32>,
15295 in_f: usize,
15296 out0: usize,
15297 out1: usize,
15298 row_bytes: usize,
15299 ws0: f32,
15300 ws1: f32,
15301 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15302 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15304 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15305 let f = self.func("qmatvec_e4m3_mmvq_fused2");
15306 let mut y0 = self.alloc_uninit::<f32>(out0)?;
15307 let mut y1 = self.alloc_uninit::<f32>(out1)?;
15308 let cfg = LaunchConfig {
15309 grid_dim: (nb0 + nb1, 1, 1),
15310 block_dim: (32, ROWS_PER_BLOCK, 1),
15311 shared_mem_bytes: 0,
15312 };
15313 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
15314 let __s_b = self.gpu.stream();
15315 let mut b = __s_b.launch_builder(&f);
15316 b.arg(b0)
15317 .arg(b1)
15318 .arg(aq)
15319 .arg(ad)
15320 .arg(&mut y0)
15321 .arg(&mut y1)
15322 .arg(&inf)
15323 .arg(&o0)
15324 .arg(&o1)
15325 .arg(&rbl)
15326 .arg(&ws0)
15327 .arg(&ws1);
15328 unsafe {
15329 b.launch(cfg)?;
15330 }
15331 Ok((y0, y1))
15332 }
15333
15334 #[allow(clippy::too_many_arguments)]
15336 fn e4m3_fused3_core(
15337 &self,
15338 b0: &CudaSlice<u8>,
15339 b1: &CudaSlice<u8>,
15340 b2: &CudaSlice<u8>,
15341 aq: &CudaSlice<i8>,
15342 ad: &CudaSlice<f32>,
15343 in_f: usize,
15344 out0: usize,
15345 out1: usize,
15346 out2: usize,
15347 row_bytes: usize,
15348 ws0: f32,
15349 ws1: f32,
15350 ws2: f32,
15351 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15352 const ROWS_PER_BLOCK: u32 = 4;
15353 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15354 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15355 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15356 let f = self.func("qmatvec_e4m3_mmvq_fused3");
15357 let mut y0 = self.alloc_uninit::<f32>(out0)?;
15358 let mut y1 = self.alloc_uninit::<f32>(out1)?;
15359 let mut y2 = self.alloc_uninit::<f32>(out2)?;
15360 let cfg = LaunchConfig {
15361 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15362 block_dim: (32, ROWS_PER_BLOCK, 1),
15363 shared_mem_bytes: 0,
15364 };
15365 let (inf, o0, o1, o2, rbl) = (
15366 in_f as i32,
15367 out0 as i32,
15368 out1 as i32,
15369 out2 as i32,
15370 row_bytes as i64,
15371 );
15372 let __s_b = self.gpu.stream();
15373 let mut b = __s_b.launch_builder(&f);
15374 b.arg(b0)
15375 .arg(b1)
15376 .arg(b2)
15377 .arg(aq)
15378 .arg(ad)
15379 .arg(&mut y0)
15380 .arg(&mut y1)
15381 .arg(&mut y2)
15382 .arg(&inf)
15383 .arg(&o0)
15384 .arg(&o1)
15385 .arg(&o2)
15386 .arg(&rbl)
15387 .arg(&ws0)
15388 .arg(&ws1)
15389 .arg(&ws2);
15390 unsafe {
15391 b.launch(cfg)?;
15392 }
15393 Ok((y0, y1, y2))
15394 }
15395
15396 #[allow(clippy::too_many_arguments)]
15400 fn e4m3_fused2_t_core(
15401 &self,
15402 b0: &CudaSlice<u8>,
15403 b1: &CudaSlice<u8>,
15404 aq: &CudaSlice<i8>,
15405 ad: &CudaSlice<f32>,
15406 m: usize,
15407 in_f: usize,
15408 out0: usize,
15409 out1: usize,
15410 row_bytes: usize,
15411 ws0: f32,
15412 ws1: f32,
15413 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15414 const ROWS_PER_BLOCK: u32 = 4;
15415 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15416 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15417 let f = self.func(match Self::batched_mcols(m) {
15418 2 => "qmatvec_e4m3_mmvq_fused2_b2",
15419 4 => "qmatvec_e4m3_mmvq_fused2_b4",
15420 _ => "qmatvec_e4m3_mmvq_fused2_b8",
15421 });
15422 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15423 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15424 let cfg = LaunchConfig {
15425 grid_dim: (nb0 + nb1, 1, 1),
15426 block_dim: (32, ROWS_PER_BLOCK, 1),
15427 shared_mem_bytes: 0,
15428 };
15429 let (inf, o0, o1, mi, rbl) = (
15430 in_f as i32,
15431 out0 as i32,
15432 out1 as i32,
15433 m as i32,
15434 row_bytes as i64,
15435 );
15436 let __s_b = self.gpu.stream();
15437 let mut b = __s_b.launch_builder(&f);
15438 b.arg(b0)
15439 .arg(b1)
15440 .arg(aq)
15441 .arg(ad)
15442 .arg(&mut y0)
15443 .arg(&mut y1)
15444 .arg(&inf)
15445 .arg(&o0)
15446 .arg(&o1)
15447 .arg(&mi)
15448 .arg(&rbl);
15449 unsafe {
15450 b.launch(cfg)?;
15451 }
15452 if ws0 != 1.0 {
15453 self.scale_inplace(&mut y0, ws0, m * out0)?;
15454 }
15455 if ws1 != 1.0 {
15456 self.scale_inplace(&mut y1, ws1, m * out1)?;
15457 }
15458 Ok((y0, y1))
15459 }
15460
15461 #[allow(clippy::too_many_arguments)]
15463 fn e4m3_fused3_t_core(
15464 &self,
15465 b0: &CudaSlice<u8>,
15466 b1: &CudaSlice<u8>,
15467 b2: &CudaSlice<u8>,
15468 aq: &CudaSlice<i8>,
15469 ad: &CudaSlice<f32>,
15470 m: usize,
15471 in_f: usize,
15472 out0: usize,
15473 out1: usize,
15474 out2: usize,
15475 row_bytes: usize,
15476 ws0: f32,
15477 ws1: f32,
15478 ws2: f32,
15479 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15480 const ROWS_PER_BLOCK: u32 = 4;
15481 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15482 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15483 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15484 let f = self.func(if Self::batched_mcols(m) == 2 {
15485 "qmatvec_e4m3_mmvq_fused3_b2"
15486 } else {
15487 "qmatvec_e4m3_mmvq_fused3_b4"
15488 });
15489 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15490 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15491 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15492 let cfg = LaunchConfig {
15493 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15494 block_dim: (32, ROWS_PER_BLOCK, 1),
15495 shared_mem_bytes: 0,
15496 };
15497 let (inf, o0, o1, o2, mi, rbl) = (
15498 in_f as i32,
15499 out0 as i32,
15500 out1 as i32,
15501 out2 as i32,
15502 m as i32,
15503 row_bytes as i64,
15504 );
15505 let __s_b = self.gpu.stream();
15506 let mut b = __s_b.launch_builder(&f);
15507 b.arg(b0)
15508 .arg(b1)
15509 .arg(b2)
15510 .arg(aq)
15511 .arg(ad)
15512 .arg(&mut y0)
15513 .arg(&mut y1)
15514 .arg(&mut y2)
15515 .arg(&inf)
15516 .arg(&o0)
15517 .arg(&o1)
15518 .arg(&o2)
15519 .arg(&mi)
15520 .arg(&rbl);
15521 unsafe {
15522 b.launch(cfg)?;
15523 }
15524 if ws0 != 1.0 {
15525 self.scale_inplace(&mut y0, ws0, m * out0)?;
15526 }
15527 if ws1 != 1.0 {
15528 self.scale_inplace(&mut y1, ws1, m * out1)?;
15529 }
15530 if ws2 != 1.0 {
15531 self.scale_inplace(&mut y2, ws2, m * out2)?;
15532 }
15533 Ok((y0, y1, y2))
15534 }
15535
15536 pub fn qmatvec_e4m3_blk_mmvq(
15546 &self,
15547 bytes: &CudaSlice<u8>,
15548 aq: &CudaSlice<i8>,
15549 ad: &CudaSlice<f32>,
15550 scales: &CudaSlice<f32>,
15551 m: usize,
15552 in_f: usize,
15553 out_f: usize,
15554 row_bytes: usize,
15555 scale_cols: usize,
15556 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15557 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
15559 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15560 )?;
15561 Ok(y)
15562 }
15563
15564 #[allow(clippy::too_many_arguments)]
15566 pub fn qmatvec_e4m3_blk_mmvq_into(
15567 &self,
15568 bytes: &CudaSlice<u8>,
15569 aq: &CudaSlice<i8>,
15570 ad: &CudaSlice<f32>,
15571 scales: &CudaSlice<f32>,
15572 m: usize,
15573 in_f: usize,
15574 out_f: usize,
15575 row_bytes: usize,
15576 scale_cols: usize,
15577 y: &mut CudaSlice<f32>,
15578 ) -> Result<(), Box<dyn std::error::Error>> {
15579 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
15581 let cfg = LaunchConfig {
15582 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15583 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15586 let (inf, outf, mi, rb, sc) = (
15587 in_f as i32,
15588 out_f as i32,
15589 m as i32,
15590 row_bytes as i64,
15591 scale_cols as i32,
15592 );
15593 let __s_b = self.gpu.stream();
15594 let mut b = __s_b.launch_builder(&f);
15595 b.arg(bytes)
15596 .arg(aq)
15597 .arg(ad)
15598 .arg(scales)
15599 .arg(&mut *y)
15600 .arg(&inf)
15601 .arg(&outf)
15602 .arg(&mi)
15603 .arg(&rb)
15604 .arg(&sc);
15605 unsafe {
15606 b.launch(cfg)?;
15607 }
15608 Ok(())
15609 }
15610
15611 #[allow(clippy::too_many_arguments)]
15617 pub fn qmatvec_e4m3_blk_mmvq_batched(
15618 &self,
15619 bytes: &CudaSlice<u8>,
15620 aq: &CudaSlice<i8>,
15621 ad: &CudaSlice<f32>,
15622 scales: &CudaSlice<f32>,
15623 m: usize,
15624 in_f: usize,
15625 out_f: usize,
15626 row_bytes: usize,
15627 scale_cols: usize,
15628 mcols: usize,
15629 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15630 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15632 let name = match mcols {
15633 2 => "qmatvec_e4m3_blk_mmvq_b2",
15634 4 => "qmatvec_e4m3_blk_mmvq_b4",
15635 8 => "qmatvec_e4m3_blk_mmvq_b8",
15636 16 => "qmatvec_e4m3_blk_mmvq_b16",
15637 _ => {
15638 return Err(
15639 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15640 );
15641 }
15642 };
15643 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15644 let f = self.func(name);
15645 let cfg = LaunchConfig {
15646 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15647 block_dim: (32, ROWS_PER_BLOCK, 1),
15648 shared_mem_bytes: 0,
15649 };
15650 let (inf, outf, mi, rb, sc) = (
15651 in_f as i32,
15652 out_f as i32,
15653 m as i32,
15654 row_bytes as i64,
15655 scale_cols as i32,
15656 );
15657 let __s_b = self.gpu.stream();
15658 let mut b = __s_b.launch_builder(&f);
15659 b.arg(bytes)
15660 .arg(aq)
15661 .arg(ad)
15662 .arg(scales)
15663 .arg(&mut y)
15664 .arg(&inf)
15665 .arg(&outf)
15666 .arg(&mi)
15667 .arg(&rb)
15668 .arg(&sc);
15669 unsafe {
15670 b.launch(cfg)?;
15671 }
15672 Ok(y)
15673 }
15674
15675 #[allow(clippy::too_many_arguments)]
15678 pub fn qmatvec_e4m3_blk_batched_raw(
15679 &self,
15680 bytes: &CudaSlice<u8>,
15681 x: &CudaSlice<f32>,
15682 scales: &CudaSlice<f32>,
15683 m: usize,
15684 in_f: usize,
15685 out_f: usize,
15686 row_bytes: usize,
15687 scale_cols: usize,
15688 mcols: usize,
15689 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15690 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15691 self.qmatvec_e4m3_blk_mmvq_batched(
15692 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15693 )
15694 }
15695
15696 #[allow(clippy::too_many_arguments)]
15699 pub fn qmatvec_e4m3_blk_mmvq_raw(
15700 &self,
15701 bytes: &CudaSlice<u8>,
15702 x: &CudaSlice<f32>,
15703 scales: &CudaSlice<f32>,
15704 m: usize,
15705 in_f: usize,
15706 out_f: usize,
15707 row_bytes: usize,
15708 scale_cols: usize,
15709 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15710 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15711 self.qmatvec_e4m3_blk_mmvq(
15712 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15713 )
15714 }
15715
15716 #[allow(clippy::too_many_arguments)]
15719 pub fn qmatvec_e4m3_fused2_raw(
15720 &self,
15721 b0: &CudaSlice<u8>,
15722 b1: &CudaSlice<u8>,
15723 x: &CudaSlice<f32>,
15724 in_f: usize,
15725 out0: usize,
15726 out1: usize,
15727 row_bytes: usize,
15728 ws0: f32,
15729 ws1: f32,
15730 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15731 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15732 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15733 }
15734
15735 #[allow(clippy::too_many_arguments)]
15736 pub fn qmatvec_e4m3_fused3_raw(
15737 &self,
15738 b0: &CudaSlice<u8>,
15739 b1: &CudaSlice<u8>,
15740 b2: &CudaSlice<u8>,
15741 x: &CudaSlice<f32>,
15742 in_f: usize,
15743 out0: usize,
15744 out1: usize,
15745 out2: usize,
15746 row_bytes: usize,
15747 ws0: f32,
15748 ws1: f32,
15749 ws2: f32,
15750 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15751 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15752 self.e4m3_fused3_core(
15753 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15754 )
15755 }
15756
15757 #[allow(clippy::too_many_arguments)]
15758 pub fn qmatvec_e4m3_fused2_t_raw(
15759 &self,
15760 b0: &CudaSlice<u8>,
15761 b1: &CudaSlice<u8>,
15762 x: &CudaSlice<f32>,
15763 m: usize,
15764 in_f: usize,
15765 out0: usize,
15766 out1: usize,
15767 row_bytes: usize,
15768 ws0: f32,
15769 ws1: f32,
15770 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15771 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15772 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15773 }
15774
15775 #[allow(clippy::too_many_arguments)]
15776 pub fn qmatvec_e4m3_fused3_t_raw(
15777 &self,
15778 b0: &CudaSlice<u8>,
15779 b1: &CudaSlice<u8>,
15780 b2: &CudaSlice<u8>,
15781 x: &CudaSlice<f32>,
15782 m: usize,
15783 in_f: usize,
15784 out0: usize,
15785 out1: usize,
15786 out2: usize,
15787 row_bytes: usize,
15788 ws0: f32,
15789 ws1: f32,
15790 ws2: f32,
15791 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15792 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15793 self.e4m3_fused3_t_core(
15794 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15795 )
15796 }
15797
15798 fn try_e4m3_blk_pre(
15809 &self,
15810 w: &crate::model::GpuTensor,
15811 aq: &CudaSlice<i8>,
15812 ad: &CudaSlice<f32>,
15813 m: usize,
15814 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15815 use crate::model::GpuTensor;
15816 if let GpuTensor::Quant {
15817 bytes,
15818 qtype,
15819 row_bytes,
15820 blk: Some(g),
15821 ..
15822 } = w
15823 {
15824 if *qtype == QT_F8_E4M3_BLK {
15825 if (2..=16).contains(&m)
15831 && std::env::var("MEMRA_NO_BATCHED").is_err()
15832 && (m <= 4 || Self::b8_enabled())
15833 {
15834 let mcols = Self::batched_mcols(m);
15835 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15836 bytes,
15837 aq,
15838 ad,
15839 &g.scales,
15840 m,
15841 w.in_features(),
15842 w.out_features(),
15843 *row_bytes,
15844 g.cols,
15845 mcols,
15846 )?));
15847 }
15848 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15849 bytes,
15850 aq,
15851 ad,
15852 &g.scales,
15853 m,
15854 w.in_features(),
15855 w.out_features(),
15856 *row_bytes,
15857 g.cols,
15858 )?));
15859 }
15860 }
15861 Ok(None)
15862 }
15863
15864 fn try_e4m3_blk_prefill(
15911 &self,
15912 w: &crate::model::GpuTensor,
15913 x: &CudaSlice<f32>,
15914 m: usize,
15915 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15916 use crate::model::GpuTensor;
15917 let GpuTensor::Quant {
15918 bytes,
15919 qtype,
15920 blk: Some(g),
15921 ..
15922 } = w
15923 else {
15924 return Ok(None);
15925 };
15926 if *qtype != QT_F8_E4M3_BLK {
15927 return Ok(None);
15928 }
15929 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15934 return Ok(Some(y));
15935 }
15936 let (in_f, out_f) = (w.in_features(), w.out_features());
15937 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15938 let tmp = GpuTensor::Quant {
15939 bytes: slab,
15940 qtype: QT_Q8_0,
15941 row_bytes: in_f / 32 * 34,
15942 ne: vec![in_f as u64, out_f as u64],
15943 scale: 1.0,
15944 rp: false,
15945 #[cfg(memra_cutlass)]
15946 cutlass: None,
15947 fp8: None,
15948 blk: None,
15949 f16: None,
15950 rp4: None,
15951 };
15952 Ok(Some(self.matmul(&tmp, x, m)?))
15954 }
15955
15956 pub fn matmul_pre_noscale(
15957 &self,
15958 w: &crate::model::GpuTensor,
15959 aq: &CudaSlice<i8>,
15960 ad: &CudaSlice<f32>,
15961 m: usize,
15962 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15963 use crate::model::GpuTensor;
15964 if m == 1 {
15968 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15969 return Ok(Some((y, 1.0)));
15970 }
15971 }
15972 if m != 1 || !self.uses_q8_1_fast(w) {
15974 return Ok(None);
15975 }
15976 let in_f = w.in_features();
15977 let out_f = w.out_features();
15978 let (bytes, qtype, row_bytes, scale, rp) = match w {
15979 GpuTensor::Quant {
15980 bytes,
15981 qtype,
15982 row_bytes,
15983 scale,
15984 rp,
15985 ..
15986 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15987 _ => return Ok(None),
15988 };
15989 if self.mmvq_supports(qtype) {
15991 let (mbytes, mrp) = match w {
15993 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15994 _ => (bytes, rp),
15995 };
15996 let y = self.qmatvec_mmvq(
15997 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
15998 )?;
15999 return Ok(Some((y, scale)));
16000 }
16001 let name = match qtype {
16003 QT_Q8_0 => "qmatvec_q8_0_dp4a",
16004 QT_Q4_K => "qmatvec_q4_K_dp4a",
16005 QT_Q6_K => "qmatvec_q6_K_dp4a",
16006 QT_Q5_K => "qmatvec_q5_K_dp4a",
16007 QT_Q3_K => "qmatvec_q3_K_dp4a",
16008 QT_NVFP4 => {
16009 if rp {
16010 "qmatvec_nvfp4_dp4a_rp"
16011 } else {
16012 "qmatvec_nvfp4_dp4a"
16013 }
16014 }
16015 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
16016 _ => return Ok(None),
16017 };
16018 let f = self.func(name);
16019 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16020 let cfg = LaunchConfig {
16021 grid_dim: (out_f as u32, m as u32, 1),
16022 block_dim: (128, 1, 1),
16023 shared_mem_bytes: 0,
16024 };
16025 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16026 let __s_b = self.gpu.stream();
16027 let mut b = __s_b.launch_builder(&f);
16028 b.arg(bytes)
16029 .arg(aq)
16030 .arg(ad)
16031 .arg(&mut y)
16032 .arg(&inf)
16033 .arg(&outf)
16034 .arg(&mi)
16035 .arg(&rb);
16036 unsafe {
16037 b.launch(cfg)?;
16038 }
16039 Ok(Some((y, scale)))
16040 }
16041
16042 pub fn mmvq_supports(&self, qtype: i32) -> bool {
16045 if qtype == QT_F8_E4M3 {
16050 return true;
16051 }
16052 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
16053 return false;
16054 }
16055 matches!(
16056 qtype,
16057 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
16058 )
16059 }
16060
16061 pub fn qmatvec_mmvq(
16066 &self,
16067 bytes: &CudaSlice<u8>,
16068 aq: &CudaSlice<i8>,
16069 ad: &CudaSlice<f32>,
16070 m: usize,
16071 in_f: usize,
16072 out_f: usize,
16073 qtype: i32,
16074 row_bytes: usize,
16075 scale: f32,
16076 rp: bool,
16077 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16078 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
16080 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
16081 )?;
16082 Ok(y)
16083 }
16084
16085 #[allow(clippy::too_many_arguments)]
16087 pub fn qmatvec_mmvq_into(
16088 &self,
16089 bytes: &CudaSlice<u8>,
16090 aq: &CudaSlice<i8>,
16091 ad: &CudaSlice<f32>,
16092 m: usize,
16093 in_f: usize,
16094 out_f: usize,
16095 qtype: i32,
16096 row_bytes: usize,
16097 scale: f32,
16098 rp: bool,
16099 y: &mut CudaSlice<f32>,
16100 ) -> Result<(), Box<dyn std::error::Error>> {
16101 debug_assert!(y.len() >= m * out_f);
16102 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
16108 && rp
16109 && m == 1
16110 && out_f >= 64
16111 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
16112 && {
16113 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16114 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
16115 }
16116 {
16117 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
16118 let cfg = LaunchConfig {
16119 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
16120 block_dim: (32, 2, 1),
16121 shared_mem_bytes: 0,
16122 };
16123 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
16124 let __s_b = self.gpu.stream();
16125 let mut b = __s_b.launch_builder(&f);
16126 b.arg(bytes)
16127 .arg(aq)
16128 .arg(ad)
16129 .arg(&mut *y)
16130 .arg(&inf)
16131 .arg(&outf)
16132 .arg(&mi)
16133 .arg(&rb);
16134 unsafe {
16135 b.launch(cfg)?;
16136 }
16137 if scale != 1.0 {
16138 self.scale_inplace(y, scale, out_f)?;
16139 }
16140 return Ok(());
16141 }
16142 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
16151 2
16152 } else {
16153 1
16154 };
16155 if m == 1 && qtype == QT_Q4_0 {
16160 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16161 mr = *Q40MR.get_or_init(|| {
16164 std::env::var("MEMRA_Q40_MR")
16165 .ok()
16166 .and_then(|v| v.parse().ok())
16167 .unwrap_or(1)
16168 });
16169 }
16170 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
16181 let q5_force = q5_mode.as_deref() == Some("2");
16182 let q5_il = qtype == QT_Q5_K
16185 && m == 1
16186 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
16187 if q5_il && !q5_force && out_f > 65536 {
16188 mr = 1;
16189 }
16190 if qtype == QT_Q4_0 && rp && mr != 1 {
16193 mr = 2;
16194 }
16195 if qtype == QT_Q8_0 && rp {
16199 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16200 mr = *Q80MR.get_or_init(|| {
16201 std::env::var("MEMRA_Q80_MR")
16202 .ok()
16203 .and_then(|v| v.parse().ok())
16204 .unwrap_or(1)
16205 });
16206 }
16207 let name = match (qtype, mr, rp) {
16208 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
16209 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
16210 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
16211 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
16212 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
16213 (QT_Q5_K, 2, _) => {
16214 if q5_il {
16215 "qmatvec_q5_K_mmvq_mr2_il"
16216 } else {
16217 "qmatvec_q5_K_mmvq_mr2"
16218 }
16219 }
16220 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
16221 (QT_Q8_0, _, true)
16226 if in_f % 1024 == 0 && {
16227 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16228 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
16229 } =>
16230 {
16231 "qmatvec_q8_0_mmvq_rpca"
16232 }
16233 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
16234 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
16235 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
16239 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
16240 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
16241 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
16242 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
16243 (QT_Q5_K, _, _) => {
16244 if q5_il {
16245 "qmatvec_q5_K_mmvq_il"
16246 } else {
16247 "qmatvec_q5_K_mmvq"
16248 }
16249 }
16250 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
16251 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
16252 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
16253 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
16254 };
16255 let f = self.func(name);
16256 let rows_per_block = ROWS_PER_BLOCK * mr;
16258 let cfg = LaunchConfig {
16259 grid_dim: (
16260 (out_f as u32 + rows_per_block - 1) / rows_per_block,
16261 m as u32,
16262 1,
16263 ),
16264 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
16267 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16268 let __s_b = self.gpu.stream();
16269 let mut b = __s_b.launch_builder(&f);
16270 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
16275 if Self::pdl_on()
16278 && Self::pdl_mmvq_on()
16279 && Self::pdl_nvfp4q8_on()
16280 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
16281 {
16282 use cudarc::driver::{DevicePtr, DevicePtrMut};
16283 let s = &self.gpu.stream();
16284 let (pw, _g0) = bytes.device_ptr(s);
16285 let (paq, _g1) = aq.device_ptr(s);
16286 let (pad, _g2) = ad.device_ptr(s);
16287 let (py, _g3) = y.device_ptr_mut(s);
16288 let mut ps = [
16289 &pw as *const _ as *mut std::ffi::c_void,
16290 &paq as *const _ as *mut _,
16291 &pad as *const _ as *mut _,
16292 &py as *const _ as *mut _,
16293 &inf as *const _ as *mut _,
16294 &outf as *const _ as *mut _,
16295 &mi as *const _ as *mut _,
16296 &rb as *const _ as *mut _,
16297 &scale as *const _ as *mut _,
16298 ];
16299 unsafe {
16300 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16301 }
16302 return Ok(());
16303 }
16304 b.arg(bytes)
16305 .arg(aq)
16306 .arg(ad)
16307 .arg(&mut *y)
16308 .arg(&inf)
16309 .arg(&outf)
16310 .arg(&mi)
16311 .arg(&rb)
16312 .arg(&scale);
16313 unsafe {
16314 b.launch(cfg)?;
16315 }
16316 } else if Self::pdl_on()
16317 && Self::pdl_mmvq_on()
16318 && (matches!(
16319 name,
16320 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
16321 ) || (Self::pdl_nvfp4q8_on()
16322 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
16323 {
16324 {
16328 use cudarc::driver::{DevicePtr, DevicePtrMut};
16329 let s = &self.gpu.stream();
16330 let (pw, _g0) = bytes.device_ptr(s);
16331 let (paq, _g1) = aq.device_ptr(s);
16332 let (pad, _g2) = ad.device_ptr(s);
16333 let (py, _g3) = y.device_ptr_mut(s);
16334 let mut ps = [
16335 &pw as *const _ as *mut std::ffi::c_void,
16336 &paq as *const _ as *mut _,
16337 &pad as *const _ as *mut _,
16338 &py as *const _ as *mut _,
16339 &inf as *const _ as *mut _,
16340 &outf as *const _ as *mut _,
16341 &mi as *const _ as *mut _,
16342 &rb as *const _ as *mut _,
16343 ];
16344 unsafe {
16345 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16346 }
16347 }
16348 if scale != 1.0 {
16349 self.scale_inplace(y, scale, m * out_f)?;
16350 }
16351 } else {
16352 b.arg(bytes)
16353 .arg(aq)
16354 .arg(ad)
16355 .arg(&mut *y)
16356 .arg(&inf)
16357 .arg(&outf)
16358 .arg(&mi)
16359 .arg(&rb);
16360 unsafe {
16361 b.launch(cfg)?;
16362 }
16363 if scale != 1.0 {
16364 self.scale_inplace(y, scale, m * out_f)?;
16365 }
16366 }
16367 Ok(())
16368 }
16369
16370 pub fn qmatvec_mmvq_raw(
16374 &self,
16375 bytes: &CudaSlice<u8>,
16376 x: &CudaSlice<f32>,
16377 m: usize,
16378 in_f: usize,
16379 out_f: usize,
16380 qtype: i32,
16381 row_bytes: usize,
16382 rp: bool,
16383 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16384 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16385 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
16386 }
16387
16388 pub fn batched_supports(&self, qtype: i32) -> bool {
16392 matches!(
16393 qtype,
16394 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16395 )
16396 }
16397
16398 pub fn iq_fast_enabled() -> bool {
16406 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16407 *ON.get_or_init(|| {
16408 std::env::var("MEMRA_IQ_FAST")
16409 .map(|v| v != "0")
16410 .unwrap_or(true)
16411 })
16412 }
16413
16414 pub fn b8_enabled() -> bool {
16417 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16418 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16419 }
16420
16421 pub fn batched_mcols(m: usize) -> usize {
16423 if m == 2 {
16424 2
16425 } else if m <= 4 {
16426 4
16427 } else if m <= 8 {
16428 8
16429 } else {
16430 16
16431 }
16432 }
16433
16434 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16439 Some(match (qtype, mcols) {
16440 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16441 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16442 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16443 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16449 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16450 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16451 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16452 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16455 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16456 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16457 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16458 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16461 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16462 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16463 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16464 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16465 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16466 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16467 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16468 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16472 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16473 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16474 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16475 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16479 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16480 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16481 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16482 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16483 _ => return None,
16484 })
16485 }
16486
16487 pub fn sm_count(&self) -> i32 {
16522 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16523 *SMS.get_or_init(|| {
16524 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16525 self.gpu
16526 .ctx
16527 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16528 .unwrap_or(82)
16529 })
16530 }
16531
16532 pub fn batched_variant(
16533 &self,
16534 _m: usize,
16535 in_f: usize,
16536 out_f: usize,
16537 qtype: i32,
16538 row_bytes: usize,
16539 mcols: usize,
16540 rp: bool,
16541 ) -> &'static str {
16542 if qtype == QT_Q8_0 {
16547 return if rp { "rp" } else { "base" };
16548 }
16549 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16550 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16551 Ok("base") => "base",
16552 Ok("pf") => "pf",
16553 Ok("r2") => "r2",
16554 Ok("r2w8") => "r2w8",
16555 Ok("pfr2") => "pfr2",
16556 Ok("ca") => "ca",
16557 Ok("car2") => "car2",
16558 Ok("rp") => "rp",
16561 Ok("rpr2") => "rpr2",
16562 Ok("rpr2w8") => "rpr2w8",
16563 Ok("rpca") => "rpca",
16566 Ok("rpcar2") => "rpcar2",
16567 Ok("rpsc") => "rpsc",
16574 Ok("rpms") => "rpms",
16575 Ok("rpmsc") => "rpmsc",
16576 Ok("rpks") => "rpks",
16577 Ok("rpksc") => "rpksc",
16578 _ => "auto",
16579 });
16580 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16584 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16589 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16590 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16591 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16592 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16593 let sms = *SMS.get_or_init(|| {
16594 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16595 self.gpu
16596 .ctx
16597 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16598 .unwrap_or(82)
16599 });
16600 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16620 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16623 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16624 Ok("base") => "base",
16625 Ok("r2") => "r2",
16626 Ok("r2w8") => "r2w8",
16627 _ => "auto",
16628 });
16629 let variant: &'static str = if qtype == QT_Q4_0 {
16630 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16634 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16635 Ok("base") => "base",
16641 Ok("r2") => "r2",
16642 Ok("ms") => "ms",
16643 Ok("sm") => "sm",
16644 Ok("la") => "la",
16645 _ => "auto",
16646 });
16647 let v = if q40 != "auto" {
16648 q40
16649 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16650 "r2"
16651 } else {
16652 "base"
16653 };
16654 if rp {
16659 match v {
16660 "ms" => "r2ms_rp",
16661 "sm" => "r2sm_rp",
16662 "la" => "r2la_rp",
16663 "r2" => "r2_rp",
16664 _ => "rp",
16665 }
16666 } else if matches!(v, "ms" | "sm" | "la") {
16667 "r2"
16668 } else {
16669 v
16670 }
16671 } else if qtype != QT_NVFP4 && !kq_r2 {
16672 "base"
16673 } else if kq_r2 && rp {
16674 "rp"
16678 } else if kq_r2 {
16679 if kq_bv != "auto" {
16682 if kq_bv == "r2w8" && mcols != 4 {
16683 "r2"
16684 } else {
16685 kq_bv
16686 }
16687 } else if bv != "auto" {
16688 match bv {
16689 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16690 "r2w8" | "rpr2w8" => {
16691 if mcols != 4 {
16692 "r2"
16693 } else {
16694 "r2w8"
16695 }
16696 }
16697 _ => "base", }
16699 } else {
16700 let blocks = (out_f + 7) / 8;
16701 let waves = blocks as f64 / (7 * sms as usize) as f64;
16702 let filled = blocks >= 4 * sms as usize;
16703 let use_r2 = if qtype == QT_Q4_K {
16704 filled
16705 } else {
16706 waves >= 2.0
16707 };
16708 if use_r2 { "r2" } else { "base" }
16709 }
16710 } else if bv != "auto" {
16711 let v = if bv == "r2w8" && mcols == 2 {
16716 "r2"
16717 } else if bv == "ca" && (!ca_ok || mcols == 8) {
16718 "pf"
16719 } else if bv == "car2" && (!ca_ok || mcols == 8) {
16720 "r2"
16721 } else if bv == "pfr2" && mcols == 8 {
16722 "r2"
16723 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16724 "rpr2"
16725 }
16726 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16728 if mcols == 8 { "rpr2w8" } else { "rpr2" }
16729 } else if bv == "rpcar2" && mcols == 2 {
16730 "rpca"
16731 }
16732 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16735 "rpr2"
16736 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16737 "rpr2"
16738 } else {
16739 bv
16740 };
16741 if rp {
16742 match v {
16743 "base" | "pf" | "ca" | "rp" => "rp",
16744 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16745 "r2w8" | "rpr2w8" => {
16746 if mcols == 2 {
16747 "rpr2"
16748 } else {
16749 "rpr2w8"
16750 }
16751 }
16752 other => other, }
16754 } else {
16755 v
16756 }
16757 } else if mcols == 8 {
16758 if rp {
16769 if sc_ok { "rpsc" } else { "rpr2w8" }
16770 } else {
16771 "r2w8"
16772 }
16773 } else if mcols >= 4 {
16774 let blocks = (out_f + 7) / 8;
16778 let r7 = 7 * sms as usize;
16779 let r8 = 8 * sms as usize;
16780 let waves = blocks as f64 / r7 as f64;
16781 let filled = blocks >= 4 * sms as usize;
16782 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16786 if rp { "rpr2w8" } else { "r2w8" }
16790 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16791 if rp { "rpr2" } else { "r2" }
16794 } else {
16795 if rp { "rp" } else { "pf" }
16799 }
16800 } else if in_f >= 6144 {
16801 if rp { "rpr2" } else { "r2" }
16805 } else if rp {
16806 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16811 if sc_ok && waves >= 0.9 && waves <= 1.1 {
16812 "rpsc"
16813 } else {
16814 "rp"
16815 }
16816 } else {
16817 "base"
16818 };
16819 variant
16820 }
16821
16822 pub fn qmatvec_mmvq_batched(
16823 &self,
16824 bytes: &CudaSlice<u8>,
16825 aq: &CudaSlice<i8>,
16826 ad: &CudaSlice<f32>,
16827 m: usize,
16828 in_f: usize,
16829 out_f: usize,
16830 qtype: i32,
16831 row_bytes: usize,
16832 mcols: usize,
16833 scale: f32,
16834 rp: bool,
16835 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16836 const ROWS_PER_BLOCK: u32 = 4;
16837 let forced: Option<&'static str> = {
16842 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16843 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16844 .as_deref()
16845 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16846 };
16847 let variant = match forced {
16848 Some(v) if !rp || v.contains("rp") => v,
16849 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16850 };
16851 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16852 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16853 })?;
16854 let variant = if mcols == 16 {
16858 if rp { "rp" } else { "base" }
16859 } else {
16860 variant
16861 };
16862 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16869 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16870 if b567
16871 && qtype == QT_NVFP4
16872 && rp
16873 && mcols == 8
16874 && (5..=7).contains(&m)
16875 && matches!(variant, "rpsc" | "rpr2w8")
16876 {
16877 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16878 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16880 let cfg = LaunchConfig {
16881 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16882 block_dim: (32, ROWS_PER_BLOCK, 1),
16883 shared_mem_bytes: 0,
16884 };
16885 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16886 let __s_b = self.gpu.stream();
16887 let mut b = __s_b.launch_builder(&f);
16888 b.arg(bytes)
16889 .arg(aq)
16890 .arg(ad)
16891 .arg(&mut y)
16892 .arg(&inf)
16893 .arg(&outf)
16894 .arg(&mi)
16895 .arg(&rb);
16896 unsafe {
16897 b.launch(cfg)?;
16898 }
16899 if scale != 1.0 {
16900 self.scale_inplace(&mut y, scale, m * out_f)?;
16901 }
16902 return Ok(y);
16903 }
16904 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16905 "base" => (base_name.into(), ROWS_PER_BLOCK),
16906 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16907 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16908 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16909 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16913 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16914 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16915 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16916 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16917 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16918 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16919 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
16921 debug_assert!(
16922 !rp || name.contains("_rp"),
16923 "rp weight dispatched to a GGUF-layout kernel"
16924 );
16925 let f = self.func(&name);
16926 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16927 let smem = if name.contains("_r2sm_rp") {
16929 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16930 } else {
16931 0
16932 };
16933 let cfg = LaunchConfig {
16934 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16935 block_dim: (32, ROWS_PER_BLOCK, 1),
16936 shared_mem_bytes: smem,
16937 };
16938 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16939 let __s_b = self.gpu.stream();
16940 let mut b = __s_b.launch_builder(&f);
16941 b.arg(bytes)
16942 .arg(aq)
16943 .arg(ad)
16944 .arg(&mut y)
16945 .arg(&inf)
16946 .arg(&outf)
16947 .arg(&mi)
16948 .arg(&rb);
16949 unsafe {
16950 b.launch(cfg)?;
16951 }
16952 if scale != 1.0 {
16953 self.scale_inplace(&mut y, scale, m * out_f)?;
16954 }
16955 Ok(y)
16956 }
16957
16958 pub fn qmatvec_batched_raw(
16962 &self,
16963 bytes: &CudaSlice<u8>,
16964 x: &CudaSlice<f32>,
16965 m: usize,
16966 in_f: usize,
16967 out_f: usize,
16968 qtype: i32,
16969 row_bytes: usize,
16970 mcols: usize,
16971 rp: bool,
16972 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16973 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16974 self.qmatvec_mmvq_batched(
16975 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16976 )
16977 }
16978
16979 pub fn qmatvec_nvfp4_batched_raw(
16981 &self,
16982 bytes: &CudaSlice<u8>,
16983 x: &CudaSlice<f32>,
16984 m: usize,
16985 in_f: usize,
16986 out_f: usize,
16987 row_bytes: usize,
16988 mcols: usize,
16989 rp: bool,
16990 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16991 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16992 }
16993
16994 fn try_fp4_gemm(
16998 &self,
16999 w: &crate::model::GpuTensor,
17000 x: &CudaSlice<f32>,
17001 m: usize,
17002 in_f: usize,
17003 out_f: usize,
17004 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17005 use crate::model::GpuTensor;
17006 if cfg!(memra_portable_cuda) {
17007 return Ok(None);
17008 }
17009 if std::env::var("MEMRA_FP4").is_ok() {
17012 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
17013 }
17014 if std::env::var("MEMRA_FP4").is_err() {
17015 return Ok(None);
17016 }
17017 #[cfg(memra_cutlass)]
17026 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
17027 if let GpuTensor::Quant {
17028 bytes,
17029 qtype,
17030 scale,
17031 row_bytes,
17032 cutlass,
17033 ..
17034 } = w
17035 {
17036 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
17037 if let Some(cw) = cutlass {
17038 let y = self.cutlass_fp4_gemm(
17040 &cw.b_packed,
17041 &cw.sfb_swizzled,
17042 x,
17043 *scale,
17044 m,
17045 out_f,
17046 in_f,
17047 )?;
17048 return Ok(Some(y));
17049 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
17050 let (b_packed, sfb_sw) =
17055 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
17056 let y =
17057 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
17058 return Ok(Some(y));
17059 }
17060 }
17061 }
17062 }
17063 if let GpuTensor::Quant {
17064 bytes,
17065 qtype,
17066 row_bytes,
17067 scale,
17068 rp,
17069 ..
17070 } = w
17071 {
17072 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
17075 let y =
17076 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
17077 return Ok(Some(y));
17078 }
17079 }
17080 Ok(None)
17081 }
17082
17083 pub fn rms_norm_f16out(
17087 &self,
17088 x: &CudaSlice<f32>,
17089 w: &CudaSlice<f32>,
17090 dst: &mut CudaSlice<f32>,
17091 dst16: &mut CudaSlice<u8>,
17092 ncols: usize,
17093 nrows: usize,
17094 eps: f32,
17095 ) -> Result<(), Box<dyn std::error::Error>> {
17096 let f = self.func("rms_norm_f16out_f32");
17097 let cfg = LaunchConfig {
17098 grid_dim: (nrows as u32, 1, 1),
17099 block_dim: (rms_block(), 1, 1),
17100 shared_mem_bytes: 0,
17101 };
17102 let (nc, e) = (ncols as i32, eps);
17103 let __s_b = self.gpu.stream();
17104 let mut b = __s_b.launch_builder(&f);
17105 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
17106 unsafe {
17107 b.launch(cfg)?;
17108 }
17109 Ok(())
17110 }
17111
17112 #[allow(clippy::too_many_arguments)]
17115 pub fn add_rms_norm_f16out(
17116 &self,
17117 a: &CudaSlice<f32>,
17118 b: &CudaSlice<f32>,
17119 w: &CudaSlice<f32>,
17120 res: &mut CudaSlice<f32>,
17121 dst: &mut CudaSlice<f32>,
17122 dst16: &mut CudaSlice<u8>,
17123 ncols: usize,
17124 nrows: usize,
17125 eps: f32,
17126 ) -> Result<(), Box<dyn std::error::Error>> {
17127 let f = self.func("add_rms_norm_f16out_f32");
17128 let cfg = LaunchConfig {
17129 grid_dim: (nrows as u32, 1, 1),
17130 block_dim: (rms_block(), 1, 1),
17131 shared_mem_bytes: 0,
17132 };
17133 let (nc, e) = (ncols as i32, eps);
17134 let __s_lb = self.gpu.stream();
17135 let mut lb = __s_lb.launch_builder(&f);
17136 lb.arg(a)
17137 .arg(b)
17138 .arg(w)
17139 .arg(res)
17140 .arg(dst)
17141 .arg(dst16)
17142 .arg(&nc)
17143 .arg(&e);
17144 unsafe {
17145 lb.launch(cfg)?;
17146 }
17147 Ok(())
17148 }
17149
17150 pub fn matmul_group_xh(
17153 &self,
17154 ws: &[&crate::model::GpuTensor],
17155 x: &CudaSlice<f32>,
17156 xh: &CudaSlice<u8>,
17157 m: usize,
17158 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17159 let mut out = Vec::with_capacity(ws.len());
17160 let in_f = ws[0].in_features();
17161 for w in ws {
17162 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
17163 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
17164 out.push(y);
17165 continue;
17166 }
17167 }
17168 out.push(self.matmul(w, x, m)?);
17169 }
17170 Ok(out)
17171 }
17172
17173 pub fn gdn_pad_mask(
17176 &self,
17177 beta: &mut CudaSlice<f32>,
17178 g_log: &mut CudaSlice<f32>,
17179 len_d: &CudaSlice<i32>,
17180 h: usize,
17181 t: usize,
17182 ) -> Result<(), Box<dyn std::error::Error>> {
17183 let f = self.func("gdn_pad_mask_f32");
17184 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
17185 let (hi, ti) = (h as i32, t as i32);
17186 let __s_b = self.gpu.stream();
17187 let mut b = __s_b.launch_builder(&f);
17188 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
17189 unsafe {
17190 b.launch(cfg)?;
17191 }
17192 Ok(())
17193 }
17194
17195 pub fn row_gather_dev(
17198 &self,
17199 src: &CudaSlice<f32>,
17200 dst: &mut CudaSlice<f32>,
17201 len_d: &CudaSlice<i32>,
17202 ncols: usize,
17203 ) -> Result<(), Box<dyn std::error::Error>> {
17204 let f = self.func("row_gather_dev_f32");
17205 let cfg = LaunchConfig::for_num_elems(ncols as u32);
17206 let nc = ncols as i32;
17207 let __s_b = self.gpu.stream();
17208 let mut b = __s_b.launch_builder(&f);
17209 b.arg(src).arg(dst).arg(len_d).arg(&nc);
17210 unsafe {
17211 b.launch(cfg)?;
17212 }
17213 Ok(())
17214 }
17215
17216 pub fn matmul_group(
17223 &self,
17224 ws: &[&crate::model::GpuTensor],
17225 x: &CudaSlice<f32>,
17226 m: usize,
17227 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17228 use crate::model::GpuTensor;
17229 let mut out = Vec::with_capacity(ws.len());
17230 let any_mirror = ws
17231 .iter()
17232 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
17233 if m >= 16 && any_mirror && !self.verify_exact_on() {
17234 let in_f = ws[0].in_features();
17235 let xh = self.f16_act(x, m * in_f, in_f)?;
17236 for w in ws {
17237 if w.in_features() == in_f {
17238 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
17239 out.push(y);
17240 continue;
17241 }
17242 }
17243 out.push(self.matmul(w, x, m)?);
17244 }
17245 return Ok(out);
17246 }
17247 for w in ws {
17248 out.push(self.matmul(w, x, m)?);
17249 }
17250 Ok(out)
17251 }
17252
17253 pub fn matmul_group_multi(
17260 &self,
17261 ws: &[&crate::model::GpuTensor],
17262 xs: &[&CudaSlice<f32>],
17263 ms: &[usize],
17264 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17265 assert_eq!(xs.len(), ms.len());
17266 let in_f = ws[0].in_features();
17267 let total: usize = ms.iter().sum();
17268 let mut xcat = self.uninit(total * in_f)?;
17269 let mut off = 0usize;
17270 for (x, &m) in xs.iter().zip(ms) {
17271 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
17272 off += m;
17273 }
17274 let ys = self.matmul_group(ws, &xcat, total)?;
17275 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
17276 for (w, y) in ws.iter().zip(ys) {
17277 let out_f = w.out_features();
17278 let mut off = 0usize;
17279 for (s, &m) in ms.iter().enumerate() {
17280 let mut ys_s = self.uninit(m * out_f)?;
17281 let src = y.slice(off * out_f..(off + m) * out_f);
17282 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
17283 out[s].push(ys_s);
17284 off += m;
17285 }
17286 }
17287 Ok(out)
17288 }
17289
17290 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
17300 use crate::model::GpuTensor;
17301 if !legacy_quant_gemm_allowed(
17302 cfg!(memra_portable_cuda),
17303 cfg!(memra_hopper_mma),
17304 std::env::var_os("MEMRA_NO_GEMM").is_some(),
17305 ) {
17306 return false;
17307 }
17308 match w {
17309 GpuTensor::Quant { qtype, .. } => {
17310 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
17311 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
17312 }
17313 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
17314 }
17315 }
17316
17317 pub fn qmatvec_gemm(
17324 &self,
17325 w: &crate::model::GpuTensor,
17326 aq: &CudaSlice<i8>,
17327 ad: &CudaSlice<f32>,
17328 m: usize,
17329 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17330 use crate::model::GpuTensor;
17331 let in_f = w.in_features();
17332 let out_f = w.out_features();
17333 let (bytes, qtype, row_bytes, scale, rp) = match w {
17334 GpuTensor::Quant {
17335 bytes,
17336 qtype,
17337 row_bytes,
17338 scale,
17339 rp,
17340 ..
17341 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17342 _ => unreachable!("gemm_supports guaranteed Quant"),
17343 };
17344 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
17350 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
17351 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
17352 if scale != 1.0 {
17353 self.scale_inplace(&mut y, scale, m * out_f)?;
17354 }
17355 return Ok(y);
17356 }
17357 }
17358 let name = match qtype {
17359 QT_Q8_0 => "qmatvec_gemm_q8_0",
17360 QT_Q4_K => "qmatvec_gemm_q4_K",
17361 QT_Q4_0 => {
17362 if rp {
17363 "qmatvec_gemm_q4_0_rp"
17364 } else {
17365 "qmatvec_gemm_q4_0"
17366 }
17367 }
17368 QT_Q5_K => "qmatvec_gemm_q5_K",
17369 QT_Q6_K => "qmatvec_gemm_q6_K",
17370 QT_NVFP4 => {
17371 if rp {
17372 "qmatvec_gemm_nvfp4_rp"
17373 } else {
17374 "qmatvec_gemm_nvfp4"
17375 }
17376 }
17377 _ => unreachable!(),
17378 };
17379 let f = self.func(name);
17380 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);
17385 let k1_tile = if is_k1 {
17387 k1_launch_override().unwrap_or((128, 128, 8))
17388 } else {
17389 (128, 128, 8)
17390 };
17391 let (bm, bn): (u32, u32) = if is_k1 {
17392 (k1_tile.0, k1_tile.1)
17393 } else {
17394 (64, 256)
17395 };
17396 let warps: u32 = if is_k1 {
17397 k1_tile.2
17398 } else {
17399 match qtype {
17400 QT_NVFP4 => 8,
17401 _ => 4,
17402 }
17403 };
17404 let cfg = LaunchConfig {
17405 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17406 block_dim: (32, warps, 1),
17407 shared_mem_bytes: 0,
17408 };
17409 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17410 let __s_b = self.gpu.stream();
17411 let mut b = __s_b.launch_builder(&f);
17412 b.arg(bytes)
17413 .arg(aq)
17414 .arg(ad)
17415 .arg(&mut y)
17416 .arg(&inf)
17417 .arg(&outf)
17418 .arg(&mi)
17419 .arg(&rb);
17420 unsafe {
17421 b.launch(cfg)?;
17422 }
17423 if scale != 1.0 {
17424 self.scale_inplace(&mut y, scale, m * out_f)?;
17425 }
17426 Ok(y)
17427 }
17428
17429 pub fn qmatvec_gemm_raw(
17434 &self,
17435 bytes: &CudaSlice<u8>,
17436 x: &CudaSlice<f32>,
17437 m: usize,
17438 in_f: usize,
17439 out_f: usize,
17440 qtype: i32,
17441 row_bytes: usize,
17442 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17443 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17444 let name = match qtype {
17445 QT_Q8_0 => "qmatvec_gemm_q8_0",
17446 QT_Q4_K => "qmatvec_gemm_q4_K",
17447 QT_Q4_0 => "qmatvec_gemm_q4_0",
17448 QT_Q5_K => "qmatvec_gemm_q5_K",
17449 QT_Q6_K => "qmatvec_gemm_q6_K",
17450 QT_NVFP4 => "qmatvec_gemm_nvfp4",
17451 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17452 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17453 };
17454 let f = self.func(name);
17455 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);
17459 let k1_tile = if is_k1 {
17461 k1_launch_override().unwrap_or((128, 128, 8))
17462 } else {
17463 (128, 128, 8)
17464 };
17465 let (bm, bn): (u32, u32) = if is_k1 {
17466 (k1_tile.0, k1_tile.1)
17467 } else {
17468 (64, 256)
17469 };
17470 let warps: u32 = if is_k1 {
17471 k1_tile.2
17472 } else {
17473 match qtype {
17474 QT_NVFP4 | QT_NVFP4_RP => 8,
17475 _ => 4,
17476 }
17477 };
17478 let cfg = LaunchConfig {
17479 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17480 block_dim: (32, warps, 1),
17481 shared_mem_bytes: 0,
17482 };
17483 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17484 let __s_b = self.gpu.stream();
17485 let mut b = __s_b.launch_builder(&f);
17486 b.arg(bytes)
17487 .arg(&aq)
17488 .arg(&ad)
17489 .arg(&mut y)
17490 .arg(&inf)
17491 .arg(&outf)
17492 .arg(&mi)
17493 .arg(&rb);
17494 unsafe {
17495 b.launch(cfg)?;
17496 }
17497 Ok(y)
17498 }
17499
17500 pub fn qmatvec_gemm_q8_0_wgmma_raw(
17507 &self,
17508 rp4: &CudaSlice<u8>,
17509 aq: &CudaSlice<i8>,
17510 ad: &CudaSlice<f32>,
17511 m: usize,
17512 in_f: usize,
17513 out_f: usize,
17514 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17515 assert!(
17516 out_f % 64 == 0 && in_f % 32 == 0,
17517 "wgmma GEMM needs out_f%64==0, in_f%32==0"
17518 );
17519 let f = self.func("qmatvec_gemm_q8_0_wgmma");
17520 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
17522 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17523 block_dim: (128, 1, 1),
17524 shared_mem_bytes: 0,
17525 };
17526 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17527 let __s_b = self.gpu.stream();
17528 let mut b = __s_b.launch_builder(&f);
17529 b.arg(rp4)
17530 .arg(aq)
17531 .arg(ad)
17532 .arg(&mut y)
17533 .arg(&inf)
17534 .arg(&outf)
17535 .arg(&mi);
17536 unsafe {
17537 b.launch(cfg)?;
17538 }
17539 Ok(y)
17540 }
17541
17542 pub fn scale_inplace(
17544 &self,
17545 y: &mut CudaSlice<f32>,
17546 s: f32,
17547 n: usize,
17548 ) -> Result<(), Box<dyn std::error::Error>> {
17549 let f = self.func("scale_f32");
17550 let cfg = LaunchConfig::for_num_elems(n as u32);
17551 let (sf, ni) = (s, n as i32);
17552 let __s_b = self.gpu.stream();
17553 let mut b = __s_b.launch_builder(&f);
17554 b.arg(y).arg(&sf).arg(&ni);
17555 unsafe {
17556 b.launch(cfg)?;
17557 }
17558 Ok(())
17559 }
17560
17561 pub fn bf16_to_f32(
17566 &self,
17567 data: &cudarc::driver::CudaView<'_, u8>,
17568 n: usize,
17569 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17570 let mut out = self.alloc_uninit::<f32>(n)?;
17571 let f = self.func("bf16_to_f32");
17572 let cfg = LaunchConfig::for_num_elems(n as u32);
17573 let ni = n as i32;
17574 let __s_b = self.gpu.stream();
17575 let mut b = __s_b.launch_builder(&f);
17576 b.arg(data).arg(&mut out).arg(&ni);
17577 unsafe {
17578 b.launch(cfg)?;
17579 }
17580 Ok(out)
17581 }
17582
17583 fn linear_bf16_chunked(
17590 &self,
17591 x: &CudaSlice<f32>,
17592 data: &CudaSlice<u8>,
17593 m: usize,
17594 in_f: usize,
17595 out_f: usize,
17596 exact: bool,
17597 canonical_chunk_rows: Option<usize>,
17598 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17599 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17603 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17604 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17605 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17606 let started = timing.then(std::time::Instant::now);
17607 let result =
17608 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17609 if let Some(started) = started {
17610 use std::sync::atomic::Ordering;
17611 self.stream().synchronize()?;
17612 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17613 + started.elapsed().as_nanos() as u64;
17614 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17615 + (in_f * out_f * 2) as u64;
17616 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17617 if calls % 1024 == 0 {
17618 eprintln!(
17619 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17620 weight_gb={:.2}",
17621 ns as f64 / 1.0e6,
17622 ns as f64 / calls as f64 / 1.0e3,
17623 wb as f64 / 1.0e9,
17624 );
17625 }
17626 }
17627 result
17628 }
17629
17630 pub(crate) fn bf16_mmv_on() -> bool {
17635 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17636 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17637 }
17638
17639 fn matvec_bf16(
17642 &self,
17643 data: &CudaSlice<u8>,
17644 x: &CudaSlice<f32>,
17645 in_f: usize,
17646 out_f: usize,
17647 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17648 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17649 return Err(format!(
17650 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17651 data.len(),
17652 x.len()
17653 )
17654 .into());
17655 }
17656 let mut y = self.alloc_uninit::<f32>(out_f)?;
17657 let f = self.func("matvec_bf16_f32acc");
17658 let cfg = LaunchConfig {
17659 grid_dim: (out_f as u32, 1, 1),
17660 block_dim: (mmv_block(), 1, 1),
17661 shared_mem_bytes: 0,
17662 };
17663 let ini = in_f as i32;
17664 let __s_bld = self.gpu.stream();
17665 let mut bld = __s_bld.launch_builder(&f);
17666 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17667 unsafe {
17668 bld.launch(cfg)?;
17669 }
17670 Ok(y)
17671 }
17672
17673 #[allow(clippy::too_many_arguments)]
17677 #[allow(clippy::too_many_arguments)]
17682 #[allow(clippy::too_many_arguments)]
17687 pub fn qk_norm_rope_append_inc_dcw_rows(
17688 &self,
17689 q_raw_t: &CudaSlice<f32>,
17690 k_raw_t: &CudaSlice<f32>,
17691 v_raw_t: &CudaSlice<f32>,
17692 qw: &CudaSlice<f32>,
17693 kw: &CudaSlice<f32>,
17694 q_out_t: &mut CudaSlice<f32>,
17695 k_out_t: &mut CudaSlice<f32>,
17696 tab: &CudaSlice<u64>,
17697 pos_t: &CudaSlice<i32>,
17698 same_session: bool,
17699 t: usize,
17700 kv_dim_k: usize,
17701 kv_dim_v: usize,
17702 k_tok_bytes: usize,
17703 v_tok_bytes: usize,
17704 head_dim: usize,
17705 n_dims: usize,
17706 nh_q: usize,
17707 nh_k: usize,
17708 eps: f32,
17709 freq_base: f32,
17710 freq_scale: f32,
17711 ff: Option<&CudaSlice<f32>>,
17712 ) -> Result<(), Box<dyn std::error::Error>> {
17713 if head_dim != 128
17714 || kv_dim_v != kv_dim_k
17715 || kv_dim_k != nh_k * head_dim
17716 || t == 0
17717 || t > 32
17718 || tab.len() < t * 6
17719 || pos_t.len() < t
17720 || q_raw_t.len() < t * nh_q * head_dim
17721 || k_raw_t.len() < t * nh_k * head_dim
17722 || v_raw_t.len() < t * kv_dim_v
17723 || q_out_t.len() < t * nh_q * head_dim
17724 || k_out_t.len() < t * nh_k * head_dim
17725 {
17726 return Err(format!(
17727 "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17728 nh_q={nh_q} nh_k={nh_k}"
17729 )
17730 .into());
17731 }
17732 let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17733 let same_t: i32 = if same_session { t as i32 } else { 0 };
17734 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17735 let cfg = LaunchConfig {
17736 grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17737 block_dim: (128, 1, 1),
17738 shared_mem_bytes: 0,
17739 };
17740 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17741 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17742 let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17743 let null: u64 = 0;
17744 let __s_b = self.gpu.stream();
17745 let mut b = __s_b.launch_builder(&f);
17746 b.arg(q_raw_t)
17747 .arg(k_raw_t)
17748 .arg(v_raw_t)
17749 .arg(qw)
17750 .arg(kw)
17751 .arg(q_out_t)
17752 .arg(k_out_t)
17753 .arg(tab)
17754 .arg(pos_t)
17755 .arg(&same_t)
17756 .arg(&kvk)
17757 .arg(&kvv)
17758 .arg(&ktb)
17759 .arg(&vtb)
17760 .arg(&hd)
17761 .arg(&nd)
17762 .arg(&nq)
17763 .arg(&nk)
17764 .arg(&eps)
17765 .arg(&theta_scale)
17766 .arg(&freq_scale);
17767 match ff {
17768 Some(freqs) => {
17769 b.arg(freqs);
17770 }
17771 None => {
17772 b.arg(&null);
17773 }
17774 }
17775 unsafe {
17776 b.launch(cfg)?;
17777 }
17778 Ok(())
17779 }
17780
17781 pub fn qk_norm_rope_append_inc_dcw(
17782 &self,
17783 q_raw: &CudaSlice<f32>,
17784 k_raw: &CudaSlice<f32>,
17785 v_raw: &CudaSlice<f32>,
17786 qw: &CudaSlice<f32>,
17787 kw: &CudaSlice<f32>,
17788 q_out: &mut CudaSlice<f32>,
17789 k_out: &mut CudaSlice<f32>,
17790 pos: &CudaSlice<i32>,
17791 k_plane: &mut CudaSlice<u8>,
17792 v_plane: &mut CudaSlice<u8>,
17793 len_dev: &CudaSlice<i32>,
17796 base_dev: Option<&CudaSlice<i32>>,
17797 done_ctr: &mut CudaSlice<u32>,
17798 kv_dim_k: usize,
17799 kv_dim_v: usize,
17800 k_tok_bytes: usize,
17801 v_tok_bytes: usize,
17802 head_dim: usize,
17803 n_dims: usize,
17804 nh_q: usize,
17805 nh_k: usize,
17806 eps: f32,
17807 freq_base: f32,
17808 freq_scale: f32,
17809 ff: Option<&CudaSlice<f32>>,
17810 ) -> Result<(), Box<dyn std::error::Error>> {
17811 if head_dim != 128
17812 || kv_dim_v != kv_dim_k
17813 || kv_dim_k != nh_k * head_dim
17814 || q_raw.len() < nh_q * head_dim
17815 || k_raw.len() < nh_k * head_dim
17816 || v_raw.len() < kv_dim_v
17817 || q_out.len() < nh_q * head_dim
17818 || k_out.len() < nh_k * head_dim
17819 || pos.is_empty()
17820 || done_ctr.is_empty()
17821 {
17822 return Err(format!(
17823 "qk_norm_rope_append_inc geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k} kv_k={kv_dim_k} kv_v={kv_dim_v}"
17824 )
17825 .into());
17826 }
17827 let f = self.func("qk_norm_rope_append_inc_dcw");
17828 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17829 let cfg = LaunchConfig {
17830 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17831 block_dim: (128, 1, 1),
17832 shared_mem_bytes: 0,
17833 };
17834 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17835 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17836 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17837 let null: u64 = 0;
17838 let __s_b = self.gpu.stream();
17839 let mut b = __s_b.launch_builder(&f);
17840 b.arg(q_raw)
17841 .arg(k_raw)
17842 .arg(v_raw)
17843 .arg(qw)
17844 .arg(kw)
17845 .arg(q_out)
17846 .arg(k_out)
17847 .arg(pos)
17848 .arg(&mut *k_plane)
17849 .arg(&mut *v_plane)
17850 .arg(len_dev);
17851 match base_dev {
17852 Some(base) => {
17853 b.arg(base);
17854 }
17855 None => {
17856 b.arg(&null);
17857 }
17858 }
17859 b.arg(&mut *done_ctr)
17860 .arg(&kvk)
17861 .arg(&kvv)
17862 .arg(&ktb)
17863 .arg(&vtb)
17864 .arg(&hd)
17865 .arg(&nd)
17866 .arg(&nq)
17867 .arg(&eps)
17868 .arg(&theta_scale)
17869 .arg(&freq_scale);
17870 match ff {
17871 Some(freqs) => {
17872 b.arg(freqs);
17873 }
17874 None => {
17875 b.arg(&null);
17876 }
17877 }
17878 unsafe {
17879 b.launch(cfg)?;
17880 }
17881 Ok(())
17882 }
17883
17884 pub fn qk_norm_rope_into(
17885 &self,
17886 q_raw: &CudaSlice<f32>,
17887 k_raw: &CudaSlice<f32>,
17888 qw: &CudaSlice<f32>,
17889 kw: &CudaSlice<f32>,
17890 q_out: &mut CudaSlice<f32>,
17891 k_out: &mut CudaSlice<f32>,
17892 pos: &CudaSlice<i32>,
17893 head_dim: usize,
17894 n_dims: usize,
17895 nh_q: usize,
17896 nh_k: usize,
17897 eps: f32,
17898 freq_base: f32,
17899 freq_scale: f32,
17900 ff: Option<&CudaSlice<f32>>,
17901 ) -> Result<(), Box<dyn std::error::Error>> {
17902 if head_dim > 512
17903 || q_raw.len() < nh_q * head_dim
17904 || k_raw.len() < nh_k * head_dim
17905 || q_out.len() < nh_q * head_dim
17906 || k_out.len() < nh_k * head_dim
17907 || qw.len() < head_dim
17908 || kw.len() < head_dim
17909 || pos.is_empty()
17910 {
17911 return Err(format!(
17912 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17913 )
17914 .into());
17915 }
17916 let f = self.func("qk_norm_rope_f32");
17917 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17918 let cfg = LaunchConfig {
17919 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17920 block_dim: (128, 1, 1),
17921 shared_mem_bytes: 0,
17922 };
17923 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17924 let __s_b = self.gpu.stream();
17925 let mut b = __s_b.launch_builder(&f);
17926 b.arg(q_raw)
17927 .arg(k_raw)
17928 .arg(qw)
17929 .arg(kw)
17930 .arg(q_out)
17931 .arg(k_out)
17932 .arg(pos)
17933 .arg(&hd)
17934 .arg(&nd)
17935 .arg(&nq)
17936 .arg(&eps)
17937 .arg(&theta_scale)
17938 .arg(&freq_scale);
17939 match ff {
17940 Some(ffv) => {
17941 b.arg(ffv);
17942 unsafe {
17943 b.launch(cfg)?;
17944 }
17945 }
17946 None => {
17947 let null: u64 = 0;
17948 b.arg(&null);
17949 unsafe {
17950 b.launch(cfg)?;
17951 }
17952 }
17953 }
17954 Ok(())
17955 }
17956
17957 #[allow(clippy::too_many_arguments)]
17960 pub fn matvec_f32_b4_into(
17961 &self,
17962 w: [&CudaSlice<f32>; 4],
17963 x: &CudaSlice<f32>,
17964 y: &mut CudaSlice<f32>,
17965 block_cols: usize,
17966 out_f: usize,
17967 ) -> Result<(), Box<dyn std::error::Error>> {
17968 if block_cols % 4 != 0
17969 || x.len() < 4 * block_cols
17970 || y.len() < out_f
17971 || w.iter().any(|w| w.len() != out_f * block_cols)
17972 {
17973 return Err(format!(
17974 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17975 x.len()
17976 )
17977 .into());
17978 }
17979 let f = self.func("matvec_f32_b4");
17980 let cfg = LaunchConfig {
17981 grid_dim: (out_f as u32, 1, 1),
17982 block_dim: (128, 1, 1),
17983 shared_mem_bytes: 0,
17984 };
17985 let (bc, of) = (block_cols as i32, out_f as i32);
17986 let __s_b = self.gpu.stream();
17987 let mut b = __s_b.launch_builder(&f);
17988 b.arg(w[0])
17989 .arg(w[1])
17990 .arg(w[2])
17991 .arg(w[3])
17992 .arg(x)
17993 .arg(y)
17994 .arg(&bc)
17995 .arg(&of);
17996 unsafe {
17997 b.launch(cfg)?;
17998 }
17999 Ok(())
18000 }
18001
18002 pub fn axpy_rows_seq_into(
18005 &self,
18006 x: &CudaSlice<f32>,
18007 w: &CudaSlice<f32>,
18008 y: &mut CudaSlice<f32>,
18009 width: usize,
18010 n_rows: usize,
18011 ) -> Result<(), Box<dyn std::error::Error>> {
18012 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
18013 return Err(format!(
18014 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
18015 x.len(),
18016 w.len(),
18017 y.len()
18018 )
18019 .into());
18020 }
18021 let f = self.func("axpy_rows_seq_f32");
18022 let cfg = LaunchConfig::for_num_elems(width as u32);
18023 let (wi, nr) = (width as i32, n_rows as i32);
18024 let __s_b = self.gpu.stream();
18025 let mut b = __s_b.launch_builder(&f);
18026 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
18027 unsafe {
18028 b.launch(cfg)?;
18029 }
18030 Ok(())
18031 }
18032
18033 #[allow(clippy::too_many_arguments)]
18037 pub fn axpy_rows_seq_md_off_into(
18038 &self,
18039 x: &CudaSlice<f32>,
18040 w_route: &CudaSlice<f32>,
18041 md: &CudaSlice<f32>,
18042 sel: &CudaSlice<i32>,
18043 y: &mut CudaSlice<f32>,
18044 width: usize,
18045 n_rows: usize,
18046 row0: usize,
18047 ) -> Result<(), Box<dyn std::error::Error>> {
18048 if x.len() < (row0 + n_rows) * width
18049 || w_route.len() < row0 + n_rows
18050 || sel.len() < row0 + n_rows
18051 || y.len() < width
18052 {
18053 return Err(format!(
18054 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
18055 rows={n_rows} row0={row0}",
18056 x.len(),
18057 w_route.len(),
18058 sel.len(),
18059 y.len()
18060 )
18061 .into());
18062 }
18063 let f = self.func("axpy_rows_seq_md_off_f32");
18064 let cfg = LaunchConfig::for_num_elems(width as u32);
18065 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
18066 let __s_b = self.gpu.stream();
18067 let mut b = __s_b.launch_builder(&f);
18068 b.arg(x)
18069 .arg(w_route)
18070 .arg(md)
18071 .arg(sel)
18072 .arg(y)
18073 .arg(&wi)
18074 .arg(&nr)
18075 .arg(&r0);
18076 unsafe {
18077 b.launch(cfg)?;
18078 }
18079 Ok(())
18080 }
18081
18082 #[allow(clippy::too_many_arguments)]
18085 pub fn axpy_rows_seq_md_into(
18086 &self,
18087 x: &CudaSlice<f32>,
18088 w_route: &CudaSlice<f32>,
18089 md: &CudaSlice<f32>,
18090 sel: &CudaSlice<i32>,
18091 y: &mut CudaSlice<f32>,
18092 width: usize,
18093 n_rows: usize,
18094 ) -> Result<(), Box<dyn std::error::Error>> {
18095 if x.len() < n_rows * width
18096 || w_route.len() < n_rows
18097 || sel.len() < n_rows
18098 || y.len() < width
18099 {
18100 return Err(format!(
18101 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
18102 x.len(),
18103 w_route.len(),
18104 sel.len(),
18105 y.len()
18106 )
18107 .into());
18108 }
18109 let f = self.func("axpy_rows_seq_md_f32");
18110 let cfg = LaunchConfig::for_num_elems(width as u32);
18111 let (wi, nr) = (width as i32, n_rows as i32);
18112 let __s_b = self.gpu.stream();
18113 let mut b = __s_b.launch_builder(&f);
18114 b.arg(x)
18115 .arg(w_route)
18116 .arg(md)
18117 .arg(sel)
18118 .arg(y)
18119 .arg(&wi)
18120 .arg(&nr);
18121 unsafe {
18122 b.launch(cfg)?;
18123 }
18124 Ok(())
18125 }
18126
18127 #[allow(clippy::too_many_arguments)]
18129 #[allow(clippy::too_many_arguments)]
18133 pub fn matvec_bf16_qkvg_tcol_into(
18134 &self,
18135 wq: &CudaSlice<u8>,
18136 wk: &CudaSlice<u8>,
18137 wv: &CudaSlice<u8>,
18138 wg: &CudaSlice<u8>,
18139 x_t: &CudaSlice<f32>,
18140 yq: &mut CudaSlice<f32>,
18141 yk: &mut CudaSlice<f32>,
18142 yv: &mut CudaSlice<f32>,
18143 yg: &mut CudaSlice<f32>,
18144 in_f: usize,
18145 out_q: usize,
18146 out_kv: usize,
18147 out_g: usize,
18148 t: usize,
18149 ) -> Result<(), Box<dyn std::error::Error>> {
18150 if t == 0
18151 || t > 8
18152 || in_f % 8 != 0
18153 || x_t.len() < t * in_f
18154 || yq.len() < t * out_q
18155 || yk.len() < t * out_kv
18156 || yv.len() < t * out_kv
18157 || (out_g > 0 && yg.len() < t * out_g)
18158 {
18159 return Err("matvec_bf16_qkvg_tcol geometry".into());
18160 }
18161 let grid = out_q + 2 * out_kv + out_g;
18162 let cfg = LaunchConfig {
18163 grid_dim: (grid as u32, 1, 1),
18164 block_dim: (mmv_block(), 1, 1),
18165 shared_mem_bytes: 0,
18166 };
18167 let (ini, oq, okv, og, ti) = (
18168 in_f as i32,
18169 out_q as i32,
18170 out_kv as i32,
18171 out_g as i32,
18172 t as i32,
18173 );
18174 let __s_b = self.gpu.stream();
18175 let f = self.func("matvec_bf16_qkvg_tcol");
18181 let mut b = __s_b.launch_builder(&f);
18182 b.arg(wq)
18183 .arg(wk)
18184 .arg(wv)
18185 .arg(wg)
18186 .arg(x_t)
18187 .arg(yq)
18188 .arg(yk)
18189 .arg(yv)
18190 .arg(yg)
18191 .arg(&ini)
18192 .arg(&oq)
18193 .arg(&okv)
18194 .arg(&og)
18195 .arg(&ti);
18196 unsafe {
18197 b.launch(cfg)?;
18198 }
18199 Ok(())
18200 }
18201
18202 pub fn matvec_bf16_qkvg_into(
18203 &self,
18204 wq: &CudaSlice<u8>,
18205 wk: &CudaSlice<u8>,
18206 wv: &CudaSlice<u8>,
18207 wg: &CudaSlice<u8>,
18208 x: &CudaSlice<f32>,
18209 yq: &mut CudaSlice<f32>,
18210 yk: &mut CudaSlice<f32>,
18211 yv: &mut CudaSlice<f32>,
18212 yg: &mut CudaSlice<f32>,
18213 in_f: usize,
18214 out_q: usize,
18215 out_kv: usize,
18216 out_g: usize,
18217 ) -> Result<(), Box<dyn std::error::Error>> {
18218 if in_f % 8 != 0
18219 || wq.len() != out_q * in_f * 2
18220 || wk.len() != out_kv * in_f * 2
18221 || wv.len() != out_kv * in_f * 2
18222 || wg.len() < out_g * in_f * 2
18223 || x.len() < in_f
18224 || yq.len() < out_q
18225 || yk.len() < out_kv
18226 || yv.len() < out_kv
18227 || (out_g > 0 && yg.len() < out_g)
18228 {
18229 return Err(format!(
18230 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
18231 )
18232 .into());
18233 }
18234 let f = self.func("matvec_bf16_qkvg");
18235 let cfg = LaunchConfig {
18236 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
18237 block_dim: (mmv_block(), 1, 1),
18238 shared_mem_bytes: 0,
18239 };
18240 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
18241 let __s_b = self.gpu.stream();
18242 let mut b = __s_b.launch_builder(&f);
18243 b.arg(wq)
18244 .arg(wk)
18245 .arg(wv)
18246 .arg(wg)
18247 .arg(x)
18248 .arg(yq)
18249 .arg(yk)
18250 .arg(yv)
18251 .arg(yg)
18252 .arg(&inf)
18253 .arg(&oq)
18254 .arg(&okv)
18255 .arg(&og);
18256 unsafe {
18257 b.launch(cfg)?;
18258 }
18259 Ok(())
18260 }
18261
18262 pub fn matvec_bf16_b4_into(
18264 &self,
18265 w: [&CudaSlice<u8>; 4],
18266 x: &CudaSlice<f32>,
18267 y: &mut CudaSlice<f32>,
18268 block_cols: usize,
18269 out_f: usize,
18270 ) -> Result<(), Box<dyn std::error::Error>> {
18271 if block_cols % 8 != 0
18272 || x.len() < 4 * block_cols
18273 || y.len() < out_f
18274 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18275 {
18276 return Err(format!(
18277 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
18278 x.len()
18279 )
18280 .into());
18281 }
18282 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18285 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
18286 let f = self.func(if x2 {
18287 "matvec_bf16_b4_x2"
18288 } else {
18289 "matvec_bf16_b4"
18290 });
18291 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
18292 let cfg = LaunchConfig {
18293 grid_dim: (grid as u32, 1, 1),
18294 block_dim: (mmv_block(), 1, 1),
18295 shared_mem_bytes: 0,
18296 };
18297 let (bc, of) = (block_cols as i32, out_f as i32);
18298 let __s_b = self.gpu.stream();
18299 let mut b = __s_b.launch_builder(&f);
18300 b.arg(w[0])
18301 .arg(w[1])
18302 .arg(w[2])
18303 .arg(w[3])
18304 .arg(x)
18305 .arg(y)
18306 .arg(&bc)
18307 .arg(&of);
18308 unsafe {
18309 b.launch(cfg)?;
18310 }
18311 Ok(())
18312 }
18313
18314 pub fn matvec_bf16_b4_tcol_into(
18320 &self,
18321 w: [&CudaSlice<u8>; 4],
18322 x_t: &CudaSlice<f32>,
18323 y_t: &mut CudaSlice<f32>,
18324 block_cols: usize,
18325 out_f: usize,
18326 t: usize,
18327 ) -> Result<(), Box<dyn std::error::Error>> {
18328 if block_cols % 8 != 0
18329 || t == 0
18330 || t > 8
18331 || x_t.len() < t * 4 * block_cols
18332 || y_t.len() < t * out_f
18333 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18334 {
18335 return Err(format!(
18336 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18337 x_t.len()
18338 )
18339 .into());
18340 }
18341 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18342 return Err(
18343 "b4 tcol verify is qualified against the plain b4 kernel only \
18344 (MEMRA_B4_X2=1 is a different t=1 program)"
18345 .into(),
18346 );
18347 }
18348 let cfg = LaunchConfig {
18352 grid_dim: (out_f as u32, 1, 1),
18353 block_dim: (mmv_block(), 1, 1),
18354 shared_mem_bytes: 0,
18355 };
18356 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18357 let __s_b = self.gpu.stream();
18358 let f = self.func("matvec_bf16_b4_tcol");
18359 let mut b = __s_b.launch_builder(&f);
18360 b.arg(w[0])
18361 .arg(w[1])
18362 .arg(w[2])
18363 .arg(w[3])
18364 .arg(x_t)
18365 .arg(y_t)
18366 .arg(&bc)
18367 .arg(&of)
18368 .arg(&ti);
18369 unsafe {
18370 b.launch(cfg)?;
18371 }
18372 Ok(())
18373 }
18374
18375 pub fn q8_0_row_bytes(in_f: usize) -> usize {
18378 in_f / 32 * 34
18379 }
18380
18381 pub fn encode_q8_0_from_bf16(
18385 &self,
18386 w_bf16: &CudaSlice<u8>,
18387 out: &mut CudaSlice<u8>,
18388 in_f: usize,
18389 out_f: usize,
18390 ) -> Result<(), Box<dyn std::error::Error>> {
18391 if in_f % 32 != 0
18392 || w_bf16.len() < in_f * out_f * 2
18393 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18394 {
18395 return Err(format!(
18396 "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
18397 w_bf16.len(),
18398 out.len()
18399 )
18400 .into());
18401 }
18402 let f = self.func("encode_q8_0_rows_from_bf16");
18403 const PAIRS_PER_BLOCK: u32 = 4;
18406 let pairs = (out_f * (in_f / 32)) as u64;
18407 let cfg = LaunchConfig {
18408 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18409 block_dim: (32, PAIRS_PER_BLOCK, 1),
18410 shared_mem_bytes: 0,
18411 };
18412 let (ini, outi) = (in_f as i32, out_f as i32);
18413 let __s_b = self.gpu.stream();
18414 let mut b = __s_b.launch_builder(&f);
18415 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18416 unsafe {
18417 b.launch(cfg)?;
18418 }
18419 Ok(())
18420 }
18421
18422 pub fn encode_q8_0_from_bf16_view(
18426 &self,
18427 w_bf16: &cudarc::driver::CudaView<'_, u8>,
18428 out: &mut CudaSlice<u8>,
18429 in_f: usize,
18430 out_f: usize,
18431 ) -> Result<(), Box<dyn std::error::Error>> {
18432 if in_f % 32 != 0
18433 || w_bf16.len() < in_f * out_f * 2
18434 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18435 {
18436 return Err(format!(
18437 "encode_q8_0_from_bf16_view geometry in={in_f} out={out_f} src={} dst={}",
18438 w_bf16.len(),
18439 out.len()
18440 )
18441 .into());
18442 }
18443 let f = self.func("encode_q8_0_rows_from_bf16");
18444 const PAIRS_PER_BLOCK: u32 = 4;
18445 let pairs = (out_f * (in_f / 32)) as u64;
18446 let cfg = LaunchConfig {
18447 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18448 block_dim: (32, PAIRS_PER_BLOCK, 1),
18449 shared_mem_bytes: 0,
18450 };
18451 let (ini, outi) = (in_f as i32, out_f as i32);
18452 let __s_b = self.gpu.stream();
18453 let mut b = __s_b.launch_builder(&f);
18454 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18455 unsafe {
18456 b.launch(cfg)?;
18457 }
18458 Ok(())
18459 }
18460
18461 #[allow(clippy::too_many_arguments)]
18466 pub fn qmatvec_q8_0_qkv_rp_into(
18467 &self,
18468 wq: &CudaSlice<u8>,
18469 wk: &CudaSlice<u8>,
18470 wv: &CudaSlice<u8>,
18471 aq: &CudaSlice<i8>,
18472 ad: &CudaSlice<f32>,
18473 yq: &mut CudaSlice<f32>,
18474 yk: &mut CudaSlice<f32>,
18475 yv: &mut CudaSlice<f32>,
18476 in_f: usize,
18477 out_q: usize,
18478 out_kv: usize,
18479 ) -> Result<(), Box<dyn std::error::Error>> {
18480 const ROWS_PER_BLOCK: u32 = 4; let rows = out_q + 2 * out_kv;
18482 let nblk = in_f / 32;
18483 if in_f % 32 != 0
18484 || aq.len() < in_f
18485 || ad.len() < nblk
18486 || yq.len() < out_q
18487 || yk.len() < out_kv
18488 || yv.len() < out_kv
18489 || wq.len() < out_q * nblk * 34
18490 || wk.len() < out_kv * nblk * 34
18491 || wv.len() < out_kv * nblk * 34
18492 {
18493 return Err(
18494 format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
18495 );
18496 }
18497 let f = self.func("qmatvec_q8_0_qkv_rp");
18498 let cfg = LaunchConfig {
18499 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18500 block_dim: (32, ROWS_PER_BLOCK, 1),
18501 shared_mem_bytes: 0,
18502 };
18503 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18504 let __s_b = self.gpu.stream();
18505 let mut b = __s_b.launch_builder(&f);
18506 b.arg(wq)
18507 .arg(wk)
18508 .arg(wv)
18509 .arg(aq)
18510 .arg(ad)
18511 .arg(yq)
18512 .arg(yk)
18513 .arg(yv)
18514 .arg(&ini)
18515 .arg(&oq)
18516 .arg(&okv);
18517 unsafe {
18518 b.launch(cfg)?;
18519 }
18520 Ok(())
18521 }
18522
18523 #[allow(clippy::too_many_arguments)]
18527 pub fn qmatvec_q8_0_b4_rp_into(
18528 &self,
18529 w: [&CudaSlice<u8>; 4],
18530 aq: &CudaSlice<i8>,
18531 ad: &CudaSlice<f32>,
18532 y: &mut CudaSlice<f32>,
18533 block_cols: usize,
18534 out_f: usize,
18535 ) -> Result<(), Box<dyn std::error::Error>> {
18536 const ROWS_PER_BLOCK: u32 = 4; let nblk = block_cols / 32;
18538 if block_cols % 32 != 0
18539 || aq.len() < 4 * block_cols
18540 || ad.len() < 4 * nblk
18541 || y.len() < out_f
18542 || w.iter().any(|p| p.len() < out_f * nblk * 34)
18543 {
18544 return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
18545 }
18546 let f = self.func("qmatvec_q8_0_b4_rp");
18547 let cfg = LaunchConfig {
18548 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18549 block_dim: (32, ROWS_PER_BLOCK, 1),
18550 shared_mem_bytes: 0,
18551 };
18552 let (bc, of) = (block_cols as i32, out_f as i32);
18553 let __s_b = self.gpu.stream();
18554 let mut b = __s_b.launch_builder(&f);
18555 b.arg(w[0])
18556 .arg(w[1])
18557 .arg(w[2])
18558 .arg(w[3])
18559 .arg(aq)
18560 .arg(ad)
18561 .arg(y)
18562 .arg(&bc)
18563 .arg(&of);
18564 unsafe {
18565 b.launch(cfg)?;
18566 }
18567 Ok(())
18568 }
18569
18570 fn matvec_bf16_via_q8_mirror_t(
18573 &self,
18574 data: &CudaSlice<u8>,
18575 x: &CudaSlice<f32>,
18576 y: &mut CudaSlice<f32>,
18577 in_f: usize,
18578 out_f: usize,
18579 t: usize,
18580 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18581 use cudarc::driver::DevicePtr;
18582 let key = {
18583 let s = self.gpu.stream();
18584 let (p, _g) = data.device_ptr(&s);
18585 (p as u64, in_f as u32, out_f as u32)
18586 };
18587 {
18588 let mut mirrors = self
18589 .w8_mirrors
18590 .lock()
18591 .map_err(|_| "w8 mirror map is poisoned")?;
18592 if !mirrors.contains_key(&key) {
18593 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18594 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18595 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18596 mirrors.insert(key, planar);
18597 }
18598 }
18599 let nblk = in_f / 32;
18600 let akey = in_f * 64 + t.min(32);
18602 {
18603 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18604 if !act.contains_key(&akey) {
18605 let aq = self.alloc_i8_uninit(32 * in_f)?;
18606 let ad = self.alloc_uninit::<f32>(32 * nblk)?;
18607 act.insert(akey, (aq, ad));
18608 }
18609 let (aq, ad) = act.get_mut(&akey).expect("just inserted");
18610 self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
18611 }
18612 let mirrors = self
18613 .w8_mirrors
18614 .lock()
18615 .map_err(|_| "w8 mirror map is poisoned")?;
18616 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18617 let mirror = mirrors.get(&key).expect("built above");
18618 let (aq, ad) = act.get(&akey).expect("built above");
18619 const ROWS_PER_BLOCK: u32 = 4;
18620 let (ini, of) = (in_f as i32, out_f as i32);
18621 if q8t_wonce_on() && t <= 32 {
18625 let f = self.func(if t <= 8 {
18626 "qmatvec_q8_0_rows_tw"
18627 } else {
18628 "qmatvec_q8_0_rows_tw32"
18629 });
18630 let cfg = LaunchConfig {
18631 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18632 block_dim: (32, ROWS_PER_BLOCK, 1),
18633 shared_mem_bytes: 0,
18634 };
18635 let ti = t as i32;
18636 let __s_b = self.gpu.stream();
18637 let mut b = __s_b.launch_builder(&f);
18638 b.arg(mirror)
18639 .arg(aq)
18640 .arg(ad)
18641 .arg(&mut *y)
18642 .arg(&ini)
18643 .arg(&of)
18644 .arg(&ti);
18645 unsafe {
18646 b.launch(cfg)?;
18647 }
18648 return Ok(Some(()));
18649 }
18650 let f = self.func("qmatvec_q8_0_rows_t");
18651 let cfg = LaunchConfig {
18652 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
18653 block_dim: (32, ROWS_PER_BLOCK, 1),
18654 shared_mem_bytes: 0,
18655 };
18656 let __s_b = self.gpu.stream();
18657 let mut b = __s_b.launch_builder(&f);
18658 b.arg(mirror)
18659 .arg(aq)
18660 .arg(ad)
18661 .arg(&mut *y)
18662 .arg(&ini)
18663 .arg(&of);
18664 unsafe {
18665 b.launch(cfg)?;
18666 }
18667 Ok(Some(()))
18668 }
18669
18670 fn matvec_bf16_via_q8_mirror(
18673 &self,
18674 data: &CudaSlice<u8>,
18675 x: &CudaSlice<f32>,
18676 y: &mut CudaSlice<f32>,
18677 in_f: usize,
18678 out_f: usize,
18679 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18680 use cudarc::driver::DevicePtr;
18681 let key = {
18682 let s = self.gpu.stream();
18683 let (p, _g) = data.device_ptr(&s);
18684 (p as u64, in_f as u32, out_f as u32)
18685 };
18686 {
18687 let mut mirrors = self
18688 .w8_mirrors
18689 .lock()
18690 .map_err(|_| "w8 mirror map is poisoned")?;
18691 if !mirrors.contains_key(&key) {
18692 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18693 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18694 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18695 mirrors.insert(key, planar);
18696 if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
18702 eprintln!(
18703 "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
18704 mirrors.len()
18705 );
18706 }
18707 }
18708 }
18709 let nblk = in_f / 32;
18710 {
18711 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18712 if !act.contains_key(&in_f) {
18713 let aq = self.alloc_uninit::<i8>(in_f)?;
18714 let ad = self.alloc_uninit::<f32>(nblk)?;
18715 act.insert(in_f, (aq, ad));
18716 }
18717 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
18718 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
18719 }
18720 let mirrors = self
18721 .w8_mirrors
18722 .lock()
18723 .map_err(|_| "w8 mirror map is poisoned")?;
18724 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18725 let mirror = mirrors.get(&key).expect("built above");
18726 let (aq, ad) = act.get(&in_f).expect("built above");
18727 self.qmatvec_mmvq_into(
18728 mirror,
18729 aq,
18730 ad,
18731 1,
18732 in_f,
18733 out_f,
18734 QT_Q8_0,
18735 Self::q8_0_row_bytes(in_f),
18736 1.0,
18737 true,
18738 y,
18739 )?;
18740 Ok(Some(()))
18741 }
18742
18743 #[allow(clippy::too_many_arguments)]
18748 pub fn qmatvec_q8_0_qkv_rp_t_into(
18749 &self,
18750 wq: &CudaSlice<u8>,
18751 wk: &CudaSlice<u8>,
18752 wv: &CudaSlice<u8>,
18753 aq: &CudaSlice<i8>,
18754 ad: &CudaSlice<f32>,
18755 yq: &mut CudaSlice<f32>,
18756 yk: &mut CudaSlice<f32>,
18757 yv: &mut CudaSlice<f32>,
18758 in_f: usize,
18759 out_q: usize,
18760 out_kv: usize,
18761 t: usize,
18762 ) -> Result<(), Box<dyn std::error::Error>> {
18763 const ROWS_PER_BLOCK: u32 = 4;
18764 let rows = out_q + 2 * out_kv;
18765 let nblk = in_f / 32;
18766 if in_f % 32 != 0
18767 || t == 0
18768 || aq.len() < t * in_f
18769 || ad.len() < t * nblk
18770 || yq.len() < t * out_q
18771 || yk.len() < t * out_kv
18772 || yv.len() < t * out_kv
18773 {
18774 return Err(format!("q8_0 qkv rp_t geometry in={in_f} t={t}").into());
18775 }
18776 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18777 if q8t_wonce_on() && t <= 32 {
18781 let f = self.func(if t <= 8 {
18782 "qmatvec_q8_0_qkv_rp_tw"
18783 } else {
18784 "qmatvec_q8_0_qkv_rp_tw32"
18785 });
18786 let cfg = LaunchConfig {
18787 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18788 block_dim: (32, ROWS_PER_BLOCK, 1),
18789 shared_mem_bytes: 0,
18790 };
18791 let ti = t as i32;
18792 let __s_b = self.gpu.stream();
18793 let mut b = __s_b.launch_builder(&f);
18794 b.arg(wq)
18795 .arg(wk)
18796 .arg(wv)
18797 .arg(aq)
18798 .arg(ad)
18799 .arg(yq)
18800 .arg(yk)
18801 .arg(yv)
18802 .arg(&ini)
18803 .arg(&oq)
18804 .arg(&okv)
18805 .arg(&ti);
18806 unsafe {
18807 b.launch(cfg)?;
18808 }
18809 return Ok(());
18810 }
18811 let f = self.func("qmatvec_q8_0_qkv_rp_t");
18812 let cfg = LaunchConfig {
18813 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
18814 block_dim: (32, ROWS_PER_BLOCK, 1),
18815 shared_mem_bytes: 0,
18816 };
18817 let __s_b = self.gpu.stream();
18818 let mut b = __s_b.launch_builder(&f);
18819 b.arg(wq)
18820 .arg(wk)
18821 .arg(wv)
18822 .arg(aq)
18823 .arg(ad)
18824 .arg(yq)
18825 .arg(yk)
18826 .arg(yv)
18827 .arg(&ini)
18828 .arg(&oq)
18829 .arg(&okv);
18830 unsafe {
18831 b.launch(cfg)?;
18832 }
18833 Ok(())
18834 }
18835
18836 #[allow(clippy::too_many_arguments)]
18839 pub fn qmatvec_q8_0_b4_rp_t_into(
18840 &self,
18841 w: [&CudaSlice<u8>; 4],
18842 aq: &CudaSlice<i8>,
18843 ad: &CudaSlice<f32>,
18844 y: &mut CudaSlice<f32>,
18845 block_cols: usize,
18846 out_f: usize,
18847 t: usize,
18848 ) -> Result<(), Box<dyn std::error::Error>> {
18849 const ROWS_PER_BLOCK: u32 = 4;
18850 let nblk = block_cols / 32;
18851 if block_cols % 32 != 0
18852 || t == 0
18853 || aq.len() < t * 4 * block_cols
18854 || ad.len() < t * 4 * nblk
18855 || y.len() < t * out_f
18856 {
18857 return Err(format!("q8_0 b4 rp_t geometry cols={block_cols} t={t}").into());
18858 }
18859 let (bc, of) = (block_cols as i32, out_f as i32);
18860 if q8t_wonce_on() && t <= 32 {
18863 let f = self.func(if t <= 8 {
18864 "qmatvec_q8_0_b4_rp_tw"
18865 } else {
18866 "qmatvec_q8_0_b4_rp_tw32"
18867 });
18868 let cfg = LaunchConfig {
18869 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18870 block_dim: (32, ROWS_PER_BLOCK, 1),
18871 shared_mem_bytes: 0,
18872 };
18873 let ti = t as i32;
18874 let __s_b = self.gpu.stream();
18875 let mut b = __s_b.launch_builder(&f);
18876 b.arg(w[0])
18877 .arg(w[1])
18878 .arg(w[2])
18879 .arg(w[3])
18880 .arg(aq)
18881 .arg(ad)
18882 .arg(y)
18883 .arg(&bc)
18884 .arg(&of)
18885 .arg(&ti);
18886 unsafe {
18887 b.launch(cfg)?;
18888 }
18889 return Ok(());
18890 }
18891 let f = self.func("qmatvec_q8_0_b4_rp_t");
18892 let cfg = LaunchConfig {
18893 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
18894 block_dim: (32, ROWS_PER_BLOCK, 1),
18895 shared_mem_bytes: 0,
18896 };
18897 let __s_b = self.gpu.stream();
18898 let mut b = __s_b.launch_builder(&f);
18899 b.arg(w[0])
18900 .arg(w[1])
18901 .arg(w[2])
18902 .arg(w[3])
18903 .arg(aq)
18904 .arg(ad)
18905 .arg(y)
18906 .arg(&bc)
18907 .arg(&of);
18908 unsafe {
18909 b.launch(cfg)?;
18910 }
18911 Ok(())
18912 }
18913
18914 fn matvec_bf16_view_via_q8_mirror(
18924 &self,
18925 data: &cudarc::driver::CudaView<'_, u8>,
18926 x: &CudaSlice<f32>,
18927 y: &mut CudaSlice<f32>,
18928 in_f: usize,
18929 out_f: usize,
18930 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18931 use cudarc::driver::DevicePtr;
18932 let key = {
18933 let s = self.gpu.stream();
18934 let (p, _g) = data.device_ptr(&s);
18935 (p as u64, in_f as u32, out_f as u32)
18936 };
18937 {
18938 let mut mirrors = self
18939 .w8_mirrors
18940 .lock()
18941 .map_err(|_| "w8 mirror map is poisoned")?;
18942 if !mirrors.contains_key(&key) {
18943 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18944 self.encode_q8_0_from_bf16_view(data, &mut interleaved, in_f, out_f)?;
18945 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18946 mirrors.insert(key, planar);
18947 eprintln!("[w8-view] mirror built in_f={in_f} out_f={out_f}");
18951 }
18952 }
18953 let nblk = in_f / 32;
18954 {
18955 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18956 if !act.contains_key(&in_f) {
18957 let aq = self.alloc_uninit::<i8>(in_f)?;
18958 let ad = self.alloc_uninit::<f32>(nblk)?;
18959 act.insert(in_f, (aq, ad));
18960 }
18961 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
18962 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
18963 }
18964 let mirrors = self
18965 .w8_mirrors
18966 .lock()
18967 .map_err(|_| "w8 mirror map is poisoned")?;
18968 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18969 let mirror = mirrors.get(&key).expect("built above");
18970 let (aq, ad) = act.get(&in_f).expect("built above");
18971 self.qmatvec_mmvq_into(
18972 mirror,
18973 aq,
18974 ad,
18975 1,
18976 in_f,
18977 out_f,
18978 QT_Q8_0,
18979 Self::q8_0_row_bytes(in_f),
18980 1.0,
18981 true,
18982 y,
18983 )?;
18984 Ok(Some(()))
18985 }
18986
18987 pub fn matvec_bf16_into(
18988 &self,
18989 data: &CudaSlice<u8>,
18990 x: &CudaSlice<f32>,
18991 y: &mut CudaSlice<f32>,
18992 in_f: usize,
18993 out_f: usize,
18994 ) -> Result<(), Box<dyn std::error::Error>> {
18995 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18996 return Err(format!(
18997 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18998 data.len(),
18999 x.len(),
19000 y.len()
19001 )
19002 .into());
19003 }
19004 if step_tp_w8_on() && w8_hybrid_on() && in_f % 32 == 0 && out_f >= 64 {
19011 if let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)? {
19012 return Ok(());
19013 }
19014 }
19015 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19019 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
19020 && in_f <= 2048;
19021 if x4 {
19022 let f = self.func("matvec_bf16_f32acc_x4");
19023 let cfg = LaunchConfig {
19024 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
19025 block_dim: (mmv_block(), 1, 1),
19026 shared_mem_bytes: 0,
19027 };
19028 let (ini, outi) = (in_f as i32, out_f as i32);
19029 let __s_b = self.gpu.stream();
19030 let mut b = __s_b.launch_builder(&f);
19031 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
19032 unsafe {
19033 b.launch(cfg)?;
19034 }
19035 return Ok(());
19036 }
19037 let f = self.func("matvec_bf16_f32acc");
19038 let cfg = LaunchConfig {
19039 grid_dim: (out_f as u32, 1, 1),
19040 block_dim: (mmv_block(), 1, 1),
19041 shared_mem_bytes: 0,
19042 };
19043 let ini = in_f as i32;
19044 let __s_b = self.gpu.stream();
19045 let mut b = __s_b.launch_builder(&f);
19046 b.arg(data).arg(x).arg(y).arg(&ini);
19047 unsafe {
19048 b.launch(cfg)?;
19049 }
19050 Ok(())
19051 }
19052
19053 pub fn matvec_bf16_view_into(
19056 &self,
19057 data: &cudarc::driver::CudaView<'_, u8>,
19058 x: &CudaSlice<f32>,
19059 y: &mut CudaSlice<f32>,
19060 in_f: usize,
19061 out_f: usize,
19062 ) -> Result<(), Box<dyn std::error::Error>> {
19063 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
19064 return Err(format!(
19065 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
19066 data.len(),
19067 x.len(),
19068 y.len()
19069 )
19070 .into());
19071 }
19072 if w8_view_on() && step_tp_w8_on() && w8_hybrid_on() && in_f % 32 == 0 && out_f >= 64 {
19073 if let Some(()) = self.matvec_bf16_view_via_q8_mirror(data, x, y, in_f, out_f)? {
19074 return Ok(());
19075 }
19076 }
19077 let f = self.func("matvec_bf16_f32acc");
19078 let cfg = LaunchConfig {
19079 grid_dim: (out_f as u32, 1, 1),
19080 block_dim: (mmv_block(), 1, 1),
19081 shared_mem_bytes: 0,
19082 };
19083 let ini = in_f as i32;
19084 let __s_b = self.gpu.stream();
19085 let mut b = __s_b.launch_builder(&f);
19086 b.arg(data).arg(x).arg(y).arg(&ini);
19087 unsafe {
19088 b.launch(cfg)?;
19089 }
19090 Ok(())
19091 }
19092
19093 pub fn matvec_bf16_raw_out(
19096 &self,
19097 w: &CudaSlice<u8>,
19098 x: &CudaSlice<f32>,
19099 y_raw: u64,
19100 in_f: usize,
19101 out_f: usize,
19102 ) -> Result<(), Box<dyn std::error::Error>> {
19103 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
19104 return Err("matvec_bf16_raw_out geometry".into());
19105 }
19106 let f = self.func("matvec_bf16_f32acc");
19107 let cfg = LaunchConfig {
19108 grid_dim: (out_f as u32, 1, 1),
19109 block_dim: (mmv_block(), 1, 1),
19110 shared_mem_bytes: 0,
19111 };
19112 let ini = in_f as i32;
19113 let __s_b = self.gpu.stream();
19114 let mut b = __s_b.launch_builder(&f);
19115 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
19116 unsafe {
19117 b.launch(cfg)?;
19118 }
19119 Ok(())
19120 }
19121
19122 pub fn add3_raw(
19126 &self,
19127 a: &CudaSlice<f32>,
19128 b: &CudaSlice<f32>,
19129 sh_raw: u64,
19130 scale_raw: u64,
19131 dst: &mut CudaSlice<f32>,
19132 n: usize,
19133 ) -> Result<(), Box<dyn std::error::Error>> {
19134 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
19135 return Err("add3_raw geometry".into());
19136 }
19137 let f = self.func("add3_f32");
19138 let cfg = LaunchConfig {
19139 grid_dim: ((n as u32).div_ceil(256), 1, 1),
19140 block_dim: (256, 1, 1),
19141 shared_mem_bytes: 0,
19142 };
19143 let ni = n as i32;
19144 let __s_b = self.gpu.stream();
19145 let mut bld = __s_b.launch_builder(&f);
19146 bld.arg(a)
19147 .arg(b)
19148 .arg(&sh_raw)
19149 .arg(&scale_raw)
19150 .arg(dst)
19151 .arg(&ni);
19152 unsafe {
19153 bld.launch(cfg)?;
19154 }
19155 Ok(())
19156 }
19157
19158 pub fn matvec_bf16_down_addscale_into(
19161 &self,
19162 w: &CudaSlice<u8>,
19163 x: &CudaSlice<f32>,
19164 scale: &CudaSlice<f32>,
19165 dst: &mut CudaSlice<f32>,
19166 in_f: usize,
19167 out_f: usize,
19168 ) -> Result<(), Box<dyn std::error::Error>> {
19169 if w.len() != in_f * out_f * 2
19170 || x.len() < in_f
19171 || in_f % 8 != 0
19172 || dst.len() < out_f
19173 || scale.is_empty()
19174 {
19175 return Err("matvec_bf16_down_addscale geometry".into());
19176 }
19177 let f = self.func("matvec_bf16_down_addscale");
19178 let cfg = LaunchConfig {
19179 grid_dim: (out_f as u32, 1, 1),
19180 block_dim: (mmv_block(), 1, 1),
19181 shared_mem_bytes: 0,
19182 };
19183 let ini = in_f as i32;
19184 let __s_b = self.gpu.stream();
19185 let mut b = __s_b.launch_builder(&f);
19186 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
19187 unsafe {
19188 b.launch(cfg)?;
19189 }
19190 Ok(())
19191 }
19192
19193 #[allow(clippy::too_many_arguments)]
19197 pub fn matvec_bf16_dual_silu_rows_into(
19198 &self,
19199 wg: &CudaSlice<u8>,
19200 wu: &CudaSlice<u8>,
19201 x: &CudaSlice<f32>,
19202 act: &mut CudaSlice<f32>,
19203 in_f: usize,
19204 out_f: usize,
19205 limit: Option<f32>,
19206 t: usize,
19207 ) -> Result<(), Box<dyn std::error::Error>> {
19208 if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
19209 return Err("matvec_bf16_dual_silu_rows geometry".into());
19210 }
19211 let f = self.func("matvec_bf16_dual_silu_rows");
19212 let cfg = LaunchConfig {
19213 grid_dim: (out_f as u32, t as u32, 1),
19214 block_dim: (mmv_block(), 1, 1),
19215 shared_mem_bytes: 0,
19216 };
19217 let (ini, outi) = (in_f as i32, out_f as i32);
19218 let lim = limit.unwrap_or(0.0);
19219 let __s_b = self.gpu.stream();
19220 let mut b = __s_b.launch_builder(&f);
19221 b.arg(wg)
19222 .arg(wu)
19223 .arg(x)
19224 .arg(&mut *act)
19225 .arg(&ini)
19226 .arg(&outi)
19227 .arg(&lim);
19228 unsafe {
19229 b.launch(cfg)?;
19230 }
19231 Ok(())
19232 }
19233
19234 pub fn matvec_bf16_rows_into(
19236 &self,
19237 w: &CudaSlice<u8>,
19238 x: &CudaSlice<f32>,
19239 y: &mut CudaSlice<f32>,
19240 in_f: usize,
19241 out_f: usize,
19242 t: usize,
19243 ) -> Result<(), Box<dyn std::error::Error>> {
19244 if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
19245 return Err("matvec_bf16_rows geometry".into());
19246 }
19247 if t >= 2 && t <= 32 && step_tp_w8_on() && w8_hybrid_on() && in_f % 32 == 0 && out_f >= 64 {
19252 if let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)? {
19253 return Ok(());
19254 }
19255 }
19256 if t == 1 && step_tp_w8_on() && w8_hybrid_on() && in_f % 32 == 0 && out_f >= 64 {
19262 if let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)? {
19263 return Ok(());
19264 }
19265 }
19266 let f = self.func("matvec_bf16_f32acc_x4_rows");
19267 let cfg = LaunchConfig {
19268 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
19269 block_dim: (mmv_block(), 1, 1),
19270 shared_mem_bytes: 0,
19271 };
19272 let (ini, outi) = (in_f as i32, out_f as i32);
19273 let __s_b = self.gpu.stream();
19274 let mut b = __s_b.launch_builder(&f);
19275 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
19276 unsafe {
19277 b.launch(cfg)?;
19278 }
19279 Ok(())
19280 }
19281
19282 pub fn matvec_bf16_dual_silu_into(
19283 &self,
19284 wg: &CudaSlice<u8>,
19285 wu: &CudaSlice<u8>,
19286 x: &CudaSlice<f32>,
19287 act: &mut CudaSlice<f32>,
19288 in_f: usize,
19289 out_f: usize,
19290 limit: Option<f32>,
19291 ) -> Result<(), Box<dyn std::error::Error>> {
19292 if wg.len() != in_f * out_f * 2
19293 || wu.len() != in_f * out_f * 2
19294 || x.len() < in_f
19295 || in_f % 8 != 0
19296 || act.len() < out_f
19297 {
19298 return Err("matvec_bf16_dual_silu geometry".into());
19299 }
19300 let f = self.func("matvec_bf16_dual_silu");
19301 let cfg = LaunchConfig {
19302 grid_dim: (out_f as u32, 1, 1),
19303 block_dim: (mmv_block(), 1, 1),
19304 shared_mem_bytes: 0,
19305 };
19306 let (ini, outi) = (in_f as i32, out_f as i32);
19307 let lim = limit.unwrap_or(0.0);
19308 let __s_b = self.gpu.stream();
19309 let mut b = __s_b.launch_builder(&f);
19310 b.arg(wg)
19311 .arg(wu)
19312 .arg(x)
19313 .arg(act)
19314 .arg(&ini)
19315 .arg(&outi)
19316 .arg(&lim);
19317 unsafe {
19318 b.launch(cfg)?;
19319 }
19320 Ok(())
19321 }
19322
19323 #[allow(clippy::too_many_arguments)]
19326 pub fn matvec_bf16_dual_view_into(
19327 &self,
19328 wg: &cudarc::driver::CudaView<'_, u8>,
19329 wu: &cudarc::driver::CudaView<'_, u8>,
19330 x: &CudaSlice<f32>,
19331 yg: &mut CudaSlice<f32>,
19332 yu: &mut CudaSlice<f32>,
19333 in_f: usize,
19334 out_f: usize,
19335 ) -> Result<(), Box<dyn std::error::Error>> {
19336 if wg.len() != in_f * out_f * 2
19337 || wu.len() != in_f * out_f * 2
19338 || x.len() < in_f
19339 || in_f % 8 != 0
19340 || yg.len() < out_f
19341 || yu.len() < out_f
19342 {
19343 return Err(format!(
19344 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
19345 wg.len(),
19346 wu.len(),
19347 x.len()
19348 )
19349 .into());
19350 }
19351 let f = self.func("matvec_bf16_dual");
19352 let cfg = LaunchConfig {
19353 grid_dim: ((2 * out_f) as u32, 1, 1),
19354 block_dim: (mmv_block(), 1, 1),
19355 shared_mem_bytes: 0,
19356 };
19357 let (ini, outi) = (in_f as i32, out_f as i32);
19358 let __s_b = self.gpu.stream();
19359 let mut b = __s_b.launch_builder(&f);
19360 b.arg(wg)
19361 .arg(wu)
19362 .arg(x)
19363 .arg(yg)
19364 .arg(yu)
19365 .arg(&ini)
19366 .arg(&outi);
19367 unsafe {
19368 b.launch(cfg)?;
19369 }
19370 Ok(())
19371 }
19372
19373 #[allow(clippy::too_many_arguments)]
19375 pub fn matvec_bf16_dual_into(
19376 &self,
19377 wg: &CudaSlice<u8>,
19378 wu: &CudaSlice<u8>,
19379 x: &CudaSlice<f32>,
19380 yg: &mut CudaSlice<f32>,
19381 yu: &mut CudaSlice<f32>,
19382 in_f: usize,
19383 out_f: usize,
19384 ) -> Result<(), Box<dyn std::error::Error>> {
19385 if wg.len() != in_f * out_f * 2
19386 || wu.len() != in_f * out_f * 2
19387 || x.len() < in_f
19388 || in_f % 8 != 0
19389 || yg.len() < out_f
19390 || yu.len() < out_f
19391 {
19392 return Err(format!(
19393 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
19394 wg.len(),
19395 wu.len(),
19396 x.len()
19397 )
19398 .into());
19399 }
19400 let f = self.func("matvec_bf16_dual");
19401 let cfg = LaunchConfig {
19402 grid_dim: ((2 * out_f) as u32, 1, 1),
19403 block_dim: (mmv_block(), 1, 1),
19404 shared_mem_bytes: 0,
19405 };
19406 let (ini, outi) = (in_f as i32, out_f as i32);
19407 let __s_b = self.gpu.stream();
19408 let mut b = __s_b.launch_builder(&f);
19409 b.arg(wg)
19410 .arg(wu)
19411 .arg(x)
19412 .arg(yg)
19413 .arg(yu)
19414 .arg(&ini)
19415 .arg(&outi);
19416 unsafe {
19417 b.launch(cfg)?;
19418 }
19419 Ok(())
19420 }
19421
19422 pub(crate) fn matvec_bf16_dual(
19425 &self,
19426 wg: &CudaSlice<u8>,
19427 wu: &CudaSlice<u8>,
19428 x: &CudaSlice<f32>,
19429 in_f: usize,
19430 out_f: usize,
19431 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19432 if wg.len() != in_f * out_f * 2
19433 || wu.len() != in_f * out_f * 2
19434 || x.len() < in_f
19435 || in_f % 8 != 0
19436 {
19437 return Err(format!(
19438 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
19439 wg.len(),
19440 wu.len(),
19441 x.len()
19442 )
19443 .into());
19444 }
19445 let mut yg = self.alloc_uninit::<f32>(out_f)?;
19446 let mut yu = self.alloc_uninit::<f32>(out_f)?;
19447 let f = self.func("matvec_bf16_dual");
19448 let cfg = LaunchConfig {
19449 grid_dim: ((2 * out_f) as u32, 1, 1),
19450 block_dim: (mmv_block(), 1, 1),
19451 shared_mem_bytes: 0,
19452 };
19453 let (ini, outi) = (in_f as i32, out_f as i32);
19454 let __s_b = self.gpu.stream();
19455 let mut b = __s_b.launch_builder(&f);
19456 b.arg(wg)
19457 .arg(wu)
19458 .arg(x)
19459 .arg(&mut yg)
19460 .arg(&mut yu)
19461 .arg(&ini)
19462 .arg(&outi);
19463 unsafe {
19464 b.launch(cfg)?;
19465 }
19466 Ok((yg, yu))
19467 }
19468
19469 #[allow(clippy::too_many_arguments)]
19470 fn linear_bf16_chunked_inner(
19471 &self,
19472 x: &CudaSlice<f32>,
19473 data: &CudaSlice<u8>,
19474 m: usize,
19475 in_f: usize,
19476 out_f: usize,
19477 exact: bool,
19478 canonical_chunk_rows: Option<usize>,
19479 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19480 const CHUNK_BYTES: usize = 256 << 20;
19481 if m == 1
19484 && !exact
19485 && canonical_chunk_rows.is_none()
19486 && in_f % 8 == 0
19487 && Self::bf16_mmv_on()
19488 {
19489 return self.matvec_bf16(data, x, in_f, out_f);
19490 }
19491 if m >= 16
19496 && !exact
19497 && canonical_chunk_rows.is_none()
19498 && data.len() == in_f * out_f * 2
19499 && crate::f16_ffi::pp_bf16_enabled()
19500 {
19501 if let Some(y) = self.bf16_tc_gemm(data, x, m, in_f, out_f)? {
19504 return Ok(y);
19505 }
19506 }
19507 let row_bytes = in_f
19508 .checked_mul(std::mem::size_of::<f32>())
19509 .ok_or("BF16 chunk row byte count overflow")?;
19510 if row_bytes == 0 || out_f == 0 {
19511 return Err("BF16 chunk dimensions must be nonzero".into());
19512 }
19513 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
19514 let chunk_rows = match canonical_chunk_rows {
19515 Some(rows) if rows == 0 => {
19516 return Err("canonical BF16 chunk rows must be nonzero".into());
19517 }
19518 Some(rows) if rows > max_chunk_rows => {
19519 return Err(format!(
19520 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
19521 )
19522 .into());
19523 }
19524 Some(rows) if out_f % rows != 0 => {
19525 return Err(format!(
19526 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
19527 )
19528 .into());
19529 }
19530 Some(rows) => rows,
19531 None => max_chunk_rows,
19532 };
19533 if chunk_rows >= out_f {
19534 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
19535 return if exact {
19536 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
19537 } else {
19538 self.linear(x, &wf32, m, in_f, out_f)
19539 };
19540 }
19541 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19542 let mut r0 = 0usize;
19543 while r0 < out_f {
19544 let rows = chunk_rows.min(out_f - r0);
19545 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
19546 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
19547 let yc = if exact {
19548 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
19549 } else {
19550 self.linear(x, &wf32, m, in_f, rows)?
19551 };
19552 for mi in 0..m {
19554 let src = yc.slice(mi * rows..(mi + 1) * rows);
19555 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
19556 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
19557 }
19558 r0 += rows;
19559 }
19560 Ok(y)
19561 }
19562
19563 pub fn linear_bf16_resident(
19567 &self,
19568 x: &CudaSlice<f32>,
19569 data: &CudaSlice<u8>,
19570 m: usize,
19571 in_f: usize,
19572 out_f: usize,
19573 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19574 if data.len() != in_f * out_f * 2 {
19575 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19576 }
19577 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
19578 }
19579
19580 pub fn linear_bf16_resident_canonical_rows(
19586 &self,
19587 x: &CudaSlice<f32>,
19588 data: &CudaSlice<u8>,
19589 m: usize,
19590 in_f: usize,
19591 out_f: usize,
19592 canonical_chunk_rows: usize,
19593 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19594 if data.len() != in_f * out_f * 2 {
19595 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19596 }
19597 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
19598 }
19599
19600 pub fn linear_f32_resident_canonical_rows(
19605 &self,
19606 x: &CudaSlice<f32>,
19607 data: &CudaSlice<f32>,
19608 m: usize,
19609 in_f: usize,
19610 out_f: usize,
19611 canonical_chunk_rows: usize,
19612 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19613 self.linear_f32_resident_canonical_rows_inner(
19614 x,
19615 data,
19616 m,
19617 in_f,
19618 out_f,
19619 canonical_chunk_rows,
19620 false,
19621 )
19622 }
19623
19624 pub fn linear_f32_resident_canonical_rows_strided(
19630 &self,
19631 x: &CudaSlice<f32>,
19632 data: &CudaSlice<f32>,
19633 m: usize,
19634 in_f: usize,
19635 out_f: usize,
19636 canonical_chunk_rows: usize,
19637 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19638 self.linear_f32_resident_canonical_rows_inner(
19639 x,
19640 data,
19641 m,
19642 in_f,
19643 out_f,
19644 canonical_chunk_rows,
19645 true,
19646 )
19647 }
19648
19649 fn linear_f32_resident_canonical_rows_inner(
19650 &self,
19651 x: &CudaSlice<f32>,
19652 data: &CudaSlice<f32>,
19653 m: usize,
19654 in_f: usize,
19655 out_f: usize,
19656 canonical_chunk_rows: usize,
19657 strided_output: bool,
19658 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19659 if data.len() != in_f * out_f {
19660 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19661 }
19662 if canonical_chunk_rows == 0
19663 || canonical_chunk_rows > out_f
19664 || out_f % canonical_chunk_rows != 0
19665 {
19666 return Err(format!(
19667 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19668 )
19669 .into());
19670 }
19671 if canonical_chunk_rows == out_f {
19672 return self.linear(x, data, m, in_f, out_f);
19673 }
19674
19675 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19676 let input = x.slice(0..x.len());
19677 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19678 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19679 if m == 1 {
19680 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19681 self.linear_device_into(
19682 &input,
19683 &weights,
19684 &mut destination,
19685 1,
19686 in_f,
19687 canonical_chunk_rows,
19688 )?;
19689 continue;
19690 }
19691 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
19692 if strided_output {
19693 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
19694 } else {
19695 for token in 0..m {
19696 let source = chunk
19697 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
19698 let mut destination =
19699 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
19700 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
19701 }
19702 }
19703 }
19704 Ok(y)
19705 }
19706
19707 pub fn linear_f32_resident_canonical_rows_t1_into(
19712 &self,
19713 x: &CudaSlice<f32>,
19714 data: &CudaSlice<f32>,
19715 y: &mut CudaSlice<f32>,
19716 in_f: usize,
19717 out_f: usize,
19718 canonical_chunk_rows: usize,
19719 ) -> Result<(), Box<dyn std::error::Error>> {
19720 if data.len() != in_f * out_f {
19721 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19722 }
19723 if y.len() != out_f || x.len() != in_f {
19724 return Err(format!(
19725 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
19726 x.len(),
19727 y.len()
19728 )
19729 .into());
19730 }
19731 if canonical_chunk_rows == 0
19732 || canonical_chunk_rows > out_f
19733 || out_f % canonical_chunk_rows != 0
19734 {
19735 return Err(format!(
19736 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19737 )
19738 .into());
19739 }
19740 let input = x.slice(0..x.len());
19741 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19742 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19743 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19744 self.linear_device_into(
19745 &input,
19746 &weights,
19747 &mut destination,
19748 1,
19749 in_f,
19750 canonical_chunk_rows,
19751 )?;
19752 }
19753 Ok(())
19754 }
19755
19756 pub fn linear_t1_into(
19759 &self,
19760 x: &cudarc::driver::CudaView<'_, f32>,
19761 w: &cudarc::driver::CudaView<'_, f32>,
19762 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
19763 in_f: usize,
19764 out_f: usize,
19765 ) -> Result<(), Box<dyn std::error::Error>> {
19766 self.linear_device_into(x, w, y, 1, in_f, out_f)
19767 }
19768
19769 pub fn linear_decode_exact(
19776 &self,
19777 x: &CudaSlice<f32>,
19778 w: &CudaSlice<f32>,
19779 m_tokens: usize,
19780 in_f: usize,
19781 out_f: usize,
19782 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19783 if m_tokens == 1 {
19784 return self.linear(x, w, 1, in_f, out_f);
19785 }
19786 let xv = self.view(x, m_tokens * in_f);
19787 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
19788 for t in 0..m_tokens {
19789 let row = xv.slice(t * in_f..(t + 1) * in_f);
19790 let mut xr = self.alloc_uninit::<f32>(in_f)?;
19791 self.copy_view_into(&mut xr, 0, &row, in_f)?;
19792 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
19793 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
19794 }
19795 Ok(y)
19796 }
19797
19798 pub fn linear(
19799 &self,
19800 x: &CudaSlice<f32>,
19801 w: &CudaSlice<f32>,
19802 m_tokens: usize,
19803 in_f: usize,
19804 out_f: usize,
19805 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19806 self.linear_device(x, w, m_tokens, in_f, out_f)
19807 }
19808
19809 fn linear_device<I>(
19810 &self,
19811 x: &I,
19812 w: &I,
19813 m_tokens: usize,
19814 in_f: usize,
19815 out_f: usize,
19816 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
19817 where
19818 I: cudarc::driver::DevicePtr<f32>,
19819 {
19820 let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; self.linear_device_into(x, w, &mut c, m_tokens, in_f, out_f)?;
19822 Ok(c)
19823 }
19824
19825 fn linear_device_into<I, O>(
19826 &self,
19827 x: &I,
19828 w: &I,
19829 c: &mut O,
19830 m_tokens: usize,
19831 in_f: usize,
19832 out_f: usize,
19833 ) -> Result<(), Box<dyn std::error::Error>>
19834 where
19835 I: cudarc::driver::DevicePtr<f32>,
19836 O: cudarc::driver::DevicePtrMut<f32>,
19837 {
19838 use cudarc::cublaslt::{Matmul, MatmulConfig};
19839 let cfg = MatmulConfig {
19840 transa: true,
19841 transb: false,
19842 transc: false,
19843 m: out_f as u64,
19844 n: m_tokens as u64,
19845 k: in_f as u64,
19846 alpha: 1.0,
19847 lda: in_f as i64,
19848 ldb: in_f as i64,
19849 beta: 0.0,
19850 ldc: out_f as i64,
19851 stride_a: None,
19852 stride_b: None,
19853 stride_c: None,
19854 stride_bias: None,
19855 batch_size: None,
19856 };
19857 let blas = self.gpu.blas();
19858 unsafe {
19859 blas.matmul(cfg, w, x, c, None, None)?;
19860 }
19861 Ok(())
19862 }
19863
19864 pub fn sdpa_naive(
19873 &self,
19874 q: &CudaSlice<f32>,
19875 k: &CudaSlice<f32>,
19876 v: &CudaSlice<f32>,
19877 o: &mut CudaSlice<f32>,
19878 head_dim: usize,
19879 n_head: usize,
19880 n_head_kv: usize,
19881 t: usize,
19882 t_kv: usize,
19883 scale: f32,
19884 causal: bool,
19885 ) -> Result<(), Box<dyn std::error::Error>> {
19886 if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
19887 return self.sdpa_naive_gmem(
19888 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19889 );
19890 }
19891 let f = self.func("sdpa_naive_f32");
19892 let cfg = LaunchConfig {
19893 grid_dim: (n_head as u32, t as u32, 1),
19894 block_dim: (128, 1, 1),
19895 shared_mem_bytes: (t_kv * 4) as u32,
19896 };
19897 let (hd, nh, nhkv, ti, tkvi, cz) = (
19898 head_dim as i32,
19899 n_head as i32,
19900 n_head_kv as i32,
19901 t as i32,
19902 t_kv as i32,
19903 causal as i32,
19904 );
19905 let __s_b = self.gpu.stream();
19906 let mut b = __s_b.launch_builder(&f);
19907 b.arg(q)
19908 .arg(k)
19909 .arg(v)
19910 .arg(o)
19911 .arg(&hd)
19912 .arg(&nh)
19913 .arg(&nhkv)
19914 .arg(&ti)
19915 .arg(&tkvi)
19916 .arg(&scale)
19917 .arg(&cz);
19918 unsafe {
19919 b.launch(cfg)?;
19920 }
19921 Ok(())
19922 }
19923
19924 #[allow(clippy::too_many_arguments)]
19933 pub fn sdpa_naive_gmem(
19934 &self,
19935 q: &CudaSlice<f32>,
19936 k: &CudaSlice<f32>,
19937 v: &CudaSlice<f32>,
19938 o: &mut CudaSlice<f32>,
19939 head_dim: usize,
19940 n_head: usize,
19941 n_head_kv: usize,
19942 t: usize,
19943 t_kv: usize,
19944 scale: f32,
19945 causal: bool,
19946 ) -> Result<(), Box<dyn std::error::Error>> {
19947 let ws_len = n_head
19948 .checked_mul(t)
19949 .and_then(|x| x.checked_mul(t_kv))
19950 .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
19951 let ws_bytes = ws_len
19952 .checked_mul(std::mem::size_of::<f32>())
19953 .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
19954 if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
19955 return Err(format!(
19956 "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
19957 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
19958 needs a tiled/flash kernel, not the naive oracle"
19959 )
19960 .into());
19961 }
19962 let mut scores = self.uninit(ws_len)?;
19963 let f = self.func("sdpa_naive_gmem_f32");
19964 let cfg = LaunchConfig {
19965 grid_dim: (n_head as u32, t as u32, 1),
19966 block_dim: (128, 1, 1),
19967 shared_mem_bytes: 0,
19968 };
19969 let (hd, nh, nhkv, ti, tkvi, cz) = (
19970 head_dim as i32,
19971 n_head as i32,
19972 n_head_kv as i32,
19973 t as i32,
19974 t_kv as i32,
19975 causal as i32,
19976 );
19977 let __s_b = self.gpu.stream();
19978 let mut b = __s_b.launch_builder(&f);
19979 b.arg(q)
19980 .arg(k)
19981 .arg(v)
19982 .arg(o)
19983 .arg(&mut scores)
19984 .arg(&hd)
19985 .arg(&nh)
19986 .arg(&nhkv)
19987 .arg(&ti)
19988 .arg(&tkvi)
19989 .arg(&scale)
19990 .arg(&cz);
19991 unsafe {
19992 b.launch(cfg)?;
19993 }
19994 Ok(())
19995 }
19996
19997 #[allow(clippy::too_many_arguments)]
20002 pub fn sdpa_naive_island(
20003 &self,
20004 q: &CudaSlice<f32>,
20005 k: &CudaSlice<f32>,
20006 v: &CudaSlice<f32>,
20007 o: &mut CudaSlice<f32>,
20008 span_id: &CudaSlice<i32>,
20009 head_dim: usize,
20010 n_head: usize,
20011 n_head_kv: usize,
20012 t: usize,
20013 t_kv: usize,
20014 scale: f32,
20015 window: usize,
20016 ) -> Result<(), Box<dyn std::error::Error>> {
20017 let f = self.func("sdpa_naive_island_f32");
20018 let cfg = LaunchConfig {
20019 grid_dim: (n_head as u32, t as u32, 1),
20020 block_dim: (128, 1, 1),
20021 shared_mem_bytes: (t_kv * 4) as u32,
20022 };
20023 let (hd, nh, nhkv, ti, tkvi, wi) = (
20024 head_dim as i32,
20025 n_head as i32,
20026 n_head_kv as i32,
20027 t as i32,
20028 t_kv as i32,
20029 window as i32,
20030 );
20031 let __s_b = self.gpu.stream();
20032 let mut b = __s_b.launch_builder(&f);
20033 b.arg(q)
20034 .arg(k)
20035 .arg(v)
20036 .arg(o)
20037 .arg(span_id)
20038 .arg(&hd)
20039 .arg(&nh)
20040 .arg(&nhkv)
20041 .arg(&ti)
20042 .arg(&tkvi)
20043 .arg(&scale)
20044 .arg(&wi);
20045 unsafe {
20046 b.launch(cfg)?;
20047 }
20048 Ok(())
20049 }
20050
20051 #[allow(clippy::too_many_arguments)]
20053 pub fn sdpa_naive_w(
20054 &self,
20055 q: &CudaSlice<f32>,
20056 k: &CudaSlice<f32>,
20057 v: &CudaSlice<f32>,
20058 o: &mut CudaSlice<f32>,
20059 head_dim: usize,
20060 n_head: usize,
20061 n_head_kv: usize,
20062 t: usize,
20063 t_kv: usize,
20064 scale: f32,
20065 causal: bool,
20066 window: usize,
20067 ) -> Result<(), Box<dyn std::error::Error>> {
20068 let f = self.func("sdpa_naive_w_f32");
20069 let cfg = LaunchConfig {
20070 grid_dim: (n_head as u32, t as u32, 1),
20071 block_dim: (128, 1, 1),
20072 shared_mem_bytes: (t_kv * 4) as u32,
20073 };
20074 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20075 head_dim as i32,
20076 n_head as i32,
20077 n_head_kv as i32,
20078 t as i32,
20079 t_kv as i32,
20080 causal as i32,
20081 window as i32,
20082 );
20083 let __s_b = self.gpu.stream();
20084 let mut b = __s_b.launch_builder(&f);
20085 b.arg(q)
20086 .arg(k)
20087 .arg(v)
20088 .arg(o)
20089 .arg(&hd)
20090 .arg(&nh)
20091 .arg(&nhkv)
20092 .arg(&ti)
20093 .arg(&tkvi)
20094 .arg(&scale)
20095 .arg(&cz)
20096 .arg(&wi);
20097 unsafe {
20098 b.launch(cfg)?;
20099 }
20100 Ok(())
20101 }
20102
20103 #[allow(clippy::too_many_arguments)]
20113 pub fn sdpa_naive_w_lo(
20114 &self,
20115 q: &CudaSlice<f32>,
20116 k: &CudaSlice<f32>,
20117 v: &CudaSlice<f32>,
20118 o: &mut CudaSlice<f32>,
20119 head_dim: usize,
20120 n_head: usize,
20121 n_head_kv: usize,
20122 t: usize,
20123 t_kv: usize,
20124 scale: f32,
20125 causal: bool,
20126 window: usize,
20127 ) -> Result<(), Box<dyn std::error::Error>> {
20128 let kv_lo = if window > 0 {
20129 (t_kv - t + 1).saturating_sub(window)
20130 } else {
20131 0
20132 };
20133 let smem = (t_kv - kv_lo) * 4;
20134 if smem > 48 * 1024 {
20135 return Err(format!(
20136 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
20137 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
20138 a window this wide needs the multi-pass long-ctx kernel"
20139 )
20140 .into());
20141 }
20142 let f = self.func("sdpa_naive_w_lo_f32");
20143 let cfg = LaunchConfig {
20144 grid_dim: (n_head as u32, t as u32, 1),
20145 block_dim: (128, 1, 1),
20146 shared_mem_bytes: smem as u32,
20147 };
20148 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
20149 head_dim as i32,
20150 n_head as i32,
20151 n_head_kv as i32,
20152 t as i32,
20153 t_kv as i32,
20154 causal as i32,
20155 window as i32,
20156 kv_lo as i32,
20157 );
20158 let __s_b = self.gpu.stream();
20159 let mut b = __s_b.launch_builder(&f);
20160 b.arg(q)
20161 .arg(k)
20162 .arg(v)
20163 .arg(o)
20164 .arg(&hd)
20165 .arg(&nh)
20166 .arg(&nhkv)
20167 .arg(&ti)
20168 .arg(&tkvi)
20169 .arg(&scale)
20170 .arg(&cz)
20171 .arg(&wi)
20172 .arg(&lo);
20173 unsafe {
20174 b.launch(cfg)?;
20175 }
20176 Ok(())
20177 }
20178
20179 pub fn sdpa_naive_view(
20181 &self,
20182 q: &CudaSlice<f32>,
20183 k: &cudarc::driver::CudaView<f32>,
20184 v: &cudarc::driver::CudaView<f32>,
20185 o: &mut CudaSlice<f32>,
20186 head_dim: usize,
20187 n_head: usize,
20188 n_head_kv: usize,
20189 t: usize,
20190 t_kv: usize,
20191 scale: f32,
20192 causal: bool,
20193 ) -> Result<(), Box<dyn std::error::Error>> {
20194 let f = self.func("sdpa_naive_f32");
20195 let cfg = LaunchConfig {
20196 grid_dim: (n_head as u32, t as u32, 1),
20197 block_dim: (128, 1, 1),
20198 shared_mem_bytes: (t_kv * 4) as u32,
20199 };
20200 let (hd, nh, nhkv, ti, tkvi, cz) = (
20201 head_dim as i32,
20202 n_head as i32,
20203 n_head_kv as i32,
20204 t as i32,
20205 t_kv as i32,
20206 causal as i32,
20207 );
20208 let __s_b = self.gpu.stream();
20209 let mut b = __s_b.launch_builder(&f);
20210 b.arg(q)
20211 .arg(k)
20212 .arg(v)
20213 .arg(o)
20214 .arg(&hd)
20215 .arg(&nh)
20216 .arg(&nhkv)
20217 .arg(&ti)
20218 .arg(&tkvi)
20219 .arg(&scale)
20220 .arg(&cz);
20221 unsafe {
20222 b.launch(cfg)?;
20223 }
20224 Ok(())
20225 }
20226
20227 #[allow(clippy::too_many_arguments)]
20235 pub fn fa_dequant_kv_view_f32(
20236 &self,
20237 k: &cudarc::driver::CudaView<u8>,
20238 v: &cudarc::driver::CudaView<u8>,
20239 kf: &mut CudaSlice<f32>,
20240 vf: &mut CudaSlice<f32>,
20241 kv_dim_k: usize,
20242 kv_dim_v: usize,
20243 t_kv: usize,
20244 k_tok_bytes: usize,
20245 v_tok_bytes: usize,
20246 g: bool,
20247 ) -> Result<(), Box<dyn std::error::Error>> {
20248 let f = if g {
20249 self.func_g("fa_dequant_kv_ws_f32")
20250 } else {
20251 self.func("fa_dequant_kv_ws_f32")
20252 };
20253 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20254 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20255 let cfg = LaunchConfig {
20256 grid_dim: (nblk.max(1), 1, 1),
20257 block_dim: (256, 1, 1),
20258 shared_mem_bytes: 0,
20259 };
20260 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20261 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20262 let __s_b = self.gpu.stream();
20263 let mut b = __s_b.launch_builder(&f);
20264 b.arg(k)
20265 .arg(v)
20266 .arg(&mut *kf)
20267 .arg(&mut *vf)
20268 .arg(&kdk)
20269 .arg(&kdv)
20270 .arg(&tkvi)
20271 .arg(&ktb)
20272 .arg(&vtb);
20273 unsafe {
20274 b.launch(cfg)?;
20275 }
20276 Ok(())
20277 }
20278
20279 #[allow(clippy::too_many_arguments)]
20280 pub fn sdpa_naive_quantized_view(
20281 &self,
20282 q: &CudaSlice<f32>,
20283 k: &cudarc::driver::CudaView<u8>,
20284 v: &cudarc::driver::CudaView<u8>,
20285 o: &mut CudaSlice<f32>,
20286 head_dim: usize,
20287 n_head: usize,
20288 n_head_kv: usize,
20289 t: usize,
20290 t_kv: usize,
20291 scale: f32,
20292 causal: bool,
20293 k_tok_bytes: usize,
20294 v_tok_bytes: usize,
20295 ) -> Result<(), Box<dyn std::error::Error>> {
20296 let kv_dim = n_head_kv * head_dim;
20297 let mut kf = self.uninit(t_kv * kv_dim)?;
20298 let mut vf = self.uninit(t_kv * kv_dim)?;
20299 let f = self.func("fa_dequant_kv_ws_f32");
20300 let total = (2 * t_kv * kv_dim) as u64;
20301 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20302 let cfg = LaunchConfig {
20303 grid_dim: (nblk.max(1), 1, 1),
20304 block_dim: (256, 1, 1),
20305 shared_mem_bytes: 0,
20306 };
20307 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
20308 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
20309 let __s_b = self.gpu.stream();
20310 let mut b = __s_b.launch_builder(&f);
20311 b.arg(k)
20312 .arg(v)
20313 .arg(&mut kf)
20314 .arg(&mut vf)
20315 .arg(&kv_dim_i)
20316 .arg(&kv_dim_i)
20317 .arg(&t_kv_i)
20318 .arg(&k_tok_bytes_i)
20319 .arg(&v_tok_bytes_i);
20320 unsafe { b.launch(cfg)? };
20321 self.sdpa_naive(
20322 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20323 )
20324 }
20325
20326 #[allow(clippy::too_many_arguments)]
20338 pub fn sdpa_naive_w_quantized_view(
20339 &self,
20340 q: &CudaSlice<f32>,
20341 k: &cudarc::driver::CudaView<u8>,
20342 v: &cudarc::driver::CudaView<u8>,
20343 o: &mut CudaSlice<f32>,
20344 head_dim: usize,
20345 n_head: usize,
20346 n_head_kv: usize,
20347 t: usize,
20348 t_kv: usize,
20349 scale: f32,
20350 causal: bool,
20351 window: usize,
20352 k_tok_bytes: usize,
20353 v_tok_bytes: usize,
20354 ) -> Result<(), Box<dyn std::error::Error>> {
20355 let kv_dim = n_head_kv * head_dim;
20356 let mut kf = self.uninit(t_kv * kv_dim)?;
20357 let mut vf = self.uninit(t_kv * kv_dim)?;
20358 let f = self.func("fa_dequant_kv_ws_f32");
20359 let total = (2 * t_kv * kv_dim) as u64;
20360 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20361 let cfg = LaunchConfig {
20362 grid_dim: (nblk.max(1), 1, 1),
20363 block_dim: (256, 1, 1),
20364 shared_mem_bytes: 0,
20365 };
20366 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
20367 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
20368 let __s_b = self.gpu.stream();
20369 let mut b = __s_b.launch_builder(&f);
20370 b.arg(k)
20371 .arg(v)
20372 .arg(&mut kf)
20373 .arg(&mut vf)
20374 .arg(&kv_dim_i)
20375 .arg(&kv_dim_i)
20376 .arg(&t_kv_i)
20377 .arg(&k_tok_bytes_i)
20378 .arg(&v_tok_bytes_i);
20379 unsafe { b.launch(cfg)? };
20380 self.sdpa_naive_w(
20381 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
20382 )
20383 }
20384
20385 pub fn fa_prefill(
20389 &self,
20390 q: &CudaSlice<f32>,
20391 k: &CudaSlice<f32>,
20392 v: &CudaSlice<f32>,
20393 o: &mut CudaSlice<f32>,
20394 head_dim: usize,
20395 n_head: usize,
20396 n_head_kv: usize,
20397 t: usize,
20398 t_kv: usize,
20399 scale: f32,
20400 causal: bool,
20401 ) -> Result<(), Box<dyn std::error::Error>> {
20402 if portable_mma_gated() {
20403 return self.sdpa_naive(
20404 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20405 );
20406 }
20407 let fa3_on = head_dim == 256
20415 && causal
20416 && t == t_kv
20417 && match std::env::var("MEMRA_FA3").as_deref() {
20418 Ok("0") => false,
20419 Ok("1") => {
20423 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
20424 true
20425 }
20426 _ => cfg!(memra_hopper_mma),
20427 };
20428 if fa3_on {
20429 let n = t * n_head * head_dim;
20430 let nkv = t * n_head_kv * head_dim;
20431 let mut q16 = self.alloc_u8_uninit(n * 2)?;
20432 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
20433 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
20434 self.f32_to_bf16_into(q, &mut q16, n)?;
20435 self.f32_to_bf16_into(k, &mut k16, nkv)?;
20436 self.f32_to_bf16_into(v, &mut v16, nkv)?;
20437 let rc = {
20438 use cudarc::driver::{DevicePtr, DevicePtrMut};
20439 let stream = self.gpu.stream();
20440 let (qp, _g1) = q16.device_ptr(&stream);
20441 let (kp, _g2) = k16.device_ptr(&stream);
20442 let (vp, _g3) = v16.device_ptr(&stream);
20443 let (op, _g4) = o.device_ptr_mut(&stream);
20444 unsafe {
20445 memra_fa3_prefill(
20446 qp as *const core::ffi::c_void,
20447 kp as *const core::ffi::c_void,
20448 vp as *const core::ffi::c_void,
20449 op as *mut f32,
20450 t as i32,
20451 n_head as i32,
20452 n_head_kv as i32,
20453 head_dim as i32,
20454 scale,
20455 stream.cu_stream() as *mut core::ffi::c_void,
20456 )
20457 }
20458 };
20459 if rc != 0 {
20460 return Err(format!("memra_fa3_prefill rc={rc}").into());
20461 }
20462 return Ok(());
20463 }
20464 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20469 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
20470 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
20471 const BLOCK_Q: usize = 64;
20472 const BKX: usize = 32;
20473 let f = self.func("fa_prefill_bf16_p1");
20474 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
20475 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
20476 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20477 f.set_attribute(
20478 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20479 shmem as i32,
20480 )?;
20481 let cfg = LaunchConfig {
20482 grid_dim: (
20483 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20484 n_head as u32,
20485 1,
20486 ),
20487 block_dim: (32, 4, 1),
20488 shared_mem_bytes: shmem,
20489 };
20490 let (hd, nh, nhkv, ti, tkvi, cz) = (
20491 head_dim as i32,
20492 n_head as i32,
20493 n_head_kv as i32,
20494 t as i32,
20495 t_kv as i32,
20496 causal as i32,
20497 );
20498 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20499 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20500 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20501 let __s_b = self.gpu.stream();
20502 let mut b = __s_b.launch_builder(&f);
20503 b.arg(&qb)
20504 .arg(&kb)
20505 .arg(&vb)
20506 .arg(o)
20507 .arg(&hd)
20508 .arg(&nh)
20509 .arg(&nhkv)
20510 .arg(&ti)
20511 .arg(&tkvi)
20512 .arg(&scale)
20513 .arg(&cz);
20514 unsafe {
20515 b.launch(cfg)?;
20516 }
20517 return Ok(());
20518 }
20519 const BK: usize = 32;
20525 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
20528 let (block_q, warps, w2_sfx): (usize, u32, &str) =
20529 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
20530 let hd_sfx = fa_hd_suffix(head_dim)?;
20534 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20535 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
20540 let (kb16, vb16) = if bf16kv {
20541 let n = t_kv * n_head_kv * head_dim;
20542 let mut kb = self.alloc_u8_uninit(n * 2)?;
20543 let mut vb = self.alloc_u8_uninit(n * 2)?;
20544 let fcv = self.func("f32_to_bf16_bulk");
20545 let ni = n as i64;
20546 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20547 let __s_b = self.gpu.stream();
20548 let mut b = __s_b.launch_builder(&fcv);
20549 b.arg(k).arg(&mut kb).arg(&ni);
20550 unsafe {
20551 b.launch(cfgc)?;
20552 }
20553 let __s_b = self.gpu.stream();
20554 let mut b = __s_b.launch_builder(&fcv);
20555 b.arg(v).arg(&mut vb).arg(&ni);
20556 unsafe {
20557 b.launch(cfgc)?;
20558 }
20559 (Some(kb), Some(vb))
20560 } else {
20561 (None, None)
20562 };
20563 let f = self.func(&if bf16kv {
20564 format!("fa_prefill_bf16kv_pp{hd_sfx}")
20565 } else {
20566 format!(
20567 "fa_prefill_f32{}{}{hd_sfx}",
20568 if floor { "" } else { "_pp" },
20569 if floor { "" } else { w2_sfx }
20570 )
20571 });
20572 let kv_stages = if bf16kv { 2 } else { 1 };
20575 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20576 + 4 * (block_q * BK + 2 * block_q)) as u32;
20577 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20578 f.set_attribute(
20579 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20580 shmem as i32,
20581 )?;
20582 let cfg = LaunchConfig {
20583 grid_dim: (
20584 (t as u32 + block_q as u32 - 1) / block_q as u32,
20585 n_head as u32,
20586 1,
20587 ),
20588 block_dim: (32, warps, 1),
20589 shared_mem_bytes: shmem,
20590 };
20591 let (hd, nh, nhkv, ti, tkvi, cz) = (
20592 head_dim as i32,
20593 n_head as i32,
20594 n_head_kv as i32,
20595 t as i32,
20596 t_kv as i32,
20597 causal as i32,
20598 );
20599 let __s_b = self.gpu.stream();
20600 let mut b = __s_b.launch_builder(&f);
20601 b.arg(q);
20602 match (&kb16, &vb16) {
20603 (Some(kb), Some(vb)) => {
20604 b.arg(kb).arg(vb);
20605 }
20606 _ => {
20607 b.arg(k).arg(v);
20608 }
20609 }
20610 b.arg(o)
20611 .arg(&hd)
20612 .arg(&nh)
20613 .arg(&nhkv)
20614 .arg(&ti)
20615 .arg(&tkvi)
20616 .arg(&scale)
20617 .arg(&cz);
20618 unsafe {
20619 b.launch(cfg)?;
20620 }
20621 Ok(())
20622 }
20623
20624 #[allow(clippy::too_many_arguments)]
20628 pub fn fa_prefill_w(
20629 &self,
20630 q: &CudaSlice<f32>,
20631 k: &CudaSlice<f32>,
20632 v: &CudaSlice<f32>,
20633 o: &mut CudaSlice<f32>,
20634 head_dim: usize,
20635 n_head: usize,
20636 n_head_kv: usize,
20637 t: usize,
20638 t_kv: usize,
20639 scale: f32,
20640 causal: bool,
20641 window: usize,
20642 ) -> Result<(), Box<dyn std::error::Error>> {
20643 if portable_mma_gated() {
20646 return self.sdpa_naive_w(
20647 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
20648 );
20649 }
20650 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20654 let faw_f32 =
20655 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
20656 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20657 self.fa_prefill_w_arm(
20658 q,
20659 k,
20660 v,
20661 o,
20662 head_dim,
20663 n_head,
20664 n_head_kv,
20665 t,
20666 t_kv,
20667 scale,
20668 causal,
20669 window,
20670 floor || faw_f32,
20671 floor,
20672 )
20673 }
20674
20675 #[allow(clippy::too_many_arguments)]
20678 pub fn fa_prefill_w_pre(
20679 &self,
20680 qb: &CudaSlice<u8>,
20681 kb: &CudaSlice<u8>,
20682 vb: &CudaSlice<u8>,
20683 o: &mut CudaSlice<f32>,
20684 head_dim: usize,
20685 n_head: usize,
20686 n_head_kv: usize,
20687 t: usize,
20688 t_kv: usize,
20689 scale: f32,
20690 causal: bool,
20691 window: usize,
20692 v_f16: bool,
20693 ) -> Result<(), Box<dyn std::error::Error>> {
20694 const BLOCK_Q: usize = 64;
20695 const BK: usize = 32;
20696 debug_assert_eq!(head_dim, 256);
20697 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20698 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
20699 if hp {
20700 const BLOCK_QH: usize = 32;
20701 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20704 let vh: &CudaSlice<u8> = if v_f16 {
20705 vb
20706 } else {
20707 let n = t_kv * n_head_kv * head_dim;
20708 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
20709 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
20710 }
20711 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
20712 vguard.as_ref().unwrap()
20713 };
20714 let f = self.func("fa_prefill_w_bf16_p1h2");
20715 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20716 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20717 f.set_attribute(
20718 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20719 shmem as i32,
20720 )?;
20721 let cfg = LaunchConfig {
20722 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20723 block_dim: (32, 4, 1),
20724 shared_mem_bytes: shmem,
20725 };
20726 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20727 head_dim as i32,
20728 n_head as i32,
20729 n_head_kv as i32,
20730 t as i32,
20731 t_kv as i32,
20732 causal as i32,
20733 window as i32,
20734 );
20735 let __s_b = self.gpu.stream();
20736 let mut b = __s_b.launch_builder(&f);
20737 b.arg(qb)
20738 .arg(kb)
20739 .arg(vh)
20740 .arg(o)
20741 .arg(&hd)
20742 .arg(&nh)
20743 .arg(&nhkv)
20744 .arg(&ti)
20745 .arg(&tkvi)
20746 .arg(&scale)
20747 .arg(&cz)
20748 .arg(&wi);
20749 unsafe {
20750 b.launch(cfg)?;
20751 }
20752 return Ok(());
20753 }
20754 let f = self.func("fa_prefill_w_bf16_p1");
20755 let shmem =
20756 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20757 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20758 f.set_attribute(
20759 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20760 shmem as i32,
20761 )?;
20762 let cfg = LaunchConfig {
20763 grid_dim: (
20764 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20765 n_head as u32,
20766 1,
20767 ),
20768 block_dim: (32, 4, 1),
20769 shared_mem_bytes: shmem,
20770 };
20771 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20772 head_dim as i32,
20773 n_head as i32,
20774 n_head_kv as i32,
20775 t as i32,
20776 t_kv as i32,
20777 causal as i32,
20778 window as i32,
20779 );
20780 let __s_b = self.gpu.stream();
20781 let mut b = __s_b.launch_builder(&f);
20782 b.arg(qb)
20783 .arg(kb)
20784 .arg(vb)
20785 .arg(o)
20786 .arg(&hd)
20787 .arg(&nh)
20788 .arg(&nhkv)
20789 .arg(&ti)
20790 .arg(&tkvi)
20791 .arg(&scale)
20792 .arg(&cz)
20793 .arg(&wi);
20794 unsafe {
20795 b.launch(cfg)?;
20796 }
20797 Ok(())
20798 }
20799
20800 #[allow(clippy::too_many_arguments)]
20802 pub fn fa_prefill_w_arm(
20803 &self,
20804 q: &CudaSlice<f32>,
20805 k: &CudaSlice<f32>,
20806 v: &CudaSlice<f32>,
20807 o: &mut CudaSlice<f32>,
20808 head_dim: usize,
20809 n_head: usize,
20810 n_head_kv: usize,
20811 t: usize,
20812 t_kv: usize,
20813 scale: f32,
20814 causal: bool,
20815 window: usize,
20816 f32_stage: bool,
20817 floor: bool,
20818 ) -> Result<(), Box<dyn std::error::Error>> {
20819 const BLOCK_Q: usize = 64;
20820 const BK: usize = 32;
20821 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
20822 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20826 let p1 = !floor
20827 && !f32_stage
20828 && *P1_ON.get_or_init(|| {
20829 std::env::var("MEMRA_FAW_P1")
20830 .map(|v| v != "0")
20831 .unwrap_or(true)
20832 });
20833 let hp =
20834 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20835 if hp {
20836 const BLOCK_QH: usize = 32;
20837 let f = self.func("fa_prefill_w_bf16_p1h2");
20838 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20839 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20840 f.set_attribute(
20841 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20842 shmem as i32,
20843 )?;
20844 let cfg = LaunchConfig {
20845 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20846 block_dim: (32, 4, 1),
20847 shared_mem_bytes: shmem,
20848 };
20849 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20850 head_dim as i32,
20851 n_head as i32,
20852 n_head_kv as i32,
20853 t as i32,
20854 t_kv as i32,
20855 causal as i32,
20856 window as i32,
20857 );
20858 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20859 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20860 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
20861 let __s_b = self.gpu.stream();
20862 let mut b = __s_b.launch_builder(&f);
20863 b.arg(&qb)
20864 .arg(&kb)
20865 .arg(&vh)
20866 .arg(o)
20867 .arg(&hd)
20868 .arg(&nh)
20869 .arg(&nhkv)
20870 .arg(&ti)
20871 .arg(&tkvi)
20872 .arg(&scale)
20873 .arg(&cz)
20874 .arg(&wi);
20875 unsafe {
20876 b.launch(cfg)?;
20877 }
20878 return Ok(());
20879 }
20880 if p1 {
20881 let f = self.func("fa_prefill_w_bf16_p1");
20882 let shmem =
20883 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20884 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20885 f.set_attribute(
20886 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20887 shmem as i32,
20888 )?;
20889 let cfg = LaunchConfig {
20890 grid_dim: (
20891 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20892 n_head as u32,
20893 1,
20894 ),
20895 block_dim: (32, 4, 1),
20896 shared_mem_bytes: shmem,
20897 };
20898 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20899 head_dim as i32,
20900 n_head as i32,
20901 n_head_kv as i32,
20902 t as i32,
20903 t_kv as i32,
20904 causal as i32,
20905 window as i32,
20906 );
20907 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20908 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20909 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20910 let __s_b = self.gpu.stream();
20911 let mut b = __s_b.launch_builder(&f);
20912 b.arg(&qb)
20913 .arg(&kb)
20914 .arg(&vb)
20915 .arg(o)
20916 .arg(&hd)
20917 .arg(&nh)
20918 .arg(&nhkv)
20919 .arg(&ti)
20920 .arg(&tkvi)
20921 .arg(&scale)
20922 .arg(&cz)
20923 .arg(&wi);
20924 unsafe {
20925 b.launch(cfg)?;
20926 }
20927 return Ok(());
20928 }
20929 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20932 let g4 = !floor
20933 && !f32_stage
20934 && n_head_kv == 1
20935 && n_head % 4 == 0
20936 && *G4_ON.get_or_init(|| {
20937 std::env::var("MEMRA_FAW_G4")
20938 .map(|v| v != "0")
20939 .unwrap_or(true)
20940 });
20941 if g4 {
20942 const SP_M: usize = 16;
20943 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20946 let o2 = *O2_ON.get_or_init(|| {
20947 std::env::var("MEMRA_FAW_O2")
20948 .map(|v| v != "0")
20949 .unwrap_or(true)
20950 });
20951 let f = self.func(if o2 {
20952 "fa_prefill_w_bf16_g4o2"
20953 } else {
20954 "fa_prefill_w_bf16_g4"
20955 });
20956 let shmem = if o2 {
20957 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
20958 } else {
20959 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
20960 as u32
20961 };
20962 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20963 f.set_attribute(
20964 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20965 shmem as i32,
20966 )?;
20967 let cfg = LaunchConfig {
20968 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
20969 block_dim: (32, 4, 1),
20970 shared_mem_bytes: shmem,
20971 };
20972 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20973 head_dim as i32,
20974 n_head as i32,
20975 n_head_kv as i32,
20976 t as i32,
20977 t_kv as i32,
20978 causal as i32,
20979 window as i32,
20980 );
20981 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20982 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20983 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20984 let __s_b = self.gpu.stream();
20985 let mut b = __s_b.launch_builder(&f);
20986 b.arg(&qb)
20987 .arg(&kb)
20988 .arg(&vb)
20989 .arg(o)
20990 .arg(&hd)
20991 .arg(&nh)
20992 .arg(&nhkv)
20993 .arg(&ti)
20994 .arg(&tkvi)
20995 .arg(&scale)
20996 .arg(&cz)
20997 .arg(&wi);
20998 unsafe {
20999 b.launch(cfg)?;
21000 }
21001 return Ok(());
21002 }
21003 let f = self.func(if floor {
21004 "fa_prefill_w_f32"
21005 } else if f32_stage {
21006 "fa_prefill_w_f32_pp"
21007 } else {
21008 "fa_prefill_w_bf16_pp"
21009 });
21010 let shmem =
21011 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21012 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21013 f.set_attribute(
21014 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21015 shmem as i32,
21016 )?;
21017 let cfg = LaunchConfig {
21018 grid_dim: (
21019 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21020 n_head as u32,
21021 1,
21022 ),
21023 block_dim: (32, 4, 1),
21024 shared_mem_bytes: shmem,
21025 };
21026 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
21027 head_dim as i32,
21028 n_head as i32,
21029 n_head_kv as i32,
21030 t as i32,
21031 t_kv as i32,
21032 causal as i32,
21033 window as i32,
21034 );
21035 if f32_stage {
21036 let __s_b = self.gpu.stream();
21037 let mut b = __s_b.launch_builder(&f);
21038 b.arg(q)
21039 .arg(k)
21040 .arg(v)
21041 .arg(o)
21042 .arg(&hd)
21043 .arg(&nh)
21044 .arg(&nhkv)
21045 .arg(&ti)
21046 .arg(&tkvi)
21047 .arg(&scale)
21048 .arg(&cz)
21049 .arg(&wi);
21050 unsafe {
21051 b.launch(cfg)?;
21052 }
21053 } else {
21054 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21055 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21056 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
21057 let __s_b = self.gpu.stream();
21058 let mut b = __s_b.launch_builder(&f);
21059 b.arg(&qb)
21060 .arg(&kb)
21061 .arg(&vb)
21062 .arg(o)
21063 .arg(&hd)
21064 .arg(&nh)
21065 .arg(&nhkv)
21066 .arg(&ti)
21067 .arg(&tkvi)
21068 .arg(&scale)
21069 .arg(&cz)
21070 .arg(&wi);
21071 unsafe {
21072 b.launch(cfg)?;
21073 }
21074 }
21075 Ok(())
21076 }
21077
21078 #[allow(clippy::too_many_arguments)]
21082 pub fn fa_prefill_hd512(
21083 &self,
21084 q: &CudaSlice<f32>,
21085 k: &CudaSlice<f32>,
21086 v: &CudaSlice<f32>,
21087 o: &mut CudaSlice<f32>,
21088 head_dim: usize,
21089 n_head: usize,
21090 n_head_kv: usize,
21091 t: usize,
21092 t_kv: usize,
21093 scale: f32,
21094 causal: bool,
21095 ) -> Result<(), Box<dyn std::error::Error>> {
21096 if portable_mma_gated() {
21098 return self.sdpa_naive(
21099 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
21100 );
21101 }
21102 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21108 let f32_stage =
21109 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
21110 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21114 let sp = !f32_stage
21115 && *SP_ON.get_or_init(|| {
21116 std::env::var("MEMRA_FA512_SP")
21117 .map(|v| v != "0")
21118 .unwrap_or(true)
21119 });
21120 self.fa_prefill_hd512_arm(
21121 q,
21122 k,
21123 v,
21124 o,
21125 head_dim,
21126 n_head,
21127 n_head_kv,
21128 t,
21129 t_kv,
21130 scale,
21131 causal,
21132 f32_stage,
21133 sp,
21134 sp && fa_f16pv_on(),
21135 )
21136 }
21137
21138 #[allow(clippy::too_many_arguments)]
21140 pub fn fa_prefill_hd512_pre(
21141 &self,
21142 qb: &CudaSlice<u8>,
21143 kb: &CudaSlice<u8>,
21144 vb: &CudaSlice<u8>,
21145 o: &mut CudaSlice<f32>,
21146 head_dim: usize,
21147 n_head: usize,
21148 n_head_kv: usize,
21149 t: usize,
21150 t_kv: usize,
21151 scale: f32,
21152 causal: bool,
21153 v_f16: bool,
21154 ) -> Result<(), Box<dyn std::error::Error>> {
21155 debug_assert_eq!(head_dim, 512);
21156 const SP_M: usize = 16;
21157 const BKS: usize = 32;
21158 let f16pv = fa_f16pv_on();
21162 let nw = if f16pv { fa512_wide_warps() } else { 2 };
21163 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
21164 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
21165 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
21166 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
21167 let n = t_kv * n_head_kv * head_dim;
21169 let need = n * 2;
21170 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
21171 *vguard = Some(self.alloc_uninit::<u8>(need)?);
21172 }
21173 let dst = vguard.as_mut().unwrap();
21174 self.bf16_to_f16_into(vb, n, dst)?;
21175 vguard.as_ref().unwrap()
21176 } else {
21177 vb
21178 };
21179 let f = self.func(if hp {
21180 "fa_prefill_bf16_hd512_sp16h2"
21181 } else {
21182 match (f16pv, nw) {
21183 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
21184 (true, _) => "fa_prefill_bf16_hd512_sp16",
21185 _ => "fa_prefill_bf16_hd512_sp",
21186 }
21187 });
21188 let (nwarp, npart) = if hp {
21189 (4usize, 4usize)
21190 } else if nw > 2 {
21191 (nw, nw)
21192 } else {
21193 (2, 1)
21194 };
21195 let shmem = if hp {
21197 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
21198 as u32
21199 } else {
21200 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
21201 + 4 * (npart * SP_M * BKS + SP_M)) as u32
21202 };
21203 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21204 f.set_attribute(
21205 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21206 shmem as i32,
21207 )?;
21208 let grid_y = if hp {
21209 (n_head / 2) as u32
21210 } else {
21211 n_head as u32
21212 };
21213 let cfg = LaunchConfig {
21214 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
21215 block_dim: (32, nwarp as u32, 1),
21216 shared_mem_bytes: shmem,
21217 };
21218 let (hd, nh, nhkv, ti, tkvi, cz) = (
21219 head_dim as i32,
21220 n_head as i32,
21221 n_head_kv as i32,
21222 t as i32,
21223 t_kv as i32,
21224 causal as i32,
21225 );
21226 let __s_b = self.gpu.stream();
21227 let mut b = __s_b.launch_builder(&f);
21228 b.arg(qb)
21229 .arg(kb)
21230 .arg(vref)
21231 .arg(o)
21232 .arg(&hd)
21233 .arg(&nh)
21234 .arg(&nhkv)
21235 .arg(&ti)
21236 .arg(&tkvi)
21237 .arg(&scale)
21238 .arg(&cz);
21239 unsafe {
21240 b.launch(cfg)?;
21241 }
21242 Ok(())
21243 }
21244
21245 #[allow(clippy::too_many_arguments)]
21248 pub fn fa_prefill_hd512_arm(
21249 &self,
21250 q: &CudaSlice<f32>,
21251 k: &CudaSlice<f32>,
21252 v: &CudaSlice<f32>,
21253 o: &mut CudaSlice<f32>,
21254 head_dim: usize,
21255 n_head: usize,
21256 n_head_kv: usize,
21257 t: usize,
21258 t_kv: usize,
21259 scale: f32,
21260 causal: bool,
21261 f32_stage: bool,
21262 sp: bool,
21263 f16pv: bool,
21264 ) -> Result<(), Box<dyn std::error::Error>> {
21265 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
21266 if sp && !f32_stage {
21267 const SP_M: usize = 16;
21271 const BKS: usize = 32;
21272 let nw = if f16pv { fa512_wide_warps() } else { 2 };
21273 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
21274 let f = self.func(if hp {
21275 "fa_prefill_bf16_hd512_sp16h2"
21276 } else {
21277 match (f16pv, nw) {
21278 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
21279 (true, _) => "fa_prefill_bf16_hd512_sp16",
21280 _ => "fa_prefill_bf16_hd512_sp",
21281 }
21282 });
21283 let (nwarp, npart) = if hp {
21284 (4usize, 4usize)
21285 } else if nw > 2 {
21286 (nw, nw)
21287 } else {
21288 (2, 1)
21289 };
21290 let shmem = if hp {
21291 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
21292 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
21293 } else {
21294 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
21295 + 4 * (npart * SP_M * BKS + SP_M)) as u32
21296 };
21297 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21298 f.set_attribute(
21299 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21300 shmem as i32,
21301 )?;
21302 let grid_y = if hp {
21303 (n_head / 2) as u32
21304 } else {
21305 n_head as u32
21306 };
21307 let cfg = LaunchConfig {
21308 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
21309 block_dim: (32, nwarp as u32, 1),
21310 shared_mem_bytes: shmem,
21311 };
21312 let (hd, nh, nhkv, ti, tkvi, cz) = (
21313 head_dim as i32,
21314 n_head as i32,
21315 n_head_kv as i32,
21316 t as i32,
21317 t_kv as i32,
21318 causal as i32,
21319 );
21320 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21321 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21322 let vb = if f16pv {
21323 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
21324 } else {
21325 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
21326 };
21327 let __s_b = self.gpu.stream();
21328 let mut b = __s_b.launch_builder(&f);
21329 b.arg(&qb)
21330 .arg(&kb)
21331 .arg(&vb)
21332 .arg(o)
21333 .arg(&hd)
21334 .arg(&nh)
21335 .arg(&nhkv)
21336 .arg(&ti)
21337 .arg(&tkvi)
21338 .arg(&scale)
21339 .arg(&cz);
21340 unsafe {
21341 b.launch(cfg)?;
21342 }
21343 return Ok(());
21344 }
21345 const BLOCK_Q: usize = 32;
21346 const BK: usize = 32;
21347 const HALF: usize = 256;
21348 let f = self.func(if f32_stage {
21349 "fa_prefill_f32_hd512"
21350 } else {
21351 "fa_prefill_bf16_hd512"
21352 });
21353 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
21355 + 4 * BLOCK_Q) as u32;
21356 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21357 f.set_attribute(
21358 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21359 shmem as i32,
21360 )?;
21361 let cfg = LaunchConfig {
21362 grid_dim: (
21363 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21364 n_head as u32,
21365 2,
21366 ),
21367 block_dim: (32, 2, 1),
21368 shared_mem_bytes: shmem,
21369 };
21370 let (hd, nh, nhkv, ti, tkvi, cz) = (
21371 head_dim as i32,
21372 n_head as i32,
21373 n_head_kv as i32,
21374 t as i32,
21375 t_kv as i32,
21376 causal as i32,
21377 );
21378 if f32_stage {
21379 let __s_b = self.gpu.stream();
21380 let mut b = __s_b.launch_builder(&f);
21381 b.arg(q)
21382 .arg(k)
21383 .arg(v)
21384 .arg(o)
21385 .arg(&hd)
21386 .arg(&nh)
21387 .arg(&nhkv)
21388 .arg(&ti)
21389 .arg(&tkvi)
21390 .arg(&scale)
21391 .arg(&cz);
21392 unsafe {
21393 b.launch(cfg)?;
21394 }
21395 } else {
21396 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21397 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21398 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
21399 let __s_b = self.gpu.stream();
21400 let mut b = __s_b.launch_builder(&f);
21401 b.arg(&qb)
21402 .arg(&kb)
21403 .arg(&vb)
21404 .arg(o)
21405 .arg(&hd)
21406 .arg(&nh)
21407 .arg(&nhkv)
21408 .arg(&ti)
21409 .arg(&tkvi)
21410 .arg(&scale)
21411 .arg(&cz);
21412 unsafe {
21413 b.launch(cfg)?;
21414 }
21415 }
21416 Ok(())
21417 }
21418
21419 #[allow(clippy::too_many_arguments)]
21423 pub fn rope_neox2_bf16e(
21424 &self,
21425 q: &mut CudaSlice<f32>,
21426 k: &mut CudaSlice<f32>,
21427 qb: &mut CudaSlice<u8>,
21428 kb: &mut CudaSlice<u8>,
21429 pos: &CudaSlice<i32>,
21430 head_dim: usize,
21431 n_dims: usize,
21432 nh_q: usize,
21433 nh_k: usize,
21434 n_tokens: usize,
21435 base: f32,
21436 freq_scale: f32,
21437 ff: Option<&CudaSlice<f32>>,
21438 ) -> Result<(), Box<dyn std::error::Error>> {
21439 let f = self.func("rope_neox2_bf16e_f32");
21440 let rows = ((nh_q + nh_k) * n_tokens) as u32;
21441 let cfg = LaunchConfig {
21442 grid_dim: (rows, 1, 1),
21443 block_dim: ((head_dim / 2) as u32, 1, 1),
21444 shared_mem_bytes: 0,
21445 };
21446 let theta_scale = base.powf(-2.0 / n_dims as f32);
21447 let (hd, nd, nhq, nhk, nt) = (
21448 head_dim as i32,
21449 n_dims as i32,
21450 nh_q as i32,
21451 nh_k as i32,
21452 n_tokens as i32,
21453 );
21454 let __s_b = self.gpu.stream();
21455 let mut b = __s_b.launch_builder(&f);
21456 match ff {
21457 Some(t) => {
21458 b.arg(&mut *q)
21459 .arg(&mut *k)
21460 .arg(&mut *qb)
21461 .arg(&mut *kb)
21462 .arg(pos)
21463 .arg(&hd)
21464 .arg(&nd)
21465 .arg(&nhq)
21466 .arg(&nhk)
21467 .arg(&nt)
21468 .arg(&theta_scale)
21469 .arg(&freq_scale)
21470 .arg(t);
21471 unsafe {
21472 b.launch(cfg)?;
21473 }
21474 }
21475 None => {
21476 let null: u64 = 0;
21477 b.arg(&mut *q)
21478 .arg(&mut *k)
21479 .arg(&mut *qb)
21480 .arg(&mut *kb)
21481 .arg(pos)
21482 .arg(&hd)
21483 .arg(&nd)
21484 .arg(&nhq)
21485 .arg(&nhk)
21486 .arg(&nt)
21487 .arg(&theta_scale)
21488 .arg(&freq_scale)
21489 .arg(&null);
21490 unsafe {
21491 b.launch(cfg)?;
21492 }
21493 }
21494 }
21495 Ok(())
21496 }
21497
21498 pub fn f32_to_bf16(
21501 &self,
21502 x: &CudaSlice<f32>,
21503 n: usize,
21504 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21505 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
21506 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21507 let f = self.func("f32_to_bf16_flat");
21508 let n_i = n as i64;
21509 let cfg = LaunchConfig {
21510 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
21511 block_dim: (256, 1, 1),
21512 shared_mem_bytes: 0,
21513 };
21514 let __s_b = self.gpu.stream();
21515 let mut b = __s_b.launch_builder(&f);
21516 b.arg(x).arg(&mut y).arg(&n_i);
21517 unsafe {
21518 b.launch(cfg)?;
21519 }
21520 Ok(y)
21521 }
21522
21523 pub fn f32_to_f16(
21524 &self,
21525 x: &CudaSlice<f32>,
21526 n: usize,
21527 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21528 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
21529 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21530 let f = self.func("f32_to_f16_flat");
21531 let n_i = n as i64;
21532 let cfg = LaunchConfig {
21533 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
21534 block_dim: (256, 1, 1),
21535 shared_mem_bytes: 0,
21536 };
21537 let __s_b = self.gpu.stream();
21538 let mut b = __s_b.launch_builder(&f);
21539 b.arg(x).arg(&mut y).arg(&n_i);
21540 unsafe {
21541 b.launch(cfg)?;
21542 }
21543 Ok(y)
21544 }
21545
21546 pub fn bf16_to_f16(
21548 &self,
21549 xb: &CudaSlice<u8>,
21550 n: usize,
21551 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21552 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21553 self.bf16_to_f16_into(xb, n, &mut y)?;
21554 Ok(y)
21555 }
21556
21557 pub fn bf16_to_f16_into(
21559 &self,
21560 xb: &CudaSlice<u8>,
21561 n: usize,
21562 y: &mut CudaSlice<u8>,
21563 ) -> Result<(), Box<dyn std::error::Error>> {
21564 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
21565 assert!(y.len() >= n * 2);
21566 let f = self.func("bf16_to_f16_flat");
21567 let n2 = (n / 2) as i64;
21568 let cfg = LaunchConfig {
21569 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
21570 block_dim: (256, 1, 1),
21571 shared_mem_bytes: 0,
21572 };
21573 let __s_b = self.gpu.stream();
21574 let mut b = __s_b.launch_builder(&f);
21575 b.arg(xb).arg(y).arg(&n2);
21576 unsafe {
21577 b.launch(cfg)?;
21578 }
21579 Ok(())
21580 }
21581
21582 #[allow(clippy::too_many_arguments)]
21587 pub fn fa_prefill_vl8(
21588 &self,
21589 seqs: &[FaSeqVl],
21590 head_dim: usize,
21591 n_head: usize,
21592 n_head_kv: usize,
21593 scale: f32,
21594 ) -> Result<(), Box<dyn std::error::Error>> {
21595 const BK: usize = 32;
21596 let b = seqs.len();
21597 assert!(b >= 1 && b <= 8);
21598 let mut packed = [FaSeqVl::default(); 8];
21599 packed[..b].copy_from_slice(seqs);
21600 let v = FaVl8(packed);
21601 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21602 let ept = (n_head_kv * head_dim) as i32;
21603 {
21604 let f = self.func("fa_mirror_vl");
21605 let max_n = (max_t as i64) * ept as i64;
21606 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21607 for which in 0..2i32 {
21608 let cfg = LaunchConfig {
21609 grid_dim: (blocks, 1, b as u32),
21610 block_dim: (256, 1, 1),
21611 shared_mem_bytes: 0,
21612 };
21613 let __s_lb = self.gpu.stream();
21614 let mut lb = __s_lb.launch_builder(&f);
21615 lb.arg(&v).arg(&ept).arg(&which);
21616 unsafe {
21617 lb.launch(cfg)?;
21618 }
21619 }
21620 }
21621 let hd_sfx = fa_hd_suffix(head_dim)?;
21622 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
21623 let block_q = 64usize;
21624 let kv_stages = 2usize;
21625 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
21626 + 4 * (block_q * BK + 2 * block_q)) as u32;
21627 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21628 f.set_attribute(
21629 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21630 shmem as i32,
21631 )?;
21632 let cfg = LaunchConfig {
21633 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
21634 block_dim: (32, 4, 1),
21635 shared_mem_bytes: shmem,
21636 };
21637 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21638 let __s_lb = self.gpu.stream();
21639 let mut lb = __s_lb.launch_builder(&f);
21640 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
21641 unsafe {
21642 lb.launch(cfg)?;
21643 }
21644 Ok(())
21645 }
21646
21647 #[allow(clippy::too_many_arguments)]
21651 pub fn attn_pre_vl8(
21652 &self,
21653 seqs: &[AttnPreVl],
21654 wq: &CudaSlice<f32>,
21655 wk: &CudaSlice<f32>,
21656 head_dim: usize,
21657 rope_dims: usize,
21658 n_head: usize,
21659 n_head_kv: usize,
21660 eps: f32,
21661 freq_base: f32,
21662 freq_scale: f32,
21663 kv_dim_k: usize,
21664 kv_dim_v: usize,
21665 k_tok_bytes: usize,
21666 v_tok_bytes: usize,
21667 ) -> Result<(), Box<dyn std::error::Error>> {
21668 let b = seqs.len();
21669 assert!(b >= 1 && b <= 8);
21670 let mut packed = [AttnPreVl::default(); 8];
21671 packed[..b].copy_from_slice(seqs);
21672 let v = AttnPreVl8(packed);
21673 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21674 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21675 {
21676 let f = self.func("q_gate_split_vl");
21677 let n = max_t * (n_head * head_dim) as u32;
21678 let cfg = LaunchConfig {
21679 grid_dim: (n.div_ceil(256), 1, b as u32),
21680 block_dim: (256, 1, 1),
21681 shared_mem_bytes: 0,
21682 };
21683 let __s_lb = self.gpu.stream();
21684 let mut lb = __s_lb.launch_builder(&f);
21685 lb.arg(&v).arg(&hd).arg(&nh);
21686 unsafe {
21687 lb.launch(cfg)?;
21688 }
21689 }
21690 {
21691 let f = self.func("attn_rms_vl");
21692 let cfg = LaunchConfig {
21693 grid_dim: (max_t * n_head as u32, 2, b as u32),
21694 block_dim: (rms_block(), 1, 1),
21695 shared_mem_bytes: 0,
21696 };
21697 let __s_lb = self.gpu.stream();
21698 let mut lb = __s_lb.launch_builder(&f);
21699 lb.arg(&v)
21700 .arg(wq)
21701 .arg(wk)
21702 .arg(&hd)
21703 .arg(&nh)
21704 .arg(&nhkv)
21705 .arg(&eps);
21706 unsafe {
21707 lb.launch(cfg)?;
21708 }
21709 }
21710 {
21711 let f = self.func("attn_rope_vl");
21712 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
21713 let nd = rope_dims as i32;
21714 let cfg = LaunchConfig {
21715 grid_dim: (max_t * n_head as u32, 2, b as u32),
21716 block_dim: ((head_dim / 2) as u32, 1, 1),
21717 shared_mem_bytes: 0,
21718 };
21719 let __s_lb = self.gpu.stream();
21720 let mut lb = __s_lb.launch_builder(&f);
21721 lb.arg(&v)
21722 .arg(&hd)
21723 .arg(&nd)
21724 .arg(&nh)
21725 .arg(&nhkv)
21726 .arg(&theta_scale)
21727 .arg(&freq_scale);
21728 unsafe {
21729 lb.launch(cfg)?;
21730 }
21731 }
21732 {
21733 let f = self.func("append_kv_vl");
21734 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21735 let cfg = LaunchConfig {
21736 grid_dim: (nblk, max_t, b as u32),
21737 block_dim: (32, 1, 1),
21738 shared_mem_bytes: 0,
21739 };
21740 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21741 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21742 let __s_lb = self.gpu.stream();
21743 let mut lb = __s_lb.launch_builder(&f);
21744 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
21745 unsafe {
21746 lb.launch(cfg)?;
21747 }
21748 }
21749 Ok(())
21750 }
21751
21752 pub fn fa_prefill_view(
21757 &self,
21758 q: &CudaSlice<f32>,
21759 k: &cudarc::driver::CudaView<u8>,
21760 v: &cudarc::driver::CudaView<u8>,
21761 o: &mut CudaSlice<f32>,
21762 head_dim: usize,
21763 n_head: usize,
21764 n_head_kv: usize,
21765 t: usize,
21766 t_kv: usize,
21767 scale: f32,
21768 causal: bool,
21769 k_tok_bytes: usize,
21770 v_tok_bytes: usize,
21771 g: bool,
21772 ) -> Result<(), Box<dyn std::error::Error>> {
21773 if portable_mma_gated() {
21774 return self.sdpa_naive_quantized_view(
21775 q,
21776 k,
21777 v,
21778 o,
21779 head_dim,
21780 n_head,
21781 n_head_kv,
21782 t,
21783 t_kv,
21784 scale,
21785 causal,
21786 k_tok_bytes,
21787 v_tok_bytes,
21788 );
21789 }
21790 const BLOCK_Q: usize = 64;
21791 const BK: usize = 32;
21792 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
21795 let f = if g {
21796 self.func_g(&name)
21797 } else {
21798 self.func(&name)
21799 };
21800 let shmem =
21801 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21802 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21803 f.set_attribute(
21804 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21805 shmem as i32,
21806 )?;
21807 let cfg = LaunchConfig {
21808 grid_dim: (
21809 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21810 n_head as u32,
21811 1,
21812 ),
21813 block_dim: (32, 4, 1),
21814 shared_mem_bytes: shmem,
21815 };
21816 let (hd, nh, nhkv, ti, tkvi, cz) = (
21817 head_dim as i32,
21818 n_head as i32,
21819 n_head_kv as i32,
21820 t as i32,
21821 t_kv as i32,
21822 causal as i32,
21823 );
21824 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21825 let __s_b = self.gpu.stream();
21826 let mut b = __s_b.launch_builder(&f);
21827 b.arg(q)
21828 .arg(k)
21829 .arg(v)
21830 .arg(o)
21831 .arg(&hd)
21832 .arg(&nh)
21833 .arg(&nhkv)
21834 .arg(&ti)
21835 .arg(&tkvi)
21836 .arg(&scale)
21837 .arg(&cz)
21838 .arg(&ktb)
21839 .arg(&vtb);
21840 unsafe {
21841 b.launch(cfg)?;
21842 }
21843 Ok(())
21844 }
21845
21846 #[allow(clippy::too_many_arguments)]
21856 pub fn fa_prefill_view_ws(
21857 &self,
21858 q: &CudaSlice<f32>,
21859 k: &cudarc::driver::CudaView<u8>,
21860 v: &cudarc::driver::CudaView<u8>,
21861 o: &mut CudaSlice<f32>,
21862 head_dim: usize,
21863 n_head: usize,
21864 n_head_kv: usize,
21865 t: usize,
21866 t_kv: usize,
21867 scale: f32,
21868 causal: bool,
21869 k_tok_bytes: usize,
21870 v_tok_bytes: usize,
21871 g: bool,
21872 ) -> Result<(), Box<dyn std::error::Error>> {
21873 if portable_mma_gated() {
21874 return self.sdpa_naive_quantized_view(
21875 q,
21876 k,
21877 v,
21878 o,
21879 head_dim,
21880 n_head,
21881 n_head_kv,
21882 t,
21883 t_kv,
21884 scale,
21885 causal,
21886 k_tok_bytes,
21887 v_tok_bytes,
21888 );
21889 }
21890 const BLOCK_Q: usize = 64;
21891 const BK: usize = 32;
21892 let kv_dim_k = n_head_kv * head_dim;
21893 let kv_dim_v = n_head_kv * head_dim;
21894 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
21896 let mut guard = self.prime_deqw_ws.lock().unwrap();
21898 let need_grow = match guard.as_ref() {
21899 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21900 None => true,
21901 };
21902 if need_grow {
21903 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21904 let (ck, cv) = guard
21905 .as_ref()
21906 .map(|(a, b)| (a.len(), b.len()))
21907 .unwrap_or((0, 0));
21908 *guard = Some((
21909 self.alloc_u8(grow(ck, k_ws_bytes))?,
21910 self.alloc_u8(grow(cv, v_ws_bytes))?,
21911 ));
21912 }
21913 let (kw, vw) = guard.as_mut().unwrap();
21914 {
21916 let f = if g {
21918 self.func_g("fa_dequant_kv_ws_bf16")
21919 } else {
21920 self.func("fa_dequant_kv_ws_bf16")
21921 };
21922 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21923 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21924 let cfg = LaunchConfig {
21925 grid_dim: (nblk.max(1), 1, 1),
21926 block_dim: (256, 1, 1),
21927 shared_mem_bytes: 0,
21928 };
21929 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21930 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21931 let __s_b = self.gpu.stream();
21932 let mut b = __s_b.launch_builder(&f);
21933 b.arg(k)
21934 .arg(v)
21935 .arg(&mut *kw)
21936 .arg(&mut *vw)
21937 .arg(&kdk)
21938 .arg(&kdv)
21939 .arg(&tkvi)
21940 .arg(&ktb)
21941 .arg(&vtb);
21942 unsafe {
21943 b.launch(cfg)?;
21944 }
21945 }
21946 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21954 .map(|v| v != "0")
21955 .unwrap_or(true);
21956 {
21957 let hd_sfx = fa_hd_suffix(head_dim)?;
21958 let f = self.func(&format!(
21959 "fa_prefill_qw{}{hd_sfx}",
21960 if db { "_db" } else { "" }
21961 ));
21962 let shmem = if db {
21963 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21965 } else {
21966 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21967 };
21968 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21969 f.set_attribute(
21970 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21971 shmem as i32,
21972 )?;
21973 let cfg = LaunchConfig {
21974 grid_dim: (
21975 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21976 n_head as u32,
21977 1,
21978 ),
21979 block_dim: (32, 4, 1),
21980 shared_mem_bytes: shmem,
21981 };
21982 let (hd, nh, nhkv, ti, tkvi, cz) = (
21983 head_dim as i32,
21984 n_head as i32,
21985 n_head_kv as i32,
21986 t as i32,
21987 t_kv as i32,
21988 causal as i32,
21989 );
21990 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21991 let __s_b = self.gpu.stream();
21992 let mut b = __s_b.launch_builder(&f);
21993 b.arg(q)
21994 .arg(&*kw)
21995 .arg(&*vw)
21996 .arg(o)
21997 .arg(&hd)
21998 .arg(&nh)
21999 .arg(&nhkv)
22000 .arg(&ti)
22001 .arg(&tkvi)
22002 .arg(&scale)
22003 .arg(&cz)
22004 .arg(&kdk)
22005 .arg(&kdv);
22006 unsafe {
22007 b.launch(cfg)?;
22008 }
22009 }
22010 Ok(())
22011 }
22012
22013 #[allow(clippy::too_many_arguments)]
22029 pub fn fa_prefill_view_ws_w_hd128(
22030 &self,
22031 q: &CudaSlice<f32>,
22032 k: &cudarc::driver::CudaView<u8>,
22033 v: &cudarc::driver::CudaView<u8>,
22034 o: &mut CudaSlice<f32>,
22035 head_dim: usize,
22036 n_head: usize,
22037 n_head_kv: usize,
22038 t: usize,
22039 t_kv: usize,
22040 scale: f32,
22041 causal: bool,
22042 window: usize,
22043 k_tok_bytes: usize,
22044 v_tok_bytes: usize,
22045 ) -> Result<(), Box<dyn std::error::Error>> {
22046 assert_eq!(
22047 head_dim, 128,
22048 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
22049 );
22050 if portable_mma_gated() {
22051 return self.sdpa_naive_w_quantized_view(
22052 q,
22053 k,
22054 v,
22055 o,
22056 head_dim,
22057 n_head,
22058 n_head_kv,
22059 t,
22060 t_kv,
22061 scale,
22062 causal,
22063 window,
22064 k_tok_bytes,
22065 v_tok_bytes,
22066 );
22067 }
22068 const BLOCK_Q: usize = 64;
22069 const BK: usize = 32;
22070 let kv_dim_k = n_head_kv * head_dim;
22071 let kv_dim_v = n_head_kv * head_dim;
22072 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
22074 let mut guard = self.prime_deqw_ws.lock().unwrap();
22075 let need_grow = match guard.as_ref() {
22076 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
22077 None => true,
22078 };
22079 if need_grow {
22080 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
22081 let (ck, cv) = guard
22082 .as_ref()
22083 .map(|(a, b)| (a.len(), b.len()))
22084 .unwrap_or((0, 0));
22085 *guard = Some((
22086 self.alloc_u8(grow(ck, k_ws_bytes))?,
22087 self.alloc_u8(grow(cv, v_ws_bytes))?,
22088 ));
22089 }
22090 let (kw, vw) = guard.as_mut().unwrap();
22091 {
22094 let f = self.func("fa_dequant_kv_ws_bf16");
22095 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
22096 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
22097 let cfg = LaunchConfig {
22098 grid_dim: (nblk.max(1), 1, 1),
22099 block_dim: (256, 1, 1),
22100 shared_mem_bytes: 0,
22101 };
22102 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
22103 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22104 let __s_b = self.gpu.stream();
22105 let mut b = __s_b.launch_builder(&f);
22106 b.arg(k)
22107 .arg(v)
22108 .arg(&mut *kw)
22109 .arg(&mut *vw)
22110 .arg(&kdk)
22111 .arg(&kdv)
22112 .arg(&tkvi)
22113 .arg(&ktb)
22114 .arg(&vtb);
22115 unsafe {
22116 b.launch(cfg)?;
22117 }
22118 }
22119 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
22121 .map(|v| v != "0")
22122 .unwrap_or(true);
22123 {
22124 let f = self.func(if db {
22125 "fa_prefill_qw_db_w_hd128"
22126 } else {
22127 "fa_prefill_qw_w_hd128"
22128 });
22129 let shmem = if db {
22130 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
22131 } else {
22132 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
22133 };
22134 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22135 f.set_attribute(
22136 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22137 shmem as i32,
22138 )?;
22139 let cfg = LaunchConfig {
22140 grid_dim: (
22141 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
22142 n_head as u32,
22143 1,
22144 ),
22145 block_dim: (32, 4, 1),
22146 shared_mem_bytes: shmem,
22147 };
22148 let (hd, nh, nhkv, ti, tkvi, cz) = (
22149 head_dim as i32,
22150 n_head as i32,
22151 n_head_kv as i32,
22152 t as i32,
22153 t_kv as i32,
22154 causal as i32,
22155 );
22156 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
22157 let __s_b = self.gpu.stream();
22158 let mut b = __s_b.launch_builder(&f);
22159 b.arg(q)
22160 .arg(&*kw)
22161 .arg(&*vw)
22162 .arg(o)
22163 .arg(&hd)
22164 .arg(&nh)
22165 .arg(&nhkv)
22166 .arg(&ti)
22167 .arg(&tkvi)
22168 .arg(&scale)
22169 .arg(&cz)
22170 .arg(&kdk)
22171 .arg(&kdv)
22172 .arg(&wnd);
22173 unsafe {
22174 b.launch(cfg)?;
22175 }
22176 }
22177 Ok(())
22178 }
22179
22180 pub fn fa_decode(
22184 &self,
22185 q: &CudaSlice<f32>,
22186 k: &cudarc::driver::CudaView<u8>,
22187 v: &cudarc::driver::CudaView<u8>,
22188 o: &mut CudaSlice<f32>,
22189 head_dim: usize,
22190 n_head: usize,
22191 n_head_kv: usize,
22192 t_kv: usize,
22193 scale: f32,
22194 k_tok_bytes: usize,
22195 v_tok_bytes: usize,
22196 ) -> Result<(), Box<dyn std::error::Error>> {
22197 self.fa_decode_kvmod(
22198 q,
22199 k,
22200 v,
22201 o,
22202 head_dim,
22203 n_head,
22204 n_head_kv,
22205 t_kv,
22206 scale,
22207 k_tok_bytes,
22208 v_tok_bytes,
22209 false,
22210 )
22211 }
22212
22213 #[allow(clippy::too_many_arguments)]
22217 #[allow(clippy::too_many_arguments)]
22221 #[allow(clippy::too_many_arguments)]
22222 fn fa_decode_scalar_unified(
22223 &self,
22224 q: &cudarc::driver::CudaView<f32>,
22225 k: &cudarc::driver::CudaView<u8>,
22226 v: &cudarc::driver::CudaView<u8>,
22227 o: &mut cudarc::driver::CudaViewMut<f32>,
22228 head_dim: usize,
22229 n_head: usize,
22230 n_head_kv: usize,
22231 t_kv_host: usize,
22232 t_kv_dev: Option<&CudaSlice<i32>>,
22233 scale: f32,
22234 n_splits: usize,
22235 split_keys: usize,
22236 k_tok_bytes: usize,
22237 v_tok_bytes: usize,
22238 g: bool,
22239 part_o: &mut CudaSlice<f32>,
22240 part_m: &mut CudaSlice<f32>,
22241 part_l: &mut CudaSlice<f32>,
22242 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22243 ) -> Result<(), Box<dyn std::error::Error>> {
22244 let f = if g {
22245 self.func_g("fa_decode_f32")
22246 } else {
22247 self.fa_func("fa_decode_f32", head_dim)
22248 };
22249 let cfg = LaunchConfig {
22250 grid_dim: (n_head as u32, n_splits as u32, 1),
22251 block_dim: (head_dim as u32, 1, 1),
22252 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
22253 };
22254 let (hd, nh, nhkv, nsp) = (
22255 head_dim as i32,
22256 n_head as i32,
22257 n_head_kv as i32,
22258 n_splits as i32,
22259 );
22260 let (ktb, vtb, tkvi, ski) = (
22261 k_tok_bytes as i64,
22262 v_tok_bytes as i64,
22263 t_kv_host as i32,
22264 split_keys as i32,
22265 );
22266 let __s_b = self.gpu.stream();
22267 let mut b = __s_b.launch_builder(&f);
22268 match t_kv_dev {
22269 Some(d) => {
22270 b.arg(q)
22271 .arg(k)
22272 .arg(v)
22273 .arg(&mut *part_o)
22274 .arg(&mut *part_m)
22275 .arg(&mut *part_l)
22276 .arg(&hd)
22277 .arg(&nh)
22278 .arg(&nhkv)
22279 .arg(&tkvi)
22280 .arg(d)
22281 .arg(&scale)
22282 .arg(&nsp)
22283 .arg(&ski)
22284 .arg(&ktb)
22285 .arg(&vtb);
22286 unsafe {
22287 b.launch(cfg)?;
22288 }
22289 }
22290 None => {
22291 let null: u64 = 0;
22292 b.arg(q)
22293 .arg(k)
22294 .arg(v)
22295 .arg(&mut *part_o)
22296 .arg(&mut *part_m)
22297 .arg(&mut *part_l)
22298 .arg(&hd)
22299 .arg(&nh)
22300 .arg(&nhkv)
22301 .arg(&tkvi)
22302 .arg(&null)
22303 .arg(&scale)
22304 .arg(&nsp)
22305 .arg(&ski)
22306 .arg(&ktb)
22307 .arg(&vtb);
22308 unsafe {
22309 b.launch(cfg)?;
22310 }
22311 }
22312 }
22313 let cfg2 = LaunchConfig {
22314 grid_dim: (n_head as u32, 1, 1),
22315 block_dim: (head_dim as u32, 1, 1),
22316 shared_mem_bytes: 0,
22317 };
22318 if let Some((oq, od)) = q8_out {
22319 let fc = if g {
22321 self.func_g("fa_decode_combine_q8_1")
22322 } else {
22323 self.fa_func("fa_decode_combine_q8_1", head_dim)
22324 };
22325 let __s_b2 = self.gpu.stream();
22326 let mut b2 = __s_b2.launch_builder(&fc);
22327 b2.arg(&*part_o)
22328 .arg(&*part_m)
22329 .arg(&*part_l)
22330 .arg(oq)
22331 .arg(od)
22332 .arg(&hd)
22333 .arg(&nh)
22334 .arg(&nsp);
22335 unsafe {
22336 b2.launch(cfg2)?;
22337 }
22338 return Ok(());
22339 }
22340 let fc = if g {
22341 self.func_g("fa_decode_combine_f32")
22342 } else {
22343 self.fa_func("fa_decode_combine_f32", head_dim)
22344 };
22345 let __s_b2 = self.gpu.stream();
22346 let mut b2 = __s_b2.launch_builder(&fc);
22347 b2.arg(&*part_o)
22348 .arg(&*part_m)
22349 .arg(&*part_l)
22350 .arg(o)
22351 .arg(&hd)
22352 .arg(&nh)
22353 .arg(&nsp);
22354 unsafe {
22355 b2.launch(cfg2)?;
22356 }
22357 Ok(())
22358 }
22359
22360 pub fn fa_decode_kvmod(
22361 &self,
22362 q: &CudaSlice<f32>,
22363 k: &cudarc::driver::CudaView<u8>,
22364 v: &cudarc::driver::CudaView<u8>,
22365 o: &mut CudaSlice<f32>,
22366 head_dim: usize,
22367 n_head: usize,
22368 n_head_kv: usize,
22369 t_kv: usize,
22370 scale: f32,
22371 k_tok_bytes: usize,
22372 v_tok_bytes: usize,
22373 g: bool,
22374 ) -> Result<(), Box<dyn std::error::Error>> {
22375 let q_view = q.as_view();
22376 let mut o_view = o.as_view_mut();
22377 self.fa_decode_kvmod_view(
22378 &q_view,
22379 k,
22380 v,
22381 &mut o_view,
22382 head_dim,
22383 n_head,
22384 n_head_kv,
22385 t_kv,
22386 scale,
22387 k_tok_bytes,
22388 v_tok_bytes,
22389 g,
22390 )
22391 }
22392
22393 #[allow(clippy::too_many_arguments)]
22398 pub fn fa_decode_kvmod_view(
22399 &self,
22400 q: &cudarc::driver::CudaView<f32>,
22401 k: &cudarc::driver::CudaView<u8>,
22402 v: &cudarc::driver::CudaView<u8>,
22403 o: &mut cudarc::driver::CudaViewMut<f32>,
22404 head_dim: usize,
22405 n_head: usize,
22406 n_head_kv: usize,
22407 t_kv: usize,
22408 scale: f32,
22409 k_tok_bytes: usize,
22410 v_tok_bytes: usize,
22411 g: bool,
22412 ) -> Result<(), Box<dyn std::error::Error>> {
22413 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
22434 if g && head_dim == 256 && !fa_v4_at(t_kv) {
22438 fa_vec = false;
22439 }
22440 let sp = fa_split_keys(t_kv, n_head_kv);
22441 let n_splits = if fa_vec {
22442 ((t_kv + sp - 1) / sp).max(1)
22443 } else {
22444 ((t_kv + 255) / 256).max(1)
22445 };
22446 let o_len = n_head * n_splits * head_dim;
22447 let ml_len = n_head * n_splits;
22448 let mut part_guard = self.fa_part_pool.lock().unwrap();
22449 if part_guard
22450 .as_ref()
22451 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22452 .unwrap_or(true)
22453 {
22454 let old = part_guard.take();
22465 let (co, cm) = old
22466 .as_ref()
22467 .map(|pp| (pp.0.len(), pp.1.len()))
22468 .unwrap_or((0, 0));
22469 if let Some(old) = old {
22470 self.fa_part_retired.lock().unwrap().push(old);
22471 }
22472 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22473 eprintln!(
22474 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22475 co, o_len, cm, ml_len
22476 );
22477 }
22478 *part_guard =
22479 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
22480 }
22481 let pg = part_guard.as_mut().unwrap();
22482 self.gpu
22483 .stream()
22484 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22485 self.gpu
22486 .stream()
22487 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22488 self.gpu
22489 .stream()
22490 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22491 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22492 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22493 let (hd, nh, nhkv, tkvi, nsp) = (
22494 head_dim as i32,
22495 n_head as i32,
22496 n_head_kv as i32,
22497 t_kv as i32,
22498 n_splits as i32,
22499 );
22500 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22501 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22505 let fa512_min = fa512_min_tkv();
22510 let deep = fa_vec
22513 && head_dim == 256
22514 && fa_v4_at(t_kv)
22515 && !g
22516 && fa_deep_at(t_kv)
22517 && !matches!(fa_v4_mode(), "noB3" | "stage");
22518 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
22519 let gqa = (n_head / n_head_kv).max(1) as u32;
22522 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
22523 (
22524 fv,
22525 LaunchConfig {
22526 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22527 block_dim: (32, gqa, 1),
22528 shared_mem_bytes: 0,
22529 },
22530 )
22531 } else if fa_vec && head_dim <= 256 {
22532 let gqa = (n_head / n_head_kv).max(1) as u32;
22533 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22544 let smem_tkv = *SMEM_TKV.get_or_init(|| {
22545 std::env::var("MEMRA_FA_SMEM_TKV")
22546 .ok()
22547 .and_then(|v| v.parse().ok())
22548 .unwrap_or_else(|| {
22549 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22550 })
22551 });
22552 if fa_v4_at(t_kv) && head_dim == 256 {
22553 let v4name = match fa_v4_mode() {
22557 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
22560 _ => "fa_decode_vec_q_v4",
22561 };
22562 let fv = if g {
22563 self.func_g(v4name)
22564 } else {
22565 self.func(v4name)
22566 };
22567 let shmem = (if deep { 12160 } else { 11520 }
22570 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22571 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22572 fv.set_attribute(
22573 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22574 shmem as i32,
22575 )?;
22576 (
22577 fv,
22578 LaunchConfig {
22579 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22580 block_dim: (32, gqa, 1),
22581 shared_mem_bytes: shmem,
22582 },
22583 )
22584 } else if fa_v3_active(head_dim) {
22585 let fv = if g {
22588 self.func_g("fa_decode_vec_q_v3")
22589 } else {
22590 self.func("fa_decode_vec_q_v3")
22591 };
22592 let shmem = (32 * head_dim * 2) as u32; (
22594 fv,
22595 LaunchConfig {
22596 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22597 block_dim: (32, gqa, 1),
22598 shared_mem_bytes: shmem,
22599 },
22600 )
22601 } else if fa_v2_on() {
22602 let fv = if g {
22606 self.func_g("fa_decode_vec_q_v2")
22607 } else {
22608 self.func("fa_decode_vec_q_v2")
22609 };
22610 let shmem = (2 * 32 * head_dim * 2) as u32; (
22612 fv,
22613 LaunchConfig {
22614 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22615 block_dim: (32, gqa, 1),
22616 shared_mem_bytes: shmem,
22617 },
22618 )
22619 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
22620 {
22621 let fv = if g {
22625 self.func_g("fa_decode_vec_q_smem")
22626 } else {
22627 self.func("fa_decode_vec_q_smem")
22628 };
22629 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
22631 fv.set_attribute(
22632 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22633 shmem as i32,
22634 )?;
22635 (
22636 fv,
22637 LaunchConfig {
22638 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22639 block_dim: (32, gqa, 1),
22640 shared_mem_bytes: shmem,
22641 },
22642 )
22643 } else {
22644 let fv = if g {
22647 self.func_g("fa_decode_vec_q")
22648 } else {
22649 self.func("fa_decode_vec_q")
22650 };
22651 (
22652 fv,
22653 LaunchConfig {
22654 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22655 block_dim: (32, gqa, 1),
22656 shared_mem_bytes: 0,
22657 },
22658 )
22659 }
22660 } else {
22661 return self.fa_decode_scalar_unified(
22664 q,
22665 k,
22666 v,
22667 o,
22668 head_dim,
22669 n_head,
22670 n_head_kv,
22671 t_kv,
22672 None,
22673 scale,
22674 n_splits,
22675 if fa_vec { sp } else { 256 },
22676 k_tok_bytes,
22677 v_tok_bytes,
22678 g,
22679 part_o,
22680 part_m,
22681 part_l,
22682 None,
22683 );
22684 };
22685 let __s_b = self.gpu.stream();
22686 let mut b = __s_b.launch_builder(&f);
22687 b.arg(q)
22688 .arg(k)
22689 .arg(v)
22690 .arg(&mut *part_o)
22691 .arg(&mut *part_m)
22692 .arg(&mut *part_l)
22693 .arg(&hd)
22694 .arg(&nh)
22695 .arg(&nhkv)
22696 .arg(&tkvi)
22697 .arg(&scale)
22698 .arg(&nsp)
22699 .arg(&ktb)
22700 .arg(&vtb);
22701 unsafe {
22702 b.launch(cfg)?;
22703 }
22704 let (fc, cfg2) = (
22707 if g {
22708 self.func_g("fa_decode_combine_f32")
22709 } else {
22710 self.fa_func("fa_decode_combine_f32", head_dim)
22711 },
22712 LaunchConfig {
22713 grid_dim: (n_head as u32, 1, 1),
22714 block_dim: (head_dim as u32, 1, 1),
22715 shared_mem_bytes: 0,
22716 },
22717 );
22718 let __s_b2 = self.gpu.stream();
22719 let mut b2 = __s_b2.launch_builder(&fc);
22720 b2.arg(&*part_o)
22721 .arg(&*part_m)
22722 .arg(&*part_l)
22723 .arg(o)
22724 .arg(&hd)
22725 .arg(&nh)
22726 .arg(&nsp);
22727 unsafe {
22728 b2.launch(cfg2)?;
22729 }
22730 Ok(())
22731 }
22732
22733 #[allow(clippy::too_many_arguments)]
22744 pub fn fa_decode_batch_seqs_v4(
22745 &self,
22746 q: &CudaSlice<f32>,
22747 kv_ptrs: &cudarc::driver::CudaView<u64>,
22748 pos_seq: &CudaSlice<i32>,
22749 o: &mut CudaSlice<f32>,
22750 head_dim: usize,
22751 n_head: usize,
22752 n_head_kv: usize,
22753 b_n: usize,
22754 t_kv_max: usize,
22755 scale: f32,
22756 split_keys: usize,
22757 k_tok_bytes: usize,
22758 v_tok_bytes: usize,
22759 ) -> Result<(), Box<dyn std::error::Error>> {
22760 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
22761 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
22762 let o_len = b_n * n_head * n_splits_max * head_dim;
22763 let ml_len = b_n * n_head * n_splits_max;
22764 let mut part_guard = self.fa_part_pool.lock().unwrap();
22765 if part_guard
22766 .as_ref()
22767 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22768 .unwrap_or(true)
22769 {
22770 let old = part_guard.take();
22781 let (co, cm) = old
22782 .as_ref()
22783 .map(|pp| (pp.0.len(), pp.1.len()))
22784 .unwrap_or((0, 0));
22785 if let Some(old) = old {
22786 self.fa_part_retired.lock().unwrap().push(old);
22787 }
22788 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22789 eprintln!(
22790 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22791 co, o_len, cm, ml_len
22792 );
22793 }
22794 *part_guard =
22795 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
22796 }
22797 let pg = part_guard.as_mut().unwrap();
22798 self.gpu
22799 .stream()
22800 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22801 self.gpu
22802 .stream()
22803 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22804 self.gpu
22805 .stream()
22806 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22807 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22808 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22809 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
22810 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22811 let gqa = (n_head / n_head_kv).max(1) as u32;
22812 let f = self.func("fa_decode_vec_q_seqs_v4");
22813 let shmem = (11520 + 32 * head_dim * 2) as u32;
22815 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22816 f.set_attribute(
22817 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22818 shmem as i32,
22819 )?;
22820 let cfg = LaunchConfig {
22821 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
22822 block_dim: (32, gqa, 1),
22823 shared_mem_bytes: shmem,
22824 };
22825 {
22826 let __s_b = self.gpu.stream();
22827 let mut b = __s_b.launch_builder(&f);
22828 b.arg(q)
22829 .arg(kv_ptrs)
22830 .arg(pos_seq)
22831 .arg(&mut *part_o)
22832 .arg(&mut *part_m)
22833 .arg(&mut *part_l)
22834 .arg(&hd)
22835 .arg(&nh)
22836 .arg(&nhkv)
22837 .arg(&scale)
22838 .arg(&nspm)
22839 .arg(&spk)
22840 .arg(&ktb)
22841 .arg(&vtb);
22842 unsafe {
22843 b.launch(cfg)?;
22844 }
22845 }
22846 let fc = self.func("fa_decode_combine_seqs");
22847 let cfg2 = LaunchConfig {
22848 grid_dim: (n_head as u32, b_n as u32, 1),
22849 block_dim: (head_dim as u32, 1, 1),
22850 shared_mem_bytes: 0,
22851 };
22852 let __s_b2 = self.gpu.stream();
22853 let mut b2 = __s_b2.launch_builder(&fc);
22854 b2.arg(&*part_o)
22855 .arg(&*part_m)
22856 .arg(&*part_l)
22857 .arg(o)
22858 .arg(&hd)
22859 .arg(&nh)
22860 .arg(pos_seq)
22861 .arg(&nspm)
22862 .arg(&spk);
22863 unsafe {
22864 b2.launch(cfg2)?;
22865 }
22866 Ok(())
22867 }
22868
22869 #[allow(clippy::too_many_arguments)]
22876 pub fn append_kv_quantized_seqs(
22877 &self,
22878 k_rows: &CudaSlice<f32>,
22879 v_rows: &CudaSlice<f32>,
22880 kv_ptrs: &cudarc::driver::CudaView<u64>,
22881 pos_seq: &CudaSlice<i32>,
22882 b_n: usize,
22883 kv_dim_k: usize,
22884 kv_dim_v: usize,
22885 k_tok_bytes: usize,
22886 v_tok_bytes: usize,
22887 ) -> Result<(), Box<dyn std::error::Error>> {
22888 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
22889 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22890 let cfg = LaunchConfig {
22891 grid_dim: (nblk, b_n as u32, 1),
22892 block_dim: (32, 1, 1),
22893 shared_mem_bytes: 0,
22894 };
22895 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22896 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22897 let __s_b = self.gpu.stream();
22898 let mut b = __s_b.launch_builder(&f);
22899 b.arg(k_rows)
22900 .arg(v_rows)
22901 .arg(kv_ptrs)
22902 .arg(pos_seq)
22903 .arg(&kdk)
22904 .arg(&kdv)
22905 .arg(&ktb)
22906 .arg(&vtb);
22907 unsafe {
22908 b.launch(cfg)?;
22909 }
22910 Ok(())
22911 }
22912
22913 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
22919 std::env::var("MEMRA_NO_FA_VEC").is_err()
22920 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
22921 && base_len + 1 >= fa_vec_min_tkv()
22922 && head_dim <= 256
22923 && head_dim % 32 == 0
22924 }
22925
22926 #[allow(clippy::too_many_arguments)]
22935 pub fn fa_decode_rows(
22936 &self,
22937 q: &CudaSlice<f32>,
22938 k: &cudarc::driver::CudaView<u8>,
22939 v: &cudarc::driver::CudaView<u8>,
22940 o: &mut CudaSlice<f32>,
22941 head_dim: usize,
22942 n_head: usize,
22943 n_head_kv: usize,
22944 base_len: usize,
22945 t: usize,
22946 scale: f32,
22947 k_tok_bytes: usize,
22948 v_tok_bytes: usize,
22949 base_dev: Option<(&CudaSlice<i32>, i32)>,
22953 kv_shared: bool,
22956 g: bool,
22960 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22963 ) -> Result<(), Box<dyn std::error::Error>> {
22964 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
22965 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
22972 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22973 let v = *SP512.get_or_init(|| {
22976 std::env::var("MEMRA_FA_SP512")
22977 .ok()
22978 .and_then(|x| x.parse().ok())
22979 .unwrap_or(0)
22980 });
22981 sp = if v >= 8 {
22982 v
22983 } else {
22984 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22985 };
22986 }
22987 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22988 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22989 let gqa = (n_head / n_head_kv).max(1) as u32;
22990 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
23001 groups.push((0, t, sp));
23002 } else {
23003 let mut r0 = 0usize;
23004 while r0 < t {
23005 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
23006 let mut r1 = r0 + 1;
23007 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
23008 r1 += 1;
23009 }
23010 groups.push((r0, r1 - r0, sp_g));
23011 r0 = r1;
23012 }
23013 }
23014 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23018 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
23019 std::env::var("MEMRA_FA_SMEM_TKV")
23020 .ok()
23021 .and_then(|v| v.parse().ok())
23022 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
23023 });
23024 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
23025 let v3 = fa_v3_active(head_dim);
23026 let smem_rows =
23027 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
23028 let _ = kv_shared;
23033 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
23036 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23050 let tb512 = head_dim == 512
23052 && sp <= 32
23053 && n_head / n_head_kv.max(1) <= 16
23054 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
23055 let fname = if tb512 {
23056 "fa_decode_vec_q_rows_v4_512_tb"
23057 } else if i2 {
23058 "fa_decode_vec_q_rows_dpl16_i2"
23059 } else if head_dim == 512 {
23060 "fa_decode_vec_q_rows_dpl16"
23061 }
23062 else if v4 {
23064 "fa_decode_vec_q_rows_v4"
23065 } else if v3 {
23066 "fa_decode_vec_q_rows_v3"
23067 } else if fa_v2_on() {
23068 "fa_decode_vec_q_rows_v2"
23069 } else if smem_rows {
23070 "fa_decode_vec_q_rows_smem"
23071 } else {
23072 "fa_decode_vec_q_rows"
23073 };
23074 let f = if head_dim == 512 {
23075 self.fa_func(fname, head_dim)
23076 } else if g {
23077 self.func_g(if smem_rows {
23085 "fa_decode_vec_q_rows"
23086 } else {
23087 fname
23088 })
23089 } else {
23090 self.func(fname)
23091 };
23092 let shmem = if tb512 {
23093 let gk = Self::gkv_on();
23095 let sh =
23096 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
23097 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23098 f.set_attribute(
23099 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23100 sh as i32,
23101 )?;
23102 sh
23103 } else if v4 || v3 || smem_rows || fa_v2_on() {
23104 let sh = (if v4 {
23106 11520 + 32 * head_dim * if g { 1 } else { 2 }
23107 } else if v3 {
23108 32 * head_dim * 2
23109 } else {
23110 2 * 32 * head_dim * 2
23111 }) as u32;
23112 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23113 f.set_attribute(
23114 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23115 sh as i32,
23116 )?;
23117 sh
23118 } else {
23119 0
23120 };
23121 for &(r0, t_g, sp_g) in &groups {
23125 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
23126 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
23127 let base_i = (base_len + r0) as i32;
23128 let o_len = t_g * n_head * n_splits_g * head_dim;
23129 let ml_len = t_g * n_head * n_splits_g;
23130 let mut part_guard = self.fa_part_pool.lock().unwrap();
23131 if part_guard
23132 .as_ref()
23133 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23134 .unwrap_or(true)
23135 {
23136 let old = part_guard.take();
23147 let (co, cm) = old
23148 .as_ref()
23149 .map(|pp| (pp.0.len(), pp.1.len()))
23150 .unwrap_or((0, 0));
23151 if let Some(old) = old {
23152 self.fa_part_retired.lock().unwrap().push(old);
23153 }
23154 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23155 eprintln!(
23156 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23157 co, o_len, cm, ml_len
23158 );
23159 }
23160 *part_guard =
23161 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
23162 }
23163 let pg = part_guard.as_mut().unwrap();
23164 self.gpu
23165 .stream()
23166 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23167 self.gpu
23168 .stream()
23169 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23170 self.gpu
23171 .stream()
23172 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23173 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23174 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
23175 let qv = self.view(q, t * n_head * head_dim);
23176 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
23177 let cfg = LaunchConfig {
23178 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
23179 block_dim: (32, gqa, 1),
23180 shared_mem_bytes: shmem,
23181 };
23182 {
23183 let __s_b = self.gpu.stream();
23184 let mut b = __s_b.launch_builder(&f);
23185 if tb512 {
23186 let (bd, plus) =
23188 base_dev.expect("hd512 rows twin requires a device base counter");
23189 let plus_g = plus + r0 as i32;
23190 let nr = t_g as i32;
23191 if Self::pdl_on() && Self::pdl_wb_on() {
23192 use cudarc::driver::{DevicePtr, DevicePtrMut};
23194 let s = &self.gpu.stream();
23195 let (pq, _b0) = q_g.device_ptr(s);
23196 let (pk, _b1) = k.device_ptr(s);
23197 let (pv, _b2) = v.device_ptr(s);
23198 let (po, _b3) = part_o.device_ptr_mut(s);
23199 let (pm, _b4) = part_m.device_ptr_mut(s);
23200 let (pl, _b5) = part_l.device_ptr_mut(s);
23201 let (pb, _b6) = bd.device_ptr(s);
23202 let mut ps = [
23203 &pq as *const _ as *mut std::ffi::c_void,
23204 &pk as *const _ as *mut _,
23205 &pv as *const _ as *mut _,
23206 &po as *const _ as *mut _,
23207 &pm as *const _ as *mut _,
23208 &pl as *const _ as *mut _,
23209 &hd as *const _ as *mut _,
23210 &nh as *const _ as *mut _,
23211 &nhkv as *const _ as *mut _,
23212 &pb as *const _ as *mut _,
23213 &plus_g as *const _ as *mut _,
23214 &scale as *const _ as *mut _,
23215 &nspm as *const _ as *mut _,
23216 &spk as *const _ as *mut _,
23217 &ktb as *const _ as *mut _,
23218 &vtb as *const _ as *mut _,
23219 &nr as *const _ as *mut _,
23220 ];
23221 unsafe {
23222 self.launch_pdl_flash(
23223 Self::gkv_on(),
23224 "fa_decode_vec_q_rows_v4_512_tb",
23225 (n_head_kv as u32, n_splits_g as u32, 1),
23226 (32, gqa, 1),
23227 shmem,
23228 &mut ps,
23229 )?;
23230 }
23231 } else {
23232 let cfg_tb = LaunchConfig {
23233 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
23234 block_dim: (32, gqa, 1),
23235 shared_mem_bytes: shmem,
23236 };
23237 b.arg(&q_g)
23238 .arg(k)
23239 .arg(v)
23240 .arg(&mut *part_o)
23241 .arg(&mut *part_m)
23242 .arg(&mut *part_l)
23243 .arg(&hd)
23244 .arg(&nh)
23245 .arg(&nhkv)
23246 .arg(bd)
23247 .arg(&plus_g)
23248 .arg(&scale)
23249 .arg(&nspm)
23250 .arg(&spk)
23251 .arg(&ktb)
23252 .arg(&vtb)
23253 .arg(&nr);
23254 unsafe {
23255 b.launch(cfg_tb)?;
23256 }
23257 }
23258 } else if head_dim == 512 {
23259 let (bd, plus) =
23260 base_dev.expect("hd512 rows twin requires a device base counter");
23261 let plus_g = plus + r0 as i32;
23262 b.arg(&q_g)
23263 .arg(k)
23264 .arg(v)
23265 .arg(&mut *part_o)
23266 .arg(&mut *part_m)
23267 .arg(&mut *part_l)
23268 .arg(&hd)
23269 .arg(&nh)
23270 .arg(&nhkv)
23271 .arg(bd)
23272 .arg(&plus_g)
23273 .arg(&scale)
23274 .arg(&nspm)
23275 .arg(&spk)
23276 .arg(&ktb)
23277 .arg(&vtb);
23278 unsafe {
23279 b.launch(cfg)?;
23280 }
23281 } else {
23282 b.arg(&q_g)
23283 .arg(k)
23284 .arg(v)
23285 .arg(&mut *part_o)
23286 .arg(&mut *part_m)
23287 .arg(&mut *part_l)
23288 .arg(&hd)
23289 .arg(&nh)
23290 .arg(&nhkv)
23291 .arg(&base_i)
23292 .arg(&scale)
23293 .arg(&nspm)
23294 .arg(&spk)
23295 .arg(&ktb)
23296 .arg(&vtb);
23297 unsafe {
23298 b.launch(cfg)?;
23299 }
23300 }
23301 }
23302 let cfg2 = LaunchConfig {
23303 grid_dim: (n_head as u32, t_g as u32, 1),
23304 block_dim: (head_dim as u32, 1, 1),
23305 shared_mem_bytes: 0,
23306 };
23307 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
23308 if head_dim == 512 {
23309 let (bd, plus) = base_dev.unwrap();
23312 let plus_g = plus + r0 as i32;
23313 if let Some((oq, od)) = q8_out.as_mut() {
23314 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
23316 if Self::pdl_on() && Self::pdl_wb_on() {
23317 use cudarc::driver::{DevicePtr, DevicePtrMut};
23319 let s = &self.gpu.stream();
23320 let (po, _g0) = part_o.device_ptr(s);
23321 let (pm, _g1) = part_m.device_ptr(s);
23322 let (pl, _g2) = part_l.device_ptr(s);
23323 let (pq, _g3) = oq.device_ptr_mut(s);
23324 let (pd, _g4) = od.device_ptr_mut(s);
23325 let (pb, _g5) = bd.device_ptr(s);
23326 let mut ps = [
23327 &po as *const _ as *mut std::ffi::c_void,
23328 &pm as *const _ as *mut _,
23329 &pl as *const _ as *mut _,
23330 &pq as *const _ as *mut _,
23331 &pd as *const _ as *mut _,
23332 &hd as *const _ as *mut _,
23333 &nh as *const _ as *mut _,
23334 &pb as *const _ as *mut _,
23335 &plus_g as *const _ as *mut _,
23336 &nspm as *const _ as *mut _,
23337 &spk as *const _ as *mut _,
23338 ];
23339 unsafe {
23340 self.launch_pdl_flash(
23341 Self::gkv_on(),
23342 "fa_decode_combine_rows_dc_q8_1",
23343 cfg2.grid_dim,
23344 cfg2.block_dim,
23345 0,
23346 &mut ps,
23347 )?;
23348 }
23349 continue;
23350 }
23351 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
23352 let __s_b2 = self.gpu.stream();
23353 let mut b2 = __s_b2.launch_builder(&fc);
23354 b2.arg(&*part_o)
23355 .arg(&*part_m)
23356 .arg(&*part_l)
23357 .arg(&mut **oq)
23358 .arg(&mut **od)
23359 .arg(&hd)
23360 .arg(&nh)
23361 .arg(bd)
23362 .arg(&plus_g)
23363 .arg(&nspm)
23364 .arg(&spk);
23365 unsafe {
23366 b2.launch(cfg2)?;
23367 }
23368 continue;
23369 }
23370 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
23371 let __s_b2 = self.gpu.stream();
23372 let mut b2 = __s_b2.launch_builder(&fc);
23373 b2.arg(&*part_o)
23374 .arg(&*part_m)
23375 .arg(&*part_l)
23376 .arg(&mut o_g)
23377 .arg(&hd)
23378 .arg(&nh)
23379 .arg(bd)
23380 .arg(&plus_g)
23381 .arg(&nspm)
23382 .arg(&spk);
23383 unsafe {
23384 b2.launch(cfg2)?;
23385 }
23386 } else {
23387 assert!(
23390 q8_out.is_none(),
23391 "rows q8 emit requires the hd512 dc combine"
23392 );
23393 let fc = self.func("fa_decode_combine_rows");
23394 let __s_b2 = self.gpu.stream();
23395 let mut b2 = __s_b2.launch_builder(&fc);
23396 b2.arg(&*part_o)
23397 .arg(&*part_m)
23398 .arg(&*part_l)
23399 .arg(&mut o_g)
23400 .arg(&hd)
23401 .arg(&nh)
23402 .arg(&base_i)
23403 .arg(&nspm)
23404 .arg(&spk);
23405 unsafe {
23406 b2.launch(cfg2)?;
23407 }
23408 }
23409 }
23410 Ok(())
23411 }
23412
23413 #[allow(clippy::too_many_arguments)]
23417 pub fn fa_decode_rows_w(
23418 &self,
23419 q: &CudaSlice<f32>,
23420 k: &cudarc::driver::CudaView<u8>,
23421 v: &cudarc::driver::CudaView<u8>,
23422 o: &mut CudaSlice<f32>,
23423 head_dim: usize,
23424 n_head: usize,
23425 n_head_kv: usize,
23426 base_dev: &CudaSlice<i32>,
23427 base_plus: i32,
23428 t: usize,
23429 scale: f32,
23430 window: usize,
23431 k_tok_bytes: usize,
23432 v_tok_bytes: usize,
23433 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23434 ) -> Result<(), Box<dyn std::error::Error>> {
23435 debug_assert!(head_dim == 256);
23440 let sp = {
23448 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23449 let v = *SPW.get_or_init(|| {
23450 std::env::var("MEMRA_FA_SPW")
23451 .ok()
23452 .and_then(|x| x.parse().ok())
23453 .unwrap_or(0)
23454 });
23455 if v >= 8 {
23456 v
23457 } else {
23458 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
23459 }
23460 };
23461 let n_splits_max = (window + sp - 1) / sp;
23462 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23463 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
23464 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23465 let gqa = (n_head / n_head_kv).max(1) as u32;
23466 let o_len = t * n_head * n_splits_max * head_dim;
23467 let ml_len = t * n_head * n_splits_max;
23468 let mut part_guard = self.fa_part_pool.lock().unwrap();
23469 if part_guard
23470 .as_ref()
23471 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23472 .unwrap_or(true)
23473 {
23474 let old = part_guard.take();
23485 let (co, cm) = old
23486 .as_ref()
23487 .map(|pp| (pp.0.len(), pp.1.len()))
23488 .unwrap_or((0, 0));
23489 if let Some(old) = old {
23490 self.fa_part_retired.lock().unwrap().push(old);
23491 }
23492 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23493 eprintln!(
23494 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23495 co, o_len, cm, ml_len
23496 );
23497 }
23498 *part_guard =
23499 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
23500 }
23501 let pg = part_guard.as_mut().unwrap();
23502 self.gpu
23503 .stream()
23504 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23505 self.gpu
23506 .stream()
23507 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23508 self.gpu
23509 .stream()
23510 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23511 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23512 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23518 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
23519 std::env::var("MEMRA_FA_SMEM_TKV")
23520 .ok()
23521 .and_then(|v| v.parse().ok())
23522 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
23523 });
23524 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23530 let wg = Self::wkv_on();
23535 let sp2 =
23538 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
23539 if sp2 {
23540 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23541 if Self::pdl_on() && Self::pdl_wb_on() {
23542 use cudarc::driver::{DevicePtr, DevicePtrMut};
23544 let s = &self.gpu.stream();
23545 let (pq, _b0) = q.device_ptr(s);
23546 let (pk, _b1) = k.device_ptr(s);
23547 let (pv, _b2) = v.device_ptr(s);
23548 let (po, _b3) = part_o.device_ptr_mut(s);
23549 let (pm, _b4) = part_m.device_ptr_mut(s);
23550 let (pl, _b5) = part_l.device_ptr_mut(s);
23551 let (pb, _b6) = base_dev.device_ptr(s);
23552 let mut ps = [
23553 &pq as *const _ as *mut std::ffi::c_void,
23554 &pk as *const _ as *mut _,
23555 &pv as *const _ as *mut _,
23556 &po as *const _ as *mut _,
23557 &pm as *const _ as *mut _,
23558 &pl as *const _ as *mut _,
23559 &hd as *const _ as *mut _,
23560 &nh as *const _ as *mut _,
23561 &nhkv as *const _ as *mut _,
23562 &pb as *const _ as *mut _,
23563 &base_plus as *const _ as *mut _,
23564 &scale as *const _ as *mut _,
23565 &nspm as *const _ as *mut _,
23566 &spk as *const _ as *mut _,
23567 &ktb as *const _ as *mut _,
23568 &vtb as *const _ as *mut _,
23569 &wini as *const _ as *mut _,
23570 ];
23571 unsafe {
23572 self.launch_pdl_flash(
23573 wg,
23574 "fa_decode_vec_q_rows_v4_w_sp",
23575 (n_head_kv as u32, n_splits_max as u32, t as u32),
23576 (32, gqa + 1, 1),
23577 sh,
23578 &mut ps,
23579 )?;
23580 }
23581 } else {
23582 let f = if wg {
23583 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
23584 } else {
23585 self.func("fa_decode_vec_q_rows_v4_w_sp")
23586 };
23587 f.set_attribute(
23588 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23589 sh as i32,
23590 )?;
23591 let cfg = LaunchConfig {
23592 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23593 block_dim: (32, gqa + 1, 1),
23594 shared_mem_bytes: sh,
23595 };
23596 let __s_b = self.gpu.stream();
23597 let mut b = __s_b.launch_builder(&f);
23598 b.arg(q)
23599 .arg(k)
23600 .arg(v)
23601 .arg(&mut *part_o)
23602 .arg(&mut *part_m)
23603 .arg(&mut *part_l)
23604 .arg(&hd)
23605 .arg(&nh)
23606 .arg(&nhkv)
23607 .arg(base_dev)
23608 .arg(&base_plus)
23609 .arg(&scale)
23610 .arg(&nspm)
23611 .arg(&spk)
23612 .arg(&ktb)
23613 .arg(&vtb)
23614 .arg(&wini);
23615 unsafe {
23616 b.launch(cfg)?;
23617 }
23618 }
23619 } else {
23620 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
23621 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23623 use cudarc::driver::{DevicePtr, DevicePtrMut};
23624 let s = &self.gpu.stream();
23625 let (pq, _b0) = q.device_ptr(s);
23626 let (pk, _b1) = k.device_ptr(s);
23627 let (pv, _b2) = v.device_ptr(s);
23628 let (po, _b3) = part_o.device_ptr_mut(s);
23629 let (pm, _b4) = part_m.device_ptr_mut(s);
23630 let (pl, _b5) = part_l.device_ptr_mut(s);
23631 let (pb, _b6) = base_dev.device_ptr(s);
23632 let mut ps = [
23633 &pq as *const _ as *mut std::ffi::c_void,
23634 &pk as *const _ as *mut _,
23635 &pv as *const _ as *mut _,
23636 &po as *const _ as *mut _,
23637 &pm as *const _ as *mut _,
23638 &pl as *const _ as *mut _,
23639 &hd as *const _ as *mut _,
23640 &nh as *const _ as *mut _,
23641 &nhkv as *const _ as *mut _,
23642 &pb as *const _ as *mut _,
23643 &base_plus as *const _ as *mut _,
23644 &scale as *const _ as *mut _,
23645 &nspm as *const _ as *mut _,
23646 &spk as *const _ as *mut _,
23647 &ktb as *const _ as *mut _,
23648 &vtb as *const _ as *mut _,
23649 &wini as *const _ as *mut _,
23650 ];
23651 unsafe {
23652 self.launch_pdl_flash(
23653 wg,
23654 "fa_decode_vec_q_rows_v4_w",
23655 (n_head_kv as u32, n_splits_max as u32, t as u32),
23656 (32, gqa, 1),
23657 sh,
23658 &mut ps,
23659 )?;
23660 }
23661 } else {
23662 let pick = |name: &str| {
23663 if wg {
23664 self.func_g(name)
23665 } else {
23666 self.func(name)
23667 }
23668 };
23669 let (f, sh) = if fa_v4_at(window) {
23670 let f = pick("fa_decode_vec_q_rows_v4_w");
23671 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
23672 } else if smem_tkv > 0 && window >= smem_tkv {
23673 (
23676 pick("fa_decode_vec_q_rows_smem_w"),
23677 (2 * 32 * head_dim * 2) as u32,
23678 )
23679 } else {
23680 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
23681 };
23682 f.set_attribute(
23683 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23684 sh as i32,
23685 )?;
23686 let cfg = LaunchConfig {
23687 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23688 block_dim: (32, gqa, 1),
23689 shared_mem_bytes: sh,
23690 };
23691 let __s_b = self.gpu.stream();
23692 let mut b = __s_b.launch_builder(&f);
23693 b.arg(q)
23694 .arg(k)
23695 .arg(v)
23696 .arg(&mut *part_o)
23697 .arg(&mut *part_m)
23698 .arg(&mut *part_l)
23699 .arg(&hd)
23700 .arg(&nh)
23701 .arg(&nhkv)
23702 .arg(base_dev)
23703 .arg(&base_plus)
23704 .arg(&scale)
23705 .arg(&nspm)
23706 .arg(&spk)
23707 .arg(&ktb)
23708 .arg(&vtb)
23709 .arg(&wini);
23710 unsafe {
23711 b.launch(cfg)?;
23712 }
23713 }
23714 }
23715 let cfg2 = LaunchConfig {
23716 grid_dim: (n_head as u32, t as u32, 1),
23717 block_dim: (head_dim as u32, 1, 1),
23718 shared_mem_bytes: 0,
23719 };
23720 if let Some((oq, od)) = q8_out {
23721 if Self::pdl_on() && Self::pdl_wb_on() {
23724 use cudarc::driver::{DevicePtr, DevicePtrMut};
23726 let s = &self.gpu.stream();
23727 let (po, _g0) = part_o.device_ptr(s);
23728 let (pm, _g1) = part_m.device_ptr(s);
23729 let (pl, _g2) = part_l.device_ptr(s);
23730 let (pq, _g3) = oq.device_ptr_mut(s);
23731 let (pd, _g4) = od.device_ptr_mut(s);
23732 let mut ps = [
23733 &po as *const _ as *mut std::ffi::c_void,
23734 &pm as *const _ as *mut _,
23735 &pl as *const _ as *mut _,
23736 &pq as *const _ as *mut _,
23737 &pd as *const _ as *mut _,
23738 &hd as *const _ as *mut _,
23739 &nh as *const _ as *mut _,
23740 &nspm as *const _ as *mut _,
23741 &spk as *const _ as *mut _,
23742 &wini as *const _ as *mut _,
23743 ];
23744 unsafe {
23745 self.launch_pdl_flash(
23746 wg,
23747 "fa_decode_combine_rows_w_q8_1",
23748 cfg2.grid_dim,
23749 cfg2.block_dim,
23750 0,
23751 &mut ps,
23752 )?;
23753 }
23754 return Ok(());
23755 }
23756 let fc = if wg {
23757 self.func_g("fa_decode_combine_rows_w_q8_1")
23758 } else {
23759 self.func("fa_decode_combine_rows_w_q8_1")
23760 };
23761 let __s_b2 = self.gpu.stream();
23762 let mut b2 = __s_b2.launch_builder(&fc);
23763 b2.arg(&*part_o)
23764 .arg(&*part_m)
23765 .arg(&*part_l)
23766 .arg(oq)
23767 .arg(od)
23768 .arg(&hd)
23769 .arg(&nh)
23770 .arg(&nspm)
23771 .arg(&spk)
23772 .arg(&wini);
23773 unsafe {
23774 b2.launch(cfg2)?;
23775 }
23776 return Ok(());
23777 }
23778 let fc = if wg {
23779 self.func_g("fa_decode_combine_rows_w")
23780 } else {
23781 self.func("fa_decode_combine_rows_w")
23782 };
23783 let __s_b2 = self.gpu.stream();
23784 let mut b2 = __s_b2.launch_builder(&fc);
23785 b2.arg(&*part_o)
23786 .arg(&*part_m)
23787 .arg(&*part_l)
23788 .arg(o)
23789 .arg(&hd)
23790 .arg(&nh)
23791 .arg(&nspm)
23792 .arg(&spk)
23793 .arg(&wini);
23794 unsafe {
23795 b2.launch(cfg2)?;
23796 }
23797 Ok(())
23798 }
23799
23800 #[allow(clippy::too_many_arguments)]
23806 pub fn fa_decode_rows_dc(
23807 &self,
23808 q: &CudaSlice<f32>,
23809 k: &cudarc::driver::CudaView<u8>,
23810 v: &cudarc::driver::CudaView<u8>,
23811 o: &mut CudaSlice<f32>,
23812 head_dim: usize,
23813 n_head: usize,
23814 n_head_kv: usize,
23815 base_dev: &CudaSlice<i32>,
23816 t_kv_upper: usize,
23817 t: usize,
23818 scale: f32,
23819 k_tok_bytes: usize,
23820 v_tok_bytes: usize,
23821 base_plus: i32,
23822 g: bool,
23823 ) -> Result<(), Box<dyn std::error::Error>> {
23824 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
23825 assert!(
23826 v4 || fa_v3_active(head_dim),
23827 "stream fa rows requires the v3 or v4 lane"
23828 );
23829 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
23830 if v4 {
23831 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23832 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23833 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23834 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23835 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23836 let gqa = (n_head / n_head_kv).max(1) as u32;
23837 let o_len = t * n_head * n_splits_max * head_dim;
23838 let ml_len = t * n_head * n_splits_max;
23839 let mut part_guard = self.fa_part_pool.lock().unwrap();
23840 if part_guard
23841 .as_ref()
23842 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23843 .unwrap_or(true)
23844 {
23845 let old = part_guard.take();
23856 let (co, cm) = old
23857 .as_ref()
23858 .map(|pp| (pp.0.len(), pp.1.len()))
23859 .unwrap_or((0, 0));
23860 if let Some(old) = old {
23861 self.fa_part_retired.lock().unwrap().push(old);
23862 }
23863 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23864 eprintln!(
23865 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23866 co, o_len, cm, ml_len
23867 );
23868 }
23869 *part_guard =
23870 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
23871 }
23872 let pg = part_guard.as_mut().unwrap();
23873 self.gpu
23874 .stream()
23875 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23876 self.gpu
23877 .stream()
23878 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23879 self.gpu
23880 .stream()
23881 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23882 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23883 let f = if g {
23884 self.func_g("fa_decode_vec_q_rows_v4_dc")
23885 } else {
23886 self.func("fa_decode_vec_q_rows_v4_dc")
23887 };
23888 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23889 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23890 f.set_attribute(
23891 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23892 sh as i32,
23893 )?;
23894 let cfg = LaunchConfig {
23895 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23896 block_dim: (32, gqa, 1),
23897 shared_mem_bytes: sh,
23898 };
23899 let __s_b = self.gpu.stream();
23900 let mut b = __s_b.launch_builder(&f);
23901 b.arg(q)
23902 .arg(k)
23903 .arg(v)
23904 .arg(&mut *part_o)
23905 .arg(&mut *part_m)
23906 .arg(&mut *part_l)
23907 .arg(&hd)
23908 .arg(&nh)
23909 .arg(&nhkv)
23910 .arg(base_dev)
23911 .arg(&base_plus)
23912 .arg(&scale)
23913 .arg(&nspm)
23914 .arg(&spk)
23915 .arg(&ktb)
23916 .arg(&vtb);
23917 unsafe {
23918 b.launch(cfg)?;
23919 }
23920 let fc = self.func("fa_decode_combine_rows_dc");
23921 let cfg2 = LaunchConfig {
23922 grid_dim: (n_head as u32, t as u32, 1),
23923 block_dim: (head_dim as u32, 1, 1),
23924 shared_mem_bytes: 0,
23925 };
23926 let __s_b2 = self.gpu.stream();
23927 let mut b2 = __s_b2.launch_builder(&fc);
23928 b2.arg(&*part_o)
23929 .arg(&*part_m)
23930 .arg(&*part_l)
23931 .arg(o)
23932 .arg(&hd)
23933 .arg(&nh)
23934 .arg(base_dev)
23935 .arg(&base_plus)
23936 .arg(&nspm)
23937 .arg(&spk);
23938 unsafe {
23939 b2.launch(cfg2)?;
23940 }
23941 return Ok(());
23942 }
23943 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23944 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23945 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23946 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23947 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23948 let gqa = (n_head / n_head_kv).max(1) as u32;
23949 let o_len = t * n_head * n_splits_max * head_dim;
23950 let ml_len = t * n_head * n_splits_max;
23951 let mut part_guard = self.fa_part_pool.lock().unwrap();
23952 if part_guard
23953 .as_ref()
23954 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23955 .unwrap_or(true)
23956 {
23957 let old = part_guard.take();
23968 let (co, cm) = old
23969 .as_ref()
23970 .map(|pp| (pp.0.len(), pp.1.len()))
23971 .unwrap_or((0, 0));
23972 if let Some(old) = old {
23973 self.fa_part_retired.lock().unwrap().push(old);
23974 }
23975 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23976 eprintln!(
23977 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23978 co, o_len, cm, ml_len
23979 );
23980 }
23981 *part_guard =
23982 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
23983 }
23984 let pg = part_guard.as_mut().unwrap();
23985 self.gpu
23986 .stream()
23987 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23988 self.gpu
23989 .stream()
23990 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23991 self.gpu
23992 .stream()
23993 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23994 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23995 let f = self.func("fa_decode_vec_q_rows_v3_dc");
23996 let sh = (32 * head_dim * 2) as u32;
23997 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23998 f.set_attribute(
23999 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24000 sh as i32,
24001 )?;
24002 let cfg = LaunchConfig {
24003 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
24004 block_dim: (32, gqa, 1),
24005 shared_mem_bytes: sh,
24006 };
24007 let __s_b = self.gpu.stream();
24008 let mut b = __s_b.launch_builder(&f);
24009 b.arg(q)
24010 .arg(k)
24011 .arg(v)
24012 .arg(&mut *part_o)
24013 .arg(&mut *part_m)
24014 .arg(&mut *part_l)
24015 .arg(&hd)
24016 .arg(&nh)
24017 .arg(&nhkv)
24018 .arg(base_dev)
24019 .arg(&scale)
24020 .arg(&nspm)
24021 .arg(&spk)
24022 .arg(&ktb)
24023 .arg(&vtb);
24024 unsafe {
24025 b.launch(cfg)?;
24026 }
24027 let fc = self.func("fa_decode_combine_rows_dc");
24028 let cfg2 = LaunchConfig {
24029 grid_dim: (n_head as u32, t as u32, 1),
24030 block_dim: (head_dim as u32, 1, 1),
24031 shared_mem_bytes: 0,
24032 };
24033 let plus0 = 0i32;
24034 let __s_b2 = self.gpu.stream();
24035 let mut b2 = __s_b2.launch_builder(&fc);
24036 b2.arg(&*part_o)
24037 .arg(&*part_m)
24038 .arg(&*part_l)
24039 .arg(o)
24040 .arg(&hd)
24041 .arg(&nh)
24042 .arg(base_dev)
24043 .arg(&plus0)
24044 .arg(&nspm)
24045 .arg(&spk);
24046 unsafe {
24047 b2.launch(cfg2)?;
24048 }
24049 Ok(())
24050 }
24051
24052 pub fn fa_decode_dc(
24063 &self,
24064 q: &CudaSlice<f32>,
24065 k: &cudarc::driver::CudaView<u8>,
24066 v: &cudarc::driver::CudaView<u8>,
24067 o: &mut CudaSlice<f32>,
24068 head_dim: usize,
24069 n_head: usize,
24070 n_head_kv: usize,
24071 t_kv_dev: &CudaSlice<i32>,
24072 bucket_max: usize,
24073 scale: f32,
24074 k_tok_bytes: usize,
24075 v_tok_bytes: usize,
24076 g: bool,
24077 ) -> Result<(), Box<dyn std::error::Error>> {
24078 self.fa_decode_dc_q8(
24079 q,
24080 k,
24081 v,
24082 o,
24083 head_dim,
24084 n_head,
24085 n_head_kv,
24086 t_kv_dev,
24087 bucket_max,
24088 scale,
24089 k_tok_bytes,
24090 v_tok_bytes,
24091 g,
24092 None,
24093 )
24094 }
24095
24096 #[allow(clippy::too_many_arguments)]
24099 pub fn fa_decode_dc_q8(
24100 &self,
24101 q: &CudaSlice<f32>,
24102 k: &cudarc::driver::CudaView<u8>,
24103 v: &cudarc::driver::CudaView<u8>,
24104 o: &mut CudaSlice<f32>,
24105 head_dim: usize,
24106 n_head: usize,
24107 n_head_kv: usize,
24108 t_kv_dev: &CudaSlice<i32>,
24109 bucket_max: usize,
24110 scale: f32,
24111 k_tok_bytes: usize,
24112 v_tok_bytes: usize,
24113 g: bool,
24114 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
24115 ) -> Result<(), Box<dyn std::error::Error>> {
24116 let mut fa_vec =
24124 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24125 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
24126 fa_vec = false;
24127 } let sp = fa_split_keys(bucket_max, n_head_kv);
24129 let n_splits = if fa_vec {
24130 ((bucket_max + sp - 1) / sp).max(1)
24131 } else {
24132 ((bucket_max + 255) / 256).max(1)
24133 };
24134 let o_len = n_head * n_splits * head_dim;
24135 let ml_len = n_head * n_splits;
24136 let mut part_guard = self.fa_part_pool.lock().unwrap();
24137 if part_guard
24138 .as_ref()
24139 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
24140 .unwrap_or(true)
24141 {
24142 let old = part_guard.take();
24153 let (co, cm) = old
24154 .as_ref()
24155 .map(|pp| (pp.0.len(), pp.1.len()))
24156 .unwrap_or((0, 0));
24157 if let Some(old) = old {
24158 self.fa_part_retired.lock().unwrap().push(old);
24159 }
24160 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
24161 eprintln!(
24162 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
24163 co, o_len, cm, ml_len
24164 );
24165 }
24166 *part_guard =
24167 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
24168 }
24169 let pg = part_guard.as_mut().unwrap();
24170 self.gpu
24171 .stream()
24172 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24173 self.gpu
24174 .stream()
24175 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24176 self.gpu
24177 .stream()
24178 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24179 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24180 let (hd, nh, nhkv, nsp) = (
24181 head_dim as i32,
24182 n_head as i32,
24183 n_head_kv as i32,
24184 n_splits as i32,
24185 );
24186 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24187 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
24188 let deep = fa_vec
24191 && head_dim == 256
24192 && fa_v4_at(bucket_max)
24193 && !g
24194 && fa_deep_at(bucket_max)
24195 && !matches!(fa_v4_mode(), "noB3" | "stage");
24196 let (f, cfg) = if fa_vec
24197 && head_dim == 512
24198 && bucket_max >= {
24199 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24200 *FA512_MIN_DC.get_or_init(|| {
24201 std::env::var("MEMRA_FA512_MIN")
24202 .ok()
24203 .and_then(|v| v.parse().ok())
24204 .unwrap_or(512)
24205 })
24206 } {
24207 let gqa = (n_head / n_head_kv).max(1) as u32;
24209 (
24210 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
24211 LaunchConfig {
24212 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24213 block_dim: (32, gqa, 1),
24214 shared_mem_bytes: 0,
24215 },
24216 )
24217 } else if fa_vec && head_dim == 512 {
24218 let q_view = q.as_view();
24221 let mut o_view = o.as_view_mut();
24222 return self.fa_decode_scalar_unified(
24223 &q_view,
24224 k,
24225 v,
24226 &mut o_view,
24227 head_dim,
24228 n_head,
24229 n_head_kv,
24230 0,
24231 Some(t_kv_dev),
24232 scale,
24233 n_splits,
24234 sp,
24235 k_tok_bytes,
24236 v_tok_bytes,
24237 g,
24238 &mut *part_o,
24239 &mut *part_m,
24240 &mut *part_l,
24241 q8_out,
24242 );
24243 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
24244 let gqa = (n_head / n_head_kv).max(1) as u32;
24247 let fv = if g {
24248 self.func_g("fa_decode_vec_q_v4_dc")
24249 } else if deep {
24250 self.func("fa_decode_vec_q_v4_deep_dc")
24251 } else {
24252 self.func("fa_decode_vec_q_v4_dc")
24253 };
24254 let shmem =
24255 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
24256 use cudarc::driver::sys::CUfunction_attribute_enum as A;
24257 fv.set_attribute(
24258 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24259 shmem as i32,
24260 )?;
24261 (
24262 fv,
24263 LaunchConfig {
24264 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24265 block_dim: (32, gqa, 1),
24266 shared_mem_bytes: shmem,
24267 },
24268 )
24269 } else if fa_vec && fa_v3_active(head_dim) {
24270 let gqa = (n_head / n_head_kv).max(1) as u32;
24273 let fv = if g {
24274 self.func_g("fa_decode_vec_q_v3_dc")
24275 } else {
24276 self.func("fa_decode_vec_q_v3_dc")
24277 };
24278 let shmem = (32 * head_dim * 2) as u32; (
24280 fv,
24281 LaunchConfig {
24282 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24283 block_dim: (32, gqa, 1),
24284 shared_mem_bytes: shmem,
24285 },
24286 )
24287 } else if fa_vec && fa_v2_on() {
24288 let gqa = (n_head / n_head_kv).max(1) as u32;
24292 let fv = if g {
24293 self.func_g("fa_decode_vec_q_v2_dc")
24294 } else {
24295 self.func("fa_decode_vec_q_v2_dc")
24296 };
24297 let shmem = (2 * 32 * head_dim * 2) as u32; (
24299 fv,
24300 LaunchConfig {
24301 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24302 block_dim: (32, gqa, 1),
24303 shared_mem_bytes: shmem,
24304 },
24305 )
24306 } else if fa_vec {
24307 let gqa = (n_head / n_head_kv).max(1) as u32;
24308 let fv = if g {
24310 self.func_g("fa_decode_vec_q_dc")
24311 } else {
24312 self.func("fa_decode_vec_q_dc")
24313 };
24314 (
24315 fv,
24316 LaunchConfig {
24317 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24318 block_dim: (32, gqa, 1),
24319 shared_mem_bytes: 0,
24320 },
24321 )
24322 } else {
24323 let q_view = q.as_view();
24324 let mut o_view = o.as_view_mut();
24325 return self.fa_decode_scalar_unified(
24326 &q_view,
24327 k,
24328 v,
24329 &mut o_view,
24330 head_dim,
24331 n_head,
24332 n_head_kv,
24333 0,
24334 Some(t_kv_dev),
24335 scale,
24336 n_splits,
24337 if fa_vec { sp } else { 256 },
24338 k_tok_bytes,
24339 v_tok_bytes,
24340 g,
24341 &mut *part_o,
24342 &mut *part_m,
24343 &mut *part_l,
24344 q8_out,
24345 );
24346 };
24347 let ski = sp as i32; let __s_b = self.gpu.stream();
24349 let mut b = __s_b.launch_builder(&f);
24350 b.arg(q)
24351 .arg(k)
24352 .arg(v)
24353 .arg(&mut *part_o)
24354 .arg(&mut *part_m)
24355 .arg(&mut *part_l)
24356 .arg(&hd)
24357 .arg(&nh)
24358 .arg(&nhkv)
24359 .arg(t_kv_dev)
24360 .arg(&scale)
24361 .arg(&nsp)
24362 .arg(&ski)
24363 .arg(&ktb)
24364 .arg(&vtb);
24365 unsafe {
24366 b.launch(cfg)?;
24367 }
24368 let cfg2 = LaunchConfig {
24369 grid_dim: (n_head as u32, 1, 1),
24370 block_dim: (head_dim as u32, 1, 1),
24371 shared_mem_bytes: 0,
24372 };
24373 if let Some((oq, od)) = q8_out {
24374 let fc = if g {
24375 self.func_g("fa_decode_combine_q8_1")
24376 } else {
24377 self.fa_func("fa_decode_combine_q8_1", head_dim)
24378 };
24379 let __s_b2 = self.gpu.stream();
24380 let mut b2 = __s_b2.launch_builder(&fc);
24381 b2.arg(&*part_o)
24382 .arg(&*part_m)
24383 .arg(&*part_l)
24384 .arg(oq)
24385 .arg(od)
24386 .arg(&hd)
24387 .arg(&nh)
24388 .arg(&nsp);
24389 unsafe {
24390 b2.launch(cfg2)?;
24391 }
24392 return Ok(());
24393 }
24394 let fc = if g {
24395 self.func_g("fa_decode_combine_f32")
24396 } else {
24397 self.fa_func("fa_decode_combine_f32", head_dim)
24398 };
24399 let __s_b2 = self.gpu.stream();
24400 let mut b2 = __s_b2.launch_builder(&fc);
24401 b2.arg(&*part_o)
24402 .arg(&*part_m)
24403 .arg(&*part_l)
24404 .arg(o)
24405 .arg(&hd)
24406 .arg(&nh)
24407 .arg(&nsp);
24408 unsafe {
24409 b2.launch(cfg2)?;
24410 }
24411 Ok(())
24412 }
24413
24414 #[allow(clippy::too_many_arguments)]
24418 pub fn append_kv_quantized_dcw(
24419 &self,
24420 k_row: &CudaSlice<f32>,
24421 v_row: &CudaSlice<f32>,
24422 kc: &mut CudaSlice<u8>,
24423 vc: &mut CudaSlice<u8>,
24424 len_dev: &CudaSlice<i32>,
24425 base_dev: Option<&CudaSlice<i32>>,
24426 kv_dim_k: usize,
24427 kv_dim_v: usize,
24428 k_tok_bytes: usize,
24429 v_tok_bytes: usize,
24430 ) -> Result<(), Box<dyn std::error::Error>> {
24431 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
24432 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
24433 let cfg = LaunchConfig {
24434 grid_dim: (nblk, 1, 1),
24435 block_dim: (32, 1, 1),
24436 shared_mem_bytes: 0,
24437 };
24438 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
24439 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24440 let null: u64 = 0;
24441 let __s_b = self.gpu.stream();
24442 let mut b = __s_b.launch_builder(&f);
24443 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
24444 match base_dev {
24445 Some(base) => {
24446 b.arg(base);
24447 }
24448 None => {
24449 b.arg(&null);
24450 }
24451 }
24452 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
24453 unsafe {
24454 b.launch(cfg)?;
24455 }
24456 Ok(())
24457 }
24458
24459 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
24461 let f = self.func("inc_i32");
24462 let cfg = LaunchConfig {
24463 grid_dim: (1, 1, 1),
24464 block_dim: (1, 1, 1),
24465 shared_mem_bytes: 0,
24466 };
24467 let __s_b = self.gpu.stream();
24468 let mut b = __s_b.launch_builder(&f);
24469 b.arg(counter);
24470 unsafe {
24471 b.launch(cfg)?;
24472 }
24473 Ok(())
24474 }
24475
24476 #[allow(clippy::too_many_arguments)]
24485 fn fa_part_alloc(
24509 &self,
24510 o_len: usize,
24511 ml_len: usize,
24512 co: usize,
24513 cm: usize,
24514 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24515 static GROWS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
24516 let n = GROWS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
24517 if n < 64 {
24518 eprintln!(
24519 "[fa-pool] grow #{n} dev={} o_len {co} -> {o_len} ml_len {cm} -> {ml_len} (retired kept, zero={})",
24520 self.ctx().ordinal(),
24521 fa_part_zero_on()
24522 );
24523 }
24524 let mut po = self.alloc_uninit::<f32>(o_len)?;
24525 let mut pm = self.alloc_uninit::<f32>(ml_len)?;
24526 let mut pl = self.alloc_uninit::<f32>(ml_len)?;
24527 if fa_part_zero_on() {
24528 self.gpu.stream().memset_zeros(&mut po)?;
24529 self.gpu.stream().memset_zeros(&mut pm)?;
24530 self.gpu.stream().memset_zeros(&mut pl)?;
24531 }
24532 Ok((po, pm, pl))
24533 }
24534
24535 fn fa_part_pool_grow(
24536 &self,
24537 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
24538 o_len: usize,
24539 ml_len: usize,
24540 ) -> Result<(), Box<dyn std::error::Error>> {
24541 if part_guard
24542 .as_ref()
24543 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
24544 .unwrap_or(true)
24545 {
24546 let old = part_guard.take();
24547 let (co, cm) = old
24548 .as_ref()
24549 .map(|pp| (pp.0.len(), pp.1.len()))
24550 .unwrap_or((0, 0));
24551 if let Some(old) = old {
24552 self.fa_part_retired.lock().unwrap().push(old);
24553 }
24554 *part_guard =
24567 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
24568 }
24569 Ok(())
24570 }
24571
24572 pub fn fa_dcw_pool_ensure(
24575 &self,
24576 head_dim: usize,
24577 n_head: usize,
24578 n_head_kv: usize,
24579 bucket_max: usize,
24580 ) -> Result<(), Box<dyn std::error::Error>> {
24581 let sp = fa_split_keys(bucket_max, n_head_kv);
24582 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24583 let o_len = n_head * n_splits * head_dim;
24584 let ml_len = n_head * n_splits;
24585 let mut part_guard = self.fa_part_pool.lock().unwrap();
24586 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
24587 }
24588
24589 #[allow(clippy::too_many_arguments)]
24597 pub fn fa_decode_dcw2(
24598 &self,
24599 q2: &CudaSlice<f32>,
24600 k_ring: &cudarc::driver::CudaView<u8>,
24601 v_ring: &cudarc::driver::CudaView<u8>,
24602 o2: &mut CudaSlice<f32>,
24603 head_dim: usize,
24604 n_head: usize,
24605 n_head_kv: usize,
24606 len_dev: &CudaSlice<i32>,
24607 base_dev: Option<&CudaSlice<i32>>,
24608 window: usize,
24609 bucket_max: usize,
24610 scale: f32,
24611 k_tok_bytes: usize,
24612 v_tok_bytes: usize,
24613 gate2: &CudaSlice<f32>,
24614 ) -> Result<(), Box<dyn std::error::Error>> {
24615 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24616 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24617 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
24618 }
24619 let sp = fa_split_keys(bucket_max, n_head_kv);
24620 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24621 let o_len = 2 * n_head * n_splits * head_dim;
24623 let ml_len = 2 * n_head * n_splits;
24624 let mut part_guard = self.fa_part_pool.lock().unwrap();
24625 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24626 let pg = part_guard.as_mut().unwrap();
24627 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24628 let (hd, nh, nhkv, nsp) = (
24629 head_dim as i32,
24630 n_head as i32,
24631 n_head_kv as i32,
24632 n_splits as i32,
24633 );
24634 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24635 let (ski, win) = (sp as i32, window as i32);
24636 let gqa = (n_head / n_head_kv).max(1) as u32;
24637 let smem = (32 * head_dim * 2) as u32;
24638 let f = self.func("fa_decode_vec_q_v3_dcw2");
24639 let cfg = LaunchConfig {
24640 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24641 block_dim: (32, gqa, 1),
24642 shared_mem_bytes: smem,
24643 };
24644 let null: u64 = 0;
24645 {
24646 let __s_b = self.gpu.stream();
24647 let mut b = __s_b.launch_builder(&f);
24648 b.arg(q2)
24649 .arg(k_ring)
24650 .arg(v_ring)
24651 .arg(&mut *part_o)
24652 .arg(&mut *part_m)
24653 .arg(&mut *part_l)
24654 .arg(&hd)
24655 .arg(&nh)
24656 .arg(&nhkv)
24657 .arg(len_dev);
24658 match base_dev {
24659 Some(base) => {
24660 b.arg(base);
24661 }
24662 None => {
24663 b.arg(&null);
24664 }
24665 }
24666 b.arg(&win)
24667 .arg(&scale)
24668 .arg(&nsp)
24669 .arg(&ski)
24670 .arg(&ktb)
24671 .arg(&vtb);
24672 unsafe {
24673 b.launch(cfg)?;
24674 }
24675 }
24676 let fc = {
24680 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24681 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24682 self.func("fa_decode_combine_gate_f32_s")
24683 } else {
24684 self.func("fa_decode_combine_gate_f32")
24685 }
24686 };
24687 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24688 let nh2 = (2 * n_head) as i32;
24689 let cfg2 = LaunchConfig {
24690 grid_dim: ((2 * n_head) as u32, 1, 1),
24691 block_dim: (head_dim as u32, 1, 1),
24692 shared_mem_bytes: if combine_shared {
24693 (2 * n_splits * 4) as u32
24694 } else {
24695 0
24696 },
24697 };
24698 let __s_b2 = self.gpu.stream();
24699 let mut b2 = __s_b2.launch_builder(&fc);
24700 b2.arg(&*part_o)
24701 .arg(&*part_m)
24702 .arg(&*part_l)
24703 .arg(gate2)
24704 .arg(o2)
24705 .arg(&hd)
24706 .arg(&nh2)
24707 .arg(&nsp);
24708 unsafe {
24709 b2.launch(cfg2)?;
24710 }
24711 Ok(())
24712 }
24713
24714 #[allow(clippy::too_many_arguments)]
24723 pub fn fa_decode_dcw_rows(
24724 &self,
24725 q_rows: &CudaSlice<f32>,
24726 tab: &CudaSlice<u64>,
24727 o_rows: &mut CudaSlice<f32>,
24728 t: usize,
24729 head_dim: usize,
24730 n_head: usize,
24731 n_head_kv: usize,
24732 window: usize,
24733 max_ns: usize,
24734 scale: f32,
24735 k_tok_bytes: usize,
24736 v_tok_bytes: usize,
24737 gate_rows: &CudaSlice<f32>,
24738 ) -> Result<(), Box<dyn std::error::Error>> {
24739 if std::env::var("MEMRA_NO_FA_VEC").is_ok()
24740 || head_dim > 256
24741 || head_dim % 32 != 0
24742 || !fa_v3_on()
24743 {
24744 return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
24745 }
24746 if fa_sm_count() < 128
24747 || std::env::var("MEMRA_FA_SPLIT").is_ok()
24748 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24749 || std::env::var("MEMRA_FA_SP16").is_ok()
24750 {
24751 return Err(
24752 "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
24753 (or a <128-SM rig) keep the per-row path"
24754 .into(),
24755 );
24756 }
24757 if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
24758 return Err("fa_decode_dcw_rows geometry".into());
24759 }
24760 let o_len = t * n_head * max_ns * head_dim;
24761 let ml_len = t * n_head * max_ns;
24762 let mut part_guard = self.fa_part_pool.lock().unwrap();
24763 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24764 let pg = part_guard.as_mut().unwrap();
24765 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24766 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
24767 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24768 let (win, mns) = (window as i32, max_ns as i32);
24769 let gqa = (n_head / n_head_kv).max(1) as u32;
24770 let smem = (32 * head_dim * 2) as u32;
24771 let f = self.func("fa_decode_vec_q_v3_dcw_rows");
24772 let cfg = LaunchConfig {
24773 grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
24774 block_dim: (32, gqa, 1),
24775 shared_mem_bytes: smem,
24776 };
24777 {
24778 let __s_b = self.gpu.stream();
24779 let mut b = __s_b.launch_builder(&f);
24780 b.arg(q_rows)
24781 .arg(tab)
24782 .arg(&mut *part_o)
24783 .arg(&mut *part_m)
24784 .arg(&mut *part_l)
24785 .arg(&hd)
24786 .arg(&nh)
24787 .arg(&nhkv)
24788 .arg(&win)
24789 .arg(&scale)
24790 .arg(&mns)
24791 .arg(&ktb)
24792 .arg(&vtb);
24793 unsafe {
24794 b.launch(cfg)?;
24795 }
24796 }
24797 let fc = {
24801 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24802 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24803 self.func("fa_decode_combine_gate_f32_s")
24804 } else {
24805 self.func("fa_decode_combine_gate_f32")
24806 }
24807 };
24808 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24809 let nht = (t * n_head) as i32;
24810 let cfg2 = LaunchConfig {
24811 grid_dim: ((t * n_head) as u32, 1, 1),
24812 block_dim: (head_dim as u32, 1, 1),
24813 shared_mem_bytes: if combine_shared {
24814 (2 * max_ns * 4) as u32
24815 } else {
24816 0
24817 },
24818 };
24819 let __s_b2 = self.gpu.stream();
24820 let mut b2 = __s_b2.launch_builder(&fc);
24821 b2.arg(&*part_o)
24822 .arg(&*part_m)
24823 .arg(&*part_l)
24824 .arg(gate_rows)
24825 .arg(o_rows)
24826 .arg(&hd)
24827 .arg(&nht)
24828 .arg(&mns);
24829 unsafe {
24830 b2.launch(cfg2)?;
24831 }
24832 Ok(())
24833 }
24834
24835 pub fn fa_decode_dcw(
24836 &self,
24837 q: &CudaSlice<f32>,
24838 k_ring: &cudarc::driver::CudaView<u8>,
24839 v_ring: &cudarc::driver::CudaView<u8>,
24840 o: &mut CudaSlice<f32>,
24841 head_dim: usize,
24842 n_head: usize,
24843 n_head_kv: usize,
24844 len_dev: &CudaSlice<i32>,
24845 base_dev: Option<&CudaSlice<i32>>,
24846 window: usize,
24847 bucket_max: usize,
24848 scale: f32,
24849 k_tok_bytes: usize,
24850 v_tok_bytes: usize,
24851 fused_gate: Option<&CudaSlice<f32>>,
24855 ) -> Result<(), Box<dyn std::error::Error>> {
24856 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24857 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24858 return Err("fa_decode_dcw supports the default v3-vec class only (bucket >= vec floor, head_dim <= 256, MEMRA_FA_V3 on); keep eager outside it"
24859 .into());
24860 }
24861 let sp = fa_split_keys(bucket_max, n_head_kv);
24862 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24863 let o_len = n_head * n_splits * head_dim;
24864 let ml_len = n_head * n_splits;
24865 let mut part_guard = self.fa_part_pool.lock().unwrap();
24866 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24867 let pg = part_guard.as_mut().unwrap();
24868 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24873 let memset_on = *MEMSET_ON
24878 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
24879 || crate::tp::token_graph_building();
24880 if memset_on {
24881 self.gpu
24882 .stream()
24883 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24884 self.gpu
24885 .stream()
24886 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24887 self.gpu
24888 .stream()
24889 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24890 }
24891 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24892 let (hd, nh, nhkv, nsp) = (
24893 head_dim as i32,
24894 n_head as i32,
24895 n_head_kv as i32,
24896 n_splits as i32,
24897 );
24898 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24899 let (ski, win) = (sp as i32, window as i32);
24900 let gqa = (n_head / n_head_kv).max(1) as u32;
24901 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24905 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
24906 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
24907 Ok("2") => 2,
24908 Ok("1") => 1,
24909 _ => 0,
24910 });
24911 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24916 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
24917 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
24918 std::sync::Mutex::new(None);
24919 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24923 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
24924 && (n_head / n_head_kv) % 2 == 0
24925 && (n_head / n_head_kv) >= 2;
24926 let f = if fprof {
24927 self.func("fa_decode_vec_q_v3_dcw_prof")
24928 } else if hs2 {
24929 self.func("fa_decode_vec_q_v3_dcw_hs2")
24930 } else if hoist == 2 {
24931 self.func("fa_decode_vec_q_v3_dcw_hc")
24933 } else if hoist == 1 {
24934 self.func("fa_decode_vec_q_v3_dcw_h")
24936 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
24937 self.func("fa_decode_vec_q_v3_dcw_u8")
24938 } else {
24939 self.func("fa_decode_vec_q_v3_dcw")
24940 };
24941 let cfg = LaunchConfig {
24942 grid_dim: if hs2 {
24943 ((2 * n_head_kv) as u32, n_splits as u32, 1)
24944 } else {
24945 (n_head_kv as u32, n_splits as u32, 1)
24946 },
24947 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
24948 shared_mem_bytes: smem,
24949 };
24950 let null: u64 = 0;
24951 let __s_b = self.gpu.stream();
24952 let mut b = __s_b.launch_builder(&f);
24953 b.arg(q)
24954 .arg(k_ring)
24955 .arg(v_ring)
24956 .arg(&mut *part_o)
24957 .arg(&mut *part_m)
24958 .arg(&mut *part_l)
24959 .arg(&hd)
24960 .arg(&nh)
24961 .arg(&nhkv)
24962 .arg(len_dev);
24963 match base_dev {
24964 Some(base) => {
24965 b.arg(base);
24966 }
24967 None => {
24968 b.arg(&null);
24969 }
24970 }
24971 b.arg(&win)
24972 .arg(&scale)
24973 .arg(&nsp)
24974 .arg(&ski)
24975 .arg(&ktb)
24976 .arg(&vtb);
24977 if fprof {
24978 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
24979 if guard
24980 .as_ref()
24981 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
24982 {
24983 *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
24984 }
24985 let (_, buf) = guard.as_mut().expect("armed above");
24986 b.arg(&*buf);
24987 unsafe {
24988 b.launch(cfg)?;
24989 }
24990 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
24991 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
24992 if n % 430 == 0 {
24993 self.stream().synchronize()?;
24994 let h = self.dtoh_u64(buf)?;
24995 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
24996 let tot: u64 = h[..6].iter().sum();
24997 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
24998 for (i, name) in phases.iter().enumerate() {
24999 let pct = if tot > 0 {
25000 h[i] as f64 / tot as f64 * 100.0
25001 } else {
25002 0.0
25003 };
25004 line.push_str(&format!(" {name}={pct:.1}%"));
25005 }
25006 if h[6] > 0 {
25007 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
25008 }
25009 eprintln!("{line}");
25010 }
25011 } else {
25012 unsafe {
25013 b.launch(cfg)?;
25014 }
25015 }
25016 let mut combine_shared = false;
25017 let fc = if fused_gate.is_some() {
25018 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25021 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
25022 combine_shared = true;
25023 self.func("fa_decode_combine_gate_f32_s")
25024 } else {
25025 self.func("fa_decode_combine_gate_f32")
25026 }
25027 } else {
25028 self.fa_func("fa_decode_combine_f32", head_dim)
25029 };
25030 let cfg2 = LaunchConfig {
25031 grid_dim: (n_head as u32, 1, 1),
25032 block_dim: (head_dim as u32, 1, 1),
25033 shared_mem_bytes: if combine_shared {
25034 (2 * n_splits * 4) as u32
25035 } else {
25036 0
25037 },
25038 };
25039 let __s_b2 = self.gpu.stream();
25040 let mut b2 = __s_b2.launch_builder(&fc);
25041 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
25042 if let Some(gate_row) = fused_gate {
25043 b2.arg(gate_row);
25044 }
25045 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
25046 unsafe {
25047 b2.launch(cfg2)?;
25048 }
25049 Ok(())
25050 }
25051
25052 pub fn fa_geom_eager(
25058 &self,
25059 t_kv: usize,
25060 head_dim: usize,
25061 n_head_kv: usize,
25062 g: bool,
25063 ) -> (bool, usize) {
25064 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
25068 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
25074 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
25075 if g && head_dim == 256 && !fa_v4_at(t_kv) {
25081 fa_vec = false;
25082 }
25083 let sp = fa_split_keys(t_kv, n_head_kv);
25084 let n_splits = if fa_vec {
25085 ((t_kv + sp - 1) / sp).max(1)
25086 } else {
25087 ((t_kv + 255) / 256).max(1)
25088 };
25089 (fa_vec, n_splits)
25090 }
25091
25092 pub fn fa_bucket_key(
25098 &self,
25099 t_kv: usize,
25100 head_dim: usize,
25101 n_head_kv: usize,
25102 g: bool,
25103 ) -> (bool, usize) {
25104 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
25105 }
25106
25107 pub fn capture_graph_retained<F>(
25119 &self,
25120 step: F,
25121 ) -> Result<
25122 (
25123 cudarc::driver::CudaGraph,
25124 Vec<Box<dyn std::any::Any + Send>>,
25125 ),
25126 Box<dyn std::error::Error>,
25127 >
25128 where
25129 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
25130 {
25131 use cudarc::driver::sys::CUgraphInstantiate_flags;
25132 self.capture_graph_retained_flags(
25133 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
25134 step,
25135 )
25136 }
25137
25138 pub fn capture_graph_retained_flags<F>(
25143 &self,
25144 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
25145 mut step: F,
25146 ) -> Result<
25147 (
25148 cudarc::driver::CudaGraph,
25149 Vec<Box<dyn std::any::Any + Send>>,
25150 ),
25151 Box<dyn std::error::Error>,
25152 >
25153 where
25154 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
25155 {
25156 use cudarc::driver::sys::CUstreamCaptureMode;
25157 self.capture_keep.lock().unwrap().clear();
25165 let was_tracking = self.gpu.ctx.is_event_tracking();
25166 if was_tracking {
25167 unsafe {
25168 self.gpu.ctx.disable_event_tracking();
25169 }
25170 }
25171 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
25172 self.capture_keep_on
25173 .store(true, std::sync::atomic::Ordering::Relaxed);
25174 let w = (|| {
25175 step(self)?;
25176 step(self)
25177 })();
25178 self.capture_keep_on
25179 .store(false, std::sync::atomic::Ordering::Relaxed);
25180 w?;
25181 self.gpu.stream().synchronize()?;
25182 self.gpu
25183 .stream()
25184 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
25185 let r = step(self);
25186 let g = self.gpu.stream().end_capture(flags);
25187 r?;
25188 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
25189 graph.upload()?;
25190 Ok(graph)
25191 };
25192 let result = run();
25193 self.capture_keep_on
25194 .store(false, std::sync::atomic::Ordering::Relaxed);
25195 if was_tracking {
25196 unsafe {
25197 self.gpu.ctx.enable_event_tracking();
25198 }
25199 }
25200 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
25201 Ok((result?, keeper))
25202 }
25203
25204 pub fn capture_graph_retained_nowarm<F>(
25210 &self,
25211 mut step: F,
25212 ) -> Result<
25213 (
25214 cudarc::driver::CudaGraph,
25215 Vec<Box<dyn std::any::Any + Send>>,
25216 ),
25217 Box<dyn std::error::Error>,
25218 >
25219 where
25220 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
25221 {
25222 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
25223 let was_tracking = self.gpu.ctx.is_event_tracking();
25224 if was_tracking {
25225 unsafe {
25226 self.gpu.ctx.disable_event_tracking();
25227 }
25228 }
25229 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
25230 self.gpu.stream().synchronize()?;
25231 self.gpu
25232 .stream()
25233 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
25234 let r = step(self);
25235 let g = self.gpu.stream().end_capture(
25236 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
25237 );
25238 r?;
25239 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
25240 graph.upload()?;
25241 Ok(graph)
25242 };
25243 let result = run();
25244 if was_tracking {
25245 unsafe {
25246 self.gpu.ctx.enable_event_tracking();
25247 }
25248 }
25249 Ok((result?, Vec::new()))
25250 }
25251
25252 pub fn capture_graph<F>(
25253 &self,
25254 mut step: F,
25255 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
25256 where
25257 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
25258 {
25259 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
25260 let was_tracking = self.gpu.ctx.is_event_tracking();
25268 if was_tracking {
25269 unsafe {
25270 self.gpu.ctx.disable_event_tracking();
25271 }
25272 }
25273 let iflag = {
25280 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
25281 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
25282 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
25285 Ok("priority") => {
25286 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
25287 }
25288 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
25289 })
25290 };
25291 let ct = {
25298 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25299 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
25300 };
25301 let warmups = {
25324 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25325 *W.get_or_init(|| {
25326 std::env::var("MEMRA_GRAPH_WARMUPS")
25327 .ok()
25328 .and_then(|v| v.parse().ok())
25329 .filter(|n| *n >= 1)
25330 .unwrap_or(1)
25331 })
25332 };
25333 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
25334 let t_w = std::time::Instant::now();
25335 for _ in 0..warmups {
25337 step(self)?;
25338 }
25339 self.gpu.stream().synchronize()?;
25340 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
25341 let t_c = std::time::Instant::now();
25343 self.gpu
25344 .stream()
25345 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
25346 let r = step(self);
25349 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
25350 let t_i = std::time::Instant::now();
25351 let g = self.gpu.stream().end_capture(iflag);
25352 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
25353 r?;
25354 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
25355 let t_u = std::time::Instant::now();
25356 graph.upload()?;
25357 if ct {
25358 println!(
25359 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
25360 instantiate {ms_inst:.2} ms upload {:.2} ms",
25361 t_u.elapsed().as_secs_f64() * 1e3
25362 );
25363 }
25364 Ok(graph)
25365 };
25366 let result = run();
25367 if was_tracking {
25368 unsafe {
25369 self.gpu.ctx.enable_event_tracking();
25370 }
25371 }
25372 result
25373 }
25374
25375 pub fn gdn_scan_s128_view(
25377 &self,
25378 q: &CudaSlice<f32>,
25379 k: &CudaSlice<f32>,
25380 v: &CudaSlice<f32>,
25381 g: &CudaSlice<f32>,
25382 beta: &CudaSlice<f32>,
25383 state_in: &cudarc::driver::CudaView<f32>,
25384 state_out: &mut cudarc::driver::CudaViewMut<f32>,
25385 o: &mut CudaSlice<f32>,
25386 n_head: usize,
25387 t: usize,
25388 scale: f32,
25389 ) -> Result<(), Box<dyn std::error::Error>> {
25390 let f = self.func("gdn_scan_s128");
25391 const S_V: u32 = 128;
25392 const WARP: u32 = 32;
25393 const COLS: u32 = 4;
25394 let cfg = LaunchConfig {
25395 grid_dim: (n_head as u32, 1, S_V / COLS),
25396 block_dim: (WARP, COLS, 1),
25397 shared_mem_bytes: 0,
25398 };
25399 let (h, ti) = (n_head as i32, t as i32);
25400 let __s_b = self.gpu.stream();
25401 let mut b = __s_b.launch_builder(&f);
25402 b.arg(q)
25403 .arg(k)
25404 .arg(v)
25405 .arg(g)
25406 .arg(beta)
25407 .arg(state_in)
25408 .arg(state_out)
25409 .arg(o)
25410 .arg(&h)
25411 .arg(&ti)
25412 .arg(&scale);
25413 unsafe {
25414 b.launch(cfg)?;
25415 }
25416 Ok(())
25417 }
25418
25419 pub fn ssm_conv1d_view(
25421 &self,
25422 x: &cudarc::driver::CudaView<f32>,
25423 w: &CudaSlice<f32>,
25424 y: &mut CudaSlice<f32>,
25425 conv_dim: usize,
25426 t: usize,
25427 d_conv: usize,
25428 silu: bool,
25429 ) -> Result<(), Box<dyn std::error::Error>> {
25430 let f = self.func("ssm_conv1d_silu_f32");
25431 let cfg = LaunchConfig {
25433 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25434 block_dim: (256, 1, 1),
25435 shared_mem_bytes: 0,
25436 };
25437 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25438 let __s_b = self.gpu.stream();
25439 let mut b = __s_b.launch_builder(&f);
25440 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25441 unsafe {
25442 b.launch(cfg)?;
25443 }
25444 Ok(())
25445 }
25446
25447 pub fn ssm_conv1d_tm(
25454 &self,
25455 qkv_tm: &CudaSlice<f32>,
25456 w: &CudaSlice<f32>,
25457 y: &mut CudaSlice<f32>,
25458 conv_dim: usize,
25459 t: usize,
25460 d_conv: usize,
25461 ) -> Result<(), Box<dyn std::error::Error>> {
25462 let f = self.func("ssm_conv1d_tm_f32");
25463 let cfg = LaunchConfig {
25464 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25465 block_dim: (256, 1, 1),
25466 shared_mem_bytes: 0,
25467 };
25468 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25469 let __s_b = self.gpu.stream();
25470 let mut b = __s_b.launch_builder(&f);
25471 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
25472 unsafe {
25473 b.launch(cfg)?;
25474 }
25475 Ok(())
25476 }
25477
25478 pub fn ssm_conv1d_tm_state(
25486 &self,
25487 qkv_tm: &CudaSlice<f32>,
25488 conv_state: &mut CudaSlice<f32>,
25489 w: &CudaSlice<f32>,
25490 y: &mut CudaSlice<f32>,
25491 conv_dim: usize,
25492 t: usize,
25493 d_conv: usize,
25494 ) -> Result<(), Box<dyn std::error::Error>> {
25495 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
25496 }
25497
25498 #[allow(clippy::too_many_arguments)]
25501 pub fn ssm_conv1d_tm_state_pad(
25502 &self,
25503 qkv_tm: &CudaSlice<f32>,
25504 conv_state: &mut CudaSlice<f32>,
25505 w: &CudaSlice<f32>,
25506 y: &mut CudaSlice<f32>,
25507 conv_dim: usize,
25508 t: usize,
25509 d_conv: usize,
25510 pad_len: Option<&CudaSlice<i32>>,
25511 ) -> Result<(), Box<dyn std::error::Error>> {
25512 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
25513 let ring_old = if t < d_conv - 1 {
25517 Some(self.clone_dtod(conv_state)?)
25518 } else {
25519 None
25520 };
25521 {
25522 let f = self.func("ssm_conv1d_tm_state_f32");
25523 let cfg = LaunchConfig {
25524 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25525 block_dim: (256, 1, 1),
25526 shared_mem_bytes: 0,
25527 };
25528 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25529 let __s_b = self.gpu.stream();
25530 let mut b = __s_b.launch_builder(&f);
25531 b.arg(qkv_tm)
25532 .arg(&*conv_state)
25533 .arg(w)
25534 .arg(y)
25535 .arg(&cd)
25536 .arg(&ti)
25537 .arg(&dc);
25538 unsafe {
25539 b.launch(cfg)?;
25540 }
25541 }
25542 match (ring_old, pad_len) {
25543 (None, Some(len_d)) => {
25544 let f = self.func("ssm_conv_ring_update_dev_f32");
25545 let n = conv_dim * (d_conv - 1);
25546 let cfg = LaunchConfig::for_num_elems(n as u32);
25547 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25548 let __s_b = self.gpu.stream();
25549 let mut b = __s_b.launch_builder(&f);
25550 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25551 unsafe {
25552 b.launch(cfg)?;
25553 }
25554 }
25555 (None, None) => {
25556 let f = self.func("ssm_conv_ring_update_f32");
25557 let n = conv_dim * (d_conv - 1);
25558 let cfg = LaunchConfig::for_num_elems(n as u32);
25559 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25560 let __s_b = self.gpu.stream();
25561 let mut b = __s_b.launch_builder(&f);
25562 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25563 unsafe {
25564 b.launch(cfg)?;
25565 }
25566 }
25567 (Some(old), _) => {
25568 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
25569 }
25570 }
25571 Ok(())
25572 }
25573
25574 pub fn ssm_conv1d_tm_state_pad_v(
25576 &self,
25577 qkv_tm: &cudarc::driver::CudaView<f32>,
25578 conv_state: &mut CudaSlice<f32>,
25579 w: &CudaSlice<f32>,
25580 y: &mut CudaSlice<f32>,
25581 conv_dim: usize,
25582 t: usize,
25583 d_conv: usize,
25584 pad_len: Option<&CudaSlice<i32>>,
25585 ) -> Result<(), Box<dyn std::error::Error>> {
25586 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
25587 let ring_old = if t < d_conv - 1 {
25591 Some(self.clone_dtod(conv_state)?)
25592 } else {
25593 None
25594 };
25595 {
25596 let f = self.func("ssm_conv1d_tm_state_f32");
25597 let cfg = LaunchConfig {
25598 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25599 block_dim: (256, 1, 1),
25600 shared_mem_bytes: 0,
25601 };
25602 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25603 let __s_b = self.gpu.stream();
25604 let mut b = __s_b.launch_builder(&f);
25605 b.arg(qkv_tm)
25606 .arg(&*conv_state)
25607 .arg(w)
25608 .arg(y)
25609 .arg(&cd)
25610 .arg(&ti)
25611 .arg(&dc);
25612 unsafe {
25613 b.launch(cfg)?;
25614 }
25615 }
25616 match (ring_old, pad_len) {
25617 (None, Some(len_d)) => {
25618 let f = self.func("ssm_conv_ring_update_dev_f32");
25619 let n = conv_dim * (d_conv - 1);
25620 let cfg = LaunchConfig::for_num_elems(n as u32);
25621 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25622 let __s_b = self.gpu.stream();
25623 let mut b = __s_b.launch_builder(&f);
25624 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25625 unsafe {
25626 b.launch(cfg)?;
25627 }
25628 }
25629 (None, None) => {
25630 let f = self.func("ssm_conv_ring_update_f32");
25631 let n = conv_dim * (d_conv - 1);
25632 let cfg = LaunchConfig::for_num_elems(n as u32);
25633 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25634 let __s_b = self.gpu.stream();
25635 let mut b = __s_b.launch_builder(&f);
25636 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25637 unsafe {
25638 b.launch(cfg)?;
25639 }
25640 }
25641 (Some(_), _) => unreachable!(
25642 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
25643 ),
25644 }
25645 Ok(())
25646 }
25647
25648 pub fn ssm_conv_ring_rebuild(
25653 &self,
25654 qkv_tm: &CudaSlice<f32>,
25655 ring_old: &CudaSlice<f32>,
25656 conv_state: &mut CudaSlice<f32>,
25657 conv_dim: usize,
25658 tc: usize,
25659 d_conv: usize,
25660 ) -> Result<(), Box<dyn std::error::Error>> {
25661 let f = self.func("ssm_conv_ring_rebuild_f32");
25662 let n = conv_dim * (d_conv - 1);
25663 let cfg = LaunchConfig::for_num_elems(n as u32);
25664 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
25665 let __s_b = self.gpu.stream();
25666 let mut b = __s_b.launch_builder(&f);
25667 b.arg(qkv_tm)
25668 .arg(ring_old)
25669 .arg(conv_state)
25670 .arg(&cd)
25671 .arg(&ti)
25672 .arg(&dc);
25673 unsafe {
25674 b.launch(cfg)?;
25675 }
25676 Ok(())
25677 }
25678
25679 #[allow(clippy::too_many_arguments)]
25684 pub fn gdn_prep_decode(
25685 &self,
25686 conv_out: &CudaSlice<f32>,
25687 beta_raw: &CudaSlice<f32>,
25688 alpha: &CudaSlice<f32>,
25689 dt_bias: &CudaSlice<f32>,
25690 a: &CudaSlice<f32>,
25691 q_l2: &mut CudaSlice<f32>,
25692 k_l2: &mut CudaSlice<f32>,
25693 v_g: &mut CudaSlice<f32>,
25694 beta: &mut CudaSlice<f32>,
25695 g_log: &mut CudaSlice<f32>,
25696 d_state: usize,
25697 num_v: usize,
25698 num_k: usize,
25699 key_dim: usize,
25700 eps: f32,
25701 ) -> Result<(), Box<dyn std::error::Error>> {
25702 let f = self.func("gdn_prep_decode_f32");
25703 let cfg = LaunchConfig {
25704 grid_dim: (num_v as u32, 1, 1),
25705 block_dim: (32, 4, 1),
25706 shared_mem_bytes: 0,
25707 };
25708 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25709 let __s_b = self.gpu.stream();
25710 let mut b = __s_b.launch_builder(&f);
25711 b.arg(conv_out)
25712 .arg(beta_raw)
25713 .arg(alpha)
25714 .arg(dt_bias)
25715 .arg(a)
25716 .arg(q_l2)
25717 .arg(k_l2)
25718 .arg(v_g)
25719 .arg(beta)
25720 .arg(g_log)
25721 .arg(&ds)
25722 .arg(&nv)
25723 .arg(&nk)
25724 .arg(&kd)
25725 .arg(&eps);
25726 unsafe {
25727 b.launch(cfg)?;
25728 }
25729 Ok(())
25730 }
25731
25732 #[allow(clippy::too_many_arguments)]
25736 pub fn ssm_conv1d_gdn(
25737 &self,
25738 qkv_tm: &CudaSlice<f32>,
25739 w: &CudaSlice<f32>,
25740 q_g: &mut CudaSlice<f32>,
25741 k_g: &mut CudaSlice<f32>,
25742 v_g: &mut CudaSlice<f32>,
25743 conv_dim: usize,
25744 t: usize,
25745 d_conv: usize,
25746 d_state: usize,
25747 num_v: usize,
25748 num_k: usize,
25749 key_dim: usize,
25750 ) -> Result<(), Box<dyn std::error::Error>> {
25751 let f = self.func("ssm_conv1d_gdn_f32");
25752 let cfg = LaunchConfig {
25753 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25754 block_dim: (256, 1, 1),
25755 shared_mem_bytes: 0,
25756 };
25757 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25758 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25759 let __s_b = self.gpu.stream();
25760 let mut b = __s_b.launch_builder(&f);
25761 b.arg(qkv_tm)
25762 .arg(w)
25763 .arg(q_g)
25764 .arg(k_g)
25765 .arg(v_g)
25766 .arg(&cd)
25767 .arg(&ti)
25768 .arg(&dc)
25769 .arg(&ds)
25770 .arg(&nv)
25771 .arg(&nk)
25772 .arg(&kd);
25773 unsafe {
25774 b.launch(cfg)?;
25775 }
25776 Ok(())
25777 }
25778
25779 pub fn ssm_conv1d(
25780 &self,
25781 x: &CudaSlice<f32>,
25782 w: &CudaSlice<f32>,
25783 y: &mut CudaSlice<f32>,
25784 conv_dim: usize,
25785 t: usize,
25786 d_conv: usize,
25787 silu: bool,
25788 ) -> Result<(), Box<dyn std::error::Error>> {
25789 let f = self.func("ssm_conv1d_silu_f32");
25790 let cfg = LaunchConfig {
25791 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25792 block_dim: (256, 1, 1),
25793 shared_mem_bytes: 0,
25794 };
25795 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25796 let __s_b = self.gpu.stream();
25797 let mut b = __s_b.launch_builder(&f);
25798 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25799 unsafe {
25800 b.launch(cfg)?;
25801 }
25802 Ok(())
25803 }
25804
25805 pub fn gdn_scan_s128(
25808 &self,
25809 q: &CudaSlice<f32>,
25810 k: &CudaSlice<f32>,
25811 v: &CudaSlice<f32>,
25812 g: &CudaSlice<f32>,
25813 beta: &CudaSlice<f32>,
25814 state_in: &CudaSlice<f32>,
25815 state_out: &mut CudaSlice<f32>,
25816 o: &mut CudaSlice<f32>,
25817 n_head: usize,
25818 t: usize,
25819 scale: f32,
25820 ) -> Result<(), Box<dyn std::error::Error>> {
25821 let f = self.func("gdn_scan_s128");
25822 const S_V: u32 = 128;
25823 const WARP: u32 = 32;
25824 const COLS_PER_BLOCK: u32 = 4;
25825 let cfg = LaunchConfig {
25826 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
25827 block_dim: (WARP, COLS_PER_BLOCK, 1),
25828 shared_mem_bytes: 0,
25829 };
25830 let (h, ti) = (n_head as i32, t as i32);
25831 let __s_b = self.gpu.stream();
25832 let mut b = __s_b.launch_builder(&f);
25833 b.arg(q)
25834 .arg(k)
25835 .arg(v)
25836 .arg(g)
25837 .arg(beta)
25838 .arg(state_in)
25839 .arg(state_out)
25840 .arg(o)
25841 .arg(&h)
25842 .arg(&ti)
25843 .arg(&scale);
25844 unsafe {
25845 b.launch(cfg)?;
25846 }
25847 Ok(())
25848 }
25849
25850 #[allow(clippy::too_many_arguments)]
25855 pub fn ssm_conv1d_fused_decode_b(
25856 &self,
25857 qkv_cols: &CudaSlice<f32>,
25858 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25859 w: &CudaSlice<f32>,
25860 conv_outs: &mut CudaSlice<f32>,
25861 conv_dim: usize,
25862 d_conv: usize,
25863 b_n: usize,
25864 ) -> Result<(), Box<dyn std::error::Error>> {
25865 let f = self.func("ssm_conv1d_fused_decode_b_f32");
25866 let cfg = LaunchConfig {
25867 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25868 block_dim: (256, 1, 1),
25869 shared_mem_bytes: 0,
25870 };
25871 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25872 let __s_b = self.gpu.stream();
25873 let mut b = __s_b.launch_builder(&f);
25874 b.arg(qkv_cols)
25875 .arg(conv_state_ptrs)
25876 .arg(w)
25877 .arg(conv_outs)
25878 .arg(&cd)
25879 .arg(&dc);
25880 unsafe {
25881 b.launch(cfg)?;
25882 }
25883 Ok(())
25884 }
25885
25886 #[allow(clippy::too_many_arguments)]
25887 pub fn gdn_prep_decode_b(
25888 &self,
25889 conv_outs: &CudaSlice<f32>,
25890 beta_raws: &CudaSlice<f32>,
25891 alphas: &CudaSlice<f32>,
25892 dt_bias: &CudaSlice<f32>,
25893 a: &CudaSlice<f32>,
25894 q_l2: &mut CudaSlice<f32>,
25895 k_l2: &mut CudaSlice<f32>,
25896 v_g: &mut CudaSlice<f32>,
25897 beta: &mut CudaSlice<f32>,
25898 g_log: &mut CudaSlice<f32>,
25899 d_state: usize,
25900 num_v: usize,
25901 num_k: usize,
25902 key_dim: usize,
25903 eps: f32,
25904 conv_dim: usize,
25905 b_n: usize,
25906 ) -> Result<(), Box<dyn std::error::Error>> {
25907 let f = self.func("gdn_prep_decode_b_f32");
25908 let cfg = LaunchConfig {
25909 grid_dim: (num_v as u32, 1, b_n as u32),
25910 block_dim: (32, 4, 1),
25911 shared_mem_bytes: 0,
25912 };
25913 let (ds, nv, nk, kd, cd) = (
25914 d_state as i32,
25915 num_v as i32,
25916 num_k as i32,
25917 key_dim as i32,
25918 conv_dim as i32,
25919 );
25920 let __s_b = self.gpu.stream();
25921 let mut b = __s_b.launch_builder(&f);
25922 b.arg(conv_outs)
25923 .arg(beta_raws)
25924 .arg(alphas)
25925 .arg(dt_bias)
25926 .arg(a)
25927 .arg(q_l2)
25928 .arg(k_l2)
25929 .arg(v_g)
25930 .arg(beta)
25931 .arg(g_log)
25932 .arg(&ds)
25933 .arg(&nv)
25934 .arg(&nk)
25935 .arg(&kd)
25936 .arg(&eps)
25937 .arg(&cd);
25938 unsafe {
25939 b.launch(cfg)?;
25940 }
25941 Ok(())
25942 }
25943
25944 #[allow(clippy::too_many_arguments)]
25945 pub fn gdn_scan_s128_batched(
25946 &self,
25947 q: &CudaSlice<f32>,
25948 k: &CudaSlice<f32>,
25949 v: &CudaSlice<f32>,
25950 g: &CudaSlice<f32>,
25951 beta: &CudaSlice<f32>,
25952 state_in_ptrs: &cudarc::driver::CudaView<u64>,
25953 state_out_ptrs: &cudarc::driver::CudaView<u64>,
25954 o: &mut CudaSlice<f32>,
25955 n_head: usize,
25956 b_n: usize,
25957 scale: f32,
25958 ) -> Result<(), Box<dyn std::error::Error>> {
25959 let f = self.func("gdn_scan_s128_b");
25960 const S_V: u32 = 128;
25961 const WARP: u32 = 32;
25962 const COLS_PER_BLOCK: u32 = 4;
25963 let cfg = LaunchConfig {
25964 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25965 block_dim: (WARP, COLS_PER_BLOCK, 1),
25966 shared_mem_bytes: 0,
25967 };
25968 let h = n_head as i32;
25969 let __s_b = self.gpu.stream();
25970 let mut b = __s_b.launch_builder(&f);
25971 b.arg(q)
25972 .arg(k)
25973 .arg(v)
25974 .arg(g)
25975 .arg(beta)
25976 .arg(state_in_ptrs)
25977 .arg(state_out_ptrs)
25978 .arg(o)
25979 .arg(&h)
25980 .arg(&scale);
25981 unsafe {
25982 b.launch(cfg)?;
25983 }
25984 Ok(())
25985 }
25986
25987 #[allow(clippy::too_many_arguments)]
25993 pub fn ssm_conv1d_fused_decode_b_view(
25994 &self,
25995 qkv_cols: &cudarc::driver::CudaView<f32>,
25996 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25997 w: &CudaSlice<f32>,
25998 conv_outs: &mut CudaSlice<f32>,
25999 conv_dim: usize,
26000 d_conv: usize,
26001 b_n: usize,
26002 ) -> Result<(), Box<dyn std::error::Error>> {
26003 let f = self.func("ssm_conv1d_fused_decode_b_f32");
26004 let cfg = LaunchConfig {
26005 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
26006 block_dim: (256, 1, 1),
26007 shared_mem_bytes: 0,
26008 };
26009 let (cd, dc) = (conv_dim as i32, d_conv as i32);
26010 let __s_b = self.gpu.stream();
26011 let mut b = __s_b.launch_builder(&f);
26012 b.arg(qkv_cols)
26013 .arg(conv_state_ptrs)
26014 .arg(w)
26015 .arg(conv_outs)
26016 .arg(&cd)
26017 .arg(&dc);
26018 unsafe {
26019 b.launch(cfg)?;
26020 }
26021 Ok(())
26022 }
26023
26024 #[allow(clippy::too_many_arguments)]
26025 pub fn gdn_prep_decode_b_view(
26026 &self,
26027 conv_outs: &CudaSlice<f32>,
26028 beta_raws: &cudarc::driver::CudaView<f32>,
26029 alphas: &cudarc::driver::CudaView<f32>,
26030 dt_bias: &CudaSlice<f32>,
26031 a: &CudaSlice<f32>,
26032 q_l2: &mut CudaSlice<f32>,
26033 k_l2: &mut CudaSlice<f32>,
26034 v_g: &mut CudaSlice<f32>,
26035 beta: &mut CudaSlice<f32>,
26036 g_log: &mut CudaSlice<f32>,
26037 d_state: usize,
26038 num_v: usize,
26039 num_k: usize,
26040 key_dim: usize,
26041 eps: f32,
26042 conv_dim: usize,
26043 b_n: usize,
26044 ) -> Result<(), Box<dyn std::error::Error>> {
26045 let f = self.func("gdn_prep_decode_b_f32");
26046 let cfg = LaunchConfig {
26047 grid_dim: (num_v as u32, 1, b_n as u32),
26048 block_dim: (32, 4, 1),
26049 shared_mem_bytes: 0,
26050 };
26051 let (ds, nv, nk, kd, cd) = (
26052 d_state as i32,
26053 num_v as i32,
26054 num_k as i32,
26055 key_dim as i32,
26056 conv_dim as i32,
26057 );
26058 let __s_b = self.gpu.stream();
26059 let mut b = __s_b.launch_builder(&f);
26060 b.arg(conv_outs)
26061 .arg(beta_raws)
26062 .arg(alphas)
26063 .arg(dt_bias)
26064 .arg(a)
26065 .arg(q_l2)
26066 .arg(k_l2)
26067 .arg(v_g)
26068 .arg(beta)
26069 .arg(g_log)
26070 .arg(&ds)
26071 .arg(&nv)
26072 .arg(&nk)
26073 .arg(&kd)
26074 .arg(&eps)
26075 .arg(&cd);
26076 unsafe {
26077 b.launch(cfg)?;
26078 }
26079 Ok(())
26080 }
26081
26082 #[allow(clippy::too_many_arguments)]
26083 pub fn gdn_scan_s128_batched_view(
26084 &self,
26085 q: &CudaSlice<f32>,
26086 k: &CudaSlice<f32>,
26087 v: &CudaSlice<f32>,
26088 g: &CudaSlice<f32>,
26089 beta: &CudaSlice<f32>,
26090 state_in_ptrs: &cudarc::driver::CudaView<u64>,
26091 state_out_ptrs: &cudarc::driver::CudaView<u64>,
26092 o: &mut cudarc::driver::CudaViewMut<f32>,
26093 n_head: usize,
26094 b_n: usize,
26095 scale: f32,
26096 ) -> Result<(), Box<dyn std::error::Error>> {
26097 let f = self.func("gdn_scan_s128_b");
26098 const S_V: u32 = 128;
26099 const WARP: u32 = 32;
26100 const COLS_PER_BLOCK: u32 = 4;
26101 let cfg = LaunchConfig {
26102 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
26103 block_dim: (WARP, COLS_PER_BLOCK, 1),
26104 shared_mem_bytes: 0,
26105 };
26106 let h = n_head as i32;
26107 let __s_b = self.gpu.stream();
26108 let mut b = __s_b.launch_builder(&f);
26109 b.arg(q)
26110 .arg(k)
26111 .arg(v)
26112 .arg(g)
26113 .arg(beta)
26114 .arg(state_in_ptrs)
26115 .arg(state_out_ptrs)
26116 .arg(o)
26117 .arg(&h)
26118 .arg(&scale);
26119 unsafe {
26120 b.launch(cfg)?;
26121 }
26122 Ok(())
26123 }
26124
26125 pub fn gdn_chunked_enabled() -> bool {
26134 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26135 *E.get_or_init(|| {
26136 std::env::var("MEMRA_GDN_CHUNKED")
26137 .map(|v| v != "0")
26138 .unwrap_or(true)
26139 })
26140 }
26141
26142 pub fn gdn_chunk_size() -> usize {
26147 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26148 *C.get_or_init(|| {
26149 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
26150 .ok()
26151 .and_then(|v| v.parse().ok())
26152 .unwrap_or(32);
26153 c.clamp(32, 128) / 32 * 32
26154 })
26155 }
26156
26157 #[allow(clippy::too_many_arguments)]
26162 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
26165 #[allow(clippy::too_many_arguments)]
26166 pub fn gdn_chunk_k123(
26167 &self,
26168 q: &CudaSlice<f32>,
26169 k: &CudaSlice<f32>,
26170 v: &CudaSlice<f32>,
26171 g: &CudaSlice<f32>,
26172 beta: &CudaSlice<f32>,
26173 wb16: Option<&mut CudaSlice<u8>>,
26174 n_head: usize,
26175 t: usize,
26176 c: usize,
26177 hk: usize,
26178 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
26179 ) -> Result<
26180 (
26181 CudaSlice<f32>,
26182 CudaSlice<f32>,
26183 CudaSlice<f32>,
26184 CudaSlice<f32>,
26185 ),
26186 Box<dyn std::error::Error>,
26187 > {
26188 const D: usize = 128;
26189 let h = n_head;
26190 let nc = (t + c - 1) / c;
26191 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
26192 let mut gcum = self.uninit(t * h)?;
26193 let mut a = self.uninit(nc * h * c * c)?;
26194 let mut p = self.uninit(nc * h * c * c)?;
26195 let mut u = self.uninit(nc * h * c * D)?;
26196 let mut w = self.uninit(nc * h * c * D)?;
26197 {
26198 let f = self.func("gdn_chunk_cumgate_f32");
26200 let cfg = LaunchConfig {
26201 grid_dim: (nc as u32, h as u32, 1),
26202 block_dim: (32, 1, 1),
26203 shared_mem_bytes: 0,
26204 };
26205 let __s_b = self.gpu.stream();
26206 let mut b = __s_b.launch_builder(&f);
26207 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
26208 unsafe {
26209 b.launch(cfg)?;
26210 }
26211 }
26212 if let Some((qb, kb, pb)) = k2w {
26213 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
26216 let f = self.func("gdn_k2_wgmma");
26217 let cfg = LaunchConfig {
26218 grid_dim: (nc as u32, h as u32, 1),
26219 block_dim: (128, 1, 1),
26220 shared_mem_bytes: 0,
26221 };
26222 let hki = hk as i32;
26223 let __s_b = self.gpu.stream();
26224 let mut b = __s_b.launch_builder(&f);
26225 b.arg(qb)
26226 .arg(kb)
26227 .arg(&gcum)
26228 .arg(beta)
26229 .arg(&mut a)
26230 .arg(&mut *pb)
26231 .arg(&hi)
26232 .arg(&ti)
26233 .arg(&ci)
26234 .arg(&hki);
26235 unsafe {
26236 b.launch(cfg)?;
26237 }
26238 } else if c <= 64 && !portable_mma_gated() {
26239 let f = self.func("gdn_chunk_attn_f32");
26241 f.set_attribute(
26242 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26243 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
26244 )?;
26245 let jt = ((c + 31) / 32) as u32;
26246 let cfg = LaunchConfig {
26247 grid_dim: (nc as u32, h as u32, jt),
26248 block_dim: (256, 1, 1),
26249 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
26250 };
26251 let hki = hk as i32;
26252 let __s_b = self.gpu.stream();
26253 let mut b = __s_b.launch_builder(&f);
26254 b.arg(q)
26255 .arg(k)
26256 .arg(&gcum)
26257 .arg(beta)
26258 .arg(&mut a)
26259 .arg(&mut p)
26260 .arg(&hi)
26261 .arg(&ti)
26262 .arg(&ci)
26263 .arg(&hki);
26264 unsafe {
26265 b.launch(cfg)?;
26266 }
26267 } else {
26268 assert!(
26270 hk == h,
26271 "generic K2 is broadcast-only (de-broadcast rides C==32)"
26272 );
26273 let f = self.func("gdn_chunk_attn_g_f32");
26274 let cfg = LaunchConfig {
26275 grid_dim: (nc as u32, h as u32, 1),
26276 block_dim: (32, 8, 1),
26277 shared_mem_bytes: 0,
26278 };
26279 let __s_b = self.gpu.stream();
26280 let mut b = __s_b.launch_builder(&f);
26281 b.arg(q)
26282 .arg(k)
26283 .arg(&gcum)
26284 .arg(beta)
26285 .arg(&mut a)
26286 .arg(&mut p)
26287 .arg(&hi)
26288 .arg(&ti)
26289 .arg(&ci);
26290 unsafe {
26291 b.launch(cfg)?;
26292 }
26293 }
26294 {
26295 let cfg = LaunchConfig {
26297 grid_dim: (nc as u32, h as u32, 1),
26298 block_dim: (256, 1, 1),
26299 shared_mem_bytes: 0,
26300 };
26301 match c {
26302 32 | 64 => {
26303 let f = self.func(if c == 32 {
26304 "gdn_chunk_solve32_f32"
26305 } else {
26306 "gdn_chunk_solve64_f32"
26307 });
26308 let wb: u64 = match wb16 {
26310 Some(d) => self.addr_u8(d),
26311 None => 0,
26312 };
26313 let hki = hk as i32;
26314 let __s_b = self.gpu.stream();
26315 let mut b = __s_b.launch_builder(&f);
26316 b.arg(v)
26317 .arg(k)
26318 .arg(&a)
26319 .arg(&gcum)
26320 .arg(&mut u)
26321 .arg(&mut w)
26322 .arg(&wb)
26323 .arg(&hi)
26324 .arg(&ti)
26325 .arg(&hki);
26326 unsafe {
26327 b.launch(cfg)?;
26328 }
26329 }
26330 _ => {
26331 assert!(hk == h, "generic K3 is broadcast-only");
26332 let f = self.func("gdn_chunk_solve_f32");
26333 let __s_b = self.gpu.stream();
26334 let mut b = __s_b.launch_builder(&f);
26335 b.arg(v)
26336 .arg(k)
26337 .arg(&a)
26338 .arg(&gcum)
26339 .arg(&mut u)
26340 .arg(&mut w)
26341 .arg(&hi)
26342 .arg(&ti)
26343 .arg(&ci);
26344 unsafe {
26345 b.launch(cfg)?;
26346 }
26347 }
26348 }
26349 }
26350 Ok((gcum, p, u, w))
26351 }
26352
26353 pub fn gdn_db_on() -> bool {
26357 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
26358 }
26359
26360 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
26369 !portable_mma_gated()
26370 && c == 32
26371 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26372 Ok("1") => true,
26373 Ok("0") => false,
26374 _ => gdn_mma_default_on(),
26375 }
26376 }
26377
26378 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
26385 cfg!(memra_hopper_mma)
26386 && self.gdn_mma_enabled(c)
26387 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
26388 }
26389
26390 #[allow(clippy::too_many_arguments)]
26395 pub fn ssm_conv1d_gdn_state_pad(
26396 &self,
26397 qkv_tm: &cudarc::driver::CudaView<f32>,
26398 conv_state: &mut CudaSlice<f32>,
26399 w: &CudaSlice<f32>,
26400 q_g: &mut CudaSlice<f32>,
26401 k_g: &mut CudaSlice<f32>,
26402 v_g: &mut CudaSlice<f32>,
26403 conv_dim: usize,
26404 t: usize,
26405 d_conv: usize,
26406 d_state: usize,
26407 num_v: usize,
26408 num_k: usize,
26409 key_dim: usize,
26410 hk: usize,
26411 pad_len: Option<&CudaSlice<i32>>,
26412 ) -> Result<(), Box<dyn std::error::Error>> {
26413 assert!(
26414 t >= d_conv - 1,
26415 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
26416 );
26417 {
26418 let f = self.func("ssm_conv1d_gdn_state_f32");
26419 let cfg = LaunchConfig {
26420 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
26421 block_dim: (256, 1, 1),
26422 shared_mem_bytes: 0,
26423 };
26424 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
26425 let (ds, nv, nk, kd, hki) = (
26426 d_state as i32,
26427 num_v as i32,
26428 num_k as i32,
26429 key_dim as i32,
26430 hk as i32,
26431 );
26432 let __s_b = self.gpu.stream();
26433 let mut b = __s_b.launch_builder(&f);
26434 b.arg(qkv_tm)
26435 .arg(&*conv_state)
26436 .arg(w)
26437 .arg(q_g)
26438 .arg(k_g)
26439 .arg(v_g)
26440 .arg(&cd)
26441 .arg(&ti)
26442 .arg(&dc)
26443 .arg(&ds)
26444 .arg(&nv)
26445 .arg(&nk)
26446 .arg(&kd)
26447 .arg(&hki);
26448 unsafe {
26449 b.launch(cfg)?;
26450 }
26451 }
26452 match pad_len {
26453 Some(len_d) => {
26454 let f = self.func("ssm_conv_ring_update_dev_f32");
26455 let n = conv_dim * (d_conv - 1);
26456 let cfg = LaunchConfig::for_num_elems(n as u32);
26457 let (cd, dc) = (conv_dim as i32, d_conv as i32);
26458 let __s_b = self.gpu.stream();
26459 let mut b = __s_b.launch_builder(&f);
26460 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
26461 unsafe {
26462 b.launch(cfg)?;
26463 }
26464 }
26465 None => {
26466 let f = self.func("ssm_conv_ring_update_f32");
26467 let n = conv_dim * (d_conv - 1);
26468 let cfg = LaunchConfig::for_num_elems(n as u32);
26469 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
26470 let __s_b = self.gpu.stream();
26471 let mut b = __s_b.launch_builder(&f);
26472 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
26473 unsafe {
26474 b.launch(cfg)?;
26475 }
26476 }
26477 }
26478 Ok(())
26479 }
26480
26481 pub fn gdn_chunk_alloc(
26485 &self,
26486 n_head: usize,
26487 t: usize,
26488 c: usize,
26489 hk: usize,
26490 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
26491 const D: usize = 128;
26492 assert!(
26493 c == 32,
26494 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
26495 );
26496 let h = n_head;
26497 let nc = (t + c - 1) / c;
26498 Ok(GdnChunkBufs {
26499 gcum: self.uninit(t * h)?,
26500 a: self.uninit(nc * h * c * c)?,
26501 p: self.uninit(nc * h * c * c)?,
26502 u: self.uninit(nc * h * c * D)?,
26503 w: self.uninit(nc * h * c * D)?,
26504 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
26505 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
26506 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
26507 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
26508 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
26509 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
26510 o: self.uninit(D * h * t)?,
26511 t,
26512 nc,
26513 })
26514 }
26515
26516 pub fn f32_to_bf16_v(
26518 &self,
26519 x: &cudarc::driver::CudaView<f32>,
26520 dst: &mut CudaSlice<u8>,
26521 n: usize,
26522 ) -> Result<(), Box<dyn std::error::Error>> {
26523 let f = self.func("f32_to_bf16_bulk");
26524 let ni = n as i64;
26525 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26526 let __s_b = self.gpu.stream();
26527 let mut b = __s_b.launch_builder(&f);
26528 b.arg(x).arg(dst).arg(&ni);
26529 unsafe {
26530 b.launch(cfg)?;
26531 }
26532 Ok(())
26533 }
26534
26535 pub fn f32_to_bf16_into(
26537 &self,
26538 x: &CudaSlice<f32>,
26539 dst: &mut CudaSlice<u8>,
26540 n: usize,
26541 ) -> Result<(), Box<dyn std::error::Error>> {
26542 let f = self.func("f32_to_bf16_bulk");
26543 let ni = n as i64;
26544 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26545 let __s_b = self.gpu.stream();
26546 let mut b = __s_b.launch_builder(&f);
26547 b.arg(x).arg(dst).arg(&ni);
26548 unsafe {
26549 b.launch(cfg)?;
26550 }
26551 Ok(())
26552 }
26553
26554 pub fn gdn_chunk_k123_vl8(
26557 &self,
26558 seqs: &[GdnSeqVl],
26559 n_head: usize,
26560 hk: usize,
26561 wq: Option<&GdnWVl8>,
26562 ) -> Result<(), Box<dyn std::error::Error>> {
26563 let b = seqs.len();
26564 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
26565 let mut packed = [GdnSeqVl::default(); 8];
26566 packed[..b].copy_from_slice(seqs);
26567 let v = GdnVl8(packed);
26568 let (hi, ci) = (n_head as i32, 32i32);
26569 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26570 {
26571 let f = self.func("gdn_chunk_cumgate_vl");
26572 let cfg = LaunchConfig {
26573 grid_dim: (max_nc, n_head as u32, b as u32),
26574 block_dim: (32, 1, 1),
26575 shared_mem_bytes: 0,
26576 };
26577 let __s_lb = self.gpu.stream();
26578 let mut lb = __s_lb.launch_builder(&f);
26579 lb.arg(&v).arg(&hi).arg(&ci);
26580 unsafe {
26581 lb.launch(cfg)?;
26582 }
26583 }
26584 let hki = hk as i32;
26585 if let Some(w) = wq {
26586 let f = self.func("gdn_k2_wgmma_vl");
26588 let cfg = LaunchConfig {
26589 grid_dim: (max_nc, n_head as u32, b as u32),
26590 block_dim: (128, 1, 1),
26591 shared_mem_bytes: 0,
26592 };
26593 let __s_lb = self.gpu.stream();
26594 let mut lb = __s_lb.launch_builder(&f);
26595 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
26596 unsafe {
26597 lb.launch(cfg)?;
26598 }
26599 } else {
26600 let f = self.func("gdn_chunk_attn_vl");
26601 f.set_attribute(
26602 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26603 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
26604 )?;
26605 let cfg = LaunchConfig {
26606 grid_dim: (max_nc, n_head as u32, b as u32),
26607 block_dim: (256, 1, 1),
26608 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
26609 };
26610 let __s_lb = self.gpu.stream();
26611 let mut lb = __s_lb.launch_builder(&f);
26612 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26613 unsafe {
26614 lb.launch(cfg)?;
26615 }
26616 }
26617 {
26618 let f = self.func("gdn_chunk_solve32_vl");
26619 let cfg = LaunchConfig {
26620 grid_dim: (max_nc, n_head as u32, b as u32),
26621 block_dim: (256, 1, 1),
26622 shared_mem_bytes: 0,
26623 };
26624 let __s_lb = self.gpu.stream();
26625 let mut lb = __s_lb.launch_builder(&f);
26626 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26627 unsafe {
26628 lb.launch(cfg)?;
26629 }
26630 }
26631 Ok(())
26632 }
26633
26634 #[allow(clippy::too_many_arguments)]
26638 pub fn gdn_prep_vl8(
26639 &self,
26640 seqs: &[GdnPrepVl],
26641 conv_w: &CudaSlice<f32>,
26642 dt_bias: &CudaSlice<f32>,
26643 a: &CudaSlice<f32>,
26644 conv_dim: usize,
26645 d_conv: usize,
26646 d_state: usize,
26647 num_v: usize,
26648 num_k: usize,
26649 key_dim: usize,
26650 hk: usize,
26651 eps: f32,
26652 ) -> Result<(), Box<dyn std::error::Error>> {
26653 let b = seqs.len();
26654 assert!(b >= 1 && b <= 8);
26655 let mut packed = [GdnPrepVl::default(); 8];
26656 packed[..b].copy_from_slice(seqs);
26657 let v = GdnPrepVl8(packed);
26658 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26659 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
26660 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
26661 assert!(
26662 conv_fuse || hk == num_v,
26663 "de-broadcast requires the fused conv"
26664 );
26665 if conv_fuse {
26666 let f = self.func("ssm_conv1d_gdn_state_vl");
26667 let cfg = LaunchConfig {
26668 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26669 block_dim: (256, 1, 1),
26670 shared_mem_bytes: 0,
26671 };
26672 let (dsi, nvi, nki, kdi, hki) = (
26673 d_state as i32,
26674 num_v as i32,
26675 num_k as i32,
26676 key_dim as i32,
26677 hk as i32,
26678 );
26679 let __s_lb = self.gpu.stream();
26680 let mut lb = __s_lb.launch_builder(&f);
26681 lb.arg(&v)
26682 .arg(conv_w)
26683 .arg(&cdi)
26684 .arg(&dci)
26685 .arg(&dsi)
26686 .arg(&nvi)
26687 .arg(&nki)
26688 .arg(&kdi)
26689 .arg(&hki);
26690 unsafe {
26691 lb.launch(cfg)?;
26692 }
26693 } else {
26694 let f = self.func("ssm_conv1d_tm_state_vl");
26695 let cfg = LaunchConfig {
26696 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26697 block_dim: (256, 1, 1),
26698 shared_mem_bytes: 0,
26699 };
26700 let __s_lb = self.gpu.stream();
26701 let mut lb = __s_lb.launch_builder(&f);
26702 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
26703 unsafe {
26704 lb.launch(cfg)?;
26705 }
26706 }
26707 {
26708 let f = self.func("ssm_conv_ring_update_vl");
26709 let n = (conv_dim * (d_conv - 1)) as u32;
26710 let cfg = LaunchConfig {
26711 grid_dim: (n.div_ceil(256), 1, b as u32),
26712 block_dim: (256, 1, 1),
26713 shared_mem_bytes: 0,
26714 };
26715 let __s_lb = self.gpu.stream();
26716 let mut lb = __s_lb.launch_builder(&f);
26717 lb.arg(&v).arg(&cdi).arg(&dci);
26718 unsafe {
26719 lb.launch(cfg)?;
26720 }
26721 }
26722 if !conv_fuse {
26723 let f = self.func("qkv_to_gdn_repack_vl");
26724 let n = max_t * (num_v * d_state) as u32;
26725 let cfg = LaunchConfig {
26726 grid_dim: (n.div_ceil(256), 1, b as u32),
26727 block_dim: (256, 1, 1),
26728 shared_mem_bytes: 0,
26729 };
26730 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
26731 let __s_lb = self.gpu.stream();
26732 let mut lb = __s_lb.launch_builder(&f);
26733 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
26734 unsafe {
26735 lb.launch(cfg)?;
26736 }
26737 }
26738 if Self::l2_v2_on(d_state) {
26739 let f = self.func("gdn_l2_v2_vl");
26740 let cfg = LaunchConfig {
26741 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
26742 block_dim: (256, 1, 1),
26743 shared_mem_bytes: 0,
26744 };
26745 let (dsi, nvi) = (d_state as i32, hk as i32);
26746 let __s_lb = self.gpu.stream();
26747 let mut lb = __s_lb.launch_builder(&f);
26748 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26749 unsafe {
26750 lb.launch(cfg)?;
26751 }
26752 } else {
26753 let f = self.func("gdn_l2_vl");
26754 let cfg = LaunchConfig {
26755 grid_dim: (max_t * hk as u32, 2, b as u32),
26756 block_dim: (256, 1, 1),
26757 shared_mem_bytes: 0,
26758 };
26759 let (dsi, nvi) = (d_state as i32, hk as i32);
26760 let __s_lb = self.gpu.stream();
26761 let mut lb = __s_lb.launch_builder(&f);
26762 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26763 unsafe {
26764 lb.launch(cfg)?;
26765 }
26766 }
26767 {
26768 let f = self.func("gdn_gate_prep_vl");
26769 let n = max_t * num_v as u32;
26770 let cfg = LaunchConfig {
26771 grid_dim: (n.div_ceil(256), 1, b as u32),
26772 block_dim: (256, 1, 1),
26773 shared_mem_bytes: 0,
26774 };
26775 let nvi = num_v as i32;
26776 let __s_lb = self.gpu.stream();
26777 let mut lb = __s_lb.launch_builder(&f);
26778 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
26779 unsafe {
26780 lb.launch(cfg)?;
26781 }
26782 }
26783 Ok(())
26784 }
26785
26786 pub fn gdn_mirror_vl8(
26788 &self,
26789 seqs: &[GdnSeqVl],
26790 n_head: usize,
26791 which: i32,
26792 hk: usize,
26793 ) -> Result<(), Box<dyn std::error::Error>> {
26794 let b = seqs.len();
26795 assert!(b >= 1 && b <= 8);
26796 let mut packed = [GdnSeqVl::default(); 8];
26797 packed[..b].copy_from_slice(seqs);
26798 let v = GdnVl8(packed);
26799 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
26800 let max_n = seqs
26801 .iter()
26802 .map(|s| {
26803 if which == 0 {
26804 s.t as i64 * ept as i64
26805 } else {
26806 s.nc as i64 * ept as i64 * 32
26807 }
26808 })
26809 .max()
26810 .unwrap();
26811 let f = self.func("gdn_mirror_vl");
26812 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
26813 let cfg = LaunchConfig {
26814 grid_dim: (blocks, 1, b as u32),
26815 block_dim: (256, 1, 1),
26816 shared_mem_bytes: 0,
26817 };
26818 let __s_lb = self.gpu.stream();
26819 let mut lb = __s_lb.launch_builder(&f);
26820 lb.arg(&v).arg(&ept).arg(&which);
26821 unsafe {
26822 lb.launch(cfg)?;
26823 }
26824 Ok(())
26825 }
26826
26827 pub fn gdn_tail_vl8(
26829 &self,
26830 seqs: &[GdnPrepVl],
26831 norm_w: &CudaSlice<f32>,
26832 d_state: usize,
26833 num_v: usize,
26834 eps: f32,
26835 ) -> Result<(), Box<dyn std::error::Error>> {
26836 let b = seqs.len();
26837 assert!(b >= 1 && b <= 8);
26838 let mut packed = [GdnPrepVl::default(); 8];
26839 packed[..b].copy_from_slice(seqs);
26840 let v = GdnPrepVl8(packed);
26841 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26842 let f = self.func("gated_rmsnorm_f16out_vl");
26843 let cfg = LaunchConfig {
26845 grid_dim: (max_t * num_v as u32, 1, b as u32),
26846 block_dim: (128, 1, 1),
26847 shared_mem_bytes: 0,
26848 };
26849 let (dsi, nvi) = (d_state as i32, num_v as i32);
26850 let __s_lb = self.gpu.stream();
26851 let mut lb = __s_lb.launch_builder(&f);
26852 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
26853 unsafe {
26854 lb.launch(cfg)?;
26855 }
26856 Ok(())
26857 }
26858
26859 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
26862 use cudarc::driver::DevicePtr;
26863 let s = self.gpu.stream();
26864 let (p, _g) = x.device_ptr(&s);
26865 p as u64
26866 }
26867 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
26868 use cudarc::driver::DevicePtrMut;
26869 let s = self.gpu.stream();
26870 let (p, _g) = x.device_ptr_mut(&s);
26871 p as u64
26872 }
26873 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
26874 use cudarc::driver::DevicePtr;
26875 let s = self.gpu.stream();
26876 let (p, _g) = x.device_ptr(&s);
26877 p as u64
26878 }
26879 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
26880 use cudarc::driver::DevicePtr;
26881 let s = self.gpu.stream();
26882 let (p, _g) = x.device_ptr(&s);
26883 p as u64
26884 }
26885
26886 pub fn gdn_chunk_vl8(
26890 &self,
26891 seqs: &[GdnSeqVl],
26892 n_head: usize,
26893 scale: f32,
26894 hk: usize,
26895 wq: Option<&GdnWVl8>,
26896 ) -> Result<(), Box<dyn std::error::Error>> {
26897 const NSPLIT: u32 = 4;
26898 let b = seqs.len();
26899 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
26900 let mut packed = [GdnSeqVl::default(); 8];
26901 packed[..b].copy_from_slice(seqs);
26902 let v = GdnVl8(packed);
26903 let (hi, ci) = (n_head as i32, 32i32);
26904 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26905 let hki = hk as i32;
26906 if let Some(w) = wq {
26907 let f = self.func("gdn_k45_wgmma_vl");
26909 let cfg = LaunchConfig {
26910 grid_dim: (n_head as u32, NSPLIT, b as u32),
26911 block_dim: (256, 1, 1),
26912 shared_mem_bytes: 0,
26913 };
26914 let __s_lb = self.gpu.stream();
26915 let mut lb = __s_lb.launch_builder(&f);
26916 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
26917 unsafe {
26918 lb.launch(cfg)?;
26919 }
26920 let _ = max_nc;
26921 return Ok(());
26922 }
26923 {
26924 let f = self.func("gdn_chunk_state_mma_vl");
26925 let cfg = LaunchConfig {
26926 grid_dim: (n_head as u32, NSPLIT, b as u32),
26927 block_dim: (256, 1, 1),
26928 shared_mem_bytes: 0,
26929 };
26930 let __s_lb = self.gpu.stream();
26931 let mut lb = __s_lb.launch_builder(&f);
26932 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26933 unsafe {
26934 lb.launch(cfg)?;
26935 }
26936 }
26937 {
26938 let f = self.func("gdn_chunk_output_mma_vl");
26939 let cfg = LaunchConfig {
26940 grid_dim: (max_nc, n_head as u32, b as u32),
26941 block_dim: (256, 1, 1),
26942 shared_mem_bytes: 0,
26943 };
26944 let __s_lb = self.gpu.stream();
26945 let mut lb = __s_lb.launch_builder(&f);
26946 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
26947 unsafe {
26948 lb.launch(cfg)?;
26949 }
26950 }
26951 Ok(())
26952 }
26953 pub fn gdn_scan_chunked(
26954 &self,
26955 q: &CudaSlice<f32>,
26956 k: &CudaSlice<f32>,
26957 v: &CudaSlice<f32>,
26958 g: &CudaSlice<f32>,
26959 beta: &CudaSlice<f32>,
26960 kb16_pre: Option<&CudaSlice<u8>>,
26961 qb16_pre: Option<&CudaSlice<u8>>,
26962 state_in: &CudaSlice<f32>,
26963 state_out: &mut CudaSlice<f32>,
26964 o: &mut CudaSlice<f32>,
26965 n_head: usize,
26966 t: usize,
26967 scale: f32,
26968 c: usize,
26969 hk: usize,
26970 ) -> Result<(), Box<dyn std::error::Error>> {
26971 const D: usize = 128;
26972 const NSPLIT: u32 = 4;
26973 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
26974 let h = n_head;
26975 let nc = (t + c - 1) / c;
26976 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
26977 let gdn_mma_pre = !portable_mma_gated()
26982 && c == 32
26983 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26984 Ok("1") => true,
26985 Ok("0") => false,
26986 _ => gdn_mma_default_on(),
26987 };
26988 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
26989 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
26990 } else {
26991 None
26992 };
26993 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
26998 && gdn_mma_pre
26999 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
27000 let nk = t * hk * D;
27001 let mut kb16_local: Option<CudaSlice<u8>> = None;
27002 if gdn_mma_pre && kb16_pre.is_none() {
27003 let mut kb = self.alloc_u8_uninit(nk * 2)?;
27004 let f = self.func("f32_to_bf16_bulk");
27005 let n2 = nk as i64;
27006 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
27007 let __s_b = self.gpu.stream();
27008 let mut b = __s_b.launch_builder(&f);
27009 b.arg(k).arg(&mut kb).arg(&n2);
27010 unsafe {
27011 b.launch(cfg2)?;
27012 }
27013 kb16_local = Some(kb);
27014 }
27015 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
27016 if let Some(kb) = kb16_pre {
27017 assert!(kb.len() >= nk * 2, "kb16_pre too small");
27018 }
27019 let mut qb16: Option<CudaSlice<u8>> = None;
27020 let mut pb16: Option<CudaSlice<u8>> = None;
27021 if gdn_wgmma_pre {
27022 if qb16_pre.is_none() {
27025 let mut qb = self.alloc_u8_uninit(nk * 2)?;
27026 let f = self.func("f32_to_bf16_bulk");
27027 let n2 = nk as i64;
27028 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
27029 let __s_b = self.gpu.stream();
27030 let mut b = __s_b.launch_builder(&f);
27031 b.arg(q).arg(&mut qb).arg(&n2);
27032 unsafe {
27033 b.launch(cfg2)?;
27034 }
27035 qb16 = Some(qb);
27036 } else if let Some(qb) = qb16_pre {
27037 assert!(qb.len() >= nk * 2, "qb16_pre too small");
27038 }
27039 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
27040 }
27041 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
27042 let k2w = if gdn_wgmma_pre {
27043 Some((
27044 *qb16_ref0.as_ref().unwrap(),
27045 *kb16_ref0.as_ref().unwrap(),
27046 pb16.as_mut().unwrap(),
27047 ))
27048 } else {
27049 None
27050 };
27051 let (gcum, p, u, w) =
27052 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
27053 let _ = &w;
27054 let mut y = self.uninit(nc * h * c * D)?;
27055 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
27069 && c == 32
27070 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
27071 Ok("1") => true,
27072 Ok("0") => false,
27073 _ => gdn_mma_default_on(),
27074 };
27075 if gdn_mma {
27076 let wb16 = wb16_pre
27077 .take()
27078 .expect("mma path pre-allocates wb16 (K3 store fold)");
27079 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
27080 if gdn_wgmma_pre {
27092 let qb16 = qb16_ref0.unwrap();
27094 let pb16 = pb16.as_ref().unwrap();
27095 {
27096 let f = self.func("gdn_k45_wgmma");
27097 let cfg = LaunchConfig {
27098 grid_dim: (h as u32, 4, 1),
27099 block_dim: (256, 1, 1),
27100 shared_mem_bytes: 0,
27101 };
27102 let hki = hk as i32;
27103 let __s_b = self.gpu.stream();
27104 let mut b = __s_b.launch_builder(&f);
27105 b.arg(kb16_ref)
27106 .arg(&gcum)
27107 .arg(beta)
27108 .arg(&u)
27109 .arg(&wb16)
27110 .arg(qb16)
27111 .arg(pb16)
27112 .arg(o)
27113 .arg(&scale)
27114 .arg(state_in)
27115 .arg(&mut *state_out)
27116 .arg(&hi)
27117 .arg(&ti)
27118 .arg(&ci)
27119 .arg(&hki);
27120 unsafe {
27121 b.launch(cfg)?;
27122 }
27123 }
27124 return Ok(());
27125 }
27126 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
27130 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
27131 {
27132 let f = self.func("gdn_chunk_state_mma");
27133 let cfg = LaunchConfig {
27134 grid_dim: (h as u32, NSPLIT, 1),
27135 block_dim: (256, 1, 1),
27136 shared_mem_bytes: 0,
27137 };
27138 let hki = hk as i32;
27139 let __s_b = self.gpu.stream();
27140 let mut b = __s_b.launch_builder(&f);
27141 b.arg(kb16_ref)
27142 .arg(&gcum)
27143 .arg(beta)
27144 .arg(&u)
27145 .arg(&wb16)
27146 .arg(&mut y16)
27147 .arg(&mut ssnap16)
27148 .arg(state_in)
27149 .arg(&mut *state_out)
27150 .arg(&hi)
27151 .arg(&ti)
27152 .arg(&ci)
27153 .arg(&hki);
27154 unsafe {
27155 b.launch(cfg)?;
27156 }
27157 }
27158 {
27159 let f = self.func("gdn_chunk_output_mma");
27161 let jt = ((c + 31) / 32) as u32;
27162 let cfg = LaunchConfig {
27163 grid_dim: (nc as u32, h as u32, jt),
27164 block_dim: (256, 1, 1),
27165 shared_mem_bytes: 0,
27166 };
27167 let hki = hk as i32;
27168 let __s_b = self.gpu.stream();
27169 let mut b = __s_b.launch_builder(&f);
27170 b.arg(q)
27171 .arg(&gcum)
27172 .arg(&p)
27173 .arg(&y16)
27174 .arg(&ssnap16)
27175 .arg(o)
27176 .arg(&hi)
27177 .arg(&ti)
27178 .arg(&ci)
27179 .arg(&scale)
27180 .arg(&hki);
27181 unsafe {
27182 b.launch(cfg)?;
27183 }
27184 }
27185 return Ok(());
27186 }
27187 {
27188 let f = self.func("gdn_chunk_state_f32");
27190 let cfg = LaunchConfig {
27191 grid_dim: (h as u32, NSPLIT, 1),
27192 block_dim: (256, 1, 1),
27193 shared_mem_bytes: 0,
27194 };
27195 let __s_b = self.gpu.stream();
27196 let mut b = __s_b.launch_builder(&f);
27197 b.arg(k)
27198 .arg(&gcum)
27199 .arg(beta)
27200 .arg(&u)
27201 .arg(&w)
27202 .arg(&mut y)
27203 .arg(&mut ssnap)
27204 .arg(state_in)
27205 .arg(&mut *state_out)
27206 .arg(&hi)
27207 .arg(&ti)
27208 .arg(&ci);
27209 unsafe {
27210 b.launch(cfg)?;
27211 }
27212 }
27213 {
27214 let f = self.func("gdn_chunk_output_f32");
27216 let jt = ((c + 31) / 32) as u32;
27217 let cfg = LaunchConfig {
27218 grid_dim: (nc as u32, h as u32, jt),
27219 block_dim: (256, 1, 1),
27220 shared_mem_bytes: 0,
27221 };
27222 let __s_b = self.gpu.stream();
27223 let mut b = __s_b.launch_builder(&f);
27224 b.arg(q)
27225 .arg(&gcum)
27226 .arg(&p)
27227 .arg(&y)
27228 .arg(&ssnap)
27229 .arg(o)
27230 .arg(&hi)
27231 .arg(&ti)
27232 .arg(&ci)
27233 .arg(&scale);
27234 unsafe {
27235 b.launch(cfg)?;
27236 }
27237 }
27238 Ok(())
27239 }
27240
27241 #[allow(clippy::too_many_arguments)]
27250 #[allow(clippy::too_many_arguments)]
27251 pub fn gdn_scan_prefill(
27252 &self,
27253 q: &CudaSlice<f32>,
27254 k: &CudaSlice<f32>,
27255 v: &CudaSlice<f32>,
27256 g: &CudaSlice<f32>,
27257 beta: &CudaSlice<f32>,
27258 kb16_pre: Option<&CudaSlice<u8>>,
27259 qb16_pre: Option<&CudaSlice<u8>>,
27260 state_in: &CudaSlice<f32>,
27261 state_out: &mut CudaSlice<f32>,
27262 o: &mut CudaSlice<f32>,
27263 n_head: usize,
27264 t: usize,
27265 scale: f32,
27266 hk: usize,
27267 ) -> Result<(), Box<dyn std::error::Error>> {
27268 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
27269 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
27270 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
27271 }
27272 if Self::gdn_chunked_enabled() && t >= 16 {
27273 self.gdn_scan_chunked(
27274 q,
27275 k,
27276 v,
27277 g,
27278 beta,
27279 kb16_pre,
27280 qb16_pre,
27281 state_in,
27282 state_out,
27283 o,
27284 n_head,
27285 t,
27286 scale,
27287 Self::gdn_chunk_size(),
27288 hk,
27289 )
27290 } else {
27291 assert!(
27292 hk == n_head,
27293 "s128 scan is broadcast-only (prep guarantees by predicate)"
27294 );
27295 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
27296 }
27297 }
27298
27299 #[allow(clippy::too_many_arguments)]
27301 fn gdn_scan_diff(
27302 &self,
27303 q: &CudaSlice<f32>,
27304 k: &CudaSlice<f32>,
27305 v: &CudaSlice<f32>,
27306 g: &CudaSlice<f32>,
27307 beta: &CudaSlice<f32>,
27308 state_in: &CudaSlice<f32>,
27309 state_out: &mut CudaSlice<f32>,
27310 o: &mut CudaSlice<f32>,
27311 n_head: usize,
27312 t: usize,
27313 scale: f32,
27314 ) -> Result<(), Box<dyn std::error::Error>> {
27315 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
27316 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
27317 let mut o_c = self.uninit(o.len())?;
27318 let mut st_c = self.uninit(state_out.len())?;
27319 self.gdn_scan_chunked(
27320 q,
27321 k,
27322 v,
27323 g,
27324 beta,
27325 None,
27326 None,
27327 state_in,
27328 &mut st_c,
27329 &mut o_c,
27330 n_head,
27331 t,
27332 scale,
27333 Self::gdn_chunk_size(),
27334 n_head,
27335 )?;
27336 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
27337 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
27338 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
27339 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
27340 let mut max_abs = 0f32;
27341 let mut max_rel = 0f32;
27342 let mut sum_rel = 0f64;
27343 for (x, y) in a.iter().zip(b) {
27344 let ad = (x - y).abs();
27345 let rel = ad / x.abs().max(y.abs()).max(1e-3);
27346 if ad > max_abs {
27347 max_abs = ad;
27348 }
27349 if rel > max_rel {
27350 max_rel = rel;
27351 }
27352 sum_rel += rel as f64;
27353 }
27354 (max_abs, max_rel, sum_rel / a.len() as f64)
27355 };
27356 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
27357 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
27358 println!(
27359 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
27360 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
27361 Self::gdn_chunk_size()
27362 );
27363 Ok(())
27364 }
27365
27366 pub fn gdn_glog(
27368 &self,
27369 alpha: &CudaSlice<f32>,
27370 dt_bias: &CudaSlice<f32>,
27371 a: &CudaSlice<f32>,
27372 g_log: &mut CudaSlice<f32>,
27373 n_head: usize,
27374 t: usize,
27375 ) -> Result<(), Box<dyn std::error::Error>> {
27376 let f = self.func("gdn_glog_f32");
27377 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
27378 let (h, ti) = (n_head as i32, t as i32);
27379 let __s_b = self.gpu.stream();
27380 let mut b = __s_b.launch_builder(&f);
27381 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
27382 unsafe {
27383 b.launch(cfg)?;
27384 }
27385 Ok(())
27386 }
27387
27388 pub fn sigmoid_v(
27391 &self,
27392 x: &cudarc::driver::CudaView<f32>,
27393 y: &mut CudaSlice<f32>,
27394 n: usize,
27395 ) -> Result<(), Box<dyn std::error::Error>> {
27396 let f = self.func("sigmoid_f32");
27397 let cfg = LaunchConfig::for_num_elems(n as u32);
27398 let ni = n as i32;
27399 let __s_b = self.gpu.stream();
27400 let mut b = __s_b.launch_builder(&f);
27401 b.arg(x).arg(y).arg(&ni);
27402 unsafe {
27403 b.launch(cfg)?;
27404 }
27405 Ok(())
27406 }
27407
27408 pub fn gdn_glog_v(
27409 &self,
27410 alpha: &cudarc::driver::CudaView<f32>,
27411 dt_bias: &CudaSlice<f32>,
27412 a: &CudaSlice<f32>,
27413 g_log: &mut CudaSlice<f32>,
27414 n_head: usize,
27415 t: usize,
27416 ) -> Result<(), Box<dyn std::error::Error>> {
27417 let f = self.func("gdn_glog_f32");
27418 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
27419 let (h, ti) = (n_head as i32, t as i32);
27420 let __s_b = self.gpu.stream();
27421 let mut b = __s_b.launch_builder(&f);
27422 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
27423 unsafe {
27424 b.launch(cfg)?;
27425 }
27426 Ok(())
27427 }
27428
27429 pub fn sigmoid(
27430 &self,
27431 x: &CudaSlice<f32>,
27432 y: &mut CudaSlice<f32>,
27433 n: usize,
27434 ) -> Result<(), Box<dyn std::error::Error>> {
27435 let f = self.func("sigmoid_f32");
27436 let cfg = LaunchConfig::for_num_elems(n as u32);
27437 let ni = n as i32;
27438 let __s_b = self.gpu.stream();
27439 let mut b = __s_b.launch_builder(&f);
27440 b.arg(x).arg(y).arg(&ni);
27441 unsafe {
27442 b.launch(cfg)?;
27443 }
27444 Ok(())
27445 }
27446
27447 pub fn sig_mul_f16out(
27450 &self,
27451 a: &CudaSlice<f32>,
27452 g: &CudaSlice<f32>,
27453 dst: &mut CudaSlice<f32>,
27454 dst16: &mut CudaSlice<u8>,
27455 n: usize,
27456 ) -> Result<(), Box<dyn std::error::Error>> {
27457 let f = self.func("sig_mul_f16out_f32");
27458 let cfg = LaunchConfig::for_num_elems(n as u32);
27459 let ni = n as i32;
27460 let __s_b = self.gpu.stream();
27461 let mut b = __s_b.launch_builder(&f);
27462 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
27463 unsafe {
27464 b.launch(cfg)?;
27465 }
27466 Ok(())
27467 }
27468
27469 #[allow(clippy::too_many_arguments)]
27478 pub fn attn_head_gate(
27479 &self,
27480 a: &CudaSlice<f32>,
27481 g: &CudaSlice<f32>,
27482 dst: &mut CudaSlice<f32>,
27483 dst16: Option<&mut CudaSlice<u8>>,
27484 head_dim: usize,
27485 n_head: usize,
27486 t: usize,
27487 ) -> Result<(), Box<dyn std::error::Error>> {
27488 let f = self.func("attn_head_gate_f32");
27489 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27490 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27491 let d16: u64 = match dst16 {
27493 Some(d) => self.addr_u8(d),
27494 None => 0,
27495 };
27496 let __s_b = self.gpu.stream();
27497 let mut b = __s_b.launch_builder(&f);
27498 b.arg(a)
27499 .arg(g)
27500 .arg(dst)
27501 .arg(&d16)
27502 .arg(&hd)
27503 .arg(&nh)
27504 .arg(&ti);
27505 unsafe {
27506 b.launch(cfg)?;
27507 }
27508 Ok(())
27509 }
27510
27511 #[allow(clippy::too_many_arguments)]
27520 pub fn swiglu_clamped_mul_scaled(
27521 &self,
27522 gate: &CudaSlice<f32>,
27523 up: &CudaSlice<f32>,
27524 gs: f32,
27525 us: f32,
27526 limit: f32,
27527 dst: &mut CudaSlice<f32>,
27528 n: usize,
27529 ) -> Result<(), Box<dyn std::error::Error>> {
27530 debug_assert!(
27531 limit > 1e-6,
27532 "swiglu_clamped needs a live limit; use silu_mul_scaled"
27533 );
27534 let f = self.func("swiglu_clamped_mul_scaled_f32");
27535 let cfg = LaunchConfig::for_num_elems(n as u32);
27536 let ni = n as i32;
27537 let __s_b = self.gpu.stream();
27538 let mut b = __s_b.launch_builder(&f);
27539 b.arg(gate)
27540 .arg(up)
27541 .arg(&gs)
27542 .arg(&us)
27543 .arg(&limit)
27544 .arg(dst)
27545 .arg(&ni);
27546 unsafe {
27547 b.launch(cfg)?;
27548 }
27549 Ok(())
27550 }
27551
27552 pub fn gated_rmsnorm(
27554 &self,
27555 o: &CudaSlice<f32>,
27556 w: &CudaSlice<f32>,
27557 z: &CudaSlice<f32>,
27558 dst: &mut CudaSlice<f32>,
27559 ncols: usize,
27560 nrows: usize,
27561 eps: f32,
27562 ) -> Result<(), Box<dyn std::error::Error>> {
27563 let f = self.func("gated_rmsnorm_f32");
27564 let cfg = LaunchConfig {
27565 grid_dim: (nrows as u32, 1, 1),
27566 block_dim: (128, 1, 1),
27567 shared_mem_bytes: 0,
27568 };
27569 let (nc, e) = (ncols as i32, eps);
27570 let __s_b = self.gpu.stream();
27571 let mut b = __s_b.launch_builder(&f);
27572 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27573 unsafe {
27574 b.launch(cfg)?;
27575 }
27576 Ok(())
27577 }
27578
27579 pub fn gated_rmsnorm_f16out(
27582 &self,
27583 o: &CudaSlice<f32>,
27584 w: &CudaSlice<f32>,
27585 z: &CudaSlice<f32>,
27586 dst: &mut CudaSlice<f32>,
27587 dst16: &mut CudaSlice<u8>,
27588 ncols: usize,
27589 nrows: usize,
27590 eps: f32,
27591 ) -> Result<(), Box<dyn std::error::Error>> {
27592 let f = self.func("gated_rmsnorm_f16out_f32");
27593 let cfg = LaunchConfig {
27595 grid_dim: (nrows as u32, 1, 1),
27596 block_dim: (128, 1, 1),
27597 shared_mem_bytes: 0,
27598 };
27599 let (nc, e) = (ncols as i32, eps);
27600 let __s_b = self.gpu.stream();
27601 let mut b = __s_b.launch_builder(&f);
27602 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27603 unsafe {
27604 b.launch(cfg)?;
27605 }
27606 Ok(())
27607 }
27608
27609 #[allow(clippy::too_many_arguments)]
27613 pub fn add_rms_norm_zq8(
27614 &self,
27615 a: &CudaSlice<f32>,
27616 b_in: &CudaSlice<f32>,
27617 w: &CudaSlice<f32>,
27618 res: &mut CudaSlice<f32>,
27619 z: &mut CudaSlice<f32>,
27620 ncols: usize,
27621 nrows: usize,
27622 eps: f32,
27623 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27624 assert!(ncols % 32 == 0);
27625 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
27626 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27627 let f = self.func("add_rms_norm_zq8");
27628 let cfg = LaunchConfig {
27629 grid_dim: (nrows as u32, 1, 1),
27630 block_dim: (1024, 1, 1),
27631 shared_mem_bytes: 0,
27632 };
27633 let (nc, ep) = (ncols as i32, eps);
27634 let __s_b = self.gpu.stream();
27635 let mut b = __s_b.launch_builder(&f);
27636 b.arg(a)
27637 .arg(b_in)
27638 .arg(w)
27639 .arg(res)
27640 .arg(z)
27641 .arg(&mut q)
27642 .arg(&mut d)
27643 .arg(&nc)
27644 .arg(&ep);
27645 unsafe {
27646 b.launch(cfg)?;
27647 }
27648 Ok((q, d))
27649 }
27650
27651 pub fn gated_rmsnorm_zv(
27656 &self,
27657 o: &CudaSlice<f32>,
27658 w: &CudaSlice<f32>,
27659 z: &cudarc::driver::CudaView<f32>,
27660 dst: &mut CudaSlice<f32>,
27661 ncols: usize,
27662 nrows: usize,
27663 eps: f32,
27664 ) -> Result<(), Box<dyn std::error::Error>> {
27665 let f = self.func("gated_rmsnorm_f32");
27666 let cfg = LaunchConfig {
27667 grid_dim: (nrows as u32, 1, 1),
27668 block_dim: (128, 1, 1),
27669 shared_mem_bytes: 0,
27670 };
27671 let (nc, e) = (ncols as i32, eps);
27672 let __s_b = self.gpu.stream();
27673 let mut b = __s_b.launch_builder(&f);
27674 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27675 unsafe {
27676 b.launch(cfg)?;
27677 }
27678 Ok(())
27679 }
27680
27681 pub fn gated_rmsnorm_f16out_zv(
27682 &self,
27683 o: &CudaSlice<f32>,
27684 w: &CudaSlice<f32>,
27685 z: &cudarc::driver::CudaView<f32>,
27686 dst: &mut CudaSlice<f32>,
27687 dst16: &mut CudaSlice<u8>,
27688 ncols: usize,
27689 nrows: usize,
27690 eps: f32,
27691 ) -> Result<(), Box<dyn std::error::Error>> {
27692 let f = self.func("gated_rmsnorm_f16out_f32");
27693 let cfg = LaunchConfig {
27695 grid_dim: (nrows as u32, 1, 1),
27696 block_dim: (128, 1, 1),
27697 shared_mem_bytes: 0,
27698 };
27699 let (nc, e) = (ncols as i32, eps);
27700 let __s_b = self.gpu.stream();
27701 let mut b = __s_b.launch_builder(&f);
27702 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27703 unsafe {
27704 b.launch(cfg)?;
27705 }
27706 Ok(())
27707 }
27708
27709 pub fn gated_rmsnorm_q8_1(
27710 &self,
27711 o: &CudaSlice<f32>,
27712 w: &CudaSlice<f32>,
27713 z: &CudaSlice<f32>,
27714 ncols: usize,
27715 nrows: usize,
27716 eps: f32,
27717 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27718 assert!(ncols % 32 == 0);
27719 let f = self.func("gated_rmsnorm_q8_1");
27720 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
27721 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27722 let cfg = LaunchConfig {
27723 grid_dim: (nrows as u32, 1, 1),
27724 block_dim: (128, 1, 1),
27725 shared_mem_bytes: 0,
27726 };
27727 let (nc, ep) = (ncols as i32, eps);
27728 let __s_b = self.gpu.stream();
27729 let mut b = __s_b.launch_builder(&f);
27730 b.arg(o)
27731 .arg(w)
27732 .arg(z)
27733 .arg(&mut out_q)
27734 .arg(&mut out_d)
27735 .arg(&nc)
27736 .arg(&ep);
27737 unsafe {
27738 b.launch(cfg)?;
27739 }
27740 Ok((out_q, out_d))
27741 }
27742
27743 pub fn transpose(
27745 &self,
27746 inp: &CudaSlice<f32>,
27747 rows: usize,
27748 cols: usize,
27749 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27750 let f = self.func("transpose_f32");
27751 let mut out = self.zeros(rows * cols)?;
27752 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
27753 let (r, c) = (rows as i32, cols as i32);
27754 let __s_b = self.gpu.stream();
27755 let mut b = __s_b.launch_builder(&f);
27756 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
27757 unsafe {
27758 b.launch(cfg)?;
27759 }
27760 Ok(out)
27761 }
27762
27763 pub fn repeat_heads(
27765 &self,
27766 inp: &CudaSlice<f32>,
27767 out: &mut CudaSlice<f32>,
27768 head_dim: usize,
27769 n_in: usize,
27770 n_out: usize,
27771 t: usize,
27772 ) -> Result<(), Box<dyn std::error::Error>> {
27773 let f = self.func("repeat_heads_f32");
27774 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
27775 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
27776 let __s_b = self.gpu.stream();
27777 let mut b = __s_b.launch_builder(&f);
27778 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
27779 unsafe {
27780 b.launch(cfg)?;
27781 }
27782 Ok(())
27783 }
27784
27785 pub fn q_gate_split(
27792 &self,
27793 qf: &CudaSlice<f32>,
27794 q_out: &mut CudaSlice<f32>,
27795 gate_out: &mut CudaSlice<f32>,
27796 head_dim: usize,
27797 n_head: usize,
27798 t: usize,
27799 ) -> Result<(), Box<dyn std::error::Error>> {
27800 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
27801 let out_need = head_dim * n_head * t;
27802 if q_out.len() < out_need || gate_out.len() < out_need {
27803 return Err(format!(
27804 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
27805 q_out.len(),
27806 gate_out.len()
27807 )
27808 .into());
27809 }
27810 let f = self.func("q_gate_split_f32");
27811 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27812 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27813 let __s_b = self.gpu.stream();
27814 let mut b = __s_b.launch_builder(&f);
27815 b.arg(qf)
27816 .arg(q_out)
27817 .arg(gate_out)
27818 .arg(&hd)
27819 .arg(&nh)
27820 .arg(&ti);
27821 unsafe {
27822 b.launch(cfg)?;
27823 }
27824 Ok(())
27825 }
27826
27827 pub fn qkv_to_gdn_repack(
27831 &self,
27832 conv_out: &CudaSlice<f32>,
27833 q_g: &mut CudaSlice<f32>,
27834 k_g: &mut CudaSlice<f32>,
27835 v_g: &mut CudaSlice<f32>,
27836 d_state: usize,
27837 num_v: usize,
27838 num_k: usize,
27839 key_dim: usize,
27840 t: usize,
27841 ) -> Result<(), Box<dyn std::error::Error>> {
27842 let f = self.func("qkv_to_gdn_repack_f32");
27843 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
27844 let (ds, nv, nk, kd, ti) = (
27845 d_state as i32,
27846 num_v as i32,
27847 num_k as i32,
27848 key_dim as i32,
27849 t as i32,
27850 );
27851 let __s_b = self.gpu.stream();
27852 let mut b = __s_b.launch_builder(&f);
27853 b.arg(conv_out)
27854 .arg(q_g)
27855 .arg(k_g)
27856 .arg(v_g)
27857 .arg(&ds)
27858 .arg(&nv)
27859 .arg(&nk)
27860 .arg(&kd)
27861 .arg(&ti);
27862 unsafe {
27863 b.launch(cfg)?;
27864 }
27865 Ok(())
27866 }
27867
27868 pub fn conv_left_pad(
27871 &self,
27872 src: &CudaSlice<f32>,
27873 dst: &mut CudaSlice<f32>,
27874 conv_dim: usize,
27875 t: usize,
27876 pad: usize,
27877 ) -> Result<(), Box<dyn std::error::Error>> {
27878 let f = self.func("conv_left_pad_f32");
27879 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
27880 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
27881 let __s_b = self.gpu.stream();
27882 let mut b = __s_b.launch_builder(&f);
27883 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
27884 unsafe {
27885 b.launch(cfg)?;
27886 }
27887 Ok(())
27888 }
27889
27890 pub fn conv_assemble_and_roll(
27894 &self,
27895 qkv_col: &CudaSlice<f32>,
27896 conv_state: &mut CudaSlice<f32>,
27897 conv_in: &mut CudaSlice<f32>,
27898 conv_dim: usize,
27899 pad: usize,
27900 ) -> Result<(), Box<dyn std::error::Error>> {
27901 let f = self.func("conv_assemble_and_roll_f32");
27902 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27903 let (cd, p) = (conv_dim as i32, pad as i32);
27904 let __s_b = self.gpu.stream();
27905 let mut b = __s_b.launch_builder(&f);
27906 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
27907 unsafe {
27908 b.launch(cfg)?;
27909 }
27910 Ok(())
27911 }
27912
27913 pub fn ssm_conv1d_fused_decode(
27919 &self,
27920 qkv_col: &CudaSlice<f32>,
27921 conv_state: &mut CudaSlice<f32>,
27922 w: &CudaSlice<f32>,
27923 conv_out: &mut CudaSlice<f32>,
27924 conv_dim: usize,
27925 d_conv: usize,
27926 ) -> Result<(), Box<dyn std::error::Error>> {
27927 let f = self.func("ssm_conv1d_fused_decode_f32");
27928 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27929 let (cd, dc) = (conv_dim as i32, d_conv as i32);
27930 let __s_b = self.gpu.stream();
27931 let mut b = __s_b.launch_builder(&f);
27932 b.arg(qkv_col)
27933 .arg(conv_state)
27934 .arg(w)
27935 .arg(conv_out)
27936 .arg(&cd)
27937 .arg(&dc);
27938 unsafe {
27939 b.launch(cfg)?;
27940 }
27941 Ok(())
27942 }
27943
27944 pub fn slice_range(
27947 &self,
27948 src: &CudaSlice<f32>,
27949 start: usize,
27950 len: usize,
27951 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27952 let host = self.gpu.stream().clone_dtoh(src)?;
27953 self.gpu.stream().synchronize()?;
27954 Ok(self.htod(&host[start..start + len])?)
27955 }
27956}
27957
27958#[cfg(test)]
27959mod target_dispatch_tests {
27960 use super::legacy_quant_gemm_allowed;
27961
27962 #[test]
27963 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
27964 assert!(legacy_quant_gemm_allowed(false, false, false));
27966 assert!(!legacy_quant_gemm_allowed(false, false, true));
27967 assert!(!legacy_quant_gemm_allowed(true, false, false));
27969 assert!(!legacy_quant_gemm_allowed(true, false, true));
27970 assert!(legacy_quant_gemm_allowed(true, true, false));
27972 assert!(!legacy_quant_gemm_allowed(true, true, true));
27973 }
27974
27975 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
27976 #[test]
27977 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
27978 assert!(!legacy_quant_gemm_allowed(
27979 cfg!(memra_portable_cuda),
27980 cfg!(memra_hopper_mma),
27981 false
27982 ));
27983 }
27984
27985 #[cfg(memra_hopper_mma)]
27986 #[test]
27987 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
27988 assert!(legacy_quant_gemm_allowed(
27989 cfg!(memra_portable_cuda),
27990 cfg!(memra_hopper_mma),
27991 false
27992 ));
27993 assert!(super::portable_mma_gated() == false);
27994 }
27995}
27996
27997impl memra_kv::KvDev for Engine {
28000 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
28001 Engine::zeros(self, n)
28002 }
28003 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
28004 Engine::uninit(self, n)
28005 }
28006 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
28007 Engine::alloc_u8(self, n)
28008 }
28009 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
28010 Engine::htod_i32(self, v)
28011 }
28012 fn clone_dtod(
28013 &self,
28014 src: &CudaSlice<f32>,
28015 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
28016 Engine::clone_dtod(self, src)
28017 }
28018 fn copy_into(
28019 &self,
28020 dst: &mut CudaSlice<f32>,
28021 off: usize,
28022 src: &CudaSlice<f32>,
28023 len: usize,
28024 ) -> Result<(), Box<dyn std::error::Error>> {
28025 Engine::copy_into(self, dst, off, src, len)
28026 }
28027 fn set_i32_one(
28028 &self,
28029 d: &mut CudaSlice<i32>,
28030 v: i32,
28031 ) -> Result<(), Box<dyn std::error::Error>> {
28032 Engine::set_i32_one(self, d, v)
28033 }
28034}
28035
28036#[cfg(test)]
28037mod fused_gate_bounds_tests {
28038 use super::*;
28039
28040 #[test]
28053 #[ignore = "requires a CUDA GPU"]
28054 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
28055 let e = Engine::new(0).unwrap();
28056 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
28057 let fused = 2 * head_dim * n_head * t;
28058 let out_n = head_dim * n_head * t;
28059
28060 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
28062 let mut q = e.uninit(out_n).unwrap();
28063 let mut gate = e.uninit(out_n).unwrap();
28064 let err = e
28065 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
28066 .expect_err("half-width wq must be refused, not read past")
28067 .to_string();
28068 assert!(err.contains("NO fused gate"), "{err}");
28069 assert!(err.contains(&format!("{fused}")), "{err}");
28070
28071 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
28074 let wide = e.htod(&host).unwrap();
28075 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
28076 .expect("full-width wq splits");
28077 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
28078 for tok in 0..t {
28079 for hh in 0..n_head {
28080 for d in 0..head_dim {
28081 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
28082 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
28083 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
28084 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
28085 }
28086 }
28087 }
28088
28089 let mut small = e.uninit(out_n - 1).unwrap();
28091 assert!(
28092 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
28093 .is_err()
28094 );
28095 }
28096}
28097
28098#[cfg(test)]
28102mod fused_rope_width_tests {
28103 use super::Engine;
28104
28105 #[test]
28108 fn full_width_is_accepted() {
28109 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
28110 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
28111 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
28112 }
28113
28114 #[test]
28127 fn gemma4_official_artifact_widths_pass() {
28128 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
28129 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
28130 }
28131
28132 #[test]
28135 fn partial_rotary_is_refused_with_the_geometry_named() {
28136 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
28138 .expect_err("partial rotary must refuse");
28139 let msg = err.to_string();
28140 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
28141 assert!(msg.contains("n_rot 64"), "{msg}");
28142 assert!(msg.contains("head_dim 256"), "{msg}");
28143 assert!(
28144 msg.contains("64..256"),
28145 "names the band it would corrupt: {msg}"
28146 );
28147 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
28149 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
28151 }
28152}