1use cudarc::driver::sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES;
4use cudarc::driver::{
5 CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut,
6 DeviceSlice, LaunchConfig, 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 hyper;
62pub mod model;
63pub mod sigrouter_contract;
64pub mod vision;
65pub mod vision_gemma;
66pub mod vision_glm5;
67pub mod vision_pre;
68pub mod vision_step;
69pub mod cache {
72 pub use memra_kv::*;
73}
74pub mod decode;
75pub mod decode_batch;
76pub mod dflash;
77pub mod eagle;
78pub mod ep_map;
83pub mod gemma_spec;
84pub mod glm5_decode_graph;
85pub mod glm5_sel_ledger;
86pub mod glm5_tp;
87pub mod glm_spec;
91pub mod graph_update;
92pub mod kda;
93pub mod mla;
97pub mod mla_ffi;
98pub mod moe_sel_dump;
99pub mod moesd;
100pub mod parallel;
101pub mod plan_backend;
102pub mod pp;
103pub mod progress;
104pub mod qwen4exp_gpu;
107pub mod round_stream;
108pub mod spec;
109pub mod spec_phase;
114pub mod tp;
115pub mod tp_transport;
116pub use memra_sampling as sampler;
117
118pub fn moe_f16g_mode() -> u8 {
162 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
163 *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
164 Ok("0") => 0,
165 Ok("2") => 2,
166 Ok("3") => 3,
167 Ok(_) => 1,
168 Err(_) => 2,
171 })
172}
173pub fn moe_f16g_sk_params() -> (i32, i32) {
187 static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
188 *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
189 Ok("0") => (-1, 0),
190 Ok("32") => (0, i32::MAX),
191 Ok("128") => (0, 1),
192 _ => {
193 let cross = std::env::var("MEMRA_F16G_SK_CROSS")
194 .ok()
195 .and_then(|v| v.parse().ok())
196 .unwrap_or(64);
197 (0, cross)
198 }
199 })
200}
201pub fn moe_f16g_direct_on(qtype: i32) -> bool {
212 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
213 let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
214 Ok("0") => 0,
215 Ok("kq") => 1,
216 _ => 2,
217 });
218 match m {
219 0 => false,
220 1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
221 _ => true,
222 }
223}
224pub fn moe_f16g_tail_on() -> bool {
233 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
234 *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
235}
236
237pub fn moe_f16g_gemma_on() -> bool {
244 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
245 *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
246}
247
248pub fn moe_fuse_actq_on() -> bool {
252 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
253 *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
254}
255
256pub fn glm5_q8_fuse_on() -> bool {
263 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
264 *ON.get_or_init(|| std::env::var("MEMRA_GLM5_Q8_FUSE").as_deref() == Ok("1"))
265}
266
267pub fn glm5_q8_fuse_attn_on() -> bool {
273 std::env::var("MEMRA_GLM5_Q8_FUSE_ATTN").as_deref() == Ok("1")
274}
275
276pub static GLM5_Q8_FUSE_ATTN_DISPATCHES: std::sync::atomic::AtomicU64 =
278 std::sync::atomic::AtomicU64::new(0);
279
280pub fn glm5_decode_graph_on() -> bool {
304 glm5_decode_graph_on_from(std::env::var("MEMRA_GLM5_DECODE_GRAPH").ok().as_deref())
305}
306
307pub fn glm5_decode_graph_on_from(v: Option<&str>) -> bool {
312 !matches!(v.map(str::trim), Some("0"))
313}
314
315pub fn glm5_graph_host_moe() -> bool {
322 std::env::var("MEMRA_GLM5_GRAPH_HOST_MOE").as_deref() == Ok("1")
323}
324
325pub fn glm5_graph_no_capture() -> bool {
339 std::env::var("MEMRA_GLM5_GRAPH_NO_CAPTURE").as_deref() == Ok("1")
340}
341
342pub fn glm5_graph_recapture_on() -> bool {
371 std::env::var("MEMRA_GLM5_GRAPH_RECAPTURE").as_deref() == Ok("1")
372}
373
374pub fn glm5_graph_mla_on() -> bool {
383 std::env::var("MEMRA_GLM5_GRAPH_MLA").as_deref() == Ok("1")
384}
385
386pub fn glm5_vrows_t1_dev_forced() -> bool {
387 glm5_vrows_t1_dev_forced_from(
388 std::env::var("MEMRA_GLM5_VROWS_T1_DEV").ok().as_deref(),
389 env!("MEMRA_BUILT_CUDA_ARCH"),
390 glm5_graph_no_capture() && glm5_decode_graph_on(),
391 )
392}
393
394pub fn glm5_vrows_t1_dev_forced_from(
404 v: Option<&str>,
405 built_arch: &str,
406 no_capture_bisect: bool,
407) -> bool {
408 match v.map(str::trim) {
409 Some("1") => true,
410 Some("0") => no_capture_bisect,
411 _ => built_arch == "100a" || no_capture_bisect,
412 }
413}
414
415#[cfg(test)]
416mod glm5_vrows_t1_dev_default_tests {
417 use super::glm5_vrows_t1_dev_forced_from;
418
419 #[test]
420 fn arch_keyed_default_with_override_and_bisect_knob() {
421 assert!(glm5_vrows_t1_dev_forced_from(None, "100a", false));
422 assert!(!glm5_vrows_t1_dev_forced_from(None, "120a", false));
423 assert!(glm5_vrows_t1_dev_forced_from(None, "120a", true));
424 assert!(glm5_vrows_t1_dev_forced_from(Some("1"), "120a", false));
425 assert!(!glm5_vrows_t1_dev_forced_from(Some("0"), "100a", false));
426 assert!(glm5_vrows_t1_dev_forced_from(Some("0"), "100a", true));
427 }
428}
429
430pub fn glm5_graph_trace_on() -> bool {
438 std::env::var("MEMRA_GLM5_GRAPH_TRACE").as_deref() == Ok("1")
439}
440
441type Glm5TraceKey = (&'static str, String, u16);
457
458fn glm5_trace_slots() -> &'static Mutex<std::collections::BTreeSet<Glm5TraceKey>> {
459 static S: std::sync::OnceLock<Mutex<std::collections::BTreeSet<Glm5TraceKey>>> =
460 std::sync::OnceLock::new();
461 S.get_or_init(Default::default)
462}
463
464pub(crate) fn glm5_trace_take_slot(kind: &'static str, arm: &str, il: u16) -> bool {
468 const GLM5_TRACE_MAX_LAYERS: usize = 8;
469 let mut s = glm5_trace_slots().lock().unwrap();
470 if s.iter().filter(|(k, a, _)| *k == kind && a == arm).count() >= GLM5_TRACE_MAX_LAYERS {
471 return false;
472 }
473 s.insert((kind, arm.to_string(), il))
474}
475
476pub fn glm5_trace_reset() {
479 glm5_trace_slots().lock().unwrap().clear();
480}
481
482pub static GLM5_DECODE_GRAPH_REPLAYS: std::sync::atomic::AtomicU64 =
485 std::sync::atomic::AtomicU64::new(0);
486pub static GLM5_DECODE_GRAPH_CAPTURES: std::sync::atomic::AtomicU64 =
487 std::sync::atomic::AtomicU64::new(0);
488pub static GLM5_DECODE_GRAPH_LAYERS: std::sync::atomic::AtomicU64 =
489 std::sync::atomic::AtomicU64::new(0);
490pub static GLM5_DECODE_GRAPH_MLA_HALVES: std::sync::atomic::AtomicU64 =
493 std::sync::atomic::AtomicU64::new(0);
494pub static GLM5_DECODE_GRAPH_RECAPTURES: std::sync::atomic::AtomicU64 =
496 std::sync::atomic::AtomicU64::new(0);
497
498pub(crate) static GLM5_GRAPH_CAPTURE_OPEN: std::sync::atomic::AtomicBool =
504 std::sync::atomic::AtomicBool::new(false);
505
506pub(crate) fn glm5_graph_capture_open() -> bool {
507 GLM5_GRAPH_CAPTURE_OPEN.load(std::sync::atomic::Ordering::Relaxed)
508}
509
510pub const ROUTER_BATCH_MIN_T: usize = 8;
525pub fn router_batch_on() -> bool {
526 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
527 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
528}
529mod cpu_experts;
530#[cfg(memra_cutlass)]
531pub mod cutlass_ffi;
532pub mod dsv4_ffi;
533pub mod dsv4_gpu;
534pub mod f16_ffi;
535pub mod fp8_ffi;
536pub mod mmq_ffi;
537pub mod moe_cache;
538pub mod prime_graph;
539pub mod spill;
540mod spill_pread;
541
542const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
549const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
550const KDA_FATBIN: &[u8] = include_bytes!(env!("MEMRA_KDA_FATBIN"));
552const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
553const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
554const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
555const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
556const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
558
559fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
565 assert!(
566 !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
567 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
568 );
569 match std::env::var("MEMRA_GEMM_FATBIN") {
570 Ok(path) => std::borrow::Cow::Owned(
571 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
572 ),
573 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
574 }
575}
576
577pub(crate) const fn portable_mma_gated() -> bool {
584 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
585}
586
587#[track_caller]
604pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
605 assert!(
606 !portable_mma_gated(),
607 "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
608 portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
609 );
610}
611
612pub(crate) const fn gdn_mma_default_on() -> bool {
621 cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
622}
623
624const fn konst_eq(a: &str, b: &str) -> bool {
626 let (a, b) = (a.as_bytes(), b.as_bytes());
627 if a.len() != b.len() {
628 return false;
629 }
630 let mut i = 0;
631 while i < a.len() {
632 if a[i] != b[i] {
633 return false;
634 }
635 i += 1;
636 }
637 true
638}
639
640const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
645 (!portable_cuda || hopper_mma) && !no_gemm
646}
647
648const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
654
655pub use memra_kv::kv_blk_bytes;
658
659fn k1_launch_override() -> Option<(u32, u32, u32)> {
666 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
667 *K1.get_or_init(|| {
668 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
669 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
670 match p.as_slice() {
671 [bm, bn, w] => Some((*bm, *bn, *w)),
672 _ => None,
673 }
674 })
675}
676
677pub(crate) fn wgmma_gemm_enabled() -> bool {
684 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
685 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
686}
687
688pub const FA_VEC_MIN_TKV: usize = 96;
703pub fn fa_vec_min_tkv() -> usize {
707 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
708 *V.get_or_init(|| {
709 std::env::var("MEMRA_FA_VEC_MIN")
710 .ok()
711 .and_then(|v| v.parse().ok())
712 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
713 })
714}
715
716pub fn fa_f16pv_on() -> bool {
727 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
728 *ON.get_or_init(|| {
729 std::env::var("MEMRA_FA_F16PV")
730 .map(|v| v != "0")
731 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
732 })
733}
734
735pub fn fa512_hp_on() -> bool {
739 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
740 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
741}
742
743pub fn faw_hp_on() -> bool {
747 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
748 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
749}
750
751pub fn fa512_min_tkv() -> usize {
754 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
755 *FA512_MIN.get_or_init(|| {
756 std::env::var("MEMRA_FA512_MIN")
757 .ok()
758 .and_then(|v| v.parse().ok())
759 .unwrap_or(512)
760 })
761}
762pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
766 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
767pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
771pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
777 std::sync::atomic::AtomicBool::new(false);
778pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
785 std::sync::atomic::AtomicBool::new(true);
786pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
787 std::sync::atomic::AtomicUsize::new(16);
788pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
793pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
795pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
800pub(crate) fn mmv_block() -> u32 {
805 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
806 *V.get_or_init(|| {
807 std::env::var("MEMRA_MMV_BLOCK")
808 .ok()
809 .and_then(|v| v.parse().ok())
810 .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
811 .unwrap_or(128)
812 })
813}
814
815pub(crate) fn b200_matvec_arm_on() -> bool {
826 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
827 *V.get_or_init(|| {
828 env!("MEMRA_BUILT_CUDA_ARCH") == "100a"
829 && std::env::var("MEMRA_B200_MATVEC_ARM").as_deref() == Ok("1")
830 })
831}
832
833pub(crate) fn b200_gemv_v2_level() -> u8 {
865 static V: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
866 *V.get_or_init(|| {
867 if env!("MEMRA_BUILT_CUDA_ARCH") != "100a" {
868 return 0;
869 }
870 match std::env::var("MEMRA_B200_GEMV_V2").as_deref() {
871 Ok("1") => 1,
872 Ok("2") => 2,
873 _ => 0,
874 }
875 })
876}
877
878pub(crate) fn b200_gemv_v2_on() -> bool {
879 b200_gemv_v2_level() >= 1
880}
881
882pub(crate) fn q8_row_ilp_on() -> bool {
889 q8_row_ilp_on_from(
890 std::env::var("MEMRA_Q8_ROW_ILP").ok().as_deref(),
891 env!("MEMRA_BUILT_CUDA_ARCH"),
892 )
893}
894
895pub fn q8_row_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
899 match v.map(str::trim) {
900 Some("1") => true,
901 Some("0") => false,
902 _ => built_arch == "100a",
903 }
904}
905
906pub static Q8_ROW_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
908 std::sync::atomic::AtomicU64::new(0);
909
910pub fn q8_row_ilp_dispatches() -> u64 {
912 Q8_ROW_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
913}
914
915fn q8_row_ilp_note(site: &str) {
916 if Q8_ROW_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
917 eprintln!(
918 "[q8-row-ilp] engaged at {site}: q8_0 row walk with four blocks' loads per lane \
919 ahead of the dp4a chains (MEMRA_Q8_ROW_ILP=1)"
920 );
921 }
922}
923
924pub(crate) fn nvfp4_row_ilp_on() -> bool {
931 nvfp4_row_ilp_on_from(
932 std::env::var("MEMRA_NVFP4_ROW_ILP").ok().as_deref(),
933 env!("MEMRA_BUILT_CUDA_ARCH"),
934 )
935}
936
937pub fn nvfp4_row_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
941 match v.map(str::trim) {
942 Some("1") => true,
943 Some("0") => false,
944 _ => built_arch == "100a",
945 }
946}
947
948pub static NVFP4_ROW_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
950 std::sync::atomic::AtomicU64::new(0);
951
952pub fn nvfp4_row_ilp_dispatches() -> u64 {
954 NVFP4_ROW_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
955}
956
957pub(crate) fn b200_mr1_fill() -> u32 {
965 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
966 *V.get_or_init(|| {
967 std::env::var("MEMRA_B200_MR1_FILL")
968 .ok()
969 .and_then(|v| v.trim().parse::<u32>().ok())
970 .filter(|&n| n >= 1)
971 .unwrap_or(16)
972 })
973}
974
975pub(crate) const GEMV_V3_STAGES: usize = 2;
979
980pub(crate) const Q8_V2_ROWS: u32 = 8;
983pub(crate) fn gemv_v3_smem_bytes(nb: usize) -> usize {
984 GEMV_V3_STAGES * GEMV_V2_ROWS * (nb * 8) * 2 + GEMV_V2_ROWS * nb * 4
985}
986
987pub(crate) fn gemv_v3_fits() -> bool {
992 gemv_v3_smem_bytes(mmv_block() as usize) <= 48 * 1024
993}
994
995pub(crate) const GEMV_V2_ROWS: usize = 8;
998
999pub(crate) static GEMV_V2_Q8_ROWS_TW_DISPATCHES: std::sync::atomic::AtomicU64 =
1010 std::sync::atomic::AtomicU64::new(0);
1011pub(crate) static GEMV_V2_Q8_RP_DISPATCHES: std::sync::atomic::AtomicU64 =
1013 std::sync::atomic::AtomicU64::new(0);
1014pub(crate) static GEMV_V2_BF16_DISPATCHES: std::sync::atomic::AtomicU64 =
1016 std::sync::atomic::AtomicU64::new(0);
1017
1018pub(crate) fn b200_bf16_gemv_lt_on() -> bool {
1041 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1042 *V.get_or_init(|| {
1043 env!("MEMRA_BUILT_CUDA_ARCH") == "100a"
1044 && std::env::var("MEMRA_B200_BF16_GEMV_LT").as_deref() == Ok("1")
1045 })
1046}
1047
1048static STEP37_SERVING_DEFAULTS: std::sync::atomic::AtomicBool =
1074 std::sync::atomic::AtomicBool::new(false);
1075
1076pub fn arm_step37_serving_defaults() {
1077 STEP37_SERVING_DEFAULTS.store(true, std::sync::atomic::Ordering::Relaxed);
1078 crate::cache::set_swa_ring_default(true);
1079 eprintln!(
1080 "[step37-defaults] serving doors armed ON for the SlidingGatedMoe program \
1081 (per-flag =0 kills, =1 forces; owner flip 2026-08-27)"
1082 );
1083}
1084
1085pub(crate) fn step37_defaults_armed() -> bool {
1086 STEP37_SERVING_DEFAULTS.load(std::sync::atomic::Ordering::Relaxed)
1087}
1088
1089pub(crate) fn step37_door(cell: &'static std::sync::OnceLock<Option<bool>>, name: &str) -> bool {
1093 match *cell.get_or_init(|| match std::env::var(name).ok().as_deref() {
1094 Some("1") => Some(true),
1095 Some("0") => Some(false),
1096 _ => None,
1097 }) {
1098 Some(forced) => forced,
1099 None => step37_defaults_armed(),
1100 }
1101}
1102
1103pub(crate) fn w8_hybrid_on() -> bool {
1104 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1105 step37_door(&ENV, "MEMRA_W8_HYBRID")
1106}
1107
1108pub(crate) fn step_tp_w8_on() -> bool {
1109 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1110 step37_door(&ENV, "MEMRA_STEP_TP_W8")
1111}
1112
1113pub(crate) fn glm5_w8_on() -> bool {
1127 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1128 *ON.get_or_init(|| std::env::var("MEMRA_GLM5_W8").as_deref() == Ok("1"))
1129}
1130
1131pub(crate) static GLM5_W8_DISPATCHES: std::sync::atomic::AtomicU64 =
1134 std::sync::atomic::AtomicU64::new(0);
1135
1136pub(crate) fn w8_view_on() -> bool {
1141 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1142 *ON.get_or_init(|| std::env::var("MEMRA_W8_VIEW").as_deref() == Ok("1"))
1143}
1144
1145pub(crate) fn step_gemm_prime_on() -> bool {
1159 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1160 step37_door(&ENV, "MEMRA_STEP_GEMM_PRIME")
1161}
1162
1163pub(crate) fn q8t_wonce_on() -> bool {
1164 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1165 step37_door(&ENV, "MEMRA_Q8T_WONCE")
1166}
1167
1168pub(crate) fn sig_expf_dev_on() -> bool {
1173 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1174 *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
1175}
1176
1177pub(crate) fn topk_fast_on() -> bool {
1178 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1179 *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
1180}
1181
1182fn sigmoid_topk_kernel(sig_expf: bool, fast: bool, n_used: usize) -> &'static str {
1187 match (sig_expf, fast && n_used <= 8) {
1188 (true, true) => "moe_router_sigmoid_topk_f32_dexp_fast",
1189 (true, false) => "moe_router_sigmoid_topk_f32_dexp",
1190 (false, true) => "moe_router_sigmoid_topk_f32_fast",
1191 (false, false) => "moe_router_sigmoid_topk_f32",
1192 }
1193}
1194
1195#[cfg(test)]
1196mod sigmoid_topk_dispatch_tests {
1197 #[test]
1198 fn fast_kernel_refuses_wide_topk_and_composes_with_dexp() {
1199 use super::sigmoid_topk_kernel;
1200
1201 assert_eq!(
1202 sigmoid_topk_kernel(false, true, 8),
1203 "moe_router_sigmoid_topk_f32_fast"
1204 );
1205 assert_eq!(
1206 sigmoid_topk_kernel(true, true, 8),
1207 "moe_router_sigmoid_topk_f32_dexp_fast"
1208 );
1209 assert_eq!(
1210 sigmoid_topk_kernel(false, true, 9),
1211 "moe_router_sigmoid_topk_f32"
1212 );
1213 assert_eq!(
1214 sigmoid_topk_kernel(true, true, 9),
1215 "moe_router_sigmoid_topk_f32_dexp"
1216 );
1217 }
1218}
1219
1220pub(crate) fn rms_block() -> u32 {
1221 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1222 *V.get_or_init(|| {
1223 std::env::var("MEMRA_RMS_BLOCK")
1224 .ok()
1225 .and_then(|v| v.parse().ok())
1226 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1227 })
1228}
1229
1230pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
1231 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1232 if let Some(forced) = *S.get_or_init(|| {
1233 std::env::var("MEMRA_FA_SPLIT")
1234 .ok()
1235 .and_then(|v| v.parse().ok())
1236 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
1237 }) {
1238 return forced;
1239 }
1240 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
1258 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
1259 {
1260 return if t_kv <= 8192 {
1261 16
1262 } else if t_kv <= 16384 {
1263 64
1264 } else {
1265 128
1266 };
1267 }
1268 let big_rig = fa_sm_count() >= 128;
1269 if big_rig {
1270 let _ = n_head_kv;
1271 if t_kv <= 2048 {
1272 static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1281 if let Some(sp) = *SHORT.get_or_init(|| {
1282 std::env::var("MEMRA_FA_SP_SHORT")
1283 .ok()
1284 .and_then(|v| v.parse().ok())
1285 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
1286 }) {
1287 return sp;
1288 }
1289 16
1290 } else if t_kv <= 16384 {
1291 64
1292 } else {
1293 128
1294 }
1295 } else if n_head_kv <= 4 {
1296 if t_kv <= 512 {
1317 8
1318 } else if t_kv <= 16384 {
1319 64
1320 } else {
1321 128
1322 }
1323 } else {
1324 if t_kv <= 8192 {
1325 32
1326 } else if t_kv <= 16384 {
1327 64
1328 } else {
1329 128
1330 }
1331 }
1332}
1333
1334pub(crate) fn fa_sm_count() -> i32 {
1337 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
1338 *N.get_or_init(|| {
1339 cudarc::driver::result::init().ok();
1340 cudarc::driver::result::device::get(0)
1341 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
1342 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
1343 .unwrap_or(82)
1344 })
1345}
1346
1347#[allow(clippy::type_complexity)] fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
1352 match head_dim {
1353 256 => Ok(""),
1354 128 => Ok("_hd128"),
1355 d => Err(format!(
1356 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
1357 callers must gate to sdpa_naive"
1358 )
1359 .into()),
1360 }
1361}
1362
1363pub const QT_Q8_0: i32 = 0;
1365pub const QT_Q4_K: i32 = 1;
1366pub const QT_Q6_K: i32 = 2;
1367pub const QT_Q5_K: i32 = 3;
1368pub const QT_Q3_K: i32 = 4;
1369pub const QT_IQ4_XS: i32 = 5;
1370pub const QT_IQ3_S: i32 = 6;
1371pub const QT_NVFP4: i32 = 7;
1372pub const QT_NVFP4_V2: i32 = 107;
1375pub const QT_F8_E4M3: i32 = 10;
1381pub const QT_NVFP4_RP: i32 = 9;
1384pub const QT_F32: i32 = 8;
1386pub const QT_BF16: i32 = 11;
1387pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
1392pub const QT_F8_E4M3_BLK: i32 = 14;
1408
1409pub struct Engine {
1411 pub gpu: memra_runtime::Gpu,
1412 module: Arc<CudaModule>,
1413 hybrid: Arc<CudaModule>,
1414 kda: Arc<CudaModule>,
1416 qmatvec: Arc<CudaModule>,
1417 flash: Arc<CudaModule>,
1418 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
1422 gemm: Arc<CudaModule>,
1423 router: Arc<CudaModule>,
1424 sample: Arc<CudaModule>,
1426 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
1430 w8_mirrors: Mutex<std::collections::HashMap<(u64, u32, u32), CudaSlice<u8>>>,
1441 #[allow(clippy::type_complexity)]
1444 w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
1446 moe_cache_layout: Mutex<Option<Vec<usize>>>,
1450 capture_keep_on: std::sync::atomic::AtomicBool,
1456 verify_exact: std::sync::atomic::AtomicBool,
1461 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
1462 pub copy_stream: Arc<CudaStream>,
1464 pub side: std::sync::OnceLock<(Arc<CudaStream>, Arc<cudarc::cublaslt::CudaBlasLT>)>,
1472 #[cfg(memra_cutlass)]
1479 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
1480 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
1484 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
1487 #[allow(clippy::type_complexity)]
1491 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
1493 #[allow(clippy::type_complexity)]
1497 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
1499 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
1501 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
1502 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
1507 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
1512 router_stage: Mutex<Option<PinnedStage>>,
1516 hyper_decode_ws: Mutex<Option<crate::hyper::HyperDecodeWs>>,
1522 mla_seg_ws: Mutex<Option<crate::hybrid_forward::MlaSegWs>>,
1525 pub(crate) mla_pre_done: std::sync::atomic::AtomicBool,
1528 verify_ws: Mutex<VerifyWs>,
1540 vrows_macro_dev: Mutex<std::collections::HashMap<(u16, u8), CudaSlice<f32>>>,
1548 shexp_ones: Mutex<Option<CudaSlice<f32>>>,
1556}
1557
1558#[derive(Default)]
1563pub struct VerifyWs {
1564 f32_pool: std::collections::HashMap<usize, Vec<CudaSlice<f32>>>,
1565 i8_pool: std::collections::HashMap<usize, Vec<CudaSlice<i8>>>,
1566 u64_pool: std::collections::HashMap<usize, Vec<CudaSlice<u64>>>,
1567 held_bytes: usize,
1568}
1569
1570const VWS_PER_CLASS_CAP: usize = 16;
1573const VWS_HELD_BYTES_CAP: usize = 256 << 20;
1576
1577impl VerifyWs {
1578 fn take<T>(
1579 pool: &mut std::collections::HashMap<usize, Vec<CudaSlice<T>>>,
1580 held: &mut usize,
1581 n: usize,
1582 ) -> Option<CudaSlice<T>> {
1583 let s = pool.get_mut(&n)?.pop()?;
1584 *held -= n * std::mem::size_of::<T>();
1585 Some(s)
1586 }
1587 fn put<T>(
1588 pool: &mut std::collections::HashMap<usize, Vec<CudaSlice<T>>>,
1589 held: &mut usize,
1590 s: CudaSlice<T>,
1591 ) {
1592 let n = s.len();
1593 let bytes = n * std::mem::size_of::<T>();
1594 if *held + bytes > VWS_HELD_BYTES_CAP {
1595 return; }
1597 let v = pool.entry(n).or_default();
1598 if v.len() >= VWS_PER_CLASS_CAP {
1599 return;
1600 }
1601 v.push(s);
1602 *held += bytes;
1603 }
1604}
1605
1606pub static SCRATCH_ALLOC_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1612
1613pub(crate) fn fa_part_zero_on() -> bool {
1616 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1617 *ON.get_or_init(|| std::env::var("MEMRA_FA_PART_ZERO").as_deref() == Ok("1"))
1618}
1619
1620pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
1628 std::sync::atomic::AtomicUsize::new(1024);
1629pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
1630 std::sync::atomic::AtomicUsize::new(usize::MAX);
1631pub fn fa_v4_at_pub(t_kv: usize) -> bool {
1632 fa_v4_at(t_kv)
1633}
1634fn fa_v4_at(t_kv: usize) -> bool {
1635 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
1636 let mx = *M.get_or_init(|| {
1637 std::env::var("MEMRA_FA_V4_MAX")
1638 .ok()
1639 .and_then(|v| v.parse().ok())
1640 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1641 });
1642 t_kv < mx
1643}
1644fn fa_v3_active(head_dim: usize) -> bool {
1650 head_dim.is_multiple_of(128)
1651}
1652
1653pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1660 std::env::var("MEMRA_NO_FA_VEC").is_err()
1661 && t_kv >= fa_vec_min_tkv()
1662 && head_dim == 256
1663 && fa_v4_at(t_kv)
1664}
1665pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1667 fa_split_keys(t_kv, n_head_kv)
1668}
1669
1670struct PinnedStage {
1675 ptr: *mut u8,
1676 cap: usize,
1677}
1678unsafe impl Send for PinnedStage {}
1679impl PinnedStage {
1680 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1681 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1682 Ok(PinnedStage { ptr, cap })
1683 }
1684}
1685impl Drop for PinnedStage {
1686 fn drop(&mut self) {
1687 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1688 }
1689}
1690
1691pub struct PinnedHostBuf {
1698 ptr: *mut u8,
1699 len: usize,
1700}
1701unsafe impl Send for PinnedHostBuf {}
1704impl PinnedHostBuf {
1705 pub fn new(len: usize) -> Result<Self, Box<dyn std::error::Error>> {
1708 let ptr = unsafe { cudarc::driver::result::malloc_host(len.max(1), 0)? } as *mut u8;
1709 Ok(PinnedHostBuf { ptr, len })
1710 }
1711 pub fn len(&self) -> usize {
1712 self.len
1713 }
1714 pub fn is_empty(&self) -> bool {
1715 self.len == 0
1716 }
1717 pub fn as_slice(&self) -> &[u8] {
1718 unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
1719 }
1720 pub fn as_mut_slice(&mut self) -> &mut [u8] {
1721 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
1722 }
1723}
1724impl Drop for PinnedHostBuf {
1725 fn drop(&mut self) {
1726 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1727 }
1728}
1729
1730pub const ARGMAX_NB: usize = 256;
1733
1734pub(crate) use memra_fa3_vl as fa3_vl_raw;
1736
1737unsafe extern "C" {
1738 fn memra_fa3_prefill(
1740 q16: *const core::ffi::c_void,
1741 k16: *const core::ffi::c_void,
1742 v16: *const core::ffi::c_void,
1743 o: *mut f32,
1744 t: i32,
1745 h: i32,
1746 hkv: i32,
1747 d: i32,
1748 scale: f32,
1749 stream: *mut core::ffi::c_void,
1750 ) -> i32;
1751 pub(crate) fn memra_fa3_vl(
1753 q16s: *const *const core::ffi::c_void,
1754 k16s: *const *const core::ffi::c_void,
1755 v16s: *const *const core::ffi::c_void,
1756 os: *const *mut f32,
1757 ts: *const i32,
1758 b: i32,
1759 h: i32,
1760 hkv: i32,
1761 d: i32,
1762 scale: f32,
1763 stream: *mut core::ffi::c_void,
1764 ) -> i32;
1765}
1766
1767#[repr(C)]
1772#[derive(Clone, Copy)]
1773pub struct WPtr8(pub [u64; 8]);
1774unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1775
1776#[repr(C)]
1781#[derive(Clone, Copy, Default)]
1782pub struct GdnSeqVl {
1783 pub kb16: u64,
1784 pub gcum: u64,
1785 pub beta: u64,
1786 pub u: u64,
1787 pub wb16: u64,
1788 pub y: u64,
1789 pub ssnap: u64,
1790 pub state_in: u64,
1791 pub state_out: u64,
1792 pub q: u64,
1793 pub p: u64,
1794 pub o: u64,
1795 pub k: u64,
1796 pub v: u64,
1797 pub g: u64,
1798 pub a: u64,
1799 pub w: u64,
1800 pub t: i32,
1801 pub nc: i32,
1802}
1803unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1804#[repr(C)]
1805#[derive(Clone, Copy)]
1806pub struct GdnVl8(pub [GdnSeqVl; 8]);
1807unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1808
1809#[repr(C)]
1812#[derive(Clone, Copy, Default)]
1813pub struct GdnWVl {
1814 pub qb16: u64,
1815 pub pb16: u64,
1816}
1817unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1818#[repr(C)]
1819#[derive(Clone, Copy)]
1820pub struct GdnWVl8(pub [GdnWVl; 8]);
1821unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1822
1823#[repr(C)]
1825#[derive(Clone, Copy, Default)]
1826pub struct GdnPrepVl {
1827 pub qkv: u64,
1828 pub conv_state: u64,
1829 pub conv_out: u64,
1830 pub q_g: u64,
1831 pub k_g: u64,
1832 pub v_g: u64,
1833 pub q_l2: u64,
1834 pub k_l2: u64,
1835 pub beta_raw: u64,
1836 pub alpha: u64,
1837 pub beta: u64,
1838 pub g_log: u64,
1839 pub o: u64,
1840 pub z: u64,
1841 pub gn: u64,
1842 pub gn16: u64,
1843 pub kb16: u64,
1844 pub qb16: u64,
1845 pub t: i32,
1846 pub pad: i32,
1847}
1848unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1849#[repr(C)]
1850#[derive(Clone, Copy)]
1851pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1852unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1853
1854#[repr(C)]
1856#[derive(Clone, Copy, Default)]
1857pub struct FaSeqVl {
1858 pub q: u64,
1859 pub k16: u64,
1860 pub v16: u64,
1861 pub o: u64,
1862 pub kf: u64,
1863 pub vf: u64,
1864 pub t: i32,
1865 pub pad: i32,
1866}
1867unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1868#[repr(C)]
1869#[derive(Clone, Copy)]
1870pub struct FaVl8(pub [FaSeqVl; 8]);
1871unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1872
1873#[repr(C)]
1875#[derive(Clone, Copy, Default)]
1876pub struct AttnPreVl {
1877 pub qf: u64,
1878 pub kf: u64,
1879 pub vf: u64,
1880 pub q: u64,
1881 pub gate: u64,
1882 pub qn: u64,
1883 pub kn: u64,
1884 pub kc: u64,
1885 pub vc: u64,
1886 pub t: i32,
1887 pub pad: i32,
1888}
1889unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1890#[repr(C)]
1891#[derive(Clone, Copy)]
1892pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1893unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1894
1895pub struct GdnChunkBufs {
1898 pub gcum: CudaSlice<f32>,
1899 pub a: CudaSlice<f32>,
1900 pub p: CudaSlice<f32>,
1901 pub u: CudaSlice<f32>,
1902 pub w: CudaSlice<f32>,
1903 pub kb16: CudaSlice<u8>,
1904 pub wb16: CudaSlice<u8>,
1905 pub y16: CudaSlice<u8>,
1906 pub ssnap16: CudaSlice<u8>,
1907 pub qb16: CudaSlice<u8>,
1908 pub pb16: CudaSlice<u8>,
1909 pub o: CudaSlice<f32>,
1910 pub t: usize,
1911 pub nc: usize,
1912}
1913
1914#[repr(C)]
1916#[derive(Clone, Copy)]
1917pub struct F32x8(pub [f32; 8]);
1918unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1919
1920pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1924
1925pub static MOE_FUSED_EPI_DISPATCHES: std::sync::atomic::AtomicU64 =
1934 std::sync::atomic::AtomicU64::new(0);
1935
1936pub fn moe_fused_epilogue_dispatches() -> u64 {
1940 MOE_FUSED_EPI_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1941}
1942
1943pub static MOE_VROWS_DISPATCHES: std::sync::atomic::AtomicU64 =
1950 std::sync::atomic::AtomicU64::new(0);
1951
1952pub fn moe_vrows_dispatches() -> u64 {
1954 MOE_VROWS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1955}
1956
1957fn bf16_tcols_wide_on() -> bool {
1966 std::env::var("MEMRA_BF16_TCOLS_WIDE").as_deref() != Ok("0")
1967}
1968
1969pub static BF16_TCOLS_WIDE_DISPATCHES: std::sync::atomic::AtomicU64 =
1973 std::sync::atomic::AtomicU64::new(0);
1974
1975pub fn bf16_tcols_wide_dispatches() -> u64 {
1977 BF16_TCOLS_WIDE_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1978}
1979
1980fn bf16_tcols_x1_on() -> bool {
1988 std::env::var("MEMRA_BF16_TCOLS_X1").as_deref() != Ok("0")
1989}
1990
1991pub static BF16_TCOLS_X1_DISPATCHES: std::sync::atomic::AtomicU64 =
1993 std::sync::atomic::AtomicU64::new(0);
1994
1995pub fn bf16_tcols_x1_dispatches() -> u64 {
1997 BF16_TCOLS_X1_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1998}
1999
2000fn bf16_tcols_red_fused_on() -> bool {
2013 std::env::var("MEMRA_BF16_TCOLS_RED_FUSED").as_deref() == Ok("1")
2014}
2015
2016pub static BF16_TCOLS_RED_FUSED_DISPATCHES: std::sync::atomic::AtomicU64 =
2019 std::sync::atomic::AtomicU64::new(0);
2020
2021pub fn bf16_tcols_red_fused_dispatches() -> u64 {
2023 BF16_TCOLS_RED_FUSED_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2024}
2025
2026pub(crate) fn moe_vrows_pack_on() -> bool {
2033 std::env::var("MEMRA_MOE_VROWS_PACK").as_deref() == Ok("1")
2034}
2035
2036pub static MOE_VROWS_PACK_DISPATCHES: std::sync::atomic::AtomicU64 =
2038 std::sync::atomic::AtomicU64::new(0);
2039
2040pub fn moe_vrows_pack_dispatches() -> u64 {
2042 MOE_VROWS_PACK_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2043}
2044
2045pub(crate) fn moe_vrows_ilp_on() -> bool {
2052 moe_vrows_ilp_on_from(
2053 std::env::var("MEMRA_MOE_VROWS_ILP").ok().as_deref(),
2054 env!("MEMRA_BUILT_CUDA_ARCH"),
2055 )
2056}
2057
2058pub fn moe_vrows_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
2063 match v.map(str::trim) {
2064 Some("1") => true,
2065 Some("0") => false,
2066 _ => built_arch == "100a",
2067 }
2068}
2069
2070pub static MOE_VROWS_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
2072 std::sync::atomic::AtomicU64::new(0);
2073
2074pub fn moe_vrows_ilp_dispatches() -> u64 {
2076 MOE_VROWS_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2077}
2078
2079fn moe_vrows_ilp_refuse(site: &str, qt: i32) {
2082 static SAID: std::sync::Once = std::sync::Once::new();
2083 SAID.call_once(|| {
2084 eprintln!(
2085 "[moe-vrows-ilp] REFUSED at {site}: MEMRA_MOE_VROWS_ILP is on but the expert qtype is \
2086 {qt}, neither QT_NVFP4 ({QT_NVFP4}) nor QT_NVFP4_V2 ({QT_NVFP4_V2}); the shipped \
2087 kernel runs (the ILP twins hoist the NVFP4 group loads, interleaved or slot-major, \
2088 and no other layout)"
2089 );
2090 });
2091}
2092
2093pub(crate) fn moe_vrows_dev_tables_on() -> bool {
2103 std::env::var("MEMRA_MOE_VROWS_DEV_TABLES").as_deref() == Ok("1")
2104}
2105
2106pub static MOE_VROWS_DEV_TABLES_DISPATCHES: std::sync::atomic::AtomicU64 =
2108 std::sync::atomic::AtomicU64::new(0);
2109
2110pub fn moe_vrows_dev_tables_dispatches() -> u64 {
2112 MOE_VROWS_DEV_TABLES_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2113}
2114
2115pub static MOE_VROWS_ROUTER_SYNCS_AVOIDED: std::sync::atomic::AtomicU64 =
2119 std::sync::atomic::AtomicU64::new(0);
2120
2121pub fn moe_vrows_router_syncs_avoided() -> u64 {
2123 MOE_VROWS_ROUTER_SYNCS_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2124}
2125
2126fn moe_vrows_dedup_stat_on() -> bool {
2136 std::env::var("MEMRA_MOE_VROWS_DEDUP_STAT").as_deref() == Ok("1")
2137}
2138
2139pub static MOE_VROWS_PAIR_VISITS: std::sync::atomic::AtomicU64 =
2141 std::sync::atomic::AtomicU64::new(0);
2142
2143pub static MOE_VROWS_PAIR_DISTINCT: std::sync::atomic::AtomicU64 =
2147 std::sync::atomic::AtomicU64::new(0);
2148
2149pub(crate) fn vrows_overlap_counts(sel_all: &[u32]) -> (u64, u64) {
2156 let mut seen = std::collections::HashSet::with_capacity(sel_all.len());
2157 for &ex in sel_all {
2158 seen.insert(ex);
2159 }
2160 (sel_all.len() as u64, seen.len() as u64)
2161}
2162
2163static MOE_VROWS_DEDUP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2165
2166fn moe_vrows_dedup_report() {
2173 let n = MOE_VROWS_DEDUP_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2174 if n != 0 && !n.is_multiple_of(42) {
2175 return;
2176 }
2177 let (visits, distinct) = moe_vrows_pair_overlap();
2178 if visits == 0 {
2179 return;
2180 }
2181 let repeat = 100.0 * (1.0 - distinct as f64 / visits as f64);
2182 eprintln!(
2183 "[moe-vrows-dedup] layer-calls={} visits={visits} distinct={distinct} \
2184 repeat={repeat:.2}% = the cross-row expert-slab dedup ceiling on the vrows pair \
2185 (MEMRA_MOE_VROWS_DEDUP_STAT=1)",
2186 n + 1
2187 );
2188}
2189
2190pub fn vrows_overlap_counts_for_test(sel_all: &[u32]) -> (u64, u64) {
2193 vrows_overlap_counts(sel_all)
2194}
2195
2196pub fn moe_vrows_pair_overlap() -> (u64, u64) {
2198 (
2199 MOE_VROWS_PAIR_VISITS.load(std::sync::atomic::Ordering::Relaxed),
2200 MOE_VROWS_PAIR_DISTINCT.load(std::sync::atomic::Ordering::Relaxed),
2201 )
2202}
2203
2204pub(crate) fn moe_vrows_dedup_order_on() -> bool {
2221 std::env::var("MEMRA_MOE_VROWS_DEDUP_ORDER").as_deref() == Ok("1")
2222}
2223
2224fn moe_vrows_down_tmaj_on() -> bool {
2233 std::env::var("MEMRA_MOE_VROWS_DOWN_TMAJ").as_deref() == Ok("1")
2234}
2235
2236pub static MOE_VROWS_DEDUP_ORDER_DISPATCHES: std::sync::atomic::AtomicU64 =
2238 std::sync::atomic::AtomicU64::new(0);
2239
2240pub fn moe_vrows_dedup_order_dispatches() -> u64 {
2242 MOE_VROWS_DEDUP_ORDER_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2243}
2244
2245pub static MOE_VROWS_DOWN_TMAJ_DISPATCHES: std::sync::atomic::AtomicU64 =
2247 std::sync::atomic::AtomicU64::new(0);
2248
2249pub fn moe_vrows_down_tmaj_dispatches() -> u64 {
2251 MOE_VROWS_DOWN_TMAJ_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2252}
2253
2254pub static MOE_VROWS_SLAB_READS_AVOIDED: std::sync::atomic::AtomicU64 =
2265 std::sync::atomic::AtomicU64::new(0);
2266
2267pub fn moe_vrows_slab_reads_avoided() -> u64 {
2269 MOE_VROWS_SLAB_READS_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2270}
2271
2272pub(crate) fn vrows_expert_major_order(sel_all: &[u32]) -> Vec<u64> {
2277 let mut ord: Vec<u64> = (0..sel_all.len() as u64).collect();
2278 ord.sort_by_key(|&p| sel_all[p as usize]);
2281 ord
2282}
2283
2284pub fn vrows_expert_major_order_for_test(sel_all: &[u32]) -> Vec<u64> {
2287 vrows_expert_major_order(sel_all)
2288}
2289
2290pub(crate) fn alias_door_from(
2311 general: (&'static str, Option<&str>),
2312 alias: (&'static str, Option<&str>),
2313) -> Result<(bool, &'static str), String> {
2314 match (general.1, alias.1) {
2315 (Some(g), Some(a)) if g != a => Err(format!(
2316 "{}={g:?} and {}={a:?} disagree — the alias and the general flag name ONE door \
2317 (unset one); refused rather than silently picking a precedence winner, and the \
2318 door falls closed to the shipped program",
2319 general.0, alias.0
2320 )),
2321 (Some(g), _) => Ok((g == "1", general.0)),
2322 (None, Some(a)) => Ok((a == "1", alias.0)),
2323 (None, None) => Ok((false, general.0)),
2324 }
2325}
2326
2327fn alias_door(
2344 general: &'static str,
2345 alias: &'static str,
2346 latch: &'static std::sync::atomic::AtomicBool,
2347) -> (bool, &'static str) {
2348 let g = std::env::var(general).ok();
2349 let a = std::env::var(alias).ok();
2350 match alias_door_from((general, g.as_deref()), (alias, a.as_deref())) {
2351 Ok(resolved) => resolved,
2352 Err(msg) => {
2353 if !latch.swap(true, std::sync::atomic::Ordering::Relaxed) {
2354 eprintln!("[flag-alias] {msg}");
2355 }
2356 (false, general)
2357 }
2358 }
2359}
2360
2361pub fn htod_diet_on() -> bool {
2381 htod_diet_armed().0
2382}
2383
2384static HTOD_DIET_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2386 std::sync::atomic::AtomicBool::new(false);
2387
2388pub(crate) fn htod_diet_armed() -> (bool, &'static str) {
2390 alias_door(
2391 "MEMRA_HTOD_DIET",
2392 "MEMRA_GLM5_HTOD_DIET",
2393 &HTOD_DIET_ALIAS_WARNED,
2394 )
2395}
2396
2397pub static HTOD_DIET_AVOIDED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2400
2401pub fn htod_diet_avoided() -> u64 {
2403 HTOD_DIET_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2404}
2405
2406pub fn ep_diet_on() -> bool {
2427 ep_diet_armed().0
2428}
2429
2430static EP_DIET_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2432 std::sync::atomic::AtomicBool::new(false);
2433
2434pub(crate) fn ep_diet_armed() -> (bool, &'static str) {
2437 alias_door("MEMRA_EP_DIET", "MEMRA_GLM5_EP_DIET", &EP_DIET_ALIAS_WARNED)
2438}
2439
2440pub fn ep_grouped_prime_on() -> bool {
2458 ep_grouped_prime_armed().0
2459}
2460
2461static EP_GROUPED_PRIME_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2463 std::sync::atomic::AtomicBool::new(false);
2464
2465pub(crate) fn ep_grouped_prime_armed() -> (bool, &'static str) {
2467 alias_door(
2468 "MEMRA_EP_GROUPED_PRIME",
2469 "MEMRA_GLM5_EP_GROUPED_PRIME",
2470 &EP_GROUPED_PRIME_ALIAS_WARNED,
2471 )
2472}
2473
2474fn topk_shards_on() -> bool {
2483 std::env::var("MEMRA_TOPK_SHARDS").as_deref() != Ok("0")
2484}
2485
2486pub static TOPK_SHARDS_DISPATCHES: std::sync::atomic::AtomicU64 =
2488 std::sync::atomic::AtomicU64::new(0);
2489
2490pub fn topk_shards_dispatches() -> u64 {
2492 TOPK_SHARDS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2493}
2494
2495pub fn alloc_trace_on() -> bool {
2502 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2503 *V.get_or_init(|| std::env::var("MEMRA_ALLOC_TRACE").as_deref() == Ok("1"))
2504}
2505
2506#[track_caller]
2507pub(crate) fn alloc_trace_hit(bytes: usize) {
2508 if alloc_trace_on() {
2509 let loc = std::panic::Location::caller();
2510 eprintln!(
2511 "[alloc-trace] {bytes} bytes from {}:{}",
2512 loc.file(),
2513 loc.line()
2514 );
2515 }
2516}
2517
2518pub fn dtoh_trace_on() -> bool {
2525 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2526 *V.get_or_init(|| std::env::var("MEMRA_DTOH_TRACE").as_deref() == Ok("1"))
2527}
2528
2529#[track_caller]
2530pub(crate) fn dtoh_trace_hit(bytes: usize) {
2531 if dtoh_trace_on() {
2532 let loc = std::panic::Location::caller();
2533 eprintln!(
2534 "[dtoh-trace] {bytes} bytes from {}:{}",
2535 loc.file(),
2536 loc.line()
2537 );
2538 }
2539}
2540
2541pub fn moe_expert_rp_on() -> bool {
2554 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2555 *V.get_or_init(|| std::env::var("MEMRA_MOE_EXPERT_RP").as_deref() == Ok("1"))
2556}
2557
2558pub fn rp_qt(rp: bool, qt: i32) -> i32 {
2562 if rp && qt == QT_NVFP4 {
2563 QT_NVFP4_V2
2564 } else {
2565 qt
2566 }
2567}
2568
2569pub fn moe_rp_refuse(rp: bool, path: &str) -> Result<(), Box<dyn std::error::Error>> {
2572 if rp {
2573 return Err(format!(
2574 "{path}: the resident expert slab is split-plane (MEMRA_MOE_EXPERT_RP=1) and this \
2575 path reads experts interleaved; it is not wired for the door (memra#147). Boot \
2576 without MEMRA_MOE_EXPERT_RP for this model or wire the path."
2577 )
2578 .into());
2579 }
2580 Ok(())
2581}
2582
2583pub(crate) fn hyper_batch_solo_on() -> bool {
2610 std::env::var("MEMRA_HYPER_BATCH_SOLO").as_deref() == Ok("1")
2611}
2612
2613pub(crate) fn hc_pre_sink_reg() -> bool {
2657 hc_pre_sink_reg_from(
2658 std::env::var("MEMRA_HC_PRE_SINK_REG").ok().as_deref(),
2659 env!("MEMRA_BUILT_CUDA_ARCH"),
2660 )
2661}
2662
2663pub fn hc_pre_sink_reg_from(v: Option<&str>, built_arch: &str) -> bool {
2667 match v.map(str::trim) {
2668 Some("1") => true,
2669 Some("0") => false,
2670 _ => built_arch == "100a",
2671 }
2672}
2673
2674pub fn hc_pre_v4_on() -> bool {
2681 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2682 *ON.get_or_init(|| std::env::var("MEMRA_HC_PRE_V4").as_deref() == Ok("1"))
2683}
2684
2685pub fn hc_pre_v4z_on() -> bool {
2691 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2692 *ON.get_or_init(|| std::env::var("MEMRA_HC_PRE_V4Z").as_deref() == Ok("1"))
2693}
2694
2695pub fn mla_absorb_bf16_on() -> bool {
2702 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2703 *ON.get_or_init(|| std::env::var("MEMRA_MLA_ABSORB_BF16").as_deref() == Ok("1"))
2704}
2705
2706pub fn moe_shexp_overlap_on() -> bool {
2713 matches!(
2714 std::env::var("MEMRA_MOE_SHEXP_OVERLAP").as_deref(),
2715 Ok("1") | Ok("2") | Ok("3")
2716 )
2717}
2718
2719pub fn moe_shexp_overlap_main_probe() -> bool {
2723 std::env::var("MEMRA_MOE_SHEXP_OVERLAP").as_deref() == Ok("3")
2724}
2725
2726pub fn moe_shexp_overlap_serial_probe() -> bool {
2731 std::env::var("MEMRA_MOE_SHEXP_OVERLAP").as_deref() == Ok("2")
2732}
2733
2734pub static MOE_SHEXP_OVERLAP_DISPATCHES: std::sync::atomic::AtomicU64 =
2736 std::sync::atomic::AtomicU64::new(0);
2737
2738pub fn hc_pre_zq8_on() -> bool {
2751 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2752 *ON.get_or_init(|| {
2756 matches!(
2757 std::env::var("MEMRA_HC_PRE_ZQ8").as_deref(),
2758 Ok("1") | Ok("2")
2759 )
2760 })
2761}
2762
2763fn q8_census_record(site: &'static std::panic::Location<'static>, m: usize, in_f: usize) {
2775 use std::collections::HashMap;
2776 use std::sync::{Mutex, OnceLock};
2777 static EVERY: OnceLock<u64> = OnceLock::new();
2782 let every = *EVERY.get_or_init(|| match std::env::var("MEMRA_Q8_CENSUS").ok().as_deref() {
2783 None | Some("0") | Some("") => 0,
2784 Some("1") => 256,
2785 Some(v) => v.parse().unwrap_or(256),
2786 });
2787 if every == 0 {
2788 return;
2789 }
2790 type Hist = (u64, HashMap<(String, u32, usize, usize), u64>);
2792 static HIST: OnceLock<Mutex<Hist>> = OnceLock::new();
2793 let mut g = HIST
2794 .get_or_init(|| Mutex::new((0, HashMap::new())))
2795 .lock()
2796 .unwrap();
2797 g.0 += 1;
2798 *g.1.entry((site.file().to_string(), site.line(), m, in_f))
2799 .or_insert(0) += 1;
2800 if g.0.is_multiple_of(every) {
2801 let mut rows: Vec<_> = g.1.iter().map(|(k, v)| (*v, k.clone())).collect();
2802 rows.sort_by_key(|r| std::cmp::Reverse(r.0));
2803 eprintln!(
2804 "[q8-census] {} quantize_q8_1 calls so far; by call site:",
2805 g.0
2806 );
2807 for (n, (file, line, m, in_f)) in rows {
2808 eprintln!(
2809 "[q8-census] {n:>8} {:5.1}% {file}:{line} m={m} in_f={in_f}",
2810 100.0 * n as f64 / g.0 as f64
2811 );
2812 }
2813 }
2814}
2815
2816pub fn hc_pre_block_default(built_arch: &str) -> usize {
2818 if built_arch == "100a" { 512 } else { 128 }
2819}
2820
2821pub(crate) fn hc_pre_block() -> usize {
2822 let default = hc_pre_block_default(env!("MEMRA_BUILT_CUDA_ARCH"));
2826 match std::env::var("MEMRA_HC_PRE_BLOCK") {
2827 Err(_) => default,
2828 Ok(v) if v.is_empty() => default,
2829 Ok(v) => match v.parse::<usize>() {
2830 Ok(n) if (32..=1024).contains(&n) && n.is_power_of_two() => n,
2831 _ => {
2832 static SAID: std::sync::Once = std::sync::Once::new();
2833 SAID.call_once(|| {
2834 eprintln!(
2835 "[hc-pre-block] MEMRA_HC_PRE_BLOCK={v:?} is not a power of two in \
2836 [32, 1024]; using the build default {default}"
2837 );
2838 });
2839 default
2840 }
2841 },
2842 }
2843}
2844
2845fn verify_ws_on() -> bool {
2846 verify_ws_on_from(
2847 std::env::var("MEMRA_VERIFY_WS").ok().as_deref(),
2848 std::env::var("MEMRA_GLM5_VERIFY_WS").ok().as_deref(),
2849 )
2850}
2851
2852fn verify_ws_on_from(general: Option<&str>, glm5_alias: Option<&str>) -> bool {
2855 general != Some("0") && glm5_alias != Some("0")
2856}
2857
2858pub static VERIFY_WS_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2862
2863pub fn verify_ws_hits() -> u64 {
2865 VERIFY_WS_HITS.load(std::sync::atomic::Ordering::Relaxed)
2866}
2867
2868pub static MLA_TC_PREFILL_DISPATCHES: std::sync::atomic::AtomicU64 =
2875 std::sync::atomic::AtomicU64::new(0);
2876
2877pub fn mla_tc_prefill_dispatches() -> u64 {
2881 MLA_TC_PREFILL_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2882}
2883
2884pub static MOE_GROUPED_PREFILL_DISPATCHES: std::sync::atomic::AtomicU64 =
2890 std::sync::atomic::AtomicU64::new(0);
2891
2892pub fn moe_grouped_prefill_dispatches() -> u64 {
2895 MOE_GROUPED_PREFILL_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2896}
2897
2898#[must_use = "dropping immediately ends the exact scope"]
2903pub struct ExactScope<'a> {
2904 flag: &'a std::sync::atomic::AtomicBool,
2905 prev: bool,
2906}
2907
2908impl<'a> ExactScope<'a> {
2909 pub(crate) fn set(flag: &'a std::sync::atomic::AtomicBool, on: bool) -> Self {
2910 let prev = flag.load(std::sync::atomic::Ordering::Relaxed);
2911 flag.store(on, std::sync::atomic::Ordering::Relaxed);
2912 ExactScope { flag, prev }
2913 }
2914}
2915
2916impl Drop for ExactScope<'_> {
2917 fn drop(&mut self) {
2918 self.flag
2919 .store(self.prev, std::sync::atomic::Ordering::Relaxed);
2920 }
2921}
2922
2923#[cfg(test)]
2924mod verify_ws_flag_tests {
2925 use super::verify_ws_on_from;
2926
2927 #[test]
2928 fn off_wins_across_general_and_alias() {
2929 assert!(verify_ws_on_from(None, None));
2931 assert!(!verify_ws_on_from(Some("0"), None));
2934 assert!(!verify_ws_on_from(None, Some("0")));
2935 assert!(!verify_ws_on_from(Some("1"), Some("0")));
2936 assert!(!verify_ws_on_from(Some("0"), Some("1")));
2937 assert!(verify_ws_on_from(Some("1"), None));
2939 assert!(verify_ws_on_from(None, Some("1")));
2940 }
2941}
2942
2943#[cfg(test)]
2944mod moe_vrows_ilp_default_tests {
2945 use super::moe_vrows_ilp_on_from;
2946
2947 #[test]
2948 fn nvfp4_row_ilp_arch_keyed_default() {
2949 use super::nvfp4_row_ilp_on_from;
2950 assert!(nvfp4_row_ilp_on_from(None, "100a"));
2951 assert!(!nvfp4_row_ilp_on_from(None, "120a"));
2952 assert!(nvfp4_row_ilp_on_from(Some("1"), "120a"));
2953 assert!(!nvfp4_row_ilp_on_from(Some("0"), "100a"));
2954 }
2955
2956 #[test]
2957 fn q8_row_ilp_arch_keyed_default() {
2958 use super::q8_row_ilp_on_from;
2959 assert!(q8_row_ilp_on_from(None, "100a"));
2960 assert!(!q8_row_ilp_on_from(None, "120a"));
2961 assert!(q8_row_ilp_on_from(Some("1"), "120a"));
2962 assert!(!q8_row_ilp_on_from(Some("0"), "100a"));
2963 }
2964
2965 #[test]
2966 fn arch_keyed_default_with_explicit_override() {
2967 assert!(moe_vrows_ilp_on_from(None, "100a"));
2968 assert!(!moe_vrows_ilp_on_from(None, "120a"));
2969 assert!(!moe_vrows_ilp_on_from(None, "90a"));
2970 assert!(moe_vrows_ilp_on_from(Some("1"), "120a"));
2971 assert!(!moe_vrows_ilp_on_from(Some("0"), "100a"));
2972 assert!(!moe_vrows_ilp_on_from(Some(" 0 "), "100a"));
2973 }
2974}
2975
2976#[cfg(test)]
2977mod hc_pre_default_tests {
2978 use super::{hc_pre_block_default, hc_pre_sink_reg_from};
2979
2980 #[test]
2981 fn arch_keyed_defaults_with_explicit_override() {
2982 assert_eq!(hc_pre_block_default("100a"), 512);
2983 assert_eq!(hc_pre_block_default("120a"), 128);
2984 assert!(hc_pre_sink_reg_from(None, "100a"));
2985 assert!(!hc_pre_sink_reg_from(None, "120a"));
2986 assert!(hc_pre_sink_reg_from(Some("1"), "120a"));
2987 assert!(!hc_pre_sink_reg_from(Some("0"), "100a"));
2988 assert_eq!(
2989 crate::hyper::hc_fused_pre_arm_from(None, "100a"),
2990 crate::hyper::HcFusedPreArm::V2
2991 );
2992 assert_eq!(
2993 crate::hyper::hc_fused_pre_arm_from(None, "120a"),
2994 crate::hyper::HcFusedPreArm::Off
2995 );
2996 assert_eq!(
2997 crate::hyper::hc_fused_pre_arm_from(Some("0"), "100a"),
2998 crate::hyper::HcFusedPreArm::Off
2999 );
3000 assert_eq!(
3001 crate::hyper::hc_fused_pre_arm_from(Some("1"), "100a"),
3002 crate::hyper::HcFusedPreArm::V1
3003 );
3004 }
3005}
3006
3007#[cfg(test)]
3008mod glm5_decode_graph_default_tests {
3009 use super::glm5_decode_graph_on_from;
3010
3011 #[test]
3014 fn unset_and_one_arm_only_zero_disarms() {
3015 assert!(glm5_decode_graph_on_from(None));
3016 assert!(glm5_decode_graph_on_from(Some("1")));
3017 assert!(glm5_decode_graph_on_from(Some("")));
3018 assert!(!glm5_decode_graph_on_from(Some("0")));
3019 assert!(!glm5_decode_graph_on_from(Some(" 0 ")));
3020 }
3021}
3022
3023#[cfg(test)]
3024mod alias_door_tests {
3025 use super::alias_door_from;
3026
3027 const G: &str = "MEMRA_EP_DIET";
3028 const A: &str = "MEMRA_GLM5_EP_DIET";
3029
3030 fn r(g: Option<&str>, a: Option<&str>) -> Result<(bool, &'static str), String> {
3031 alias_door_from((G, g), (A, a))
3032 }
3033
3034 #[test]
3035 fn default_off_and_either_name_arms() {
3036 assert_eq!(r(None, None).unwrap(), (false, G));
3038 assert_eq!(r(Some("1"), None).unwrap(), (true, G));
3040 assert_eq!(r(None, Some("1")).unwrap(), (true, A));
3041 assert_eq!(r(Some("0"), None).unwrap(), (false, G));
3043 assert_eq!(r(None, Some("0")).unwrap(), (false, A));
3044 assert_eq!(r(None, Some("on")).unwrap(), (false, A));
3046 assert_eq!(r(Some(""), None).unwrap(), (false, G));
3047 }
3048
3049 #[test]
3050 fn agreeing_pair_resolves_to_the_general_name() {
3051 assert_eq!(r(Some("1"), Some("1")).unwrap(), (true, G));
3052 assert_eq!(r(Some("0"), Some("0")).unwrap(), (false, G));
3053 }
3054
3055 #[test]
3056 fn disagreeing_pair_refuses_and_names_both() {
3057 for (g, a) in [("1", "0"), ("0", "1")] {
3058 let err = r(Some(g), Some(a)).expect_err("a disagreeing pair must refuse");
3059 assert!(
3060 err.contains(G),
3061 "the refusal must name the general flag: {err}"
3062 );
3063 assert!(err.contains(A), "the refusal must name the alias: {err}");
3064 assert!(err.contains("falls closed"), "{err}");
3067 }
3068 }
3069}
3070
3071#[cfg(test)]
3072mod exact_scope_tests {
3073 use std::sync::atomic::{AtomicBool, Ordering};
3074
3075 #[test]
3076 fn error_path_restores_verify_exact() {
3077 let flag = AtomicBool::new(false);
3082 let failing = |flag: &AtomicBool| -> Result<(), &'static str> {
3083 let _scope = super::ExactScope::set(flag, true);
3084 assert!(flag.load(Ordering::Relaxed), "scope arms the flag");
3085 Err("draft forward failed")? };
3087 assert!(failing(&flag).is_err());
3088 assert!(
3089 !flag.load(Ordering::Relaxed),
3090 "error propagation must restore the pre-scope value"
3091 );
3092 let flag = AtomicBool::new(true);
3094 {
3095 let _scope = super::ExactScope::set(&flag, true);
3096 }
3097 assert!(flag.load(Ordering::Relaxed));
3098 let flag = AtomicBool::new(false);
3100 let scope = super::ExactScope::set(&flag, true);
3101 drop(scope);
3102 assert!(!flag.load(Ordering::Relaxed));
3103 }
3104}
3105
3106impl Engine {
3107 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
3108 let gpu = memra_runtime::Gpu::new(ordinal)?;
3109 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
3113 use cudarc::driver::sys::CUdevice_attribute_enum as A;
3114 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
3115 .and_then(|d| unsafe {
3116 Ok((
3117 cudarc::driver::result::device::get_attribute(
3118 d,
3119 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
3120 )?,
3121 cudarc::driver::result::device::get_attribute(
3122 d,
3123 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
3124 )?,
3125 ))
3126 })
3127 .unwrap_or((0, 0));
3128 let built = env!("MEMRA_BUILT_CUDA_ARCH");
3129 let ok = matches!(
3130 (built, maj, min),
3131 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
3132 );
3133 if !ok {
3134 return Err(format!(
3135 "memra was built for sm_{built} but device {ordinal} reports compute \
3136 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
3137 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
3138 )
3139 .into());
3140 }
3141 }
3142 unsafe {
3147 use cudarc::driver::sys;
3148 let dev: sys::CUdevice = ordinal as sys::CUdevice;
3149 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3150 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
3151 let mut thresh: u64 = u64::MAX;
3152 let _ = sys::cuMemPoolSetAttribute(
3153 pool,
3154 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
3155 &mut thresh as *mut u64 as *mut core::ffi::c_void,
3156 );
3157 }
3158 }
3159 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
3160 let hybrid = gpu
3161 .ctx
3162 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
3163 let kda = gpu.ctx.load_module(Ptx::from_binary(KDA_FATBIN.to_vec()))?;
3164 let qmatvec = gpu
3165 .ctx
3166 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
3167 let flash = gpu
3168 .ctx
3169 .load_module(Ptx::from_binary(FLASH_FATBIN.to_vec()))?;
3170 let gemm = gpu
3171 .ctx
3172 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
3173 let router = gpu
3174 .ctx
3175 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
3176 let sample = gpu
3177 .ctx
3178 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
3179 let copy_stream = gpu.ctx.new_stream()?;
3180 if std::env::var("MEMRA_EVT")
3196 .map(|v| v == "1")
3197 .unwrap_or(false)
3198 {
3199 } else {
3201 unsafe {
3202 gpu.ctx.disable_event_tracking();
3203 }
3204 }
3205 Ok(Self {
3206 gpu,
3207 module,
3208 hybrid,
3209 kda,
3210 qmatvec,
3211 flash,
3212 flash_g: std::sync::OnceLock::new(),
3213 gemm,
3214 router,
3215 sample,
3216 moe_cache: Mutex::new(None),
3217 w8_mirrors: Mutex::new(std::collections::HashMap::new()),
3218 w8_act: Mutex::new(std::collections::HashMap::new()),
3219 moe_cache_layout: Mutex::new(None),
3220 copy_stream,
3221 side: std::sync::OnceLock::new(),
3222 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
3223 verify_exact: std::sync::atomic::AtomicBool::new(false),
3224 capture_keep: Mutex::new(Vec::new()),
3225 argmax_partials: Mutex::new(None),
3226 prime_deqw_ws: Mutex::new(None),
3227 router_stage: Mutex::new(None),
3228 hyper_decode_ws: Mutex::new(None),
3229 mla_seg_ws: Mutex::new(None),
3230 mla_pre_done: std::sync::atomic::AtomicBool::new(false),
3231 verify_ws: Mutex::new(VerifyWs::default()),
3232 vrows_macro_dev: Mutex::new(std::collections::HashMap::new()),
3233 shexp_ones: Mutex::new(None),
3234 fp8_scratch: Mutex::new(None),
3235 fa_vf16_scratch: Mutex::new(None),
3236 fa_part_pool: Mutex::new(None),
3237 fa_part_retired: Mutex::new(Vec::new()),
3238 fn_cache: Mutex::new(Default::default()),
3239 f16_scratch: Mutex::new(None),
3240 #[cfg(memra_cutlass)]
3241 cutlass_scratch: Mutex::new(None),
3242 })
3243 }
3244
3245 pub fn ctx(&self) -> &Arc<CudaContext> {
3246 &self.gpu.ctx
3247 }
3248
3249 pub fn pool_cached_bytes(&self) -> usize {
3267 let (reserved, used) = self.pool_reserved_used();
3268 reserved.saturating_sub(used)
3269 }
3270
3271 pub fn device_graph_mem_reserved(&self) -> usize {
3281 use cudarc::driver::sys as cus;
3282 let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
3283 return 0;
3284 };
3285 let mut bytes: u64 = 0;
3286 let rc = unsafe {
3287 cus::cuDeviceGetGraphMemAttribute(
3288 dev,
3289 cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
3290 &mut bytes as *mut u64 as *mut std::ffi::c_void,
3291 )
3292 };
3293 if rc == cus::cudaError_enum::CUDA_SUCCESS {
3294 bytes as usize
3295 } else {
3296 0
3297 }
3298 }
3299
3300 pub fn pool_trim_to_zero(&self) -> usize {
3314 self.pool_trim_to(0)
3315 }
3316 pub fn pool_trim_to(&self, keep: usize) -> usize {
3325 use cudarc::driver::sys;
3326 let (before, _) = self.pool_reserved_used();
3327 unsafe {
3328 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3329 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3330 != sys::CUresult::CUDA_SUCCESS
3331 {
3332 return 0;
3333 }
3334 let _ = sys::cuMemPoolTrimTo(pool, keep);
3335 }
3336 let (after, _) = self.pool_reserved_used();
3337 before.saturating_sub(after)
3338 }
3339
3340 pub fn pool_reserved_used(&self) -> (usize, usize) {
3341 use cudarc::driver::sys;
3342 unsafe {
3343 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3344 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3345 != sys::CUresult::CUDA_SUCCESS
3346 {
3347 return (0, 0);
3348 }
3349 let (mut reserved, mut used) = (0u64, 0u64);
3350 if sys::cuMemPoolGetAttribute(
3351 pool,
3352 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
3353 &mut reserved as *mut u64 as *mut core::ffi::c_void,
3354 ) != sys::CUresult::CUDA_SUCCESS
3355 {
3356 return (0, 0);
3357 }
3358 if sys::cuMemPoolGetAttribute(
3359 pool,
3360 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
3361 &mut used as *mut u64 as *mut core::ffi::c_void,
3362 ) != sys::CUresult::CUDA_SUCCESS
3363 {
3364 return (0, 0);
3365 }
3366 (reserved as usize, used as usize)
3367 }
3368 }
3369
3370 pub fn pool_high_water_reset(&self) -> (usize, usize) {
3379 use cudarc::driver::sys;
3380 unsafe {
3381 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3382 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3383 != sys::CUresult::CUDA_SUCCESS
3384 {
3385 return (0, 0);
3386 }
3387 let (mut reserved, mut used) = (0u64, 0u64);
3388 if sys::cuMemPoolGetAttribute(
3389 pool,
3390 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
3391 &mut reserved as *mut u64 as *mut core::ffi::c_void,
3392 ) != sys::CUresult::CUDA_SUCCESS
3393 {
3394 return (0, 0);
3395 }
3396 if sys::cuMemPoolGetAttribute(
3397 pool,
3398 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
3399 &mut used as *mut u64 as *mut core::ffi::c_void,
3400 ) != sys::CUresult::CUDA_SUCCESS
3401 {
3402 return (0, 0);
3403 }
3404 let mut zero: u64 = 0;
3407 let _ = sys::cuMemPoolSetAttribute(
3408 pool,
3409 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
3410 &mut zero as *mut u64 as *mut core::ffi::c_void,
3411 );
3412 let mut zero2: u64 = 0;
3413 let _ = sys::cuMemPoolSetAttribute(
3414 pool,
3415 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
3416 &mut zero2 as *mut u64 as *mut core::ffi::c_void,
3417 );
3418 (reserved as usize, used as usize)
3419 }
3420 }
3421
3422 pub fn stream(&self) -> Arc<CudaStream> {
3425 self.gpu.stream()
3426 }
3427 pub fn gkv_on() -> bool {
3430 memra_kv::gkv_on()
3431 }
3432
3433 pub fn wkv_on() -> bool {
3445 memra_kv::wkv_on()
3446 }
3447
3448 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
3451 if head_dim == 512 && Self::gkv_on() {
3452 self.func_g(name)
3453 } else {
3454 self.func(name)
3455 }
3456 }
3457
3458 fn func_g(&self, name: &str) -> CudaFunction {
3462 let m = self.flash_g.get_or_init(|| {
3463 self.gpu
3464 .ctx
3465 .load_module(cudarc::nvrtc::Ptx::from_binary(
3466 FLASH_FATBIN_KF8VF8.to_vec(),
3467 ))
3468 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
3469 });
3470 let key = format!("g:{name}");
3471 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
3472 return f.clone();
3473 }
3474 let f = match m.load_function(name) {
3475 Ok(f) => f,
3476 Err(_) => self.func(name),
3477 };
3478 self.fn_cache.lock().unwrap().insert(key, f.clone());
3479 f
3480 }
3481
3482 fn func(&self, name: &str) -> CudaFunction {
3483 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
3486 return f.clone();
3487 }
3488 let f = self
3489 .module
3490 .load_function(name)
3491 .or_else(|_| self.hybrid.load_function(name))
3492 .or_else(|_| self.kda.load_function(name))
3493 .or_else(|_| self.qmatvec.load_function(name))
3494 .or_else(|_| self.flash.load_function(name))
3495 .or_else(|_| self.gemm.load_function(name))
3496 .or_else(|_| self.router.load_function(name))
3497 .or_else(|_| self.sample.load_function(name))
3498 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
3499 self.fn_cache
3500 .lock()
3501 .unwrap()
3502 .insert(name.to_string(), f.clone());
3503 f
3504 }
3505
3506 pub fn scatter_trim_logits(
3509 &self,
3510 src: &CudaSlice<f32>,
3511 d2t: &CudaSlice<u32>,
3512 dst: &mut CudaSlice<f32>,
3513 d_vocab: usize,
3514 n_vocab: usize,
3515 ) -> Result<(), Box<dyn std::error::Error>> {
3516 let f1 = self.func("scatter_trim_logits_f32");
3517 let f2 = self.func("scatter_trim_logits_pass2_f32");
3518 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
3519 let cfg1 = LaunchConfig {
3520 grid_dim: (256, 1, 1),
3521 block_dim: (256, 1, 1),
3522 shared_mem_bytes: 0,
3523 };
3524 let __s_b1 = self.gpu.stream();
3525 let mut b1 = __s_b1.launch_builder(&f1);
3526 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
3527 unsafe {
3528 b1.launch(cfg1)?;
3529 }
3530 let cfg2 = LaunchConfig {
3531 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
3532 block_dim: (256, 1, 1),
3533 shared_mem_bytes: 0,
3534 };
3535 let __s_b2 = self.gpu.stream();
3536 let mut b2 = __s_b2.launch_builder(&f2);
3537 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
3538 unsafe {
3539 b2.launch(cfg2)?;
3540 }
3541 Ok(())
3542 }
3543
3544 #[allow(clippy::too_many_arguments)]
3550 pub fn filter_stats(
3551 &self,
3552 x: &CudaSlice<f32>,
3553 row_stride: usize,
3554 rows: &CudaSlice<i32>,
3555 out_th: &mut CudaSlice<f32>,
3556 out_z: &mut CudaSlice<f32>,
3557 out_max: &mut CudaSlice<f32>,
3558 n: usize,
3559 nrow: usize,
3560 temp: f32,
3561 top_k: i32,
3562 top_p: f32,
3563 min_p: f32,
3564 ) -> Result<(), Box<dyn std::error::Error>> {
3565 if self.sm_count() >= 16 {
3590 let cap = self.sm_count() as usize / 16;
3591 let mut done = 0usize;
3592 while done < nrow {
3593 let chunk = cap.min(nrow - done);
3594 self.filter_stats_coop_chunk(
3595 x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
3596 top_p, min_p,
3597 )?;
3598 done += chunk;
3599 }
3600 return Ok(());
3601 }
3602 self.filter_stats_plain_program(
3603 x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
3604 )
3605 }
3606
3607 #[allow(clippy::too_many_arguments)]
3612 pub fn filter_stats_coop_chunk(
3613 &self,
3614 x: &CudaSlice<f32>,
3615 row_stride: usize,
3616 rows: &CudaSlice<i32>,
3617 row0: usize,
3618 out_th: &mut CudaSlice<f32>,
3619 out_z: &mut CudaSlice<f32>,
3620 out_max: &mut CudaSlice<f32>,
3621 n: usize,
3622 chunk: usize,
3623 temp: f32,
3624 top_k: i32,
3625 top_p: f32,
3626 min_p: f32,
3627 ) -> Result<(), Box<dyn std::error::Error>> {
3628 let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
3629 let f = self.func("filter_stats_coop_f32");
3630 let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
3631 let cfg = LaunchConfig {
3632 grid_dim: (16, chunk as u32, 1),
3633 block_dim: (512, 1, 1),
3634 shared_mem_bytes: 0,
3635 };
3636 let rows_v = rows.slice(row0..row0 + chunk);
3637 let mut th_v = out_th.slice_mut(row0..row0 + chunk);
3638 let mut z_v = out_z.slice_mut(row0..row0 + chunk);
3639 let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
3640 let __s_b = self.gpu.stream();
3641 let mut b = __s_b.launch_builder(&f);
3642 b.arg(x)
3643 .arg(&rs)
3644 .arg(&rows_v)
3645 .arg(&mut th_v)
3646 .arg(&mut z_v)
3647 .arg(&mut mx_v)
3648 .arg(&mut ws)
3649 .arg(&ni)
3650 .arg(&nr)
3651 .arg(&temp)
3652 .arg(&top_k)
3653 .arg(&top_p)
3654 .arg(&min_p);
3655 unsafe {
3656 b.launch_cooperative(cfg)?;
3657 }
3658 Ok(())
3659 }
3660
3661 #[allow(clippy::too_many_arguments)]
3665 pub fn filter_stats_plain_program(
3666 &self,
3667 x: &CudaSlice<f32>,
3668 row_stride: usize,
3669 rows: &CudaSlice<i32>,
3670 out_th: &mut CudaSlice<f32>,
3671 out_z: &mut CudaSlice<f32>,
3672 out_max: &mut CudaSlice<f32>,
3673 n: usize,
3674 nrow: usize,
3675 temp: f32,
3676 top_k: i32,
3677 top_p: f32,
3678 min_p: f32,
3679 ) -> Result<(), Box<dyn std::error::Error>> {
3680 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
3681 let f = self.func("filter_stats_f32");
3682 let cfg = LaunchConfig {
3683 grid_dim: (nrow as u32, 1, 1),
3684 block_dim: (1024, 1, 1),
3685 shared_mem_bytes: 0,
3686 };
3687 let __s_b = self.gpu.stream();
3688 let mut b = __s_b.launch_builder(&f);
3689 b.arg(x)
3690 .arg(&rs)
3691 .arg(rows)
3692 .arg(&mut *out_th)
3693 .arg(&mut *out_z)
3694 .arg(&mut *out_max)
3695 .arg(&ni)
3696 .arg(&nr)
3697 .arg(&temp)
3698 .arg(&top_k)
3699 .arg(&top_p)
3700 .arg(&min_p);
3701 unsafe {
3702 b.launch(cfg)?;
3703 }
3704 Ok(())
3705 }
3706
3707 #[allow(clippy::too_many_arguments)]
3709 pub fn softmax_gather_filtered(
3710 &self,
3711 x: &CudaSlice<f32>,
3712 row_stride: usize,
3713 ids: &CudaSlice<u32>,
3714 rows: &CudaSlice<i32>,
3715 th: &CudaSlice<f32>,
3716 z: &CudaSlice<f32>,
3717 out: &mut CudaSlice<f32>,
3718 n: usize,
3719 npair: usize,
3720 temp: f32,
3721 ) -> Result<(), Box<dyn std::error::Error>> {
3722 let f = self.func("softmax_gather_filtered_f32");
3723 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
3724 let cfg = LaunchConfig {
3725 grid_dim: (npair as u32, 1, 1),
3726 block_dim: (256, 1, 1),
3727 shared_mem_bytes: 0,
3728 };
3729 let __s_b = self.gpu.stream();
3730 let mut b = __s_b.launch_builder(&f);
3731 b.arg(x)
3732 .arg(&rs)
3733 .arg(ids)
3734 .arg(rows)
3735 .arg(th)
3736 .arg(z)
3737 .arg(&mut *out)
3738 .arg(&ni)
3739 .arg(&np)
3740 .arg(&temp);
3741 unsafe {
3742 b.launch(cfg)?;
3743 }
3744 Ok(())
3745 }
3746
3747 #[allow(clippy::too_many_arguments)]
3749 pub fn residual_sample_filtered(
3750 &self,
3751 p: &CudaSlice<f32>,
3752 q: Option<&CudaSlice<f32>>,
3753 n: usize,
3754 temp: f32,
3755 seed: u64,
3756 stream_pos: u32,
3757 p_stats: (f32, f32, f32),
3758 q_stats: (f32, f32, f32),
3759 out_tok: &mut CudaSlice<u32>,
3760 ) -> Result<(), Box<dyn std::error::Error>> {
3761 let f = self.func("residual_sample_filtered_f32");
3762 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3763 let has_q: i32 = q.is_some() as i32;
3764 let qbuf = q.unwrap_or(p);
3765 let (pm, pth, pz) = p_stats;
3766 let (qm, qth, qz) = q_stats;
3767 let cfg = LaunchConfig {
3768 grid_dim: (1, 1, 1),
3769 block_dim: (1024, 1, 1),
3770 shared_mem_bytes: 0,
3771 };
3772 let __s_b = self.gpu.stream();
3773 let mut b = __s_b.launch_builder(&f);
3774 b.arg(p)
3775 .arg(qbuf)
3776 .arg(&has_q)
3777 .arg(&ni)
3778 .arg(&temp)
3779 .arg(&slo)
3780 .arg(&shi)
3781 .arg(&stream_pos)
3782 .arg(&pm)
3783 .arg(&pth)
3784 .arg(&pz)
3785 .arg(&qm)
3786 .arg(&qth)
3787 .arg(&qz)
3788 .arg(&mut *out_tok);
3789 unsafe {
3790 b.launch(cfg)?;
3791 }
3792 Ok(())
3793 }
3794
3795 #[allow(clippy::too_many_arguments)]
3801 pub fn residual_sample_sparse_q(
3802 &self,
3803 p: &CudaSlice<f32>,
3804 cand_ids: &CudaSlice<u32>,
3805 q_probs: &CudaSlice<f32>,
3806 n_cand: usize,
3807 n: usize,
3808 temp: f32,
3809 seed: u64,
3810 stream_pos: u32,
3811 p_stats: (f32, f32, f32),
3812 out_tok: &mut CudaSlice<u32>,
3813 ) -> Result<(), Box<dyn std::error::Error>> {
3814 assert!(
3815 (1..=32).contains(&n_cand),
3816 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
3817 );
3818 let f = self.func("residual_sample_sparse_q_f32");
3819 let (ni, nc) = (n as i32, n_cand as i32);
3820 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3821 let (pm, pth, pz) = p_stats;
3822 let cfg = LaunchConfig {
3823 grid_dim: (1, 1, 1),
3824 block_dim: (1024, 1, 1),
3825 shared_mem_bytes: 0,
3826 };
3827 let __s_b = self.gpu.stream();
3828 let mut b = __s_b.launch_builder(&f);
3829 b.arg(p)
3830 .arg(cand_ids)
3831 .arg(q_probs)
3832 .arg(&nc)
3833 .arg(&ni)
3834 .arg(&temp)
3835 .arg(&slo)
3836 .arg(&shi)
3837 .arg(&stream_pos)
3838 .arg(&pm)
3839 .arg(&pth)
3840 .arg(&pz)
3841 .arg(&mut *out_tok);
3842 unsafe {
3843 b.launch(cfg)?;
3844 }
3845 Ok(())
3846 }
3847
3848 #[allow(clippy::too_many_arguments)]
3850 pub fn gumbel_perturb_filtered(
3851 &self,
3852 x: &CudaSlice<f32>,
3853 y: &mut CudaSlice<f32>,
3854 n: usize,
3855 seed: u64,
3856 stream_pos: u32,
3857 temp: f32,
3858 row_max: f32,
3859 th: f32,
3860 ) -> Result<(), Box<dyn std::error::Error>> {
3861 let f = self.func("gumbel_perturb_filtered_f32");
3862 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3863 let cfg = LaunchConfig {
3864 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3865 block_dim: (256, 1, 1),
3866 shared_mem_bytes: 0,
3867 };
3868 let __s_b = self.gpu.stream();
3869 let mut b = __s_b.launch_builder(&f);
3870 b.arg(x)
3871 .arg(&mut *y)
3872 .arg(&ni)
3873 .arg(&slo)
3874 .arg(&shi)
3875 .arg(&stream_pos)
3876 .arg(&temp)
3877 .arg(&row_max)
3878 .arg(&th);
3879 unsafe {
3880 b.launch(cfg)?;
3881 }
3882 Ok(())
3883 }
3884
3885 #[allow(clippy::too_many_arguments)]
3889 pub fn penalize_logits(
3890 &self,
3891 x: &mut CudaSlice<f32>,
3892 hist: &CudaSlice<u32>,
3893 n_hist: usize,
3894 rep: f32,
3895 freq: f32,
3896 present: f32,
3897 n: usize,
3898 ) -> Result<(), Box<dyn std::error::Error>> {
3899 if n_hist == 0 {
3900 return Ok(());
3901 }
3902 let f = self.func("penalize_logits_f32");
3903 let (nh, ni) = (n_hist as i32, n as i32);
3904 let cfg = LaunchConfig {
3905 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
3906 block_dim: (128, 1, 1),
3907 shared_mem_bytes: 0,
3908 };
3909 let __s_b = self.gpu.stream();
3910 let mut b = __s_b.launch_builder(&f);
3911 b.arg(&mut *x)
3912 .arg(hist)
3913 .arg(&nh)
3914 .arg(&rep)
3915 .arg(&freq)
3916 .arg(&present)
3917 .arg(&ni);
3918 unsafe {
3919 b.launch(cfg)?;
3920 }
3921 Ok(())
3922 }
3923
3924 #[allow(clippy::too_many_arguments)]
3926 pub fn penalize_logits_rows(
3927 &self,
3928 x: &mut CudaSlice<f32>,
3929 hist: &CudaSlice<u32>,
3930 n_hist: usize,
3931 rep: f32,
3932 freq: f32,
3933 present: f32,
3934 n: usize,
3935 nrow: usize,
3936 ) -> Result<(), Box<dyn std::error::Error>> {
3937 if n_hist == 0 || nrow == 0 {
3938 return Ok(());
3939 }
3940 let f = self.func("penalize_logits_rows_f32");
3941 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
3942 let cfg = LaunchConfig {
3943 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
3944 block_dim: (128, 1, 1),
3945 shared_mem_bytes: 0,
3946 };
3947 let __s_b = self.gpu.stream();
3948 let mut b = __s_b.launch_builder(&f);
3949 b.arg(&mut *x)
3950 .arg(hist)
3951 .arg(&nh)
3952 .arg(&rep)
3953 .arg(&freq)
3954 .arg(&present)
3955 .arg(&ni)
3956 .arg(&nr);
3957 unsafe {
3958 b.launch(cfg)?;
3959 }
3960 Ok(())
3961 }
3962
3963 #[allow(clippy::too_many_arguments)]
3969 pub fn penalize_logits_sparse_rows(
3970 &self,
3971 x: &mut CudaSlice<f32>,
3972 ids: &[u32],
3973 counts: &[u32],
3974 offsets: &[i32],
3975 rows: &[i32],
3976 reps: &[f32],
3977 freqs: &[f32],
3978 presents: &[f32],
3979 n: usize,
3980 ) -> Result<(), Box<dyn std::error::Error>> {
3981 let nrow = rows.len();
3982 if nrow == 0 {
3983 return Ok(());
3984 }
3985 let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
3986 let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
3987 let entry_count =
3988 i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
3989 if ids.len() != counts.len()
3990 || offsets.len() != nrow + 1
3991 || reps.len() != nrow
3992 || freqs.len() != nrow
3993 || presents.len() != nrow
3994 || offsets.first().copied() != Some(0)
3995 || offsets.last().copied() != Some(entry_count)
3996 {
3997 return Err("sparse penalty row metadata shape mismatch".into());
3998 }
3999 if counts.contains(&0) {
4000 return Err("sparse penalty counts must be positive".into());
4001 }
4002 let mut max_len = 0usize;
4003 for pair in offsets.windows(2) {
4004 if pair[0] < 0 || pair[1] < pair[0] {
4005 return Err("sparse penalty offsets must be monotonic".into());
4006 }
4007 max_len = max_len.max((pair[1] - pair[0]) as usize);
4008 }
4009 if max_len == 0 {
4010 return Ok(());
4011 }
4012
4013 let mut seen = std::collections::HashSet::with_capacity(ids.len());
4014 for (r, &row) in rows.iter().enumerate() {
4015 if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
4016 return Err("sparse penalty row index exceeds logits shape".into());
4017 }
4018 let begin = offsets[r] as usize;
4019 let end = offsets[r + 1] as usize;
4020 for &id in &ids[begin..end] {
4021 if id as usize >= n {
4022 return Err("sparse penalty token id exceeds logits row".into());
4023 }
4024 if !seen.insert((row, id)) {
4025 return Err("sparse penalty entries must be unique per logits row".into());
4026 }
4027 }
4028 }
4029
4030 unsafe {
4032 self.penalize_logits_sparse_rows_unchecked(
4033 x, ids, counts, offsets, rows, reps, freqs, presents, n,
4034 )
4035 }
4036 }
4037
4038 #[allow(clippy::too_many_arguments)]
4047 pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
4048 &self,
4049 x: &mut CudaSlice<f32>,
4050 ids: &[u32],
4051 counts: &[u32],
4052 offsets: &[i32],
4053 rows: &[i32],
4054 reps: &[f32],
4055 freqs: &[f32],
4056 presents: &[f32],
4057 n: usize,
4058 ) -> Result<(), Box<dyn std::error::Error>> {
4059 let nrow = rows.len();
4060 if nrow == 0 {
4061 return Ok(());
4062 }
4063 let max_len = offsets
4064 .windows(2)
4065 .map(|pair| (pair[1] - pair[0]) as usize)
4066 .max()
4067 .unwrap_or(0);
4068 if max_len == 0 {
4069 return Ok(());
4070 }
4071 let ids_d = self.htod_u32_v(ids)?;
4072 let counts_d = self.htod_u32_v(counts)?;
4073 let offsets_d = self.htod_i32(offsets)?;
4074 let rows_d = self.htod_i32(rows)?;
4075 let reps_d = self.htod(reps)?;
4076 let freqs_d = self.htod(freqs)?;
4077 let presents_d = self.htod(presents)?;
4078 let f = self.func("penalize_logits_sparse_rows_f32");
4079 let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
4080 let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
4081 let cfg = LaunchConfig {
4082 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
4083 block_dim: (128, 1, 1),
4084 shared_mem_bytes: 0,
4085 };
4086 let __s_b = self.gpu.stream();
4087 let mut b = __s_b.launch_builder(&f);
4088 b.arg(&mut *x)
4089 .arg(&ids_d)
4090 .arg(&counts_d)
4091 .arg(&offsets_d)
4092 .arg(&rows_d)
4093 .arg(&reps_d)
4094 .arg(&freqs_d)
4095 .arg(&presents_d)
4096 .arg(&ni)
4097 .arg(&nr);
4098 unsafe {
4099 b.launch(cfg)?;
4100 }
4101 Ok(())
4102 }
4103
4104 #[allow(clippy::too_many_arguments)]
4112 pub fn penalize_logits_rows_inc(
4113 &self,
4114 x: &mut CudaSlice<f32>,
4115 hist: &CudaSlice<u32>,
4116 n_hist0: usize,
4117 rep: f32,
4118 freq: f32,
4119 present: f32,
4120 n: usize,
4121 nrow: usize,
4122 win: usize,
4123 ) -> Result<(), Box<dyn std::error::Error>> {
4124 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
4125 return Ok(());
4126 }
4127 debug_assert!(
4128 hist.len() >= n_hist0 + nrow - 1,
4129 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
4130 );
4131 let f = self.func("penalize_logits_rows_inc_f32");
4132 let max_len = win.min(n_hist0 + nrow - 1).max(1);
4133 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
4134 let cfg = LaunchConfig {
4135 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
4136 block_dim: (128, 1, 1),
4137 shared_mem_bytes: 0,
4138 };
4139 let __s_b = self.gpu.stream();
4140 let mut b = __s_b.launch_builder(&f);
4141 b.arg(&mut *x)
4142 .arg(hist)
4143 .arg(&nh)
4144 .arg(&rep)
4145 .arg(&freq)
4146 .arg(&present)
4147 .arg(&ni)
4148 .arg(&nr)
4149 .arg(&wi);
4150 unsafe {
4151 b.launch(cfg)?;
4152 }
4153 Ok(())
4154 }
4155
4156 pub fn wpf_level() -> u32 {
4164 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
4165 *ON.get_or_init(|| {
4166 std::env::var("MEMRA_WPF")
4167 .ok()
4168 .and_then(|v| v.parse().ok())
4169 .unwrap_or(1)
4170 })
4171 }
4172
4173 pub fn set_verify_exact(&self, on: bool) {
4189 self.verify_exact
4190 .store(on, std::sync::atomic::Ordering::Relaxed);
4191 }
4192 pub(crate) fn verify_exact_on(&self) -> bool {
4193 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
4194 }
4195
4196 pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
4202 ExactScope::set(&self.verify_exact, on)
4203 }
4204
4205 pub fn qkv_append_on() -> bool {
4208 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4209 *ON.get_or_init(|| {
4210 std::env::var("MEMRA_QKV_APPEND")
4211 .map(|v| v != "0")
4212 .unwrap_or(true)
4213 })
4214 }
4215
4216 pub fn pdl_wb_on() -> bool {
4219 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4220 *ON.get_or_init(|| {
4221 std::env::var("MEMRA_PDL_WB")
4222 .map(|v| v != "0")
4223 .unwrap_or(true)
4224 })
4225 }
4226
4227 pub fn norm_ilp_on() -> bool {
4235 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4236 *ON.get_or_init(|| {
4237 std::env::var("MEMRA_NORM_ILP")
4238 .map(|v| v != "0")
4239 .unwrap_or(true)
4240 })
4241 }
4242
4243 pub fn norm_ilp_zq8_on() -> bool {
4248 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4249 *ON.get_or_init(|| {
4250 Self::norm_ilp_on()
4251 && std::env::var("MEMRA_NORM_ILP_ZQ8")
4252 .map(|v| v != "0")
4253 .unwrap_or(true)
4254 })
4255 }
4256
4257 pub fn tk_ffn_dual_on() -> bool {
4265 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4266 *ON.get_or_init(|| {
4267 std::env::var("MEMRA_TK_FFN_DUAL")
4268 .map(|v| v != "0")
4269 .unwrap_or(true)
4270 })
4271 }
4272
4273 pub fn pdl_mmvq_on() -> bool {
4277 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4278 *ON.get_or_init(|| {
4279 std::env::var("MEMRA_PDL_MMVQ")
4280 .map(|v| v != "0")
4281 .unwrap_or(true)
4282 })
4283 }
4284
4285 pub fn pdl_on() -> bool {
4286 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4287 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
4288 }
4289
4290 pub fn pdl_nvfp4q8_on() -> bool {
4296 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4297 *ON.get_or_init(|| {
4298 std::env::var("MEMRA_PDL_NVFP4")
4299 .map(|v| v != "0")
4300 .unwrap_or(true)
4301 })
4302 }
4303
4304 fn q40_mr1_on() -> bool {
4310 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
4311 match *Q40MR.get_or_init(|| {
4312 std::env::var("MEMRA_Q40_MR")
4313 .ok()
4314 .and_then(|v| v.parse().ok())
4315 }) {
4316 Some(v) => v == 1,
4317 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4318 }
4319 }
4320
4321 fn pdl_func_flash(
4326 &self,
4327 g: bool,
4328 name: &'static str,
4329 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4330 use cudarc::driver::sys as cu;
4331 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
4338 std::sync::Mutex::new(None);
4339 #[allow(clippy::type_complexity)] static FNS: std::sync::Mutex<
4341 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
4342 > = std::sync::Mutex::new(None);
4343 let ctx_key = self.ctx().cu_ctx() as usize;
4344 if let Some(&f) = FNS
4345 .lock()
4346 .unwrap()
4347 .get_or_insert_with(Default::default)
4348 .get(&(ctx_key, g, name))
4349 {
4350 return Ok(f as cu::CUfunction);
4351 }
4352 let module = {
4353 let mut mods = MODS.lock().unwrap();
4354 let map = mods.get_or_insert_with(Default::default);
4355 match map.get(&(ctx_key, g)) {
4356 Some(&m) => m,
4357 None => {
4358 let m = self.pdl_load_module_in_ctx(if g {
4359 FLASH_FATBIN_KF8VF8
4360 } else {
4361 FLASH_FATBIN
4362 })?;
4363 map.insert((ctx_key, g), m);
4364 m
4365 }
4366 }
4367 };
4368 let cname = std::ffi::CString::new(name)?;
4369 let mut f: cu::CUfunction = std::ptr::null_mut();
4370 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
4371 if r != cu::CUresult::CUDA_SUCCESS {
4372 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
4373 }
4374 FNS.lock()
4375 .unwrap()
4376 .get_or_insert_with(Default::default)
4377 .insert((ctx_key, g, name), f as usize);
4378 Ok(f)
4379 }
4380
4381 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
4386 use cudarc::driver::sys as cu;
4387 let mut prev: cu::CUcontext = std::ptr::null_mut();
4388 unsafe {
4389 cu::cuCtxGetCurrent(&mut prev).result()?;
4390 }
4391 self.ctx().bind_to_thread()?;
4392 let mut m: cu::CUmodule = std::ptr::null_mut();
4393 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
4394 let restore = if prev.is_null() {
4395 cu::CUresult::CUDA_SUCCESS
4396 } else {
4397 unsafe { cu::cuCtxSetCurrent(prev) }
4398 };
4399 if r != cu::CUresult::CUDA_SUCCESS {
4400 return Err(format!("pdl module load: {r:?}").into());
4401 }
4402 if restore != cu::CUresult::CUDA_SUCCESS {
4403 return Err(format!("pdl module load: ctx restore {restore:?}").into());
4404 }
4405 Ok(m as usize)
4406 }
4407
4408 pub fn raw_kernel_function(
4411 &self,
4412 name: &'static str,
4413 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4414 self.pdl_func(name)
4415 }
4416
4417 fn pdl_func(
4418 &self,
4419 name: &'static str,
4420 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4421 use cudarc::driver::sys as cu;
4422 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
4425 std::sync::Mutex::new(None);
4426 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
4429 std::sync::Mutex::new(None);
4430 static FNS: std::sync::Mutex<
4431 Option<std::collections::HashMap<(usize, &'static str), usize>>,
4432 > = std::sync::Mutex::new(None);
4433 let ctx_key = self.ctx().cu_ctx() as usize;
4434 if let Some(&f) = FNS
4435 .lock()
4436 .unwrap()
4437 .get_or_insert_with(Default::default)
4438 .get(&(ctx_key, name))
4439 {
4440 return Ok(f as cu::CUfunction);
4441 }
4442 let module = {
4443 let mut mods = MODULES.lock().unwrap();
4444 let map = mods.get_or_insert_with(Default::default);
4445 match map.get(&ctx_key) {
4446 Some(&m) => m,
4447 None => {
4448 let m = self.pdl_load_module_in_ctx(FATBIN)?;
4449 map.insert(ctx_key, m);
4450 m
4451 }
4452 }
4453 };
4454 let cname = std::ffi::CString::new(name)?;
4455 let mut f: cu::CUfunction = std::ptr::null_mut();
4456 let mut r =
4457 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
4458 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
4459 let qmodule = {
4460 let mut mods = QMODULES.lock().unwrap();
4461 let map = mods.get_or_insert_with(Default::default);
4462 match map.get(&ctx_key) {
4463 Some(&m) => m,
4464 None => {
4465 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
4466 map.insert(ctx_key, m);
4467 m
4468 }
4469 }
4470 };
4471 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
4472 }
4473 if r != cu::CUresult::CUDA_SUCCESS {
4474 return Err(format!("pdl_func {name}: {r:?}").into());
4475 }
4476 FNS.lock()
4477 .unwrap()
4478 .get_or_insert_with(Default::default)
4479 .insert((ctx_key, name), f as usize);
4480 Ok(f)
4481 }
4482
4483 unsafe fn launch_pdl_flash(
4495 &self,
4496 g: bool,
4497 name: &'static str,
4498 grid: (u32, u32, u32),
4499 block: (u32, u32, u32),
4500 smem: u32,
4501 params: &mut [*mut std::ffi::c_void],
4502 ) -> Result<(), Box<dyn std::error::Error>> {
4503 use cudarc::driver::sys as cu;
4504 let f = self.pdl_func_flash(g, name)?;
4505 if smem > 0 {
4506 let r =
4508 unsafe {
4509 cu::cuFuncSetAttribute(f,
4510 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
4511 smem as i32)
4512 };
4513 if r != cu::CUresult::CUDA_SUCCESS {
4514 return Err(format!("pdl smem attr {name}: {r:?}").into());
4515 }
4516 }
4517 let mut attr = cu::CUlaunchAttribute {
4518 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
4519 pad: [0; 4],
4520 value: cu::CUlaunchAttributeValue {
4521 programmaticStreamSerializationAllowed: 1,
4522 },
4523 };
4524 let cfg = cu::CUlaunchConfig {
4525 gridDimX: grid.0,
4526 gridDimY: grid.1,
4527 gridDimZ: grid.2,
4528 blockDimX: block.0,
4529 blockDimY: block.1,
4530 blockDimZ: block.2,
4531 sharedMemBytes: smem,
4532 hStream: self.gpu.stream().cu_stream(),
4533 attrs: &mut attr,
4534 numAttrs: 1,
4535 };
4536 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
4537 if r != cu::CUresult::CUDA_SUCCESS {
4538 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
4539 }
4540 Ok(())
4541 }
4542
4543 unsafe fn launch_pdl(
4544 &self,
4545 name: &'static str,
4546 grid: (u32, u32, u32),
4547 block: (u32, u32, u32),
4548 params: &mut [*mut std::ffi::c_void],
4549 ) -> Result<(), Box<dyn std::error::Error>> {
4550 use cudarc::driver::sys as cu;
4551 let f = self.pdl_func(name)?;
4552 let mut attr = cu::CUlaunchAttribute {
4553 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
4554 pad: [0; 4],
4555 value: cu::CUlaunchAttributeValue {
4556 programmaticStreamSerializationAllowed: 1,
4557 },
4558 };
4559 let cfg = cu::CUlaunchConfig {
4560 gridDimX: grid.0,
4561 gridDimY: grid.1,
4562 gridDimZ: grid.2,
4563 blockDimX: block.0,
4564 blockDimY: block.1,
4565 blockDimZ: block.2,
4566 sharedMemBytes: 0,
4567 hStream: self.gpu.stream().cu_stream(),
4568 attrs: &mut attr,
4569 numAttrs: 1,
4570 };
4571 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
4572 if r != cu::CUresult::CUDA_SUCCESS {
4573 return Err(format!("launch_pdl {name}: {r:?}").into());
4574 }
4575 Ok(())
4576 }
4577
4578 pub fn prefetch_weight_l2(
4581 &self,
4582 w: &crate::model::GpuTensor,
4583 ) -> Result<(), Box<dyn std::error::Error>> {
4584 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
4585 let p = rp4.as_ref().unwrap_or(bytes);
4586 self.prefetch_l2(p, p.len())?;
4587 }
4588 Ok(())
4589 }
4590
4591 pub fn gather_row_bf16(
4594 &self,
4595 table: &CudaSlice<u8>,
4596 tok: &CudaSlice<u32>,
4597 idx: usize,
4598 dst: &mut CudaSlice<f32>,
4599 ncols: usize,
4600 ) -> Result<(), Box<dyn std::error::Error>> {
4601 let f = self.func("gather_row_bf16_f32");
4602 let cfg = LaunchConfig {
4603 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
4604 block_dim: (256, 1, 1),
4605 shared_mem_bytes: 0,
4606 };
4607 let (nc, ix) = (ncols as i32, idx as i32);
4608 let __s_b = self.gpu.stream();
4609 let mut b = __s_b.launch_builder(&f);
4610 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
4611 unsafe {
4612 b.launch(cfg)?;
4613 }
4614 Ok(())
4615 }
4616
4617 #[allow(clippy::too_many_arguments)]
4623 pub fn dflash2_dynconv(
4624 &self,
4625 x: &CudaSlice<f32>,
4626 dyn_: &CudaSlice<f32>,
4627 base: &CudaSlice<f32>,
4628 out: &mut CudaSlice<f32>,
4629 rows: usize,
4630 hidden: usize,
4631 group_size: usize,
4632 ksize: usize,
4633 half: usize,
4634 ) -> Result<(), Box<dyn std::error::Error>> {
4635 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
4636 let f = self.func("dflash2_dynconv_f32");
4637 let n = rows * hidden;
4638 let cfg = LaunchConfig {
4639 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4640 block_dim: (256, 1, 1),
4641 shared_mem_bytes: 0,
4642 };
4643 let (ri, hi, gi, ki, hf) = (
4644 rows as i32,
4645 hidden as i32,
4646 group_size as i32,
4647 ksize as i32,
4648 half as i32,
4649 );
4650 let __s_b = self.gpu.stream();
4651 let mut b = __s_b.launch_builder(&f);
4652 b.arg(x)
4653 .arg(dyn_)
4654 .arg(base)
4655 .arg(out)
4656 .arg(&ri)
4657 .arg(&hi)
4658 .arg(&gi)
4659 .arg(&ki)
4660 .arg(&hf);
4661 unsafe {
4662 b.launch(cfg)?;
4663 }
4664 Ok(())
4665 }
4666
4667 pub fn topk_rows(
4671 &self,
4672 logits: &CudaSlice<f32>,
4673 n_rows: usize,
4674 n_cols: usize,
4675 k: usize,
4676 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
4677 assert!((1..=32).contains(&k), "topk_rows supports 1..=32, got {k}");
4678 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
4679 if topk_shards_on() && n_cols >= 16 * 1024 && k <= n_cols / 16 {
4686 if TOPK_SHARDS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
4687 eprintln!(
4688 "[topk-shards] engaged: rows={n_rows} cols={n_cols} k={k} shards=16 \
4689 (MEMRA_TOPK_SHARDS=1)"
4690 );
4691 }
4692 return self.topk_rows_sharded(logits, n_rows, n_cols, k, 16);
4693 }
4694 let f = self.func("topk_rows_f32");
4695 let nth = 256usize;
4696 let mut vals = self.uninit(n_rows * k)?;
4697 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
4698 let cfg = LaunchConfig {
4699 grid_dim: (n_rows as u32, 1, 1),
4700 block_dim: (nth as u32, 1, 1),
4701 shared_mem_bytes: (nth * k * 8) as u32,
4702 };
4703 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
4704 let __s_b = self.gpu.stream();
4705 let mut b = __s_b.launch_builder(&f);
4706 b.arg(logits)
4707 .arg(&nr)
4708 .arg(&nc)
4709 .arg(&ki)
4710 .arg(&mut vals)
4711 .arg(&mut idxs);
4712 unsafe {
4713 b.launch(cfg)?;
4714 }
4715 Ok((vals, idxs))
4716 }
4717
4718 fn topk_rows_sharded(
4724 &self,
4725 logits: &CudaSlice<f32>,
4726 n_rows: usize,
4727 n_cols: usize,
4728 k: usize,
4729 n_shards: usize,
4730 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
4731 assert!((1..=64).contains(&n_shards), "shard merge head cap is 64");
4732 let nth = 256usize;
4733 let mut pvals = self.uninit(n_rows * n_shards * k)?;
4734 let mut pidxs = self.alloc_uninit::<u32>(n_rows * n_shards * k)?;
4735 let f1 = self.func("topk_rows_shard_f32");
4736 let cfg1 = LaunchConfig {
4737 grid_dim: (n_rows as u32, n_shards as u32, 1),
4738 block_dim: (nth as u32, 1, 1),
4739 shared_mem_bytes: (nth * k * 8) as u32,
4740 };
4741 let (nr, nc, ki, ns) = (n_rows as i32, n_cols as i32, k as i32, n_shards as i32);
4742 {
4743 let __s_b = self.gpu.stream();
4744 let mut b = __s_b.launch_builder(&f1);
4745 b.arg(logits)
4746 .arg(&nr)
4747 .arg(&nc)
4748 .arg(&ki)
4749 .arg(&ns)
4750 .arg(&mut pvals)
4751 .arg(&mut pidxs);
4752 unsafe {
4753 b.launch(cfg1)?;
4754 }
4755 }
4756 let mut vals = self.uninit(n_rows * k)?;
4757 let mut idxs = self.alloc_uninit::<u32>(n_rows * k)?;
4758 let f2 = self.func("topk_rows_shard_merge_f32");
4759 let cfg2 = LaunchConfig {
4760 grid_dim: (n_rows as u32, 1, 1),
4761 block_dim: (32, 1, 1),
4762 shared_mem_bytes: 0,
4763 };
4764 let __s_b = self.gpu.stream();
4765 let mut b = __s_b.launch_builder(&f2);
4766 b.arg(&pvals)
4767 .arg(&pidxs)
4768 .arg(&nr)
4769 .arg(&ns)
4770 .arg(&ki)
4771 .arg(&mut vals)
4772 .arg(&mut idxs);
4773 unsafe {
4774 b.launch(cfg2)?;
4775 }
4776 Ok((vals, idxs))
4777 }
4778
4779 pub fn add_row_inplace(
4781 &self,
4782 logits: &mut CudaSlice<f32>,
4783 bias: &CudaSlice<f32>,
4784 n: usize,
4785 row_off: usize,
4786 ) -> Result<(), Box<dyn std::error::Error>> {
4787 let f = self.func("add_row_inplace_f32");
4788 let cfg = LaunchConfig {
4789 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4790 block_dim: (256, 1, 1),
4791 shared_mem_bytes: 0,
4792 };
4793 let (ni, off) = (n as i32, row_off as i64);
4794 let __s_b = self.gpu.stream();
4795 let mut b = __s_b.launch_builder(&f);
4796 b.arg(logits).arg(bias).arg(&ni).arg(&off);
4797 unsafe {
4798 b.launch(cfg)?;
4799 }
4800 Ok(())
4801 }
4802
4803 pub fn prefetch_l2(
4805 &self,
4806 p: &CudaSlice<u8>,
4807 n: usize,
4808 ) -> Result<(), Box<dyn std::error::Error>> {
4809 let f = self.func("prefetch_l2_bytes");
4810 let lines = n.div_ceil(128);
4811 let ni = n as i64;
4812 let cfg = LaunchConfig {
4813 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
4814 block_dim: (256, 1, 1),
4815 shared_mem_bytes: 0,
4816 };
4817 let __s_b = self.gpu.stream();
4818 let mut b = __s_b.launch_builder(&f);
4819 b.arg(p).arg(&ni);
4820 unsafe {
4821 b.launch(cfg)?;
4822 }
4823 Ok(())
4824 }
4825
4826 pub fn router_gemv(
4829 &self,
4830 w: &CudaSlice<f32>,
4831 x: &CudaSlice<f32>,
4832 n_embd: usize,
4833 n_experts: usize,
4834 t: usize,
4835 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4836 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
4842 Ok("0") => false,
4843 Ok(_) => true,
4844 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4845 };
4846 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
4856 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
4857 }
4858
4859 #[allow(clippy::too_many_arguments)] pub fn router_gemv_form(
4863 &self,
4864 w: &CudaSlice<f32>,
4865 x: &CudaSlice<f32>,
4866 n_embd: usize,
4867 n_experts: usize,
4868 t: usize,
4869 w8: bool,
4870 batch: bool,
4871 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4872 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
4873 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
4874 let f = if batch {
4875 self.func("router_gemv_f32_w8_batch")
4876 } else if w8 {
4877 self.func("router_gemv_f32_w8")
4878 } else {
4879 self.func("router_gemv_f32")
4880 };
4881 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
4882 let cfg = if batch {
4883 LaunchConfig {
4884 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
4885 block_dim: (32, 8, 1),
4886 shared_mem_bytes: 0,
4887 }
4888 } else {
4889 LaunchConfig {
4890 grid_dim: (n_experts as u32, t as u32, 1),
4891 block_dim: (32, if w8 { 8 } else { 1 }, 1),
4892 shared_mem_bytes: 0,
4893 }
4894 };
4895 let __s_b = self.gpu.stream();
4896 let mut b = __s_b.launch_builder(&f);
4897 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
4898 unsafe {
4899 b.launch(cfg)?;
4900 }
4901 Ok(y)
4902 }
4903
4904 pub fn router_gemv_into(
4907 &self,
4908 w: &CudaSlice<f32>,
4909 x: &CudaSlice<f32>,
4910 y: &mut CudaSlice<f32>,
4911 n_embd: usize,
4912 n_experts: usize,
4913 t: usize,
4914 ) -> Result<(), Box<dyn std::error::Error>> {
4915 if y.len() < t * n_experts {
4916 return Err("router_gemv_into output too small".into());
4917 }
4918 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
4919 Ok("0") => false,
4920 Ok(_) => true,
4921 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4922 };
4923 let f = if w8 {
4924 self.func("router_gemv_f32_w8")
4925 } else {
4926 self.func("router_gemv_f32")
4927 };
4928 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
4929 let cfg = LaunchConfig {
4930 grid_dim: (n_experts as u32, t as u32, 1),
4931 block_dim: (32, if w8 { 8 } else { 1 }, 1),
4932 shared_mem_bytes: 0,
4933 };
4934 let __s_b = self.gpu.stream();
4935 let mut b = __s_b.launch_builder(&f);
4936 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
4937 unsafe {
4938 b.launch(cfg)?;
4939 }
4940 Ok(())
4941 }
4942
4943 pub fn rows_permute(
4945 &self,
4946 src: &CudaSlice<f32>,
4947 idx: &CudaSlice<i32>,
4948 nrows: usize,
4949 ncols: usize,
4950 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4951 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
4952 let f = self.func("rows_permute_f32");
4953 let (nc, nr) = (ncols as i32, nrows as i32);
4954 let cfg = LaunchConfig {
4955 grid_dim: (nrows as u32, 1, 1),
4956 block_dim: (256, 1, 1),
4957 shared_mem_bytes: 0,
4958 };
4959 let __s_b = self.gpu.stream();
4960 let mut b = __s_b.launch_builder(&f);
4961 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
4962 unsafe {
4963 b.launch(cfg)?;
4964 }
4965 Ok(dst)
4966 }
4967
4968 pub fn sigmoid_dot_rows(
4973 &self,
4974 x: &CudaSlice<f32>,
4975 w: &CudaSlice<f32>,
4976 n_embd: usize,
4977 t: usize,
4978 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4979 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4982 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
4983 let gs = self.linear(x, w, t, n_embd, 1)?;
4984 let mut g = self.uninit(t)?;
4985 self.sigmoid(&gs, &mut g, t)?;
4986 return Ok(g);
4987 }
4988 let mut g = self.alloc_uninit::<f32>(t)?;
4994 let f = self.func("sigmoid_dot_rows_f32");
4995 let (ne, ti) = (n_embd as i32, t as i32);
4996 let cfg = LaunchConfig {
4997 grid_dim: (t as u32, 1, 1),
4998 block_dim: (32, 8, 1),
4999 shared_mem_bytes: 0,
5000 };
5001 let __s_b = self.gpu.stream();
5002 let mut b = __s_b.launch_builder(&f);
5003 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
5004 unsafe {
5005 b.launch(cfg)?;
5006 }
5007 Ok(g)
5008 }
5009
5010 pub fn sigmoid_dot_rows_into(
5012 &self,
5013 x: &CudaSlice<f32>,
5014 w: &CudaSlice<f32>,
5015 g: &mut CudaSlice<f32>,
5016 n_embd: usize,
5017 t: usize,
5018 ) -> Result<(), Box<dyn std::error::Error>> {
5019 if g.len() < t {
5020 return Err("sigmoid_dot_rows_into output too small".into());
5021 }
5022 let f = self.func("sigmoid_dot_rows_f32");
5023 let (ne, ti) = (n_embd as i32, t as i32);
5024 let cfg = LaunchConfig {
5025 grid_dim: (t as u32, 1, 1),
5026 block_dim: (32, 8, 1),
5027 shared_mem_bytes: 0,
5028 };
5029 let __s_b = self.gpu.stream();
5030 let mut b = __s_b.launch_builder(&f);
5031 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
5032 unsafe {
5033 b.launch(cfg)?;
5034 }
5035 Ok(())
5036 }
5037
5038 pub fn spec_rollback_stream(
5040 &self,
5041 len_ptrs: &CudaSlice<u64>,
5042 pos_start: &CudaSlice<i32>,
5043 acc: &CudaSlice<u32>,
5044 base: usize,
5045 n_rows: usize,
5046 ) -> Result<(), Box<dyn std::error::Error>> {
5047 let f = self.func("spec_rollback_stream");
5048 let (b, nr) = (base as i32, n_rows as i32);
5049 let cfg = LaunchConfig {
5050 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
5051 block_dim: (64, 1, 1),
5052 shared_mem_bytes: 0,
5053 };
5054 let __s_bl = self.gpu.stream();
5055 let mut bl = __s_bl.launch_builder(&f);
5056 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
5057 unsafe {
5058 bl.launch(cfg)?;
5059 }
5060 Ok(())
5061 }
5062
5063 pub fn plain_tok_ring(
5065 &self,
5066 vam: &CudaSlice<u32>,
5067 pos_start: &CudaSlice<i32>,
5068 base: usize,
5069 ring: &mut CudaSlice<u32>,
5070 ) -> Result<(), Box<dyn std::error::Error>> {
5071 let f = self.func("plain_tok_ring");
5072 let (b, cap) = (base as i32, ring.len() as i32);
5073 let cfg = LaunchConfig {
5074 grid_dim: (1, 1, 1),
5075 block_dim: (32, 1, 1),
5076 shared_mem_bytes: 0,
5077 };
5078 let __s_bl = self.gpu.stream();
5079 let mut bl = __s_bl.launch_builder(&f);
5080 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
5081 unsafe {
5082 bl.launch(cfg)?;
5083 }
5084 Ok(())
5085 }
5086
5087 pub fn spec_ring_commit(
5089 &self,
5090 vtok: &CudaSlice<u32>,
5091 acc: &CudaSlice<u32>,
5092 brk: &CudaSlice<u32>,
5093 ring: &mut CudaSlice<u32>,
5094 pend: &mut CudaSlice<u32>,
5095 ) -> Result<(), Box<dyn std::error::Error>> {
5096 let f = self.func("spec_ring_commit");
5097 let cfg = LaunchConfig {
5098 grid_dim: (1, 1, 1),
5099 block_dim: (32, 1, 1),
5100 shared_mem_bytes: 0,
5101 };
5102 let __s_b = self.gpu.stream();
5103 let mut b = __s_b.launch_builder(&f);
5104 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
5105 unsafe {
5106 b.launch(cfg)?;
5107 }
5108 Ok(())
5109 }
5110 pub fn i32_copy_add(
5111 &self,
5112 src: &CudaSlice<i32>,
5113 dst: &mut CudaSlice<i32>,
5114 delta: i32,
5115 ) -> Result<(), Box<dyn std::error::Error>> {
5116 let f = self.func("i32_copy_add");
5117 let cfg = LaunchConfig {
5118 grid_dim: (1, 1, 1),
5119 block_dim: (32, 1, 1),
5120 shared_mem_bytes: 0,
5121 };
5122 let __s_b = self.gpu.stream();
5123 let mut b = __s_b.launch_builder(&f);
5124 b.arg(src).arg(dst).arg(&delta);
5125 unsafe {
5126 b.launch(cfg)?;
5127 }
5128 Ok(())
5129 }
5130 pub fn u32_copy(
5131 &self,
5132 src: &CudaSlice<u32>,
5133 dst: &mut CudaSlice<u32>,
5134 ) -> Result<(), Box<dyn std::error::Error>> {
5135 let f = self.func("u32_copy");
5136 let cfg = LaunchConfig {
5137 grid_dim: (1, 1, 1),
5138 block_dim: (32, 1, 1),
5139 shared_mem_bytes: 0,
5140 };
5141 let __s_b = self.gpu.stream();
5142 let mut b = __s_b.launch_builder(&f);
5143 b.arg(src).arg(dst);
5144 unsafe {
5145 b.launch(cfg)?;
5146 }
5147 Ok(())
5148 }
5149
5150 pub fn spec_adapt_k(
5154 &self,
5155 acc: &CudaSlice<u32>,
5156 brk: &mut CudaSlice<u32>,
5157 floor: usize,
5158 cap: usize,
5159 ) -> Result<(), Box<dyn std::error::Error>> {
5160 let f = self.func("spec_adapt_k");
5161 let (fl, cp) = (floor as i32, cap as i32);
5162 let cfg = LaunchConfig {
5163 grid_dim: (1, 1, 1),
5164 block_dim: (32, 1, 1),
5165 shared_mem_bytes: 0,
5166 };
5167 let __s_b = self.gpu.stream();
5168 let mut b = __s_b.launch_builder(&f);
5169 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
5170 unsafe {
5171 b.launch(cfg)?;
5172 }
5173 Ok(())
5174 }
5175
5176 pub fn spec_accept_greedy_dc(
5178 &self,
5179 preds: &CudaSlice<u32>,
5180 vtok: &CudaSlice<u32>,
5181 last_pred: &CudaSlice<u32>,
5182 brk: &CudaSlice<u32>,
5183 out: &mut CudaSlice<u32>,
5184 ) -> Result<(), Box<dyn std::error::Error>> {
5185 let f = self.func("spec_accept_greedy_dc");
5186 let cfg = LaunchConfig {
5187 grid_dim: (1, 1, 1),
5188 block_dim: (32, 1, 1),
5189 shared_mem_bytes: 0,
5190 };
5191 let __s_b = self.gpu.stream();
5192 let mut b = __s_b.launch_builder(&f);
5193 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
5194 unsafe {
5195 b.launch(cfg)?;
5196 }
5197 Ok(())
5198 }
5199
5200 pub fn pos_iota(
5202 &self,
5203 pos0: &CudaSlice<i32>,
5204 out: &mut CudaSlice<i32>,
5205 t: usize,
5206 ) -> Result<(), Box<dyn std::error::Error>> {
5207 let f = self.func("pos_iota_i32");
5208 let ti = t as i32;
5209 let cfg = LaunchConfig {
5210 grid_dim: (1, 1, 1),
5211 block_dim: (t.max(1) as u32, 1, 1),
5212 shared_mem_bytes: 0,
5213 };
5214 let __s_b = self.gpu.stream();
5215 let mut b = __s_b.launch_builder(&f);
5216 b.arg(pos0).arg(out).arg(&ti);
5217 unsafe {
5218 b.launch(cfg)?;
5219 }
5220 Ok(())
5221 }
5222 #[allow(clippy::too_many_arguments)]
5223 pub fn append_kv_quantized_rows_dc(
5224 &self,
5225 k_rows: &CudaSlice<f32>,
5226 v_rows: &CudaSlice<f32>,
5227 kc: &mut CudaSlice<u8>,
5228 vc: &mut CudaSlice<u8>,
5229 t0_dev: &CudaSlice<i32>,
5230 t: usize,
5231 kv_dim_k: usize,
5232 kv_dim_v: usize,
5233 k_tok_bytes: usize,
5234 v_tok_bytes: usize,
5235 g: bool,
5236 ) -> Result<(), Box<dyn std::error::Error>> {
5237 let f = if g {
5238 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
5239 } else {
5240 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
5241 };
5242 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
5243 let cfg = LaunchConfig {
5244 grid_dim: (nblk, t as u32, 1),
5245 block_dim: (32, 1, 1),
5246 shared_mem_bytes: 0,
5247 };
5248 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
5249 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5250 let __s_b = self.gpu.stream();
5251 let mut b = __s_b.launch_builder(&f);
5252 b.arg(k_rows)
5253 .arg(v_rows)
5254 .arg(kc)
5255 .arg(vc)
5256 .arg(t0_dev)
5257 .arg(&kdk)
5258 .arg(&kdv)
5259 .arg(&ktb)
5260 .arg(&vtb);
5261 unsafe {
5262 b.launch(cfg)?;
5263 }
5264 Ok(())
5265 }
5266
5267 #[allow(clippy::too_many_arguments)]
5270 pub fn append_kv_quantized_row_dc_inc(
5271 &self,
5272 k_row: &CudaSlice<f32>,
5273 v_row: &CudaSlice<f32>,
5274 kc: &mut CudaSlice<u8>,
5275 vc: &mut CudaSlice<u8>,
5276 t0_dev: &mut CudaSlice<i32>,
5277 kv_dim_k: usize,
5278 kv_dim_v: usize,
5279 k_tok_bytes: usize,
5280 v_tok_bytes: usize,
5281 g: bool,
5282 ) -> Result<(), Box<dyn std::error::Error>> {
5283 let f = if g {
5284 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
5285 } else {
5286 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
5287 };
5288 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
5289 let cfg = LaunchConfig {
5290 grid_dim: (1, 1, 1),
5291 block_dim: (nthreads, 1, 1),
5292 shared_mem_bytes: 0,
5293 };
5294 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
5295 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5296 let __s_b = self.gpu.stream();
5297 let mut b = __s_b.launch_builder(&f);
5298 b.arg(k_row)
5299 .arg(v_row)
5300 .arg(kc)
5301 .arg(vc)
5302 .arg(t0_dev)
5303 .arg(&kdk)
5304 .arg(&kdv)
5305 .arg(&ktb)
5306 .arg(&vtb);
5307 unsafe {
5308 b.launch(cfg)?;
5309 }
5310 Ok(())
5311 }
5312
5313 pub fn pack_tok_p(
5315 &self,
5316 tok: &CudaSlice<u32>,
5317 p: &CudaSlice<f32>,
5318 out: &mut CudaSlice<u32>,
5319 slot: usize,
5320 ) -> Result<(), Box<dyn std::error::Error>> {
5321 let f = self.func("pack_tok_p");
5322 let sl = slot as i32;
5323 let cfg = LaunchConfig {
5324 grid_dim: (1, 1, 1),
5325 block_dim: (32, 1, 1),
5326 shared_mem_bytes: 0,
5327 };
5328 let __s_b = self.gpu.stream();
5329 let mut b = __s_b.launch_builder(&f);
5330 b.arg(tok).arg(p).arg(out).arg(&sl);
5331 unsafe {
5332 b.launch(cfg)?;
5333 }
5334 Ok(())
5335 }
5336 pub fn tok_map_u32(
5337 &self,
5338 tok: &mut CudaSlice<u32>,
5339 map: &CudaSlice<u32>,
5340 ) -> Result<(), Box<dyn std::error::Error>> {
5341 let f = self.func("tok_map_u32");
5342 let cfg = LaunchConfig {
5343 grid_dim: (1, 1, 1),
5344 block_dim: (32, 1, 1),
5345 shared_mem_bytes: 0,
5346 };
5347 let __s_b = self.gpu.stream();
5348 let mut b = __s_b.launch_builder(&f);
5349 b.arg(tok).arg(map);
5350 unsafe {
5351 b.launch(cfg)?;
5352 }
5353 Ok(())
5354 }
5355
5356 #[allow(clippy::too_many_arguments)]
5358 pub fn spec_assemble_verify(
5359 &self,
5360 tokp: &CudaSlice<u32>,
5361 pend: &CudaSlice<u32>,
5362 d2t: Option<&CudaSlice<u32>>,
5363 vtok: &mut CudaSlice<u32>,
5364 brk: &mut CudaSlice<u32>,
5365 p_min: f32,
5366 k: usize,
5367 pmin0: bool,
5368 ) -> Result<(), Box<dyn std::error::Error>> {
5369 let f = self.func("spec_assemble_verify");
5370 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
5371 let cfg = LaunchConfig {
5372 grid_dim: (1, 1, 1),
5373 block_dim: (32, 1, 1),
5374 shared_mem_bytes: 0,
5375 };
5376 let __s_b = self.gpu.stream();
5377 let mut b = __s_b.launch_builder(&f);
5378 match d2t {
5379 Some(m) => {
5380 b.arg(tokp)
5381 .arg(pend)
5382 .arg(m)
5383 .arg(vtok)
5384 .arg(brk)
5385 .arg(&p_min)
5386 .arg(&ki)
5387 .arg(&pm);
5388 unsafe {
5389 b.launch(cfg)?;
5390 }
5391 }
5392 None => {
5393 let null: u64 = 0;
5394 b.arg(tokp)
5395 .arg(pend)
5396 .arg(&null)
5397 .arg(vtok)
5398 .arg(brk)
5399 .arg(&p_min)
5400 .arg(&ki)
5401 .arg(&pm);
5402 unsafe {
5403 b.launch(cfg)?;
5404 }
5405 }
5406 }
5407 Ok(())
5408 }
5409
5410 #[allow(clippy::too_many_arguments)]
5412 pub fn ssm_conv_ring_rebuild_dc(
5413 &self,
5414 qkv_tm: &CudaSlice<f32>,
5415 ring_old: &CudaSlice<f32>,
5416 conv_state: &mut CudaSlice<f32>,
5417 conv_dim: usize,
5418 acc: &CudaSlice<u32>,
5419 base: usize,
5420 t_v: usize,
5421 d_conv: usize,
5422 ) -> Result<(), Box<dyn std::error::Error>> {
5423 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
5424 let n = conv_dim * (d_conv - 1);
5425 let cfg = LaunchConfig::for_num_elems(n as u32);
5426 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
5427 let __s_b = self.gpu.stream();
5428 let mut b = __s_b.launch_builder(&f);
5429 b.arg(qkv_tm)
5430 .arg(ring_old)
5431 .arg(conv_state)
5432 .arg(&cd)
5433 .arg(acc)
5434 .arg(&b0)
5435 .arg(&tv)
5436 .arg(&dc);
5437 unsafe {
5438 b.launch(cfg)?;
5439 }
5440 Ok(())
5441 }
5442 #[allow(clippy::too_many_arguments)]
5443 pub fn gdn_scan_s128_dc(
5444 &self,
5445 q: &CudaSlice<f32>,
5446 k: &CudaSlice<f32>,
5447 v: &CudaSlice<f32>,
5448 g: &CudaSlice<f32>,
5449 beta: &CudaSlice<f32>,
5450 state_in: &CudaSlice<f32>,
5451 state_out: &mut CudaSlice<f32>,
5452 o: &mut CudaSlice<f32>,
5453 n_head: usize,
5454 acc: &CudaSlice<u32>,
5455 base: usize,
5456 t_v: usize,
5457 scale: f32,
5458 ) -> Result<(), Box<dyn std::error::Error>> {
5459 let f = self.func("gdn_scan_s128_dc");
5460 const S_V: u32 = 128;
5461 const WARP: u32 = 32;
5462 const COLS_PER_BLOCK: u32 = 4;
5463 let cfg = LaunchConfig {
5464 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
5465 block_dim: (WARP, COLS_PER_BLOCK, 1),
5466 shared_mem_bytes: 0,
5467 };
5468 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
5469 let __s_b = self.gpu.stream();
5470 let mut b = __s_b.launch_builder(&f);
5471 b.arg(q)
5472 .arg(k)
5473 .arg(v)
5474 .arg(g)
5475 .arg(beta)
5476 .arg(state_in)
5477 .arg(state_out)
5478 .arg(o)
5479 .arg(&h)
5480 .arg(acc)
5481 .arg(&b0)
5482 .arg(&tv)
5483 .arg(&scale);
5484 unsafe {
5485 b.launch(cfg)?;
5486 }
5487 Ok(())
5488 }
5489
5490 pub fn spec_seed_gather(
5493 &self,
5494 vx: &CudaSlice<f32>,
5495 fill_prev: &CudaSlice<f32>,
5496 acc: &CudaSlice<u32>,
5497 h_seed: &mut CudaSlice<f32>,
5498 base: usize,
5499 n_embd: usize,
5500 ) -> Result<(), Box<dyn std::error::Error>> {
5501 let f = self.func("spec_seed_gather");
5502 let (b, ne) = (base as i32, n_embd as i32);
5503 let cfg = LaunchConfig {
5504 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
5505 block_dim: (256, 1, 1),
5506 shared_mem_bytes: 0,
5507 };
5508 let __s_bl = self.gpu.stream();
5509 let mut bl = __s_bl.launch_builder(&f);
5510 bl.arg(vx)
5511 .arg(fill_prev)
5512 .arg(acc)
5513 .arg(h_seed)
5514 .arg(&b)
5515 .arg(&ne);
5516 unsafe {
5517 bl.launch(cfg)?;
5518 }
5519 Ok(())
5520 }
5521
5522 pub fn gumbel_perturb(
5529 &self,
5530 x: &CudaSlice<f32>,
5531 y: &mut CudaSlice<f32>,
5532 n: usize,
5533 seed: u64,
5534 stream_pos: u32,
5535 temp: f32,
5536 ) -> Result<(), Box<dyn std::error::Error>> {
5537 let f = self.func("gumbel_perturb_f32");
5538 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5539 let cfg = LaunchConfig {
5540 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5541 block_dim: (256, 1, 1),
5542 shared_mem_bytes: 0,
5543 };
5544 let __s_b = self.gpu.stream();
5545 let mut b = __s_b.launch_builder(&f);
5546 b.arg(x)
5547 .arg(&mut *y)
5548 .arg(&ni)
5549 .arg(&slo)
5550 .arg(&shi)
5551 .arg(&stream_pos)
5552 .arg(&temp);
5553 unsafe {
5554 b.launch(cfg)?;
5555 }
5556 Ok(())
5557 }
5558
5559 pub fn mask_logits_col(
5567 &self,
5568 logits: &mut CudaSlice<f32>,
5569 mask: &CudaSlice<u32>,
5570 col: usize,
5571 n: usize,
5572 mask_words: usize,
5573 ) -> Result<(), Box<dyn std::error::Error>> {
5574 let f = self.func("mask_logits_f32");
5575 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
5576 let cfg = LaunchConfig {
5577 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
5578 block_dim: (256, 1, 1),
5579 shared_mem_bytes: 0,
5580 };
5581 let __s_b = self.gpu.stream();
5582 let mut b = __s_b.launch_builder(&f);
5583 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
5584 unsafe {
5585 b.launch(cfg)?;
5586 }
5587 Ok(())
5588 }
5589
5590 #[allow(clippy::too_many_arguments)] pub fn gumbel_perturb_col(
5598 &self,
5599 x: &CudaSlice<f32>,
5600 col: usize,
5601 y: &mut CudaSlice<f32>,
5602 n: usize,
5603 seed: u64,
5604 stream_pos: u32,
5605 temp: f32,
5606 ) -> Result<(), Box<dyn std::error::Error>> {
5607 let f = self.func("gumbel_perturb_f32");
5608 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5609 let col_view = x.slice(col * n..(col + 1) * n);
5610 let cfg = LaunchConfig {
5611 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5612 block_dim: (256, 1, 1),
5613 shared_mem_bytes: 0,
5614 };
5615 let __s_b = self.gpu.stream();
5616 let mut b = __s_b.launch_builder(&f);
5617 b.arg(&col_view)
5618 .arg(&mut *y)
5619 .arg(&ni)
5620 .arg(&slo)
5621 .arg(&shi)
5622 .arg(&stream_pos)
5623 .arg(&temp);
5624 unsafe {
5625 b.launch(cfg)?;
5626 }
5627 Ok(())
5628 }
5629
5630 #[allow(clippy::too_many_arguments)]
5636 pub fn gumbel_perturb_filtered_col(
5637 &self,
5638 x: &CudaSlice<f32>,
5639 col: usize,
5640 y: &mut CudaSlice<f32>,
5641 n: usize,
5642 seed: u64,
5643 stream_pos: u32,
5644 temp: f32,
5645 stat_max: &CudaSlice<f32>,
5646 stat_th: &CudaSlice<f32>,
5647 stat_idx: usize,
5648 ) -> Result<(), Box<dyn std::error::Error>> {
5649 let f = self.func("gumbel_perturb_filtered_col_f32");
5650 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5651 let (ci, si) = (col as i32, stat_idx as i32);
5652 let cfg = LaunchConfig {
5653 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5654 block_dim: (256, 1, 1),
5655 shared_mem_bytes: 0,
5656 };
5657 let __s_b = self.gpu.stream();
5658 let mut b = __s_b.launch_builder(&f);
5659 b.arg(x)
5660 .arg(&ci)
5661 .arg(&mut *y)
5662 .arg(&ni)
5663 .arg(&slo)
5664 .arg(&shi)
5665 .arg(&stream_pos)
5666 .arg(&temp)
5667 .arg(stat_max)
5668 .arg(stat_th)
5669 .arg(&si);
5670 unsafe {
5671 b.launch(cfg)?;
5672 }
5673 Ok(())
5674 }
5675
5676 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
5681 let f = self.func("memra_sctr_inc");
5682 let cfg = LaunchConfig {
5683 grid_dim: (1, 1, 1),
5684 block_dim: (1, 1, 1),
5685 shared_mem_bytes: 0,
5686 };
5687 let __s_b = self.gpu.stream();
5688 let mut b = __s_b.launch_builder(&f);
5689 b.arg(&mut *ctr);
5690 unsafe {
5691 b.launch(cfg)?;
5692 }
5693 Ok(())
5694 }
5695
5696 pub fn gumbel_perturb_ctr(
5701 &self,
5702 x: &CudaSlice<f32>,
5703 y: &mut CudaSlice<f32>,
5704 n: usize,
5705 seed: u64,
5706 ctr: &CudaSlice<u32>,
5707 temp: f32,
5708 ) -> Result<(), Box<dyn std::error::Error>> {
5709 let f = self.func("gumbel_perturb_ctr_f32");
5710 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5711 let cfg = LaunchConfig {
5712 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5713 block_dim: (256, 1, 1),
5714 shared_mem_bytes: 0,
5715 };
5716 let __s_b = self.gpu.stream();
5717 let mut b = __s_b.launch_builder(&f);
5718 b.arg(x)
5719 .arg(&mut *y)
5720 .arg(&ni)
5721 .arg(&slo)
5722 .arg(&shi)
5723 .arg(ctr)
5724 .arg(&temp);
5725 unsafe {
5726 b.launch(cfg)?;
5727 }
5728 Ok(())
5729 }
5730
5731 #[allow(clippy::too_many_arguments)]
5739 pub fn gumbel_perturb_filtered_ctr(
5740 &self,
5741 x: &CudaSlice<f32>,
5742 y: &mut CudaSlice<f32>,
5743 n: usize,
5744 seed: u64,
5745 ctr: &CudaSlice<u32>,
5746 temp: f32,
5747 stat_max: &CudaSlice<f32>,
5748 stat_th: &CudaSlice<f32>,
5749 ) -> Result<(), Box<dyn std::error::Error>> {
5750 let f = self.func("gumbel_perturb_filtered_ctr_f32");
5751 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5752 let cfg = LaunchConfig {
5753 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5754 block_dim: (256, 1, 1),
5755 shared_mem_bytes: 0,
5756 };
5757 let __s_b = self.gpu.stream();
5758 let mut b = __s_b.launch_builder(&f);
5759 b.arg(x)
5760 .arg(&mut *y)
5761 .arg(&ni)
5762 .arg(&slo)
5763 .arg(&shi)
5764 .arg(ctr)
5765 .arg(&temp)
5766 .arg(stat_max)
5767 .arg(stat_th);
5768 unsafe {
5769 b.launch(cfg)?;
5770 }
5771 Ok(())
5772 }
5773
5774 #[allow(clippy::too_many_arguments)] pub fn softmax_gather(
5779 &self,
5780 x: &CudaSlice<f32>,
5781 row_stride: usize,
5782 ids: &CudaSlice<u32>,
5783 rows: &CudaSlice<i32>,
5784 out: &mut CudaSlice<f32>,
5785 n: usize,
5786 npair: usize,
5787 temp: f32,
5788 ) -> Result<(), Box<dyn std::error::Error>> {
5789 let f = self.func("softmax_gather_f32");
5790 let (ni, rs) = (n as i32, row_stride as i64);
5791 let np = npair as i32;
5792 let cfg = LaunchConfig {
5793 grid_dim: (npair as u32, 1, 1),
5794 block_dim: (256, 1, 1),
5795 shared_mem_bytes: 0,
5796 };
5797 let __s_b = self.gpu.stream();
5798 let mut b = __s_b.launch_builder(&f);
5799 b.arg(x)
5800 .arg(&rs)
5801 .arg(ids)
5802 .arg(rows)
5803 .arg(&mut *out)
5804 .arg(&ni)
5805 .arg(&np)
5806 .arg(&temp);
5807 unsafe {
5808 b.launch(cfg)?;
5809 }
5810 Ok(())
5811 }
5812
5813 #[allow(clippy::too_many_arguments)] pub fn residual_sample(
5818 &self,
5819 p: &CudaSlice<f32>,
5820 q: Option<&CudaSlice<f32>>,
5821 n: usize,
5822 temp: f32,
5823 seed: u64,
5824 stream_pos: u32,
5825 out_tok: &mut CudaSlice<u32>,
5826 ) -> Result<(), Box<dyn std::error::Error>> {
5827 let f = self.func("residual_sample_f32");
5828 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5829 let nth = 1024u32;
5830 let cfg = LaunchConfig {
5831 grid_dim: (1, 1, 1),
5832 block_dim: (nth, 1, 1),
5833 shared_mem_bytes: 0,
5834 };
5835 let has_q: i32 = q.is_some() as i32;
5836 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
5838 let mut b = __s_b.launch_builder(&f);
5839 b.arg(p)
5840 .arg(qbuf)
5841 .arg(&has_q)
5842 .arg(&ni)
5843 .arg(&temp)
5844 .arg(&slo)
5845 .arg(&shi)
5846 .arg(&stream_pos)
5847 .arg(&mut *out_tok);
5848 unsafe {
5849 b.launch(cfg)?;
5850 }
5851 Ok(())
5852 }
5853
5854 pub fn with_moe_cache<R>(
5859 &self,
5860 max_block_bytes: usize,
5861 f: impl FnOnce(
5862 &mut crate::moe_cache::MoeSlotCache,
5863 &Engine,
5864 ) -> Result<R, Box<dyn std::error::Error>>,
5865 ) -> Result<R, Box<dyn std::error::Error>> {
5866 let mut guard = self.moe_cache.lock().unwrap();
5867 if guard.is_none() {
5868 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
5869 }
5870 let cache = guard.as_mut().unwrap();
5871 f(cache, self)
5872 }
5873
5874 pub fn freeze_moe_cache(&self) {
5877 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
5878 cache.freeze();
5879 }
5880 }
5881
5882 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
5885 self.moe_cache
5886 .lock()
5887 .unwrap()
5888 .as_ref()
5889 .map(crate::moe_cache::MoeSlotCache::export_residency)
5890 }
5891
5892 pub(crate) fn moe_cache_frozen(&self) -> bool {
5893 self.moe_cache
5894 .lock()
5895 .unwrap()
5896 .as_ref()
5897 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
5898 }
5899
5900 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
5907 crate::cpu_experts::configured()
5908 && self.moe_cache_frozen()
5909 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
5910 }
5911
5912 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
5914 assert!(
5915 self.moe_cache.lock().unwrap().is_none(),
5916 "MoE cache layout configured after cache construction"
5917 );
5918 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
5919 }
5920
5921 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
5922 self.moe_cache_layout.lock().unwrap().clone()
5923 }
5924
5925 pub fn moe_cache_enabled() -> bool {
5927 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
5928 }
5929
5930 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
5933 let guard = self.moe_cache.lock().unwrap();
5934 guard
5935 .as_ref()
5936 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
5937 }
5938
5939 #[allow(clippy::type_complexity)] pub fn cpu_expert_stats(
5944 &self,
5945 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
5946 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
5947 }
5948
5949 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
5952 crate::cpu_experts::predictor_stats()
5953 }
5954
5955 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
5956 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
5957 }
5958
5959 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
5962 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
5963 }
5964
5965 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
5968 let guard = self.moe_cache.lock().unwrap();
5969 guard
5970 .as_ref()
5971 .and_then(|cache| cache.pread_stats())
5972 .map(|stats| {
5973 (
5974 stats.reads,
5975 stats.bytes,
5976 stats.read_errors,
5977 stats.short_reads,
5978 stats.fallbacks,
5979 stats.buffer_waits,
5980 stats.ring_full,
5981 )
5982 })
5983 }
5984
5985 pub fn spill_config_fallbacks(&self) -> u64 {
5987 crate::spill_pread::config_fallbacks()
5988 }
5989
5990 pub fn moe_cache_reset_counters(&self) {
5992 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
5993 c.reset_counters();
5994 }
5995 }
5996
5997 #[track_caller]
5998 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5999 crate::alloc_trace_hit(v.len());
6000 Ok(self.gpu.stream().clone_htod(v)?)
6001 }
6002
6003 pub fn htod_bytes_padded(
6007 &self,
6008 v: &[u8],
6009 pad: usize,
6010 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6011 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
6012 {
6013 let mut view = d.slice_mut(0..v.len());
6014 self.gpu.stream().memcpy_htod(v, &mut view)?;
6015 }
6016 Ok(d)
6017 }
6018
6019 pub fn copy_into(
6021 &self,
6022 dst: &mut CudaSlice<f32>,
6023 off: usize,
6024 src: &CudaSlice<f32>,
6025 len: usize,
6026 ) -> Result<(), Box<dyn std::error::Error>> {
6027 let mut view = dst.slice_mut(off..off + len);
6028 self.gpu
6029 .stream()
6030 .memcpy_dtod(&src.slice(0..len), &mut view)?;
6031 Ok(())
6032 }
6033
6034 pub fn copy_range_into(
6038 &self,
6039 dst: &mut CudaSlice<f32>,
6040 dst_off: usize,
6041 src: &CudaSlice<f32>,
6042 src_off: usize,
6043 len: usize,
6044 ) -> Result<(), Box<dyn std::error::Error>> {
6045 let mut view = dst.slice_mut(dst_off..dst_off + len);
6046 self.gpu
6047 .stream()
6048 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut view)?;
6049 Ok(())
6050 }
6051
6052 pub fn copy_u8_into(
6055 &self,
6056 dst: &mut CudaSlice<u8>,
6057 off: usize,
6058 src: &CudaSlice<u8>,
6059 len: usize,
6060 ) -> Result<(), Box<dyn std::error::Error>> {
6061 let cap = dst.len();
6065 let mut view = dst.try_slice_mut(off..off + len).ok_or_else(|| {
6066 format!(
6067 "copy_u8_into dst range [{off},{}) exceeds capacity {cap}",
6068 off + len,
6069 )
6070 })?;
6071 self.gpu
6072 .stream()
6073 .memcpy_dtod(&src.slice(0..len), &mut view)?;
6074 Ok(())
6075 }
6076
6077 pub fn copy_u8_range_into(
6079 &self,
6080 dst: &mut CudaSlice<u8>,
6081 dst_off: usize,
6082 src: &CudaSlice<u8>,
6083 src_off: usize,
6084 len: usize,
6085 ) -> Result<(), Box<dyn std::error::Error>> {
6086 let cap = dst.len();
6089 let mut dst_view = dst.try_slice_mut(dst_off..dst_off + len).ok_or_else(|| {
6090 format!(
6091 "copy_u8_range_into dst range [{dst_off},{}) exceeds capacity {cap}",
6092 dst_off + len,
6093 )
6094 })?;
6095 self.gpu
6096 .stream()
6097 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
6098 Ok(())
6099 }
6100
6101 #[track_caller]
6111 pub fn prepare_kv_append(
6112 &self,
6113 kv: &mut crate::cache::KvLayer,
6114 retain_from: usize,
6115 append_rows: usize,
6116 ) -> Result<usize, Box<dyn std::error::Error>> {
6117 let caller = std::panic::Location::caller();
6118 let base_before = kv.ring.as_ref().map(|r| r.base());
6119 let Some(plan) = kv
6120 .ring
6121 .as_ref()
6122 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
6123 .transpose()
6124 .map_err(|err| -> Box<dyn std::error::Error> {
6125 format!(
6126 "{err} [append len={} retain_from={retain_from} append_rows={append_rows} base={base_before:?} called from {caller}]",
6127 kv.len
6128 )
6129 .into()
6130 })?
6131 else {
6132 return Ok(kv.len);
6133 };
6134 match plan {
6135 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
6136 crate::cache::KvRingAppend::Rebase {
6137 src_row,
6138 keep_rows,
6139 new_base,
6140 write_row,
6141 } => {
6142 if keep_rows > 0 {
6143 let k_len = keep_rows * kv.k_tok_bytes;
6144 let v_len = keep_rows * kv.v_tok_bytes;
6145 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
6146 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
6147 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
6148 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
6149 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
6150 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
6151 }
6152 if std::env::var("MEMRA_KV_REBASE_TRACE").as_deref() == Ok("1") {
6155 eprintln!(
6156 "[kv-rebase] new_base={new_base} keep_rows={keep_rows} len={} \
6157 retain_from={retain_from} called from {caller}",
6158 kv.len
6159 );
6160 }
6161 kv.ring.as_mut().unwrap().apply_rebase(new_base);
6162 if let Some(base_d) = kv.base_d.as_mut() {
6166 self.set_i32_one(base_d, new_base as i32)?;
6167 }
6168 Ok(write_row)
6169 }
6170 }
6171 }
6172
6173 pub fn htod_u8_into(
6176 &self,
6177 dst: &mut CudaSlice<u8>,
6178 off: usize,
6179 src: &[u8],
6180 ) -> Result<(), Box<dyn std::error::Error>> {
6181 let mut view = dst.slice_mut(off..off + src.len());
6182 self.gpu.stream().memcpy_htod(src, &mut view)?;
6183 Ok(())
6184 }
6185
6186 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
6187 b.slice(0..len)
6188 }
6189
6190 pub fn view_u8_range<'a>(
6193 &self,
6194 b: &'a CudaSlice<u8>,
6195 start: usize,
6196 end: usize,
6197 ) -> cudarc::driver::CudaView<'a, u8> {
6198 b.slice(start..end)
6199 }
6200 pub fn view_u8<'a>(
6201 &self,
6202 b: &'a CudaSlice<u8>,
6203 len: usize,
6204 ) -> cudarc::driver::CudaView<'a, u8> {
6205 b.slice(0..len)
6206 }
6207
6208 #[allow(clippy::too_many_arguments)] pub fn append_kv_quantized(
6213 &self,
6214 k_row: &CudaSlice<f32>,
6215 v_row: &CudaSlice<f32>,
6216 kc: &mut CudaSlice<u8>,
6217 vc: &mut CudaSlice<u8>,
6218 t: usize,
6219 kv_dim_k: usize,
6220 kv_dim_v: usize,
6221 k_tok_bytes: usize,
6222 v_tok_bytes: usize,
6223 g: bool,
6224 ) -> Result<(), Box<dyn std::error::Error>> {
6225 let f = if g {
6226 self.func_g("append_quantize_kv_q8_0_q5_1")
6227 } else {
6228 self.func("append_quantize_kv_q8_0_q5_1")
6229 };
6230 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6231 let cfg = LaunchConfig {
6232 grid_dim: (nblk, 1, 1),
6233 block_dim: (32, 1, 1),
6234 shared_mem_bytes: 0,
6235 };
6236 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
6237 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6238 let __s_b = self.gpu.stream();
6239 let mut b = __s_b.launch_builder(&f);
6240 b.arg(k_row)
6241 .arg(v_row)
6242 .arg(kc)
6243 .arg(vc)
6244 .arg(&ti)
6245 .arg(&kdk)
6246 .arg(&kdv)
6247 .arg(&ktb)
6248 .arg(&vtb);
6249 unsafe {
6250 b.launch(cfg)?;
6251 }
6252 Ok(())
6253 }
6254
6255 #[allow(clippy::too_many_arguments)] pub fn append_kv_quantized_dc(
6260 &self,
6261 k_row: &CudaSlice<f32>,
6262 v_row: &CudaSlice<f32>,
6263 kc: &mut CudaSlice<u8>,
6264 vc: &mut CudaSlice<u8>,
6265 t_dev: &CudaSlice<i32>,
6266 kv_dim_k: usize,
6267 kv_dim_v: usize,
6268 k_tok_bytes: usize,
6269 v_tok_bytes: usize,
6270 g: bool,
6271 ) -> Result<(), Box<dyn std::error::Error>> {
6272 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6273 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
6274 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6275 if Self::pdl_on() && Self::pdl_wb_on() {
6277 use cudarc::driver::{DevicePtr, DevicePtrMut};
6278 let s = &self.gpu.stream();
6279 let (pk, _g0) = k_row.device_ptr(s);
6280 let (pv, _g1) = v_row.device_ptr(s);
6281 let (pkc, _g2) = kc.device_ptr_mut(s);
6282 let (pvc, _g3) = vc.device_ptr_mut(s);
6283 let (pt, _g4) = t_dev.device_ptr(s);
6284 let mut ps = [
6285 &pk as *const _ as *mut std::ffi::c_void,
6286 &pv as *const _ as *mut _,
6287 &pkc as *const _ as *mut _,
6288 &pvc as *const _ as *mut _,
6289 &pt as *const _ as *mut _,
6290 &kdk as *const _ as *mut _,
6291 &kdv as *const _ as *mut _,
6292 &ktb as *const _ as *mut _,
6293 &vtb as *const _ as *mut _,
6294 ];
6295 unsafe {
6296 self.launch_pdl_flash(
6297 g,
6298 "append_quantize_kv_q8_0_q5_1_dc",
6299 (nblk, 1, 1),
6300 (32, 1, 1),
6301 0,
6302 &mut ps,
6303 )?;
6304 }
6305 return Ok(());
6306 }
6307 let f = if g {
6308 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
6309 } else {
6310 self.func("append_quantize_kv_q8_0_q5_1_dc")
6311 };
6312 let cfg = LaunchConfig {
6313 grid_dim: (nblk, 1, 1),
6314 block_dim: (32, 1, 1),
6315 shared_mem_bytes: 0,
6316 };
6317 let __s_b = self.gpu.stream();
6318 let mut b = __s_b.launch_builder(&f);
6319 b.arg(k_row)
6320 .arg(v_row)
6321 .arg(kc)
6322 .arg(vc)
6323 .arg(t_dev)
6324 .arg(&kdk)
6325 .arg(&kdv)
6326 .arg(&ktb)
6327 .arg(&vtb);
6328 unsafe {
6329 b.launch(cfg)?;
6330 }
6331 Ok(())
6332 }
6333
6334 #[allow(clippy::too_many_arguments)]
6341 pub fn append_kv_quantized_rows(
6342 &self,
6343 k_rows: &CudaSlice<f32>,
6344 v_rows: &CudaSlice<f32>,
6345 kc: &mut CudaSlice<u8>,
6346 vc: &mut CudaSlice<u8>,
6347 t0: usize,
6348 t: usize,
6349 kv_dim_k: usize,
6350 kv_dim_v: usize,
6351 k_tok_bytes: usize,
6352 v_tok_bytes: usize,
6353 g: bool,
6354 ) -> Result<(), Box<dyn std::error::Error>> {
6355 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
6356 for i in 0..t {
6357 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
6358 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
6359 self.append_kv_quantized_view(
6360 &k_row,
6361 &v_row,
6362 kc,
6363 vc,
6364 t0 + i,
6365 kv_dim_k,
6366 kv_dim_v,
6367 k_tok_bytes,
6368 v_tok_bytes,
6369 g,
6370 )?;
6371 }
6372 return Ok(());
6373 }
6374 let f = if g {
6375 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
6376 } else {
6377 self.func("append_quantize_kv_q8_0_q5_1_rows")
6378 };
6379 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6380 let cfg = LaunchConfig {
6381 grid_dim: (nblk, t as u32, 1),
6382 block_dim: (32, 1, 1),
6383 shared_mem_bytes: 0,
6384 };
6385 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
6386 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6387 let __s_b = self.gpu.stream();
6388 let mut b = __s_b.launch_builder(&f);
6389 b.arg(k_rows)
6390 .arg(v_rows)
6391 .arg(kc)
6392 .arg(vc)
6393 .arg(&t0i)
6394 .arg(&kdk)
6395 .arg(&kdv)
6396 .arg(&ktb)
6397 .arg(&vtb);
6398 unsafe {
6399 b.launch(cfg)?;
6400 }
6401 Ok(())
6402 }
6403
6404 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
6408 let f = self.func("inc_i32");
6409 let cfg = LaunchConfig {
6410 grid_dim: (1, 1, 1),
6411 block_dim: (1, 1, 1),
6412 shared_mem_bytes: 0,
6413 };
6414 let __s_b = self.gpu.stream();
6415 let mut b = __s_b.launch_builder(&f);
6416 b.arg(p);
6417 unsafe {
6418 b.launch(cfg)?;
6419 }
6420 Ok(())
6421 }
6422
6423 #[allow(clippy::too_many_arguments)] pub fn append_kv_quantized_view(
6427 &self,
6428 k_row: &cudarc::driver::CudaView<f32>,
6429 v_row: &cudarc::driver::CudaView<f32>,
6430 kc: &mut CudaSlice<u8>,
6431 vc: &mut CudaSlice<u8>,
6432 t: usize,
6433 kv_dim_k: usize,
6434 kv_dim_v: usize,
6435 k_tok_bytes: usize,
6436 v_tok_bytes: usize,
6437 g: bool,
6438 ) -> Result<(), Box<dyn std::error::Error>> {
6439 let stream = self.gpu.stream();
6440 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
6441 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
6442 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
6443 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
6444 let f = if g {
6445 self.func_g("append_quantize_kv_q8_0_q5_1")
6446 } else {
6447 self.func("append_quantize_kv_q8_0_q5_1")
6448 };
6449 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6450 let cfg = LaunchConfig {
6451 grid_dim: (nblk, 1, 1),
6452 block_dim: (32, 1, 1),
6453 shared_mem_bytes: 0,
6454 };
6455 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
6456 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6457 let mut b = stream.launch_builder(&f);
6458 b.arg(k_row)
6459 .arg(v_row)
6460 .arg(kc)
6461 .arg(vc)
6462 .arg(&ti)
6463 .arg(&kdk)
6464 .arg(&kdv)
6465 .arg(&ktb)
6466 .arg(&vtb);
6467 unsafe {
6468 b.launch(cfg)?;
6469 }
6470 Ok(())
6471 }
6472
6473 pub fn copy_view_into(
6476 &self,
6477 dst: &mut CudaSlice<f32>,
6478 off: usize,
6479 src: &cudarc::driver::CudaView<f32>,
6480 len: usize,
6481 ) -> Result<(), Box<dyn std::error::Error>> {
6482 let mut view = dst.slice_mut(off..off + len);
6483 self.gpu
6484 .stream()
6485 .memcpy_dtod(&src.slice(0..len), &mut view)?;
6486 Ok(())
6487 }
6488
6489 pub fn clone_dtod(
6510 &self,
6511 src: &CudaSlice<f32>,
6512 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6513 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
6514 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
6515 Ok(dst)
6516 }
6517
6518 pub fn dtod_copy_view(
6521 &self,
6522 src: &cudarc::driver::CudaView<f32>,
6523 dst: &mut CudaSlice<f32>,
6524 ) -> Result<(), Box<dyn std::error::Error>> {
6525 self.gpu.stream().memcpy_dtod(src, dst)?;
6526 Ok(())
6527 }
6528
6529 pub fn dtod_copy_view_i8(
6531 &self,
6532 src: &cudarc::driver::CudaView<i8>,
6533 dst: &mut CudaSlice<i8>,
6534 ) -> Result<(), Box<dyn std::error::Error>> {
6535 self.gpu.stream().memcpy_dtod(src, dst)?;
6536 Ok(())
6537 }
6538
6539 pub fn dtod_copy_into(
6541 &self,
6542 src: &CudaSlice<f32>,
6543 dst: &mut CudaSlice<f32>,
6544 offset: usize,
6545 ) -> Result<(), Box<dyn std::error::Error>> {
6546 let n = src.len();
6547 let mut dv = dst.slice_mut(offset..offset + n);
6548 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
6549 Ok(())
6550 }
6551
6552 pub fn copy_batch_uniform_f32(
6558 &self,
6559 table: &CudaSlice<u64>,
6560 n: usize,
6561 words: usize,
6562 ) -> Result<(), Box<dyn std::error::Error>> {
6563 if n == 0 || words == 0 {
6564 return Ok(());
6565 }
6566 debug_assert!(
6567 table.len() >= 2 * n,
6568 "pointer table must hold n srcs + n dsts"
6569 );
6570 let f = self.func("copy_batch_uniform_f32");
6571 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
6574 let (ni, wi) = (n as i32, words as i32);
6575 let cfg = LaunchConfig {
6576 grid_dim: (chunks, n as u32, 1),
6577 block_dim: (256, 1, 1),
6578 shared_mem_bytes: 0,
6579 };
6580 let __s = self.gpu.stream();
6581 let mut b = __s.launch_builder(&f);
6582 b.arg(table).arg(&ni).arg(&wi);
6583 unsafe {
6584 b.launch(cfg)?;
6585 }
6586 Ok(())
6587 }
6588
6589 #[allow(clippy::too_many_arguments)] pub fn copy_batch_uniform_kv_u8_set_len(
6594 &self,
6595 table: &CudaSlice<u64>,
6596 n: usize,
6597 rows: usize,
6598 k_row_bytes: usize,
6599 v_row_bytes: usize,
6600 k_src_stride: usize,
6601 v_src_stride: usize,
6602 logical_len: usize,
6603 ) -> Result<(), Box<dyn std::error::Error>> {
6604 if n == 0 || rows == 0 || (k_row_bytes == 0 && v_row_bytes == 0) {
6605 return Ok(());
6606 }
6607 if table.len() < 5 * n {
6608 return Err(format!(
6609 "TP KV repair table has {} words, expected at least {}",
6610 table.len(),
6611 5 * n
6612 )
6613 .into());
6614 }
6615 let ni = i32::try_from(n).map_err(|_| "TP KV repair layer count exceeds i32")?;
6616 let rows = i32::try_from(rows).map_err(|_| "TP KV repair rows exceed i32")?;
6617 let kb = i32::try_from(k_row_bytes).map_err(|_| "TP KV repair K bytes exceed i32")?;
6618 let vb = i32::try_from(v_row_bytes).map_err(|_| "TP KV repair V bytes exceed i32")?;
6619 let ks = i32::try_from(k_src_stride).map_err(|_| "TP KV repair K stride exceeds i32")?;
6620 let vs = i32::try_from(v_src_stride).map_err(|_| "TP KV repair V stride exceeds i32")?;
6621 let len = i32::try_from(logical_len).map_err(|_| "TP KV repair length exceeds i32")?;
6622 let f = self.func("copy_batch_uniform_kv_u8_set_len");
6623 let cfg = LaunchConfig {
6624 grid_dim: (n as u32, 1, 1),
6625 block_dim: (256, 1, 1),
6626 shared_mem_bytes: 0,
6627 };
6628 let stream = self.gpu.stream();
6629 let mut builder = stream.launch_builder(&f);
6630 builder
6631 .arg(table)
6632 .arg(&ni)
6633 .arg(&rows)
6634 .arg(&kb)
6635 .arg(&vb)
6636 .arg(&ks)
6637 .arg(&vs)
6638 .arg(&len);
6639 unsafe {
6640 builder.launch(cfg)?;
6641 }
6642 Ok(())
6643 }
6644
6645 pub fn htod_u64_into(
6648 &self,
6649 v: &[u64],
6650 dst: &mut CudaSlice<u64>,
6651 ) -> Result<(), Box<dyn std::error::Error>> {
6652 let mut view = dst.slice_mut(0..v.len());
6653 self.gpu.stream().memcpy_htod(v, &mut view)?;
6654 Ok(())
6655 }
6656
6657 pub fn htod_f32_into(
6660 &self,
6661 v: &[f32],
6662 dst: &mut CudaSlice<f32>,
6663 ) -> Result<(), Box<dyn std::error::Error>> {
6664 let mut view = dst.slice_mut(0..v.len());
6665 self.gpu.stream().memcpy_htod(v, &mut view)?;
6666 Ok(())
6667 }
6668
6669 pub fn htod_f32_into_at(
6673 &self,
6674 v: &[f32],
6675 dst: &mut CudaSlice<f32>,
6676 off: usize,
6677 ) -> Result<(), Box<dyn std::error::Error>> {
6678 if off + v.len() > dst.len() {
6679 return Err(format!(
6680 "htod_f32_into_at range {}..{} exceeds dst {}",
6681 off,
6682 off + v.len(),
6683 dst.len()
6684 )
6685 .into());
6686 }
6687 let mut view = dst.slice_mut(off..off + v.len());
6688 self.gpu.stream().memcpy_htod(v, &mut view)?;
6689 Ok(())
6690 }
6691
6692 pub fn copy_indirect_src_f32(
6697 &self,
6698 src_entry: &cudarc::driver::CudaView<u64>,
6699 dst: &mut CudaSlice<f32>,
6700 dst_off: usize,
6701 words: usize,
6702 ) -> Result<(), Box<dyn std::error::Error>> {
6703 let f = self.func("copy_indirect_src_f32");
6704 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
6705 let wi = words as i32;
6706 let cfg = LaunchConfig {
6707 grid_dim: (chunks, 1, 1),
6708 block_dim: (256, 1, 1),
6709 shared_mem_bytes: 0,
6710 };
6711 let mut dv = dst.slice_mut(dst_off..dst_off + words);
6712 let __s = self.gpu.stream();
6713 let mut b = __s.launch_builder(&f);
6714 b.arg(src_entry).arg(&mut dv).arg(&wi);
6715 unsafe {
6716 b.launch(cfg)?;
6717 }
6718 Ok(())
6719 }
6720
6721 #[track_caller]
6723 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6724 self.alloc_uninit::<i8>(n)
6725 }
6726
6727 #[allow(clippy::too_many_arguments)] pub fn qmatvec(
6730 &self,
6731 w: &CudaSlice<u8>,
6732 x: &CudaSlice<f32>,
6733 m: usize,
6734 in_f: usize,
6735 out_f: usize,
6736 qtype: i32,
6737 row_bytes: usize,
6738 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6739 let f = self.func("qmatvec_f32");
6740 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6742 grid_dim: (out_f as u32, m as u32, 1),
6743 block_dim: (256, 1, 1),
6744 shared_mem_bytes: 0,
6745 };
6746 let (inf, outf, mi, qt, rb) =
6747 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
6748 let __s_b = self.gpu.stream();
6749 let mut b = __s_b.launch_builder(&f);
6750 b.arg(w)
6751 .arg(x)
6752 .arg(&mut y)
6753 .arg(&inf)
6754 .arg(&outf)
6755 .arg(&mi)
6756 .arg(&qt)
6757 .arg(&rb);
6758 unsafe {
6759 b.launch(cfg)?;
6760 }
6761 Ok(y)
6762 }
6763
6764 #[track_caller]
6766 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6767 crate::alloc_trace_hit(n);
6768 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
6769 self.keep_if_capturing(&s);
6770 Ok(s)
6771 }
6772
6773 #[track_caller]
6777 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6778 crate::alloc_trace_hit(n);
6779 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
6780 self.keep_if_capturing(&s);
6781 Ok(s)
6782 }
6783
6784 pub fn memset_zeros_view(
6787 &self,
6788 dst: &mut cudarc::driver::CudaViewMut<f32>,
6789 ) -> Result<(), Box<dyn std::error::Error>> {
6790 self.gpu.stream().memset_zeros(dst)?;
6791 Ok(())
6792 }
6793
6794 pub fn stage_expert(
6800 &self,
6801 host_bytes: &[u8],
6802 scratch: &mut CudaSlice<u8>,
6803 off: usize,
6804 ) -> Result<(), Box<dyn std::error::Error>> {
6805 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
6808 }
6809
6810 pub fn nvfp4_expert_split_repack(
6814 &self,
6815 src: &CudaSlice<u8>,
6816 n_expert: usize,
6817 rows: usize,
6818 nsb64: usize,
6819 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6820 let need = n_expert * rows * nsb64 * 36;
6821 if src.len() < need {
6822 return Err(format!(
6823 "nvfp4_expert_split_repack: slab holds {} bytes, {n_expert} x {rows} x {nsb64} x 36 = {need} needed",
6824 src.len()
6825 )
6826 .into());
6827 }
6828 let f = self.func("nvfp4_expert_split_repack");
6829 let mut dst = self.alloc_u8_uninit(src.len())?; let nblk = (n_expert * rows * nsb64) as u32;
6831 let cfg = LaunchConfig {
6832 grid_dim: (nblk.div_ceil(256), 1, 1),
6833 block_dim: (256, 1, 1),
6834 shared_mem_bytes: 0,
6835 };
6836 let (ne, nr, ns) = (n_expert as i32, rows as i32, nsb64 as i32);
6837 let __s_b = self.gpu.stream();
6838 let mut b = __s_b.launch_builder(&f);
6839 b.arg(src).arg(&mut dst).arg(&ne).arg(&nr).arg(&ns);
6840 unsafe {
6841 b.launch(cfg)?;
6842 }
6843 Ok(dst)
6844 }
6845
6846 pub fn moe_router_topk(
6852 &self,
6853 logits: &CudaSlice<f32>,
6854 t: usize,
6855 n_expert: usize,
6856 n_used: usize,
6857 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6858 let f = self.func("moe_router_topk_f32");
6859 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 {
6862 grid_dim: (t as u32, 1, 1),
6863 block_dim: (n_expert as u32, 1, 1),
6864 shared_mem_bytes: 0,
6865 };
6866 let (ne, nu) = (n_expert as i32, n_used as i32);
6867 let __s_b = self.gpu.stream();
6868 let mut b = __s_b.launch_builder(&f);
6869 b.arg(logits)
6870 .arg(&mut sel_idx)
6871 .arg(&mut sel_w)
6872 .arg(&ne)
6873 .arg(&nu);
6874 unsafe {
6875 b.launch(cfg)?;
6876 }
6877 Ok((sel_idx, sel_w))
6878 }
6879
6880 pub fn moe_router_topk_scaled(
6883 &self,
6884 logits: &CudaSlice<f32>,
6885 t: usize,
6886 n_expert: usize,
6887 n_used: usize,
6888 ex_scale: &CudaSlice<f32>,
6889 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6890 let f = self.func("moe_router_topk_scaled_f32");
6895 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
6896 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
6897 let cfg = LaunchConfig {
6898 grid_dim: (t as u32, 1, 1),
6899 block_dim: (n_expert as u32, 1, 1),
6900 shared_mem_bytes: 0,
6901 };
6902 let (ne, nu) = (n_expert as i32, n_used as i32);
6903 let __s_b = self.gpu.stream();
6904 let mut b = __s_b.launch_builder(&f);
6905 b.arg(logits)
6906 .arg(&mut sel_idx)
6907 .arg(&mut sel_w)
6908 .arg(&ne)
6909 .arg(&nu)
6910 .arg(ex_scale);
6911 unsafe {
6912 b.launch(cfg)?;
6913 }
6914 Ok((sel_idx, sel_w))
6915 }
6916
6917 pub fn moe_router_topk_host(
6925 &self,
6926 logits: &CudaSlice<f32>,
6927 t: usize,
6928 n_expert: usize,
6929 n_used: usize,
6930 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6931 let f = self.func("moe_router_topk_f32");
6932 let n = t * n_used;
6933 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
6934 let mut sel_w = self.alloc_uninit::<f32>(n)?;
6935 let cfg = LaunchConfig {
6936 grid_dim: (t as u32, 1, 1),
6937 block_dim: (n_expert as u32, 1, 1),
6938 shared_mem_bytes: 0,
6939 };
6940 let (ne, nu) = (n_expert as i32, n_used as i32);
6941 let __s_b = self.gpu.stream();
6942 let mut b = __s_b.launch_builder(&f);
6943 b.arg(logits)
6944 .arg(&mut sel_idx)
6945 .arg(&mut sel_w)
6946 .arg(&ne)
6947 .arg(&nu);
6948 unsafe {
6949 b.launch(cfg)?;
6950 }
6951 let bytes = n * 8;
6953 let mut guard = self.router_stage.lock().unwrap();
6954 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
6955 *guard = Some(PinnedStage::new(bytes.max(4096))?);
6956 }
6957 let stage = guard.as_mut().unwrap();
6958 let (si, sw) = unsafe {
6959 (
6960 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
6961 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
6962 )
6963 };
6964 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()))
6968 }
6969
6970 #[allow(clippy::too_many_arguments)]
6974 pub fn moe_router_sigmoid_topk(
6975 &self,
6976 logits: &CudaSlice<f32>,
6977 t: usize,
6978 n_expert: usize,
6979 n_used: usize,
6980 active_count: usize,
6981 correction_bias: &CudaSlice<f32>,
6982 active: &CudaSlice<u8>,
6983 scaling_factor: f32,
6984 route_norm: bool,
6985 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6986 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
6987 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
6988 return Err(format!(
6989 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
6990 )
6991 .into());
6992 }
6993 if logits.len() < t * n_expert
6994 || correction_bias.len() != n_expert
6995 || active.len() != n_expert
6996 {
6997 return Err(format!(
6998 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
6999 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
7000 ).into());
7001 }
7002 let f = self.func(crate::sigmoid_topk_kernel(
7003 crate::sig_expf_dev_on(),
7004 crate::topk_fast_on(),
7005 n_used,
7006 ));
7007 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
7008 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
7009 let threads = n_expert.div_ceil(32) * 32;
7010 let cfg = LaunchConfig {
7011 grid_dim: (t as u32, 1, 1),
7012 block_dim: (threads as u32, 1, 1),
7013 shared_mem_bytes: 0,
7014 };
7015 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
7016 let __s_b = self.gpu.stream();
7017 let mut b = __s_b.launch_builder(&f);
7018 b.arg(logits)
7019 .arg(correction_bias)
7020 .arg(active)
7021 .arg(&mut sel_idx)
7022 .arg(&mut sel_w)
7023 .arg(&ne)
7024 .arg(&nu)
7025 .arg(&scaling_factor)
7026 .arg(&rn);
7027 unsafe {
7028 b.launch(cfg)?;
7029 }
7030 Ok((sel_idx, sel_w))
7031 }
7032
7033 #[allow(clippy::too_many_arguments)]
7036 pub fn moe_sel_w_mirror(
7039 &self,
7040 sel_src: &CudaSlice<i32>,
7041 w_src: &CudaSlice<f32>,
7042 sel_dst: &mut CudaSlice<i32>,
7043 w_dst: &mut CudaSlice<f32>,
7044 n: usize,
7045 ) -> Result<(), Box<dyn std::error::Error>> {
7046 if n == 0
7047 || n > i32::MAX as usize
7048 || sel_src.len() < n
7049 || w_src.len() < n
7050 || sel_dst.len() < n
7051 || w_dst.len() < n
7052 {
7053 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
7054 }
7055 let f = self.func("moe_sel_w_mirror");
7056 let threads = if n <= 32 { 32 } else { 128 };
7057 let cfg = LaunchConfig {
7058 grid_dim: ((n as u32).div_ceil(threads), 1, 1),
7059 block_dim: (threads, 1, 1),
7060 shared_mem_bytes: 0,
7061 };
7062 let ni = n as i32;
7063 let __s_b = self.gpu.stream();
7064 let mut b = __s_b.launch_builder(&f);
7065 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
7066 unsafe {
7067 b.launch(cfg)?;
7068 }
7069 Ok(())
7070 }
7071
7072 #[allow(clippy::too_many_arguments)]
7076 pub fn nvfp4_ep_stage_inputs(
7077 &self,
7078 input_src: &CudaSlice<f32>,
7079 sel_src: &CudaSlice<i32>,
7080 w_src: &CudaSlice<f32>,
7081 input_bf16_dst: &mut CudaSlice<u8>,
7082 sel_dst: &mut CudaSlice<i32>,
7083 w_dst: &mut CudaSlice<f32>,
7084 input_values: usize,
7085 pairs: usize,
7086 copy_weights: bool,
7087 ) -> Result<(), Box<dyn std::error::Error>> {
7088 if input_values == 0
7089 || pairs == 0
7090 || input_src.len() < input_values
7091 || sel_src.len() < pairs
7092 || w_src.len() < pairs
7093 || input_bf16_dst.len() < 2 * input_values
7094 || sel_dst.len() < pairs
7095 || w_dst.len() < pairs
7096 {
7097 return Err(format!(
7098 "W4A16 EP stage geometry input={} sel={} weights={} input_bf16={} \
7099 sel_dst={} weights_dst={} active={input_values} pairs={pairs}",
7100 input_src.len(),
7101 sel_src.len(),
7102 w_src.len(),
7103 input_bf16_dst.len(),
7104 sel_dst.len(),
7105 w_dst.len(),
7106 )
7107 .into());
7108 }
7109 let f = self.func("nvfp4_ep_stage_inputs");
7110 let n = input_values.max(pairs);
7111 let cfg = LaunchConfig::for_num_elems(n as u32);
7112 let (input_values, pairs, copy_weights) =
7113 (input_values as i32, pairs as i32, i32::from(copy_weights));
7114 let __s_b = self.gpu.stream();
7115 let mut b = __s_b.launch_builder(&f);
7116 b.arg(input_src)
7117 .arg(sel_src)
7118 .arg(w_src)
7119 .arg(input_bf16_dst)
7120 .arg(sel_dst)
7121 .arg(w_dst)
7122 .arg(&input_values)
7123 .arg(&pairs)
7124 .arg(©_weights);
7125 unsafe {
7126 b.launch(cfg)?;
7127 }
7128 Ok(())
7129 }
7130
7131 #[allow(clippy::too_many_arguments)]
7134 pub fn nvfp4_ep_stage_inputs_raw(
7135 &self,
7136 input_src: u64,
7137 sel_src: u64,
7138 w_src: u64,
7139 input_bf16_dst: &mut CudaSlice<u8>,
7140 sel_dst: &mut CudaSlice<i32>,
7141 w_dst: &mut CudaSlice<f32>,
7142 input_values: usize,
7143 pairs: usize,
7144 copy_weights: bool,
7145 ) -> Result<(), Box<dyn std::error::Error>> {
7146 if input_src == 0
7147 || sel_src == 0
7148 || w_src == 0
7149 || input_values == 0
7150 || pairs == 0
7151 || input_bf16_dst.len() < 2 * input_values
7152 || sel_dst.len() < pairs
7153 || w_dst.len() < pairs
7154 {
7155 return Err(format!(
7156 "W4A16 EP raw stage geometry input={input_src:#x} sel={sel_src:#x} \
7157 weights={w_src:#x} input_bf16={} sel_dst={} weights_dst={} \
7158 active={input_values} pairs={pairs}",
7159 input_bf16_dst.len(),
7160 sel_dst.len(),
7161 w_dst.len(),
7162 )
7163 .into());
7164 }
7165 let f = self.func("nvfp4_ep_stage_inputs");
7166 let n = input_values.max(pairs);
7167 let cfg = LaunchConfig::for_num_elems(n as u32);
7168 let (input_values, pairs, copy_weights) =
7169 (input_values as i32, pairs as i32, i32::from(copy_weights));
7170 let __s_b = self.gpu.stream();
7171 let mut b = __s_b.launch_builder(&f);
7172 b.arg(&input_src)
7173 .arg(&sel_src)
7174 .arg(&w_src)
7175 .arg(input_bf16_dst)
7176 .arg(sel_dst)
7177 .arg(w_dst)
7178 .arg(&input_values)
7179 .arg(&pairs)
7180 .arg(©_weights);
7181 unsafe {
7182 b.launch(cfg)?;
7183 }
7184 Ok(())
7185 }
7186
7187 #[allow(clippy::too_many_arguments)] pub fn moe_router_sigmoid_topk_into(
7189 &self,
7190 logits: &CudaSlice<f32>,
7191 t: usize,
7192 n_expert: usize,
7193 n_used: usize,
7194 active_count: usize,
7195 correction_bias: &CudaSlice<f32>,
7196 active: &CudaSlice<u8>,
7197 scaling_factor: f32,
7198 route_norm: bool,
7199 sel_idx: &mut CudaSlice<i32>,
7200 sel_w: &mut CudaSlice<f32>,
7201 ) -> Result<(), Box<dyn std::error::Error>> {
7202 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
7203 if n_expert == 0
7204 || n_expert > 1024
7205 || n_used == 0
7206 || n_used > 32 || n_used > n_expert
7208 || logits.len() < t * n_expert
7209 || correction_bias.len() != n_expert
7210 || active.len() != n_expert
7211 || sel_idx.len() < t * n_used
7212 || sel_w.len() < t * n_used
7213 {
7214 return Err("sigmoid router _into geometry mismatch".into());
7215 }
7216 let f = self.func(crate::sigmoid_topk_kernel(
7217 crate::sig_expf_dev_on(),
7218 crate::topk_fast_on(),
7219 n_used,
7220 ));
7221 let threads = n_expert.div_ceil(32) * 32;
7222 let cfg = LaunchConfig {
7223 grid_dim: (t as u32, 1, 1),
7224 block_dim: (threads as u32, 1, 1),
7225 shared_mem_bytes: 0,
7226 };
7227 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
7228 let __s_b = self.gpu.stream();
7229 let mut b = __s_b.launch_builder(&f);
7230 b.arg(logits)
7231 .arg(correction_bias)
7232 .arg(active)
7233 .arg(&mut *sel_idx)
7234 .arg(&mut *sel_w)
7235 .arg(&ne)
7236 .arg(&nu)
7237 .arg(&scaling_factor)
7238 .arg(&rn);
7239 unsafe {
7240 b.launch(cfg)?;
7241 }
7242 Ok(())
7243 }
7244
7245 #[allow(clippy::too_many_arguments)]
7248 pub fn moe_router_sigmoid_topk_host(
7249 &self,
7250 logits: &CudaSlice<f32>,
7251 t: usize,
7252 n_expert: usize,
7253 n_used: usize,
7254 active_count: usize,
7255 correction_bias: &CudaSlice<f32>,
7256 active: &CudaSlice<u8>,
7257 scaling_factor: f32,
7258 route_norm: bool,
7259 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
7260 if glm5_graph_capture_open() {
7274 return Err("moe_router_sigmoid_topk_host was reached inside an open CUDA graph capture: the host readback would record an unexecuted DtoH and route on uninitialised memory. The captured layer's device-table MoE arm did not fire, so the capture-admission predicate (glm5_t1_dev_moe_ready) and the dispatch predicate (vrows_fires) disagree"
7275 .into());
7276 }
7277 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
7278 logits,
7279 t,
7280 n_expert,
7281 n_used,
7282 active_count,
7283 correction_bias,
7284 active,
7285 scaling_factor,
7286 route_norm,
7287 )?;
7288 let n = t * n_used;
7289 let bytes = n * 8;
7290 let mut guard = self.router_stage.lock().unwrap();
7291 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
7292 *guard = Some(PinnedStage::new(bytes.max(4096))?);
7293 }
7294 let stage = guard.as_mut().unwrap();
7295 let (si, sw) = unsafe {
7296 (
7297 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
7298 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
7299 )
7300 };
7301 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
7302 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
7303 self.gpu.stream().synchronize()?;
7304 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
7305 }
7306
7307 pub fn stage_expert_async(
7311 &self,
7312 host_bytes: &[u8],
7313 scratch: &mut CudaSlice<u8>,
7314 off: usize,
7315 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
7316 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
7317 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
7318 Ok(self.copy_stream.record_event(None)?)
7319 }
7320
7321 pub fn compute_wait(
7323 &self,
7324 ev: &cudarc::driver::CudaEvent,
7325 ) -> Result<(), Box<dyn std::error::Error>> {
7326 self.gpu.stream().wait(ev)?;
7327 Ok(())
7328 }
7329
7330 #[allow(clippy::too_many_arguments)] pub fn qmatvec_view(
7336 &self,
7337 w: &CudaSlice<u8>,
7338 range: std::ops::Range<usize>,
7339 x: &cudarc::driver::CudaView<f32>,
7340 m: usize,
7341 in_f: usize,
7342 out_f: usize,
7343 qtype: i32,
7344 row_bytes: usize,
7345 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7346 self.qmatvec_view_inner(w, range, x, m, in_f, out_f, qtype, row_bytes)
7347 }
7348
7349 #[allow(clippy::too_many_arguments)]
7353 pub fn qmatvec_view_bf16_activation(
7354 &self,
7355 w: &CudaSlice<u8>,
7356 range: std::ops::Range<usize>,
7357 x: &cudarc::driver::CudaView<f32>,
7358 m: usize,
7359 in_f: usize,
7360 out_f: usize,
7361 qtype: i32,
7362 row_bytes: usize,
7363 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7364 let n = m * in_f;
7365 if x.len() != n {
7366 return Err(format!(
7367 "W4A16 BF16 activation input length {} != {m}x{in_f}",
7368 x.len()
7369 )
7370 .into());
7371 }
7372 let mut x_bf16 = self.alloc_u8_uninit(n * 2)?;
7373 self.f32_to_bf16_v(x, &mut x_bf16, n)?;
7374 let x_f32 = self.bf16_to_f32(&x_bf16.slice(0..n * 2), n)?;
7375 self.qmatvec_view_inner(
7376 w,
7377 range,
7378 &x_f32.slice(0..n),
7379 m,
7380 in_f,
7381 out_f,
7382 qtype,
7383 row_bytes,
7384 )
7385 }
7386
7387 #[allow(clippy::too_many_arguments)]
7388 fn qmatvec_view_inner(
7389 &self,
7390 w: &CudaSlice<u8>,
7391 range: std::ops::Range<usize>,
7392 x: &cudarc::driver::CudaView<f32>,
7393 m: usize,
7394 in_f: usize,
7395 out_f: usize,
7396 qtype: i32,
7397 row_bytes: usize,
7398 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7399 let f = self.func("qmatvec_f32");
7400 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7403 grid_dim: (out_f as u32, m as u32, 1),
7404 block_dim: (256, 1, 1),
7405 shared_mem_bytes: 0,
7406 };
7407 let (inf, outf, mi, qt, rb) =
7408 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
7409 let __s_b = self.gpu.stream();
7410 let mut b = __s_b.launch_builder(&f);
7411 b.arg(&wv)
7412 .arg(x)
7413 .arg(&mut y)
7414 .arg(&inf)
7415 .arg(&outf)
7416 .arg(&mi)
7417 .arg(&qt)
7418 .arg(&rb);
7419 unsafe {
7420 b.launch(cfg)?;
7421 }
7422 Ok(y)
7423 }
7424
7425 #[allow(clippy::too_many_arguments)]
7432 pub fn moe_gate_up_silu8_q8(
7436 &self,
7437 gp: WPtr8,
7438 up: WPtr8,
7439 aq: &CudaSlice<i8>,
7440 ad: &CudaSlice<f32>,
7441 in_f: usize,
7442 n_ff: usize,
7443 n_used: usize,
7444 qt_g: i32,
7445 qt_u: i32,
7446 rb_g: usize,
7447 rb_u: usize,
7448 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7449 let f = self.func("moe_gate_up_silu8_q8");
7450 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7451 let cfg = LaunchConfig {
7452 grid_dim: (n_ff as u32, n_used as u32, 1),
7453 block_dim: (32, 1, 1),
7454 shared_mem_bytes: 0,
7455 };
7456 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7457 let __s_b = self.gpu.stream();
7458 let mut b = __s_b.launch_builder(&f);
7459 b.arg(&gp)
7460 .arg(&up)
7461 .arg(aq)
7462 .arg(ad)
7463 .arg(&mut act)
7464 .arg(&inf)
7465 .arg(&nff)
7466 .arg(&qt_g)
7467 .arg(&qt_u)
7468 .arg(&rbg)
7469 .arg(&rbu);
7470 unsafe {
7471 b.launch(cfg)?;
7472 }
7473 Ok(act)
7474 }
7475
7476 #[allow(clippy::too_many_arguments)]
7489 pub fn moe_gate_up_preclamp8_q8(
7490 &self,
7491 gp: WPtr8,
7492 up: WPtr8,
7493 aq: &CudaSlice<i8>,
7494 ad: &CudaSlice<f32>,
7495 gs: F32x8,
7496 us: F32x8,
7497 limit: f32,
7498 in_f: usize,
7499 n_ff: usize,
7500 n_used: usize,
7501 qt_g: i32,
7502 qt_u: i32,
7503 rb_g: usize,
7504 rb_u: usize,
7505 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7506 debug_assert!(
7507 limit > 1e-6,
7508 "moe_gate_up_preclamp8_q8 needs a live limit; use moe_gate_up_silu8_q8"
7509 );
7510 let f = self.func("moe_gate_up_preclamp8_q8");
7511 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7512 let cfg = LaunchConfig {
7513 grid_dim: (n_ff as u32, n_used as u32, 1),
7514 block_dim: (32, 1, 1),
7515 shared_mem_bytes: 0,
7516 };
7517 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7518 let __s_b = self.gpu.stream();
7519 let mut b = __s_b.launch_builder(&f);
7520 b.arg(&gp)
7521 .arg(&up)
7522 .arg(aq)
7523 .arg(ad)
7524 .arg(&gs)
7525 .arg(&us)
7526 .arg(&limit)
7527 .arg(&mut act)
7528 .arg(&inf)
7529 .arg(&nff)
7530 .arg(&qt_g)
7531 .arg(&qt_u)
7532 .arg(&rbg)
7533 .arg(&rbu);
7534 unsafe {
7535 b.launch(cfg)?;
7536 }
7537 Ok(act)
7538 }
7539
7540 #[allow(clippy::too_many_arguments)]
7541 pub fn moe_down8_fma_q8(
7542 &self,
7543 dp: WPtr8,
7544 w: F32x8,
7545 aq2: &CudaSlice<i8>,
7546 ad2: &CudaSlice<f32>,
7547 dst: &mut cudarc::driver::CudaViewMut<f32>,
7548 in_f: usize,
7549 out_f: usize,
7550 n_used: usize,
7551 qt: i32,
7552 rb: usize,
7553 ) -> Result<(), Box<dyn std::error::Error>> {
7554 let f = self.func("moe_down8_fma_q8");
7555 let cfg = LaunchConfig {
7556 grid_dim: (out_f as u32, 1, 1),
7557 block_dim: (32, 1, 1),
7558 shared_mem_bytes: 0,
7559 };
7560 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
7561 let __s_b = self.gpu.stream();
7562 let mut b = __s_b.launch_builder(&f);
7563 b.arg(&dp)
7564 .arg(&w)
7565 .arg(aq2)
7566 .arg(ad2)
7567 .arg(dst)
7568 .arg(&inf)
7569 .arg(&outf)
7570 .arg(&nu)
7571 .arg(&qt)
7572 .arg(&rbi);
7573 unsafe {
7574 b.launch(cfg)?;
7575 }
7576 Ok(())
7577 }
7578
7579 #[allow(clippy::too_many_arguments)]
7584 pub fn moe_gate_up_preclamp8_q8_w4(
7585 &self,
7586 gp: WPtr8,
7587 up: WPtr8,
7588 aq: &CudaSlice<i8>,
7589 ad: &CudaSlice<f32>,
7590 gs: F32x8,
7591 us: F32x8,
7592 limit: f32,
7593 in_f: usize,
7594 n_ff: usize,
7595 n_used: usize,
7596 qt_g: i32,
7597 qt_u: i32,
7598 rb_g: usize,
7599 rb_u: usize,
7600 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7601 debug_assert!(
7602 limit > 1e-6,
7603 "moe_gate_up_preclamp8_q8_w4 needs a live limit; use moe_gate_up_silu8_q8"
7604 );
7605 const ROWS: u32 = 4; let f = self.func("moe_gate_up_preclamp8_q8_w4");
7607 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7608 let cfg = LaunchConfig {
7609 grid_dim: ((n_ff as u32).div_ceil(ROWS), n_used as u32, 1),
7610 block_dim: (32, ROWS, 1),
7611 shared_mem_bytes: 0,
7612 };
7613 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7614 let __s_b = self.gpu.stream();
7615 let mut b = __s_b.launch_builder(&f);
7616 b.arg(&gp)
7617 .arg(&up)
7618 .arg(aq)
7619 .arg(ad)
7620 .arg(&gs)
7621 .arg(&us)
7622 .arg(&limit)
7623 .arg(&mut act)
7624 .arg(&inf)
7625 .arg(&nff)
7626 .arg(&qt_g)
7627 .arg(&qt_u)
7628 .arg(&rbg)
7629 .arg(&rbu);
7630 unsafe {
7631 b.launch(cfg)?;
7632 }
7633 Ok(act)
7634 }
7635
7636 #[allow(clippy::too_many_arguments)]
7639 pub fn moe_down8_fma_q8_w4(
7640 &self,
7641 dp: WPtr8,
7642 w: F32x8,
7643 aq2: &CudaSlice<i8>,
7644 ad2: &CudaSlice<f32>,
7645 dst: &mut cudarc::driver::CudaViewMut<f32>,
7646 in_f: usize,
7647 out_f: usize,
7648 n_used: usize,
7649 qt: i32,
7650 rb: usize,
7651 ) -> Result<(), Box<dyn std::error::Error>> {
7652 const ROWS: u32 = 4; let f = self.func("moe_down8_fma_q8_w4");
7654 let cfg = LaunchConfig {
7655 grid_dim: ((out_f as u32).div_ceil(ROWS), 1, 1),
7656 block_dim: (32, ROWS, 1),
7657 shared_mem_bytes: 0,
7658 };
7659 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
7660 let __s_b = self.gpu.stream();
7661 let mut b = __s_b.launch_builder(&f);
7662 b.arg(&dp)
7663 .arg(&w)
7664 .arg(aq2)
7665 .arg(ad2)
7666 .arg(dst)
7667 .arg(&inf)
7668 .arg(&outf)
7669 .arg(&nu)
7670 .arg(&qt)
7671 .arg(&rbi);
7672 unsafe {
7673 b.launch(cfg)?;
7674 }
7675 Ok(())
7676 }
7677
7678 #[allow(clippy::too_many_arguments)]
7690 pub fn moe_vrows_tables_from_sel(
7692 &self,
7693 sel: &CudaSlice<i32>,
7694 selw: &CudaSlice<f32>,
7695 il: u16,
7696 macros: Option<(&[f32], &[f32], &[f32])>,
7697 (pg, pu, pd): (u64, u64, u64),
7698 (sg, su, sd): (usize, usize, usize),
7699 n_pairs: usize,
7700 ptrs: &mut CudaSlice<u64>,
7701 scl: &mut CudaSlice<f32>,
7702 ) -> Result<(), Box<dyn std::error::Error>> {
7703 debug_assert!(sel.len() >= n_pairs && selw.len() >= n_pairs);
7704 debug_assert!(ptrs.len() >= 3 * n_pairs);
7706 debug_assert_eq!(scl.len(), 3 * n_pairs);
7707 let mut mac = self
7710 .vrows_macro_dev
7711 .lock()
7712 .map_err(|_| "vrows macro mirror map is poisoned")?;
7713 if let Some((hg, hu, hd)) = macros {
7714 for (plane, host) in [(0u8, hg), (1u8, hu), (2u8, hd)] {
7715 if let std::collections::hash_map::Entry::Vacant(slot) = mac.entry((il, plane)) {
7718 slot.insert(self.htod(host)?);
7719 }
7720 }
7721 }
7722 let (mg, mu, md, have) = match macros {
7725 Some(_) => (
7726 mac.get(&(il, 0)).expect("gate macro mirror built above"),
7727 mac.get(&(il, 1)).expect("up macro mirror built above"),
7728 mac.get(&(il, 2)).expect("down macro mirror built above"),
7729 1i32,
7730 ),
7731 None => (selw, selw, selw, 0i32),
7732 };
7733 let f = self.func("moe_vrows_tables_from_sel");
7734 let threads = 128u32;
7735 let cfg = LaunchConfig {
7736 grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
7737 block_dim: (threads, 1, 1),
7738 shared_mem_bytes: 0,
7739 };
7740 let (sgi, sui, sdi) = (sg as i64, su as i64, sd as i64);
7741 let (npi, havei) = (n_pairs as i32, have);
7742 let __s_b = self.gpu.stream();
7743 let mut b = __s_b.launch_builder(&f);
7744 b.arg(sel)
7745 .arg(selw)
7746 .arg(mg)
7747 .arg(mu)
7748 .arg(md)
7749 .arg(&mut *ptrs)
7750 .arg(&mut *scl)
7751 .arg(&pg)
7752 .arg(&pu)
7753 .arg(&pd)
7754 .arg(&sgi)
7755 .arg(&sui)
7756 .arg(&sdi)
7757 .arg(&npi)
7758 .arg(&havei);
7759 unsafe {
7760 b.launch(cfg)?;
7761 }
7762 Ok(())
7763 }
7764
7765 pub fn moe_vrows_order_from_sel(
7777 &self,
7778 sel: &CudaSlice<i32>,
7779 n_pairs: usize,
7780 ptrs: &mut CudaSlice<u64>,
7781 ) -> Result<(), Box<dyn std::error::Error>> {
7782 debug_assert!(sel.len() >= n_pairs);
7783 debug_assert!(
7784 ptrs.len() >= 4 * n_pairs,
7785 "the order plane lives at ptrs[3*n_pairs .. 4*n_pairs)"
7786 );
7787 let f = self.func("moe_vrows_order_from_sel");
7788 let threads = 128u32;
7789 let cfg = LaunchConfig {
7790 grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
7791 block_dim: (threads, 1, 1),
7792 shared_mem_bytes: 0,
7793 };
7794 let np = n_pairs as i32;
7795 let __s_b = self.gpu.stream();
7796 let mut b = __s_b.launch_builder(&f);
7797 b.arg(sel).arg(&mut *ptrs).arg(&np);
7798 unsafe {
7799 b.launch(cfg)?;
7800 }
7801 Ok(())
7802 }
7803
7804 #[allow(clippy::too_many_arguments)]
7810 pub fn moe_gate_up_preclamp8_q8_rows(
7812 &self,
7813 ptrs: &CudaSlice<u64>,
7814 scl: &CudaSlice<f32>,
7815 aq: &CudaSlice<i8>,
7816 ad: &CudaSlice<f32>,
7817 limit: f32,
7818 in_f: usize,
7819 n_ff: usize,
7820 n_used: usize,
7821 n_pairs: usize,
7822 qt_g: i32,
7823 qt_u: i32,
7824 rb_g: usize,
7825 rb_u: usize,
7826 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7827 debug_assert!(
7828 limit > 1e-6,
7829 "moe_gate_up_preclamp8_q8_rows needs a live limit; the kernel collapses every gate \
7830 to silu(0) at limit 0"
7831 );
7832 debug_assert!(ptrs.len() >= 3 * n_pairs);
7833 debug_assert_eq!(scl.len(), 3 * n_pairs);
7834 let packed = moe_vrows_pack_on();
7842 let ordered =
7843 !packed && moe_vrows_dedup_order_on() && ptrs.len() >= 4 * n_pairs && n_ff <= 65535;
7844 let ilp = moe_vrows_ilp_on()
7847 && if (qt_g == QT_NVFP4 && qt_u == QT_NVFP4)
7848 || (qt_g == QT_NVFP4_V2 && qt_u == QT_NVFP4_V2)
7849 {
7850 true
7851 } else {
7852 moe_vrows_ilp_refuse("gate/up", if qt_g == QT_NVFP4 { qt_u } else { qt_g });
7853 false
7854 };
7855 if ilp && MOE_VROWS_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
7856 eprintln!(
7857 "[moe-vrows-ilp] engaged: verify-rows MoE pair with four groups' loads per lane \
7858 hoisted ahead of their math (MEMRA_MOE_VROWS_ILP=1, packed={packed})"
7859 );
7860 }
7861 let ilp2 = ilp && crate::moe_gateup_ilp2_on();
7863 if ilp2
7864 && MOE_GATEUP_ILP2_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0
7865 {
7866 eprintln!(
7867 "[moe-gateup-ilp2] engaged: verify-rows gate/up gives a warp two pairs at one \
7868 row (16 groups in flight per lane, shared activation at t=1; same per-pair \
7869 order; MEMRA_MOE_GATEUP_ILP2=1)"
7870 );
7871 }
7872 let (f, cfg) = if packed {
7873 if MOE_VROWS_PACK_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
7874 eprintln!(
7875 "[moe-vrows-pack] engaged: 4-warp blocks on the verify-rows MoE pair \
7876 (MEMRA_MOE_VROWS_PACK=1)"
7877 );
7878 }
7879 (
7880 self.func(if ilp2 {
7881 "moe_gate_up_preclamp8_q8_rows_w4_ilp2"
7882 } else if ilp {
7883 "moe_gate_up_preclamp8_q8_rows_w4_ilp"
7884 } else {
7885 "moe_gate_up_preclamp8_q8_rows_w4"
7886 }),
7887 LaunchConfig {
7888 grid_dim: (
7889 (n_ff as u32).div_ceil(4),
7890 if ilp2 {
7891 (n_pairs as u32).div_ceil(2)
7892 } else {
7893 n_pairs as u32
7894 },
7895 1,
7896 ),
7897 block_dim: (32, 4, 1),
7898 shared_mem_bytes: 0,
7899 },
7900 )
7901 } else if ordered {
7902 if MOE_VROWS_DEDUP_ORDER_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
7903 == 0
7904 {
7905 eprintln!(
7906 "[moe-vrows-dedup-order] engaged: verify-rows gate/up walks the pair union \
7907 EXPERT-MAJOR with the pair index as the fastest grid dimension, so the \
7908 21.96%-measured repeat visits read a shared expert slab's rows in adjacent \
7909 blocks (MEMRA_MOE_VROWS_DEDUP_ORDER=1)"
7910 );
7911 }
7912 (
7913 self.func("moe_gate_up_preclamp8_q8_rows_ord"),
7914 LaunchConfig {
7915 grid_dim: (n_pairs as u32, n_ff as u32, 1),
7916 block_dim: (32, 1, 1),
7917 shared_mem_bytes: 0,
7918 },
7919 )
7920 } else {
7921 (
7922 self.func(if ilp2 {
7923 "moe_gate_up_preclamp8_q8_rows_ilp2"
7924 } else if ilp {
7925 "moe_gate_up_preclamp8_q8_rows_ilp"
7926 } else {
7927 "moe_gate_up_preclamp8_q8_rows"
7928 }),
7929 LaunchConfig {
7930 grid_dim: (
7931 n_ff as u32,
7932 if ilp2 {
7933 (n_pairs as u32).div_ceil(2)
7934 } else {
7935 n_pairs as u32
7936 },
7937 1,
7938 ),
7939 block_dim: (32, 1, 1),
7940 shared_mem_bytes: 0,
7941 },
7942 )
7943 };
7944 let mut act = self.vws_uninit(n_pairs * n_ff)?;
7946 let (inf, nff, nu, np) = (in_f as i32, n_ff as i32, n_used as i32, n_pairs as i32);
7947 let (rbg, rbu) = (rb_g as i64, rb_u as i64);
7948 let __s_b = self.gpu.stream();
7949 let mut b = __s_b.launch_builder(&f);
7950 b.arg(ptrs)
7951 .arg(scl)
7952 .arg(aq)
7953 .arg(ad)
7954 .arg(&limit)
7955 .arg(&mut act)
7956 .arg(&inf)
7957 .arg(&nff)
7958 .arg(&nu)
7959 .arg(&np)
7960 .arg(&qt_g)
7961 .arg(&qt_u)
7962 .arg(&rbg)
7963 .arg(&rbu);
7964 unsafe {
7965 b.launch(cfg)?;
7966 }
7967 Ok(act)
7968 }
7969
7970 #[allow(clippy::too_many_arguments)]
7974 pub fn moe_down8_fma_q8_rows(
7976 &self,
7977 ptrs: &CudaSlice<u64>,
7978 scl: &CudaSlice<f32>,
7979 aq2: &CudaSlice<i8>,
7980 ad2: &CudaSlice<f32>,
7981 dst: &mut CudaSlice<f32>,
7982 in_f: usize,
7983 out_f: usize,
7984 n_used: usize,
7985 n_pairs: usize,
7986 qt: i32,
7987 rb: usize,
7988 ) -> Result<(), Box<dyn std::error::Error>> {
7989 debug_assert!(ptrs.len() >= 3 * n_pairs);
7990 debug_assert_eq!(scl.len(), 3 * n_pairs);
7991 debug_assert_eq!(n_pairs % n_used, 0, "pairs are dense slot-major");
7992 let t = n_pairs / n_used;
7993 debug_assert!(dst.len() >= t * out_f);
7994 let packed = moe_vrows_pack_on();
7996 let tmaj = !packed && moe_vrows_down_tmaj_on() && out_f <= 65535;
8002 let ilp = moe_vrows_ilp_on()
8004 && if qt == QT_NVFP4 || qt == QT_NVFP4_V2 {
8005 true
8006 } else {
8007 moe_vrows_ilp_refuse("down", qt);
8008 false
8009 };
8010 if ilp {
8011 MOE_VROWS_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8012 }
8013 let ilp2 = ilp && crate::moe_down_ilp2_on();
8015 if ilp2 && MOE_DOWN_ILP2_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0
8016 {
8017 eprintln!(
8018 "[moe-down-ilp2] engaged: verify-rows down/FMA walks two experts per warp \
8019 (8 groups in flight per lane, same per-expert order and slot chain; \
8020 MEMRA_MOE_DOWN_ILP2=1)"
8021 );
8022 }
8023 let (f, cfg) = if packed {
8024 (
8025 self.func(if ilp2 {
8026 "moe_down8_fma_q8_rows_w4_ilp2"
8027 } else if ilp {
8028 "moe_down8_fma_q8_rows_w4_ilp"
8029 } else {
8030 "moe_down8_fma_q8_rows_w4"
8031 }),
8032 LaunchConfig {
8033 grid_dim: ((out_f as u32).div_ceil(4), t as u32, 1),
8034 block_dim: (32, 4, 1),
8035 shared_mem_bytes: 0,
8036 },
8037 )
8038 } else if tmaj {
8039 if MOE_VROWS_DOWN_TMAJ_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8040 == 0
8041 {
8042 eprintln!(
8043 "[moe-vrows-down-tmaj] engaged: verify-rows down/FMA grid transposed to \
8044 (t, out_f) so the verify rows at one output row are adjacent blocks; the \
8045 slot-ordered FMA chain is unchanged (MEMRA_MOE_VROWS_DOWN_TMAJ=1)"
8046 );
8047 }
8048 (
8049 self.func("moe_down8_fma_q8_rows_tmaj"),
8050 LaunchConfig {
8051 grid_dim: (t as u32, out_f as u32, 1),
8052 block_dim: (32, 1, 1),
8053 shared_mem_bytes: 0,
8054 },
8055 )
8056 } else {
8057 (
8058 self.func(if ilp2 {
8059 "moe_down8_fma_q8_rows_ilp2"
8060 } else if ilp {
8061 "moe_down8_fma_q8_rows_ilp"
8062 } else {
8063 "moe_down8_fma_q8_rows"
8064 }),
8065 LaunchConfig {
8066 grid_dim: (out_f as u32, t as u32, 1),
8067 block_dim: (32, 1, 1),
8068 shared_mem_bytes: 0,
8069 },
8070 )
8071 };
8072 let (inf, outf, nu, np, rbi) = (
8073 in_f as i32,
8074 out_f as i32,
8075 n_used as i32,
8076 n_pairs as i32,
8077 rb as i64,
8078 );
8079 let __s_b = self.gpu.stream();
8080 let mut b = __s_b.launch_builder(&f);
8081 b.arg(ptrs)
8082 .arg(scl)
8083 .arg(aq2)
8084 .arg(ad2)
8085 .arg(dst)
8086 .arg(&inf)
8087 .arg(&outf)
8088 .arg(&nu)
8089 .arg(&np)
8090 .arg(&qt)
8091 .arg(&rbi);
8092 unsafe {
8093 b.launch(cfg)?;
8094 }
8095 Ok(())
8096 }
8097
8098 #[allow(clippy::too_many_arguments)]
8100 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_expert_q8(
8103 &self,
8104 w: &CudaSlice<u8>,
8105 range: std::ops::Range<usize>,
8106 aq: &CudaSlice<i8>,
8107 ad: &CudaSlice<f32>,
8108 m: usize,
8109 in_f: usize,
8110 out_f: usize,
8111 qtype: i32,
8112 row_bytes: usize,
8113 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8114 let f = self.func("qmatvec_expert_q8");
8115 let wv = w.slice(range);
8116 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
8117 const ROWS: u32 = 4; let cfg = LaunchConfig {
8119 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
8120 block_dim: (32, ROWS, 1),
8121 shared_mem_bytes: 0,
8122 };
8123 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8124 let __s_b = self.gpu.stream();
8125 let mut b = __s_b.launch_builder(&f);
8126 b.arg(&wv)
8127 .arg(aq)
8128 .arg(ad)
8129 .arg(&mut y)
8130 .arg(&inf)
8131 .arg(&outf)
8132 .arg(&mi)
8133 .arg(&qtype)
8134 .arg(&rbi);
8135 unsafe {
8136 b.launch(cfg)?;
8137 }
8138 Ok(y)
8139 }
8140
8141 #[allow(clippy::too_many_arguments)] pub fn moe_gate_up_silu8(
8143 &self,
8144 gp: WPtr8,
8145 up: WPtr8,
8146 x: &cudarc::driver::CudaView<f32>,
8147 in_f: usize,
8148 n_ff: usize,
8149 n_used: usize,
8150 qt_g: i32,
8151 qt_u: i32,
8152 rb_g: usize,
8153 rb_u: usize,
8154 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8155 let f = self.func("moe_gate_up_silu8_f32");
8156 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
8158 grid_dim: (n_ff as u32, n_used as u32, 1),
8159 block_dim: (256, 1, 1),
8160 shared_mem_bytes: 0,
8161 };
8162 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
8163 let __s_b = self.gpu.stream();
8164 let mut b = __s_b.launch_builder(&f);
8165 b.arg(&gp)
8166 .arg(&up)
8167 .arg(x)
8168 .arg(&mut act)
8169 .arg(&inf)
8170 .arg(&nff)
8171 .arg(&qt_g)
8172 .arg(&qt_u)
8173 .arg(&rbg)
8174 .arg(&rbu);
8175 unsafe {
8176 b.launch(cfg)?;
8177 }
8178 Ok(act)
8179 }
8180
8181 #[allow(clippy::too_many_arguments)]
8187 pub fn moe_down8_fma_into(
8188 &self,
8189 dp: WPtr8,
8190 w: F32x8,
8191 act: &CudaSlice<f32>,
8192 dst: &mut cudarc::driver::CudaViewMut<f32>,
8193 in_f: usize,
8194 out_f: usize,
8195 n_used: usize,
8196 qt: i32,
8197 rb: usize,
8198 ) -> Result<(), Box<dyn std::error::Error>> {
8199 let f = self.func("moe_down8_fma_f32");
8200 let cfg = LaunchConfig {
8201 grid_dim: (out_f as u32, 1, 1),
8202 block_dim: (256, 1, 1),
8203 shared_mem_bytes: 0,
8204 };
8205 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
8206 let __s_b = self.gpu.stream();
8207 let mut b = __s_b.launch_builder(&f);
8208 b.arg(&dp)
8209 .arg(&w)
8210 .arg(act)
8211 .arg(dst)
8212 .arg(&inf)
8213 .arg(&outf)
8214 .arg(&nu)
8215 .arg(&qt)
8216 .arg(&rbv);
8217 unsafe {
8218 b.launch(cfg)?;
8219 }
8220 Ok(())
8221 }
8222
8223 #[allow(clippy::too_many_arguments)]
8228 #[allow(clippy::too_many_arguments)]
8243 #[allow(clippy::too_many_arguments)]
8245 #[allow(clippy::manual_div_ceil)] pub fn moe_pairs_matvec_q8(
8247 &self,
8248 table: &CudaSlice<u64>,
8249 proj: i32,
8250 pair_tok: &CudaSlice<i32>,
8251 pair_ex: &CudaSlice<i32>,
8252 aq: &CudaSlice<i8>,
8253 ad: &CudaSlice<f32>,
8254 in_f: usize,
8255 out_f: usize,
8256 n_expert: usize,
8257 n_pairs: usize,
8258 qtype: i32,
8259 row_bytes: usize,
8260 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8261 let f = self.func("moe_pairs_matvec_q8");
8262 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8263 const ROWS: u32 = 4;
8264 let cfg = LaunchConfig {
8265 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
8266 block_dim: (32, ROWS, 1),
8267 shared_mem_bytes: 0,
8268 };
8269 let (inf, outf, ne, np, rbi) = (
8270 in_f as i32,
8271 out_f as i32,
8272 n_expert as i32,
8273 n_pairs as i32,
8274 row_bytes as i64,
8275 );
8276 let __s_b = self.gpu.stream();
8277 let mut b = __s_b.launch_builder(&f);
8278 b.arg(table)
8279 .arg(&proj)
8280 .arg(pair_tok)
8281 .arg(pair_ex)
8282 .arg(aq)
8283 .arg(ad)
8284 .arg(&mut y)
8285 .arg(&inf)
8286 .arg(&outf)
8287 .arg(&ne)
8288 .arg(&np)
8289 .arg(&qtype)
8290 .arg(&rbi);
8291 unsafe {
8292 b.launch(cfg)?;
8293 }
8294 Ok(y)
8295 }
8296
8297 #[allow(clippy::too_many_arguments)]
8299 #[allow(clippy::manual_div_ceil)] pub fn moe_pairs_matvec_q8_em(
8301 &self,
8302 table: &CudaSlice<u64>,
8303 proj: i32,
8304 ex_ids: &CudaSlice<i32>,
8305 ex_off: &CudaSlice<i32>,
8306 ex_pairs: &CudaSlice<i32>,
8307 pair_tok: &CudaSlice<i32>,
8308 aq: &CudaSlice<i8>,
8309 ad: &CudaSlice<f32>,
8310 in_f: usize,
8311 out_f: usize,
8312 n_expert: usize,
8313 n_active: usize,
8314 n_pairs: usize,
8315 qtype: i32,
8316 row_bytes: usize,
8317 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8318 let f = self.func("moe_pairs_matvec_q8_em");
8319 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8320 const ROWS: u32 = 4;
8321 let cfg = LaunchConfig {
8322 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
8323 block_dim: (32, ROWS, 1),
8324 shared_mem_bytes: 0,
8325 };
8326 let (inf, outf, ne, na, rbi) = (
8327 in_f as i32,
8328 out_f as i32,
8329 n_expert as i32,
8330 n_active as i32,
8331 row_bytes as i64,
8332 );
8333 let __s_b = self.gpu.stream();
8334 let mut b = __s_b.launch_builder(&f);
8335 b.arg(table)
8336 .arg(&proj)
8337 .arg(ex_ids)
8338 .arg(ex_off)
8339 .arg(ex_pairs)
8340 .arg(pair_tok)
8341 .arg(aq)
8342 .arg(ad)
8343 .arg(&mut y)
8344 .arg(&inf)
8345 .arg(&outf)
8346 .arg(&ne)
8347 .arg(&na)
8348 .arg(&qtype)
8349 .arg(&rbi);
8350 unsafe {
8351 b.launch(cfg)?;
8352 }
8353 Ok(y)
8354 }
8355
8356 #[allow(clippy::too_many_arguments)]
8359 #[allow(clippy::manual_div_ceil)] pub fn moe_pairs_matvec_q8_dec(
8361 &self,
8362 table: &CudaSlice<u64>,
8363 proj: i32,
8364 ex_ids: &CudaSlice<i32>,
8365 ex_off: &CudaSlice<i32>,
8366 ex_pairs: &CudaSlice<i32>,
8367 pair_tok: &CudaSlice<i32>,
8368 aq: &CudaSlice<i8>,
8369 ad: &CudaSlice<f32>,
8370 in_f: usize,
8371 out_f: usize,
8372 n_expert: usize,
8373 n_active: usize,
8374 n_pairs: usize,
8375 qtype: i32,
8376 row_bytes: usize,
8377 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8378 let f = self.func("moe_pairs_matvec_q8_dec");
8379 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8380 const ROWS: u32 = 4;
8381 let cfg = LaunchConfig {
8382 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
8383 block_dim: (32, ROWS, 1),
8384 shared_mem_bytes: 0,
8385 };
8386 let (inf, outf, ne, na, rbi) = (
8387 in_f as i32,
8388 out_f as i32,
8389 n_expert as i32,
8390 n_active as i32,
8391 row_bytes as i64,
8392 );
8393 let __s_b = self.gpu.stream();
8394 let mut b = __s_b.launch_builder(&f);
8395 b.arg(table)
8396 .arg(&proj)
8397 .arg(ex_ids)
8398 .arg(ex_off)
8399 .arg(ex_pairs)
8400 .arg(pair_tok)
8401 .arg(aq)
8402 .arg(ad)
8403 .arg(&mut y)
8404 .arg(&inf)
8405 .arg(&outf)
8406 .arg(&ne)
8407 .arg(&na)
8408 .arg(&qtype)
8409 .arg(&rbi);
8410 unsafe {
8411 b.launch(cfg)?;
8412 }
8413 Ok(y)
8414 }
8415
8416 pub fn moe_pairs_gelu_mul(
8417 &self,
8418 gate: &CudaSlice<f32>,
8419 up: &CudaSlice<f32>,
8420 n: usize,
8421 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8422 let f = self.func("moe_pairs_gelu_mul");
8423 let mut act = self.alloc_uninit::<f32>(n)?;
8424 let cfg = LaunchConfig::for_num_elems(n as u32);
8425 let nl = n as i64;
8426 let __s_b = self.gpu.stream();
8427 let mut b = __s_b.launch_builder(&f);
8428 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
8429 unsafe {
8430 b.launch(cfg)?;
8431 }
8432 Ok(act)
8433 }
8434
8435 pub fn moe_pairs_silu_mul(
8436 &self,
8437 gate: &CudaSlice<f32>,
8438 up: &CudaSlice<f32>,
8439 n: usize,
8440 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8441 let f = self.func("moe_pairs_silu_mul");
8442 let mut act = self.alloc_uninit::<f32>(n)?;
8443 let cfg = LaunchConfig::for_num_elems(n as u32);
8444 let nl = n as i64;
8445 let __s_b = self.gpu.stream();
8446 let mut b = __s_b.launch_builder(&f);
8447 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
8448 unsafe {
8449 b.launch(cfg)?;
8450 }
8451 Ok(act)
8452 }
8453
8454 #[allow(clippy::too_many_arguments)]
8455 #[allow(clippy::manual_div_ceil)] pub fn moe_pairs_scatter(
8457 &self,
8458 y_down: &CudaSlice<f32>,
8459 pair_w: &CudaSlice<f32>,
8460 tok_pair_off: &CudaSlice<i32>,
8461 tok_pair_ids: &CudaSlice<i32>,
8462 moe_out: &mut CudaSlice<f32>,
8463 t: usize,
8464 n_embd: usize,
8465 ) -> Result<(), Box<dyn std::error::Error>> {
8466 let f = self.func("moe_pairs_scatter");
8467 let cfg = LaunchConfig {
8468 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
8469 block_dim: (256, 1, 1),
8470 shared_mem_bytes: 0,
8471 };
8472 let ne = n_embd as i32;
8473 let __s_b = self.gpu.stream();
8474 let mut b = __s_b.launch_builder(&f);
8475 b.arg(y_down)
8476 .arg(pair_w)
8477 .arg(tok_pair_off)
8478 .arg(tok_pair_ids)
8479 .arg(moe_out)
8480 .arg(&ne);
8481 unsafe {
8482 b.launch(cfg)?;
8483 }
8484 Ok(())
8485 }
8486
8487 #[allow(clippy::too_many_arguments)]
8491 pub fn moe_gate_up_gelu8_dev_q8(
8492 &self,
8493 table: &CudaSlice<u64>,
8494 sel: &cudarc::driver::CudaView<i32>,
8495 aq: &CudaSlice<i8>,
8496 ad: &CudaSlice<f32>,
8497 in_f: usize,
8498 n_ff: usize,
8499 n_used: usize,
8500 n_expert: usize,
8501 qt_g: i32,
8502 qt_u: i32,
8503 rb_g: usize,
8504 rb_u: usize,
8505 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8506 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
8507 let (inf, nff, ne, rbg, rbu) = (
8508 in_f as i32,
8509 n_ff as i32,
8510 n_expert as i32,
8511 rb_g as i64,
8512 rb_u as i64,
8513 );
8514 let f = self.func("moe_gate_up_gelu8_dev_q8");
8515 let cfg = LaunchConfig {
8516 grid_dim: (n_ff as u32, n_used as u32, 1),
8517 block_dim: (32, 1, 1),
8518 shared_mem_bytes: 0,
8519 };
8520 let __s_b = self.gpu.stream();
8521 let mut b = __s_b.launch_builder(&f);
8522 b.arg(table)
8523 .arg(sel)
8524 .arg(aq)
8525 .arg(ad)
8526 .arg(&mut act)
8527 .arg(&inf)
8528 .arg(&nff)
8529 .arg(&ne)
8530 .arg(&qt_g)
8531 .arg(&qt_u)
8532 .arg(&rbg)
8533 .arg(&rbu);
8534 unsafe {
8535 b.launch(cfg)?;
8536 }
8537 Ok(act)
8538 }
8539
8540 #[allow(clippy::too_many_arguments)]
8542 pub fn moe_gate_up_gelu8_dev_q8_rows(
8543 &self,
8544 table: &CudaSlice<u64>,
8545 sel: &CudaSlice<i32>,
8546 aq: &CudaSlice<i8>,
8547 ad: &CudaSlice<f32>,
8548 t: usize,
8549 in_f: usize,
8550 n_ff: usize,
8551 n_used: usize,
8552 n_expert: usize,
8553 qt_g: i32,
8554 qt_u: i32,
8555 rb_g: usize,
8556 rb_u: usize,
8557 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8558 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
8559 let (inf, nff, ne, rbg, rbu, nu) = (
8560 in_f as i32,
8561 n_ff as i32,
8562 n_expert as i32,
8563 rb_g as i64,
8564 rb_u as i64,
8565 n_used as i32,
8566 );
8567 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
8568 let cfg = LaunchConfig {
8569 grid_dim: (n_ff as u32, n_used as u32, t as u32),
8570 block_dim: (32, 1, 1),
8571 shared_mem_bytes: 0,
8572 };
8573 let __s_b = self.gpu.stream();
8574 let mut b = __s_b.launch_builder(&f);
8575 b.arg(table)
8576 .arg(sel)
8577 .arg(aq)
8578 .arg(ad)
8579 .arg(&mut act)
8580 .arg(&inf)
8581 .arg(&nff)
8582 .arg(&ne)
8583 .arg(&qt_g)
8584 .arg(&qt_u)
8585 .arg(&rbg)
8586 .arg(&rbu)
8587 .arg(&nu);
8588 unsafe {
8589 b.launch(cfg)?;
8590 }
8591 Ok(act)
8592 }
8593
8594 #[allow(clippy::too_many_arguments)]
8596 pub fn moe_gate_up_gelu8_dev_q8_csr(
8597 &self,
8598 table: &CudaSlice<u64>,
8599 sel: &CudaSlice<i32>,
8600 aq: &CudaSlice<i8>,
8601 ad: &CudaSlice<f32>,
8602 n_pairs: usize,
8603 in_f: usize,
8604 n_ff: usize,
8605 n_used: usize,
8606 n_expert: usize,
8607 qt_g: i32,
8608 qt_u: i32,
8609 rb_g: usize,
8610 rb_u: usize,
8611 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8612 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
8613 let (inf, nff, ne, rbg, rbu, nu, npi) = (
8614 in_f as i32,
8615 n_ff as i32,
8616 n_expert as i32,
8617 rb_g as i64,
8618 rb_u as i64,
8619 n_used as i32,
8620 n_pairs as i32,
8621 );
8622 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
8623 let cfg = LaunchConfig {
8624 grid_dim: (n_ff as u32, n_pairs as u32, 1),
8625 block_dim: (32, 1, 1),
8626 shared_mem_bytes: 0,
8627 };
8628 let __s_b = self.gpu.stream();
8629 let mut b = __s_b.launch_builder(&f);
8630 b.arg(table)
8631 .arg(sel)
8632 .arg(aq)
8633 .arg(ad)
8634 .arg(&mut act)
8635 .arg(&inf)
8636 .arg(&nff)
8637 .arg(&ne)
8638 .arg(&qt_g)
8639 .arg(&qt_u)
8640 .arg(&rbg)
8641 .arg(&rbu)
8642 .arg(&nu)
8643 .arg(&npi);
8644 unsafe {
8645 b.launch(cfg)?;
8646 }
8647 Ok(act)
8648 }
8649
8650 #[allow(clippy::too_many_arguments)]
8652 pub fn moe_down8_fma_dev_q8_rows_g(
8653 &self,
8654 table: &CudaSlice<u64>,
8655 sel: &CudaSlice<i32>,
8656 w: &CudaSlice<f32>,
8657 aq2: &CudaSlice<i8>,
8658 ad2: &CudaSlice<f32>,
8659 dst: &mut CudaSlice<f32>,
8660 t: usize,
8661 in_f: usize,
8662 out_f: usize,
8663 n_used: usize,
8664 n_expert: usize,
8665 qt: i32,
8666 rb: usize,
8667 ) -> Result<(), Box<dyn std::error::Error>> {
8668 let (inf, outf, nu, ne, rbi) = (
8669 in_f as i32,
8670 out_f as i32,
8671 n_used as i32,
8672 n_expert as i32,
8673 rb as i64,
8674 );
8675 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
8679 let f = self.func(if step_b1_w8 {
8680 "moe_down8_fma_dev_q8_rows_w8"
8681 } else {
8682 "moe_down8_fma_dev_q8_rows_g"
8683 });
8684 let cfg = LaunchConfig {
8685 grid_dim: (out_f as u32, 1, t as u32),
8686 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
8687 shared_mem_bytes: 0,
8688 };
8689 let __s_b = self.gpu.stream();
8690 let mut b = __s_b.launch_builder(&f);
8691 b.arg(table)
8692 .arg(sel)
8693 .arg(w)
8694 .arg(aq2)
8695 .arg(ad2)
8696 .arg(dst)
8697 .arg(&inf)
8698 .arg(&outf)
8699 .arg(&nu)
8700 .arg(&ne)
8701 .arg(&qt)
8702 .arg(&rbi);
8703 unsafe {
8704 b.launch(cfg)?;
8705 }
8706 Ok(())
8707 }
8708
8709 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
8713 let (out_f, in_f) = (2048usize, 2816usize);
8714 let nblk = in_f / 32;
8715 let mut seed = 0x9E3779B97F4A7C15u64;
8716 let mut rng = move || {
8717 seed = seed
8718 .wrapping_mul(6364136223846793005)
8719 .wrapping_add(1442695040888963407);
8720 (seed >> 33) as u8
8721 };
8722 let mut w = vec![0u8; out_f * nblk * 18];
8723 for b in w.iter_mut() {
8724 *b = rng();
8725 }
8726 for r in 0..out_f {
8727 for g in 0..nblk {
8728 let off = (r * nblk + g) * 18;
8729 w[off] = 0x00;
8730 w[off + 1] = 0x2C; }
8732 }
8733 let qplane = out_f * nblk * 16;
8734 let mut wrp = vec![0u8; w.len()];
8735 for r in 0..out_f {
8736 for g in 0..nblk {
8737 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
8738 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
8739 .copy_from_slice(&src[0..2]);
8740 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
8741 }
8742 }
8743 let w_d = self.htod_bytes(&w)?;
8744 let wrp_d = self.htod_bytes(&wrp)?;
8745 let mut aq = vec![0i8; m * in_f];
8746 for v in aq.iter_mut() {
8747 *v = rng() as i8;
8748 }
8749 let aq_d = self.htod_i8(&aq)?;
8750 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
8751 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
8752 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
8753 const RPB: u32 = 4;
8754 let cfg = LaunchConfig {
8755 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
8756 block_dim: (32, RPB, 1),
8757 shared_mem_bytes: 0,
8758 };
8759 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
8760 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
8761 let fb = self.func("qmatvec_q4_0_mmvq_b4");
8762 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
8763 {
8764 let __s_b = self.gpu.stream();
8765 let mut b = __s_b.launch_builder(&fb);
8766 b.arg(&w_d)
8767 .arg(&aq_d)
8768 .arg(&ad_d)
8769 .arg(&mut y0)
8770 .arg(&inf)
8771 .arg(&outf)
8772 .arg(&mi)
8773 .arg(&rb);
8774 unsafe {
8775 b.launch(cfg)?;
8776 }
8777 let __s_b = self.gpu.stream();
8778 let mut b = __s_b.launch_builder(&fr);
8779 b.arg(&wrp_d)
8780 .arg(&aq_d)
8781 .arg(&ad_d)
8782 .arg(&mut y1)
8783 .arg(&inf)
8784 .arg(&outf)
8785 .arg(&mi)
8786 .arg(&qp);
8787 unsafe {
8788 b.launch(cfg)?;
8789 }
8790 }
8791 self.gpu.stream().synchronize()?;
8792 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
8793 let nd = h0
8794 .iter()
8795 .zip(&h1)
8796 .filter(|(a, b)| a.to_bits() != b.to_bits())
8797 .count();
8798 if nd != 0 {
8799 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
8800 }
8801 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
8802 self.gpu.stream().synchronize()?;
8803 let t0 = std::time::Instant::now();
8804 for _ in 0..500 {
8805 if rp {
8806 let __s_b = self.gpu.stream();
8807 let mut b = __s_b.launch_builder(&fr);
8808 b.arg(&wrp_d)
8809 .arg(&aq_d)
8810 .arg(&ad_d)
8811 .arg(&mut y1)
8812 .arg(&inf)
8813 .arg(&outf)
8814 .arg(&mi)
8815 .arg(&qp);
8816 unsafe {
8817 b.launch(cfg)?;
8818 }
8819 } else {
8820 let __s_b = self.gpu.stream();
8821 let mut b = __s_b.launch_builder(&fb);
8822 b.arg(&w_d)
8823 .arg(&aq_d)
8824 .arg(&ad_d)
8825 .arg(&mut y0)
8826 .arg(&inf)
8827 .arg(&outf)
8828 .arg(&mi)
8829 .arg(&rb);
8830 unsafe {
8831 b.launch(cfg)?;
8832 }
8833 }
8834 }
8835 self.gpu.stream().synchronize()?;
8836 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
8837 };
8838 let _ = time(false)?;
8839 let _ = time(true)?; Ok((time(false)?, time(true)?))
8841 }
8842
8843 pub fn build_q4_rp4(
8848 &self,
8849 t: &mut crate::model::GpuTensor,
8850 ) -> Result<(), Box<dyn std::error::Error>> {
8851 use crate::model::GpuTensor;
8852 let GpuTensor::Quant {
8853 bytes,
8854 qtype,
8855 row_bytes,
8856 ne,
8857 rp4,
8858 ..
8859 } = t
8860 else {
8861 return Ok(());
8862 };
8863 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
8864 return Ok(());
8865 }
8866 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
8867 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
8868 return Ok(());
8869 }
8870 let nblk = in_f / 32;
8871 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
8872 let f = self.func("q4_0_split_rp_build");
8873 let n = (out_f * nblk) as i32;
8874 let cfg = LaunchConfig {
8875 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
8876 block_dim: (256, 1, 1),
8877 shared_mem_bytes: 0,
8878 };
8879 let (of, nb) = (out_f as i32, nblk as i32);
8880 let _ = n;
8881 let __s_b = self.gpu.stream();
8882 let mut b = __s_b.launch_builder(&f);
8883 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
8884 unsafe {
8885 b.launch(cfg)?;
8886 }
8887 *rp4 = Some(dst);
8888 Ok(())
8889 }
8890
8891 pub fn build_q8_rp4(
8896 &self,
8897 t: &mut crate::model::GpuTensor,
8898 ) -> Result<(), Box<dyn std::error::Error>> {
8899 use crate::model::GpuTensor;
8900 let GpuTensor::Quant {
8901 bytes,
8902 qtype,
8903 row_bytes,
8904 ne,
8905 rp4,
8906 ..
8907 } = t
8908 else {
8909 return Ok(());
8910 };
8911 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
8912 return Ok(());
8913 }
8914 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
8915 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
8916 return Ok(());
8917 }
8918 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
8919 Ok(())
8920 }
8921
8922 pub fn build_q8_rp4_raw(
8925 &self,
8926 bytes: &CudaSlice<u8>,
8927 in_f: usize,
8928 out_f: usize,
8929 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8930 assert!(in_f.is_multiple_of(32));
8931 let nblk = in_f / 32;
8932 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
8933 let f = self.func("q8_0_split_rp_build");
8934 let cfg = LaunchConfig {
8935 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
8936 block_dim: (256, 1, 1),
8937 shared_mem_bytes: 0,
8938 };
8939 let (of, nb) = (out_f as i32, nblk as i32);
8940 let __s_b = self.gpu.stream();
8941 let mut b = __s_b.launch_builder(&f);
8942 b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
8943 unsafe {
8944 b.launch(cfg)?;
8945 }
8946 Ok(dst)
8947 }
8948
8949 pub fn build_q4k_rp4(
8957 &self,
8958 t: &mut crate::model::GpuTensor,
8959 ) -> Result<(), Box<dyn std::error::Error>> {
8960 use crate::model::GpuTensor;
8961 let GpuTensor::Quant {
8962 bytes,
8963 qtype,
8964 row_bytes,
8965 ne,
8966 rp4,
8967 ..
8968 } = t
8969 else {
8970 return Ok(());
8971 };
8972 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
8973 return Ok(());
8974 }
8975 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
8976 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
8977 return Ok(());
8978 }
8979 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
8980 Ok(())
8981 }
8982
8983 pub fn build_q6k_rp4(
8984 &self,
8985 t: &mut crate::model::GpuTensor,
8986 ) -> Result<(), Box<dyn std::error::Error>> {
8987 use crate::model::GpuTensor;
8988 let GpuTensor::Quant {
8989 bytes,
8990 qtype,
8991 row_bytes,
8992 ne,
8993 rp4,
8994 ..
8995 } = t
8996 else {
8997 return Ok(());
8998 };
8999 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
9000 return Ok(());
9001 }
9002 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9003 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
9004 return Ok(());
9005 }
9006 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
9007 Ok(())
9008 }
9009
9010 pub fn build_kq_rp4_raw(
9012 &self,
9013 bytes: &CudaSlice<u8>,
9014 in_f: usize,
9015 out_f: usize,
9016 qtype: i32,
9017 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9018 assert!(in_f.is_multiple_of(256));
9019 let nsbk = in_f / 256;
9020 let (sb_bytes, kname) = match qtype {
9021 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
9022 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
9023 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
9024 };
9025 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
9026 let f = self.func(kname);
9027 let cfg = LaunchConfig {
9028 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
9029 block_dim: (256, 1, 1),
9030 shared_mem_bytes: 0,
9031 };
9032 let (of, nb) = (out_f as i32, nsbk as i32);
9033 let __s_b = self.gpu.stream();
9034 let mut b = __s_b.launch_builder(&f);
9035 b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
9036 unsafe {
9037 b.launch(cfg)?;
9038 }
9039 Ok(dst)
9040 }
9041
9042 pub fn kqrp_enabled() -> bool {
9046 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9047 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
9048 Ok("0") => false,
9049 Ok(_) => true,
9050 Err(_) => cfg!(memra_hopper_mma),
9051 })
9052 }
9053
9054 pub fn build_q4_rp_swap(
9060 &self,
9061 t: &mut crate::model::GpuTensor,
9062 ) -> Result<bool, Box<dyn std::error::Error>> {
9063 use crate::model::GpuTensor;
9064 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
9074 return Ok(false);
9075 }
9076 self.build_q4_rp4(t)?;
9077 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
9079 return Ok(false);
9080 };
9081 match rp4.take() {
9082 Some(split) => {
9083 *bytes = split; *rp = true;
9085 Ok(true)
9086 }
9087 None => Ok(false),
9088 }
9089 }
9090
9091 pub fn q4rp_enabled() -> bool {
9093 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9094 *ON.get_or_init(|| {
9095 std::env::var("MEMRA_Q4RP")
9096 .map(|v| v != "0")
9097 .unwrap_or(true)
9098 })
9099 }
9100
9101 #[allow(clippy::manual_div_ceil)] pub fn copy_rows_strided(
9105 &self,
9106 src: &CudaSlice<f32>,
9107 dst: &mut CudaSlice<f32>,
9108 row_elems: usize,
9109 n_rows: usize,
9110 src_stride: usize,
9111 src_off: usize,
9112 ) -> Result<(), Box<dyn std::error::Error>> {
9113 let f = self.func("copy_rows_strided_f32");
9114 let cfg = LaunchConfig {
9115 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
9116 block_dim: (256, 1, 1),
9117 shared_mem_bytes: 0,
9118 };
9119 let (re, nr) = (row_elems as i32, n_rows as i32);
9120 let (st, off) = (src_stride as i64, src_off as i64);
9121 let __s_b = self.gpu.stream();
9122 let mut b = __s_b.launch_builder(&f);
9123 b.arg(src)
9124 .arg(&mut *dst)
9125 .arg(&re)
9126 .arg(&nr)
9127 .arg(&st)
9128 .arg(&off);
9129 unsafe {
9130 b.launch(cfg)?;
9131 }
9132 Ok(())
9133 }
9134
9135 #[allow(clippy::manual_div_ceil)] pub fn place_rows_strided(
9142 &self,
9143 src: &CudaSlice<f32>,
9144 dst: &mut CudaSlice<f32>,
9145 row_elems: usize,
9146 n_rows: usize,
9147 dst_stride: usize,
9148 dst_off: usize,
9149 ) -> Result<(), Box<dyn std::error::Error>> {
9150 if row_elems == 0 || n_rows == 0 {
9151 return Err("strided row placement requires nonzero rows and row width".into());
9152 }
9153 let src_len = n_rows
9154 .checked_mul(row_elems)
9155 .ok_or("strided row placement source size overflow")?;
9156 let dst_len = n_rows
9157 .checked_sub(1)
9158 .and_then(|rows| rows.checked_mul(dst_stride))
9159 .and_then(|base| base.checked_add(dst_off))
9160 .and_then(|base| base.checked_add(row_elems))
9161 .ok_or("strided row placement destination size overflow")?;
9162 let row_end = dst_off
9163 .checked_add(row_elems)
9164 .ok_or("strided row placement row size overflow")?;
9165 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
9166 return Err(format!(
9167 "strided row placement geometry mismatch: src={} need_src={src_len} \
9168 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
9169 dst_stride={dst_stride} dst_off={dst_off}",
9170 src.len(),
9171 dst.len(),
9172 )
9173 .into());
9174 }
9175 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
9176 return Err("strided row placement exceeds CUDA kernel geometry".into());
9177 }
9178 let f = self.func("place_rows_strided_f32");
9179 let cfg = LaunchConfig {
9180 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
9181 block_dim: (256, 1, 1),
9182 shared_mem_bytes: 0,
9183 };
9184 let (re, nr) = (row_elems as i32, n_rows as i32);
9185 let (st, off) = (dst_stride as i64, dst_off as i64);
9186 let __s_b = self.gpu.stream();
9187 let mut b = __s_b.launch_builder(&f);
9188 b.arg(src)
9189 .arg(&mut *dst)
9190 .arg(&re)
9191 .arg(&nr)
9192 .arg(&st)
9193 .arg(&off);
9194 unsafe {
9195 b.launch(cfg)?;
9196 }
9197 Ok(())
9198 }
9199
9200 pub fn u32_set_k(
9202 &self,
9203 dst: &mut CudaSlice<u32>,
9204 v: u32,
9205 idx: usize,
9206 ) -> Result<(), Box<dyn std::error::Error>> {
9207 let f = self.func("u32_set_k");
9208 let cfg = LaunchConfig {
9209 grid_dim: (1, 1, 1),
9210 block_dim: (1, 1, 1),
9211 shared_mem_bytes: 0,
9212 };
9213 let ii = idx as i32;
9214 let __s_b = self.gpu.stream();
9215 let mut b = __s_b.launch_builder(&f);
9216 b.arg(dst).arg(&v).arg(&ii);
9217 unsafe {
9218 b.launch(cfg)?;
9219 }
9220 Ok(())
9221 }
9222
9223 pub fn i32_add_k(
9225 &self,
9226 d: &mut CudaSlice<i32>,
9227 v: i32,
9228 ) -> Result<(), Box<dyn std::error::Error>> {
9229 let f = self.func("i32_add_k");
9230 let cfg = LaunchConfig {
9231 grid_dim: (1, 1, 1),
9232 block_dim: (32, 1, 1),
9233 shared_mem_bytes: 0,
9234 };
9235 let __s_b = self.gpu.stream();
9236 let mut b = __s_b.launch_builder(&f);
9237 b.arg(d).arg(&v);
9238 unsafe {
9239 b.launch(cfg)?;
9240 }
9241 Ok(())
9242 }
9243
9244 pub fn i32_iota_from(
9246 &self,
9247 ctr: &CudaSlice<i32>,
9248 dst: &mut CudaSlice<i32>,
9249 n: usize,
9250 ) -> Result<(), Box<dyn std::error::Error>> {
9251 let f = self.func("i32_iota_from");
9252 let cfg = LaunchConfig::for_num_elems(n as u32);
9253 let ni = n as i32;
9254 let __s_b = self.gpu.stream();
9255 let mut b = __s_b.launch_builder(&f);
9256 b.arg(ctr).arg(dst).arg(&ni);
9257 unsafe {
9258 b.launch(cfg)?;
9259 }
9260 Ok(())
9261 }
9262
9263 pub fn u32_map_k(
9265 &self,
9266 buf: &mut CudaSlice<u32>,
9267 map: &CudaSlice<u32>,
9268 idx: usize,
9269 ) -> Result<(), Box<dyn std::error::Error>> {
9270 let f = self.func("u32_map_k");
9271 let cfg = LaunchConfig {
9272 grid_dim: (1, 1, 1),
9273 block_dim: (1, 1, 1),
9274 shared_mem_bytes: 0,
9275 };
9276 let ii = idx as i32;
9277 let __s_b = self.gpu.stream();
9278 let mut b = __s_b.launch_builder(&f);
9279 b.arg(buf).arg(map).arg(&ii);
9280 unsafe {
9281 b.launch(cfg)?;
9282 }
9283 Ok(())
9284 }
9285
9286 #[allow(clippy::too_many_arguments)]
9288 pub fn u32_pack2(
9289 &self,
9290 a: &CudaSlice<u32>,
9291 off_a: usize,
9292 n1: usize,
9293 b_in: &CudaSlice<u32>,
9294 n2: usize,
9295 out: &mut CudaSlice<u32>,
9296 ) -> Result<(), Box<dyn std::error::Error>> {
9297 let f = self.func("u32_pack2");
9298 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
9299 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
9300 let __s_b = self.gpu.stream();
9301 let mut b = __s_b.launch_builder(&f);
9302 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
9303 unsafe {
9304 b.launch(cfg)?;
9305 }
9306 Ok(())
9307 }
9308
9309 pub fn moe_w_exscale(
9311 &self,
9312 w: &mut CudaSlice<f32>,
9313 sel: &CudaSlice<i32>,
9314 s: &CudaSlice<f32>,
9315 n: usize,
9316 ) -> Result<(), Box<dyn std::error::Error>> {
9317 let f = self.func("moe_w_exscale");
9318 let cfg = LaunchConfig::for_num_elems(n as u32);
9319 let ni = n as i32;
9320 let __s_b = self.gpu.stream();
9321 let mut b = __s_b.launch_builder(&f);
9322 b.arg(w).arg(sel).arg(s).arg(&ni);
9323 unsafe {
9324 b.launch(cfg)?;
9325 }
9326 Ok(())
9327 }
9328
9329 pub fn moe_w_scale_by_expert(
9332 &self,
9333 w: &mut CudaSlice<f32>,
9334 sel: &CudaSlice<i32>,
9335 macros: &CudaSlice<f32>,
9336 n_expert: usize,
9337 n: usize,
9338 ) -> Result<(), Box<dyn std::error::Error>> {
9339 let f = self.func("moe_w_scale_by_expert");
9340 let cfg = LaunchConfig {
9341 grid_dim: (n.div_ceil(64) as u32, 1, 1),
9342 block_dim: (64, 1, 1),
9343 shared_mem_bytes: 0,
9344 };
9345 let (ne, nn) = (n_expert as i32, n as i32);
9346 let __s_b = self.gpu.stream();
9347 let mut b = __s_b.launch_builder(&f);
9348 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
9349 unsafe {
9350 b.launch(cfg)?;
9351 }
9352 Ok(())
9353 }
9354
9355 #[allow(clippy::too_many_arguments)] pub fn moe_gate_up_silu8_dev_q8(
9357 &self,
9358 table: &CudaSlice<u64>,
9359 sel: &cudarc::driver::CudaView<i32>,
9360 aq: &CudaSlice<i8>,
9361 ad: &CudaSlice<f32>,
9362 in_f: usize,
9363 n_ff: usize,
9364 n_used: usize,
9365 n_expert: usize,
9366 qt_g: i32,
9367 qt_u: i32,
9368 rb_g: usize,
9369 rb_u: usize,
9370 macros: &CudaSlice<f32>,
9371 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9372 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
9373 let (mode, wpb) = GU.get_or_init(|| {
9374 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
9375 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
9376 .ok()
9377 .and_then(|v| v.parse().ok())
9378 .unwrap_or(4u32)
9379 .clamp(1, 16);
9380 (mode, wpb)
9381 });
9382 let (mode, wpb) = (mode.as_str(), *wpb);
9383 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
9384 let (inf, nff, ne, rbg, rbu) = (
9385 in_f as i32,
9386 n_ff as i32,
9387 n_expert as i32,
9388 rb_g as i64,
9389 rb_u as i64,
9390 );
9391 let (f, cfg) = match mode {
9392 "1" | "2" | "4" => {
9393 let rpw: u32 = mode.parse().unwrap();
9394 let f = self.func(match rpw {
9395 1 => "moe_gate_up_silu8_dev_q8_r1",
9396 2 => "moe_gate_up_silu8_dev_q8_r2",
9397 _ => "moe_gate_up_silu8_dev_q8_r4",
9398 });
9399 let rows_per_block = (rpw * wpb) as usize;
9400 let gx = n_ff.div_ceil(rows_per_block) as u32;
9401 (
9402 f,
9403 LaunchConfig {
9404 grid_dim: (gx, n_used as u32, 1),
9405 block_dim: (32, wpb, 1),
9406 shared_mem_bytes: 0,
9407 },
9408 )
9409 }
9410 "j8" if n_used <= 32 => (
9411 self.func("moe_gate_up_silu8_dev_q8_j8"),
9412 LaunchConfig {
9413 grid_dim: (n_ff as u32, 1, 1),
9414 block_dim: (32, n_used as u32, 1),
9415 shared_mem_bytes: 0,
9416 },
9417 ),
9418 "vsm2" => {
9420 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
9421 let sh = (rb_g + rb_u) as u32;
9422 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9423 f.set_attribute(
9424 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
9425 sh as i32,
9426 )?;
9427 (
9428 f,
9429 LaunchConfig {
9430 grid_dim: (n_ff as u32, n_used as u32, 1),
9431 block_dim: (32, 1, 1),
9432 shared_mem_bytes: sh,
9433 },
9434 )
9435 }
9436 "vsm" => {
9437 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
9438 let sh = (rb_g + rb_u) as u32;
9439 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9440 f.set_attribute(
9441 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
9442 sh as i32,
9443 )?;
9444 (
9445 f,
9446 LaunchConfig {
9447 grid_dim: (n_ff as u32, n_used as u32, 1),
9448 block_dim: (32, 1, 1),
9449 shared_mem_bytes: sh,
9450 },
9451 )
9452 }
9453 "sg" => (
9454 self.func("moe_gate_up_silu8_dev_q8_sg"),
9455 LaunchConfig {
9456 grid_dim: (n_ff as u32, n_used as u32, 1),
9457 block_dim: (32, 1, 1),
9458 shared_mem_bytes: 0,
9459 },
9460 ),
9461 "j8sg" if n_used <= 32 => (
9462 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
9463 LaunchConfig {
9464 grid_dim: (n_ff as u32, 1, 1),
9465 block_dim: (32, n_used as u32, 1),
9466 shared_mem_bytes: 0,
9467 },
9468 ),
9469 "u64" if in_f == 2048 => (
9470 self.func("moe_gate_up_silu8_dev_q8_u64"),
9471 LaunchConfig {
9472 grid_dim: (n_ff as u32, n_used as u32, 1),
9473 block_dim: (32, 1, 1),
9474 shared_mem_bytes: 0,
9475 },
9476 ),
9477 "gs4" if in_f == 2048 => (
9478 self.func("moe_gate_up_silu8_dev_q8_gs4"),
9479 LaunchConfig {
9480 grid_dim: (n_ff as u32, n_used as u32, 1),
9481 block_dim: (32, 4, 1),
9482 shared_mem_bytes: 0,
9483 },
9484 ),
9485 "v" | "" => (
9487 self.func("moe_gate_up_silu8_dev_q8_v"),
9488 LaunchConfig {
9489 grid_dim: (n_ff as u32, n_used as u32, 1),
9490 block_dim: (32, 1, 1),
9491 shared_mem_bytes: 0,
9492 },
9493 ),
9494 "s2" => (
9495 self.func("moe_gate_up_silu8_dev_q8_s2"),
9496 LaunchConfig {
9497 grid_dim: (n_ff as u32, n_used as u32, 1),
9498 block_dim: (32, 2, 1),
9499 shared_mem_bytes: 0,
9500 },
9501 ),
9502 "s2z" => {
9503 let rz = wpb.min(16); (
9505 self.func("moe_gate_up_silu8_dev_q8_s2z"),
9506 LaunchConfig {
9507 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
9508 block_dim: (32, 2, rz),
9509 shared_mem_bytes: 0,
9510 },
9511 )
9512 }
9513 _ => (
9514 self.func("moe_gate_up_silu8_dev_q8"),
9515 LaunchConfig {
9516 grid_dim: (n_ff as u32, n_used as u32, 1),
9517 block_dim: (32, 1, 1),
9518 shared_mem_bytes: 0,
9519 },
9520 ),
9521 };
9522 let __s_b = self.gpu.stream();
9523 let mut b = __s_b.launch_builder(&f);
9524 b.arg(table)
9525 .arg(sel)
9526 .arg(aq)
9527 .arg(ad)
9528 .arg(&mut act)
9529 .arg(&inf)
9530 .arg(&nff)
9531 .arg(&ne)
9532 .arg(&qt_g)
9533 .arg(&qt_u)
9534 .arg(&rbg)
9535 .arg(&rbu)
9536 .arg(macros);
9537 unsafe {
9538 b.launch(cfg)?;
9539 }
9540 Ok(act)
9541 }
9542
9543 #[allow(clippy::too_many_arguments)]
9544 pub fn moe_down8_fma_dev_q8(
9545 &self,
9546 table: &CudaSlice<u64>,
9547 sel: &cudarc::driver::CudaView<i32>,
9548 w: &cudarc::driver::CudaView<f32>,
9549 aq2: &CudaSlice<i8>,
9550 ad2: &CudaSlice<f32>,
9551 dst: &mut cudarc::driver::CudaViewMut<f32>,
9552 in_f: usize,
9553 out_f: usize,
9554 n_used: usize,
9555 n_expert: usize,
9556 qt: i32,
9557 rb: usize,
9558 ) -> Result<(), Box<dyn std::error::Error>> {
9559 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
9560 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
9561 let (inf, outf, nu, ne, rbi) = (
9562 in_f as i32,
9563 out_f as i32,
9564 n_used as i32,
9565 n_expert as i32,
9566 rb as i64,
9567 );
9568 let (f, cfg) = match mode.as_str() {
9571 m @ ("1" | "2" | "4") if n_used <= 8 => {
9572 let rpw: usize = m.parse().unwrap();
9573 let f = self.func(match rpw {
9574 1 => "moe_down8_fma_dev_q8_w8r1",
9575 2 => "moe_down8_fma_dev_q8_w8r2",
9576 _ => "moe_down8_fma_dev_q8_w8r4",
9577 });
9578 (
9579 f,
9580 LaunchConfig {
9581 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
9582 block_dim: (32, n_used as u32, 1),
9583 shared_mem_bytes: 0,
9584 },
9585 )
9586 }
9587 "h2" if in_f == 512 => (
9588 self.func("moe_down8_fma_dev_q8_h2"),
9589 LaunchConfig {
9590 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9591 block_dim: (32, 1, 1),
9592 shared_mem_bytes: 0,
9593 },
9594 ),
9595 "" if in_f == 704 && n_used <= 8 => (
9598 self.func("moe_down8_fma_dev_q8_w8r2"),
9599 LaunchConfig {
9600 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9601 block_dim: (32, n_used as u32, 1),
9602 shared_mem_bytes: 0,
9603 },
9604 ),
9605 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
9609 self.func("moe_down8_fma_dev_q8_w8h2v"),
9610 LaunchConfig {
9611 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9612 block_dim: (32, n_used as u32, 1),
9613 shared_mem_bytes: 0,
9614 },
9615 ),
9616 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
9617 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
9618 LaunchConfig {
9619 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
9620 block_dim: (32, n_used as u32, 1),
9621 shared_mem_bytes: 0,
9622 },
9623 ),
9624 "w8h2r2" if in_f == 512 && n_used <= 8 => (
9625 self.func("moe_down8_fma_dev_q8_w8h2r2"),
9626 LaunchConfig {
9627 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
9628 block_dim: (32, n_used as u32, 1),
9629 shared_mem_bytes: 0,
9630 },
9631 ),
9632 "w8h2" if in_f == 512 && n_used <= 8 => (
9633 self.func("moe_down8_fma_dev_q8_w8h2"),
9634 LaunchConfig {
9635 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9636 block_dim: (32, n_used as u32, 1),
9637 shared_mem_bytes: 0,
9638 },
9639 ),
9640 _ => (
9641 self.func("moe_down8_fma_dev_q8"),
9642 LaunchConfig {
9643 grid_dim: (out_f as u32, 1, 1),
9644 block_dim: (32, 1, 1),
9645 shared_mem_bytes: 0,
9646 },
9647 ),
9648 };
9649 let __s_b = self.gpu.stream();
9650 let mut b = __s_b.launch_builder(&f);
9651 b.arg(table)
9652 .arg(sel)
9653 .arg(w)
9654 .arg(aq2)
9655 .arg(ad2)
9656 .arg(dst)
9657 .arg(&inf)
9658 .arg(&outf)
9659 .arg(&nu)
9660 .arg(&ne)
9661 .arg(&qt)
9662 .arg(&rbi);
9663 unsafe {
9664 b.launch(cfg)?;
9665 }
9666 Ok(())
9667 }
9668
9669 #[allow(clippy::too_many_arguments)]
9676 pub fn moe_gate_up_silu8_dev_q8_rows(
9677 &self,
9678 table: &CudaSlice<u64>,
9679 sel: &CudaSlice<i32>,
9680 aq: &CudaSlice<i8>,
9681 ad: &CudaSlice<f32>,
9682 t: usize,
9683 in_f: usize,
9684 n_ff: usize,
9685 n_used: usize,
9686 n_expert: usize,
9687 qt_g: i32,
9688 qt_u: i32,
9689 rb_g: usize,
9690 rb_u: usize,
9691 macros: &CudaSlice<f32>,
9692 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9693 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
9694 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
9695 let cfg = LaunchConfig {
9696 grid_dim: (n_ff as u32, n_used as u32, t as u32),
9697 block_dim: (32, 1, 1),
9698 shared_mem_bytes: 0,
9699 };
9700 let (inf, nff, ne, nu, rbg, rbu) = (
9701 in_f as i32,
9702 n_ff as i32,
9703 n_expert as i32,
9704 n_used as i32,
9705 rb_g as i64,
9706 rb_u as i64,
9707 );
9708 let __s_b = self.gpu.stream();
9709 let mut b = __s_b.launch_builder(&f);
9710 b.arg(table)
9711 .arg(sel)
9712 .arg(aq)
9713 .arg(ad)
9714 .arg(&mut act)
9715 .arg(&inf)
9716 .arg(&nff)
9717 .arg(&ne)
9718 .arg(&qt_g)
9719 .arg(&qt_u)
9720 .arg(&rbg)
9721 .arg(&rbu)
9722 .arg(&nu)
9723 .arg(macros);
9724 unsafe {
9725 b.launch(cfg)?;
9726 }
9727 Ok(act)
9728 }
9729
9730 #[allow(clippy::too_many_arguments)]
9735 pub fn moe_down8_fma_dev_q8_rows(
9736 &self,
9737 table: &CudaSlice<u64>,
9738 sel: &CudaSlice<i32>,
9739 w: &CudaSlice<f32>,
9740 aq2: &CudaSlice<i8>,
9741 ad2: &CudaSlice<f32>,
9742 dst: &mut CudaSlice<f32>,
9743 t: usize,
9744 in_f: usize,
9745 out_f: usize,
9746 n_used: usize,
9747 n_expert: usize,
9748 qt: i32,
9749 rb: usize,
9750 ) -> Result<(), Box<dyn std::error::Error>> {
9751 assert!(
9752 in_f == 512 && n_used <= 8,
9753 "down rows twin is w8h2v shape-gated"
9754 );
9755 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
9756 let cfg = LaunchConfig {
9757 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
9758 block_dim: (32, n_used as u32, 1),
9759 shared_mem_bytes: 0,
9760 };
9761 let (inf, outf, nu, ne, rbi) = (
9762 in_f as i32,
9763 out_f as i32,
9764 n_used as i32,
9765 n_expert as i32,
9766 rb as i64,
9767 );
9768 let __s_b = self.gpu.stream();
9769 let mut b = __s_b.launch_builder(&f);
9770 b.arg(table)
9771 .arg(sel)
9772 .arg(w)
9773 .arg(aq2)
9774 .arg(ad2)
9775 .arg(dst)
9776 .arg(&inf)
9777 .arg(&outf)
9778 .arg(&nu)
9779 .arg(&ne)
9780 .arg(&qt)
9781 .arg(&rbi);
9782 unsafe {
9783 b.launch(cfg)?;
9784 }
9785 Ok(())
9786 }
9787
9788 #[allow(clippy::too_many_arguments)]
9792 pub fn moe_gate_up_silu8_dev_q8_csr(
9793 &self,
9794 table: &CudaSlice<u64>,
9795 sel: &CudaSlice<i32>,
9796 aq: &CudaSlice<i8>,
9797 ad: &CudaSlice<f32>,
9798 n_pairs: usize,
9799 in_f: usize,
9800 n_ff: usize,
9801 n_used: usize,
9802 n_expert: usize,
9803 qt_g: i32,
9804 qt_u: i32,
9805 rb_g: usize,
9806 rb_u: usize,
9807 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9808 let f = if qt_g == crate::QT_NVFP4 {
9811 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
9812 } else {
9813 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
9814 };
9815 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
9816 let cfg = LaunchConfig {
9817 grid_dim: (n_ff as u32, n_pairs as u32, 1),
9818 block_dim: (32, 1, 1),
9819 shared_mem_bytes: 0,
9820 };
9821 let (inf, nff, ne, nu, npi, rbg, rbu) = (
9822 in_f as i32,
9823 n_ff as i32,
9824 n_expert as i32,
9825 n_used as i32,
9826 n_pairs as i32,
9827 rb_g as i64,
9828 rb_u as i64,
9829 );
9830 let __s_b = self.gpu.stream();
9831 let mut b = __s_b.launch_builder(&f);
9832 b.arg(table)
9833 .arg(sel)
9834 .arg(aq)
9835 .arg(ad)
9836 .arg(&mut act)
9837 .arg(&inf)
9838 .arg(&nff)
9839 .arg(&ne)
9840 .arg(&qt_g)
9841 .arg(&qt_u)
9842 .arg(&rbg)
9843 .arg(&rbu)
9844 .arg(&nu)
9845 .arg(&npi);
9846 unsafe {
9847 b.launch(cfg)?;
9848 }
9849 Ok(act)
9850 }
9851
9852 #[allow(clippy::too_many_arguments)]
9856 pub fn moe_down8_fma_dev_q8_variant(
9857 &self,
9858 variant: &str,
9859 table: &CudaSlice<u64>,
9860 sel: &cudarc::driver::CudaView<i32>,
9861 w: &cudarc::driver::CudaView<f32>,
9862 aq2: &CudaSlice<i8>,
9863 ad2: &CudaSlice<f32>,
9864 dst: &mut cudarc::driver::CudaViewMut<f32>,
9865 in_f: usize,
9866 out_f: usize,
9867 n_used: usize,
9868 n_expert: usize,
9869 qt: i32,
9870 rb: usize,
9871 ) -> Result<(), Box<dyn std::error::Error>> {
9872 let (inf, outf, nu, ne, rbi) = (
9873 in_f as i32,
9874 out_f as i32,
9875 n_used as i32,
9876 n_expert as i32,
9877 rb as i64,
9878 );
9879 let (f, cfg) = match variant {
9880 "w8h2" | "w8h2v" => (
9881 self.func(if variant == "w8h2" {
9882 "moe_down8_fma_dev_q8_w8h2"
9883 } else {
9884 "moe_down8_fma_dev_q8_w8h2v"
9885 }),
9886 LaunchConfig {
9887 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9888 block_dim: (32, n_used as u32, 1),
9889 shared_mem_bytes: 0,
9890 },
9891 ),
9892 "w8h2r2" | "w8h2r2v" => (
9893 self.func(if variant == "w8h2r2" {
9894 "moe_down8_fma_dev_q8_w8h2r2"
9895 } else {
9896 "moe_down8_fma_dev_q8_w8h2r2v"
9897 }),
9898 LaunchConfig {
9899 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
9900 block_dim: (32, n_used as u32, 1),
9901 shared_mem_bytes: 0,
9902 },
9903 ),
9904 _ => (
9905 self.func("moe_down8_fma_dev_q8"),
9906 LaunchConfig {
9907 grid_dim: (out_f as u32, 1, 1),
9908 block_dim: (32, 1, 1),
9909 shared_mem_bytes: 0,
9910 },
9911 ),
9912 };
9913 let __s_b = self.gpu.stream();
9914 let mut b = __s_b.launch_builder(&f);
9915 b.arg(table)
9916 .arg(sel)
9917 .arg(w)
9918 .arg(aq2)
9919 .arg(ad2)
9920 .arg(dst)
9921 .arg(&inf)
9922 .arg(&outf)
9923 .arg(&nu)
9924 .arg(&ne)
9925 .arg(&qt)
9926 .arg(&rbi);
9927 unsafe {
9928 b.launch(cfg)?;
9929 }
9930 Ok(())
9931 }
9932
9933 #[allow(clippy::too_many_arguments)]
9935 pub fn moe_gate_up_silu8_dev_q8_variant(
9936 &self,
9937 variant: &str,
9938 table: &CudaSlice<u64>,
9939 sel: &cudarc::driver::CudaView<i32>,
9940 aq: &CudaSlice<i8>,
9941 ad: &CudaSlice<f32>,
9942 in_f: usize,
9943 n_ff: usize,
9944 n_used: usize,
9945 n_expert: usize,
9946 qt_g: i32,
9947 qt_u: i32,
9948 rb_g: usize,
9949 rb_u: usize,
9950 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9951 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
9952 let (inf, nff, ne, rbg, rbu) = (
9953 in_f as i32,
9954 n_ff as i32,
9955 n_expert as i32,
9956 rb_g as i64,
9957 rb_u as i64,
9958 );
9959 let f = self.func(if variant == "v" {
9960 "moe_gate_up_silu8_dev_q8_v"
9961 } else {
9962 "moe_gate_up_silu8_dev_q8"
9963 });
9964 let cfg = LaunchConfig {
9965 grid_dim: (n_ff as u32, n_used as u32, 1),
9966 block_dim: (32, 1, 1),
9967 shared_mem_bytes: 0,
9968 };
9969 let __s_b = self.gpu.stream();
9970 let mut b = __s_b.launch_builder(&f);
9971 b.arg(table)
9972 .arg(sel)
9973 .arg(aq)
9974 .arg(ad)
9975 .arg(&mut act)
9976 .arg(&inf)
9977 .arg(&nff)
9978 .arg(&ne)
9979 .arg(&qt_g)
9980 .arg(&qt_u)
9981 .arg(&rbg)
9982 .arg(&rbu);
9983 unsafe {
9984 b.launch(cfg)?;
9985 }
9986 Ok(act)
9987 }
9988
9989 #[allow(clippy::too_many_arguments)] pub fn moe_gate_up_silu8_dev(
9991 &self,
9992 table: &CudaSlice<u64>,
9993 sel: &cudarc::driver::CudaView<i32>,
9994 x: &cudarc::driver::CudaView<f32>,
9995 in_f: usize,
9996 n_ff: usize,
9997 n_used: usize,
9998 n_expert: usize,
9999 qt_g: i32,
10000 qt_u: i32,
10001 rb_g: usize,
10002 rb_u: usize,
10003 macros: &CudaSlice<f32>,
10004 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10005 let f = self.func("moe_gate_up_silu8_dev");
10006 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
10008 grid_dim: (n_ff as u32, n_used as u32, 1),
10009 block_dim: (256, 1, 1),
10010 shared_mem_bytes: 0,
10011 };
10012 let (inf, nff, ne, rbg, rbu) = (
10013 in_f as i32,
10014 n_ff as i32,
10015 n_expert as i32,
10016 rb_g as i64,
10017 rb_u as i64,
10018 );
10019 let __s_b = self.gpu.stream();
10020 let mut b = __s_b.launch_builder(&f);
10021 b.arg(table)
10022 .arg(sel)
10023 .arg(x)
10024 .arg(&mut act)
10025 .arg(&inf)
10026 .arg(&nff)
10027 .arg(&ne)
10028 .arg(&qt_g)
10029 .arg(&qt_u)
10030 .arg(&rbg)
10031 .arg(&rbu)
10032 .arg(macros);
10033 unsafe {
10034 b.launch(cfg)?;
10035 }
10036 Ok(act)
10037 }
10038
10039 #[allow(clippy::too_many_arguments)]
10042 pub fn moe_down8_fma_dev(
10043 &self,
10044 table: &CudaSlice<u64>,
10045 sel: &cudarc::driver::CudaView<i32>,
10046 w: &cudarc::driver::CudaView<f32>,
10047 act: &CudaSlice<f32>,
10048 dst: &mut cudarc::driver::CudaViewMut<f32>,
10049 in_f: usize,
10050 out_f: usize,
10051 n_used: usize,
10052 n_expert: usize,
10053 qt: i32,
10054 rb: usize,
10055 ) -> Result<(), Box<dyn std::error::Error>> {
10056 let f = self.func("moe_down8_fma_dev");
10057 let cfg = LaunchConfig {
10058 grid_dim: (out_f as u32, 1, 1),
10059 block_dim: (256, 1, 1),
10060 shared_mem_bytes: 0,
10061 };
10062 let (inf, outf, nu, ne, rbv) = (
10063 in_f as i32,
10064 out_f as i32,
10065 n_used as i32,
10066 n_expert as i32,
10067 rb as i64,
10068 );
10069 let __s_b = self.gpu.stream();
10070 let mut b = __s_b.launch_builder(&f);
10071 b.arg(table)
10072 .arg(sel)
10073 .arg(w)
10074 .arg(act)
10075 .arg(dst)
10076 .arg(&inf)
10077 .arg(&outf)
10078 .arg(&nu)
10079 .arg(&ne)
10080 .arg(&qt)
10081 .arg(&rbv);
10082 unsafe {
10083 b.launch(cfg)?;
10084 }
10085 Ok(())
10086 }
10087
10088 pub fn axpy_into(
10090 &self,
10091 src: &CudaSlice<f32>,
10092 alpha: f32,
10093 dst: &mut cudarc::driver::CudaViewMut<f32>,
10094 n: usize,
10095 ) -> Result<(), Box<dyn std::error::Error>> {
10096 let f = self.func("axpy_f32");
10097 let cfg = LaunchConfig::for_num_elems(n as u32);
10098 let (a, ni) = (alpha, n as i32);
10099 let __s_b = self.gpu.stream();
10100 let mut b = __s_b.launch_builder(&f);
10101 b.arg(src).arg(dst).arg(&a).arg(&ni);
10102 unsafe {
10103 b.launch(cfg)?;
10104 }
10105 Ok(())
10106 }
10107
10108 pub fn axpy_host_into(
10110 &self,
10111 src: &cudarc::driver::CudaView<'_, f32>,
10112 alpha: f32,
10113 dst: &mut cudarc::driver::CudaViewMut<f32>,
10114 n: usize,
10115 ) -> Result<(), Box<dyn std::error::Error>> {
10116 let f = self.func("axpy_host_f32");
10117 let cfg = LaunchConfig::for_num_elems(n as u32);
10118 let (a, ni) = (alpha, n as i32);
10119 let __s_b = self.gpu.stream();
10120 let mut b = __s_b.launch_builder(&f);
10121 b.arg(src).arg(dst).arg(&a).arg(&ni);
10122 unsafe {
10123 b.launch(cfg)?;
10124 }
10125 Ok(())
10126 }
10127
10128 pub fn add_scaled_rows(
10130 &self,
10131 src: &CudaSlice<f32>,
10132 scale: &CudaSlice<f32>,
10133 dst: &mut CudaSlice<f32>,
10134 ncols: usize,
10135 nrows: usize,
10136 ) -> Result<(), Box<dyn std::error::Error>> {
10137 let f = self.func("add_scaled_rows_f32");
10138 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10139 let (nc, nr) = (ncols as i32, nrows as i32);
10140 let __s_b = self.gpu.stream();
10141 let mut b = __s_b.launch_builder(&f);
10142 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
10143 unsafe {
10144 b.launch(cfg)?;
10145 }
10146 Ok(())
10147 }
10148
10149 pub fn add_scaled_rows_ones(
10154 &self,
10155 src: &CudaSlice<f32>,
10156 dst: &mut CudaSlice<f32>,
10157 ncols: usize,
10158 nrows: usize,
10159 ) -> Result<(), Box<dyn std::error::Error>> {
10160 let mut guard = self
10161 .shexp_ones
10162 .lock()
10163 .map_err(|_| "shexp ones buffer is poisoned")?;
10164 if guard.as_ref().map(|b| b.len() < nrows).unwrap_or(true) {
10165 *guard = Some(self.htod(&vec![1.0f32; nrows.max(64)])?);
10168 }
10169 let ones = guard.as_ref().expect("just ensured");
10170 let f = self.func("add_scaled_rows_f32");
10171 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10172 let (nc, nr) = (ncols as i32, nrows as i32);
10173 let __s_b = self.gpu.stream();
10174 let mut b = __s_b.launch_builder(&f);
10175 b.arg(src).arg(ones).arg(&mut *dst).arg(&nc).arg(&nr);
10176 unsafe {
10177 b.launch(cfg)?;
10178 }
10179 Ok(())
10180 }
10181
10182 pub fn i32_mirror_store(
10186 &self,
10187 dst: &mut CudaSlice<i32>,
10188 v: i32,
10189 ) -> Result<(), Box<dyn std::error::Error>> {
10190 if crate::htod_diet_on() {
10191 HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10192 return self.i32_set_k(dst, v);
10193 }
10194 self.gpu.stream().memcpy_htod(&[v], dst)?;
10195 Ok(())
10196 }
10197
10198 pub fn scale_rows(
10201 &self,
10202 y: &mut CudaSlice<f32>,
10203 s: &CudaSlice<f32>,
10204 ncols: usize,
10205 nrows: usize,
10206 ) -> Result<(), Box<dyn std::error::Error>> {
10207 let f = self.func("scale_rows_f32");
10208 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10209 let (nc, nr) = (ncols as i32, nrows as i32);
10210 let __s_b = self.gpu.stream();
10211 let mut b = __s_b.launch_builder(&f);
10212 b.arg(&mut *y).arg(s).arg(&nc).arg(&nr);
10213 unsafe {
10214 b.launch(cfg)?;
10215 }
10216 Ok(())
10217 }
10218
10219 #[allow(clippy::too_many_arguments)]
10223 pub fn moe_prime_join_scatter(
10224 &self,
10225 y0: &CudaSlice<f32>,
10226 y1: &CudaSlice<f32>,
10227 inv: &CudaSlice<i32>,
10228 w: &CudaSlice<f32>,
10229 out: &mut CudaSlice<f32>,
10230 ncols: usize,
10231 n_used: usize,
10232 t: usize,
10233 ) -> Result<(), Box<dyn std::error::Error>> {
10234 let f = self.func("moe_prime_join_scatter_f32");
10235 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10236 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10237 let __s_b = self.gpu.stream();
10238 let mut b = __s_b.launch_builder(&f);
10239 b.arg(y0)
10240 .arg(y1)
10241 .arg(inv)
10242 .arg(w)
10243 .arg(&mut *out)
10244 .arg(&nc)
10245 .arg(&nu)
10246 .arg(&ti);
10247 unsafe {
10248 b.launch(cfg)?;
10249 }
10250 Ok(())
10251 }
10252
10253 pub fn moe_pairs_weighted_scatter(
10256 &self,
10257 y: &CudaSlice<f32>,
10258 w: &CudaSlice<f32>,
10259 out: &mut CudaSlice<f32>,
10260 ncols: usize,
10261 n_used: usize,
10262 t: usize,
10263 ) -> Result<(), Box<dyn std::error::Error>> {
10264 let f = self.func("moe_pairs_weighted_scatter_f32");
10265 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10266 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10267 let __s_b = self.gpu.stream();
10268 let mut b = __s_b.launch_builder(&f);
10269 b.arg(y).arg(w).arg(&mut *out).arg(&nc).arg(&nu).arg(&ti);
10270 unsafe {
10271 b.launch(cfg)?;
10272 }
10273 Ok(())
10274 }
10275
10276 pub fn gather_rows(
10280 &self,
10281 src: &CudaSlice<f32>,
10282 idx: &CudaSlice<i32>,
10283 dst: &mut CudaSlice<f32>,
10284 ncols: usize,
10285 m_e: usize,
10286 ) -> Result<(), Box<dyn std::error::Error>> {
10287 let f = self.func("gather_rows_f32");
10288 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
10289 let (nc, me) = (ncols as i32, m_e as i32);
10290 let __s_b = self.gpu.stream();
10291 let mut b = __s_b.launch_builder(&f);
10292 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
10293 unsafe {
10294 b.launch(cfg)?;
10295 }
10296 Ok(())
10297 }
10298
10299 #[allow(clippy::too_many_arguments)] pub fn scatter_slot(
10305 &self,
10306 src: &CudaSlice<f32>,
10307 tok_idx: &CudaSlice<i32>,
10308 slot_idx: &CudaSlice<i32>,
10309 weight: &CudaSlice<f32>,
10310 dst: &mut CudaSlice<f32>,
10311 wbuf: &mut CudaSlice<f32>,
10312 ncols: usize,
10313 n_used: usize,
10314 m_e: usize,
10315 ) -> Result<(), Box<dyn std::error::Error>> {
10316 let f = self.func("scatter_add_slot_f32");
10317 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
10318 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
10319 let __s_b = self.gpu.stream();
10320 let mut b = __s_b.launch_builder(&f);
10321 b.arg(src)
10322 .arg(tok_idx)
10323 .arg(slot_idx)
10324 .arg(weight)
10325 .arg(dst)
10326 .arg(wbuf)
10327 .arg(&nc)
10328 .arg(&nu)
10329 .arg(&me);
10330 unsafe {
10331 b.launch(cfg)?;
10332 }
10333 Ok(())
10334 }
10335
10336 pub fn reduce_slots(
10340 &self,
10341 slots: &CudaSlice<f32>,
10342 wbuf: &CudaSlice<f32>,
10343 dst: &mut CudaSlice<f32>,
10344 ncols: usize,
10345 n_used: usize,
10346 t: usize,
10347 ) -> Result<(), Box<dyn std::error::Error>> {
10348 let f = self.func("reduce_slots_f32");
10349 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10350 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10351 let __s_b = self.gpu.stream();
10352 let mut b = __s_b.launch_builder(&f);
10353 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
10354 unsafe {
10355 b.launch(cfg)?;
10356 }
10357 Ok(())
10358 }
10359
10360 pub fn reduce_slots_host(
10365 &self,
10366 slots: &CudaSlice<f32>,
10367 wbuf: &CudaSlice<f32>,
10368 dst: &mut CudaSlice<f32>,
10369 ncols: usize,
10370 n_used: usize,
10371 t: usize,
10372 ) -> Result<(), Box<dyn std::error::Error>> {
10373 let f = self.func("reduce_slots_host_f32");
10374 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10375 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10376 let __s_b = self.gpu.stream();
10377 let mut b = __s_b.launch_builder(&f);
10378 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
10379 unsafe {
10380 b.launch(cfg)?;
10381 }
10382 Ok(())
10383 }
10384
10385 #[track_caller]
10392 pub fn quantize_q8_1_view(
10393 &self,
10394 x: &cudarc::driver::CudaView<f32>,
10395 m: usize,
10396 in_f: usize,
10397 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10398 q8_census_record(std::panic::Location::caller(), m, in_f);
10399 let f = self.func("quantize_q8_1");
10400 let nblk = in_f / 32;
10401 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
10402 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
10403 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10404 let (inf, mi) = (in_f as i32, m as i32);
10405 let __s_b = self.gpu.stream();
10406 let mut b = __s_b.launch_builder(&f);
10407 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
10408 unsafe {
10409 b.launch(cfg)?;
10410 }
10411 Ok((q, d))
10412 }
10413
10414 #[track_caller]
10415 pub fn quantize_q8_1(
10416 &self,
10417 x: &CudaSlice<f32>,
10418 m: usize,
10419 in_f: usize,
10420 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10421 q8_census_record(std::panic::Location::caller(), m, in_f);
10422 let nblk = in_f / 32;
10423 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);
10427 let (inf, mi) = (in_f as i32, m as i32);
10428 if Self::pdl_on() && Self::pdl_wb_on() {
10429 {
10430 use cudarc::driver::{DevicePtr, DevicePtrMut};
10431 let s = &self.gpu.stream();
10432 let (px, _g0) = x.device_ptr(s);
10433 let (pq, _g1) = q.device_ptr_mut(s);
10434 let (pd, _g2) = d.device_ptr_mut(s);
10435 let mut ps = [
10436 &px as *const _ as *mut std::ffi::c_void,
10437 &pq as *const _ as *mut _,
10438 &pd as *const _ as *mut _,
10439 &inf as *const _ as *mut _,
10440 &mi as *const _ as *mut _,
10441 ];
10442 unsafe {
10443 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10444 }
10445 }
10446 return Ok((q, d));
10447 }
10448 let f = self.func("quantize_q8_1");
10449 let __s_b = self.gpu.stream();
10450 let mut b = __s_b.launch_builder(&f);
10451 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
10452 unsafe {
10453 b.launch(cfg)?;
10454 }
10455 Ok((q, d))
10456 }
10457
10458 pub fn quantize_fp4_act(
10462 &self,
10463 x: &CudaSlice<f32>,
10464 m: usize,
10465 in_f: usize,
10466 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
10467 let f = self.func("quantize_fp4_act");
10468 let nb16 = in_f / 16;
10469 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);
10472 let (inf, mi) = (in_f as i32, m as i32);
10473 let __s_b = self.gpu.stream();
10474 let mut b = __s_b.launch_builder(&f);
10475 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
10476 unsafe {
10477 b.launch(cfg)?;
10478 }
10479 Ok((aq4, ad4))
10480 }
10481
10482 #[allow(clippy::too_many_arguments)] pub fn qmatvec_gemm_nvfp4_fp4(
10488 &self,
10489 bytes: &CudaSlice<u8>,
10490 x: &CudaSlice<f32>,
10491 m: usize,
10492 in_f: usize,
10493 out_f: usize,
10494 row_bytes: usize,
10495 scale: f32,
10496 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10497 assert!(
10498 in_f.is_multiple_of(64),
10499 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
10500 );
10501 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
10502 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
10503 if scale != 1.0 {
10504 self.scale_inplace(&mut y, scale, m * out_f)?;
10505 }
10506 Ok(y)
10507 }
10508
10509 #[allow(clippy::too_many_arguments)]
10512 #[allow(clippy::manual_div_ceil)] fn fp4_gemm_launch(
10515 &self,
10516 bytes: &CudaSlice<u8>,
10517 aq4: &CudaSlice<u32>,
10518 ad4: &CudaSlice<u8>,
10519 m: usize,
10520 in_f: usize,
10521 out_f: usize,
10522 row_bytes: usize,
10523 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10524 let f = self.func("qmatvec_gemm_nvfp4_fp4");
10525 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
10527 const BN: u32 = 256;
10528 let cfg = LaunchConfig {
10529 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
10530 block_dim: (32, 4, 1),
10531 shared_mem_bytes: 0,
10532 };
10533 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10534 let __s_b = self.gpu.stream();
10535 let mut b = __s_b.launch_builder(&f);
10536 b.arg(bytes)
10537 .arg(aq4)
10538 .arg(ad4)
10539 .arg(&mut y)
10540 .arg(&inf)
10541 .arg(&outf)
10542 .arg(&mi)
10543 .arg(&rb);
10544 unsafe {
10545 b.launch(cfg)?;
10546 }
10547 Ok(y)
10548 }
10549
10550 pub fn qmatvec_gemm_nvfp4_fp4_raw(
10552 &self,
10553 bytes: &CudaSlice<u8>,
10554 x: &CudaSlice<f32>,
10555 m: usize,
10556 in_f: usize,
10557 out_f: usize,
10558 row_bytes: usize,
10559 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10560 assert!(
10561 in_f.is_multiple_of(64),
10562 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
10563 );
10564 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
10565 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
10566 }
10567
10568 pub fn qmatvec_q8_0_fast(
10570 &self,
10571 w: &CudaSlice<u8>,
10572 x: &CudaSlice<f32>,
10573 m: usize,
10574 in_f: usize,
10575 out_f: usize,
10576 row_bytes: usize,
10577 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10578 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10579 let f = self.func("qmatvec_q8_0_dp4a");
10580 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10582 grid_dim: (out_f as u32, m as u32, 1),
10583 block_dim: (128, 1, 1),
10584 shared_mem_bytes: 0,
10585 };
10586 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10587 let __s_b = self.gpu.stream();
10588 let mut b = __s_b.launch_builder(&f);
10589 b.arg(w)
10590 .arg(&aq)
10591 .arg(&ad)
10592 .arg(&mut y)
10593 .arg(&inf)
10594 .arg(&outf)
10595 .arg(&mi)
10596 .arg(&rb);
10597 unsafe {
10598 b.launch(cfg)?;
10599 }
10600 Ok(y)
10601 }
10602
10603 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
10606 &self,
10607 w: &CudaSlice<u8>,
10608 x: &CudaSlice<f32>,
10609 m: usize,
10610 in_f: usize,
10611 out_f: usize,
10612 row_bytes: usize,
10613 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10614 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10615 let f = self.func("qmatvec_q4_K_dp4a");
10616 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10618 grid_dim: (out_f as u32, m as u32, 1),
10619 block_dim: (128, 1, 1),
10620 shared_mem_bytes: 0,
10621 };
10622 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10623 let __s_b = self.gpu.stream();
10624 let mut b = __s_b.launch_builder(&f);
10625 b.arg(w)
10626 .arg(&aq)
10627 .arg(&ad)
10628 .arg(&mut y)
10629 .arg(&inf)
10630 .arg(&outf)
10631 .arg(&mi)
10632 .arg(&rb);
10633 unsafe {
10634 b.launch(cfg)?;
10635 }
10636 Ok(y)
10637 }
10638
10639 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
10642 &self,
10643 w: &CudaSlice<u8>,
10644 x: &CudaSlice<f32>,
10645 m: usize,
10646 in_f: usize,
10647 out_f: usize,
10648 row_bytes: usize,
10649 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10650 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10651 let f = self.func("qmatvec_q6_K_dp4a");
10652 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10654 grid_dim: (out_f as u32, m as u32, 1),
10655 block_dim: (128, 1, 1),
10656 shared_mem_bytes: 0,
10657 };
10658 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10659 let __s_b = self.gpu.stream();
10660 let mut b = __s_b.launch_builder(&f);
10661 b.arg(w)
10662 .arg(&aq)
10663 .arg(&ad)
10664 .arg(&mut y)
10665 .arg(&inf)
10666 .arg(&outf)
10667 .arg(&mi)
10668 .arg(&rb);
10669 unsafe {
10670 b.launch(cfg)?;
10671 }
10672 Ok(y)
10673 }
10674
10675 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
10678 &self,
10679 w: &CudaSlice<u8>,
10680 x: &CudaSlice<f32>,
10681 m: usize,
10682 in_f: usize,
10683 out_f: usize,
10684 row_bytes: usize,
10685 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10686 self.qmatvec_dp4a_named(
10687 "qmatvec_q5_K_dp4a",
10688 &w.slice(0..w.len()),
10689 x,
10690 m,
10691 in_f,
10692 out_f,
10693 row_bytes,
10694 )
10695 }
10696 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
10699 &self,
10700 w: &CudaSlice<u8>,
10701 x: &CudaSlice<f32>,
10702 m: usize,
10703 in_f: usize,
10704 out_f: usize,
10705 row_bytes: usize,
10706 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10707 self.qmatvec_dp4a_named(
10708 "qmatvec_q3_K_dp4a",
10709 &w.slice(0..w.len()),
10710 x,
10711 m,
10712 in_f,
10713 out_f,
10714 row_bytes,
10715 )
10716 }
10717 pub fn qmatvec_nvfp4_fast_rp(
10719 &self,
10720 w: &CudaSlice<u8>,
10721 x: &CudaSlice<f32>,
10722 m: usize,
10723 in_f: usize,
10724 out_f: usize,
10725 row_bytes: usize,
10726 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10727 assert!(
10728 in_f.is_multiple_of(64),
10729 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10730 );
10731 self.qmatvec_dp4a_named(
10732 "qmatvec_nvfp4_dp4a_rp",
10733 &w.slice(0..w.len()),
10734 x,
10735 m,
10736 in_f,
10737 out_f,
10738 row_bytes,
10739 )
10740 }
10741 pub fn qmatvec_nvfp4_fast(
10743 &self,
10744 w: &cudarc::driver::CudaView<'_, u8>,
10745 x: &CudaSlice<f32>,
10746 m: usize,
10747 in_f: usize,
10748 out_f: usize,
10749 row_bytes: usize,
10750 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10751 assert!(
10754 in_f.is_multiple_of(64),
10755 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10756 );
10757 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
10758 }
10759 pub fn qmatvec_nvfp4_fast_v2(
10764 &self,
10765 w: &cudarc::driver::CudaView<'_, u8>,
10766 x: &CudaSlice<f32>,
10767 m: usize,
10768 in_f: usize,
10769 out_f: usize,
10770 row_bytes: usize,
10771 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10772 assert!(
10773 in_f.is_multiple_of(64),
10774 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10775 );
10776 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
10777 }
10778 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
10781 &self,
10782 w: &CudaSlice<u8>,
10783 x: &CudaSlice<f32>,
10784 m: usize,
10785 in_f: usize,
10786 out_f: usize,
10787 row_bytes: usize,
10788 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10789 self.qmatvec_dp4a_named(
10790 "qmatvec_iq4_XS_dp4a",
10791 &w.slice(0..w.len()),
10792 x,
10793 m,
10794 in_f,
10795 out_f,
10796 row_bytes,
10797 )
10798 }
10799
10800 #[allow(clippy::too_many_arguments)] fn qmatvec_dp4a_named(
10803 &self,
10804 name: &str,
10805 w: &cudarc::driver::CudaView<'_, u8>,
10806 x: &CudaSlice<f32>,
10807 m: usize,
10808 in_f: usize,
10809 out_f: usize,
10810 row_bytes: usize,
10811 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10812 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10813 let f = self.func(name);
10814 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10816 grid_dim: (out_f as u32, m as u32, 1),
10817 block_dim: (128, 1, 1),
10818 shared_mem_bytes: 0,
10819 };
10820 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10821 let __s_b = self.gpu.stream();
10822 let mut b = __s_b.launch_builder(&f);
10823 b.arg(w)
10824 .arg(&aq)
10825 .arg(&ad)
10826 .arg(&mut y)
10827 .arg(&inf)
10828 .arg(&outf)
10829 .arg(&mi)
10830 .arg(&rb);
10831 unsafe {
10832 b.launch(cfg)?;
10833 }
10834 Ok(y)
10835 }
10836
10837 #[allow(clippy::too_many_arguments)]
10843 pub fn qmatvec_nvfp4_fast_prequant_into(
10844 &self,
10845 w: &CudaSlice<u8>,
10846 aq: &CudaSlice<i8>,
10847 ad: &CudaSlice<f32>,
10848 y: &mut CudaSlice<f32>,
10849 m: usize,
10850 in_f: usize,
10851 out_f: usize,
10852 row_bytes: usize,
10853 ) -> Result<(), Box<dyn std::error::Error>> {
10854 assert!(
10855 in_f.is_multiple_of(64),
10856 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10857 );
10858 if y.len() < m * out_f {
10859 return Err(format!(
10860 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
10861 y.len()
10862 )
10863 .into());
10864 }
10865 let f = self.func("qmatvec_nvfp4_dp4a");
10866 let cfg = LaunchConfig {
10867 grid_dim: (out_f as u32, m as u32, 1),
10868 block_dim: (128, 1, 1),
10869 shared_mem_bytes: 0,
10870 };
10871 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10872 let __s_b = self.gpu.stream();
10873 let mut b = __s_b.launch_builder(&f);
10874 b.arg(w)
10875 .arg(aq)
10876 .arg(ad)
10877 .arg(y)
10878 .arg(&inf)
10879 .arg(&outf)
10880 .arg(&mi)
10881 .arg(&rb);
10882 unsafe {
10883 b.launch(cfg)?;
10884 }
10885 Ok(())
10886 }
10887
10888 #[allow(clippy::too_many_arguments)]
10891 pub fn matvec_f32_qkv_into(
10892 &self,
10893 wq: &CudaSlice<f32>,
10894 wk: &CudaSlice<f32>,
10895 wv: &CudaSlice<f32>,
10896 wg: &CudaSlice<f32>,
10897 x: &CudaSlice<f32>,
10898 yq: &mut CudaSlice<f32>,
10899 yk: &mut CudaSlice<f32>,
10900 yv: &mut CudaSlice<f32>,
10901 yg: &mut CudaSlice<f32>,
10902 in_f: usize,
10903 out_q: usize,
10904 out_kv: usize,
10905 out_g: usize,
10906 ) -> Result<(), Box<dyn std::error::Error>> {
10907 if !in_f.is_multiple_of(4)
10908 || wq.len() != out_q * in_f
10909 || wk.len() != out_kv * in_f
10910 || wv.len() != out_kv * in_f
10911 || wg.len() < out_g * in_f
10912 || x.len() < in_f
10913 || yq.len() < out_q
10914 || yk.len() < out_kv
10915 || yv.len() < out_kv
10916 || (out_g > 0 && yg.len() < out_g)
10917 {
10918 return Err(format!(
10919 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
10920 wq={} wk={} wv={} wg={}",
10921 wq.len(),
10922 wk.len(),
10923 wv.len(),
10924 wg.len()
10925 )
10926 .into());
10927 }
10928 let f = self.func("matvec_f32_qkv");
10929 let cfg = LaunchConfig {
10930 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
10931 block_dim: (128, 1, 1),
10932 shared_mem_bytes: 0,
10933 };
10934 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
10935 let __s_b = self.gpu.stream();
10936 let mut b = __s_b.launch_builder(&f);
10937 b.arg(wq)
10938 .arg(wk)
10939 .arg(wv)
10940 .arg(wg)
10941 .arg(x)
10942 .arg(yq)
10943 .arg(yk)
10944 .arg(yv)
10945 .arg(yg)
10946 .arg(&inf)
10947 .arg(&oq)
10948 .arg(&okv)
10949 .arg(&og);
10950 unsafe {
10951 b.launch(cfg)?;
10952 }
10953 Ok(())
10954 }
10955
10956 #[allow(clippy::too_many_arguments)]
10968 pub fn qmatvec_nvfp4_sel_down8_into(
10969 &self,
10970 bank: &CudaSlice<u8>,
10971 sel: &CudaSlice<i32>,
10972 aq: &CudaSlice<i8>,
10973 ad: &CudaSlice<f32>,
10974 route_w: &CudaSlice<f32>,
10975 md: &CudaSlice<f32>,
10976 dst: &mut CudaSlice<f32>,
10977 n_sel: usize,
10978 in_f: usize,
10979 out_f: usize,
10980 row_bytes: usize,
10981 expert_stride: usize,
10982 act_row_stride: usize,
10983 ad_row_stride: usize,
10984 slot_major: bool,
10985 ) -> Result<(), Box<dyn std::error::Error>> {
10986 if !in_f.is_multiple_of(64)
10987 || n_sel == 0
10988 || n_sel > 8
10989 || (in_f >> 5) > 32
10990 || dst.len() < out_f
10991 || sel.len() < n_sel
10992 || route_w.len() < n_sel
10993 {
10994 return Err(format!(
10995 "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
10996 dst.len()
10997 )
10998 .into());
10999 }
11000 if !slot_major {
11001 return Err(
11002 "NVFP4 sel down8 reads slot-major rows: this shard is block_nvfp4 v1 \
11003 (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
11004 .into(),
11005 );
11006 }
11007 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
11008 let cfg = LaunchConfig {
11009 grid_dim: (out_f as u32, 1, 1),
11010 block_dim: (32, n_sel as u32, 1),
11011 shared_mem_bytes: 0,
11012 };
11013 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11014 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11015 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
11016 let __s_b = self.gpu.stream();
11017 let mut b = __s_b.launch_builder(&f);
11018 b.arg(bank)
11019 .arg(sel)
11020 .arg(aq)
11021 .arg(ad)
11022 .arg(route_w)
11023 .arg(md)
11024 .arg(dst)
11025 .arg(&inf)
11026 .arg(&outf)
11027 .arg(&ns)
11028 .arg(&rb)
11029 .arg(&es)
11030 .arg(&ars)
11031 .arg(&adrs);
11032 unsafe {
11033 b.launch(cfg)?;
11034 }
11035 Ok(())
11036 }
11037
11038 #[allow(clippy::too_many_arguments)] pub fn qmatvec_nvfp4_sel_into(
11053 &self,
11054 bank: &CudaSlice<u8>,
11055 sel: &CudaSlice<i32>,
11056 aq: &CudaSlice<i8>,
11057 ad: &CudaSlice<f32>,
11058 y: &mut CudaSlice<f32>,
11059 n_sel: usize,
11060 in_f: usize,
11061 out_f: usize,
11062 row_bytes: usize,
11063 expert_stride: usize,
11064 act_row_stride: usize,
11065 ad_row_stride: usize,
11066 slot_major: bool,
11067 ) -> Result<(), Box<dyn std::error::Error>> {
11068 assert!(
11069 in_f.is_multiple_of(64),
11070 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
11071 );
11072 if y.len() < n_sel * out_f || sel.len() < n_sel {
11073 return Err(format!(
11074 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
11075 y.len(),
11076 sel.len()
11077 )
11078 .into());
11079 }
11080 let mode: u8 = if slot_major { 3 } else { 0 };
11083 static SM_STREAM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11087 let sm_stream = mode == 3
11088 && *SM_STREAM
11089 .get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_SM_STREAM").as_deref() == Ok("1"))
11090 && row_bytes.is_multiple_of(16)
11091 && in_f <= 4096;
11092 let kname = match (mode, sm_stream) {
11093 (3, true) => "qmatvec_nvfp4_dp4a_sel_v2s",
11094 (3, false) => "qmatvec_nvfp4_dp4a_sel_v2",
11095 _ => "qmatvec_nvfp4_dp4a_sel",
11096 };
11097 {
11103 static SEEN_SEL: std::sync::Mutex<Vec<(&'static str, usize, usize)>> =
11104 std::sync::Mutex::new(Vec::new());
11105 let combo = (kname, in_f, out_f);
11106 let mut seen = SEEN_SEL.lock().unwrap();
11107 if !seen.contains(&combo) {
11108 seen.push(combo);
11109 eprintln!(
11110 "[nvfp4-sel] kernel={kname} slot_major={slot_major} in_f={in_f} \
11111 out_f={out_f} nsb={} row_bytes={row_bytes}",
11112 in_f >> 5
11113 );
11114 }
11115 }
11116 let f = self.func(kname);
11117 let nsb = in_f >> 5;
11122 let fit_block: u32 = if (mode == 0 || mode == 3) && !sm_stream && nsb <= 32 {
11123 32
11124 } else if mode == 1 {
11125 512
11126 } else {
11127 128
11128 };
11129 let cfg = LaunchConfig {
11130 grid_dim: (
11131 if sm_stream {
11132 (out_f as u32).div_ceil(8)
11133 } else {
11134 match mode {
11135 2 => (out_f as u32).div_ceil(16),
11136 1 => (out_f as u32).div_ceil(4),
11137 _ => out_f as u32,
11138 }
11139 },
11140 n_sel as u32,
11141 1,
11142 ),
11143 block_dim: (fit_block, 1, 1),
11144 shared_mem_bytes: 0,
11145 };
11146 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11147 let (rb, es, ars, adrs) = (
11148 row_bytes as i64,
11149 expert_stride as i64,
11150 act_row_stride as i64,
11151 ad_row_stride as i64,
11152 );
11153 let __s_b = self.gpu.stream();
11154 let mut b = __s_b.launch_builder(&f);
11155 b.arg(bank)
11156 .arg(sel)
11157 .arg(aq)
11158 .arg(ad)
11159 .arg(y)
11160 .arg(&inf)
11161 .arg(&outf)
11162 .arg(&ns)
11163 .arg(&rb)
11164 .arg(&es)
11165 .arg(&ars)
11166 .arg(&adrs);
11167 unsafe {
11168 b.launch(cfg)?;
11169 }
11170 Ok(())
11171 }
11172
11173 #[allow(clippy::too_many_arguments)]
11176 pub fn qmatvec_nvfp4_bf16_sel_dual_rows_into(
11177 &self,
11178 gate_bank: &CudaSlice<u8>,
11179 up_bank: &CudaSlice<u8>,
11180 sel: &CudaSlice<i32>,
11181 token_rows: &CudaSlice<i32>,
11182 x_bf16: &CudaSlice<u8>,
11183 gate_out: &mut CudaSlice<f32>,
11184 up_out: &mut CudaSlice<f32>,
11185 n_sel: usize,
11186 in_f: usize,
11187 out_f: usize,
11188 row_bytes: usize,
11189 expert_stride: usize,
11190 tokens: usize,
11191 ) -> Result<(), Box<dyn std::error::Error>> {
11192 if !in_f.is_multiple_of(64)
11193 || sel.len() < n_sel
11194 || token_rows.len() < n_sel
11195 || gate_out.len() < n_sel * out_f
11196 || up_out.len() < n_sel * out_f
11197 || x_bf16.len() < 2 * in_f * tokens
11198 {
11199 return Err(format!(
11200 "W4A16 NVFP4 dual selected rows geometry sel={} token_rows={} x={} gate={} up={} \
11201 n_sel={n_sel} tokens={tokens} in={in_f} out={out_f}",
11202 sel.len(),
11203 token_rows.len(),
11204 x_bf16.len(),
11205 gate_out.len(),
11206 up_out.len(),
11207 )
11208 .into());
11209 }
11210 let adjacent_rows = tokens > 1;
11211 let f = if adjacent_rows {
11212 self.func("qmatvec_nvfp4_bf16_sel_quad_rows")
11213 } else {
11214 self.func("qmatvec_nvfp4_bf16_sel_dual_rows")
11215 };
11216 let cfg = LaunchConfig {
11217 grid_dim: (
11218 if adjacent_rows {
11219 out_f.div_ceil(2) as u32
11220 } else {
11221 (2 * out_f) as u32
11222 },
11223 n_sel as u32,
11224 1,
11225 ),
11226 block_dim: (256, 1, 1),
11227 shared_mem_bytes: 0,
11228 };
11229 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11230 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11231 let __s_b = self.gpu.stream();
11232 let mut b = __s_b.launch_builder(&f);
11233 b.arg(gate_bank)
11234 .arg(up_bank)
11235 .arg(sel)
11236 .arg(token_rows)
11237 .arg(x_bf16)
11238 .arg(gate_out)
11239 .arg(up_out)
11240 .arg(&inf)
11241 .arg(&outf)
11242 .arg(&ns)
11243 .arg(&rb)
11244 .arg(&es);
11245 unsafe {
11246 b.launch(cfg)?;
11247 }
11248 Ok(())
11249 }
11250
11251 #[allow(clippy::too_many_arguments)]
11254 pub fn qmatvec_nvfp4_bf16_ep_dual_slots_into(
11255 &self,
11256 gate_bank: &CudaSlice<u8>,
11257 up_bank: &CudaSlice<u8>,
11258 sel: &CudaSlice<i32>,
11259 x_bf16: &CudaSlice<u8>,
11260 gate_out: &mut CudaSlice<f32>,
11261 up_out: &mut CudaSlice<f32>,
11262 n_pairs: usize,
11263 top_k: usize,
11264 in_f: usize,
11265 out_f: usize,
11266 owner_start: usize,
11267 owner_end: usize,
11268 row_bytes: usize,
11269 expert_stride: usize,
11270 ) -> Result<(), Box<dyn std::error::Error>> {
11271 let tokens = n_pairs.div_ceil(top_k);
11272 if top_k == 0
11273 || owner_start >= owner_end
11274 || !in_f.is_multiple_of(64)
11275 || sel.len() < n_pairs
11276 || gate_out.len() < n_pairs * out_f
11277 || up_out.len() < n_pairs * out_f
11278 || x_bf16.len() < 2 * in_f * tokens
11279 {
11280 return Err(format!(
11281 "W4A16 NVFP4 device EP dual-slot geometry sel={} x={} gate={} up={} \
11282 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11283 owner={owner_start}..{owner_end}",
11284 sel.len(),
11285 x_bf16.len(),
11286 gate_out.len(),
11287 up_out.len(),
11288 )
11289 .into());
11290 }
11291 let pair_parallel = tokens > 1;
11292 let f = if pair_parallel {
11293 self.func("qmatvec_nvfp4_bf16_ep_quad_pairs")
11294 } else {
11295 self.func("qmatvec_nvfp4_bf16_ep_dual_slots")
11296 };
11297 let cfg = LaunchConfig {
11298 grid_dim: (
11299 if pair_parallel {
11300 out_f.div_ceil(2) as u32
11301 } else {
11302 (2 * out_f) as u32
11303 },
11304 if pair_parallel { n_pairs as u32 } else { 1 },
11305 1,
11306 ),
11307 block_dim: (256, 1, 1),
11308 shared_mem_bytes: 0,
11309 };
11310 let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11311 let (os, oe) = (owner_start as i32, owner_end as i32);
11312 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11313 let __s_b = self.gpu.stream();
11314 let mut b = __s_b.launch_builder(&f);
11315 b.arg(gate_bank)
11316 .arg(up_bank)
11317 .arg(sel)
11318 .arg(x_bf16)
11319 .arg(gate_out)
11320 .arg(up_out)
11321 .arg(&inf)
11322 .arg(&outf)
11323 .arg(&np)
11324 .arg(&tk)
11325 .arg(&os)
11326 .arg(&oe)
11327 .arg(&rb)
11328 .arg(&es);
11329 unsafe {
11330 b.launch(cfg)?;
11331 }
11332 Ok(())
11333 }
11334
11335 #[allow(clippy::too_many_arguments)]
11337 pub fn qmatvec_nvfp4_q8_ep_dual_slots_into(
11338 &self,
11339 gate_bank: &CudaSlice<u8>,
11340 up_bank: &CudaSlice<u8>,
11341 sel: &CudaSlice<i32>,
11342 aq: &CudaSlice<i8>,
11343 ad: &CudaSlice<f32>,
11344 gate_out: &mut CudaSlice<f32>,
11345 up_out: &mut CudaSlice<f32>,
11346 n_pairs: usize,
11347 top_k: usize,
11348 in_f: usize,
11349 out_f: usize,
11350 owner_start: usize,
11351 owner_end: usize,
11352 row_bytes: usize,
11353 expert_stride: usize,
11354 ) -> Result<(), Box<dyn std::error::Error>> {
11355 let tokens = n_pairs.div_ceil(top_k);
11356 if top_k == 0
11357 || owner_start >= owner_end
11358 || !in_f.is_multiple_of(64)
11359 || sel.len() < n_pairs
11360 || aq.len() < tokens * in_f
11361 || ad.len() < tokens * (in_f / 32)
11362 || gate_out.len() < n_pairs * out_f
11363 || up_out.len() < n_pairs * out_f
11364 {
11365 return Err(format!(
11366 "W4A8 NVFP4 device EP gate/up geometry sel={} aq={} ad={} gate={} up={} \
11367 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11368 owner={owner_start}..{owner_end}",
11369 sel.len(),
11370 aq.len(),
11371 ad.len(),
11372 gate_out.len(),
11373 up_out.len(),
11374 )
11375 .into());
11376 }
11377 let f = self.func("qmatvec_nvfp4_q8_ep_dual_slots");
11378 let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
11379 let cfg = LaunchConfig {
11380 grid_dim: (out_f as u32, 1, 1),
11381 block_dim: (threads, 1, 1),
11382 shared_mem_bytes: 0,
11383 };
11384 let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11385 let (os, oe) = (owner_start as i32, owner_end as i32);
11386 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11387 let __s_b = self.gpu.stream();
11388 let mut b = __s_b.launch_builder(&f);
11389 b.arg(gate_bank)
11390 .arg(up_bank)
11391 .arg(sel)
11392 .arg(aq)
11393 .arg(ad)
11394 .arg(gate_out)
11395 .arg(up_out)
11396 .arg(&inf)
11397 .arg(&outf)
11398 .arg(&np)
11399 .arg(&tk)
11400 .arg(&os)
11401 .arg(&oe)
11402 .arg(&rb)
11403 .arg(&es);
11404 unsafe {
11405 b.launch(cfg)?;
11406 }
11407 Ok(())
11408 }
11409
11410 #[allow(clippy::too_many_arguments)]
11413 pub fn qmatvec_nvfp4_q8_ep_paired_slots_into(
11414 &self,
11415 gate_bank: &CudaSlice<u8>,
11416 up_bank: &CudaSlice<u8>,
11417 sel: &CudaSlice<i32>,
11418 aq: &CudaSlice<i8>,
11419 ad: &CudaSlice<f32>,
11420 gate_out: &mut CudaSlice<f32>,
11421 up_out: &mut CudaSlice<f32>,
11422 n_pairs: usize,
11423 top_k: usize,
11424 in_f: usize,
11425 out_f: usize,
11426 owner_start: usize,
11427 owner_end: usize,
11428 row_bytes: usize,
11429 expert_stride: usize,
11430 ) -> Result<(), Box<dyn std::error::Error>> {
11431 let tokens = n_pairs.div_ceil(top_k);
11432 if top_k == 0
11433 || owner_start >= owner_end
11434 || !in_f.is_multiple_of(64)
11435 || sel.len() < n_pairs
11436 || aq.len() < tokens * in_f
11437 || ad.len() < tokens * (in_f / 32)
11438 || gate_out.len() < n_pairs * out_f
11439 || up_out.len() < n_pairs * out_f
11440 {
11441 return Err(format!(
11442 "W4A8 NVFP4 paired gate/up geometry sel={} aq={} ad={} gate={} up={} \
11443 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11444 owner={owner_start}..{owner_end}",
11445 sel.len(),
11446 aq.len(),
11447 ad.len(),
11448 gate_out.len(),
11449 up_out.len(),
11450 )
11451 .into());
11452 }
11453 let f = self.func("qmatvec_nvfp4_q8_ep_paired_slots");
11454 let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
11455 let cfg = LaunchConfig {
11456 grid_dim: (out_f as u32, 1, 1),
11457 block_dim: (threads, 1, 1),
11458 shared_mem_bytes: 0,
11459 };
11460 let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11461 let (os, oe) = (owner_start as i32, owner_end as i32);
11462 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11463 let __s_b = self.gpu.stream();
11464 let mut b = __s_b.launch_builder(&f);
11465 b.arg(gate_bank)
11466 .arg(up_bank)
11467 .arg(sel)
11468 .arg(aq)
11469 .arg(ad)
11470 .arg(gate_out)
11471 .arg(up_out)
11472 .arg(&inf)
11473 .arg(&outf)
11474 .arg(&np)
11475 .arg(&tk)
11476 .arg(&os)
11477 .arg(&oe)
11478 .arg(&rb)
11479 .arg(&es);
11480 unsafe {
11481 b.launch(cfg)?;
11482 }
11483 Ok(())
11484 }
11485
11486 #[allow(clippy::too_many_arguments)]
11488 pub fn qmatvec_nvfp4_bf16_sel_down_rows_raw(
11489 &self,
11490 bank: &CudaSlice<u8>,
11491 sel: &CudaSlice<i32>,
11492 global_pairs: &CudaSlice<i32>,
11493 activation_bf16: &CudaSlice<u8>,
11494 macros_down: &CudaSlice<f32>,
11495 dst_raw: u64,
11496 n_sel: usize,
11497 in_f: usize,
11498 out_f: usize,
11499 row_bytes: usize,
11500 expert_stride: usize,
11501 total_pairs: usize,
11502 ) -> Result<(), Box<dyn std::error::Error>> {
11503 if !in_f.is_multiple_of(64)
11504 || sel.len() < n_sel
11505 || global_pairs.len() < n_sel
11506 || activation_bf16.len() < 2 * n_sel * in_f
11507 || dst_raw == 0
11508 {
11509 return Err(format!(
11510 "W4A16 NVFP4 down rows geometry sel={} pairs={} act={} dst_raw={dst_raw:#x} \
11511 n_sel={n_sel} total_pairs={total_pairs} in={in_f} out={out_f}",
11512 sel.len(),
11513 global_pairs.len(),
11514 activation_bf16.len(),
11515 )
11516 .into());
11517 }
11518 let f = self.func("qmatvec_nvfp4_bf16_sel_down_rows");
11519 let cfg = LaunchConfig {
11520 grid_dim: (out_f.div_ceil(2) as u32, n_sel as u32, 1),
11521 block_dim: (256, 1, 1),
11522 shared_mem_bytes: 0,
11523 };
11524 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11525 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11526 let __s_b = self.gpu.stream();
11527 let mut b = __s_b.launch_builder(&f);
11528 b.arg(bank)
11529 .arg(sel)
11530 .arg(global_pairs)
11531 .arg(activation_bf16)
11532 .arg(macros_down)
11533 .arg(&dst_raw)
11534 .arg(&inf)
11535 .arg(&outf)
11536 .arg(&ns)
11537 .arg(&rb)
11538 .arg(&es);
11539 unsafe {
11540 b.launch(cfg)?;
11541 }
11542 Ok(())
11543 }
11544
11545 #[allow(clippy::too_many_arguments)]
11548 pub fn qmatvec_nvfp4_bf16_ep_down_slots_raw(
11549 &self,
11550 bank: &CudaSlice<u8>,
11551 sel: &CudaSlice<i32>,
11552 activation_bf16: &CudaSlice<u8>,
11553 macros_down: &CudaSlice<f32>,
11554 dst_raw: u64,
11555 n_pairs: usize,
11556 in_f: usize,
11557 out_f: usize,
11558 owner_start: usize,
11559 owner_end: usize,
11560 row_bytes: usize,
11561 expert_stride: usize,
11562 ) -> Result<(), Box<dyn std::error::Error>> {
11563 if owner_start >= owner_end
11564 || !in_f.is_multiple_of(64)
11565 || sel.len() < n_pairs
11566 || activation_bf16.len() < 2 * n_pairs * in_f
11567 || dst_raw == 0
11568 {
11569 return Err(format!(
11570 "W4A16 NVFP4 device EP down-slot geometry sel={} act={} dst_raw={dst_raw:#x} \
11571 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
11572 sel.len(),
11573 activation_bf16.len(),
11574 )
11575 .into());
11576 }
11577 let f = self.func("qmatvec_nvfp4_bf16_ep_down_slots");
11578 let cfg = LaunchConfig {
11579 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
11580 block_dim: (256, 1, 1),
11581 shared_mem_bytes: 0,
11582 };
11583 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11584 let (os, oe) = (owner_start as i32, owner_end as i32);
11585 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11586 let __s_b = self.gpu.stream();
11587 let mut b = __s_b.launch_builder(&f);
11588 b.arg(bank)
11589 .arg(sel)
11590 .arg(activation_bf16)
11591 .arg(macros_down)
11592 .arg(&dst_raw)
11593 .arg(&inf)
11594 .arg(&outf)
11595 .arg(&np)
11596 .arg(&os)
11597 .arg(&oe)
11598 .arg(&rb)
11599 .arg(&es);
11600 unsafe {
11601 b.launch(cfg)?;
11602 }
11603 Ok(())
11604 }
11605
11606 #[allow(clippy::too_many_arguments)]
11608 pub fn qmatvec_nvfp4_bf16_ep_down_pairs_raw(
11609 &self,
11610 bank: &CudaSlice<u8>,
11611 sel: &CudaSlice<i32>,
11612 activation_bf16: &CudaSlice<u8>,
11613 macros_down: &CudaSlice<f32>,
11614 dst_raw: u64,
11615 n_pairs: usize,
11616 in_f: usize,
11617 out_f: usize,
11618 owner_start: usize,
11619 owner_end: usize,
11620 row_bytes: usize,
11621 expert_stride: usize,
11622 ) -> Result<(), Box<dyn std::error::Error>> {
11623 if owner_start >= owner_end
11624 || !in_f.is_multiple_of(64)
11625 || sel.len() < n_pairs
11626 || activation_bf16.len() < 2 * n_pairs * in_f
11627 || dst_raw == 0
11628 {
11629 return Err(format!(
11630 "W4A16 NVFP4 device EP down-pair geometry sel={} act={} dst_raw={dst_raw:#x} \
11631 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
11632 sel.len(),
11633 activation_bf16.len(),
11634 )
11635 .into());
11636 }
11637 let f = self.func("qmatvec_nvfp4_bf16_ep_down_pairs");
11638 let cfg = LaunchConfig {
11639 grid_dim: (out_f.div_ceil(2) as u32, n_pairs as u32, 1),
11640 block_dim: (256, 1, 1),
11641 shared_mem_bytes: 0,
11642 };
11643 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11644 let (os, oe) = (owner_start as i32, owner_end as i32);
11645 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11646 let __s_b = self.gpu.stream();
11647 let mut b = __s_b.launch_builder(&f);
11648 b.arg(bank)
11649 .arg(sel)
11650 .arg(activation_bf16)
11651 .arg(macros_down)
11652 .arg(&dst_raw)
11653 .arg(&inf)
11654 .arg(&outf)
11655 .arg(&np)
11656 .arg(&os)
11657 .arg(&oe)
11658 .arg(&rb)
11659 .arg(&es);
11660 unsafe {
11661 b.launch(cfg)?;
11662 }
11663 Ok(())
11664 }
11665
11666 #[allow(clippy::too_many_arguments)]
11668 pub fn silu_mul_scaled_host_expf_bf16_sel_into(
11669 &self,
11670 gate: &CudaSlice<f32>,
11671 up: &CudaSlice<f32>,
11672 gate_macros: &CudaSlice<f32>,
11673 up_macros: &CudaSlice<f32>,
11674 sel: &CudaSlice<i32>,
11675 limit: Option<f32>,
11676 output_bf16: &mut CudaSlice<u8>,
11677 n_per: usize,
11678 n_sel: usize,
11679 ) -> Result<(), Box<dyn std::error::Error>> {
11680 let n = n_per * n_sel;
11681 if sel.len() < n_sel || gate.len() < n || up.len() < n || output_bf16.len() < 2 * n {
11682 return Err(format!(
11683 "W4A16 selected activation geometry sel={} gate={} up={} out={} \
11684 n_per={n_per} n_sel={n_sel}",
11685 sel.len(),
11686 gate.len(),
11687 up.len(),
11688 output_bf16.len(),
11689 )
11690 .into());
11691 }
11692 let (limit, has_limit) = match limit {
11693 Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
11694 Some(limit) => {
11695 return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
11696 }
11697 None => (0.0f32, 0i32),
11698 };
11699 let f = self.func("silu_mul_scaled_host_expf_bf16_sel");
11700 let cfg = LaunchConfig::for_num_elems(n as u32);
11701 let (np, ns) = (n_per as i32, n_sel as i32);
11702 let __s_b = self.gpu.stream();
11703 let mut b = __s_b.launch_builder(&f);
11704 b.arg(gate)
11705 .arg(up)
11706 .arg(gate_macros)
11707 .arg(up_macros)
11708 .arg(sel)
11709 .arg(&limit)
11710 .arg(&has_limit)
11711 .arg(output_bf16)
11712 .arg(&np)
11713 .arg(&ns);
11714 unsafe {
11715 b.launch(cfg)?;
11716 }
11717 Ok(())
11718 }
11719
11720 #[allow(clippy::too_many_arguments)]
11723 pub fn silu_mul_scaled_host_expf_bf16_ep_slots_into(
11724 &self,
11725 gate: &CudaSlice<f32>,
11726 up: &CudaSlice<f32>,
11727 gate_macros: &CudaSlice<f32>,
11728 up_macros: &CudaSlice<f32>,
11729 sel: &CudaSlice<i32>,
11730 owner_start: usize,
11731 owner_end: usize,
11732 limit: Option<f32>,
11733 output_bf16: &mut CudaSlice<u8>,
11734 n_per: usize,
11735 n_pairs: usize,
11736 ) -> Result<(), Box<dyn std::error::Error>> {
11737 let n = n_per * n_pairs;
11738 if owner_start >= owner_end
11739 || sel.len() < n_pairs
11740 || gate.len() < n
11741 || up.len() < n
11742 || output_bf16.len() < 2 * n
11743 {
11744 return Err(format!(
11745 "W4A16 device EP activation geometry sel={} gate={} up={} out={} \
11746 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
11747 sel.len(),
11748 gate.len(),
11749 up.len(),
11750 output_bf16.len(),
11751 )
11752 .into());
11753 }
11754 let (limit, has_limit) = match limit {
11755 Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
11756 Some(limit) => {
11757 return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
11758 }
11759 None => (0.0f32, 0i32),
11760 };
11761 let f = self.func("silu_mul_scaled_host_expf_bf16_ep_slots");
11762 let cfg = LaunchConfig::for_num_elems(n as u32);
11763 let (np, pairs) = (n_per as i32, n_pairs as i32);
11764 let (os, oe) = (owner_start as i32, owner_end as i32);
11765 let __s_b = self.gpu.stream();
11766 let mut b = __s_b.launch_builder(&f);
11767 b.arg(gate)
11768 .arg(up)
11769 .arg(gate_macros)
11770 .arg(up_macros)
11771 .arg(sel)
11772 .arg(&limit)
11773 .arg(&has_limit)
11774 .arg(output_bf16)
11775 .arg(&np)
11776 .arg(&pairs)
11777 .arg(&os)
11778 .arg(&oe);
11779 unsafe {
11780 b.launch(cfg)?;
11781 }
11782 Ok(())
11783 }
11784
11785 #[allow(clippy::too_many_arguments)]
11787 pub fn silu_mul_scaled_host_expf_q8_ep_slots_into(
11788 &self,
11789 gate: &CudaSlice<f32>,
11790 up: &CudaSlice<f32>,
11791 gate_macros: &CudaSlice<f32>,
11792 up_macros: &CudaSlice<f32>,
11793 sel: &CudaSlice<i32>,
11794 owner_start: usize,
11795 owner_end: usize,
11796 limit: Option<f32>,
11797 output_q8: &mut CudaSlice<i8>,
11798 output_scales: &mut CudaSlice<f32>,
11799 n_per: usize,
11800 n_pairs: usize,
11801 ) -> Result<(), Box<dyn std::error::Error>> {
11802 let n = n_per * n_pairs;
11803 if owner_start >= owner_end
11804 || !n_per.is_multiple_of(32)
11805 || sel.len() < n_pairs
11806 || gate.len() < n
11807 || up.len() < n
11808 || output_q8.len() < n
11809 || output_scales.len() < n / 32
11810 {
11811 return Err(format!(
11812 "W4A8 device EP activation geometry sel={} gate={} up={} q8={} scales={} \
11813 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
11814 sel.len(),
11815 gate.len(),
11816 up.len(),
11817 output_q8.len(),
11818 output_scales.len(),
11819 )
11820 .into());
11821 }
11822 let (limit, has_limit) = match limit {
11823 Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
11824 Some(limit) => {
11825 return Err(format!("W4A8 selected activation limit {limit} is invalid").into());
11826 }
11827 None => (0.0f32, 0i32),
11828 };
11829 let f = self.func("silu_mul_scaled_host_expf_q8_ep_slots");
11830 let warps = n / 32;
11831 let cfg = LaunchConfig {
11832 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
11833 block_dim: (128, 1, 1),
11834 shared_mem_bytes: 0,
11835 };
11836 let (np, pairs) = (n_per as i32, n_pairs as i32);
11837 let (os, oe) = (owner_start as i32, owner_end as i32);
11838 let __s_b = self.gpu.stream();
11839 let mut b = __s_b.launch_builder(&f);
11840 b.arg(gate)
11841 .arg(up)
11842 .arg(gate_macros)
11843 .arg(up_macros)
11844 .arg(sel)
11845 .arg(&limit)
11846 .arg(&has_limit)
11847 .arg(output_q8)
11848 .arg(output_scales)
11849 .arg(&np)
11850 .arg(&pairs)
11851 .arg(&os)
11852 .arg(&oe);
11853 unsafe {
11854 b.launch(cfg)?;
11855 }
11856 Ok(())
11857 }
11858
11859 #[allow(clippy::too_many_arguments)]
11862 pub fn qmatvec_nvfp4_bf16_sel_down_fma_into(
11863 &self,
11864 bank: &CudaSlice<u8>,
11865 sel: &CudaSlice<i32>,
11866 activation_bf16: &CudaSlice<u8>,
11867 route_weights: &CudaSlice<f32>,
11868 macros_down: &CudaSlice<f32>,
11869 dst: &mut cudarc::driver::CudaViewMut<f32>,
11870 n_sel: usize,
11871 in_f: usize,
11872 out_f: usize,
11873 row_bytes: usize,
11874 expert_stride: usize,
11875 ) -> Result<(), Box<dyn std::error::Error>> {
11876 if !in_f.is_multiple_of(64)
11877 || sel.len() < n_sel
11878 || route_weights.len() < n_sel
11879 || activation_bf16.len() < 2 * n_sel * in_f
11880 || dst.len() < out_f
11881 {
11882 return Err(format!(
11883 "W4A16 NVFP4 down selected geometry sel={} act={} weights={} dst={} \
11884 n_sel={n_sel} in={in_f} out={out_f}",
11885 sel.len(),
11886 activation_bf16.len(),
11887 route_weights.len(),
11888 dst.len(),
11889 )
11890 .into());
11891 }
11892 let f = self.func("qmatvec_nvfp4_bf16_sel_down_fma");
11893 let cfg = LaunchConfig {
11894 grid_dim: (out_f as u32, 1, 1),
11895 block_dim: (256, 1, 1),
11896 shared_mem_bytes: 0,
11897 };
11898 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11899 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11900 let __s_b = self.gpu.stream();
11901 let mut b = __s_b.launch_builder(&f);
11902 b.arg(bank)
11903 .arg(sel)
11904 .arg(activation_bf16)
11905 .arg(route_weights)
11906 .arg(macros_down)
11907 .arg(dst)
11908 .arg(&inf)
11909 .arg(&outf)
11910 .arg(&ns)
11911 .arg(&rb)
11912 .arg(&es);
11913 unsafe {
11914 b.launch(cfg)?;
11915 }
11916 Ok(())
11917 }
11918
11919 #[allow(clippy::too_many_arguments)]
11921 pub fn qmatvec_nvfp4_bf16_ep_down_fma_into(
11922 &self,
11923 bank: &CudaSlice<u8>,
11924 sel: &CudaSlice<i32>,
11925 activation_bf16: &CudaSlice<u8>,
11926 route_weights: &CudaSlice<f32>,
11927 macros_down: &CudaSlice<f32>,
11928 dst: &mut cudarc::driver::CudaViewMut<f32>,
11929 n_pairs: usize,
11930 in_f: usize,
11931 out_f: usize,
11932 owner_start: usize,
11933 owner_end: usize,
11934 row_bytes: usize,
11935 expert_stride: usize,
11936 ) -> Result<(), Box<dyn std::error::Error>> {
11937 if owner_start >= owner_end
11938 || !in_f.is_multiple_of(64)
11939 || sel.len() < n_pairs
11940 || route_weights.len() < n_pairs
11941 || activation_bf16.len() < 2 * n_pairs * in_f
11942 || dst.len() < out_f
11943 {
11944 return Err(format!(
11945 "W4A16 device EP down-FMA geometry sel={} act={} weights={} dst={} \
11946 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
11947 sel.len(),
11948 activation_bf16.len(),
11949 route_weights.len(),
11950 dst.len(),
11951 )
11952 .into());
11953 }
11954 let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
11955 let cfg = LaunchConfig {
11956 grid_dim: (out_f as u32, 1, 1),
11957 block_dim: (256, 1, 1),
11958 shared_mem_bytes: 0,
11959 };
11960 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
11961 let (os, oe) = (owner_start as i32, owner_end as i32);
11962 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11963 let __s_b = self.gpu.stream();
11964 let mut b = __s_b.launch_builder(&f);
11965 b.arg(bank)
11966 .arg(sel)
11967 .arg(activation_bf16)
11968 .arg(route_weights)
11969 .arg(macros_down)
11970 .arg(dst)
11971 .arg(&inf)
11972 .arg(&outf)
11973 .arg(&np)
11974 .arg(&os)
11975 .arg(&oe)
11976 .arg(&rb)
11977 .arg(&es);
11978 unsafe {
11979 b.launch(cfg)?;
11980 }
11981 Ok(())
11982 }
11983
11984 #[allow(clippy::too_many_arguments)]
11987 pub fn qmatvec_nvfp4_bf16_ep_down_fma_raw(
11988 &self,
11989 bank: &CudaSlice<u8>,
11990 sel: &CudaSlice<i32>,
11991 activation_bf16: &CudaSlice<u8>,
11992 route_weights: &CudaSlice<f32>,
11993 macros_down: &CudaSlice<f32>,
11994 dst_raw: u64,
11995 n_pairs: usize,
11996 in_f: usize,
11997 out_f: usize,
11998 owner_start: usize,
11999 owner_end: usize,
12000 row_bytes: usize,
12001 expert_stride: usize,
12002 ) -> Result<(), Box<dyn std::error::Error>> {
12003 if dst_raw == 0
12004 || owner_start >= owner_end
12005 || !in_f.is_multiple_of(64)
12006 || sel.len() < n_pairs
12007 || route_weights.len() < n_pairs
12008 || activation_bf16.len() < 2 * n_pairs * in_f
12009 {
12010 return Err(format!(
12011 "W4A16 device EP raw down-FMA geometry sel={} act={} weights={} \
12012 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12013 owner={owner_start}..{owner_end}",
12014 sel.len(),
12015 activation_bf16.len(),
12016 route_weights.len(),
12017 )
12018 .into());
12019 }
12020 let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
12021 let cfg = LaunchConfig {
12022 grid_dim: (out_f as u32, 1, 1),
12023 block_dim: (256, 1, 1),
12024 shared_mem_bytes: 0,
12025 };
12026 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12027 let (os, oe) = (owner_start as i32, owner_end as i32);
12028 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12029 let __s_b = self.gpu.stream();
12030 let mut b = __s_b.launch_builder(&f);
12031 b.arg(bank)
12032 .arg(sel)
12033 .arg(activation_bf16)
12034 .arg(route_weights)
12035 .arg(macros_down)
12036 .arg(&dst_raw)
12037 .arg(&inf)
12038 .arg(&outf)
12039 .arg(&np)
12040 .arg(&os)
12041 .arg(&oe)
12042 .arg(&rb)
12043 .arg(&es);
12044 unsafe {
12045 b.launch(cfg)?;
12046 }
12047 Ok(())
12048 }
12049
12050 #[allow(clippy::too_many_arguments)]
12053 pub fn qmatvec_nvfp4_q8_ep_down_slots_raw(
12054 &self,
12055 bank: &CudaSlice<u8>,
12056 sel: &CudaSlice<i32>,
12057 aq: &CudaSlice<i8>,
12058 ad: &CudaSlice<f32>,
12059 macros_down: &CudaSlice<f32>,
12060 dst_raw: u64,
12061 n_pairs: usize,
12062 in_f: usize,
12063 out_f: usize,
12064 owner_start: usize,
12065 owner_end: usize,
12066 row_bytes: usize,
12067 expert_stride: usize,
12068 ) -> Result<(), Box<dyn std::error::Error>> {
12069 if dst_raw == 0
12070 || owner_start >= owner_end
12071 || !in_f.is_multiple_of(64)
12072 || sel.len() < n_pairs
12073 || aq.len() < n_pairs * in_f
12074 || ad.len() < n_pairs * (in_f / 32)
12075 {
12076 return Err(format!(
12077 "W4A8 device EP raw down-slot geometry sel={} aq={} ad={} \
12078 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12079 owner={owner_start}..{owner_end}",
12080 sel.len(),
12081 aq.len(),
12082 ad.len(),
12083 )
12084 .into());
12085 }
12086 let f = self.func("qmatvec_nvfp4_q8_ep_down_slots");
12087 let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
12088 let cfg = LaunchConfig {
12089 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
12090 block_dim: (threads, 1, 1),
12091 shared_mem_bytes: 0,
12092 };
12093 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12094 let (os, oe) = (owner_start as i32, owner_end as i32);
12095 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12096 let __s_b = self.gpu.stream();
12097 let mut b = __s_b.launch_builder(&f);
12098 b.arg(bank)
12099 .arg(sel)
12100 .arg(aq)
12101 .arg(ad)
12102 .arg(macros_down)
12103 .arg(&dst_raw)
12104 .arg(&inf)
12105 .arg(&outf)
12106 .arg(&np)
12107 .arg(&os)
12108 .arg(&oe)
12109 .arg(&rb)
12110 .arg(&es);
12111 unsafe {
12112 b.launch(cfg)?;
12113 }
12114 Ok(())
12115 }
12116
12117 #[allow(clippy::too_many_arguments)]
12119 pub fn qmatvec_nvfp4_q8_ep_down_fma_raw(
12120 &self,
12121 bank: &CudaSlice<u8>,
12122 sel: &CudaSlice<i32>,
12123 aq: &CudaSlice<i8>,
12124 ad: &CudaSlice<f32>,
12125 route_weights: &CudaSlice<f32>,
12126 macros_down: &CudaSlice<f32>,
12127 dst_raw: u64,
12128 n_pairs: usize,
12129 in_f: usize,
12130 out_f: usize,
12131 owner_start: usize,
12132 owner_end: usize,
12133 row_bytes: usize,
12134 expert_stride: usize,
12135 ) -> Result<(), Box<dyn std::error::Error>> {
12136 if dst_raw == 0
12137 || owner_start >= owner_end
12138 || !in_f.is_multiple_of(64)
12139 || sel.len() < n_pairs
12140 || aq.len() < n_pairs * in_f
12141 || ad.len() < n_pairs * (in_f / 32)
12142 || route_weights.len() < n_pairs
12143 {
12144 return Err(format!(
12145 "W4A8 device EP raw down-FMA geometry sel={} aq={} ad={} weights={} \
12146 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12147 owner={owner_start}..{owner_end}",
12148 sel.len(),
12149 aq.len(),
12150 ad.len(),
12151 route_weights.len(),
12152 )
12153 .into());
12154 }
12155 let f = self.func("qmatvec_nvfp4_q8_ep_down_fma");
12156 let cfg = LaunchConfig {
12157 grid_dim: (out_f as u32, 1, 1),
12158 block_dim: (256, 1, 1),
12159 shared_mem_bytes: 0,
12160 };
12161 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12162 let (os, oe) = (owner_start as i32, owner_end as i32);
12163 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12164 let __s_b = self.gpu.stream();
12165 let mut b = __s_b.launch_builder(&f);
12166 b.arg(bank)
12167 .arg(sel)
12168 .arg(aq)
12169 .arg(ad)
12170 .arg(route_weights)
12171 .arg(macros_down)
12172 .arg(&dst_raw)
12173 .arg(&inf)
12174 .arg(&outf)
12175 .arg(&np)
12176 .arg(&os)
12177 .arg(&oe)
12178 .arg(&rb)
12179 .arg(&es);
12180 unsafe {
12181 b.launch(cfg)?;
12182 }
12183 Ok(())
12184 }
12185
12186 #[allow(clippy::too_many_arguments)]
12191 pub fn silu_mul_scaled_q8_1_sel_into(
12192 &self,
12193 gate: &CudaSlice<f32>,
12194 up: &CudaSlice<f32>,
12195 gmac: &CudaSlice<f32>,
12196 umac: &CudaSlice<f32>,
12197 sel: &CudaSlice<i32>,
12198 limit: Option<f32>,
12199 out_q: &mut CudaSlice<i8>,
12200 out_d: &mut CudaSlice<f32>,
12201 n_per: usize,
12202 n_sel: usize,
12203 ) -> Result<(), Box<dyn std::error::Error>> {
12204 let n = n_per * n_sel;
12205 if !n_per.is_multiple_of(32) || out_q.len() < n || out_d.len() < n / 32 {
12206 return Err(format!(
12207 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
12208 out_q.len(),
12209 out_d.len()
12210 )
12211 .into());
12212 }
12213 if let Some(limit) = limit {
12214 if limit <= 1e-6 {
12215 return Err(format!(
12216 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
12217 )
12218 .into());
12219 }
12220 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
12221 let cfg = LaunchConfig::for_num_elems(n as u32);
12222 let (np, ns) = (n_per as i32, n_sel as i32);
12223 let __s_b = self.gpu.stream();
12224 let mut b = __s_b.launch_builder(&f);
12225 b.arg(gate)
12226 .arg(up)
12227 .arg(gmac)
12228 .arg(umac)
12229 .arg(sel)
12230 .arg(&limit)
12231 .arg(out_q)
12232 .arg(out_d)
12233 .arg(&np)
12234 .arg(&ns);
12235 unsafe {
12236 b.launch(cfg)?;
12237 }
12238 return Ok(());
12239 }
12240 let f = self.func("silu_mul_scaled_q8_1_sel");
12241 let cfg = LaunchConfig::for_num_elems(n as u32);
12242 let (np, ns) = (n_per as i32, n_sel as i32);
12243 let __s_b = self.gpu.stream();
12244 let mut b = __s_b.launch_builder(&f);
12245 b.arg(gate)
12246 .arg(up)
12247 .arg(gmac)
12248 .arg(umac)
12249 .arg(sel)
12250 .arg(out_q)
12251 .arg(out_d)
12252 .arg(&np)
12253 .arg(&ns);
12254 unsafe {
12255 b.launch(cfg)?;
12256 }
12257 Ok(())
12258 }
12259
12260 #[track_caller]
12261 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12262 crate::alloc_trace_hit(v.len() * 4);
12263 Ok(self.gpu.stream().clone_htod(v)?)
12264 }
12265 #[track_caller]
12266 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
12267 crate::alloc_trace_hit(v.len() * 4);
12268 Ok(self.gpu.stream().clone_htod(v)?)
12269 }
12270
12271 pub fn htod_u16(&self, v: &[u16]) -> Result<CudaSlice<u16>, Box<dyn std::error::Error>> {
12272 crate::alloc_trace_hit(v.len() * 4);
12273 Ok(self.gpu.stream().clone_htod(v)?)
12274 }
12275 #[track_caller]
12277 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
12278 crate::alloc_trace_hit(v.len());
12279 Ok(self.gpu.stream().clone_htod(v)?)
12280 }
12281 #[track_caller]
12282 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
12283 crate::alloc_trace_hit(v.len() * 8);
12284 Ok(self.gpu.stream().clone_htod(v)?)
12285 }
12286 #[track_caller]
12288 pub fn dtoh_view(
12289 &self,
12290 d: &cudarc::driver::CudaView<f32>,
12291 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
12292 crate::dtoh_trace_hit(d.len() * 4);
12293 let v = self.gpu.stream().clone_dtoh(d)?;
12294 self.gpu.stream().synchronize()?;
12295 Ok(v)
12296 }
12297 #[track_caller]
12298 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
12299 crate::dtoh_trace_hit(d.len() * 4);
12300 let v = self.gpu.stream().clone_dtoh(d)?;
12301 self.gpu.stream().synchronize()?;
12302 Ok(v)
12303 }
12304 #[track_caller]
12308 pub fn dtoh_pair(
12309 &self,
12310 a: &CudaSlice<f32>,
12311 b: &CudaSlice<f32>,
12312 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12313 crate::dtoh_trace_hit((a.len() + b.len()) * 4);
12314 let av = self.gpu.stream().clone_dtoh(a)?;
12315 let bv = self.gpu.stream().clone_dtoh(b)?;
12316 self.gpu.stream().synchronize()?;
12317 Ok((av, bv))
12318 }
12319 #[track_caller]
12322 pub fn dtoh_pair_views(
12323 &self,
12324 a: &cudarc::driver::CudaView<f32>,
12325 b: &cudarc::driver::CudaView<f32>,
12326 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12327 crate::dtoh_trace_hit((a.len() + b.len()) * 4);
12328 let av = self.gpu.stream().clone_dtoh(a)?;
12329 let bv = self.gpu.stream().clone_dtoh(b)?;
12330 self.gpu.stream().synchronize()?;
12331 Ok((av, bv))
12332 }
12333 #[track_caller]
12335 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
12336 crate::dtoh_trace_hit(d.len() * 4);
12337 let v = self.gpu.stream().clone_dtoh(d)?;
12338 self.gpu.stream().synchronize()?;
12339 Ok(v)
12340 }
12341 #[track_caller]
12345 pub fn dtoh_i8(&self, d: &CudaSlice<i8>) -> Result<Vec<i8>, Box<dyn std::error::Error>> {
12346 crate::dtoh_trace_hit(d.len());
12347 let v = self.gpu.stream().clone_dtoh(d)?;
12348 self.gpu.stream().synchronize()?;
12349 Ok(v)
12350 }
12351 #[track_caller]
12352 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
12353 crate::dtoh_trace_hit(d.len());
12354 let v = self.gpu.stream().clone_dtoh(d)?;
12355 self.gpu.stream().synchronize()?;
12356 Ok(v)
12357 }
12358 #[track_caller]
12359 pub fn dtoh_u8_view(
12360 &self,
12361 d: &cudarc::driver::CudaView<u8>,
12362 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
12363 crate::dtoh_trace_hit(d.len());
12364 let v = self.gpu.stream().clone_dtoh(d)?;
12365 self.gpu.stream().synchronize()?;
12366 Ok(v)
12367 }
12368 #[track_caller]
12376 pub fn dtoh_u8_into_pinned(
12377 &self,
12378 d: &CudaSlice<u8>,
12379 dst: &mut PinnedHostBuf,
12380 n: usize,
12381 ) -> Result<(), Box<dyn std::error::Error>> {
12382 crate::dtoh_trace_hit(d.len());
12383 if n > d.len() || n > dst.len() {
12384 return Err(format!(
12385 "dtoh_u8_into_pinned range {n} exceeds src {} or pinned dst {}",
12386 d.len(),
12387 dst.len(),
12388 )
12389 .into());
12390 }
12391 if n == 0 {
12392 return Ok(());
12393 }
12394 let host = &mut dst.as_mut_slice()[..n];
12395 self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
12396 self.gpu.stream().synchronize()?;
12397 Ok(())
12398 }
12399 #[track_caller]
12403 pub fn dtoh_f32_into_pinned(
12404 &self,
12405 d: &CudaSlice<f32>,
12406 dst: &mut PinnedHostBuf,
12407 n: usize,
12408 ) -> Result<(), Box<dyn std::error::Error>> {
12409 crate::dtoh_trace_hit(d.len() * 4);
12410 if n > d.len() || n * std::mem::size_of::<f32>() > dst.len() {
12411 return Err(format!(
12412 "dtoh_f32_into_pinned range {n} floats exceeds src {} or pinned dst {} bytes",
12413 d.len(),
12414 dst.len(),
12415 )
12416 .into());
12417 }
12418 if n == 0 {
12419 return Ok(());
12420 }
12421 let host: &mut [f32] = unsafe {
12423 std::slice::from_raw_parts_mut(dst.as_mut_slice().as_mut_ptr() as *mut f32, n)
12424 };
12425 self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
12426 self.gpu.stream().synchronize()?;
12427 Ok(())
12428 }
12429
12430 pub fn htod_f32_from_pinned_async(
12438 &self,
12439 src: &PinnedHostBuf,
12440 n: usize,
12441 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12442 if n * std::mem::size_of::<f32>() > src.len() {
12443 return Err(format!(
12444 "htod_f32_from_pinned_async range {n} floats exceeds pinned src {} bytes",
12445 src.len(),
12446 )
12447 .into());
12448 }
12449 let mut d = self.uninit(n)?;
12450 if n == 0 {
12451 return Ok(d);
12452 }
12453 let host: &[f32] =
12455 unsafe { std::slice::from_raw_parts(src.as_slice().as_ptr() as *const f32, n) };
12456 let stream = self.gpu.stream();
12457 let s: &CudaStream = &stream;
12458 {
12459 let (pd, _guard) = d.device_ptr_mut(s);
12460 unsafe { cudarc::driver::result::memcpy_htod_async(pd, host, s.cu_stream())? };
12464 }
12465 Ok(d)
12466 }
12467
12468 #[allow(clippy::too_many_arguments)]
12475 pub fn copy_2d_dtod_async(
12478 &self,
12479 dst: &mut CudaSlice<f32>,
12480 dst_off_floats: usize,
12481 dst_pitch_floats: usize,
12482 src: &CudaSlice<f32>,
12483 src_pitch_floats: usize,
12484 width_floats: usize,
12485 height: usize,
12486 ) -> Result<(), Box<dyn std::error::Error>> {
12487 if height == 0 || width_floats == 0 {
12488 return Ok(());
12489 }
12490 let f = std::mem::size_of::<f32>();
12491 let dst_need = dst_off_floats + (height - 1) * dst_pitch_floats + width_floats;
12492 let src_need = (height - 1) * src_pitch_floats + width_floats;
12493 if width_floats > src_pitch_floats
12494 || width_floats > dst_pitch_floats
12495 || dst_need > dst.len()
12496 || src_need > src.len()
12497 {
12498 return Err(format!(
12499 "copy_2d_dtod_async out of bounds: dst needs {dst_need} of {}, src needs \
12500 {src_need} of {}, width {width_floats} pitches {src_pitch_floats}/{dst_pitch_floats}",
12501 dst.len(),
12502 src.len(),
12503 )
12504 .into());
12505 }
12506 let stream = self.gpu.stream();
12507 let s: &CudaStream = &stream;
12508 let (sp, _g0) = src.device_ptr(s);
12509 let (dp, _g1) = dst.device_ptr_mut(s);
12510 let copy = cudarc::driver::sys::CUDA_MEMCPY2D_st {
12511 srcXInBytes: 0,
12512 srcY: 0,
12513 srcMemoryType: cudarc::driver::sys::CUmemorytype_enum::CU_MEMORYTYPE_DEVICE,
12514 srcHost: std::ptr::null(),
12515 srcDevice: sp,
12516 srcArray: std::ptr::null_mut(),
12517 srcPitch: src_pitch_floats * f,
12518 dstXInBytes: dst_off_floats * f,
12519 dstY: 0,
12520 dstMemoryType: cudarc::driver::sys::CUmemorytype_enum::CU_MEMORYTYPE_DEVICE,
12521 dstHost: std::ptr::null_mut(),
12522 dstDevice: dp,
12523 dstArray: std::ptr::null_mut(),
12524 dstPitch: dst_pitch_floats * f,
12525 WidthInBytes: width_floats * f,
12526 Height: height,
12527 };
12528 unsafe { cudarc::driver::sys::cuMemcpy2DAsync_v2(©, s.cu_stream()).result()? };
12531 Ok(())
12532 }
12533
12534 pub fn copy_peer_from_async(
12541 &self,
12542 dst: &mut CudaSlice<f32>,
12543 src_engine: &Engine,
12544 src: &CudaSlice<f32>,
12545 n: usize,
12546 ) -> Result<(), Box<dyn std::error::Error>> {
12547 if n > src.len() || n > dst.len() {
12548 return Err(format!(
12549 "copy_peer_from_async range {n} exceeds src {} or dst {}",
12550 src.len(),
12551 dst.len(),
12552 )
12553 .into());
12554 }
12555 if n == 0 {
12556 return Ok(());
12557 }
12558 let stream = src_engine.gpu.stream();
12559 let s_src: &CudaStream = &stream;
12560 let (sp, _g0) = src.device_ptr(s_src);
12561 let (dp, _g1) = dst.device_ptr_mut(s_src);
12562 src_engine.ctx().bind_to_thread()?;
12563 let r = unsafe {
12566 cudarc::driver::result::memcpy_peer_async(
12567 self.ctx().cu_ctx(),
12568 dp,
12569 src_engine.ctx().cu_ctx(),
12570 sp,
12571 n * std::mem::size_of::<f32>(),
12572 s_src.cu_stream(),
12573 )
12574 };
12575 self.ctx().bind_to_thread()?;
12576 r?;
12577 Ok(())
12578 }
12579
12580 pub fn free_mem_mb(&self) -> u64 {
12582 self.ctx()
12583 .mem_get_info()
12584 .map(|(f, _)| f as u64 >> 20)
12585 .unwrap_or(0)
12586 }
12587
12588 #[track_caller]
12589 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12590 crate::alloc_trace_hit(n * 4);
12591 SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12592 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
12593 self.keep_if_capturing(&s);
12594 Ok(s)
12595 }
12596
12597 pub(crate) fn hyper_ws_take(&self) -> Option<crate::hyper::HyperDecodeWs> {
12601 self.hyper_decode_ws.lock().unwrap().take()
12602 }
12603
12604 pub(crate) fn hyper_ws_put(&self, ws: crate::hyper::HyperDecodeWs) {
12605 *self.hyper_decode_ws.lock().unwrap() = Some(ws);
12606 }
12607
12608 pub(crate) fn mla_seg_ws_take(
12611 &self,
12612 nh: usize,
12613 dn: usize,
12614 dr: usize,
12615 r: usize,
12616 q_lora: usize,
12617 ) -> Result<crate::hybrid_forward::MlaSegWs, Box<dyn std::error::Error>> {
12618 let held = self.mla_seg_ws.lock().unwrap().take();
12619 match held {
12620 Some(ws) if ws.sig == (nh, dn, dr, r, q_lora) => Ok(ws),
12621 _ => crate::hybrid_forward::MlaSegWs::new(self, nh, dn, dr, r, q_lora),
12622 }
12623 }
12624
12625 pub(crate) fn mla_seg_ws_put(&self, ws: crate::hybrid_forward::MlaSegWs) {
12626 *self.mla_seg_ws.lock().unwrap() = Some(ws);
12627 }
12628
12629 pub(crate) fn mla_seg_ws_on() -> bool {
12634 std::env::var("MEMRA_MLA_SEG_WS").as_deref() == Ok("1")
12635 }
12636
12637 pub(crate) fn vws_uninit(
12646 &self,
12647 n: usize,
12648 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12649 if verify_ws_on() {
12650 let mut ws = self.verify_ws.lock().unwrap();
12651 let ws = &mut *ws;
12652 if let Some(s) = VerifyWs::take(&mut ws.f32_pool, &mut ws.held_bytes, n) {
12653 if VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
12654 eprintln!(
12655 "[glm5-verify-ws] engaged: verify-walk buffers recycling through \
12656 the size-keyed pool (MEMRA_GLM5_VERIFY_WS=1)"
12657 );
12658 }
12659 return Ok(s);
12660 }
12661 }
12662 self.alloc_uninit::<f32>(n)
12663 }
12664
12665 pub(crate) fn vws_uninit_i8(
12667 &self,
12668 n: usize,
12669 ) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
12670 if verify_ws_on() {
12671 let mut ws = self.verify_ws.lock().unwrap();
12672 let ws = &mut *ws;
12673 if let Some(s) = VerifyWs::take(&mut ws.i8_pool, &mut ws.held_bytes, n) {
12674 VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12675 return Ok(s);
12676 }
12677 }
12678 self.alloc_uninit::<i8>(n)
12679 }
12680
12681 pub(crate) fn vws_uninit_u64(
12683 &self,
12684 n: usize,
12685 ) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
12686 if verify_ws_on() {
12687 let mut ws = self.verify_ws.lock().unwrap();
12688 let ws = &mut *ws;
12689 if let Some(s) = VerifyWs::take(&mut ws.u64_pool, &mut ws.held_bytes, n) {
12690 VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12691 return Ok(s);
12692 }
12693 }
12694 self.alloc_uninit::<u64>(n)
12695 }
12696
12697 #[allow(clippy::type_complexity)] pub(crate) fn glm5_graph_keep(&self) -> &Mutex<Vec<Box<dyn std::any::Any + Send>>> {
12703 &self.capture_keep
12704 }
12705
12706 pub(crate) fn vws_recycle(&self, s: CudaSlice<f32>) {
12709 if glm5_graph_capture_open() {
12710 self.capture_keep.lock().unwrap().push(Box::new(s));
12711 return;
12712 }
12713 if verify_ws_on() {
12714 let mut ws = self.verify_ws.lock().unwrap();
12715 let ws = &mut *ws;
12716 VerifyWs::put(&mut ws.f32_pool, &mut ws.held_bytes, s);
12717 }
12718 }
12719
12720 pub(crate) fn vws_recycle_i8(&self, s: CudaSlice<i8>) {
12722 if glm5_graph_capture_open() {
12723 self.capture_keep.lock().unwrap().push(Box::new(s));
12724 return;
12725 }
12726 if verify_ws_on() {
12727 let mut ws = self.verify_ws.lock().unwrap();
12728 let ws = &mut *ws;
12729 VerifyWs::put(&mut ws.i8_pool, &mut ws.held_bytes, s);
12730 }
12731 }
12732
12733 pub(crate) fn vws_recycle_u64(&self, s: CudaSlice<u64>) {
12735 if glm5_graph_capture_open() {
12736 self.capture_keep.lock().unwrap().push(Box::new(s));
12737 return;
12738 }
12739 if verify_ws_on() {
12740 let mut ws = self.verify_ws.lock().unwrap();
12741 let ws = &mut *ws;
12742 VerifyWs::put(&mut ws.u64_pool, &mut ws.held_bytes, s);
12743 }
12744 }
12745
12746 pub fn prob_of_token_device(
12755 &self,
12756 logits: &CudaSlice<f32>,
12757 tok: &CudaSlice<u32>,
12758 n_vocab: usize,
12759 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12760 let nb = ARGMAX_NB;
12761 let mut part = self.alloc_uninit::<f32>(nb)?;
12762 let mut p = self.alloc_uninit::<f32>(1)?;
12763 let f1 = self.func("prob_of_token_partial_f32");
12764 let cfg1 = LaunchConfig {
12765 grid_dim: (nb as u32, 1, 1),
12766 block_dim: (256, 1, 1),
12767 shared_mem_bytes: 0,
12768 };
12769 let nv = n_vocab as i32;
12770 let __s_b1 = self.gpu.stream();
12771 let mut b1 = __s_b1.launch_builder(&f1);
12772 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
12773 unsafe {
12774 b1.launch(cfg1)?;
12775 }
12776 let f2 = self.func("prob_of_token_final_f32");
12777 let cfg2 = LaunchConfig {
12778 grid_dim: (1, 1, 1),
12779 block_dim: (256, 1, 1),
12780 shared_mem_bytes: 0,
12781 };
12782 let nbi = nb as i32;
12783 let __s_b2 = self.gpu.stream();
12784 let mut b2 = __s_b2.launch_builder(&f2);
12785 b2.arg(&part).arg(&mut p).arg(&nbi);
12786 unsafe {
12787 b2.launch(cfg2)?;
12788 }
12789 Ok(p)
12790 }
12791
12792 pub fn prob_of_token_device_col(
12799 &self,
12800 logits: &CudaSlice<f32>,
12801 tok_all: &CudaSlice<u32>,
12802 tok_idx: usize,
12803 p_out: &mut CudaSlice<f32>,
12804 p_idx: usize,
12805 n_vocab: usize,
12806 ) -> Result<(), Box<dyn std::error::Error>> {
12807 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
12808 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
12809 let nb = ARGMAX_NB;
12810 let mut part = self.alloc_uninit::<f32>(nb)?;
12811 let f1 = self.func("prob_of_token_partial_f32");
12812 let cfg1 = LaunchConfig {
12813 grid_dim: (nb as u32, 1, 1),
12814 block_dim: (256, 1, 1),
12815 shared_mem_bytes: 0,
12816 };
12817 let nv = n_vocab as i32;
12818 let __s_b1 = self.gpu.stream();
12819 let mut b1 = __s_b1.launch_builder(&f1);
12820 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
12821 unsafe {
12822 b1.launch(cfg1)?;
12823 }
12824 let f2 = self.func("prob_of_token_final_f32");
12825 let cfg2 = LaunchConfig {
12826 grid_dim: (1, 1, 1),
12827 block_dim: (256, 1, 1),
12828 shared_mem_bytes: 0,
12829 };
12830 let nbi = nb as i32;
12831 let __s_b2 = self.gpu.stream();
12832 let mut b2 = __s_b2.launch_builder(&f2);
12833 b2.arg(&part).arg(&mut p_v).arg(&nbi);
12834 unsafe {
12835 b2.launch(cfg2)?;
12836 }
12837 Ok(())
12838 }
12839
12840 pub fn prob_of_token_device_into(
12841 &self,
12842 logits: &CudaSlice<f32>,
12843 tok: &CudaSlice<u32>,
12844 p_out: &mut CudaSlice<f32>,
12845 n_vocab: usize,
12846 ) -> Result<(), Box<dyn std::error::Error>> {
12847 let nb = ARGMAX_NB;
12848 let mut part = self.alloc_uninit::<f32>(nb)?;
12849 let f1 = self.func("prob_of_token_partial_f32");
12850 let cfg1 = LaunchConfig {
12851 grid_dim: (nb as u32, 1, 1),
12852 block_dim: (256, 1, 1),
12853 shared_mem_bytes: 0,
12854 };
12855 let nv = n_vocab as i32;
12856 let __s_b1 = self.gpu.stream();
12857 let mut b1 = __s_b1.launch_builder(&f1);
12858 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
12859 unsafe {
12860 b1.launch(cfg1)?;
12861 }
12862 let f2 = self.func("prob_of_token_final_f32");
12863 let cfg2 = LaunchConfig {
12864 grid_dim: (1, 1, 1),
12865 block_dim: (256, 1, 1),
12866 shared_mem_bytes: 0,
12867 };
12868 let nbi = nb as i32;
12869 let __s_b2 = self.gpu.stream();
12870 let mut b2 = __s_b2.launch_builder(&f2);
12871 b2.arg(&part).arg(p_out).arg(&nbi);
12872 unsafe {
12873 b2.launch(cfg2)?;
12874 }
12875 Ok(())
12876 }
12877
12878 pub fn u32_hist_append(
12881 &self,
12882 tok: &CudaSlice<u32>,
12883 hist: &mut CudaSlice<u32>,
12884 idx: &mut CudaSlice<i32>,
12885 ) -> Result<(), Box<dyn std::error::Error>> {
12886 let f = self.func("u32_hist_append");
12887 let cfg = LaunchConfig {
12888 grid_dim: (1, 1, 1),
12889 block_dim: (32, 1, 1),
12890 shared_mem_bytes: 0,
12891 };
12892 let __s_b = self.gpu.stream();
12893 let mut b = __s_b.launch_builder(&f);
12894 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
12895 unsafe {
12896 b.launch(cfg)?;
12897 }
12898 Ok(())
12899 }
12900
12901 pub fn argmax_token_device(
12902 &self,
12903 logits: &CudaSlice<f32>,
12904 n_vocab: usize,
12905 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
12906 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
12907 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
12908 Ok(tok)
12909 }
12910 pub fn argmax_token_device_into(
12917 &self,
12918 logits: &CudaSlice<f32>,
12919 tok: &mut CudaSlice<u32>,
12920 n_vocab: usize,
12921 ) -> Result<(), Box<dyn std::error::Error>> {
12922 let nb = ARGMAX_NB;
12923 let f1 = self.func("argmax_partial_f32");
12924 let f2 = self.func("argmax_final_f32");
12925 let mut guard = self.argmax_partials.lock().unwrap();
12926 if guard.is_none() {
12927 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
12930 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
12931 *guard = Some((pv, pi));
12932 }
12933 let (part_v, part_i) = guard.as_mut().unwrap();
12934 let nv = n_vocab as i32;
12935 let nbi = nb as i32;
12936 let cfg1 = LaunchConfig {
12938 grid_dim: (nb as u32, 1, 1),
12939 block_dim: (256, 1, 1),
12940 shared_mem_bytes: 0,
12941 };
12942 let __s_b1 = self.gpu.stream();
12943 let mut b1 = __s_b1.launch_builder(&f1);
12944 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
12945 unsafe {
12946 b1.launch(cfg1)?;
12947 }
12948 let cfg2 = LaunchConfig {
12950 grid_dim: (1, 1, 1),
12951 block_dim: (256, 1, 1),
12952 shared_mem_bytes: 0,
12953 };
12954 let __s_b2 = self.gpu.stream();
12955 let mut b2 = __s_b2.launch_builder(&f2);
12956 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
12957 unsafe {
12958 b2.launch(cfg2)?;
12959 }
12960 Ok(())
12961 }
12962 pub fn argmax_token_device_col(
12968 &self,
12969 logits: &CudaSlice<f32>,
12970 col: usize,
12971 n_vocab: usize,
12972 toks: &mut CudaSlice<u32>,
12973 out_idx: usize,
12974 ) -> Result<(), Box<dyn std::error::Error>> {
12975 let nb = ARGMAX_NB;
12976 let f1 = self.func("argmax_partial_f32");
12977 let f2 = self.func("argmax_final_f32");
12978 let mut guard = self.argmax_partials.lock().unwrap();
12979 if guard.is_none() {
12980 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
12981 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
12982 *guard = Some((pv, pi));
12983 }
12984 let (part_v, part_i) = guard.as_mut().unwrap();
12985 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
12986 let nv = n_vocab as i32;
12987 let nbi = nb as i32;
12988 let cfg1 = LaunchConfig {
12989 grid_dim: (nb as u32, 1, 1),
12990 block_dim: (256, 1, 1),
12991 shared_mem_bytes: 0,
12992 };
12993 let __s_b1 = self.gpu.stream();
12994 let mut b1 = __s_b1.launch_builder(&f1);
12995 b1.arg(&col_view)
12996 .arg(&mut *part_v)
12997 .arg(&mut *part_i)
12998 .arg(&nv);
12999 unsafe {
13000 b1.launch(cfg1)?;
13001 }
13002 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
13003 let cfg2 = LaunchConfig {
13004 grid_dim: (1, 1, 1),
13005 block_dim: (256, 1, 1),
13006 shared_mem_bytes: 0,
13007 };
13008 let __s_b2 = self.gpu.stream();
13009 let mut b2 = __s_b2.launch_builder(&f2);
13010 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
13011 unsafe {
13012 b2.launch(cfg2)?;
13013 }
13014 Ok(())
13015 }
13016 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13018 Ok(self.gpu.stream().clone_htod(v)?)
13019 }
13020 #[track_caller]
13021 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
13022 crate::dtoh_trace_hit(d.len() * 8);
13023 let v = self.gpu.stream().clone_dtoh(d)?;
13024 self.gpu.stream().synchronize()?;
13025 Ok(v)
13026 }
13027
13028 #[track_caller]
13029 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
13030 crate::dtoh_trace_hit(d.len() * 4);
13031 let v = self.gpu.stream().clone_dtoh(d)?;
13032 self.gpu.stream().synchronize()?;
13033 Ok(v)
13034 }
13035 pub fn htod_u32_into(
13039 &self,
13040 dst: &mut CudaSlice<u32>,
13041 src: &[u32],
13042 ) -> Result<(), Box<dyn std::error::Error>> {
13043 let mut view = dst.slice_mut(0..src.len());
13044 self.gpu.stream().memcpy_htod(src, &mut view)?;
13045 Ok(())
13046 }
13047
13048 pub fn htod_i32_into(
13051 &self,
13052 dst: &mut CudaSlice<i32>,
13053 src: &[i32],
13054 ) -> Result<(), Box<dyn std::error::Error>> {
13055 let mut view = dst.slice_mut(0..src.len());
13056 self.gpu.stream().memcpy_htod(src, &mut view)?;
13057 Ok(())
13058 }
13059
13060 #[track_caller]
13061 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13062 crate::alloc_trace_hit(n * 4);
13063 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
13064 self.keep_if_capturing(&s);
13065 Ok(s)
13066 }
13067 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device_into(
13071 &self,
13072 embd: &CudaSlice<u8>,
13073 token_d: &CudaSlice<u32>,
13074 x_out: &mut CudaSlice<f32>,
13075 n_embd: usize,
13076 qtype: i32,
13077 row_bytes: usize,
13078 ) -> Result<(), Box<dyn std::error::Error>> {
13079 let f = self.func("embed_gather_u32");
13080 let cfg = LaunchConfig {
13081 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
13082 block_dim: (256, 1, 1),
13083 shared_mem_bytes: 0,
13084 };
13085 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
13086 let __s_b = self.gpu.stream();
13087 let mut b = __s_b.launch_builder(&f);
13088 b.arg(embd)
13089 .arg(token_d)
13090 .arg(x_out)
13091 .arg(&ne)
13092 .arg(&qt)
13093 .arg(&rb);
13094 unsafe {
13095 b.launch(cfg)?;
13096 }
13097 Ok(())
13098 }
13099 #[track_caller]
13101 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
13102 crate::dtoh_trace_hit(d.len() * 4);
13103 let v = self.gpu.stream().clone_dtoh(d)?;
13104 self.gpu.stream().synchronize()?;
13105 Ok(v[0])
13106 }
13107 pub fn i32_set_k(
13114 &self,
13115 dst: &mut CudaSlice<i32>,
13116 v: i32,
13117 ) -> Result<(), Box<dyn std::error::Error>> {
13118 let f = self.func("i32_set_k");
13119 let cfg = LaunchConfig {
13120 grid_dim: (1, 1, 1),
13121 block_dim: (1, 1, 1),
13122 shared_mem_bytes: 0,
13123 };
13124 let idx = 0i32;
13125 let __s_b = self.gpu.stream();
13126 let mut b = __s_b.launch_builder(&f);
13127 b.arg(dst).arg(&v).arg(&idx);
13128 unsafe {
13129 b.launch(cfg)?;
13130 }
13131 Ok(())
13132 }
13133
13134 pub fn set_i32_one(
13135 &self,
13136 d: &mut CudaSlice<i32>,
13137 v: i32,
13138 ) -> Result<(), Box<dyn std::error::Error>> {
13139 self.gpu.stream().memcpy_htod(&[v], d)?;
13140 Ok(())
13141 }
13142 pub fn set_u32_one(
13145 &self,
13146 d: &mut CudaSlice<u32>,
13147 v: u32,
13148 ) -> Result<(), Box<dyn std::error::Error>> {
13149 self.gpu.stream().memcpy_htod(&[v], d)?;
13150 Ok(())
13151 }
13152 #[track_caller]
13154 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
13155 crate::dtoh_trace_hit(d.len() * 4);
13156 let v = self.gpu.stream().clone_dtoh(d)?;
13157 self.gpu.stream().synchronize()?;
13158 Ok(v[0])
13159 }
13160 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
13162 Ok(self.gpu.stream().clone_htod(bytes)?)
13163 }
13164 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device(
13169 &self,
13170 embd: &CudaSlice<u8>,
13171 token_d: &CudaSlice<u32>,
13172 n_embd: usize,
13173 qtype: i32,
13174 row_bytes: usize,
13175 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13176 let f = self.func("embed_gather_u32");
13177 let mut x = self.alloc_uninit::<f32>(n_embd)?;
13178 let cfg = LaunchConfig {
13179 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
13180 block_dim: (256, 1, 1),
13181 shared_mem_bytes: 0,
13182 };
13183 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
13184 let __s_b = self.gpu.stream();
13185 let mut b = __s_b.launch_builder(&f);
13186 b.arg(embd)
13187 .arg(token_d)
13188 .arg(&mut x)
13189 .arg(&ne)
13190 .arg(&qt)
13191 .arg(&rb);
13192 unsafe {
13193 b.launch(cfg)?;
13194 }
13195 Ok(x)
13196 }
13197
13198 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device_t(
13203 &self,
13204 embd: &CudaSlice<u8>,
13205 tokens: &[u32],
13206 n_embd: usize,
13207 qtype: i32,
13208 row_bytes: usize,
13209 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13210 let t = tokens.len();
13211 let tok_d = self.gpu.stream().clone_htod(tokens)?;
13212 let f = self.func("embed_gather_u32_t");
13213 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13214 let cfg = LaunchConfig {
13215 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13216 block_dim: (256, 1, 1),
13217 shared_mem_bytes: 0,
13218 };
13219 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13220 let __s_b = self.gpu.stream();
13221 let mut b = __s_b.launch_builder(&f);
13222 b.arg(embd)
13223 .arg(&tok_d)
13224 .arg(&mut x)
13225 .arg(&ne)
13226 .arg(&qt)
13227 .arg(&rb)
13228 .arg(&ti);
13229 unsafe {
13230 b.launch(cfg)?;
13231 }
13232 Ok(x)
13233 }
13234
13235 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device_tv(
13241 &self,
13242 embd: &CudaSlice<u8>,
13243 tok_v: &cudarc::driver::CudaView<u32>,
13244 t: usize,
13245 n_embd: usize,
13246 qtype: i32,
13247 row_bytes: usize,
13248 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13249 let f = self.func("embed_gather_u32_t");
13250 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13251 let cfg = LaunchConfig {
13252 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13253 block_dim: (256, 1, 1),
13254 shared_mem_bytes: 0,
13255 };
13256 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13257 let __s_b = self.gpu.stream();
13258 let mut b = __s_b.launch_builder(&f);
13259 b.arg(embd)
13260 .arg(tok_v)
13261 .arg(&mut x)
13262 .arg(&ne)
13263 .arg(&qt)
13264 .arg(&rb)
13265 .arg(&ti);
13266 unsafe {
13267 b.launch(cfg)?;
13268 }
13269 Ok(x)
13270 }
13271
13272 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device_td(
13274 &self,
13275 embd: &CudaSlice<u8>,
13276 tok_d: &CudaSlice<u32>,
13277 t: usize,
13278 n_embd: usize,
13279 qtype: i32,
13280 row_bytes: usize,
13281 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13282 let f = self.func("embed_gather_u32_t");
13283 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13284 let cfg = LaunchConfig {
13285 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13286 block_dim: (256, 1, 1),
13287 shared_mem_bytes: 0,
13288 };
13289 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13290 let __s_b = self.gpu.stream();
13291 let mut b = __s_b.launch_builder(&f);
13292 b.arg(embd)
13293 .arg(tok_d)
13294 .arg(&mut x)
13295 .arg(&ne)
13296 .arg(&qt)
13297 .arg(&rb)
13298 .arg(&ti);
13299 unsafe {
13300 b.launch(cfg)?;
13301 }
13302 Ok(x)
13303 }
13304
13305 #[inline]
13311 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
13313 if self
13314 .capture_keep_on
13315 .load(std::sync::atomic::Ordering::Relaxed)
13316 {
13317 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
13318 }
13319 }
13320
13321 #[track_caller]
13322 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
13323 &self,
13324 n: usize,
13325 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
13326 SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13327 crate::alloc_trace_hit(n * std::mem::size_of::<T>());
13328 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
13329 {
13332 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13333 if *ON.get_or_init(|| std::env::var("MEMRA_ALLOC_SITES").as_deref() == Ok("1")) {
13334 use cudarc::driver::DevicePtr;
13335 type Sites = std::collections::HashMap<u64, (&'static str, u32, usize, usize)>;
13336 static MAP: std::sync::OnceLock<std::sync::Mutex<Sites>> =
13337 std::sync::OnceLock::new();
13338 static LINES: std::sync::atomic::AtomicUsize =
13339 std::sync::atomic::AtomicUsize::new(0);
13340 let stream = self.gpu.stream();
13341 let (ptr, _g) = s.device_ptr(&stream);
13342 let sh = stream.cu_stream() as usize;
13343 let loc = std::panic::Location::caller();
13344 let mut m = MAP.get_or_init(Default::default).lock().unwrap();
13345 if let Some((pf, pl, ps, pb)) = m.get(&ptr)
13346 && *ps != sh
13347 && LINES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) < 24
13348 {
13349 eprintln!(
13350 "[alloc-sites] XSTREAM reuse ptr={ptr:#x} now {}:{} {} B on stream {sh:#x} \
13351 <- last {pf}:{pl} {pb} B on stream {ps:#x}",
13352 loc.file(),
13353 loc.line(),
13354 n * std::mem::size_of::<T>()
13355 );
13356 }
13357 m.insert(
13358 ptr,
13359 (loc.file(), loc.line(), sh, n * std::mem::size_of::<T>()),
13360 );
13361 }
13362 }
13363 {
13367 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13368 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
13369 use cudarc::driver::DevicePtrMut;
13371 let n_bytes = s.len() * std::mem::size_of::<T>();
13372 let stream = self.gpu.stream();
13373 let (p_, _g) = s.device_ptr_mut(&stream);
13374 unsafe {
13375 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
13376 .result()?;
13377 }
13378 }
13379 }
13380 self.keep_if_capturing(&s);
13381 Ok(s)
13382 }
13383
13384 #[track_caller]
13389 pub fn uninit_q8_pair(
13390 &self,
13391 n: usize,
13392 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13393 Ok((
13394 self.alloc_uninit::<i8>(n)?,
13395 self.alloc_uninit::<f32>(n / 32)?,
13396 ))
13397 }
13398
13399 #[track_caller]
13400 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13401 self.alloc_uninit::<f32>(n)
13402 }
13403
13404 #[track_caller]
13406 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
13407 self.alloc_uninit::<i8>(n)
13408 }
13409
13410 #[track_caller]
13412 pub fn uninit_i32(&self, n: usize) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
13413 self.alloc_uninit::<i32>(n)
13414 }
13415
13416 #[allow(clippy::too_many_arguments)]
13420 pub fn rms_norm3(
13421 &self,
13422 x: &CudaSlice<f32>,
13423 w0: &CudaSlice<f32>,
13424 w1: &CudaSlice<f32>,
13425 w2: &CudaSlice<f32>,
13426 d0: &mut CudaSlice<f32>,
13427 d1: &mut CudaSlice<f32>,
13428 d2: &mut CudaSlice<f32>,
13429 ncols: usize,
13430 nrows: usize,
13431 eps: f32,
13432 ) -> Result<(), Box<dyn std::error::Error>> {
13433 let f = self.func("rms_norm3_f32");
13434 let cfg = LaunchConfig {
13435 grid_dim: (nrows as u32, 1, 1),
13436 block_dim: (rms_block(), 1, 1),
13437 shared_mem_bytes: 0,
13438 };
13439 let (nc, e) = (ncols as i32, eps);
13440 let __s_b = self.gpu.stream();
13441 let mut b = __s_b.launch_builder(&f);
13442 b.arg(x)
13443 .arg(w0)
13444 .arg(w1)
13445 .arg(w2)
13446 .arg(d0)
13447 .arg(d1)
13448 .arg(d2)
13449 .arg(&nc)
13450 .arg(&e);
13451 unsafe {
13452 b.launch(cfg)?;
13453 }
13454 Ok(())
13455 }
13456
13457 #[allow(clippy::too_many_arguments)]
13459 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
13462 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13463 *WARP_ON.get_or_init(|| {
13464 std::env::var("MEMRA_QKVNORM_W")
13465 .map(|v| v != "0")
13466 .unwrap_or(true)
13467 }) && ncols.is_multiple_of(4)
13468 && rows >= 64
13469 }
13470
13471 #[allow(clippy::too_many_arguments)]
13474 pub fn rms_norm_qkv_w4b(
13475 &self,
13476 q: &CudaSlice<f32>,
13477 k: &CudaSlice<f32>,
13478 v: &CudaSlice<f32>,
13479 wq: &CudaSlice<f32>,
13480 wk: &CudaSlice<f32>,
13481 wv: &CudaSlice<f32>,
13482 dq: &mut CudaSlice<f32>,
13483 dk: &mut CudaSlice<f32>,
13484 dv: &mut CudaSlice<f32>,
13485 dvb: &mut CudaSlice<u8>,
13486 ncols: usize,
13487 rq: usize,
13488 rk: usize,
13489 eps: f32,
13490 vf16: bool,
13491 ) -> Result<(), Box<dyn std::error::Error>> {
13492 assert!(ncols.is_multiple_of(4) && rq + 2 * rk >= 64);
13493 let f = self.func("rms_norm_qkv_w4b_f32");
13494 let rows = (rq + 2 * rk) as u32;
13495 let cfg = LaunchConfig {
13496 grid_dim: (rows.div_ceil(8), 1, 1),
13497 block_dim: (256, 1, 1),
13498 shared_mem_bytes: 0,
13499 };
13500 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
13501 let vf = vf16 as i32;
13502 let __s_b = self.gpu.stream();
13503 let mut b = __s_b.launch_builder(&f);
13504 b.arg(q)
13505 .arg(k)
13506 .arg(v)
13507 .arg(wq)
13508 .arg(wk)
13509 .arg(wv)
13510 .arg(dq)
13511 .arg(dk)
13512 .arg(dv)
13513 .arg(&mut *dvb)
13514 .arg(&nc)
13515 .arg(&rqi)
13516 .arg(&rki)
13517 .arg(&rvi)
13518 .arg(&e)
13519 .arg(&vf);
13520 unsafe {
13521 b.launch(cfg)?;
13522 }
13523 Ok(())
13524 }
13525
13526 #[allow(clippy::too_many_arguments)] pub fn rms_norm_qkv(
13528 &self,
13529 q: &CudaSlice<f32>,
13530 k: &CudaSlice<f32>,
13531 v: &CudaSlice<f32>,
13532 wq: &CudaSlice<f32>,
13533 wk: &CudaSlice<f32>,
13534 wv: &CudaSlice<f32>,
13535 dq: &mut CudaSlice<f32>,
13536 dk: &mut CudaSlice<f32>,
13537 dv: &mut CudaSlice<f32>,
13538 ncols: usize,
13539 rq: usize,
13540 rk: usize,
13541 eps: f32,
13542 ) -> Result<(), Box<dyn std::error::Error>> {
13543 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13547 let warp_on = *WARP_ON.get_or_init(|| {
13548 std::env::var("MEMRA_QKVNORM_W")
13549 .map(|v| v != "0")
13550 .unwrap_or(true)
13551 });
13552 if warp_on && ncols.is_multiple_of(4) && rq + 2 * rk >= 64 {
13555 let f = self.func("rms_norm_qkv_w4_f32");
13556 let rows = (rq + 2 * rk) as u32;
13557 let cfg = LaunchConfig {
13558 grid_dim: (rows.div_ceil(8), 1, 1),
13559 block_dim: (256, 1, 1),
13560 shared_mem_bytes: 0,
13561 };
13562 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
13563 let __s_b = self.gpu.stream();
13564 let mut b = __s_b.launch_builder(&f);
13565 b.arg(q)
13566 .arg(k)
13567 .arg(v)
13568 .arg(wq)
13569 .arg(wk)
13570 .arg(wv)
13571 .arg(dq)
13572 .arg(dk)
13573 .arg(dv)
13574 .arg(&nc)
13575 .arg(&rqi)
13576 .arg(&rki)
13577 .arg(&rvi)
13578 .arg(&e);
13579 unsafe {
13580 b.launch(cfg)?;
13581 }
13582 return Ok(());
13583 }
13584 let f = self.func("rms_norm_qkv_f32");
13585 let grid = (rq + 2 * rk) as u32;
13586 let cfg = LaunchConfig {
13587 grid_dim: (grid, 1, 1),
13588 block_dim: (rms_block(), 1, 1),
13589 shared_mem_bytes: 0,
13590 };
13591 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
13592 let __s_b = self.gpu.stream();
13593 let mut b = __s_b.launch_builder(&f);
13594 b.arg(q)
13595 .arg(k)
13596 .arg(v)
13597 .arg(wq)
13598 .arg(wk)
13599 .arg(wv)
13600 .arg(dq)
13601 .arg(dk)
13602 .arg(dv)
13603 .arg(&nc)
13604 .arg(&rqi)
13605 .arg(&rki)
13606 .arg(&e);
13607 unsafe {
13608 b.launch(cfg)?;
13609 }
13610 Ok(())
13611 }
13612
13613 #[allow(clippy::too_many_arguments)]
13615 pub fn rms_norm2x(
13616 &self,
13617 a: &CudaSlice<f32>,
13618 bb: &CudaSlice<f32>,
13619 wa: &CudaSlice<f32>,
13620 wb: &CudaSlice<f32>,
13621 da: &mut CudaSlice<f32>,
13622 db: &mut CudaSlice<f32>,
13623 ncols: usize,
13624 nrows: usize,
13625 eps: f32,
13626 ) -> Result<(), Box<dyn std::error::Error>> {
13627 let f = self.func("rms_norm2x_f32");
13628 let cfg = LaunchConfig {
13629 grid_dim: (2 * nrows as u32, 1, 1),
13630 block_dim: (rms_block(), 1, 1),
13631 shared_mem_bytes: 0,
13632 };
13633 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
13634 let __s_b = self.gpu.stream();
13635 let mut b = __s_b.launch_builder(&f);
13636 b.arg(a)
13637 .arg(bb)
13638 .arg(wa)
13639 .arg(wb)
13640 .arg(da)
13641 .arg(db)
13642 .arg(&nc)
13643 .arg(&nr)
13644 .arg(&e);
13645 unsafe {
13646 b.launch(cfg)?;
13647 }
13648 Ok(())
13649 }
13650
13651 pub fn softcap(
13653 &self,
13654 y: &mut CudaSlice<f32>,
13655 cap: f32,
13656 n: usize,
13657 ) -> Result<(), Box<dyn std::error::Error>> {
13658 let f = self.func("softcap_f32");
13659 let cfg = LaunchConfig::for_num_elems(n as u32);
13660 let ni = n as i32;
13661 let __s_b = self.gpu.stream();
13662 let mut b = __s_b.launch_builder(&f);
13663 b.arg(y).arg(&cap).arg(&ni);
13664 unsafe {
13665 b.launch(cfg)?;
13666 }
13667 Ok(())
13668 }
13669
13670 pub fn mask_ids_rows(
13673 &self,
13674 y: &mut CudaSlice<f32>,
13675 ids: &CudaSlice<i32>,
13676 n_ids: usize,
13677 n_vocab: usize,
13678 t: usize,
13679 ) -> Result<(), Box<dyn std::error::Error>> {
13680 let f = self.func("mask_ids_rows_f32");
13681 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
13682 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
13683 let __s_b = self.gpu.stream();
13684 let mut b = __s_b.launch_builder(&f);
13685 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
13686 unsafe {
13687 b.launch(cfg)?;
13688 }
13689 Ok(())
13690 }
13691
13692 #[allow(clippy::too_many_arguments)]
13694 pub fn add_scale_rms_norm(
13695 &self,
13696 a: &CudaSlice<f32>,
13697 b_in: &CudaSlice<f32>,
13698 c: f32,
13699 w: &CudaSlice<f32>,
13700 res: &mut CudaSlice<f32>,
13701 dst: &mut CudaSlice<f32>,
13702 ncols: usize,
13703 nrows: usize,
13704 eps: f32,
13705 ) -> Result<(), Box<dyn std::error::Error>> {
13706 let f = self.func("add_scale_rms_norm_f32");
13707 let cfg = LaunchConfig {
13708 grid_dim: (nrows as u32, 1, 1),
13709 block_dim: (rms_block(), 1, 1),
13710 shared_mem_bytes: 0,
13711 };
13712 let (nc, e2) = (ncols as i32, eps);
13713 let __s_b = self.gpu.stream();
13714 let mut b = __s_b.launch_builder(&f);
13715 b.arg(a)
13716 .arg(b_in)
13717 .arg(&c)
13718 .arg(w)
13719 .arg(res)
13720 .arg(dst)
13721 .arg(&nc)
13722 .arg(&e2);
13723 unsafe {
13724 b.launch(cfg)?;
13725 }
13726 Ok(())
13727 }
13728
13729 #[allow(clippy::too_many_arguments)]
13732 pub fn add_scale_rms_norm_q8_1(
13733 &self,
13734 a: &CudaSlice<f32>,
13735 b_in: &CudaSlice<f32>,
13736 c: f32,
13737 w: &CudaSlice<f32>,
13738 res: &mut CudaSlice<f32>,
13739 ncols: usize,
13740 nrows: usize,
13741 eps: f32,
13742 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13743 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
13744 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13745 let (nc, e2) = (ncols as i32, eps);
13746 if Self::pdl_on() && Self::pdl_wb_on() {
13747 {
13748 use cudarc::driver::{DevicePtr, DevicePtrMut};
13749 let s = &self.gpu.stream();
13750 let (pa, _g0) = a.device_ptr(s);
13751 let (pb, _g1) = b_in.device_ptr(s);
13752 let (pw, _g2) = w.device_ptr(s);
13753 let (pr, _g3) = res.device_ptr_mut(s);
13754 let (pq, _g4) = out_q.device_ptr_mut(s);
13755 let (pd, _g5) = out_d.device_ptr_mut(s);
13756 let mut ps = [
13757 &pa as *const _ as *mut std::ffi::c_void,
13758 &pb as *const _ as *mut _,
13759 &c as *const _ as *mut _,
13760 &pw as *const _ as *mut _,
13761 &pr as *const _ as *mut _,
13762 &pq as *const _ as *mut _,
13763 &pd as *const _ as *mut _,
13764 &nc as *const _ as *mut _,
13765 &e2 as *const _ as *mut _,
13766 ];
13767 unsafe {
13768 self.launch_pdl(
13769 "add_scale_rms_norm_q8_1",
13770 (nrows as u32, 1, 1),
13771 (rms_block(), 1, 1),
13772 &mut ps,
13773 )?;
13774 }
13775 }
13776 return Ok((out_q, out_d));
13777 }
13778 let f = self.func("add_scale_rms_norm_q8_1");
13779 let cfg = LaunchConfig {
13780 grid_dim: (nrows as u32, 1, 1),
13781 block_dim: (rms_block(), 1, 1),
13782 shared_mem_bytes: 0,
13783 };
13784 let __s_b = self.gpu.stream();
13785 let mut b = __s_b.launch_builder(&f);
13786 b.arg(a)
13787 .arg(b_in)
13788 .arg(&c)
13789 .arg(w)
13790 .arg(res)
13791 .arg(&mut out_q)
13792 .arg(&mut out_d)
13793 .arg(&nc)
13794 .arg(&e2);
13795 unsafe {
13796 b.launch(cfg)?;
13797 }
13798 Ok((out_q, out_d))
13799 }
13800
13801 #[allow(clippy::too_many_arguments)]
13803 pub fn add_scale_rms_norm_q8_1_into(
13804 &self,
13805 a: &CudaSlice<f32>,
13806 b_in: &CudaSlice<f32>,
13807 c: f32,
13808 w: &CudaSlice<f32>,
13809 res: &mut CudaSlice<f32>,
13810 ncols: usize,
13811 nrows: usize,
13812 eps: f32,
13813 out_q: &mut CudaSlice<i8>,
13814 out_d: &mut CudaSlice<f32>,
13815 ) -> Result<(), Box<dyn std::error::Error>> {
13816 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
13817 let (nc, e2) = (ncols as i32, eps);
13818 if Self::pdl_on() && Self::pdl_wb_on() {
13819 use cudarc::driver::{DevicePtr, DevicePtrMut};
13820 let s = &self.gpu.stream();
13821 let (pa, _g0) = a.device_ptr(s);
13822 let (pb, _g1) = b_in.device_ptr(s);
13823 let (pw, _g2) = w.device_ptr(s);
13824 let (pr, _g3) = res.device_ptr_mut(s);
13825 let (pq, _g4) = out_q.device_ptr_mut(s);
13826 let (pd, _g5) = out_d.device_ptr_mut(s);
13827 let mut ps = [
13828 &pa as *const _ as *mut std::ffi::c_void,
13829 &pb as *const _ as *mut _,
13830 &c as *const _ as *mut _,
13831 &pw as *const _ as *mut _,
13832 &pr as *const _ as *mut _,
13833 &pq as *const _ as *mut _,
13834 &pd as *const _ as *mut _,
13835 &nc as *const _ as *mut _,
13836 &e2 as *const _ as *mut _,
13837 ];
13838 unsafe {
13839 self.launch_pdl(
13840 "add_scale_rms_norm_q8_1",
13841 (nrows as u32, 1, 1),
13842 (rms_block(), 1, 1),
13843 &mut ps,
13844 )?;
13845 }
13846 return Ok(());
13847 }
13848 let f = self.func("add_scale_rms_norm_q8_1");
13849 let cfg = LaunchConfig {
13850 grid_dim: (nrows as u32, 1, 1),
13851 block_dim: (rms_block(), 1, 1),
13852 shared_mem_bytes: 0,
13853 };
13854 let __s_b = self.gpu.stream();
13855 let mut b = __s_b.launch_builder(&f);
13856 b.arg(a)
13857 .arg(b_in)
13858 .arg(&c)
13859 .arg(w)
13860 .arg(res)
13861 .arg(&mut *out_q)
13862 .arg(&mut *out_d)
13863 .arg(&nc)
13864 .arg(&e2);
13865 unsafe {
13866 b.launch(cfg)?;
13867 }
13868 Ok(())
13869 }
13870
13871 #[allow(clippy::too_many_arguments)]
13874 pub fn rms_pre_add_scale_rms_norm_q8_1(
13875 &self,
13876 a: &CudaSlice<f32>,
13877 wa: &CudaSlice<f32>,
13878 b_in: &CudaSlice<f32>,
13879 c: f32,
13880 w: &CudaSlice<f32>,
13881 res: &mut CudaSlice<f32>,
13882 ncols: usize,
13883 nrows: usize,
13884 eps: f32,
13885 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13886 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
13887 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13888 let (nc, e2) = (ncols as i32, eps);
13889 if Self::pdl_on() {
13890 {
13891 use cudarc::driver::{DevicePtr, DevicePtrMut};
13892 let s = &self.gpu.stream();
13893 let (pa, _g0) = a.device_ptr(s);
13894 let (pwa, _g1) = wa.device_ptr(s);
13895 let (pb, _g2) = b_in.device_ptr(s);
13896 let (pw, _g3) = w.device_ptr(s);
13897 let (pr, _g4) = res.device_ptr_mut(s);
13898 let (pq, _g5) = out_q.device_ptr_mut(s);
13899 let (pd, _g6) = out_d.device_ptr_mut(s);
13900 let mut ps = [
13901 &pa as *const _ as *mut std::ffi::c_void,
13902 &pwa as *const _ as *mut _,
13903 &pb as *const _ as *mut _,
13904 &c as *const _ as *mut _,
13905 &pw as *const _ as *mut _,
13906 &pr as *const _ as *mut _,
13907 &pq as *const _ as *mut _,
13908 &pd as *const _ as *mut _,
13909 &nc as *const _ as *mut _,
13910 &e2 as *const _ as *mut _,
13911 ];
13912 unsafe {
13913 self.launch_pdl(
13914 "rms_pre_add_scale_rms_norm_q8_1",
13915 (nrows as u32, 1, 1),
13916 (rms_block(), 1, 1),
13917 &mut ps,
13918 )?;
13919 }
13920 }
13921 return Ok((out_q, out_d));
13922 }
13923 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
13924 let cfg = LaunchConfig {
13925 grid_dim: (nrows as u32, 1, 1),
13926 block_dim: (rms_block(), 1, 1),
13927 shared_mem_bytes: 0,
13928 };
13929 let __s_b = self.gpu.stream();
13930 let mut b = __s_b.launch_builder(&f);
13931 b.arg(a)
13932 .arg(wa)
13933 .arg(b_in)
13934 .arg(&c)
13935 .arg(w)
13936 .arg(res)
13937 .arg(&mut out_q)
13938 .arg(&mut out_d)
13939 .arg(&nc)
13940 .arg(&e2);
13941 unsafe {
13942 b.launch(cfg)?;
13943 }
13944 Ok((out_q, out_d))
13945 }
13946
13947 pub fn gelu_tanh_mul_q8_1(
13950 &self,
13951 gate: &CudaSlice<f32>,
13952 up: &cudarc::driver::CudaView<f32>,
13953 act: &mut CudaSlice<f32>,
13954 ncols: usize,
13955 nrows: usize,
13956 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13957 debug_assert!(ncols.is_multiple_of(128));
13958 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
13959 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
13960 let nc = ncols as i32;
13961 if Self::pdl_on() {
13962 {
13963 use cudarc::driver::{DevicePtr, DevicePtrMut};
13964 let s = &self.gpu.stream();
13965 let (pg, _g0) = gate.device_ptr(s);
13966 let (pu, _g1) = up.device_ptr(s);
13967 let (pact, _g2) = act.device_ptr_mut(s);
13968 let (pq, _g3) = out_q.device_ptr_mut(s);
13969 let (pd, _g4) = out_d.device_ptr_mut(s);
13970 let mut ps = [
13971 &pg as *const _ as *mut std::ffi::c_void,
13972 &pu as *const _ as *mut _,
13973 &pact as *const _ as *mut _,
13974 &pq as *const _ as *mut _,
13975 &pd as *const _ as *mut _,
13976 &nc as *const _ as *mut _,
13977 ];
13978 unsafe {
13979 self.launch_pdl(
13980 "gelu_tanh_mul_q8_1",
13981 (nrows as u32, 1, 1),
13982 (rms_block(), 1, 1),
13983 &mut ps,
13984 )?;
13985 }
13986 }
13987 return Ok((out_q, out_d));
13988 }
13989 let f = self.func("gelu_tanh_mul_q8_1");
13990 let cfg = LaunchConfig {
13991 grid_dim: (nrows as u32, 1, 1),
13992 block_dim: (rms_block(), 1, 1),
13993 shared_mem_bytes: 0,
13994 };
13995 let __s_b = self.gpu.stream();
13996 let mut b = __s_b.launch_builder(&f);
13997 b.arg(gate)
13998 .arg(up)
13999 .arg(act)
14000 .arg(&mut out_q)
14001 .arg(&mut out_d)
14002 .arg(&nc);
14003 unsafe {
14004 b.launch(cfg)?;
14005 }
14006 Ok((out_q, out_d))
14007 }
14008
14009 #[allow(clippy::too_many_arguments)]
14011 pub fn gelu_tanh_mul_q8_1_into(
14012 &self,
14013 gate: &CudaSlice<f32>,
14014 up: &cudarc::driver::CudaView<f32>,
14015 act: &mut CudaSlice<f32>,
14016 ncols: usize,
14017 nrows: usize,
14018 out_q: &mut CudaSlice<i8>,
14019 out_d: &mut CudaSlice<f32>,
14020 ) -> Result<(), Box<dyn std::error::Error>> {
14021 debug_assert!(ncols.is_multiple_of(128));
14022 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
14023 let nc = ncols as i32;
14024 if Self::pdl_on() {
14025 use cudarc::driver::{DevicePtr, DevicePtrMut};
14026 let s = &self.gpu.stream();
14027 let (pg, _g0) = gate.device_ptr(s);
14028 let (pu, _g1) = up.device_ptr(s);
14029 let (pact, _g2) = act.device_ptr_mut(s);
14030 let (pq, _g3) = out_q.device_ptr_mut(s);
14031 let (pd, _g4) = out_d.device_ptr_mut(s);
14032 let mut ps = [
14033 &pg as *const _ as *mut std::ffi::c_void,
14034 &pu as *const _ as *mut _,
14035 &pact as *const _ as *mut _,
14036 &pq as *const _ as *mut _,
14037 &pd as *const _ as *mut _,
14038 &nc as *const _ as *mut _,
14039 ];
14040 unsafe {
14041 self.launch_pdl(
14042 "gelu_tanh_mul_q8_1",
14043 (nrows as u32, 1, 1),
14044 (rms_block(), 1, 1),
14045 &mut ps,
14046 )?;
14047 }
14048 return Ok(());
14049 }
14050 let f = self.func("gelu_tanh_mul_q8_1");
14051 let cfg = LaunchConfig {
14052 grid_dim: (nrows as u32, 1, 1),
14053 block_dim: (rms_block(), 1, 1),
14054 shared_mem_bytes: 0,
14055 };
14056 let __s_b = self.gpu.stream();
14057 let mut b = __s_b.launch_builder(&f);
14058 b.arg(gate)
14059 .arg(up)
14060 .arg(&mut *act)
14061 .arg(&mut *out_q)
14062 .arg(&mut *out_d)
14063 .arg(&nc);
14064 unsafe {
14065 b.launch(cfg)?;
14066 }
14067 Ok(())
14068 }
14069
14070 #[allow(clippy::too_many_arguments)]
14072 #[allow(clippy::type_complexity)] pub fn add_rms_norm3_q8z(
14074 &self,
14075 a: &CudaSlice<f32>,
14076 b_in: &CudaSlice<f32>,
14077 w0: &CudaSlice<f32>,
14078 w1: &CudaSlice<f32>,
14079 w2: &CudaSlice<f32>,
14080 res: &mut CudaSlice<f32>,
14081 out1: &mut CudaSlice<f32>,
14082 ncols: usize,
14083 nrows: usize,
14084 eps: f32,
14085 ) -> Result<
14086 (
14087 (CudaSlice<i8>, CudaSlice<f32>),
14088 (CudaSlice<i8>, CudaSlice<f32>),
14089 ),
14090 Box<dyn std::error::Error>,
14091 > {
14092 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
14093 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14094 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
14095 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14096 let f = self.func("add_rms_norm3_q8z_f32");
14097 let cfg = LaunchConfig {
14098 grid_dim: (nrows as u32, 1, 1),
14099 block_dim: (rms_block(), 1, 1),
14100 shared_mem_bytes: 0,
14101 };
14102 let (nc, e2) = (ncols as i32, eps);
14103 let __s_b = self.gpu.stream();
14104 let mut b = __s_b.launch_builder(&f);
14105 b.arg(a)
14106 .arg(b_in)
14107 .arg(w0)
14108 .arg(w1)
14109 .arg(w2)
14110 .arg(res)
14111 .arg(&mut q0)
14112 .arg(&mut d0)
14113 .arg(out1)
14114 .arg(&mut q2)
14115 .arg(&mut d2)
14116 .arg(&nc)
14117 .arg(&e2);
14118 unsafe {
14119 b.launch(cfg)?;
14120 }
14121 Ok(((q0, d0), (q2, d2)))
14122 }
14123
14124 #[allow(clippy::too_many_arguments)]
14126 pub fn add_rms_norm3(
14127 &self,
14128 a: &CudaSlice<f32>,
14129 b_in: &CudaSlice<f32>,
14130 w0: &CudaSlice<f32>,
14131 w1: &CudaSlice<f32>,
14132 w2: &CudaSlice<f32>,
14133 res: &mut CudaSlice<f32>,
14134 d0: &mut CudaSlice<f32>,
14135 d1: &mut CudaSlice<f32>,
14136 d2: &mut CudaSlice<f32>,
14137 ncols: usize,
14138 nrows: usize,
14139 eps: f32,
14140 ) -> Result<(), Box<dyn std::error::Error>> {
14141 let f = self.func("add_rms_norm3_f32");
14142 let cfg = LaunchConfig {
14143 grid_dim: (nrows as u32, 1, 1),
14144 block_dim: (rms_block(), 1, 1),
14145 shared_mem_bytes: 0,
14146 };
14147 let (nc, e2) = (ncols as i32, eps);
14148 let __s_b = self.gpu.stream();
14149 let mut b = __s_b.launch_builder(&f);
14150 b.arg(a)
14151 .arg(b_in)
14152 .arg(w0)
14153 .arg(w1)
14154 .arg(w2)
14155 .arg(res)
14156 .arg(d0)
14157 .arg(d1)
14158 .arg(d2)
14159 .arg(&nc)
14160 .arg(&e2);
14161 unsafe {
14162 b.launch(cfg)?;
14163 }
14164 Ok(())
14165 }
14166
14167 pub fn add_scale(
14169 &self,
14170 a: &CudaSlice<f32>,
14171 b_in: &CudaSlice<f32>,
14172 c: f32,
14173 dst: &mut CudaSlice<f32>,
14174 n: usize,
14175 ) -> Result<(), Box<dyn std::error::Error>> {
14176 let f = self.func("add_scale_f32");
14177 let cfg = LaunchConfig::for_num_elems(n as u32);
14178 let ni = n as i32;
14179 let __s_b = self.gpu.stream();
14180 let mut b = __s_b.launch_builder(&f);
14181 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
14182 unsafe {
14183 b.launch(cfg)?;
14184 }
14185 Ok(())
14186 }
14187
14188 #[allow(clippy::too_many_arguments)] pub fn layer_norm_bias(
14191 &self,
14192 x: &CudaSlice<f32>,
14193 w: &CudaSlice<f32>,
14194 b: &CudaSlice<f32>,
14195 dst: &mut CudaSlice<f32>,
14196 ncols: usize,
14197 nrows: usize,
14198 eps: f32,
14199 ) -> Result<(), Box<dyn std::error::Error>> {
14200 let f = self.func("layer_norm_bias_f32");
14201 let (nc, e) = (ncols as i32, eps);
14202 let cfg = LaunchConfig {
14203 grid_dim: (nrows as u32, 1, 1),
14204 block_dim: (256, 1, 1),
14205 shared_mem_bytes: 0,
14206 };
14207 let __s_b = self.gpu.stream();
14208 let mut lb = __s_b.launch_builder(&f);
14209 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
14210 unsafe {
14211 lb.launch(cfg)?;
14212 }
14213 Ok(())
14214 }
14215
14216 pub fn gelu_tanh(
14218 &self,
14219 x: &CudaSlice<f32>,
14220 dst: &mut CudaSlice<f32>,
14221 n: usize,
14222 ) -> Result<(), Box<dyn std::error::Error>> {
14223 let f = self.func("gelu_tanh_f32");
14224 let ni = n as i64;
14225 let cfg = LaunchConfig {
14226 grid_dim: (n.div_ceil(256) as u32, 1, 1),
14227 block_dim: (256, 1, 1),
14228 shared_mem_bytes: 0,
14229 };
14230 let __s_b = self.gpu.stream();
14231 let mut lb = __s_b.launch_builder(&f);
14232 lb.arg(x).arg(&mut *dst).arg(&ni);
14233 unsafe {
14234 lb.launch(cfg)?;
14235 }
14236 Ok(())
14237 }
14238
14239 pub fn row_softmax(
14241 &self,
14242 x: &mut CudaSlice<f32>,
14243 ncols: usize,
14244 nrows: usize,
14245 ) -> Result<(), Box<dyn std::error::Error>> {
14246 let f = self.func("row_softmax_f32");
14247 let nc = ncols as i32;
14248 let cfg = LaunchConfig {
14249 grid_dim: (nrows as u32, 1, 1),
14250 block_dim: (256, 1, 1),
14251 shared_mem_bytes: 0,
14252 };
14253 let __s_b = self.gpu.stream();
14254 let mut lb = __s_b.launch_builder(&f);
14255 lb.arg(&mut *x).arg(&nc);
14256 unsafe {
14257 lb.launch(cfg)?;
14258 }
14259 Ok(())
14260 }
14261
14262 pub fn rms_norm(
14263 &self,
14264 x: &CudaSlice<f32>,
14265 w: &CudaSlice<f32>,
14266 dst: &mut CudaSlice<f32>,
14267 ncols: usize,
14268 nrows: usize,
14269 eps: f32,
14270 ) -> Result<(), Box<dyn std::error::Error>> {
14271 let (nc, e) = (ncols as i32, eps);
14272 let kname = if Self::norm_ilp_on() {
14273 "rms_norm_f32_v2"
14274 } else {
14275 "rms_norm_f32"
14276 };
14277 if Self::pdl_on() && Self::pdl_wb_on() {
14278 use cudarc::driver::{DevicePtr, DevicePtrMut};
14279 let s = &self.gpu.stream();
14280 let (px, _g0) = x.device_ptr(s);
14281 let (pw, _g1) = w.device_ptr(s);
14282 let (pd, _g2) = dst.device_ptr_mut(s);
14283 let mut ps = [
14284 &px as *const _ as *mut std::ffi::c_void,
14285 &pw as *const _ as *mut _,
14286 &pd as *const _ as *mut _,
14287 &nc as *const _ as *mut _,
14288 &e as *const _ as *mut _,
14289 ];
14290 unsafe {
14291 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
14292 }
14293 return Ok(());
14294 }
14295 let f = self.func(kname);
14296 let cfg = LaunchConfig {
14297 grid_dim: (nrows as u32, 1, 1),
14298 block_dim: (rms_block(), 1, 1),
14299 shared_mem_bytes: 0,
14300 };
14301 let __s_b = self.gpu.stream();
14302 let mut b = __s_b.launch_builder(&f);
14303 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
14304 unsafe {
14305 b.launch(cfg)?;
14306 }
14307 Ok(())
14308 }
14309
14310 pub fn rms_norm_decode(
14318 &self,
14319 x: &CudaSlice<f32>,
14320 w: &CudaSlice<f32>,
14321 dst: &mut CudaSlice<f32>,
14322 ncols: usize,
14323 nrows: usize,
14324 eps: f32,
14325 ) -> Result<(), Box<dyn std::error::Error>> {
14326 let f = self.func(if Self::norm_ilp_on() {
14327 "rms_norm_f32_v2"
14328 } else {
14329 "rms_norm_f32"
14330 });
14331 let cfg = LaunchConfig {
14332 grid_dim: (nrows as u32, 1, 1),
14333 block_dim: (1024, 1, 1),
14334 shared_mem_bytes: 0,
14335 };
14336 let (nc, e) = (ncols as i32, eps);
14337 let __s_b = self.gpu.stream();
14338 let mut b = __s_b.launch_builder(&f);
14339 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
14340 unsafe {
14341 b.launch(cfg)?;
14342 }
14343 Ok(())
14344 }
14345
14346 pub fn rms_norm_q8_1(
14350 &self,
14351 x: &CudaSlice<f32>,
14352 w: &CudaSlice<f32>,
14353 ncols: usize,
14354 nrows: usize,
14355 eps: f32,
14356 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14357 let nblk = ncols / 32;
14358 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
14359 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
14360 let (nc, e) = (ncols as i32, eps);
14361 if Self::pdl_on() {
14362 {
14363 use cudarc::driver::{DevicePtr, DevicePtrMut};
14364 let s = &self.gpu.stream();
14365 let (px, _g0) = x.device_ptr(s);
14366 let (pw, _g1) = w.device_ptr(s);
14367 let (pq, _g2) = q.device_ptr_mut(s);
14368 let (pd, _g3) = d.device_ptr_mut(s);
14369 let mut ps = [
14370 &px as *const _ as *mut std::ffi::c_void,
14371 &pw as *const _ as *mut _,
14372 &pq as *const _ as *mut _,
14373 &pd as *const _ as *mut _,
14374 &nc as *const _ as *mut _,
14375 &e as *const _ as *mut _,
14376 ];
14377 unsafe {
14378 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
14379 }
14380 }
14381 return Ok((q, d));
14382 }
14383 let f = self.func("rms_norm_q8_1");
14384 let cfg = LaunchConfig {
14387 grid_dim: (nrows as u32, 1, 1),
14388 block_dim: (1024, 1, 1),
14389 shared_mem_bytes: 0,
14390 };
14391 let __s_b = self.gpu.stream();
14392 let mut b = __s_b.launch_builder(&f);
14393 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
14394 unsafe {
14395 b.launch(cfg)?;
14396 }
14397 Ok((q, d))
14398 }
14399
14400 #[allow(clippy::too_many_arguments)] pub fn rms_norm_q8_1_into(
14404 &self,
14405 x: &CudaSlice<f32>,
14406 w: &CudaSlice<f32>,
14407 ncols: usize,
14408 nrows: usize,
14409 eps: f32,
14410 q: &mut CudaSlice<i8>,
14411 d: &mut CudaSlice<f32>,
14412 ) -> Result<(), Box<dyn std::error::Error>> {
14413 let nblk = ncols / 32;
14414 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
14415 let (nc, e) = (ncols as i32, eps);
14416 if Self::pdl_on() {
14417 use cudarc::driver::{DevicePtr, DevicePtrMut};
14418 let s = &self.gpu.stream();
14419 let (px, _g0) = x.device_ptr(s);
14420 let (pw, _g1) = w.device_ptr(s);
14421 let (pq, _g2) = q.device_ptr_mut(s);
14422 let (pd, _g3) = d.device_ptr_mut(s);
14423 let mut ps = [
14424 &px as *const _ as *mut std::ffi::c_void,
14425 &pw as *const _ as *mut _,
14426 &pq as *const _ as *mut _,
14427 &pd as *const _ as *mut _,
14428 &nc as *const _ as *mut _,
14429 &e as *const _ as *mut _,
14430 ];
14431 unsafe {
14432 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
14433 }
14434 return Ok(());
14435 }
14436 let f = self.func("rms_norm_q8_1");
14437 let cfg = LaunchConfig {
14438 grid_dim: (nrows as u32, 1, 1),
14439 block_dim: (1024, 1, 1),
14440 shared_mem_bytes: 0,
14441 };
14442 let __s_b = self.gpu.stream();
14443 let mut b = __s_b.launch_builder(&f);
14444 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
14445 unsafe {
14446 b.launch(cfg)?;
14447 }
14448 Ok(())
14449 }
14450
14451 pub fn rms_norm_zq8_f32(
14464 &self,
14465 x: &CudaSlice<f32>,
14466 w: &CudaSlice<f32>,
14467 z: &mut CudaSlice<f32>,
14468 ncols: usize,
14469 nrows: usize,
14470 eps: f32,
14471 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14472 self.rms_norm_zq8_f32_arm(x, w, z, ncols, nrows, eps, Self::norm_ilp_zq8_on())
14473 }
14474
14475 pub fn hc_mixes_kernel_on() -> bool {
14480 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14481 *ON.get_or_init(|| std::env::var("MEMRA_HC_MIXES_KERNEL").as_deref() == Ok("1"))
14482 }
14483
14484 pub fn hc_mixes_gemv_into(
14488 &self,
14489 x: &cudarc::driver::CudaView<'_, f32>,
14490 w: &cudarc::driver::CudaView<'_, f32>,
14491 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
14492 in_f: usize,
14493 out_f: usize,
14494 ) -> Result<bool, Box<dyn std::error::Error>> {
14495 const BLOCK: usize = 1024;
14496 if in_f != 16 * BLOCK || out_f == 0 || out_f > 65535 {
14497 return Ok(false);
14498 }
14499 let f = self.func("hc_mixes_gemv_f32");
14500 let cfg = LaunchConfig {
14501 grid_dim: (out_f as u32, 1, 1),
14502 block_dim: (BLOCK as u32, 1, 1),
14503 shared_mem_bytes: 0,
14504 };
14505 let (inf, outf) = (in_f as i32, out_f as i32);
14506 let __s_b = self.gpu.stream();
14507 let mut b = __s_b.launch_builder(&f);
14508 b.arg(x).arg(w).arg(&mut *y).arg(&inf).arg(&outf);
14509 unsafe {
14510 b.launch(cfg)?;
14511 }
14512 Ok(true)
14513 }
14514
14515 pub fn gemv_f32_rows_into<I, O>(
14521 &self,
14522 x: &I,
14523 w: &I,
14524 y: &mut O,
14525 m: usize,
14526 in_f: usize,
14527 out_f: usize,
14528 ) -> Result<bool, Box<dyn std::error::Error>>
14529 where
14530 I: cudarc::driver::DevicePtr<f32>,
14531 O: cudarc::driver::DevicePtrMut<f32>,
14532 {
14533 const BLOCK: usize = 256;
14534 if in_f == 0
14535 || !in_f.is_multiple_of(4 * BLOCK)
14536 || out_f == 0
14537 || out_f > 65535
14538 || m == 0
14539 || m > 16
14540 {
14541 return Ok(false);
14542 }
14543 let f = self.func("gemv_f32_rows");
14544 let cfg = LaunchConfig {
14545 grid_dim: (out_f as u32, m as u32, 1),
14546 block_dim: (BLOCK as u32, 1, 1),
14547 shared_mem_bytes: 0,
14548 };
14549 let (inf, outf) = (in_f as i32, out_f as i32);
14550 let stream = self.gpu.stream();
14551 let (px, _gx) = x.device_ptr(&stream);
14553 let (pw, _gw) = w.device_ptr(&stream);
14554 let (py, _gy) = y.device_ptr_mut(&stream);
14555 let mut b = stream.launch_builder(&f);
14556 b.arg(&px).arg(&pw).arg(&py).arg(&inf).arg(&outf);
14557 unsafe {
14558 b.launch(cfg)?;
14559 }
14560 if F32_GEMV_KERNEL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
14561 eprintln!(
14562 "[f32-gemv-kernel] engaged: f32 linear at m<=16 rides gemv_f32_rows instead of \
14563 cuBLASLt (MEMRA_F32_GEMV_KERNEL=1; first shape {in_f}->{out_f} m={m})"
14564 );
14565 }
14566 Ok(true)
14567 }
14568
14569 pub fn side_pair(
14571 &self,
14572 ) -> Result<(Arc<CudaStream>, Arc<cudarc::cublaslt::CudaBlasLT>), Box<dyn std::error::Error>>
14573 {
14574 if let Some(p) = self.side.get() {
14575 return Ok(p.clone());
14576 }
14577 let stream = self.gpu.ctx.new_stream()?;
14578 let blas = Arc::new(cudarc::cublaslt::CudaBlasLT::new(stream.clone())?);
14579 let _ = self.side.set((stream, blas));
14580 Ok(self.side.get().expect("set above").clone())
14581 }
14582
14583 #[allow(clippy::too_many_arguments)]
14589 pub fn rms_norm_zq8_f32_arm(
14590 &self,
14591 x: &CudaSlice<f32>,
14592 w: &CudaSlice<f32>,
14593 z: &mut CudaSlice<f32>,
14594 ncols: usize,
14595 nrows: usize,
14596 eps: f32,
14597 ilp: bool,
14598 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14599 assert!(ncols.is_multiple_of(32));
14600 let nblk = ncols / 32;
14601 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?; let mut d = self.alloc_uninit::<f32>(nrows * nblk)?; let f = self.func(if ilp {
14604 "rms_norm_zq8_f32_v2"
14605 } else {
14606 "rms_norm_zq8_f32"
14607 });
14608 let cfg = LaunchConfig {
14612 grid_dim: (nrows as u32, 1, 1),
14613 block_dim: (rms_block(), 1, 1),
14614 shared_mem_bytes: 0,
14615 };
14616 let (nc, ep) = (ncols as i32, eps);
14617 let __s_b = self.gpu.stream();
14618 let mut b = __s_b.launch_builder(&f);
14619 b.arg(x)
14620 .arg(w)
14621 .arg(&mut *z)
14622 .arg(&mut q)
14623 .arg(&mut d)
14624 .arg(&nc)
14625 .arg(&ep);
14626 unsafe {
14627 b.launch(cfg)?;
14628 }
14629 Ok((q, d))
14630 }
14631
14632 #[track_caller]
14634 pub fn quantize_q8_1_into(
14635 &self,
14636 x: &CudaSlice<f32>,
14637 m: usize,
14638 in_f: usize,
14639 q: &mut CudaSlice<i8>,
14640 d: &mut CudaSlice<f32>,
14641 ) -> Result<(), Box<dyn std::error::Error>> {
14642 q8_census_record(std::panic::Location::caller(), m, in_f);
14643 let nblk = in_f / 32;
14644 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
14645 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
14646 let (inf, mi) = (in_f as i32, m as i32);
14647 if Self::pdl_on() && Self::pdl_wb_on() {
14648 use cudarc::driver::{DevicePtr, DevicePtrMut};
14649 let s = &self.gpu.stream();
14650 let (px, _g0) = x.device_ptr(s);
14651 let (pq, _g1) = q.device_ptr_mut(s);
14652 let (pd, _g2) = d.device_ptr_mut(s);
14653 let mut ps = [
14654 &px as *const _ as *mut std::ffi::c_void,
14655 &pq as *const _ as *mut _,
14656 &pd as *const _ as *mut _,
14657 &inf as *const _ as *mut _,
14658 &mi as *const _ as *mut _,
14659 ];
14660 unsafe {
14661 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
14662 }
14663 return Ok(());
14664 }
14665 let f = self.func("quantize_q8_1");
14666 let __s_b = self.gpu.stream();
14667 let mut b = __s_b.launch_builder(&f);
14668 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
14669 unsafe {
14670 b.launch(cfg)?;
14671 }
14672 Ok(())
14673 }
14674
14675 #[allow(clippy::too_many_arguments)] pub fn add_rms_norm_q8_1(
14680 &self,
14681 a: &CudaSlice<f32>,
14682 b_in: &CudaSlice<f32>,
14683 w: &CudaSlice<f32>,
14684 res: &mut CudaSlice<f32>,
14685 ncols: usize,
14686 nrows: usize,
14687 eps: f32,
14688 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14689 let nblk = ncols / 32;
14690 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
14691 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
14692 let f = self.func("add_rms_norm_q8_1");
14693 let cfg = LaunchConfig {
14695 grid_dim: (nrows as u32, 1, 1),
14696 block_dim: (1024, 1, 1),
14697 shared_mem_bytes: 0,
14698 };
14699 let (nc, e) = (ncols as i32, eps);
14700 let __s_bld = self.gpu.stream();
14701 let mut bld = __s_bld.launch_builder(&f);
14702 bld.arg(a)
14703 .arg(b_in)
14704 .arg(w)
14705 .arg(res)
14706 .arg(&mut q)
14707 .arg(&mut d)
14708 .arg(&nc)
14709 .arg(&e);
14710 unsafe {
14711 bld.launch(cfg)?;
14712 }
14713 Ok((q, d))
14714 }
14715
14716 #[allow(clippy::too_many_arguments)]
14722 pub fn join_add_rms_norm_raw(
14723 &self,
14724 a0_raw: u64,
14725 a1_raw: u64,
14726 x: &CudaSlice<f32>,
14727 w: &CudaSlice<f32>,
14728 res: &mut CudaSlice<f32>,
14729 dst: &mut CudaSlice<f32>,
14730 ncols: usize,
14731 eps: f32,
14732 ) -> Result<(), Box<dyn std::error::Error>> {
14733 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
14734 return Err("join_add_rms_norm geometry".into());
14735 }
14736 let f = self.func("join_add_rms_norm_f32");
14737 let cfg = LaunchConfig {
14738 grid_dim: (1, 1, 1),
14739 block_dim: (rms_block(), 1, 1),
14740 shared_mem_bytes: 0,
14741 };
14742 let (nc, e) = (ncols as i32, eps);
14743 let __s_b = self.gpu.stream();
14744 let mut b = __s_b.launch_builder(&f);
14745 b.arg(&a0_raw)
14746 .arg(&a1_raw)
14747 .arg(x)
14748 .arg(w)
14749 .arg(&mut *res)
14750 .arg(&mut *dst)
14751 .arg(&nc)
14752 .arg(&e);
14753 unsafe {
14754 b.launch(cfg)?;
14755 }
14756 Ok(())
14757 }
14758
14759 #[allow(clippy::too_many_arguments)] pub fn add_rms_norm(
14761 &self,
14762 a: &CudaSlice<f32>,
14763 b: &CudaSlice<f32>,
14764 w: &CudaSlice<f32>,
14765 res: &mut CudaSlice<f32>,
14766 dst: &mut CudaSlice<f32>,
14767 ncols: usize,
14768 nrows: usize,
14769 eps: f32,
14770 ) -> Result<(), Box<dyn std::error::Error>> {
14771 let (nc, e) = (ncols as i32, eps);
14772 let kname = if Self::norm_ilp_on() {
14773 "add_rms_norm_f32_v2"
14774 } else {
14775 "add_rms_norm_f32"
14776 };
14777 if Self::pdl_on() && Self::pdl_wb_on() {
14778 use cudarc::driver::{DevicePtr, DevicePtrMut};
14779 let s = &self.gpu.stream();
14780 let (pa, _g0) = a.device_ptr(s);
14781 let (pb, _g1) = b.device_ptr(s);
14782 let (pw, _g2) = w.device_ptr(s);
14783 let (pr, _g3) = res.device_ptr_mut(s);
14784 let (pd, _g4) = dst.device_ptr_mut(s);
14785 let mut ps = [
14786 &pa as *const _ as *mut std::ffi::c_void,
14787 &pb as *const _ as *mut _,
14788 &pw as *const _ as *mut _,
14789 &pr as *const _ as *mut _,
14790 &pd as *const _ as *mut _,
14791 &nc as *const _ as *mut _,
14792 &e as *const _ as *mut _,
14793 ];
14794 unsafe {
14795 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
14796 }
14797 return Ok(());
14798 }
14799 let f = self.func(kname);
14800 let cfg = LaunchConfig {
14801 grid_dim: (nrows as u32, 1, 1),
14802 block_dim: (rms_block(), 1, 1),
14803 shared_mem_bytes: 0,
14804 };
14805 let __s_b2 = self.gpu.stream();
14806 let mut b2 = __s_b2.launch_builder(&f);
14807 b2.arg(a)
14808 .arg(b)
14809 .arg(w)
14810 .arg(&mut *res)
14811 .arg(&mut *dst)
14812 .arg(&nc)
14813 .arg(&e);
14814 unsafe {
14815 b2.launch(cfg)?;
14816 }
14817 Ok(())
14818 }
14819
14820 #[allow(clippy::too_many_arguments)]
14823 pub fn rms_pre_add_rms_norm(
14824 &self,
14825 a: &CudaSlice<f32>,
14826 wa: &CudaSlice<f32>,
14827 b: &CudaSlice<f32>,
14828 w: &CudaSlice<f32>,
14829 res: &mut CudaSlice<f32>,
14830 dst: &mut CudaSlice<f32>,
14831 ncols: usize,
14832 nrows: usize,
14833 eps: f32,
14834 ) -> Result<(), Box<dyn std::error::Error>> {
14835 let f = self.func("rms_pre_add_rms_norm_f32");
14836 let cfg = LaunchConfig {
14837 grid_dim: (nrows as u32, 1, 1),
14838 block_dim: (rms_block(), 1, 1),
14839 shared_mem_bytes: 0,
14840 };
14841 let (nc, e) = (ncols as i32, eps);
14842 let __s_b2 = self.gpu.stream();
14843 let mut b2 = __s_b2.launch_builder(&f);
14844 b2.arg(a)
14845 .arg(wa)
14846 .arg(b)
14847 .arg(w)
14848 .arg(&mut *res)
14849 .arg(&mut *dst)
14850 .arg(&nc)
14851 .arg(&e);
14852 unsafe {
14853 b2.launch(cfg)?;
14854 }
14855 Ok(())
14856 }
14857
14858 #[allow(clippy::too_many_arguments)]
14860 pub fn rms_pre_add_rms_norm_q8z(
14861 &self,
14862 a: &CudaSlice<f32>,
14863 wa: &CudaSlice<f32>,
14864 b: &CudaSlice<f32>,
14865 w: &CudaSlice<f32>,
14866 res: &mut CudaSlice<f32>,
14867 dst: &mut CudaSlice<f32>,
14868 ncols: usize,
14869 nrows: usize,
14870 eps: f32,
14871 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14872 debug_assert!(ncols.is_multiple_of(128));
14873 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14874 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14875 let (nc, e) = (ncols as i32, eps);
14876 if Self::pdl_on() {
14877 {
14878 use cudarc::driver::{DevicePtr, DevicePtrMut};
14879 let s = &self.gpu.stream();
14880 let (pa, _g0) = a.device_ptr(s);
14881 let (pwa, _g1) = wa.device_ptr(s);
14882 let (pb, _g2) = b.device_ptr(s);
14883 let (pw, _g3) = w.device_ptr(s);
14884 let (pr, _g4) = res.device_ptr_mut(s);
14885 let (pdst, _g5) = dst.device_ptr_mut(s);
14886 let (pq, _g6) = out_q.device_ptr_mut(s);
14887 let (pd, _g7) = out_d.device_ptr_mut(s);
14888 let mut ps = [
14889 &pa as *const _ as *mut std::ffi::c_void,
14890 &pwa as *const _ as *mut _,
14891 &pb as *const _ as *mut _,
14892 &pw as *const _ as *mut _,
14893 &pr as *const _ as *mut _,
14894 &pdst as *const _ as *mut _,
14895 &pq as *const _ as *mut _,
14896 &pd as *const _ as *mut _,
14897 &nc as *const _ as *mut _,
14898 &e as *const _ as *mut _,
14899 ];
14900 unsafe {
14901 self.launch_pdl(
14902 "rms_pre_add_rms_norm_q8z_f32",
14903 (nrows as u32, 1, 1),
14904 (rms_block(), 1, 1),
14905 &mut ps,
14906 )?;
14907 }
14908 }
14909 return Ok((out_q, out_d));
14910 }
14911 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
14912 let cfg = LaunchConfig {
14913 grid_dim: (nrows as u32, 1, 1),
14914 block_dim: (rms_block(), 1, 1),
14915 shared_mem_bytes: 0,
14916 };
14917 let __s_b2 = self.gpu.stream();
14918 let mut b2 = __s_b2.launch_builder(&f);
14919 b2.arg(a)
14920 .arg(wa)
14921 .arg(b)
14922 .arg(w)
14923 .arg(&mut *res)
14924 .arg(&mut *dst)
14925 .arg(&mut out_q)
14926 .arg(&mut out_d)
14927 .arg(&nc)
14928 .arg(&e);
14929 unsafe {
14930 b2.launch(cfg)?;
14931 }
14932 Ok((out_q, out_d))
14933 }
14934
14935 #[allow(clippy::too_many_arguments)]
14939 pub fn rms_pre_add_rms_norm_q8z_into(
14940 &self,
14941 a: &CudaSlice<f32>,
14942 wa: &CudaSlice<f32>,
14943 b: &CudaSlice<f32>,
14944 w: &CudaSlice<f32>,
14945 res: &mut CudaSlice<f32>,
14946 dst: &mut CudaSlice<f32>,
14947 ncols: usize,
14948 nrows: usize,
14949 eps: f32,
14950 out_q: &mut CudaSlice<i8>,
14951 out_d: &mut CudaSlice<f32>,
14952 ) -> Result<(), Box<dyn std::error::Error>> {
14953 debug_assert!(ncols.is_multiple_of(128));
14954 let (nc, e) = (ncols as i32, eps);
14955 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
14956 let cfg = LaunchConfig {
14957 grid_dim: (nrows as u32, 1, 1),
14958 block_dim: (rms_block(), 1, 1),
14959 shared_mem_bytes: 0,
14960 };
14961 let __s_b = self.gpu.stream();
14962 let mut b2 = __s_b.launch_builder(&f);
14963 b2.arg(a)
14964 .arg(wa)
14965 .arg(b)
14966 .arg(w)
14967 .arg(&mut *res)
14968 .arg(&mut *dst)
14969 .arg(&mut *out_q)
14970 .arg(&mut *out_d)
14971 .arg(&nc)
14972 .arg(&e);
14973 unsafe {
14974 b2.launch(cfg)?;
14975 }
14976 Ok(())
14977 }
14978
14979 #[allow(clippy::too_many_arguments)]
14982 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
14983 &self,
14984 a: &CudaSlice<f32>,
14985 wa: &CudaSlice<f32>,
14986 b_in: &CudaSlice<f32>,
14987 c: f32,
14988 w: &CudaSlice<f32>,
14989 res: &mut CudaSlice<f32>,
14990 ncols: usize,
14991 nrows: usize,
14992 eps: f32,
14993 out_q: &mut CudaSlice<i8>,
14994 out_d: &mut CudaSlice<f32>,
14995 ) -> Result<(), Box<dyn std::error::Error>> {
14996 debug_assert!(ncols.is_multiple_of(128));
14997 let (nc, e2) = (ncols as i32, eps);
14998 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
14999 let cfg = LaunchConfig {
15000 grid_dim: (nrows as u32, 1, 1),
15001 block_dim: (rms_block(), 1, 1),
15002 shared_mem_bytes: 0,
15003 };
15004 let __s_b = self.gpu.stream();
15005 let mut b2 = __s_b.launch_builder(&f);
15006 b2.arg(a)
15007 .arg(wa)
15008 .arg(b_in)
15009 .arg(&c)
15010 .arg(w)
15011 .arg(&mut *res)
15012 .arg(&mut *out_q)
15013 .arg(&mut *out_d)
15014 .arg(&nc)
15015 .arg(&e2);
15016 unsafe {
15017 b2.launch(cfg)?;
15018 }
15019 Ok(())
15020 }
15021
15022 pub fn g4_pnfold_on() -> bool {
15030 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15031 *ON.get_or_init(|| {
15032 std::env::var("MEMRA_G4_PNFOLD")
15033 .map(|v| v != "0")
15034 .unwrap_or(true)
15035 })
15036 }
15037
15038 pub fn build_q4_out_concat3(
15042 &self,
15043 w0: &crate::model::GpuTensor,
15044 w1: &crate::model::GpuTensor,
15045 w2: &crate::model::GpuTensor,
15046 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
15047 use crate::model::GpuTensor;
15048 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
15049 match w {
15050 GpuTensor::Quant {
15051 qtype,
15052 row_bytes,
15053 rp,
15054 ..
15055 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
15056 _ => None,
15057 }
15058 };
15059 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
15060 else {
15061 return Ok(None);
15062 };
15063 if rb0 != rb1
15064 || rb0 != rb2
15065 || w0.in_features() != w1.in_features()
15066 || w0.in_features() != w2.in_features()
15067 {
15068 return Ok(None);
15069 }
15070 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
15071 match w {
15072 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
15073 _ => unreachable!(),
15074 }
15075 }
15076 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
15077 let total = rb0 * (o0 + o1 + o2);
15078 let mut cat = self.alloc_u8(total)?;
15079 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
15080 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
15081 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
15082 Ok(Some(GpuTensor::Quant {
15083 bytes: cat,
15084 qtype: QT_Q4_0,
15085 row_bytes: rb0,
15086 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
15087 scale: 1.0,
15088 rp: false,
15089 #[cfg(memra_cutlass)]
15090 cutlass: None,
15091 fp8: None,
15092 blk: None,
15093 rp4: None,
15094 f16: None,
15095 }))
15096 }
15097
15098 fn full_width_rope_only(
15116 kernel: &str,
15117 n_rot: usize,
15118 head_dim: usize,
15119 ) -> Result<(), Box<dyn std::error::Error>> {
15120 if n_rot == head_dim {
15121 return Ok(());
15122 }
15123 Err(format!(
15124 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
15125 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
15126 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
15127 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
15128 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
15129 )
15130 .into())
15131 }
15132
15133 #[allow(clippy::too_many_arguments)]
15137 pub fn rms_norm_qkv_rope_cat(
15138 &self,
15139 qkv: &CudaSlice<f32>,
15140 wq: &CudaSlice<f32>,
15141 wk: &CudaSlice<f32>,
15142 wv: &CudaSlice<f32>,
15143 q: &mut CudaSlice<f32>,
15144 k: &mut CudaSlice<f32>,
15145 v: &mut CudaSlice<f32>,
15146 head_dim: usize,
15147 n_rot: usize,
15148 rq: usize,
15149 rk: usize,
15150 pos: &CudaSlice<i32>,
15151 nh_q: usize,
15152 nh_k: usize,
15153 base: f32,
15154 freq_scale: f32,
15155 ff: Option<&CudaSlice<f32>>,
15156 eps: f32,
15157 ) -> Result<(), Box<dyn std::error::Error>> {
15158 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
15159 let rows = rq + rk + rk;
15160 let theta_scale = base.powf(-2.0 / head_dim as f32);
15161 let (nc, rqi, rki, nhq, nhk) = (
15162 head_dim as i32,
15163 rq as i32,
15164 rk as i32,
15165 nh_q as i32,
15166 nh_k as i32,
15167 );
15168 if Self::pdl_on() {
15169 use cudarc::driver::{DevicePtr, DevicePtrMut};
15170 let s = &self.gpu.stream();
15171 let (pqkv, _g0) = qkv.device_ptr(s);
15172 let (pwq, _g1) = wq.device_ptr(s);
15173 let (pwk, _g2) = wk.device_ptr(s);
15174 let (pwv, _g3) = wv.device_ptr(s);
15175 let (pq, _g4) = q.device_ptr_mut(s);
15176 let (pk, _g5) = k.device_ptr_mut(s);
15177 let (pv, _g6) = v.device_ptr_mut(s);
15178 let (ppos, _g7) = pos.device_ptr(s);
15179 let (pff, _g8) = match ff {
15180 Some(t) => {
15181 let (p, g) = t.device_ptr(s);
15182 (p, Some(g))
15183 }
15184 None => (0, None),
15185 };
15186 let mut ps = [
15187 &pqkv as *const _ as *mut std::ffi::c_void,
15188 &pwq as *const _ as *mut _,
15189 &pwk as *const _ as *mut _,
15190 &pwv as *const _ as *mut _,
15191 &pq as *const _ as *mut _,
15192 &pk as *const _ as *mut _,
15193 &pv as *const _ as *mut _,
15194 &nc as *const _ as *mut _,
15195 &rqi as *const _ as *mut _,
15196 &rki as *const _ as *mut _,
15197 &ppos as *const _ as *mut _,
15198 &nhq as *const _ as *mut _,
15199 &nhk as *const _ as *mut _,
15200 &theta_scale as *const _ as *mut _,
15201 &freq_scale as *const _ as *mut _,
15202 &pff as *const _ as *mut _,
15203 &eps as *const _ as *mut _,
15204 ];
15205 unsafe {
15206 self.launch_pdl(
15207 "rms_norm_qkv_rope_cat_f32",
15208 (rows as u32, 1, 1),
15209 (rms_block(), 1, 1),
15210 &mut ps,
15211 )?;
15212 }
15213 return Ok(());
15214 }
15215 let f = self.func("rms_norm_qkv_rope_cat_f32");
15216 let cfg = LaunchConfig {
15217 grid_dim: (rows as u32, 1, 1),
15218 block_dim: (rms_block(), 1, 1),
15219 shared_mem_bytes: 0,
15220 };
15221 let __s_b = self.gpu.stream();
15222 let mut b = __s_b.launch_builder(&f);
15223 match ff {
15224 Some(t) => {
15225 b.arg(qkv)
15226 .arg(wq)
15227 .arg(wk)
15228 .arg(wv)
15229 .arg(&mut *q)
15230 .arg(&mut *k)
15231 .arg(&mut *v)
15232 .arg(&nc)
15233 .arg(&rqi)
15234 .arg(&rki)
15235 .arg(pos)
15236 .arg(&nhq)
15237 .arg(&nhk)
15238 .arg(&theta_scale)
15239 .arg(&freq_scale)
15240 .arg(t)
15241 .arg(&eps);
15242 unsafe {
15243 b.launch(cfg)?;
15244 }
15245 }
15246 None => {
15247 let null: u64 = 0;
15248 b.arg(qkv)
15249 .arg(wq)
15250 .arg(wk)
15251 .arg(wv)
15252 .arg(&mut *q)
15253 .arg(&mut *k)
15254 .arg(&mut *v)
15255 .arg(&nc)
15256 .arg(&rqi)
15257 .arg(&rki)
15258 .arg(pos)
15259 .arg(&nhq)
15260 .arg(&nhk)
15261 .arg(&theta_scale)
15262 .arg(&freq_scale)
15263 .arg(&null)
15264 .arg(&eps);
15265 unsafe {
15266 b.launch(cfg)?;
15267 }
15268 }
15269 }
15270 Ok(())
15271 }
15272
15273 #[allow(clippy::too_many_arguments)]
15277 pub fn rms_norm_qkv_rope(
15278 &self,
15279 q0: &CudaSlice<f32>,
15280 k0: &CudaSlice<f32>,
15281 v0: &CudaSlice<f32>,
15282 wq: &CudaSlice<f32>,
15283 wk: &CudaSlice<f32>,
15284 wv: &CudaSlice<f32>,
15285 q: &mut CudaSlice<f32>,
15286 k: &mut CudaSlice<f32>,
15287 v: &mut CudaSlice<f32>,
15288 head_dim: usize,
15289 n_rot: usize,
15290 rq: usize,
15291 rk: usize,
15292 pos: &CudaSlice<i32>,
15293 nh_q: usize,
15294 nh_k: usize,
15295 base: f32,
15296 freq_scale: f32,
15297 ff: Option<&CudaSlice<f32>>,
15298 eps: f32,
15299 ) -> Result<(), Box<dyn std::error::Error>> {
15300 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
15301 let f = self.func("rms_norm_qkv_rope_f32");
15302 let rows = rq + rk + rk; let cfg = LaunchConfig {
15304 grid_dim: (rows as u32, 1, 1),
15305 block_dim: (rms_block(), 1, 1),
15306 shared_mem_bytes: 0,
15307 };
15308 let theta_scale = base.powf(-2.0 / head_dim as f32);
15309 let (nc, rqi, rki, nhq, nhk) = (
15310 head_dim as i32,
15311 rq as i32,
15312 rk as i32,
15313 nh_q as i32,
15314 nh_k as i32,
15315 );
15316 let __s_b = self.gpu.stream();
15317 let mut b = __s_b.launch_builder(&f);
15318 match ff {
15319 Some(t) => {
15320 b.arg(q0)
15321 .arg(k0)
15322 .arg(v0)
15323 .arg(wq)
15324 .arg(wk)
15325 .arg(wv)
15326 .arg(&mut *q)
15327 .arg(&mut *k)
15328 .arg(&mut *v)
15329 .arg(&nc)
15330 .arg(&rqi)
15331 .arg(&rki)
15332 .arg(pos)
15333 .arg(&nhq)
15334 .arg(&nhk)
15335 .arg(&theta_scale)
15336 .arg(&freq_scale)
15337 .arg(t)
15338 .arg(&eps);
15339 unsafe {
15340 b.launch(cfg)?;
15341 }
15342 }
15343 None => {
15344 let null: u64 = 0;
15345 b.arg(q0)
15346 .arg(k0)
15347 .arg(v0)
15348 .arg(wq)
15349 .arg(wk)
15350 .arg(wv)
15351 .arg(&mut *q)
15352 .arg(&mut *k)
15353 .arg(&mut *v)
15354 .arg(&nc)
15355 .arg(&rqi)
15356 .arg(&rki)
15357 .arg(pos)
15358 .arg(&nhq)
15359 .arg(&nhk)
15360 .arg(&theta_scale)
15361 .arg(&freq_scale)
15362 .arg(&null)
15363 .arg(&eps);
15364 unsafe {
15365 b.launch(cfg)?;
15366 }
15367 }
15368 }
15369 Ok(())
15370 }
15371
15372 #[allow(clippy::too_many_arguments)]
15378 pub fn rms_norm_qkv_rope_append_dc(
15379 &self,
15380 q0: &CudaSlice<f32>,
15381 k0: &CudaSlice<f32>,
15382 v0: &CudaSlice<f32>,
15383 wq: &CudaSlice<f32>,
15384 wk: &CudaSlice<f32>,
15385 wv: &CudaSlice<f32>,
15386 q: &mut CudaSlice<f32>,
15387 k: &mut CudaSlice<f32>,
15388 v: &mut CudaSlice<f32>,
15389 head_dim: usize,
15390 n_rot: usize,
15391 rq: usize,
15392 rk: usize,
15393 pos: &CudaSlice<i32>,
15394 nh_q: usize,
15395 nh_k: usize,
15396 base: f32,
15397 freq_scale: f32,
15398 ff: Option<&CudaSlice<f32>>,
15399 eps: f32,
15400 kc: &mut CudaSlice<u8>,
15401 vc: &mut CudaSlice<u8>,
15402 t_dev: &CudaSlice<i32>,
15403 k_tok_bytes: usize,
15404 v_tok_bytes: usize,
15405 g: bool,
15406 ) -> Result<(), Box<dyn std::error::Error>> {
15407 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
15408 let rows = rq + rk + rk;
15409 let theta_scale = base.powf(-2.0 / head_dim as f32);
15410 let (nc, rqi, rki, nhq, nhk) = (
15411 head_dim as i32,
15412 rq as i32,
15413 rk as i32,
15414 nh_q as i32,
15415 nh_k as i32,
15416 );
15417 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15418 if Self::pdl_on() && Self::pdl_wb_on() {
15419 use cudarc::driver::{DevicePtr, DevicePtrMut};
15420 let s = &self.gpu.stream();
15421 let (p0, _a0) = q0.device_ptr(s);
15422 let (p1, _a1) = k0.device_ptr(s);
15423 let (p2, _a2) = v0.device_ptr(s);
15424 let (pwq, _a3) = wq.device_ptr(s);
15425 let (pwk, _a4) = wk.device_ptr(s);
15426 let (pwv, _a5) = wv.device_ptr(s);
15427 let (pq, _a6) = q.device_ptr_mut(s);
15428 let (pk, _a7) = k.device_ptr_mut(s);
15429 let (pv, _a8) = v.device_ptr_mut(s);
15430 let (pp, _a9) = pos.device_ptr(s);
15431 let pff: u64 = match ff {
15432 Some(t) => {
15433 let (p, _gg) = t.device_ptr(s);
15434 p
15435 }
15436 None => 0,
15437 };
15438 let (pkc, _a10) = kc.device_ptr_mut(s);
15439 let (pvc, _a11) = vc.device_ptr_mut(s);
15440 let (pt, _a12) = t_dev.device_ptr(s);
15441 let mut ps = [
15442 &p0 as *const _ as *mut std::ffi::c_void,
15443 &p1 as *const _ as *mut _,
15444 &p2 as *const _ as *mut _,
15445 &pwq as *const _ as *mut _,
15446 &pwk as *const _ as *mut _,
15447 &pwv as *const _ as *mut _,
15448 &pq as *const _ as *mut _,
15449 &pk as *const _ as *mut _,
15450 &pv as *const _ as *mut _,
15451 &nc as *const _ as *mut _,
15452 &rqi as *const _ as *mut _,
15453 &rki as *const _ as *mut _,
15454 &pp as *const _ as *mut _,
15455 &nhq as *const _ as *mut _,
15456 &nhk as *const _ as *mut _,
15457 &theta_scale as *const _ as *mut _,
15458 &freq_scale as *const _ as *mut _,
15459 &pff as *const _ as *mut _,
15460 &eps as *const _ as *mut _,
15461 &pkc as *const _ as *mut _,
15462 &pvc as *const _ as *mut _,
15463 &pt as *const _ as *mut _,
15464 &ktb as *const _ as *mut _,
15465 &vtb as *const _ as *mut _,
15466 ];
15467 unsafe {
15468 self.launch_pdl_flash(
15469 g,
15470 "rms_norm_qkv_rope_append_dc_f32",
15471 (rows as u32, 1, 1),
15472 (rms_block(), 1, 1),
15473 0,
15474 &mut ps,
15475 )?;
15476 }
15477 return Ok(());
15478 }
15479 let f = if g {
15480 self.func_g("rms_norm_qkv_rope_append_dc_f32")
15481 } else {
15482 self.func("rms_norm_qkv_rope_append_dc_f32")
15483 };
15484 let cfg = LaunchConfig {
15485 grid_dim: (rows as u32, 1, 1),
15486 block_dim: (rms_block(), 1, 1),
15487 shared_mem_bytes: 0,
15488 };
15489 let __s_b = self.gpu.stream();
15490 let mut b = __s_b.launch_builder(&f);
15491 match ff {
15492 Some(t) => {
15493 b.arg(q0)
15494 .arg(k0)
15495 .arg(v0)
15496 .arg(wq)
15497 .arg(wk)
15498 .arg(wv)
15499 .arg(&mut *q)
15500 .arg(&mut *k)
15501 .arg(&mut *v)
15502 .arg(&nc)
15503 .arg(&rqi)
15504 .arg(&rki)
15505 .arg(pos)
15506 .arg(&nhq)
15507 .arg(&nhk)
15508 .arg(&theta_scale)
15509 .arg(&freq_scale)
15510 .arg(t)
15511 .arg(&eps)
15512 .arg(&mut *kc)
15513 .arg(&mut *vc)
15514 .arg(t_dev)
15515 .arg(&ktb)
15516 .arg(&vtb);
15517 unsafe {
15518 b.launch(cfg)?;
15519 }
15520 }
15521 None => {
15522 let null: u64 = 0;
15523 b.arg(q0)
15524 .arg(k0)
15525 .arg(v0)
15526 .arg(wq)
15527 .arg(wk)
15528 .arg(wv)
15529 .arg(&mut *q)
15530 .arg(&mut *k)
15531 .arg(&mut *v)
15532 .arg(&nc)
15533 .arg(&rqi)
15534 .arg(&rki)
15535 .arg(pos)
15536 .arg(&nhq)
15537 .arg(&nhk)
15538 .arg(&theta_scale)
15539 .arg(&freq_scale)
15540 .arg(&null)
15541 .arg(&eps)
15542 .arg(&mut *kc)
15543 .arg(&mut *vc)
15544 .arg(t_dev)
15545 .arg(&ktb)
15546 .arg(&vtb);
15547 unsafe {
15548 b.launch(cfg)?;
15549 }
15550 }
15551 }
15552 Ok(())
15553 }
15554
15555 #[allow(clippy::too_many_arguments)]
15563 pub fn rms_norm_qkv_rope_append(
15564 &self,
15565 q0: &CudaSlice<f32>,
15566 k0: &CudaSlice<f32>,
15567 v0: &CudaSlice<f32>,
15568 wq: &CudaSlice<f32>,
15569 wk: &CudaSlice<f32>,
15570 wv: &CudaSlice<f32>,
15571 q: &mut CudaSlice<f32>,
15572 k: &mut CudaSlice<f32>,
15573 v: &mut CudaSlice<f32>,
15574 head_dim: usize,
15575 n_rot: usize,
15576 rq: usize,
15577 rk: usize,
15578 pos: &CudaSlice<i32>,
15579 nh_q: usize,
15580 nh_k: usize,
15581 base: f32,
15582 freq_scale: f32,
15583 ff: Option<&CudaSlice<f32>>,
15584 eps: f32,
15585 kc: &mut CudaSlice<u8>,
15586 vc: &mut CudaSlice<u8>,
15587 t: usize,
15588 k_tok_bytes: usize,
15589 v_tok_bytes: usize,
15590 g: bool,
15591 ) -> Result<(), Box<dyn std::error::Error>> {
15592 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
15593 let rows = rq + rk + rk;
15594 let theta_scale = base.powf(-2.0 / head_dim as f32);
15595 let (nc, rqi, rki, nhq, nhk) = (
15596 head_dim as i32,
15597 rq as i32,
15598 rk as i32,
15599 nh_q as i32,
15600 nh_k as i32,
15601 );
15602 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15603 let ti = t as i32;
15604 if Self::pdl_on() && Self::pdl_wb_on() {
15605 use cudarc::driver::{DevicePtr, DevicePtrMut};
15606 let s = &self.gpu.stream();
15607 let (p0, _a0) = q0.device_ptr(s);
15608 let (p1, _a1) = k0.device_ptr(s);
15609 let (p2, _a2) = v0.device_ptr(s);
15610 let (pwq, _a3) = wq.device_ptr(s);
15611 let (pwk, _a4) = wk.device_ptr(s);
15612 let (pwv, _a5) = wv.device_ptr(s);
15613 let (pq, _a6) = q.device_ptr_mut(s);
15614 let (pk, _a7) = k.device_ptr_mut(s);
15615 let (pv, _a8) = v.device_ptr_mut(s);
15616 let (pp, _a9) = pos.device_ptr(s);
15617 let pff: u64 = match ff {
15618 Some(t) => {
15619 let (p, _gg) = t.device_ptr(s);
15620 p
15621 }
15622 None => 0,
15623 };
15624 let (pkc, _a10) = kc.device_ptr_mut(s);
15625 let (pvc, _a11) = vc.device_ptr_mut(s);
15626 let mut ps = [
15627 &p0 as *const _ as *mut std::ffi::c_void,
15628 &p1 as *const _ as *mut _,
15629 &p2 as *const _ as *mut _,
15630 &pwq as *const _ as *mut _,
15631 &pwk as *const _ as *mut _,
15632 &pwv as *const _ as *mut _,
15633 &pq as *const _ as *mut _,
15634 &pk as *const _ as *mut _,
15635 &pv as *const _ as *mut _,
15636 &nc as *const _ as *mut _,
15637 &rqi as *const _ as *mut _,
15638 &rki as *const _ as *mut _,
15639 &pp as *const _ as *mut _,
15640 &nhq as *const _ as *mut _,
15641 &nhk as *const _ as *mut _,
15642 &theta_scale as *const _ as *mut _,
15643 &freq_scale as *const _ as *mut _,
15644 &pff as *const _ as *mut _,
15645 &eps as *const _ as *mut _,
15646 &pkc as *const _ as *mut _,
15647 &pvc as *const _ as *mut _,
15648 &ti as *const _ as *mut _,
15649 &ktb as *const _ as *mut _,
15650 &vtb as *const _ as *mut _,
15651 ];
15652 unsafe {
15653 self.launch_pdl_flash(
15654 g,
15655 "rms_norm_qkv_rope_append_f32",
15656 (rows as u32, 1, 1),
15657 (rms_block(), 1, 1),
15658 0,
15659 &mut ps,
15660 )?;
15661 }
15662 return Ok(());
15663 }
15664 let f = if g {
15665 self.func_g("rms_norm_qkv_rope_append_f32")
15666 } else {
15667 self.func("rms_norm_qkv_rope_append_f32")
15668 };
15669 let cfg = LaunchConfig {
15670 grid_dim: (rows as u32, 1, 1),
15671 block_dim: (rms_block(), 1, 1),
15672 shared_mem_bytes: 0,
15673 };
15674 let __s_b = self.gpu.stream();
15675 let mut b = __s_b.launch_builder(&f);
15676 let null: u64 = 0;
15677 b.arg(q0)
15678 .arg(k0)
15679 .arg(v0)
15680 .arg(wq)
15681 .arg(wk)
15682 .arg(wv)
15683 .arg(&mut *q)
15684 .arg(&mut *k)
15685 .arg(&mut *v)
15686 .arg(&nc)
15687 .arg(&rqi)
15688 .arg(&rki)
15689 .arg(pos)
15690 .arg(&nhq)
15691 .arg(&nhk)
15692 .arg(&theta_scale)
15693 .arg(&freq_scale);
15694 match ff {
15695 Some(t) => {
15696 b.arg(t);
15697 }
15698 None => {
15699 b.arg(&null);
15700 }
15701 }
15702 b.arg(&eps)
15703 .arg(&mut *kc)
15704 .arg(&mut *vc)
15705 .arg(&ti)
15706 .arg(&ktb)
15707 .arg(&vtb);
15708 unsafe {
15709 b.launch(cfg)?;
15710 }
15711 Ok(())
15712 }
15713
15714 pub fn add_q8_1(
15715 &self,
15716 a: &CudaSlice<f32>,
15717 b: &CudaSlice<f32>,
15718 res: &mut CudaSlice<f32>,
15719 ncols: usize,
15720 nrows: usize,
15721 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15722 debug_assert!(ncols.is_multiple_of(128));
15723 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
15724 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
15725 let f = self.func("add_q8_1_f32");
15726 let cfg = LaunchConfig {
15727 grid_dim: (nrows as u32, 1, 1),
15728 block_dim: (rms_block(), 1, 1),
15729 shared_mem_bytes: 0,
15730 };
15731 let nc = ncols as i32;
15732 let __s_b2 = self.gpu.stream();
15733 let mut b2 = __s_b2.launch_builder(&f);
15734 b2.arg(a)
15735 .arg(b)
15736 .arg(&mut *res)
15737 .arg(&mut out_q)
15738 .arg(&mut out_d)
15739 .arg(&nc);
15740 unsafe {
15741 b2.launch(cfg)?;
15742 }
15743 Ok((out_q, out_d))
15744 }
15745
15746 #[allow(clippy::too_many_arguments)] pub fn rms_pre_add_q8_1(
15751 &self,
15752 a: &CudaSlice<f32>,
15753 wa: &CudaSlice<f32>,
15754 b: &CudaSlice<f32>,
15755 res: &mut CudaSlice<f32>,
15756 ncols: usize,
15757 nrows: usize,
15758 eps: f32,
15759 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15760 debug_assert!(ncols.is_multiple_of(128));
15761 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
15762 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
15763 let f = self.func("rms_pre_add_q8_1_f32");
15764 let cfg = LaunchConfig {
15765 grid_dim: (nrows as u32, 1, 1),
15766 block_dim: (rms_block(), 1, 1),
15767 shared_mem_bytes: 0,
15768 };
15769 let (nc, ep) = (ncols as i32, eps);
15770 let __s_b2 = self.gpu.stream();
15771 let mut b2 = __s_b2.launch_builder(&f);
15772 b2.arg(a)
15773 .arg(wa)
15774 .arg(b)
15775 .arg(&mut *res)
15776 .arg(&mut out_q)
15777 .arg(&mut out_d)
15778 .arg(&nc)
15779 .arg(&ep);
15780 unsafe {
15781 b2.launch(cfg)?;
15782 }
15783 Ok((out_q, out_d))
15784 }
15785
15786 pub fn l2_v2_on(ncols: usize) -> bool {
15790 ncols == 128
15791 }
15792
15793 pub fn l2_norm_pp(
15794 &self,
15795 x: &CudaSlice<f32>,
15796 dst: &mut CudaSlice<f32>,
15797 dst16: Option<&mut CudaSlice<u8>>,
15798 ncols: usize,
15799 nrows: usize,
15800 eps: f32,
15801 ) -> Result<(), Box<dyn std::error::Error>> {
15802 if Self::l2_v2_on(ncols) {
15803 let f = self.func("l2_norm_pp_v2_f32");
15804 let rows_per_block = 8u32; let cfg = LaunchConfig {
15806 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
15807 block_dim: (256, 1, 1),
15808 shared_mem_bytes: 0,
15809 };
15810 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
15811 let d16: u64 = match dst16 {
15813 Some(d) => self.addr_u8(d),
15814 None => 0,
15815 };
15816 let __s_b = self.gpu.stream();
15817 let mut b = __s_b.launch_builder(&f);
15818 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
15819 unsafe {
15820 b.launch(cfg)?;
15821 }
15822 return Ok(());
15823 }
15824 self.l2_norm(x, dst, ncols, nrows, eps)
15825 }
15826
15827 pub fn l2_norm(
15828 &self,
15829 x: &CudaSlice<f32>,
15830 dst: &mut CudaSlice<f32>,
15831 ncols: usize,
15832 nrows: usize,
15833 eps: f32,
15834 ) -> Result<(), Box<dyn std::error::Error>> {
15835 let f = self.func("l2_norm_f32");
15836 let cfg = LaunchConfig {
15837 grid_dim: (nrows as u32, 1, 1),
15838 block_dim: (256, 1, 1),
15839 shared_mem_bytes: 0,
15840 };
15841 let (nc, e) = (ncols as i32, eps);
15842 let __s_b = self.gpu.stream();
15843 let mut b = __s_b.launch_builder(&f);
15844 b.arg(x).arg(dst).arg(&nc).arg(&e);
15845 unsafe {
15846 b.launch(cfg)?;
15847 }
15848 Ok(())
15849 }
15850
15851 pub fn l2_norm_decode(
15857 &self,
15858 x: &CudaSlice<f32>,
15859 dst: &mut CudaSlice<f32>,
15860 ncols: usize,
15861 nrows: usize,
15862 eps: f32,
15863 ) -> Result<(), Box<dyn std::error::Error>> {
15864 let f = self.func("l2_norm_f32");
15865 let cfg = LaunchConfig {
15866 grid_dim: (nrows as u32, 1, 1),
15867 block_dim: (32, 1, 1),
15868 shared_mem_bytes: 0,
15869 };
15870 let (nc, e) = (ncols as i32, eps);
15871 let __s_b = self.gpu.stream();
15872 let mut b = __s_b.launch_builder(&f);
15873 b.arg(x).arg(dst).arg(&nc).arg(&e);
15874 unsafe {
15875 b.launch(cfg)?;
15876 }
15877 Ok(())
15878 }
15879
15880 #[allow(clippy::too_many_arguments)] pub fn rope_neox(
15883 &self,
15884 x: &mut CudaSlice<f32>,
15885 pos: &CudaSlice<i32>,
15886 head_dim: usize,
15887 n_dims: usize,
15888 n_heads: usize,
15889 n_tokens: usize,
15890 freq_base: f32,
15891 freq_scale: f32,
15892 ) -> Result<(), Box<dyn std::error::Error>> {
15893 let f = self.func("rope_neox_f32");
15894 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
15895 let grid = (n_heads * n_tokens) as u32;
15896 let cfg = LaunchConfig {
15897 grid_dim: (grid, 1, 1),
15898 block_dim: ((head_dim / 2) as u32, 1, 1),
15899 shared_mem_bytes: 0,
15900 };
15901 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
15902 let __s_b = self.gpu.stream();
15903 let mut b = __s_b.launch_builder(&f);
15904 b.arg(x)
15905 .arg(pos)
15906 .arg(&hd)
15907 .arg(&nd)
15908 .arg(&nh)
15909 .arg(&theta_scale)
15910 .arg(&freq_scale);
15911 unsafe {
15912 b.launch(cfg)?;
15913 }
15914 Ok(())
15915 }
15916
15917 #[allow(clippy::too_many_arguments)] pub fn rope_neox_ff(
15920 &self,
15921 x: &mut CudaSlice<f32>,
15922 pos: &CudaSlice<i32>,
15923 head_dim: usize,
15924 n_dims: usize,
15925 n_heads: usize,
15926 n_tokens: usize,
15927 freq_base: f32,
15928 freq_scale: f32,
15929 ff: &CudaSlice<f32>,
15930 ) -> Result<(), Box<dyn std::error::Error>> {
15931 let f = self.func("rope_neox_ff_f32");
15932 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
15933 let grid = (n_heads * n_tokens) as u32;
15934 let cfg = LaunchConfig {
15935 grid_dim: (grid, 1, 1),
15936 block_dim: ((head_dim / 2) as u32, 1, 1),
15937 shared_mem_bytes: 0,
15938 };
15939 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
15940 let __s_b = self.gpu.stream();
15941 let mut b = __s_b.launch_builder(&f);
15942 b.arg(x)
15943 .arg(pos)
15944 .arg(&hd)
15945 .arg(&nd)
15946 .arg(&nh)
15947 .arg(&theta_scale)
15948 .arg(&freq_scale)
15949 .arg(ff);
15950 unsafe {
15951 b.launch(cfg)?;
15952 }
15953 Ok(())
15954 }
15955
15956 #[allow(clippy::too_many_arguments)]
15960 pub fn rope_neox_ffm(
15961 &self,
15962 x: &mut CudaSlice<f32>,
15963 pos: &CudaSlice<i32>,
15964 head_dim: usize,
15965 n_dims: usize,
15966 n_heads: usize,
15967 n_tokens: usize,
15968 freq_base: f32,
15969 freq_scale: f32,
15970 ff: &CudaSlice<f32>,
15971 mscale: f32,
15972 ) -> Result<(), Box<dyn std::error::Error>> {
15973 let f = self.func("rope_neox_ffm_f32");
15974 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
15975 let grid = (n_heads * n_tokens) as u32;
15976 let cfg = LaunchConfig {
15977 grid_dim: (grid, 1, 1),
15978 block_dim: ((head_dim / 2) as u32, 1, 1),
15979 shared_mem_bytes: 0,
15980 };
15981 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
15982 let __s_b = self.gpu.stream();
15983 let mut b = __s_b.launch_builder(&f);
15984 b.arg(x)
15985 .arg(pos)
15986 .arg(&hd)
15987 .arg(&nd)
15988 .arg(&nh)
15989 .arg(&theta_scale)
15990 .arg(&freq_scale)
15991 .arg(ff)
15992 .arg(&mscale);
15993 unsafe {
15994 b.launch(cfg)?;
15995 }
15996 Ok(())
15997 }
15998
15999 #[allow(clippy::too_many_arguments)]
16001 pub fn rope_neox2(
16002 &self,
16003 q: &mut CudaSlice<f32>,
16004 k: &mut CudaSlice<f32>,
16005 pos: &CudaSlice<i32>,
16006 head_dim: usize,
16007 n_dims: usize,
16008 nh_q: usize,
16009 nh_k: usize,
16010 n_tokens: usize,
16011 freq_base: f32,
16012 freq_scale: f32,
16013 ff: Option<&CudaSlice<f32>>,
16014 ) -> Result<(), Box<dyn std::error::Error>> {
16015 let f = self.func("rope_neox2_f32");
16016 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16017 let grid = ((nh_q + nh_k) * n_tokens) as u32;
16018 let cfg = LaunchConfig {
16019 grid_dim: (grid, 1, 1),
16020 block_dim: ((head_dim / 2) as u32, 1, 1),
16021 shared_mem_bytes: 0,
16022 };
16023 let (hd, nd, nq, nk, nt) = (
16024 head_dim as i32,
16025 n_dims as i32,
16026 nh_q as i32,
16027 nh_k as i32,
16028 n_tokens as i32,
16029 );
16030 let __s_b = self.gpu.stream();
16031 let mut b = __s_b.launch_builder(&f);
16032 b.arg(q)
16033 .arg(k)
16034 .arg(pos)
16035 .arg(&hd)
16036 .arg(&nd)
16037 .arg(&nq)
16038 .arg(&nk)
16039 .arg(&nt)
16040 .arg(&theta_scale)
16041 .arg(&freq_scale);
16042 match ff {
16043 Some(ffv) => {
16044 b.arg(ffv);
16045 unsafe {
16046 b.launch(cfg)?;
16047 }
16048 }
16049 None => {
16050 let null: u64 = 0;
16051 b.arg(&null);
16052 unsafe {
16053 b.launch(cfg)?;
16054 }
16055 }
16056 }
16057 Ok(())
16058 }
16059
16060 pub fn gelu_tanh_mul(
16062 &self,
16063 gate: &CudaSlice<f32>,
16064 up: &CudaSlice<f32>,
16065 dst: &mut CudaSlice<f32>,
16066 n: usize,
16067 ) -> Result<(), Box<dyn std::error::Error>> {
16068 let f = self.func("gelu_tanh_mul_f32");
16069 let cfg = LaunchConfig::for_num_elems(n as u32);
16070 let ni = n as i32;
16071 let __s_b = self.gpu.stream();
16072 let mut b = __s_b.launch_builder(&f);
16073 b.arg(gate).arg(up).arg(dst).arg(&ni);
16074 unsafe {
16075 b.launch(cfg)?;
16076 }
16077 Ok(())
16078 }
16079
16080 pub fn silu_mul(
16081 &self,
16082 gate: &CudaSlice<f32>,
16083 up: &CudaSlice<f32>,
16084 dst: &mut CudaSlice<f32>,
16085 n: usize,
16086 ) -> Result<(), Box<dyn std::error::Error>> {
16087 let f = self.func("silu_mul_f32");
16088 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16090 let ni = n as i32;
16091 let __s_b = self.gpu.stream();
16092 let mut b = __s_b.launch_builder(&f);
16093 b.arg(gate).arg(up).arg(dst).arg(&ni);
16094 unsafe {
16095 b.launch(cfg)?;
16096 }
16097 Ok(())
16098 }
16099
16100 pub fn silu_mul_host_expf(
16102 &self,
16103 gate: &CudaSlice<f32>,
16104 up: &CudaSlice<f32>,
16105 dst: &mut CudaSlice<f32>,
16106 n: usize,
16107 ) -> Result<(), Box<dyn std::error::Error>> {
16108 let f = self.func("silu_mul_host_expf_f32");
16109 let cfg = LaunchConfig::for_num_elems(n as u32);
16110 let ni = n as i32;
16111 let __s_b = self.gpu.stream();
16112 let mut b = __s_b.launch_builder(&f);
16113 b.arg(gate).arg(up).arg(dst).arg(&ni);
16114 unsafe {
16115 b.launch(cfg)?;
16116 }
16117 Ok(())
16118 }
16119
16120 pub fn silu_clamped_mul_host_expf(
16122 &self,
16123 gate: &CudaSlice<f32>,
16124 up: &CudaSlice<f32>,
16125 limit: f32,
16126 dst: &mut CudaSlice<f32>,
16127 n: usize,
16128 ) -> Result<(), Box<dyn std::error::Error>> {
16129 if !limit.is_finite() || limit <= 0.0 {
16130 return Err(
16131 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
16132 );
16133 }
16134 let f = self.func("silu_clamped_mul_host_expf_f32");
16135 let cfg = LaunchConfig::for_num_elems(n as u32);
16136 let ni = n as i32;
16137 let __s_b = self.gpu.stream();
16138 let mut b = __s_b.launch_builder(&f);
16139 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
16140 unsafe {
16141 b.launch(cfg)?;
16142 }
16143 Ok(())
16144 }
16145
16146 pub fn silu_mul_f16out(
16149 &self,
16150 gate: &CudaSlice<f32>,
16151 up: &CudaSlice<f32>,
16152 dst: &mut CudaSlice<f32>,
16153 dst16: &mut CudaSlice<u8>,
16154 n: usize,
16155 ) -> Result<(), Box<dyn std::error::Error>> {
16156 let f = self.func("silu_mul_f16out_f32");
16157 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16158 let ni = n as i32;
16159 let __s_b = self.gpu.stream();
16160 let mut b = __s_b.launch_builder(&f);
16161 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
16162 unsafe {
16163 b.launch(cfg)?;
16164 }
16165 Ok(())
16166 }
16167
16168 pub fn silu_mul_scaled(
16175 &self,
16176 gate: &CudaSlice<f32>,
16177 up: &CudaSlice<f32>,
16178 gs: f32,
16179 us: f32,
16180 dst: &mut CudaSlice<f32>,
16181 n: usize,
16182 ) -> Result<(), Box<dyn std::error::Error>> {
16183 let f = self.func("silu_mul_scaled_f32");
16184 let cfg = LaunchConfig::for_num_elems(n as u32);
16185 let ni = n as i32;
16186 let (gsf, usf) = (gs, us);
16187 let __s_b = self.gpu.stream();
16188 let mut b = __s_b.launch_builder(&f);
16189 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
16190 unsafe {
16191 b.launch(cfg)?;
16192 }
16193 Ok(())
16194 }
16195
16196 #[allow(clippy::too_many_arguments)]
16200 pub fn swigluoai_mul_scaled(
16201 &self,
16202 gate: &CudaSlice<f32>,
16203 up: &CudaSlice<f32>,
16204 gs: f32,
16205 us: f32,
16206 alpha: f32,
16207 limit: f32,
16208 dst: &mut CudaSlice<f32>,
16209 n: usize,
16210 ) -> Result<(), Box<dyn std::error::Error>> {
16211 let f = self.func("swigluoai_mul_scaled_f32");
16212 let cfg = LaunchConfig::for_num_elems(n as u32);
16213 let ni = n as i32;
16214 let __s_b = self.gpu.stream();
16215 let mut b = __s_b.launch_builder(&f);
16216 b.arg(gate)
16217 .arg(up)
16218 .arg(&gs)
16219 .arg(&us)
16220 .arg(&alpha)
16221 .arg(&limit)
16222 .arg(dst)
16223 .arg(&ni);
16224 unsafe {
16225 b.launch(cfg)?;
16226 }
16227 Ok(())
16228 }
16229
16230 pub fn silu_mul_scaled_q8_1(
16238 &self,
16239 gate: &CudaSlice<f32>,
16240 up: &CudaSlice<f32>,
16241 gs: f32,
16242 us: f32,
16243 n: usize,
16244 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16245 let f = self.func("silu_mul_scaled_q8_1");
16246 let nblk = n / 32;
16247 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);
16251 let (gsf, usf, ni) = (gs, us, n as i32);
16252 let __s_b = self.gpu.stream();
16253 let mut b = __s_b.launch_builder(&f);
16254 b.arg(gate)
16255 .arg(up)
16256 .arg(&gsf)
16257 .arg(&usf)
16258 .arg(&mut aq)
16259 .arg(&mut ad)
16260 .arg(&ni);
16261 unsafe {
16262 b.launch(cfg)?;
16263 }
16264 Ok((aq, ad))
16265 }
16266
16267 pub fn add(
16268 &self,
16269 a: &CudaSlice<f32>,
16270 b_in: &CudaSlice<f32>,
16271 dst: &mut CudaSlice<f32>,
16272 n: usize,
16273 ) -> Result<(), Box<dyn std::error::Error>> {
16274 let f = self.func("add_f32");
16275 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16277 let ni = n as i32;
16278 let __s_bld = self.gpu.stream();
16279 let mut bld = __s_bld.launch_builder(&f);
16280 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
16281 unsafe {
16282 bld.launch(cfg)?;
16283 }
16284 Ok(())
16285 }
16286
16287 pub fn mul(
16288 &self,
16289 a: &CudaSlice<f32>,
16290 b_in: &CudaSlice<f32>,
16291 dst: &mut CudaSlice<f32>,
16292 n: usize,
16293 ) -> Result<(), Box<dyn std::error::Error>> {
16294 let f = self.func("mul_f32");
16295 let cfg = LaunchConfig::for_num_elems(n as u32);
16296 let ni = n as i32;
16297 let __s_bld = self.gpu.stream();
16298 let mut bld = __s_bld.launch_builder(&f);
16299 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
16300 unsafe {
16301 bld.launch(cfg)?;
16302 }
16303 Ok(())
16304 }
16305
16306 pub fn matmul(
16309 &self,
16310 w: &crate::model::GpuTensor,
16311 x: &CudaSlice<f32>,
16312 m: usize,
16313 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16314 use crate::model::GpuTensor;
16315 let in_f = w.in_features();
16316 let out_f = w.out_features();
16317 #[allow(non_snake_case)]
16325 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
16328 usize::MAX
16329 } else {
16330 16usize
16331 };
16332
16333 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
16358 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
16359 return Ok(y);
16360 }
16361 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
16368 return Ok(y);
16369 }
16370 if let Some(y) = self.try_f16_gemm(w, x, m)? {
16373 return Ok(y);
16374 }
16375 }
16376 if let GpuTensor::Quant { qtype, .. } = w
16391 && *qtype == QT_F8_E4M3_BLK
16392 {
16393 if m >= GEMM_M_THRESHOLD
16394 && let Some(y) = self.try_e4m3_blk_prefill(w, x, m)?
16395 {
16396 return Ok(y);
16397 }
16398 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16399 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
16400 return Ok(y);
16401 }
16402 }
16403 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
16404 return self.qmatvec_mmq(w, x, m);
16405 }
16406 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
16407 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16408 return self.qmatvec_gemm(w, &aq, &ad, m);
16409 }
16410 if m >= GEMM_M_THRESHOLD
16413 && let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)?
16414 {
16415 return Ok(y);
16416 }
16417 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
16421 if m == 1
16426 && fast
16427 && let GpuTensor::Quant {
16428 bytes,
16429 qtype,
16430 row_bytes,
16431 rp,
16432 rp4,
16433 scale,
16434 ..
16435 } = w
16436 && self.mmvq_supports(*qtype)
16437 {
16438 let (bytes, rp) = match rp4 {
16442 Some(m4) => (m4, true),
16443 None => (bytes, *rp),
16444 };
16445 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16446 return self.qmatvec_mmvq(
16447 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
16448 );
16449 }
16450 if (2..=16).contains(&m)
16466 && fast
16467 && std::env::var("MEMRA_NO_BATCHED").is_err()
16468 && (m <= 4 || Self::b8_enabled())
16469 {
16470 let m_ok = m <= 8
16480 || matches!(w, GpuTensor::Quant { qtype, .. }
16481 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
16482 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
16483 if m_ok
16484 && let GpuTensor::Quant {
16485 bytes,
16486 qtype,
16487 row_bytes,
16488 rp,
16489 rp4,
16490 ..
16491 } = w
16492 && self.batched_supports(*qtype)
16493 && self.mmvq_supports(*qtype)
16494 {
16495 let (bytes, rp) = match rp4 {
16496 Some(m4) => (m4, true),
16497 None => (bytes, *rp),
16498 };
16499 let mcols = Self::batched_mcols(m);
16500 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16501 let mut y = self.qmatvec_mmvq_batched(
16502 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
16503 )?;
16504 if let GpuTensor::Quant { scale, .. } = w
16505 && *scale != 1.0
16506 {
16507 self.scale_inplace(&mut y, *scale, m * out_f)?;
16508 }
16509 return Ok(y);
16510 }
16511 }
16512 if fast
16518 && let GpuTensor::Quant {
16519 bytes,
16520 qtype,
16521 row_bytes,
16522 scale,
16523 ..
16524 } = w
16525 && *qtype == QT_F8_E4M3
16526 {
16527 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16528 return self.qmatvec_mmvq(
16529 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
16530 );
16531 }
16532 let mut y = match w {
16533 GpuTensor::Quant {
16534 bytes,
16535 qtype,
16536 row_bytes,
16537 ..
16538 } if fast && *qtype == QT_Q8_0 => {
16539 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16540 }
16541 GpuTensor::Quant {
16542 bytes,
16543 qtype,
16544 row_bytes,
16545 ..
16546 } if fast && *qtype == QT_Q4_K => {
16547 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16548 }
16549 GpuTensor::Quant {
16550 bytes,
16551 qtype,
16552 row_bytes,
16553 ..
16554 } if fast && *qtype == QT_Q6_K => {
16555 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16556 }
16557 GpuTensor::Quant {
16558 bytes,
16559 qtype,
16560 row_bytes,
16561 ..
16562 } if fast && *qtype == QT_Q5_K => {
16563 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16564 }
16565 GpuTensor::Quant {
16566 bytes,
16567 qtype,
16568 row_bytes,
16569 ..
16570 } if fast && *qtype == QT_Q3_K => {
16571 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16572 }
16573 GpuTensor::Quant {
16574 bytes,
16575 qtype,
16576 row_bytes,
16577 rp,
16578 ..
16579 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
16580 if *rp {
16581 "qmatvec_nvfp4_dp4a_rp"
16582 } else {
16583 "qmatvec_nvfp4_dp4a"
16584 },
16585 &bytes.slice(0..bytes.len()),
16586 x,
16587 m,
16588 in_f,
16589 out_f,
16590 *row_bytes,
16591 )?,
16592 GpuTensor::Quant {
16596 bytes,
16597 qtype,
16598 row_bytes,
16599 ..
16600 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
16601 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16602 }
16603 GpuTensor::Quant {
16608 bytes,
16609 qtype,
16610 row_bytes,
16611 rp,
16612 ..
16613 } =>
16614 {
16617 self.qmatvec(
16618 bytes,
16619 x,
16620 m,
16621 in_f,
16622 out_f,
16623 if *rp && *qtype == QT_NVFP4 {
16624 QT_NVFP4_RP
16625 } else {
16626 *qtype
16627 },
16628 *row_bytes,
16629 )?
16630 }
16631 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
16632 GpuTensor::FloatBf16 { data, .. } => {
16635 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f.is_multiple_of(8) {
16642 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16643 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
16644 y
16645 } else {
16646 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
16647 }
16648 }
16649 };
16650 if let GpuTensor::Quant { scale, .. } = w
16652 && *scale != 1.0
16653 {
16654 self.scale_inplace(&mut y, *scale, m * out_f)?;
16655 }
16656 Ok(y)
16657 }
16658
16659 pub fn stage_a_raw_needed() -> bool {
16669 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16670 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
16671 }
16672
16673 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
16676 use crate::model::GpuTensor;
16677 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
16678 return false;
16679 }
16680 match w {
16681 GpuTensor::Quant { qtype, .. } => {
16688 matches!(
16689 *qtype,
16690 QT_Q8_0
16691 | QT_Q4_K
16692 | QT_Q6_K
16693 | QT_Q5_K
16694 | QT_Q3_K
16695 | QT_NVFP4
16696 | QT_F8_E4M3
16697 | QT_F8_E4M3_BLK
16698 | QT_Q4_0
16699 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
16700 }
16701 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16702 }
16703 }
16704
16705 pub fn matmul_pre(
16710 &self,
16711 w: &crate::model::GpuTensor,
16712 aq: &CudaSlice<i8>,
16713 ad: &CudaSlice<f32>,
16714 x_fallback: &CudaSlice<f32>,
16715 m: usize,
16716 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16717 use crate::model::GpuTensor;
16718 let x_raw_ok = x_fallback.len() >= m * w.in_features();
16724 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
16727 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
16728 return Ok(y);
16729 }
16730 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
16733 return Ok(y);
16734 }
16735 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
16737 return Ok(y);
16738 }
16739 }
16740 if m >= 16
16746 && x_raw_ok
16747 && !self.verify_exact_on()
16748 && let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)?
16749 {
16750 return Ok(y);
16751 }
16752 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
16753 return Ok(y);
16754 }
16755 if m >= 16
16760 && w.out_features() >= 128
16761 && self.mmq_supports(w)
16762 && !self.verify_exact_on()
16763 && x_raw_ok
16764 {
16765 return self.qmatvec_mmq(w, x_fallback, m);
16766 }
16767 if m >= 16
16770 && x_raw_ok
16771 && !self.verify_exact_on()
16772 && let Some(y) =
16773 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
16774 {
16775 return Ok(y);
16776 }
16777 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
16780 return self.qmatvec_gemm(w, aq, ad, m);
16781 }
16782 if !self.uses_q8_1_fast(w) {
16801 if !x_raw_ok {
16802 return Err(format!(
16803 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
16804 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
16805 dropped the f32, so there is nothing to fall back to — pass the real f32 \
16806 activation (see Engine::rms_norm_decode, which is bit-identical to \
16807 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
16808 x_fallback.len(),
16809 m,
16810 w.in_features(),
16811 m * w.in_features()
16812 )
16813 .into());
16814 }
16815 return self.matmul(w, x_fallback, m);
16816 }
16817 let in_f = w.in_features();
16818 let out_f = w.out_features();
16819 let (bytes, qtype, row_bytes, scale, rp) = match w {
16820 GpuTensor::Quant {
16821 bytes,
16822 qtype,
16823 row_bytes,
16824 scale,
16825 rp,
16826 ..
16827 } => (bytes, *qtype, *row_bytes, *scale, *rp),
16828 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
16829 };
16830 let (mbytes, mrp) = match w {
16833 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
16834 _ => (bytes, rp),
16835 };
16836 if m == 1 && self.mmvq_supports(qtype) {
16840 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
16841 }
16842 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
16855 && std::env::var("MEMRA_NO_BATCHED").is_err()
16856 && (m <= 4 || Self::b8_enabled())
16857 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
16861 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
16862 {
16863 let mcols = Self::batched_mcols(m);
16864 return self.qmatvec_mmvq_batched(
16865 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
16866 );
16867 }
16868 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
16874 let (b2, r2) = if qtype == QT_Q4_0 {
16875 (mbytes, mrp)
16876 } else {
16877 (bytes, rp)
16878 };
16879 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
16880 }
16881 let name = match qtype {
16882 QT_Q8_0 => "qmatvec_q8_0_dp4a",
16883 QT_Q4_K => "qmatvec_q4_K_dp4a",
16884 QT_Q6_K => "qmatvec_q6_K_dp4a",
16885 QT_Q5_K => "qmatvec_q5_K_dp4a",
16886 QT_Q3_K => "qmatvec_q3_K_dp4a",
16887 QT_NVFP4 => {
16888 if rp {
16889 "qmatvec_nvfp4_dp4a_rp"
16890 } else {
16891 "qmatvec_nvfp4_dp4a"
16892 }
16893 }
16894 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
16895 _ => unreachable!(),
16896 };
16897 let f = self.func(name);
16898 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
16900 grid_dim: (out_f as u32, m as u32, 1),
16901 block_dim: (128, 1, 1),
16902 shared_mem_bytes: 0,
16903 };
16904 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16905 let __s_b = self.gpu.stream();
16906 let mut b = __s_b.launch_builder(&f);
16907 b.arg(bytes)
16908 .arg(aq)
16909 .arg(ad)
16910 .arg(&mut y)
16911 .arg(&inf)
16912 .arg(&outf)
16913 .arg(&mi)
16914 .arg(&rb);
16915 unsafe {
16916 b.launch(cfg)?;
16917 }
16918 if scale != 1.0 {
16919 self.scale_inplace(&mut y, scale, m * out_f)?;
16920 }
16921 Ok(y)
16922 }
16923
16924 pub fn matmul_decode_exact(
16932 &self,
16933 w: &crate::model::GpuTensor,
16934 x: &CudaSlice<f32>,
16935 m: usize,
16936 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16937 use crate::model::GpuTensor;
16938 if let GpuTensor::Float { data, .. } = w {
16946 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
16947 }
16948 if let GpuTensor::FloatBf16 { data, .. } = w {
16951 let (in_f, out_f) = (w.in_features(), w.out_features());
16952 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
16955 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16956 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
16957 return Ok(y);
16958 }
16959 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
16960 }
16961 if !self.uses_q8_1_fast(w) {
16962 return self.matmul(w, x, m);
16963 }
16964 let in_f = w.in_features();
16965 let out_f = w.out_features();
16966 let (bytes, qtype, row_bytes, scale, rp) = match w {
16967 GpuTensor::Quant {
16968 bytes,
16969 qtype,
16970 row_bytes,
16971 scale,
16972 rp,
16973 ..
16974 } => (bytes, *qtype, *row_bytes, *scale, *rp),
16975 _ => return self.matmul(w, x, m),
16976 };
16977 let (bytes, rp) = match w {
16980 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
16981 _ => (bytes, rp),
16982 };
16983 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16984 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
16988 return Ok(y);
16989 }
16990 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
16999 && std::env::var("MEMRA_NO_BATCHED").is_err()
17000 && (m <= 4 || Self::b8_enabled())
17001 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
17004 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
17005 {
17006 let mcols = Self::batched_mcols(m);
17007 return self.qmatvec_mmvq_batched(
17008 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
17009 );
17010 }
17011 if self.mmvq_supports(qtype) {
17012 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
17015 }
17016 self.matmul_pre(w, &aq, &ad, x, m)
17019 }
17020
17021 pub fn matmul_decode_exact_pre(
17031 &self,
17032 w: &crate::model::GpuTensor,
17033 aq: &CudaSlice<i8>,
17034 ad: &CudaSlice<f32>,
17035 m: usize,
17036 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17037 use crate::model::GpuTensor;
17038 debug_assert!(
17039 self.uses_q8_1_fast(w),
17040 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
17041 );
17042 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
17044 return Ok(y);
17045 }
17046 let in_f = w.in_features();
17047 let out_f = w.out_features();
17048 let (bytes, qtype, row_bytes, scale, rp) = match w {
17049 GpuTensor::Quant {
17050 bytes,
17051 qtype,
17052 row_bytes,
17053 scale,
17054 rp,
17055 ..
17056 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17057 _ => {
17058 return Err(
17059 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
17060 );
17061 }
17062 };
17063 let (bytes, rp) = match w {
17065 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
17066 _ => (bytes, rp),
17067 };
17068 if (2..=16).contains(&m)
17070 && self.batched_supports(qtype)
17071 && self.mmvq_supports(qtype)
17072 && std::env::var("MEMRA_NO_BATCHED").is_err()
17073 && (m <= 4 || Self::b8_enabled())
17074 && (m <= 8
17075 || qtype == QT_Q4_0
17076 || qtype == QT_Q6_K
17077 || qtype == QT_F8_E4M3
17078 || qtype == QT_NVFP4
17079 || qtype == QT_Q4_K
17080 || qtype == QT_Q5_K
17081 || qtype == QT_Q8_0)
17082 {
17083 let mcols = Self::batched_mcols(m);
17084 return self.qmatvec_mmvq_batched(
17085 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
17086 );
17087 }
17088 if self.mmvq_supports(qtype) {
17089 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
17090 }
17091 let x0 = self.zeros(0)?;
17094 self.matmul_pre(w, aq, ad, &x0, m)
17095 }
17096
17097 #[allow(clippy::type_complexity)] pub fn matmul_decode_exact_dual_pre(
17107 &self,
17108 w0: &crate::model::GpuTensor,
17109 w1: &crate::model::GpuTensor,
17110 aq: &CudaSlice<i8>,
17111 ad: &CudaSlice<f32>,
17112 m: usize,
17113 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
17114 {
17115 use crate::model::GpuTensor;
17116 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17117 let on = *ON.get_or_init(|| {
17118 std::env::var("MEMRA_SPEC_DUAL_T")
17119 .map(|v| v != "0")
17120 .unwrap_or(true)
17121 });
17122 if !on
17123 || !(2..=7).contains(&m)
17124 || std::env::var("MEMRA_NO_BATCHED").is_ok()
17125 || !self.uses_q8_1_fast(w0)
17126 || !self.uses_q8_1_fast(w1)
17127 {
17128 return Ok(None);
17129 }
17130 if !self.mmvq_supports(QT_NVFP4) {
17135 return Ok(None);
17136 }
17137 let (in_f, out_f) = (w0.in_features(), w0.out_features());
17138 if w1.in_features() != in_f || w1.out_features() != out_f {
17139 return Ok(None);
17140 }
17141 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
17142 (
17143 GpuTensor::Quant {
17144 bytes: b0,
17145 qtype: q0,
17146 row_bytes: rb0,
17147 scale: s0,
17148 rp: rp0,
17149 rp4: None,
17150 ..
17151 },
17152 GpuTensor::Quant {
17153 bytes: b1,
17154 qtype: q1,
17155 row_bytes: rb1,
17156 scale: s1,
17157 rp: rp1,
17158 rp4: None,
17159 ..
17160 },
17161 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
17162 (b0, b1, *rb0, *s0, *s1, *rp0)
17163 }
17164 _ => return Ok(None),
17165 };
17166 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
17169 {
17170 return Ok(None);
17171 }
17172 let (y0, y1) =
17173 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
17174 Ok(Some(((y0, s0), (y1, s1))))
17175 }
17176
17177 pub fn matmul_decode_exact_group4_pre(
17189 &self,
17190 ws: [&crate::model::GpuTensor; 4],
17191 aq: &CudaSlice<i8>,
17192 ad: &CudaSlice<f32>,
17193 m: usize,
17194 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17195 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17196 let on = *ON.get_or_init(|| {
17197 std::env::var("MEMRA_TK_GDN_GROUP")
17198 .map(|v| v != "0")
17199 .unwrap_or(true)
17200 });
17201 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
17202 }
17203
17204 pub fn matmul_decode_exact_group3_pre(
17209 &self,
17210 ws: [&crate::model::GpuTensor; 3],
17211 aq: &CudaSlice<i8>,
17212 ad: &CudaSlice<f32>,
17213 m: usize,
17214 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17215 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17216 let on = *ON.get_or_init(|| {
17217 std::env::var("MEMRA_TK_FA_GROUP")
17218 .map(|v| v != "0")
17219 .unwrap_or(true)
17220 });
17221 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
17222 }
17223
17224 #[allow(clippy::manual_div_ceil)] fn matmul_decode_exact_group_pre(
17229 &self,
17230 ws: &[&crate::model::GpuTensor],
17231 aq: &CudaSlice<i8>,
17232 ad: &CudaSlice<f32>,
17233 m: usize,
17234 on: bool,
17235 tag: &'static str,
17236 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17237 use crate::model::GpuTensor;
17238 if !on
17239 || !(2..=16).contains(&m)
17240 || std::env::var("MEMRA_NO_BATCHED").is_ok()
17241 || (m > 4 && !Self::b8_enabled())
17242 || !self.mmvq_supports(QT_NVFP4)
17243 || !self.batched_supports(QT_NVFP4)
17244 {
17245 return Ok(None);
17246 }
17247 let in_f = ws[0].in_features();
17248 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
17249 for w in ws {
17250 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
17251 return Ok(None);
17252 }
17253 match w {
17254 GpuTensor::Quant {
17255 bytes,
17256 qtype,
17257 scale,
17258 rp: true,
17259 rp4: None,
17260 ..
17261 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
17262 parts.push((bytes, w.out_features(), *scale));
17263 }
17264 _ => return Ok(None),
17265 }
17266 }
17267 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17269 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
17270 let mcols = if (5..=7).contains(&m) && b567 {
17271 m
17272 } else {
17273 Self::batched_mcols(m)
17274 };
17275 let kname: &'static str = match mcols {
17276 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
17277 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
17278 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
17279 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
17280 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
17281 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
17282 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
17283 _ => return Ok(None),
17284 };
17285 if std::env::var("MEMRA_DEBUG").is_ok() {
17288 use std::sync::Mutex;
17289 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
17290 let mut seen = SEEN.lock().unwrap();
17291 if !seen.contains(&tag) {
17292 seen.push(tag);
17293 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
17294 }
17295 }
17296 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
17299 let three = parts.len() == 3;
17300 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
17301 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
17302 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
17303 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
17306 let cfg = LaunchConfig {
17307 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
17308 block_dim: (32, ROWS_PER_BLOCK, 1),
17309 shared_mem_bytes: 0,
17310 };
17311 let (inf, mi) = (in_f as i32, m as i32);
17312 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
17313 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
17314 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
17315 let s3 = if three { 1.0f32 } else { parts[3].2 };
17316 let w3 = if three { parts[0].0 } else { parts[3].0 };
17317 let f = self.func(kname);
17318 let __s_b = self.gpu.stream();
17319 let mut b = __s_b.launch_builder(&f);
17320 b.arg(parts[0].0)
17321 .arg(parts[1].0)
17322 .arg(parts[2].0)
17323 .arg(w3)
17324 .arg(aq)
17325 .arg(ad)
17326 .arg(&mut y0)
17327 .arg(&mut y1)
17328 .arg(&mut y2)
17329 .arg(&mut y3)
17330 .arg(&inf)
17331 .arg(&n0)
17332 .arg(&n1)
17333 .arg(&n2)
17334 .arg(&n3)
17335 .arg(&mi)
17336 .arg(&s0)
17337 .arg(&s1)
17338 .arg(&s2)
17339 .arg(&s3);
17340 unsafe {
17341 b.launch(cfg)?;
17342 }
17343 Ok(Some(if three {
17344 vec![y0, y1, y2]
17345 } else {
17346 vec![y0, y1, y2, y3]
17347 }))
17348 }
17349
17350 #[allow(clippy::type_complexity)] pub fn matmul_decode_exact_dual(
17367 &self,
17368 w0: &crate::model::GpuTensor,
17369 w1: &crate::model::GpuTensor,
17370 x: &CudaSlice<f32>,
17371 m: usize,
17372 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17373 use crate::model::GpuTensor;
17374 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17375 let on = *ON.get_or_init(|| {
17376 std::env::var("MEMRA_SPEC_DUAL_T")
17377 .map(|v| v != "0")
17378 .unwrap_or(true)
17379 });
17380 if !on
17381 || !(2..=4).contains(&m)
17382 || std::env::var("MEMRA_NO_BATCHED").is_ok()
17383 || !self.uses_q8_1_fast(w0)
17384 || !self.uses_q8_1_fast(w1)
17385 {
17386 return Ok(None);
17387 }
17388 if !self.mmvq_supports(QT_NVFP4) {
17393 return Ok(None);
17394 }
17395 let (in_f, out_f) = (w0.in_features(), w0.out_features());
17396 if w1.in_features() != in_f || w1.out_features() != out_f {
17397 return Ok(None);
17398 }
17399 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
17400 (
17401 GpuTensor::Quant {
17402 bytes: b0,
17403 qtype: q0,
17404 row_bytes: rb0,
17405 scale: s0,
17406 rp: rp0,
17407 rp4: None,
17408 ..
17409 },
17410 GpuTensor::Quant {
17411 bytes: b1,
17412 qtype: q1,
17413 row_bytes: rb1,
17414 scale: s1,
17415 rp: rp1,
17416 rp4: None,
17417 ..
17418 },
17419 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
17420 (b0, b1, *rb0, *s0, *s1, *rp0)
17421 }
17422 _ => return Ok(None),
17423 };
17424 if std::env::var("MEMRA_DEBUG").is_ok() {
17427 static ONCE: std::sync::Once = std::sync::Once::new();
17428 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
17429 }
17430 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17431 let (y0, y1) =
17432 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
17433 let mut y0 = y0;
17434 let mut y1 = y1;
17435 if s0 != 1.0 {
17436 self.scale_inplace(&mut y0, s0, m * out_f)?;
17437 }
17438 if s1 != 1.0 {
17439 self.scale_inplace(&mut y1, s1, m * out_f)?;
17440 }
17441 Ok(Some((y0, y1)))
17442 }
17443
17444 #[allow(clippy::too_many_arguments)]
17449 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_batched_dual_raw(
17451 &self,
17452 b0: &CudaSlice<u8>,
17453 b1: &CudaSlice<u8>,
17454 aq: &CudaSlice<i8>,
17455 ad: &CudaSlice<f32>,
17456 m: usize,
17457 in_f: usize,
17458 out_f: usize,
17459 row_bytes: usize,
17460 rp: bool,
17461 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17462 const ROWS_PER_BLOCK: u32 = 4;
17463 let mcols = Self::batched_mcols(m);
17464 let tiny_rp1 = rp
17467 && mcols == 4
17468 && out_f <= 128
17469 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
17470 let (name, rows_per_block) = if tiny_rp1 {
17471 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
17472 } else {
17473 match (mcols, rp, m) {
17474 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
17475 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
17476 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
17477 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
17478 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
17479 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
17480 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
17481 _ => {
17482 return Err(
17483 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
17484 );
17485 }
17486 }
17487 };
17488 let f = self.func(name);
17489 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
17490 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
17491 let cfg = LaunchConfig {
17492 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
17493 block_dim: (32, ROWS_PER_BLOCK, 1),
17494 shared_mem_bytes: 0,
17495 };
17496 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17497 let __s_b = self.gpu.stream();
17498 let mut b = __s_b.launch_builder(&f);
17499 b.arg(b0)
17500 .arg(b1)
17501 .arg(aq)
17502 .arg(ad)
17503 .arg(&mut y0)
17504 .arg(&mut y1)
17505 .arg(&inf)
17506 .arg(&outf)
17507 .arg(&mi)
17508 .arg(&rb);
17509 unsafe {
17510 b.launch(cfg)?;
17511 }
17512 Ok((y0, y1))
17513 }
17514
17515 #[allow(clippy::type_complexity)] #[allow(clippy::manual_div_ceil)] pub fn matmul_pre_dual_noscale(
17529 &self,
17530 w0: &crate::model::GpuTensor,
17531 w1: &crate::model::GpuTensor,
17532 aq: &CudaSlice<i8>,
17533 ad: &CudaSlice<f32>,
17534 m: usize,
17535 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
17536 {
17537 use crate::model::GpuTensor;
17538 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
17539 return Ok(None);
17540 }
17541 if !self.mmvq_supports(QT_NVFP4) {
17551 return Ok(None);
17552 }
17553 let (in_f, out_f) = (w0.in_features(), w0.out_features());
17554 if w1.in_features() != in_f || w1.out_features() != out_f {
17555 return Ok(None);
17556 }
17557 let no_mirror =
17570 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
17571 if self.q8_ffn_fuse2_on()
17572 && no_mirror(w0)
17573 && no_mirror(w1)
17574 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
17575 {
17576 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
17577 return Ok(Some(((y0, 1.0), (y1, 1.0))));
17578 }
17579 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
17589 let (y0, y1) =
17590 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
17591 return Ok(Some(((y0, p0.3), (y1, p1.3))));
17592 }
17593 let (b0, q0, rb0, s0, rp0) = match w0 {
17594 GpuTensor::Quant {
17595 bytes,
17596 qtype,
17597 row_bytes,
17598 scale,
17599 rp,
17600 ..
17601 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17602 _ => return Ok(None),
17603 };
17604 let (b1, q1, rb1, s1, rp1) = match w1 {
17605 GpuTensor::Quant {
17606 bytes,
17607 qtype,
17608 row_bytes,
17609 scale,
17610 rp,
17611 ..
17612 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17613 _ => return Ok(None),
17614 };
17615 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
17616 return Ok(None);
17617 }
17618 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
17620 let rows_per_block = ROWS_PER_BLOCK * RPW;
17621 let f = self.func(if rp0 {
17622 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
17623 } else {
17624 "qmatvec_nvfp4_mmvq_dual_mr2"
17625 });
17626 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
17627 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
17628 let cfg = LaunchConfig {
17629 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
17630 block_dim: (32, ROWS_PER_BLOCK, 1),
17631 shared_mem_bytes: 0,
17632 };
17633 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
17634 let one = 1.0f32;
17637 let __s_b = self.gpu.stream();
17638 let mut b = __s_b.launch_builder(&f);
17639 b.arg(b0)
17640 .arg(b1)
17641 .arg(aq)
17642 .arg(ad)
17643 .arg(&mut y0)
17644 .arg(&mut y1)
17645 .arg(&inf)
17646 .arg(&outf)
17647 .arg(&mi)
17648 .arg(&rb)
17649 .arg(&one)
17650 .arg(&one);
17651 unsafe {
17652 b.launch(cfg)?;
17653 }
17654 Ok(Some(((y0, s0), (y1, s1))))
17655 }
17656
17657 #[allow(clippy::too_many_arguments)]
17665 #[allow(clippy::type_complexity)] pub fn matmul_nvfp4_fused3(
17667 &self,
17668 w0: &crate::model::GpuTensor,
17669 w1: &crate::model::GpuTensor,
17670 w2: &crate::model::GpuTensor,
17671 aq: &CudaSlice<i8>,
17672 ad: &CudaSlice<f32>,
17673 m: usize,
17674 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17675 {
17676 use crate::model::GpuTensor;
17677 if !self.mmvq_supports(QT_NVFP4)
17684 || !self.uses_q8_1_fast(w0)
17685 || !self.uses_q8_1_fast(w1)
17686 || !self.uses_q8_1_fast(w2)
17687 {
17688 return Ok(None);
17689 }
17690 if (9..=16).contains(&m) {
17693 return Ok(
17694 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
17695 Some(mut ys) => {
17696 let y2 = ys.pop().unwrap();
17697 let y1 = ys.pop().unwrap();
17698 let y0 = ys.pop().unwrap();
17699 Some((y0, y1, y2))
17700 }
17701 None => None,
17702 },
17703 );
17704 }
17705 if !(1..=8).contains(&m) {
17706 return Ok(None);
17707 }
17708 if m > 1 {
17709 let in_f = w0.in_features();
17710 if !self.batched_supports(QT_NVFP4)
17711 || std::env::var("MEMRA_NO_BATCHED").is_ok()
17712 || (m > 4 && !Self::b8_enabled())
17713 || !in_f.is_multiple_of(512)
17714 || in_f / 64 > 272
17715 {
17716 return Ok(None);
17717 }
17718 }
17719 let unpack = |w: &crate::model::GpuTensor| match w {
17720 GpuTensor::Quant {
17721 bytes,
17722 qtype,
17723 scale,
17724 rp,
17725 ..
17726 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
17727 _ => None,
17728 };
17729 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
17730 return Ok(None);
17731 };
17732 let in_f = w0.in_features();
17733 if w1.in_features() != in_f || w2.in_features() != in_f {
17734 return Ok(None);
17735 }
17736 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
17737 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
17739 let rows_pb = ROWS_PER_BLOCK * RPW;
17740 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
17741 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
17742 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
17743 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
17744 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
17745 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
17748 if m > 1 {
17749 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
17751 return Ok(None);
17752 }
17753 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
17754 let cfg = LaunchConfig {
17755 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
17756 block_dim: (32, ROWS_PER_BLOCK, 1),
17757 shared_mem_bytes: 0,
17758 };
17759 let __s_b = self.gpu.stream();
17760 let mut b = __s_b.launch_builder(&f);
17761 b.arg(b0)
17762 .arg(b1)
17763 .arg(b2)
17764 .arg(aq)
17765 .arg(ad)
17766 .arg(&mut y0)
17767 .arg(&mut y1)
17768 .arg(&mut y2)
17769 .arg(&inf)
17770 .arg(&oi0)
17771 .arg(&oi1)
17772 .arg(&oi2)
17773 .arg(&mi);
17774 unsafe {
17775 b.launch(cfg)?;
17776 }
17777 return Ok(Some((y0, y1, y2)));
17778 }
17779 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
17780 let cfg = LaunchConfig {
17781 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
17782 block_dim: (32, ROWS_PER_BLOCK, 1),
17783 shared_mem_bytes: 0,
17784 };
17785 let __s_b = self.gpu.stream();
17786 let mut b = __s_b.launch_builder(&f);
17787 b.arg(b0)
17788 .arg(b1)
17789 .arg(b2)
17790 .arg(aq)
17791 .arg(ad)
17792 .arg(&mut y0)
17793 .arg(&mut y1)
17794 .arg(&mut y2)
17795 .arg(&inf)
17796 .arg(&oi0)
17797 .arg(&oi1)
17798 .arg(&oi2)
17799 .arg(&mi)
17800 .arg(&p0.1)
17801 .arg(&p1.1)
17802 .arg(&p2.1);
17803 unsafe {
17804 b.launch(cfg)?;
17805 }
17806 Ok(Some((y0, y1, y2)))
17807 }
17808
17809 #[allow(clippy::type_complexity)] pub fn matmul_nvfp4_fused2(
17819 &self,
17820 w0: &crate::model::GpuTensor,
17821 w1: &crate::model::GpuTensor,
17822 aq: &CudaSlice<i8>,
17823 ad: &CudaSlice<f32>,
17824 m: usize,
17825 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17826 use crate::model::GpuTensor;
17827 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17828 let off =
17829 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
17830 if off
17833 || m != 1
17834 || !self.mmvq_supports(QT_NVFP4)
17835 || !self.uses_q8_1_fast(w0)
17836 || !self.uses_q8_1_fast(w1)
17837 {
17838 return Ok(None);
17839 }
17840 let unpack = |w: &crate::model::GpuTensor| match w {
17841 GpuTensor::Quant {
17842 bytes,
17843 qtype,
17844 scale,
17845 rp,
17846 ..
17847 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
17848 _ => None,
17849 };
17850 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
17851 return Ok(None);
17852 };
17853 let in_f = w0.in_features();
17854 if w1.in_features() != in_f {
17855 return Ok(None);
17856 }
17857 let (o0, o1) = (w0.out_features(), w1.out_features());
17858 const ROWS_PER_BLOCK: u32 = 4; let mut rpw: u32 = 2;
17860 let mut kname = "qmatvec_nvfp4_mmvq_fused2_rp";
17861 if b200_matvec_arm_on() {
17868 let waves_at_rpw2 =
17869 (o0 as u32).div_ceil(ROWS_PER_BLOCK * 2) + (o1 as u32).div_ceil(ROWS_PER_BLOCK * 2);
17870 if waves_at_rpw2 < 2 * self.sm_count() as u32 {
17871 rpw = 1;
17872 kname = "qmatvec_nvfp4_mmvq_fused2_rp_g2";
17873 }
17874 }
17875 let rows_pb = ROWS_PER_BLOCK * rpw;
17876 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
17877 let f = self.func(kname);
17878 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
17879 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
17880 let cfg = LaunchConfig {
17881 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
17882 block_dim: (32, ROWS_PER_BLOCK, 1),
17883 shared_mem_bytes: 0,
17884 };
17885 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
17886 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
17889 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
17892 {
17893 use cudarc::driver::{DevicePtr, DevicePtrMut};
17894 let s = &self.gpu.stream();
17895 let (pw0, _g0) = b0.device_ptr(s);
17896 let (pw1, _g1) = b1.device_ptr(s);
17897 let (paq, _g2) = aq.device_ptr(s);
17898 let (pad, _g3) = ad.device_ptr(s);
17899 let (py0, _g4) = y0.device_ptr_mut(s);
17900 let (py1, _g5) = y1.device_ptr_mut(s);
17901 let (s0, s1) = (p0.1, p1.1);
17902 let mut ps = [
17903 &pw0 as *const _ as *mut std::ffi::c_void,
17904 &pw1 as *const _ as *mut _,
17905 &paq as *const _ as *mut _,
17906 &pad as *const _ as *mut _,
17907 &py0 as *const _ as *mut _,
17908 &py1 as *const _ as *mut _,
17909 &inf as *const _ as *mut _,
17910 &oi0 as *const _ as *mut _,
17911 &oi1 as *const _ as *mut _,
17912 &mi as *const _ as *mut _,
17913 &s0 as *const _ as *mut _,
17914 &s1 as *const _ as *mut _,
17915 ];
17916 unsafe {
17917 self.launch_pdl(kname, cfg.grid_dim, cfg.block_dim, &mut ps)?;
17918 }
17919 }
17920 return Ok(Some((y0, y1)));
17921 }
17922 let __s_b = self.gpu.stream();
17923 let mut b = __s_b.launch_builder(&f);
17924 b.arg(b0)
17925 .arg(b1)
17926 .arg(aq)
17927 .arg(ad)
17928 .arg(&mut y0)
17929 .arg(&mut y1)
17930 .arg(&inf)
17931 .arg(&oi0)
17932 .arg(&oi1)
17933 .arg(&mi)
17934 .arg(&p0.1)
17935 .arg(&p1.1);
17936 unsafe {
17937 b.launch(cfg)?;
17938 }
17939 Ok(Some((y0, y1)))
17940 }
17941
17942 pub fn matmul_nvfp4_fused2_into(
17947 &self,
17948 w0: &crate::model::GpuTensor,
17949 w1: &crate::model::GpuTensor,
17950 aq: &CudaSlice<i8>,
17951 ad: &CudaSlice<f32>,
17952 y0: &mut CudaSlice<f32>,
17953 y1: &mut CudaSlice<f32>,
17954 ) -> Result<bool, Box<dyn std::error::Error>> {
17955 use crate::model::GpuTensor;
17956 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17957 let off =
17958 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
17959 if off
17960 || !self.mmvq_supports(QT_NVFP4)
17961 || !self.uses_q8_1_fast(w0)
17962 || !self.uses_q8_1_fast(w1)
17963 {
17964 return Ok(false);
17965 }
17966 let unpack = |w: &crate::model::GpuTensor| match w {
17967 GpuTensor::Quant {
17968 bytes,
17969 qtype,
17970 scale,
17971 rp,
17972 ..
17973 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
17974 _ => None,
17975 };
17976 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
17977 return Ok(false);
17978 };
17979 let in_f = w0.in_features();
17980 if w1.in_features() != in_f {
17981 return Ok(false);
17982 }
17983 let (o0, o1) = (w0.out_features(), w1.out_features());
17984 if y0.len() < o0 || y1.len() < o1 {
17985 return Ok(false);
17986 }
17987 const ROWS_PER_BLOCK: u32 = 4; let mut rpw: u32 = 2;
17989 let mut kname = "qmatvec_nvfp4_mmvq_fused2_rp";
17990 if b200_matvec_arm_on() {
17994 let waves_at_rpw2 =
17995 (o0 as u32).div_ceil(ROWS_PER_BLOCK * 2) + (o1 as u32).div_ceil(ROWS_PER_BLOCK * 2);
17996 if waves_at_rpw2 < 2 * self.sm_count() as u32 {
17997 rpw = 1;
17998 kname = "qmatvec_nvfp4_mmvq_fused2_rp_g2";
17999 }
18000 }
18001 let rows_pb = ROWS_PER_BLOCK * rpw;
18002 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18003 let f = self.func(kname);
18004 let cfg = LaunchConfig {
18005 grid_dim: (nb(o0) + nb(o1), 1, 1),
18006 block_dim: (32, ROWS_PER_BLOCK, 1),
18007 shared_mem_bytes: 0,
18008 };
18009 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
18010 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
18013 let __s_b = self.gpu.stream();
18014 let mut b = __s_b.launch_builder(&f);
18015 b.arg(b0)
18016 .arg(b1)
18017 .arg(aq)
18018 .arg(ad)
18019 .arg(&mut *y0)
18020 .arg(&mut *y1)
18021 .arg(&inf)
18022 .arg(&oi0)
18023 .arg(&oi1)
18024 .arg(&mi)
18025 .arg(&p0.1)
18026 .arg(&p1.1);
18027 unsafe {
18028 b.launch(cfg)?;
18029 }
18030 Ok(true)
18031 }
18032
18033 #[allow(clippy::type_complexity)]
18038 #[allow(clippy::too_many_arguments)] pub fn matmul_nvfp4_fused4(
18040 &self,
18041 w0: &crate::model::GpuTensor,
18042 w1: &crate::model::GpuTensor,
18043 w2: &crate::model::GpuTensor,
18044 w3: &crate::model::GpuTensor,
18045 aq: &CudaSlice<i8>,
18046 ad: &CudaSlice<f32>,
18047 m: usize,
18048 ) -> Result<
18049 Option<(
18050 CudaSlice<f32>,
18051 CudaSlice<f32>,
18052 CudaSlice<f32>,
18053 CudaSlice<f32>,
18054 )>,
18055 Box<dyn std::error::Error>,
18056 > {
18057 use crate::model::GpuTensor;
18058 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
18065 || !self.mmvq_supports(QT_NVFP4)
18066 || !self.uses_q8_1_fast(w0)
18067 || !self.uses_q8_1_fast(w1)
18068 || !self.uses_q8_1_fast(w2)
18069 || !self.uses_q8_1_fast(w3)
18070 {
18071 return Ok(None);
18072 }
18073 if (9..=16).contains(&m) {
18078 return Ok(
18079 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
18080 Some(mut ys) => {
18081 let y3 = ys.pop().unwrap();
18082 let y2 = ys.pop().unwrap();
18083 let y1 = ys.pop().unwrap();
18084 let y0 = ys.pop().unwrap();
18085 Some((y0, y1, y2, y3))
18086 }
18087 None => None,
18088 },
18089 );
18090 }
18091 if !(1..=8).contains(&m) {
18092 return Ok(None);
18093 }
18094 if m > 1 {
18095 let in_f = w0.in_features();
18098 if !self.batched_supports(QT_NVFP4)
18099 || std::env::var("MEMRA_NO_BATCHED").is_ok()
18100 || (m > 4 && !Self::b8_enabled())
18101 || !in_f.is_multiple_of(512)
18102 || in_f / 64 > 272
18103 {
18104 return Ok(None);
18105 }
18106 }
18107 let unpack = |w: &crate::model::GpuTensor| match w {
18108 GpuTensor::Quant {
18109 bytes,
18110 qtype,
18111 scale,
18112 rp,
18113 ..
18114 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18115 _ => None,
18116 };
18117 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
18118 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
18119 else {
18120 return Ok(None);
18121 };
18122 let in_f = w0.in_features();
18123 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
18124 return Ok(None);
18125 }
18126 let (o0, o1, o2, o3) = (
18127 w0.out_features(),
18128 w1.out_features(),
18129 w2.out_features(),
18130 w3.out_features(),
18131 );
18132 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
18134 let rows_pb = ROWS_PER_BLOCK * RPW;
18135 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18136 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
18137 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
18138 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
18139 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
18140 let (inf, oi0, oi1, oi2, oi3, mi) = (
18141 in_f as i32,
18142 o0 as i32,
18143 o1 as i32,
18144 o2 as i32,
18145 o3 as i32,
18146 m as i32,
18147 );
18148 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
18151 if m > 1 {
18152 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
18155 return Ok(None);
18156 }
18157 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
18158 let cfg = LaunchConfig {
18159 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
18160 block_dim: (32, ROWS_PER_BLOCK, 1),
18161 shared_mem_bytes: 0,
18162 };
18163 let __s_b = self.gpu.stream();
18164 let mut b = __s_b.launch_builder(&f);
18165 b.arg(b0)
18166 .arg(b1)
18167 .arg(b2)
18168 .arg(b3)
18169 .arg(aq)
18170 .arg(ad)
18171 .arg(&mut y0)
18172 .arg(&mut y1)
18173 .arg(&mut y2)
18174 .arg(&mut y3)
18175 .arg(&inf)
18176 .arg(&oi0)
18177 .arg(&oi1)
18178 .arg(&oi2)
18179 .arg(&oi3)
18180 .arg(&mi);
18181 unsafe {
18182 b.launch(cfg)?;
18183 }
18184 return Ok(Some((y0, y1, y2, y3)));
18185 }
18186 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
18187 let cfg = LaunchConfig {
18188 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
18189 block_dim: (32, ROWS_PER_BLOCK, 1),
18190 shared_mem_bytes: 0,
18191 };
18192 let __s_b = self.gpu.stream();
18193 let mut b = __s_b.launch_builder(&f);
18194 b.arg(b0)
18195 .arg(b1)
18196 .arg(b2)
18197 .arg(b3)
18198 .arg(aq)
18199 .arg(ad)
18200 .arg(&mut y0)
18201 .arg(&mut y1)
18202 .arg(&mut y2)
18203 .arg(&mut y3)
18204 .arg(&inf)
18205 .arg(&oi0)
18206 .arg(&oi1)
18207 .arg(&oi2)
18208 .arg(&oi3)
18209 .arg(&mi)
18210 .arg(&p0.1)
18211 .arg(&p1.1)
18212 .arg(&p2.1)
18213 .arg(&p3.1);
18214 unsafe {
18215 b.launch(cfg)?;
18216 }
18217 Ok(Some((y0, y1, y2, y3)))
18218 }
18219
18220 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused2(
18229 &self,
18230 w0: &crate::model::GpuTensor,
18231 w1: &crate::model::GpuTensor,
18232 aq: &CudaSlice<i8>,
18233 ad: &CudaSlice<f32>,
18234 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18235 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
18241 return Ok(Some(self.e4m3_fused2_core(
18242 p0.0,
18243 p1.0,
18244 aq,
18245 ad,
18246 w0.in_features(),
18247 p0.1,
18248 p1.1,
18249 p0.2,
18250 p0.3,
18251 p1.3,
18252 )?));
18253 }
18254 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
18255 return Ok(None);
18256 };
18257 Ok(Some(self.q8_fused2_core(
18258 p0.0,
18259 p1.0,
18260 aq,
18261 ad,
18262 w0.in_features(),
18263 p0.1,
18264 p1.1,
18265 p0.2,
18266 )?))
18267 }
18268
18269 #[allow(clippy::too_many_arguments)]
18270 fn q8_fused2_core(
18271 &self,
18272 b0: &CudaSlice<u8>,
18273 b1: &CudaSlice<u8>,
18274 aq: &CudaSlice<i8>,
18275 ad: &CudaSlice<f32>,
18276 in_f: usize,
18277 out0: usize,
18278 out1: usize,
18279 row_bytes: usize,
18280 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18281 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18283 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18284 let f = self.func("qmatvec_q8_0_mmvq_fused2");
18285 let mut y0 = self.alloc_uninit::<f32>(out0)?;
18286 let mut y1 = self.alloc_uninit::<f32>(out1)?;
18287 let cfg = LaunchConfig {
18288 grid_dim: (nb0 + nb1, 1, 1),
18289 block_dim: (32, ROWS_PER_BLOCK, 1),
18290 shared_mem_bytes: 0,
18291 };
18292 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
18293 let __s_b = self.gpu.stream();
18294 let mut b = __s_b.launch_builder(&f);
18295 b.arg(b0)
18296 .arg(b1)
18297 .arg(aq)
18298 .arg(ad)
18299 .arg(&mut y0)
18300 .arg(&mut y1)
18301 .arg(&inf)
18302 .arg(&o0)
18303 .arg(&o1)
18304 .arg(&rbl);
18305 unsafe {
18306 b.launch(cfg)?;
18307 }
18308 Ok((y0, y1))
18309 }
18310
18311 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused2_x(
18318 &self,
18319 w0: &crate::model::GpuTensor,
18320 w1: &crate::model::GpuTensor,
18321 x: &CudaSlice<f32>,
18322 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18323 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
18324 return Ok(None);
18325 }
18326 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
18327 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
18328 return Ok(Some(self.e4m3_fused2_core(
18329 p0.0,
18330 p1.0,
18331 &aq,
18332 &ad,
18333 w0.in_features(),
18334 p0.1,
18335 p1.1,
18336 p0.2,
18337 p0.3,
18338 p1.3,
18339 )?));
18340 }
18341 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
18342 return Ok(None);
18343 };
18344 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
18345 Ok(Some(self.q8_fused2_core(
18346 p0.0,
18347 p1.0,
18348 &aq,
18349 &ad,
18350 w0.in_features(),
18351 p0.1,
18352 p1.1,
18353 p0.2,
18354 )?))
18355 }
18356
18357 #[allow(clippy::too_many_arguments)]
18360 pub fn qmatvec_q8_fused2_raw(
18361 &self,
18362 b0: &CudaSlice<u8>,
18363 b1: &CudaSlice<u8>,
18364 x: &CudaSlice<f32>,
18365 in_f: usize,
18366 out0: usize,
18367 out1: usize,
18368 row_bytes: usize,
18369 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18370 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
18371 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
18372 }
18373
18374 #[allow(clippy::type_complexity)] pub fn matmul_q4_fused3(
18381 &self,
18382 w0: &crate::model::GpuTensor,
18383 w1: &crate::model::GpuTensor,
18384 w2: &crate::model::GpuTensor,
18385 aq: &CudaSlice<i8>,
18386 ad: &CudaSlice<f32>,
18387 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
18388 {
18389 use crate::model::GpuTensor;
18390 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18391 match w {
18392 GpuTensor::Quant {
18393 qtype, row_bytes, ..
18394 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18395 _ => None,
18396 }
18397 };
18398 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
18399 return Ok(None);
18400 };
18401 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
18402 return Ok(None);
18403 }
18404 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18408 match w {
18409 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18410 Some(m) => (m, true),
18411 None => (bytes, *rp),
18412 },
18413 _ => unreachable!(),
18414 }
18415 }
18416 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
18417 if rp0 != rp1 || rp1 != rp2 {
18418 return Ok(None);
18419 }
18420 let rp = rp0;
18421 let rpb: u32 = 4;
18422 let mr1 = rp && Self::q40_mr1_on();
18426 let nb = |o: usize| {
18427 if mr1 {
18428 (o as u32).div_ceil(rpb)
18429 } else {
18430 (o as u32).div_ceil(2).div_ceil(rpb)
18431 }
18432 };
18433 let grid = nb(o0) + nb(o1) + nb(o2);
18434 let mut y0 = self.alloc_uninit::<f32>(o0)?;
18435 let mut y1 = self.alloc_uninit::<f32>(o1)?;
18436 let mut y2 = self.alloc_uninit::<f32>(o2)?;
18437 let f = self.func(if mr1 {
18438 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
18439 } else if rp {
18440 "qmatvec_q4_0_mmvq_fused3_rp"
18441 } else {
18442 "qmatvec_q4_0_mmvq_fused3"
18443 });
18444 let cfg = LaunchConfig {
18445 grid_dim: (grid, 1, 1),
18446 block_dim: (32, rpb, 1),
18447 shared_mem_bytes: 0,
18448 };
18449 let inf = w0.in_features() as i32;
18450 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
18451 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
18452 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18455 {
18456 use cudarc::driver::{DevicePtr, DevicePtrMut};
18457 let s = &self.gpu.stream();
18458 let (p0, _g0) = b0.device_ptr(s);
18459 let (p1, _g1) = b1.device_ptr(s);
18460 let (p2, _g2) = b2.device_ptr(s);
18461 let (paq, _g3) = aq.device_ptr(s);
18462 let (pad, _g4) = ad.device_ptr(s);
18463 let (py0, _g5) = y0.device_ptr_mut(s);
18464 let (py1, _g6) = y1.device_ptr_mut(s);
18465 let (py2, _g7) = y2.device_ptr_mut(s);
18466 let mut ps = [
18467 &p0 as *const _ as *mut std::ffi::c_void,
18468 &p1 as *const _ as *mut _,
18469 &p2 as *const _ as *mut _,
18470 &paq as *const _ as *mut _,
18471 &pad as *const _ as *mut _,
18472 &py0 as *const _ as *mut _,
18473 &py1 as *const _ as *mut _,
18474 &py2 as *const _ as *mut _,
18475 &inf as *const _ as *mut _,
18476 &oo0 as *const _ as *mut _,
18477 &oo1 as *const _ as *mut _,
18478 &oo2 as *const _ as *mut _,
18479 &r0 as *const _ as *mut _,
18480 &r1 as *const _ as *mut _,
18481 &r2 as *const _ as *mut _,
18482 ];
18483 unsafe {
18484 self.launch_pdl(
18485 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
18486 (grid, 1, 1),
18487 (32, rpb, 1),
18488 &mut ps,
18489 )?;
18490 }
18491 }
18492 return Ok(Some((y0, y1, y2)));
18493 }
18494 let __s_b = self.gpu.stream();
18495 let mut b = __s_b.launch_builder(&f);
18496 b.arg(b0)
18497 .arg(b1)
18498 .arg(b2)
18499 .arg(aq)
18500 .arg(ad)
18501 .arg(&mut y0)
18502 .arg(&mut y1)
18503 .arg(&mut y2)
18504 .arg(&inf)
18505 .arg(&oo0)
18506 .arg(&oo1)
18507 .arg(&oo2)
18508 .arg(&r0)
18509 .arg(&r1)
18510 .arg(&r2);
18511 unsafe {
18512 b.launch(cfg)?;
18513 }
18514 Ok(Some((y0, y1, y2)))
18515 }
18516
18517 #[allow(clippy::too_many_arguments)]
18520 pub fn matmul_q4_fused3_into(
18521 &self,
18522 w0: &crate::model::GpuTensor,
18523 w1: &crate::model::GpuTensor,
18524 w2: &crate::model::GpuTensor,
18525 aq: &CudaSlice<i8>,
18526 ad: &CudaSlice<f32>,
18527 y0: &mut CudaSlice<f32>,
18528 y1: &mut CudaSlice<f32>,
18529 y2: &mut CudaSlice<f32>,
18530 ) -> Result<bool, Box<dyn std::error::Error>> {
18531 use crate::model::GpuTensor;
18532 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18533 match w {
18534 GpuTensor::Quant {
18535 qtype, row_bytes, ..
18536 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18537 _ => None,
18538 }
18539 };
18540 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
18541 return Ok(false);
18542 };
18543 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
18544 return Ok(false);
18545 }
18546 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18547 match w {
18548 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18549 Some(m) => (m, true),
18550 None => (bytes, *rp),
18551 },
18552 _ => unreachable!(),
18553 }
18554 }
18555 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
18556 if rp0 != rp1 || rp1 != rp2 {
18557 return Ok(false);
18558 }
18559 let rp = rp0;
18560 let rpb: u32 = 4;
18561 let mr1 = rp && Self::q40_mr1_on();
18562 let nb = |o: usize| {
18563 if mr1 {
18564 (o as u32).div_ceil(rpb)
18565 } else {
18566 (o as u32).div_ceil(2).div_ceil(rpb)
18567 }
18568 };
18569 let grid = nb(o0) + nb(o1) + nb(o2);
18570 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
18571 let f = self.func(if mr1 {
18572 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
18573 } else if rp {
18574 "qmatvec_q4_0_mmvq_fused3_rp"
18575 } else {
18576 "qmatvec_q4_0_mmvq_fused3"
18577 });
18578 let cfg = LaunchConfig {
18579 grid_dim: (grid, 1, 1),
18580 block_dim: (32, rpb, 1),
18581 shared_mem_bytes: 0,
18582 };
18583 let inf = w0.in_features() as i32;
18584 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
18585 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
18586 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18588 use cudarc::driver::{DevicePtr, DevicePtrMut};
18589 let s = &self.gpu.stream();
18590 let (p0, _g0) = b0.device_ptr(s);
18591 let (p1, _g1) = b1.device_ptr(s);
18592 let (p2, _g2) = b2.device_ptr(s);
18593 let (paq, _g3) = aq.device_ptr(s);
18594 let (pad, _g4) = ad.device_ptr(s);
18595 let (py0, _g5) = y0.device_ptr_mut(s);
18596 let (py1, _g6) = y1.device_ptr_mut(s);
18597 let (py2, _g7) = y2.device_ptr_mut(s);
18598 let mut ps = [
18599 &p0 as *const _ as *mut std::ffi::c_void,
18600 &p1 as *const _ as *mut _,
18601 &p2 as *const _ as *mut _,
18602 &paq as *const _ as *mut _,
18603 &pad as *const _ as *mut _,
18604 &py0 as *const _ as *mut _,
18605 &py1 as *const _ as *mut _,
18606 &py2 as *const _ as *mut _,
18607 &inf as *const _ as *mut _,
18608 &oo0 as *const _ as *mut _,
18609 &oo1 as *const _ as *mut _,
18610 &oo2 as *const _ as *mut _,
18611 &r0 as *const _ as *mut _,
18612 &r1 as *const _ as *mut _,
18613 &r2 as *const _ as *mut _,
18614 ];
18615 unsafe {
18616 self.launch_pdl(
18617 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
18618 (grid, 1, 1),
18619 (32, rpb, 1),
18620 &mut ps,
18621 )?;
18622 }
18623 return Ok(true);
18624 }
18625 let __s_b = self.gpu.stream();
18626 let mut b = __s_b.launch_builder(&f);
18627 b.arg(b0)
18628 .arg(b1)
18629 .arg(b2)
18630 .arg(aq)
18631 .arg(ad)
18632 .arg(&mut *y0)
18633 .arg(&mut *y1)
18634 .arg(&mut *y2)
18635 .arg(&inf)
18636 .arg(&oo0)
18637 .arg(&oo1)
18638 .arg(&oo2)
18639 .arg(&r0)
18640 .arg(&r1)
18641 .arg(&r2);
18642 unsafe {
18643 b.launch(cfg)?;
18644 }
18645 Ok(true)
18646 }
18647
18648 #[allow(clippy::type_complexity)] pub fn matmul_q4_fused2(
18651 &self,
18652 w0: &crate::model::GpuTensor,
18653 w1: &crate::model::GpuTensor,
18654 aq: &CudaSlice<i8>,
18655 ad: &CudaSlice<f32>,
18656 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18657 use crate::model::GpuTensor;
18658 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18659 match w {
18660 GpuTensor::Quant {
18661 qtype, row_bytes, ..
18662 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18663 _ => None,
18664 }
18665 };
18666 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
18667 return Ok(None);
18668 };
18669 if w0.in_features() != w1.in_features() {
18670 return Ok(None);
18671 }
18672 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18674 match w {
18675 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18676 Some(m) => (m, true),
18677 None => (bytes, *rp),
18678 },
18679 _ => unreachable!(),
18680 }
18681 }
18682 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
18683 if rp0 != rp1 {
18684 return Ok(None);
18685 }
18686 let rp = rp0;
18687 let rpb: u32 = 4;
18688 let mr1 = rp && Self::q40_mr1_on();
18690 let nb = |o: usize| {
18691 if mr1 {
18692 (o as u32).div_ceil(rpb)
18693 } else {
18694 (o as u32).div_ceil(2).div_ceil(rpb)
18695 }
18696 };
18697 let grid = nb(o0) + nb(o1);
18698 let mut y0 = self.alloc_uninit::<f32>(o0)?;
18699 let mut y1 = self.alloc_uninit::<f32>(o1)?;
18700 let f = self.func(if mr1 {
18701 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
18702 } else if rp {
18703 "qmatvec_q4_0_mmvq_fused2_rp"
18704 } else {
18705 "qmatvec_q4_0_mmvq_fused2"
18706 });
18707 let cfg = LaunchConfig {
18708 grid_dim: (grid, 1, 1),
18709 block_dim: (32, rpb, 1),
18710 shared_mem_bytes: 0,
18711 };
18712 let inf = w0.in_features() as i32;
18713 let (oo0, oo1) = (o0 as i32, o1 as i32);
18714 let (r0, r1) = (rb0 as i64, rb1 as i64);
18715 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18717 {
18718 use cudarc::driver::{DevicePtr, DevicePtrMut};
18719 let s = &self.gpu.stream();
18720 let (p0, _g0) = b0.device_ptr(s);
18721 let (p1, _g1) = b1.device_ptr(s);
18722 let (paq, _g2) = aq.device_ptr(s);
18723 let (pad, _g3) = ad.device_ptr(s);
18724 let (py0, _g4) = y0.device_ptr_mut(s);
18725 let (py1, _g5) = y1.device_ptr_mut(s);
18726 let mut ps = [
18727 &p0 as *const _ as *mut std::ffi::c_void,
18728 &p1 as *const _ as *mut _,
18729 &paq as *const _ as *mut _,
18730 &pad as *const _ as *mut _,
18731 &py0 as *const _ as *mut _,
18732 &py1 as *const _ as *mut _,
18733 &inf as *const _ as *mut _,
18734 &oo0 as *const _ as *mut _,
18735 &oo1 as *const _ as *mut _,
18736 &r0 as *const _ as *mut _,
18737 &r1 as *const _ as *mut _,
18738 ];
18739 unsafe {
18740 self.launch_pdl(
18741 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
18742 (grid, 1, 1),
18743 (32, rpb, 1),
18744 &mut ps,
18745 )?;
18746 }
18747 }
18748 return Ok(Some((y0, y1)));
18749 }
18750 let __s_b = self.gpu.stream();
18751 let mut b = __s_b.launch_builder(&f);
18752 b.arg(b0)
18753 .arg(b1)
18754 .arg(aq)
18755 .arg(ad)
18756 .arg(&mut y0)
18757 .arg(&mut y1)
18758 .arg(&inf)
18759 .arg(&oo0)
18760 .arg(&oo1)
18761 .arg(&r0)
18762 .arg(&r1);
18763 unsafe {
18764 b.launch(cfg)?;
18765 }
18766 Ok(Some((y0, y1)))
18767 }
18768
18769 pub fn matmul_q4_fused2_into(
18771 &self,
18772 w0: &crate::model::GpuTensor,
18773 w1: &crate::model::GpuTensor,
18774 aq: &CudaSlice<i8>,
18775 ad: &CudaSlice<f32>,
18776 y0: &mut CudaSlice<f32>,
18777 y1: &mut CudaSlice<f32>,
18778 ) -> Result<bool, Box<dyn std::error::Error>> {
18779 use crate::model::GpuTensor;
18780 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18781 match w {
18782 GpuTensor::Quant {
18783 qtype, row_bytes, ..
18784 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18785 _ => None,
18786 }
18787 };
18788 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
18789 return Ok(false);
18790 };
18791 if w0.in_features() != w1.in_features() {
18792 return Ok(false);
18793 }
18794 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18795 match w {
18796 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18797 Some(m) => (m, true),
18798 None => (bytes, *rp),
18799 },
18800 _ => unreachable!(),
18801 }
18802 }
18803 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
18804 if rp0 != rp1 {
18805 return Ok(false);
18806 }
18807 let rp = rp0;
18808 let rpb: u32 = 4;
18809 let mr1 = rp && Self::q40_mr1_on();
18810 let nb = |o: usize| {
18811 if mr1 {
18812 (o as u32).div_ceil(rpb)
18813 } else {
18814 (o as u32).div_ceil(2).div_ceil(rpb)
18815 }
18816 };
18817 let grid = nb(o0) + nb(o1);
18818 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
18819 let f = self.func(if mr1 {
18820 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
18821 } else if rp {
18822 "qmatvec_q4_0_mmvq_fused2_rp"
18823 } else {
18824 "qmatvec_q4_0_mmvq_fused2"
18825 });
18826 let cfg = LaunchConfig {
18827 grid_dim: (grid, 1, 1),
18828 block_dim: (32, rpb, 1),
18829 shared_mem_bytes: 0,
18830 };
18831 let inf = w0.in_features() as i32;
18832 let (oo0, oo1) = (o0 as i32, o1 as i32);
18833 let (r0, r1) = (rb0 as i64, rb1 as i64);
18834 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18836 use cudarc::driver::{DevicePtr, DevicePtrMut};
18837 let s = &self.gpu.stream();
18838 let (p0, _g0) = b0.device_ptr(s);
18839 let (p1, _g1) = b1.device_ptr(s);
18840 let (paq, _g2) = aq.device_ptr(s);
18841 let (pad, _g3) = ad.device_ptr(s);
18842 let (py0, _g4) = y0.device_ptr_mut(s);
18843 let (py1, _g5) = y1.device_ptr_mut(s);
18844 let mut ps = [
18845 &p0 as *const _ as *mut std::ffi::c_void,
18846 &p1 as *const _ as *mut _,
18847 &paq as *const _ as *mut _,
18848 &pad as *const _ as *mut _,
18849 &py0 as *const _ as *mut _,
18850 &py1 as *const _ as *mut _,
18851 &inf as *const _ as *mut _,
18852 &oo0 as *const _ as *mut _,
18853 &oo1 as *const _ as *mut _,
18854 &r0 as *const _ as *mut _,
18855 &r1 as *const _ as *mut _,
18856 ];
18857 unsafe {
18858 self.launch_pdl(
18859 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
18860 (grid, 1, 1),
18861 (32, rpb, 1),
18862 &mut ps,
18863 )?;
18864 }
18865 return Ok(true);
18866 }
18867 let __s_b = self.gpu.stream();
18868 let mut b = __s_b.launch_builder(&f);
18869 b.arg(b0)
18870 .arg(b1)
18871 .arg(aq)
18872 .arg(ad)
18873 .arg(&mut *y0)
18874 .arg(&mut *y1)
18875 .arg(&inf)
18876 .arg(&oo0)
18877 .arg(&oo1)
18878 .arg(&r0)
18879 .arg(&r1);
18880 unsafe {
18881 b.launch(cfg)?;
18882 }
18883 Ok(true)
18884 }
18885
18886 #[allow(clippy::type_complexity)] pub fn matmul_q4_fused2_batched(
18892 &self,
18893 w0: &crate::model::GpuTensor,
18894 w1: &crate::model::GpuTensor,
18895 aq: &CudaSlice<i8>,
18896 ad: &CudaSlice<f32>,
18897 m: usize,
18898 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18899 use crate::model::GpuTensor;
18900 if !(2..=8).contains(&m) {
18901 return Ok(None);
18902 }
18903 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18904 match w {
18905 GpuTensor::Quant {
18906 qtype, row_bytes, ..
18907 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18908 _ => None,
18909 }
18910 };
18911 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
18912 return Ok(None);
18913 };
18914 if w0.in_features() != w1.in_features() {
18915 return Ok(None);
18916 }
18917 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18918 match w {
18919 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18920 Some(mr) => (mr, true),
18921 None => (bytes, *rp),
18922 },
18923 _ => unreachable!(),
18924 }
18925 }
18926 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
18927 if !rp0 || !rp1 {
18928 return Ok(None);
18929 }
18930 let mcols = Self::batched_mcols(m);
18931 let rpb: u32 = 4;
18932 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
18933 let grid = nb(o0) + nb(o1);
18934 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
18935 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
18936 let f = self.func(match mcols {
18937 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
18938 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
18939 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
18940 });
18941 let cfg = LaunchConfig {
18942 grid_dim: (grid, 1, 1),
18943 block_dim: (32, rpb, 1),
18944 shared_mem_bytes: 0,
18945 };
18946 let inf = w0.in_features() as i32;
18947 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
18948 let rb = rb0 as i64;
18949 let __s_b = self.gpu.stream();
18950 let mut b = __s_b.launch_builder(&f);
18951 b.arg(b0)
18952 .arg(b1)
18953 .arg(aq)
18954 .arg(ad)
18955 .arg(&mut y0)
18956 .arg(&mut y1)
18957 .arg(&inf)
18958 .arg(&oo0)
18959 .arg(&oo1)
18960 .arg(&mi)
18961 .arg(&rb);
18962 unsafe {
18963 b.launch(cfg)?;
18964 }
18965 Ok(Some((y0, y1)))
18966 }
18967
18968 #[allow(clippy::too_many_arguments)]
18971 #[allow(clippy::type_complexity)] pub fn matmul_q4_fused3_batched(
18973 &self,
18974 w0: &crate::model::GpuTensor,
18975 w1: &crate::model::GpuTensor,
18976 w2: &crate::model::GpuTensor,
18977 aq: &CudaSlice<i8>,
18978 ad: &CudaSlice<f32>,
18979 m: usize,
18980 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
18981 {
18982 use crate::model::GpuTensor;
18983 if !(2..=8).contains(&m) {
18984 return Ok(None);
18985 }
18986 let q4 = |w: &GpuTensor| -> Option<usize> {
18987 match w {
18988 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
18989 _ => None,
18990 }
18991 };
18992 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
18993 return Ok(None);
18994 };
18995 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
18996 return Ok(None);
18997 }
18998 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18999 match w {
19000 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
19001 Some(mr) => (mr, true),
19002 None => (bytes, *rp),
19003 },
19004 _ => unreachable!(),
19005 }
19006 }
19007 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
19008 if !rp0 || !rp1 || !rp2 {
19009 return Ok(None);
19010 }
19011 let mcols = Self::batched_mcols(m);
19012 let rpb: u32 = 4;
19013 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
19014 let grid = nb(o0) + nb(o1) + nb(o2);
19015 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
19016 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
19017 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
19018 let f = self.func(match mcols {
19019 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
19020 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
19021 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
19022 });
19023 let cfg = LaunchConfig {
19024 grid_dim: (grid, 1, 1),
19025 block_dim: (32, rpb, 1),
19026 shared_mem_bytes: 0,
19027 };
19028 let inf = w0.in_features() as i32;
19029 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
19030 let rb = 0i64;
19031 let __s_b = self.gpu.stream();
19032 let mut b = __s_b.launch_builder(&f);
19033 b.arg(b0)
19034 .arg(b1)
19035 .arg(b2)
19036 .arg(aq)
19037 .arg(ad)
19038 .arg(&mut y0)
19039 .arg(&mut y1)
19040 .arg(&mut y2)
19041 .arg(&inf)
19042 .arg(&oo0)
19043 .arg(&oo1)
19044 .arg(&oo2)
19045 .arg(&mi)
19046 .arg(&rb);
19047 unsafe {
19048 b.launch(cfg)?;
19049 }
19050 Ok(Some((y0, y1, y2)))
19051 }
19052
19053 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused3(
19055 &self,
19056 w0: &crate::model::GpuTensor,
19057 w1: &crate::model::GpuTensor,
19058 w2: &crate::model::GpuTensor,
19059 aq: &CudaSlice<i8>,
19060 ad: &CudaSlice<f32>,
19061 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
19062 {
19063 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
19066 return Ok(Some(self.e4m3_fused3_core(
19067 p0.0,
19068 p1.0,
19069 p2.0,
19070 aq,
19071 ad,
19072 w0.in_features(),
19073 p0.1,
19074 p1.1,
19075 p2.1,
19076 p0.2,
19077 p0.3,
19078 p1.3,
19079 p2.3,
19080 )?));
19081 }
19082 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
19083 return Ok(None);
19084 };
19085 Ok(Some(self.q8_fused3_core(
19086 p0.0,
19087 p1.0,
19088 p2.0,
19089 aq,
19090 ad,
19091 w0.in_features(),
19092 p0.1,
19093 p1.1,
19094 p2.1,
19095 p0.2,
19096 )?))
19097 }
19098
19099 #[allow(clippy::too_many_arguments)]
19100 #[allow(clippy::type_complexity)] fn q8_fused3_core(
19102 &self,
19103 b0: &CudaSlice<u8>,
19104 b1: &CudaSlice<u8>,
19105 b2: &CudaSlice<u8>,
19106 aq: &CudaSlice<i8>,
19107 ad: &CudaSlice<f32>,
19108 in_f: usize,
19109 out0: usize,
19110 out1: usize,
19111 out2: usize,
19112 row_bytes: usize,
19113 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19114 const ROWS_PER_BLOCK: u32 = 4;
19115 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19116 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19117 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19118 let f = self.func("qmatvec_q8_0_mmvq_fused3");
19119 let mut y0 = self.alloc_uninit::<f32>(out0)?;
19120 let mut y1 = self.alloc_uninit::<f32>(out1)?;
19121 let mut y2 = self.alloc_uninit::<f32>(out2)?;
19122 let cfg = LaunchConfig {
19123 grid_dim: (nb0 + nb1 + nb2, 1, 1),
19124 block_dim: (32, ROWS_PER_BLOCK, 1),
19125 shared_mem_bytes: 0,
19126 };
19127 let (inf, o0, o1, o2, rbl) = (
19128 in_f as i32,
19129 out0 as i32,
19130 out1 as i32,
19131 out2 as i32,
19132 row_bytes as i64,
19133 );
19134 let __s_b = self.gpu.stream();
19135 let mut b = __s_b.launch_builder(&f);
19136 b.arg(b0)
19137 .arg(b1)
19138 .arg(b2)
19139 .arg(aq)
19140 .arg(ad)
19141 .arg(&mut y0)
19142 .arg(&mut y1)
19143 .arg(&mut y2)
19144 .arg(&inf)
19145 .arg(&o0)
19146 .arg(&o1)
19147 .arg(&o2)
19148 .arg(&rbl);
19149 unsafe {
19150 b.launch(cfg)?;
19151 }
19152 Ok((y0, y1, y2))
19153 }
19154
19155 #[allow(clippy::too_many_arguments)]
19157 #[allow(clippy::type_complexity)] pub fn qmatvec_q8_fused3_raw(
19159 &self,
19160 b0: &CudaSlice<u8>,
19161 b1: &CudaSlice<u8>,
19162 b2: &CudaSlice<u8>,
19163 x: &CudaSlice<f32>,
19164 in_f: usize,
19165 out0: usize,
19166 out1: usize,
19167 out2: usize,
19168 row_bytes: usize,
19169 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19170 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
19171 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
19172 }
19173
19174 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused2_t(
19186 &self,
19187 w0: &crate::model::GpuTensor,
19188 w1: &crate::model::GpuTensor,
19189 aq: &CudaSlice<i8>,
19190 ad: &CudaSlice<f32>,
19191 m: usize,
19192 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
19193 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
19197 return Ok(None);
19198 }
19199 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
19202 if m > 4 && !Self::b8_enabled() {
19203 return Ok(None);
19204 }
19205 return Ok(Some(self.e4m3_fused2_t_core(
19206 p0.0,
19207 p1.0,
19208 aq,
19209 ad,
19210 m,
19211 w0.in_features(),
19212 p0.1,
19213 p1.1,
19214 p0.2,
19215 p0.3,
19216 p1.3,
19217 )?));
19218 }
19219 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
19220 return Ok(None);
19221 };
19222 Ok(Some(self.q8_fused2_t_core(
19223 p0.0,
19224 p1.0,
19225 aq,
19226 ad,
19227 m,
19228 w0.in_features(),
19229 p0.1,
19230 p1.1,
19231 p0.2,
19232 )?))
19233 }
19234
19235 #[allow(clippy::too_many_arguments)]
19236 fn q8_fused2_t_core(
19237 &self,
19238 b0: &CudaSlice<u8>,
19239 b1: &CudaSlice<u8>,
19240 aq: &CudaSlice<i8>,
19241 ad: &CudaSlice<f32>,
19242 m: usize,
19243 in_f: usize,
19244 out0: usize,
19245 out1: usize,
19246 row_bytes: usize,
19247 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19248 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19250 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19251 let f = self.func(match Self::batched_mcols(m) {
19252 2 => "qmatvec_q8_0_mmvq_fused2_b2",
19253 4 => "qmatvec_q8_0_mmvq_fused2_b4",
19254 _ => "qmatvec_q8_0_mmvq_fused2_b8",
19256 });
19257 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
19258 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
19259 let cfg = LaunchConfig {
19260 grid_dim: (nb0 + nb1, 1, 1),
19261 block_dim: (32, ROWS_PER_BLOCK, 1),
19262 shared_mem_bytes: 0,
19263 };
19264 let (inf, o0, o1, mi, rbl) = (
19265 in_f as i32,
19266 out0 as i32,
19267 out1 as i32,
19268 m as i32,
19269 row_bytes as i64,
19270 );
19271 let __s_b = self.gpu.stream();
19272 let mut b = __s_b.launch_builder(&f);
19273 b.arg(b0)
19274 .arg(b1)
19275 .arg(aq)
19276 .arg(ad)
19277 .arg(&mut y0)
19278 .arg(&mut y1)
19279 .arg(&inf)
19280 .arg(&o0)
19281 .arg(&o1)
19282 .arg(&mi)
19283 .arg(&rbl);
19284 unsafe {
19285 b.launch(cfg)?;
19286 }
19287 Ok((y0, y1))
19288 }
19289
19290 #[allow(clippy::too_many_arguments)]
19293 pub fn qmatvec_q8_fused2_t_raw(
19294 &self,
19295 b0: &CudaSlice<u8>,
19296 b1: &CudaSlice<u8>,
19297 x: &CudaSlice<f32>,
19298 m: usize,
19299 in_f: usize,
19300 out0: usize,
19301 out1: usize,
19302 row_bytes: usize,
19303 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19304 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19305 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
19306 }
19307
19308 #[allow(clippy::too_many_arguments)]
19311 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused3_t(
19313 &self,
19314 w0: &crate::model::GpuTensor,
19315 w1: &crate::model::GpuTensor,
19316 w2: &crate::model::GpuTensor,
19317 aq: &CudaSlice<i8>,
19318 ad: &CudaSlice<f32>,
19319 m: usize,
19320 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
19321 {
19322 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
19323 return Ok(None);
19324 }
19325 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
19326 return Ok(Some(self.e4m3_fused3_t_core(
19327 p0.0,
19328 p1.0,
19329 p2.0,
19330 aq,
19331 ad,
19332 m,
19333 w0.in_features(),
19334 p0.1,
19335 p1.1,
19336 p2.1,
19337 p0.2,
19338 p0.3,
19339 p1.3,
19340 p2.3,
19341 )?));
19342 }
19343 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
19344 return Ok(None);
19345 };
19346 Ok(Some(self.q8_fused3_t_core(
19347 p0.0,
19348 p1.0,
19349 p2.0,
19350 aq,
19351 ad,
19352 m,
19353 w0.in_features(),
19354 p0.1,
19355 p1.1,
19356 p2.1,
19357 p0.2,
19358 )?))
19359 }
19360
19361 #[allow(clippy::too_many_arguments)]
19362 #[allow(clippy::type_complexity)] fn q8_fused3_t_core(
19364 &self,
19365 b0: &CudaSlice<u8>,
19366 b1: &CudaSlice<u8>,
19367 b2: &CudaSlice<u8>,
19368 aq: &CudaSlice<i8>,
19369 ad: &CudaSlice<f32>,
19370 m: usize,
19371 in_f: usize,
19372 out0: usize,
19373 out1: usize,
19374 out2: usize,
19375 row_bytes: usize,
19376 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19377 const ROWS_PER_BLOCK: u32 = 4;
19378 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19379 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19380 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19381 let f = self.func(if Self::batched_mcols(m) == 2 {
19382 "qmatvec_q8_0_mmvq_fused3_b2"
19383 } else {
19384 "qmatvec_q8_0_mmvq_fused3_b4"
19385 });
19386 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
19387 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
19388 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
19389 let cfg = LaunchConfig {
19390 grid_dim: (nb0 + nb1 + nb2, 1, 1),
19391 block_dim: (32, ROWS_PER_BLOCK, 1),
19392 shared_mem_bytes: 0,
19393 };
19394 let (inf, o0, o1, o2, mi, rbl) = (
19395 in_f as i32,
19396 out0 as i32,
19397 out1 as i32,
19398 out2 as i32,
19399 m as i32,
19400 row_bytes as i64,
19401 );
19402 let __s_b = self.gpu.stream();
19403 let mut b = __s_b.launch_builder(&f);
19404 b.arg(b0)
19405 .arg(b1)
19406 .arg(b2)
19407 .arg(aq)
19408 .arg(ad)
19409 .arg(&mut y0)
19410 .arg(&mut y1)
19411 .arg(&mut y2)
19412 .arg(&inf)
19413 .arg(&o0)
19414 .arg(&o1)
19415 .arg(&o2)
19416 .arg(&mi)
19417 .arg(&rbl);
19418 unsafe {
19419 b.launch(cfg)?;
19420 }
19421 Ok((y0, y1, y2))
19422 }
19423
19424 #[allow(clippy::too_many_arguments)]
19426 #[allow(clippy::type_complexity)] pub fn qmatvec_q8_fused3_t_raw(
19428 &self,
19429 b0: &CudaSlice<u8>,
19430 b1: &CudaSlice<u8>,
19431 b2: &CudaSlice<u8>,
19432 x: &CudaSlice<f32>,
19433 m: usize,
19434 in_f: usize,
19435 out0: usize,
19436 out1: usize,
19437 out2: usize,
19438 row_bytes: usize,
19439 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19440 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19441 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
19442 }
19443
19444 pub fn q8_ffn_fuse2_on(&self) -> bool {
19448 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19449 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
19450 }
19451
19452 #[allow(clippy::type_complexity)]
19458 fn q8_fused_params<'w, const N: usize>(
19459 &self,
19460 ws: &[&'w crate::model::GpuTensor; N],
19461 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
19462 use crate::model::GpuTensor;
19463 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
19464 return None;
19465 }
19466 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
19467 return None;
19468 }
19469 let in_f = ws[0].in_features();
19470 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
19471 for (i, w) in ws.iter().enumerate() {
19472 match w {
19473 GpuTensor::Quant {
19474 bytes,
19475 qtype,
19476 row_bytes,
19477 scale,
19478 ..
19479 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
19480 out[i] = Some((bytes, w.out_features(), *row_bytes))
19481 }
19482 _ => return None,
19483 }
19484 }
19485 Some(out.map(|o| o.unwrap()))
19486 }
19487
19488 pub fn e4m3_dual_on(&self) -> bool {
19491 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19492 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
19493 }
19494
19495 #[allow(clippy::type_complexity)]
19507 fn e4m3_fused_params<'w, const N: usize>(
19508 &self,
19509 ws: &[&'w crate::model::GpuTensor; N],
19510 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
19511 use crate::model::GpuTensor;
19512 if !self.e4m3_dual_on() {
19513 return None;
19514 }
19515 let in_f = ws[0].in_features();
19516 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
19517 for (i, w) in ws.iter().enumerate() {
19518 match w {
19519 GpuTensor::Quant {
19520 bytes,
19521 qtype,
19522 row_bytes,
19523 scale,
19524 rp,
19525 rp4,
19526 ..
19527 } if *qtype == QT_F8_E4M3
19528 && w.in_features() == in_f
19529 && *row_bytes == in_f
19530 && !*rp
19531 && rp4.is_none() =>
19532 {
19533 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
19534 }
19535 _ => return None,
19536 }
19537 }
19538 Some(out.map(|o| o.unwrap()))
19539 }
19540
19541 #[allow(clippy::too_many_arguments)]
19545 fn e4m3_fused2_core(
19546 &self,
19547 b0: &CudaSlice<u8>,
19548 b1: &CudaSlice<u8>,
19549 aq: &CudaSlice<i8>,
19550 ad: &CudaSlice<f32>,
19551 in_f: usize,
19552 out0: usize,
19553 out1: usize,
19554 row_bytes: usize,
19555 ws0: f32,
19556 ws1: f32,
19557 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19558 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19560 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19561 let f = self.func("qmatvec_e4m3_mmvq_fused2");
19562 let mut y0 = self.alloc_uninit::<f32>(out0)?;
19563 let mut y1 = self.alloc_uninit::<f32>(out1)?;
19564 let cfg = LaunchConfig {
19565 grid_dim: (nb0 + nb1, 1, 1),
19566 block_dim: (32, ROWS_PER_BLOCK, 1),
19567 shared_mem_bytes: 0,
19568 };
19569 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
19570 let __s_b = self.gpu.stream();
19571 let mut b = __s_b.launch_builder(&f);
19572 b.arg(b0)
19573 .arg(b1)
19574 .arg(aq)
19575 .arg(ad)
19576 .arg(&mut y0)
19577 .arg(&mut y1)
19578 .arg(&inf)
19579 .arg(&o0)
19580 .arg(&o1)
19581 .arg(&rbl)
19582 .arg(&ws0)
19583 .arg(&ws1);
19584 unsafe {
19585 b.launch(cfg)?;
19586 }
19587 Ok((y0, y1))
19588 }
19589
19590 #[allow(clippy::too_many_arguments)]
19592 #[allow(clippy::type_complexity)] fn e4m3_fused3_core(
19594 &self,
19595 b0: &CudaSlice<u8>,
19596 b1: &CudaSlice<u8>,
19597 b2: &CudaSlice<u8>,
19598 aq: &CudaSlice<i8>,
19599 ad: &CudaSlice<f32>,
19600 in_f: usize,
19601 out0: usize,
19602 out1: usize,
19603 out2: usize,
19604 row_bytes: usize,
19605 ws0: f32,
19606 ws1: f32,
19607 ws2: f32,
19608 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19609 const ROWS_PER_BLOCK: u32 = 4;
19610 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19611 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19612 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19613 let f = self.func("qmatvec_e4m3_mmvq_fused3");
19614 let mut y0 = self.alloc_uninit::<f32>(out0)?;
19615 let mut y1 = self.alloc_uninit::<f32>(out1)?;
19616 let mut y2 = self.alloc_uninit::<f32>(out2)?;
19617 let cfg = LaunchConfig {
19618 grid_dim: (nb0 + nb1 + nb2, 1, 1),
19619 block_dim: (32, ROWS_PER_BLOCK, 1),
19620 shared_mem_bytes: 0,
19621 };
19622 let (inf, o0, o1, o2, rbl) = (
19623 in_f as i32,
19624 out0 as i32,
19625 out1 as i32,
19626 out2 as i32,
19627 row_bytes as i64,
19628 );
19629 let __s_b = self.gpu.stream();
19630 let mut b = __s_b.launch_builder(&f);
19631 b.arg(b0)
19632 .arg(b1)
19633 .arg(b2)
19634 .arg(aq)
19635 .arg(ad)
19636 .arg(&mut y0)
19637 .arg(&mut y1)
19638 .arg(&mut y2)
19639 .arg(&inf)
19640 .arg(&o0)
19641 .arg(&o1)
19642 .arg(&o2)
19643 .arg(&rbl)
19644 .arg(&ws0)
19645 .arg(&ws1)
19646 .arg(&ws2);
19647 unsafe {
19648 b.launch(cfg)?;
19649 }
19650 Ok((y0, y1, y2))
19651 }
19652
19653 #[allow(clippy::too_many_arguments)]
19660 pub fn e4m3_fused6_into(
19661 &self,
19662 w: [&CudaSlice<u8>; 6],
19663 aq: &CudaSlice<i8>,
19664 ad: &CudaSlice<f32>,
19665 in_f: usize,
19666 dims: [usize; 6],
19667 row_bytes: usize,
19668 ws: [f32; 6],
19669 outs: &mut [CudaSlice<f32>; 6],
19670 ) -> Result<(), Box<dyn std::error::Error>> {
19671 self.e4m3_fused6_into_arm(w, aq, ad, in_f, dims, row_bytes, ws, outs, 0)
19672 }
19673
19674 #[allow(clippy::too_many_arguments)]
19679 pub fn e4m3_fused6_into_arm(
19680 &self,
19681 w: [&CudaSlice<u8>; 6],
19682 aq: &CudaSlice<i8>,
19683 ad: &CudaSlice<f32>,
19684 in_f: usize,
19685 dims: [usize; 6],
19686 row_bytes: usize,
19687 ws: [f32; 6],
19688 outs: &mut [CudaSlice<f32>; 6],
19689 arm: u32,
19690 ) -> Result<(), Box<dyn std::error::Error>> {
19691 const ROWS_PER_BLOCK: u32 = 4;
19692 let blocks: u32 = dims
19693 .iter()
19694 .map(|&o| (o as u32).div_ceil(ROWS_PER_BLOCK))
19695 .sum();
19696 if arm != 0 {
19700 return Err(format!(
19701 "e4m3 six-group arm {arm} was removed 2026-09-05 (the ILP twins measured neutral)"
19702 )
19703 .into());
19704 }
19705 let f = self.func("qmatvec_e4m3_mmvq_fused6");
19706 let cfg = LaunchConfig {
19707 grid_dim: (blocks, 1, 1),
19708 block_dim: (32, ROWS_PER_BLOCK, 1),
19709 shared_mem_bytes: 0,
19710 };
19711 let inf = in_f as i32;
19712 let o: [i32; 6] = std::array::from_fn(|i| dims[i] as i32);
19713 let rbl = row_bytes as i64;
19714 let (o0, o1) = outs.split_at_mut(1);
19715 let (o1, o2) = o1.split_at_mut(1);
19716 let (o2, o3) = o2.split_at_mut(1);
19717 let (o3, o4) = o3.split_at_mut(1);
19718 let (o4, o5) = o4.split_at_mut(1);
19719 let __s_b = self.gpu.stream();
19720 let mut b = __s_b.launch_builder(&f);
19721 b.arg(w[0])
19722 .arg(w[1])
19723 .arg(w[2])
19724 .arg(w[3])
19725 .arg(w[4])
19726 .arg(w[5])
19727 .arg(aq)
19728 .arg(ad)
19729 .arg(&mut o0[0])
19730 .arg(&mut o1[0])
19731 .arg(&mut o2[0])
19732 .arg(&mut o3[0])
19733 .arg(&mut o4[0])
19734 .arg(&mut o5[0])
19735 .arg(&inf)
19736 .arg(&o[0])
19737 .arg(&o[1])
19738 .arg(&o[2])
19739 .arg(&o[3])
19740 .arg(&o[4])
19741 .arg(&o[5])
19742 .arg(&rbl)
19743 .arg(&ws[0])
19744 .arg(&ws[1])
19745 .arg(&ws[2])
19746 .arg(&ws[3])
19747 .arg(&ws[4])
19748 .arg(&ws[5]);
19749 unsafe {
19750 b.launch(cfg)?;
19751 }
19752 Ok(())
19753 }
19754
19755 #[allow(clippy::too_many_arguments)]
19759 fn e4m3_fused2_t_core(
19760 &self,
19761 b0: &CudaSlice<u8>,
19762 b1: &CudaSlice<u8>,
19763 aq: &CudaSlice<i8>,
19764 ad: &CudaSlice<f32>,
19765 m: usize,
19766 in_f: usize,
19767 out0: usize,
19768 out1: usize,
19769 row_bytes: usize,
19770 ws0: f32,
19771 ws1: f32,
19772 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19773 const ROWS_PER_BLOCK: u32 = 4;
19774 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19775 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19776 let f = self.func(match Self::batched_mcols(m) {
19777 2 => "qmatvec_e4m3_mmvq_fused2_b2",
19778 4 => "qmatvec_e4m3_mmvq_fused2_b4",
19779 _ => "qmatvec_e4m3_mmvq_fused2_b8",
19780 });
19781 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
19782 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
19783 let cfg = LaunchConfig {
19784 grid_dim: (nb0 + nb1, 1, 1),
19785 block_dim: (32, ROWS_PER_BLOCK, 1),
19786 shared_mem_bytes: 0,
19787 };
19788 let (inf, o0, o1, mi, rbl) = (
19789 in_f as i32,
19790 out0 as i32,
19791 out1 as i32,
19792 m as i32,
19793 row_bytes as i64,
19794 );
19795 let __s_b = self.gpu.stream();
19796 let mut b = __s_b.launch_builder(&f);
19797 b.arg(b0)
19798 .arg(b1)
19799 .arg(aq)
19800 .arg(ad)
19801 .arg(&mut y0)
19802 .arg(&mut y1)
19803 .arg(&inf)
19804 .arg(&o0)
19805 .arg(&o1)
19806 .arg(&mi)
19807 .arg(&rbl);
19808 unsafe {
19809 b.launch(cfg)?;
19810 }
19811 if ws0 != 1.0 {
19812 self.scale_inplace(&mut y0, ws0, m * out0)?;
19813 }
19814 if ws1 != 1.0 {
19815 self.scale_inplace(&mut y1, ws1, m * out1)?;
19816 }
19817 Ok((y0, y1))
19818 }
19819
19820 #[allow(clippy::too_many_arguments)]
19822 #[allow(clippy::type_complexity)] fn e4m3_fused3_t_core(
19824 &self,
19825 b0: &CudaSlice<u8>,
19826 b1: &CudaSlice<u8>,
19827 b2: &CudaSlice<u8>,
19828 aq: &CudaSlice<i8>,
19829 ad: &CudaSlice<f32>,
19830 m: usize,
19831 in_f: usize,
19832 out0: usize,
19833 out1: usize,
19834 out2: usize,
19835 row_bytes: usize,
19836 ws0: f32,
19837 ws1: f32,
19838 ws2: f32,
19839 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19840 const ROWS_PER_BLOCK: u32 = 4;
19841 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19842 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19843 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19844 let f = self.func(if Self::batched_mcols(m) == 2 {
19845 "qmatvec_e4m3_mmvq_fused3_b2"
19846 } else {
19847 "qmatvec_e4m3_mmvq_fused3_b4"
19848 });
19849 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
19850 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
19851 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
19852 let cfg = LaunchConfig {
19853 grid_dim: (nb0 + nb1 + nb2, 1, 1),
19854 block_dim: (32, ROWS_PER_BLOCK, 1),
19855 shared_mem_bytes: 0,
19856 };
19857 let (inf, o0, o1, o2, mi, rbl) = (
19858 in_f as i32,
19859 out0 as i32,
19860 out1 as i32,
19861 out2 as i32,
19862 m as i32,
19863 row_bytes as i64,
19864 );
19865 let __s_b = self.gpu.stream();
19866 let mut b = __s_b.launch_builder(&f);
19867 b.arg(b0)
19868 .arg(b1)
19869 .arg(b2)
19870 .arg(aq)
19871 .arg(ad)
19872 .arg(&mut y0)
19873 .arg(&mut y1)
19874 .arg(&mut y2)
19875 .arg(&inf)
19876 .arg(&o0)
19877 .arg(&o1)
19878 .arg(&o2)
19879 .arg(&mi)
19880 .arg(&rbl);
19881 unsafe {
19882 b.launch(cfg)?;
19883 }
19884 if ws0 != 1.0 {
19885 self.scale_inplace(&mut y0, ws0, m * out0)?;
19886 }
19887 if ws1 != 1.0 {
19888 self.scale_inplace(&mut y1, ws1, m * out1)?;
19889 }
19890 if ws2 != 1.0 {
19891 self.scale_inplace(&mut y2, ws2, m * out2)?;
19892 }
19893 Ok((y0, y1, y2))
19894 }
19895
19896 #[allow(clippy::too_many_arguments)] pub fn qmatvec_e4m3_blk_mmvq(
19907 &self,
19908 bytes: &CudaSlice<u8>,
19909 aq: &CudaSlice<i8>,
19910 ad: &CudaSlice<f32>,
19911 scales: &CudaSlice<f32>,
19912 m: usize,
19913 in_f: usize,
19914 out_f: usize,
19915 row_bytes: usize,
19916 scale_cols: usize,
19917 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19918 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
19920 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
19921 )?;
19922 Ok(y)
19923 }
19924
19925 #[allow(clippy::too_many_arguments)]
19927 pub fn qmatvec_e4m3_blk_mmvq_into(
19928 &self,
19929 bytes: &CudaSlice<u8>,
19930 aq: &CudaSlice<i8>,
19931 ad: &CudaSlice<f32>,
19932 scales: &CudaSlice<f32>,
19933 m: usize,
19934 in_f: usize,
19935 out_f: usize,
19936 row_bytes: usize,
19937 scale_cols: usize,
19938 y: &mut CudaSlice<f32>,
19939 ) -> Result<(), Box<dyn std::error::Error>> {
19940 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
19942 let cfg = LaunchConfig {
19943 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
19944 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
19947 let (inf, outf, mi, rb, sc) = (
19948 in_f as i32,
19949 out_f as i32,
19950 m as i32,
19951 row_bytes as i64,
19952 scale_cols as i32,
19953 );
19954 let __s_b = self.gpu.stream();
19955 let mut b = __s_b.launch_builder(&f);
19956 b.arg(bytes)
19957 .arg(aq)
19958 .arg(ad)
19959 .arg(scales)
19960 .arg(&mut *y)
19961 .arg(&inf)
19962 .arg(&outf)
19963 .arg(&mi)
19964 .arg(&rb)
19965 .arg(&sc);
19966 unsafe {
19967 b.launch(cfg)?;
19968 }
19969 Ok(())
19970 }
19971
19972 #[allow(clippy::too_many_arguments)]
19978 pub fn qmatvec_e4m3_blk_mmvq_batched(
19979 &self,
19980 bytes: &CudaSlice<u8>,
19981 aq: &CudaSlice<i8>,
19982 ad: &CudaSlice<f32>,
19983 scales: &CudaSlice<f32>,
19984 m: usize,
19985 in_f: usize,
19986 out_f: usize,
19987 row_bytes: usize,
19988 scale_cols: usize,
19989 mcols: usize,
19990 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19991 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
19993 let name = match mcols {
19994 2 => "qmatvec_e4m3_blk_mmvq_b2",
19995 4 => "qmatvec_e4m3_blk_mmvq_b4",
19996 8 => "qmatvec_e4m3_blk_mmvq_b8",
19997 16 => "qmatvec_e4m3_blk_mmvq_b16",
19998 _ => {
19999 return Err(
20000 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
20001 );
20002 }
20003 };
20004 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20005 let f = self.func(name);
20006 let cfg = LaunchConfig {
20007 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
20008 block_dim: (32, ROWS_PER_BLOCK, 1),
20009 shared_mem_bytes: 0,
20010 };
20011 let (inf, outf, mi, rb, sc) = (
20012 in_f as i32,
20013 out_f as i32,
20014 m as i32,
20015 row_bytes as i64,
20016 scale_cols as i32,
20017 );
20018 let __s_b = self.gpu.stream();
20019 let mut b = __s_b.launch_builder(&f);
20020 b.arg(bytes)
20021 .arg(aq)
20022 .arg(ad)
20023 .arg(scales)
20024 .arg(&mut y)
20025 .arg(&inf)
20026 .arg(&outf)
20027 .arg(&mi)
20028 .arg(&rb)
20029 .arg(&sc);
20030 unsafe {
20031 b.launch(cfg)?;
20032 }
20033 Ok(y)
20034 }
20035
20036 #[allow(clippy::too_many_arguments)]
20039 pub fn qmatvec_e4m3_blk_batched_raw(
20040 &self,
20041 bytes: &CudaSlice<u8>,
20042 x: &CudaSlice<f32>,
20043 scales: &CudaSlice<f32>,
20044 m: usize,
20045 in_f: usize,
20046 out_f: usize,
20047 row_bytes: usize,
20048 scale_cols: usize,
20049 mcols: usize,
20050 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20051 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20052 self.qmatvec_e4m3_blk_mmvq_batched(
20053 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
20054 )
20055 }
20056
20057 #[allow(clippy::too_many_arguments)]
20060 pub fn qmatvec_e4m3_blk_mmvq_raw(
20061 &self,
20062 bytes: &CudaSlice<u8>,
20063 x: &CudaSlice<f32>,
20064 scales: &CudaSlice<f32>,
20065 m: usize,
20066 in_f: usize,
20067 out_f: usize,
20068 row_bytes: usize,
20069 scale_cols: usize,
20070 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20071 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20072 self.qmatvec_e4m3_blk_mmvq(
20073 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
20074 )
20075 }
20076
20077 #[allow(clippy::too_many_arguments)]
20080 pub fn qmatvec_e4m3_fused2_raw(
20081 &self,
20082 b0: &CudaSlice<u8>,
20083 b1: &CudaSlice<u8>,
20084 x: &CudaSlice<f32>,
20085 in_f: usize,
20086 out0: usize,
20087 out1: usize,
20088 row_bytes: usize,
20089 ws0: f32,
20090 ws1: f32,
20091 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20092 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20093 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
20094 }
20095
20096 #[allow(clippy::too_many_arguments)]
20097 #[allow(clippy::type_complexity)] pub fn qmatvec_e4m3_fused3_raw(
20099 &self,
20100 b0: &CudaSlice<u8>,
20101 b1: &CudaSlice<u8>,
20102 b2: &CudaSlice<u8>,
20103 x: &CudaSlice<f32>,
20104 in_f: usize,
20105 out0: usize,
20106 out1: usize,
20107 out2: usize,
20108 row_bytes: usize,
20109 ws0: f32,
20110 ws1: f32,
20111 ws2: f32,
20112 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20113 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20114 self.e4m3_fused3_core(
20115 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
20116 )
20117 }
20118
20119 #[allow(clippy::too_many_arguments)]
20123 pub fn qmatvec_e4m3_fused6_raw(
20124 &self,
20125 w: [&CudaSlice<u8>; 6],
20126 x: &CudaSlice<f32>,
20127 in_f: usize,
20128 dims: [usize; 6],
20129 row_bytes: usize,
20130 ws: [f32; 6],
20131 arm: u32,
20132 ) -> Result<[CudaSlice<f32>; 6], Box<dyn std::error::Error>> {
20133 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20134 let mut outs = [
20135 self.alloc_uninit::<f32>(dims[0])?,
20136 self.alloc_uninit::<f32>(dims[1])?,
20137 self.alloc_uninit::<f32>(dims[2])?,
20138 self.alloc_uninit::<f32>(dims[3])?,
20139 self.alloc_uninit::<f32>(dims[4])?,
20140 self.alloc_uninit::<f32>(dims[5])?,
20141 ];
20142 self.e4m3_fused6_into_arm(w, &aq, &ad, in_f, dims, row_bytes, ws, &mut outs, arm)?;
20143 Ok(outs)
20144 }
20145
20146 #[allow(clippy::too_many_arguments)]
20147 pub fn qmatvec_e4m3_fused2_t_raw(
20148 &self,
20149 b0: &CudaSlice<u8>,
20150 b1: &CudaSlice<u8>,
20151 x: &CudaSlice<f32>,
20152 m: usize,
20153 in_f: usize,
20154 out0: usize,
20155 out1: usize,
20156 row_bytes: usize,
20157 ws0: f32,
20158 ws1: f32,
20159 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20160 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20161 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
20162 }
20163
20164 #[allow(clippy::too_many_arguments)]
20165 #[allow(clippy::type_complexity)] pub fn qmatvec_e4m3_fused3_t_raw(
20167 &self,
20168 b0: &CudaSlice<u8>,
20169 b1: &CudaSlice<u8>,
20170 b2: &CudaSlice<u8>,
20171 x: &CudaSlice<f32>,
20172 m: usize,
20173 in_f: usize,
20174 out0: usize,
20175 out1: usize,
20176 out2: usize,
20177 row_bytes: usize,
20178 ws0: f32,
20179 ws1: f32,
20180 ws2: f32,
20181 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20182 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20183 self.e4m3_fused3_t_core(
20184 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
20185 )
20186 }
20187
20188 fn try_e4m3_blk_pre(
20199 &self,
20200 w: &crate::model::GpuTensor,
20201 aq: &CudaSlice<i8>,
20202 ad: &CudaSlice<f32>,
20203 m: usize,
20204 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20205 use crate::model::GpuTensor;
20206 if let GpuTensor::Quant {
20207 bytes,
20208 qtype,
20209 row_bytes,
20210 blk: Some(g),
20211 ..
20212 } = w
20213 && *qtype == QT_F8_E4M3_BLK
20214 {
20215 if (2..=16).contains(&m)
20221 && std::env::var("MEMRA_NO_BATCHED").is_err()
20222 && (m <= 4 || Self::b8_enabled())
20223 {
20224 let mcols = Self::batched_mcols(m);
20225 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
20226 bytes,
20227 aq,
20228 ad,
20229 &g.scales,
20230 m,
20231 w.in_features(),
20232 w.out_features(),
20233 *row_bytes,
20234 g.cols,
20235 mcols,
20236 )?));
20237 }
20238 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
20239 bytes,
20240 aq,
20241 ad,
20242 &g.scales,
20243 m,
20244 w.in_features(),
20245 w.out_features(),
20246 *row_bytes,
20247 g.cols,
20248 )?));
20249 }
20250 Ok(None)
20251 }
20252
20253 fn try_e4m3_blk_prefill(
20300 &self,
20301 w: &crate::model::GpuTensor,
20302 x: &CudaSlice<f32>,
20303 m: usize,
20304 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20305 use crate::model::GpuTensor;
20306 let GpuTensor::Quant {
20307 bytes,
20308 qtype,
20309 blk: Some(g),
20310 ..
20311 } = w
20312 else {
20313 return Ok(None);
20314 };
20315 if *qtype != QT_F8_E4M3_BLK {
20316 return Ok(None);
20317 }
20318 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
20323 return Ok(Some(y));
20324 }
20325 let (in_f, out_f) = (w.in_features(), w.out_features());
20326 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
20327 let tmp = GpuTensor::Quant {
20328 bytes: slab,
20329 qtype: QT_Q8_0,
20330 row_bytes: in_f / 32 * 34,
20331 ne: vec![in_f as u64, out_f as u64],
20332 scale: 1.0,
20333 rp: false,
20334 #[cfg(memra_cutlass)]
20335 cutlass: None,
20336 fp8: None,
20337 blk: None,
20338 f16: None,
20339 rp4: None,
20340 };
20341 Ok(Some(self.matmul(&tmp, x, m)?))
20343 }
20344
20345 #[allow(clippy::type_complexity)] pub fn matmul_pre_noscale(
20347 &self,
20348 w: &crate::model::GpuTensor,
20349 aq: &CudaSlice<i8>,
20350 ad: &CudaSlice<f32>,
20351 m: usize,
20352 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
20353 use crate::model::GpuTensor;
20354 if m == 1
20358 && let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)?
20359 {
20360 return Ok(Some((y, 1.0)));
20361 }
20362 if m != 1 || !self.uses_q8_1_fast(w) {
20364 return Ok(None);
20365 }
20366 let in_f = w.in_features();
20367 let out_f = w.out_features();
20368 let (bytes, qtype, row_bytes, scale, rp) = match w {
20369 GpuTensor::Quant {
20370 bytes,
20371 qtype,
20372 row_bytes,
20373 scale,
20374 rp,
20375 ..
20376 } => (bytes, *qtype, *row_bytes, *scale, *rp),
20377 _ => return Ok(None),
20378 };
20379 if self.mmvq_supports(qtype) {
20381 let (mbytes, mrp) = match w {
20383 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
20384 _ => (bytes, rp),
20385 };
20386 let y = self.qmatvec_mmvq(
20387 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
20388 )?;
20389 return Ok(Some((y, scale)));
20390 }
20391 let name = match qtype {
20393 QT_Q8_0 => "qmatvec_q8_0_dp4a",
20394 QT_Q4_K => "qmatvec_q4_K_dp4a",
20395 QT_Q6_K => "qmatvec_q6_K_dp4a",
20396 QT_Q5_K => "qmatvec_q5_K_dp4a",
20397 QT_Q3_K => "qmatvec_q3_K_dp4a",
20398 QT_NVFP4 => {
20399 if rp {
20400 "qmatvec_nvfp4_dp4a_rp"
20401 } else {
20402 "qmatvec_nvfp4_dp4a"
20403 }
20404 }
20405 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
20406 _ => return Ok(None),
20407 };
20408 let f = self.func(name);
20409 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20410 let cfg = LaunchConfig {
20411 grid_dim: (out_f as u32, m as u32, 1),
20412 block_dim: (128, 1, 1),
20413 shared_mem_bytes: 0,
20414 };
20415 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20416 let __s_b = self.gpu.stream();
20417 let mut b = __s_b.launch_builder(&f);
20418 b.arg(bytes)
20419 .arg(aq)
20420 .arg(ad)
20421 .arg(&mut y)
20422 .arg(&inf)
20423 .arg(&outf)
20424 .arg(&mi)
20425 .arg(&rb);
20426 unsafe {
20427 b.launch(cfg)?;
20428 }
20429 Ok(Some((y, scale)))
20430 }
20431
20432 pub fn mmvq_fast_eligible(&self, w: &crate::model::GpuTensor, m: usize) -> bool {
20441 use crate::model::GpuTensor;
20442 if m != 1 || std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
20443 return false;
20444 }
20445 match w {
20446 GpuTensor::Quant { qtype, .. } => {
20447 *qtype != QT_F8_E4M3_BLK && self.mmvq_supports(*qtype)
20448 }
20449 _ => false,
20450 }
20451 }
20452
20453 pub fn matmul_q8_fast(
20457 &self,
20458 w: &crate::model::GpuTensor,
20459 aq: &CudaSlice<i8>,
20460 ad: &CudaSlice<f32>,
20461 m: usize,
20462 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20463 use crate::model::GpuTensor;
20464 if !self.mmvq_fast_eligible(w, m) {
20465 return Ok(None);
20466 }
20467 let GpuTensor::Quant {
20468 bytes,
20469 qtype,
20470 row_bytes,
20471 rp,
20472 rp4,
20473 scale,
20474 ..
20475 } = w
20476 else {
20477 return Ok(None);
20478 };
20479 let (bytes, rp) = match rp4 {
20480 Some(m4) => (m4, true),
20481 None => (bytes, *rp),
20482 };
20483 let y = self.qmatvec_mmvq(
20484 bytes,
20485 aq,
20486 ad,
20487 m,
20488 w.in_features(),
20489 w.out_features(),
20490 *qtype,
20491 *row_bytes,
20492 *scale,
20493 rp,
20494 )?;
20495 Ok(Some(y))
20496 }
20497
20498 pub fn matmul_q8_narrow_f32in(
20504 &self,
20505 w: &crate::model::GpuTensor,
20506 x: &CudaSlice<f32>,
20507 m: usize,
20508 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20509 use crate::model::GpuTensor;
20510 let GpuTensor::Quant {
20511 bytes,
20512 qtype,
20513 row_bytes,
20514 rp,
20515 rp4,
20516 scale,
20517 ..
20518 } = w
20519 else {
20520 return Ok(None);
20521 };
20522 if *qtype != QT_Q8_0
20523 || *rp
20524 || rp4.is_some()
20525 || *scale != 1.0
20526 || !self.mmvq_fast_eligible(w, m)
20527 {
20528 return Ok(None);
20529 }
20530 self.q8_narrow_f32in_raw(bytes, x, m, w.in_features(), w.out_features(), *row_bytes)
20531 }
20532
20533 pub fn q8_narrow_f32in_raw(
20536 &self,
20537 w: &CudaSlice<u8>,
20538 x: &CudaSlice<f32>,
20539 m: usize,
20540 in_f: usize,
20541 out_f: usize,
20542 row_bytes: usize,
20543 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20544 const ROWS_PER_BLOCK: u32 = 4; if in_f == 0 || !in_f.is_multiple_of(32) || in_f > 256 || m == 0 || m > 16 || out_f == 0 {
20546 return Ok(None);
20547 }
20548 debug_assert_eq!(row_bytes, in_f / 32 * 34);
20549 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20550 let f = self.func("qmatvec_q8_0_mmvq_f32in_narrow");
20551 let cfg = LaunchConfig {
20552 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
20553 block_dim: (32, ROWS_PER_BLOCK, 1),
20554 shared_mem_bytes: 0,
20555 };
20556 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20557 let stream = self.gpu.stream();
20558 let mut b = stream.launch_builder(&f);
20559 b.arg(w)
20560 .arg(x)
20561 .arg(&mut y)
20562 .arg(&inf)
20563 .arg(&outf)
20564 .arg(&mi)
20565 .arg(&rb);
20566 unsafe {
20567 b.launch(cfg)?;
20568 }
20569 Ok(Some(y))
20570 }
20571
20572 pub fn mmvq_supports(&self, qtype: i32) -> bool {
20573 if qtype == QT_F8_E4M3 {
20578 return true;
20579 }
20580 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
20581 return false;
20582 }
20583 matches!(
20584 qtype,
20585 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
20586 )
20587 }
20588
20589 #[allow(clippy::too_many_arguments)] pub fn qmatvec_mmvq(
20595 &self,
20596 bytes: &CudaSlice<u8>,
20597 aq: &CudaSlice<i8>,
20598 ad: &CudaSlice<f32>,
20599 m: usize,
20600 in_f: usize,
20601 out_f: usize,
20602 qtype: i32,
20603 row_bytes: usize,
20604 scale: f32,
20605 rp: bool,
20606 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20607 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
20609 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
20610 )?;
20611 Ok(y)
20612 }
20613
20614 #[allow(clippy::too_many_arguments)]
20616 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_mmvq_into(
20618 &self,
20619 bytes: &CudaSlice<u8>,
20620 aq: &CudaSlice<i8>,
20621 ad: &CudaSlice<f32>,
20622 m: usize,
20623 in_f: usize,
20624 out_f: usize,
20625 qtype: i32,
20626 row_bytes: usize,
20627 scale: f32,
20628 rp: bool,
20629 y: &mut CudaSlice<f32>,
20630 ) -> Result<(), Box<dyn std::error::Error>> {
20631 debug_assert!(y.len() >= m * out_f);
20632 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
20638 && rp
20639 && m == 1
20640 && out_f >= 64
20641 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
20642 && {
20643 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20644 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
20645 }
20646 {
20647 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
20648 let cfg = LaunchConfig {
20649 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
20650 block_dim: (32, 2, 1),
20651 shared_mem_bytes: 0,
20652 };
20653 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
20654 let __s_b = self.gpu.stream();
20655 let mut b = __s_b.launch_builder(&f);
20656 b.arg(bytes)
20657 .arg(aq)
20658 .arg(ad)
20659 .arg(&mut *y)
20660 .arg(&inf)
20661 .arg(&outf)
20662 .arg(&mi)
20663 .arg(&rb);
20664 unsafe {
20665 b.launch(cfg)?;
20666 }
20667 if scale != 1.0 {
20668 self.scale_inplace(y, scale, out_f)?;
20669 }
20670 return Ok(());
20671 }
20672 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
20681 2
20682 } else {
20683 1
20684 };
20685 if m == 1 && qtype == QT_Q4_0 {
20690 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
20691 mr = *Q40MR.get_or_init(|| {
20694 std::env::var("MEMRA_Q40_MR")
20695 .ok()
20696 .and_then(|v| v.parse().ok())
20697 .unwrap_or(1)
20698 });
20699 }
20700 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
20711 let q5_force = q5_mode.as_deref() == Some("2");
20712 let q5_il = qtype == QT_Q5_K
20715 && m == 1
20716 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
20717 if q5_il && !q5_force && out_f > 65536 {
20718 mr = 1;
20719 }
20720 if qtype == QT_Q4_0 && rp && mr != 1 {
20723 mr = 2;
20724 }
20725 if qtype == QT_Q8_0 && rp {
20729 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
20730 mr = *Q80MR.get_or_init(|| {
20731 std::env::var("MEMRA_Q80_MR")
20732 .ok()
20733 .and_then(|v| v.parse().ok())
20734 .unwrap_or(1)
20735 });
20736 }
20737 if qtype == QT_NVFP4
20749 && mr == 2
20750 && rp
20751 && m == 1
20752 && b200_matvec_arm_on()
20753 && (out_f as u32).div_ceil(ROWS_PER_BLOCK * 2)
20754 < b200_mr1_fill() * self.sm_count() as u32
20755 {
20756 mr = 1;
20757 }
20758 let nv_ilp = qtype == QT_NVFP4 && rp && nvfp4_row_ilp_on();
20761 if nv_ilp
20762 && NVFP4_ROW_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0
20763 {
20764 eprintln!(
20765 "[nvfp4-row-ilp] engaged: split-plane NVFP4 trunk matvec with four groups' loads \
20766 per lane ahead of the dp4a chains (MEMRA_NVFP4_ROW_ILP=1, mr={mr})"
20767 );
20768 }
20769 let name = match (qtype, mr, rp) {
20770 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
20771 (QT_NVFP4, 2, true) if nv_ilp => "qmatvec_nvfp4_mmvq_mr2_rp_ilp",
20772 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
20773 (QT_NVFP4, _, true) if nv_ilp => "qmatvec_nvfp4_mmvq_rp_ilp",
20774 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
20775 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
20776 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
20777 (QT_Q5_K, 2, _) => {
20778 if q5_il {
20779 "qmatvec_q5_K_mmvq_mr2_il"
20780 } else {
20781 "qmatvec_q5_K_mmvq_mr2"
20782 }
20783 }
20784 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
20785 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
20786 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
20787 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
20791 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
20792 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
20793 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
20794 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
20795 (QT_Q5_K, _, _) => {
20796 if q5_il {
20797 "qmatvec_q5_K_mmvq_il"
20798 } else {
20799 "qmatvec_q5_K_mmvq"
20800 }
20801 }
20802 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
20803 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
20804 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
20805 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
20806 };
20807 let f = self.func(name);
20808 let rows_per_block = ROWS_PER_BLOCK * mr;
20810 let cfg = LaunchConfig {
20811 grid_dim: (
20812 (out_f as u32 + rows_per_block - 1) / rows_per_block,
20813 m as u32,
20814 1,
20815 ),
20816 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
20819 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20820 let __s_b = self.gpu.stream();
20821 let mut b = __s_b.launch_builder(&f);
20822 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
20827 if Self::pdl_on()
20833 && Self::pdl_mmvq_on()
20834 && Self::pdl_nvfp4q8_on()
20835 && matches!(
20836 name,
20837 "qmatvec_nvfp4_mmvq_mr2_rp"
20838 | "qmatvec_nvfp4_mmvq_mr2_rp_ilp"
20839 | "qmatvec_nvfp4_mmvq_rp"
20840 | "qmatvec_nvfp4_mmvq_rp_ilp"
20841 )
20842 {
20843 use cudarc::driver::{DevicePtr, DevicePtrMut};
20844 let s = &self.gpu.stream();
20845 let (pw, _g0) = bytes.device_ptr(s);
20846 let (paq, _g1) = aq.device_ptr(s);
20847 let (pad, _g2) = ad.device_ptr(s);
20848 let (py, _g3) = y.device_ptr_mut(s);
20849 let mut ps = [
20850 &pw as *const _ as *mut std::ffi::c_void,
20851 &paq as *const _ as *mut _,
20852 &pad as *const _ as *mut _,
20853 &py as *const _ as *mut _,
20854 &inf as *const _ as *mut _,
20855 &outf as *const _ as *mut _,
20856 &mi as *const _ as *mut _,
20857 &rb as *const _ as *mut _,
20858 &scale as *const _ as *mut _,
20859 ];
20860 unsafe {
20861 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
20862 }
20863 return Ok(());
20864 }
20865 b.arg(bytes)
20866 .arg(aq)
20867 .arg(ad)
20868 .arg(&mut *y)
20869 .arg(&inf)
20870 .arg(&outf)
20871 .arg(&mi)
20872 .arg(&rb)
20873 .arg(&scale);
20874 unsafe {
20875 b.launch(cfg)?;
20876 }
20877 } else if Self::pdl_on()
20878 && Self::pdl_mmvq_on()
20879 && (matches!(
20880 name,
20881 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
20882 ) || (Self::pdl_nvfp4q8_on()
20883 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
20884 {
20885 {
20889 use cudarc::driver::{DevicePtr, DevicePtrMut};
20890 let s = &self.gpu.stream();
20891 let (pw, _g0) = bytes.device_ptr(s);
20892 let (paq, _g1) = aq.device_ptr(s);
20893 let (pad, _g2) = ad.device_ptr(s);
20894 let (py, _g3) = y.device_ptr_mut(s);
20895 let mut ps = [
20896 &pw as *const _ as *mut std::ffi::c_void,
20897 &paq as *const _ as *mut _,
20898 &pad as *const _ as *mut _,
20899 &py as *const _ as *mut _,
20900 &inf as *const _ as *mut _,
20901 &outf as *const _ as *mut _,
20902 &mi as *const _ as *mut _,
20903 &rb as *const _ as *mut _,
20904 ];
20905 unsafe {
20906 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
20907 }
20908 }
20909 if scale != 1.0 {
20910 self.scale_inplace(y, scale, m * out_f)?;
20911 }
20912 } else {
20913 b.arg(bytes)
20914 .arg(aq)
20915 .arg(ad)
20916 .arg(&mut *y)
20917 .arg(&inf)
20918 .arg(&outf)
20919 .arg(&mi)
20920 .arg(&rb);
20921 unsafe {
20922 b.launch(cfg)?;
20923 }
20924 if scale != 1.0 {
20925 self.scale_inplace(y, scale, m * out_f)?;
20926 }
20927 }
20928 Ok(())
20929 }
20930
20931 #[allow(clippy::too_many_arguments)] pub fn qmatvec_mmvq_raw(
20936 &self,
20937 bytes: &CudaSlice<u8>,
20938 x: &CudaSlice<f32>,
20939 m: usize,
20940 in_f: usize,
20941 out_f: usize,
20942 qtype: i32,
20943 row_bytes: usize,
20944 rp: bool,
20945 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20946 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20947 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
20948 }
20949
20950 #[allow(clippy::too_many_arguments)]
20959 pub fn qmatvec_nvfp4_rp_arm_raw(
20960 &self,
20961 bytes: &CudaSlice<u8>,
20962 aq: &CudaSlice<i8>,
20963 ad: &CudaSlice<f32>,
20964 m: usize,
20965 in_f: usize,
20966 out_f: usize,
20967 row_bytes: usize,
20968 yscale: f32,
20969 use_arm: bool,
20970 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20971 const ROWS_PER_BLOCK: u32 = 4; let mr: u32 = if use_arm { 1 } else { 2 };
20973 let name = if use_arm {
20974 "qmatvec_nvfp4_mmvq_rp"
20975 } else {
20976 "qmatvec_nvfp4_mmvq_mr2_rp"
20977 };
20978 let f = self.func(name);
20979 let rows_per_block = ROWS_PER_BLOCK * mr;
20980 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20981 let cfg = LaunchConfig {
20982 grid_dim: ((out_f as u32).div_ceil(rows_per_block), m as u32, 1),
20983 block_dim: (32, ROWS_PER_BLOCK, 1),
20984 shared_mem_bytes: 0,
20985 };
20986 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20987 let __s_b = self.gpu.stream();
20988 let mut b = __s_b.launch_builder(&f);
20989 b.arg(bytes)
20990 .arg(aq)
20991 .arg(ad)
20992 .arg(&mut y)
20993 .arg(&inf)
20994 .arg(&outf)
20995 .arg(&mi)
20996 .arg(&rb)
20997 .arg(&yscale);
20998 unsafe {
20999 b.launch(cfg)?;
21000 }
21001 Ok(y)
21002 }
21003
21004 #[allow(clippy::too_many_arguments)]
21007 pub fn qmatvec_nvfp4_rp_arm_raw_ilp(
21008 &self,
21009 bytes: &CudaSlice<u8>,
21010 aq: &CudaSlice<i8>,
21011 ad: &CudaSlice<f32>,
21012 m: usize,
21013 in_f: usize,
21014 out_f: usize,
21015 row_bytes: usize,
21016 yscale: f32,
21017 use_arm: bool,
21018 ilp: bool,
21019 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21020 const ROWS_PER_BLOCK: u32 = 4; let mr: u32 = if use_arm { 1 } else { 2 };
21022 let name = match (use_arm, ilp) {
21023 (true, false) => "qmatvec_nvfp4_mmvq_rp",
21024 (true, true) => "qmatvec_nvfp4_mmvq_rp_ilp",
21025 (false, false) => "qmatvec_nvfp4_mmvq_mr2_rp",
21026 (false, true) => "qmatvec_nvfp4_mmvq_mr2_rp_ilp",
21027 };
21028 let f = self.func(name);
21029 let rows_per_block = ROWS_PER_BLOCK * mr;
21030 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
21031 let cfg = LaunchConfig {
21032 grid_dim: ((out_f as u32).div_ceil(rows_per_block), m as u32, 1),
21033 block_dim: (32, ROWS_PER_BLOCK, 1),
21034 shared_mem_bytes: 0,
21035 };
21036 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
21037 let __s_b = self.gpu.stream();
21038 let mut b = __s_b.launch_builder(&f);
21039 b.arg(bytes)
21040 .arg(aq)
21041 .arg(ad)
21042 .arg(&mut y)
21043 .arg(&inf)
21044 .arg(&outf)
21045 .arg(&mi)
21046 .arg(&rb)
21047 .arg(&yscale);
21048 unsafe {
21049 b.launch(cfg)?;
21050 }
21051 Ok(y)
21052 }
21053
21054 #[allow(clippy::too_many_arguments)]
21059 pub fn qmatvec_nvfp4_fused2_rp_arm_raw(
21060 &self,
21061 w0: &CudaSlice<u8>,
21062 w1: &CudaSlice<u8>,
21063 aq: &CudaSlice<i8>,
21064 ad: &CudaSlice<f32>,
21065 in_f: usize,
21066 out0: usize,
21067 out1: usize,
21068 s0: f32,
21069 s1: f32,
21070 use_arm: bool,
21071 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21072 const ROWS_PER_BLOCK: u32 = 4; let rpw: u32 = if use_arm { 1 } else { 2 };
21074 let name = if use_arm {
21075 "qmatvec_nvfp4_mmvq_fused2_rp_g2"
21076 } else {
21077 "qmatvec_nvfp4_mmvq_fused2_rp"
21078 };
21079 let rows_pb = ROWS_PER_BLOCK * rpw;
21080 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
21081 let f = self.func(name);
21082 let mut y0 = self.alloc_uninit::<f32>(out0)?;
21083 let mut y1 = self.alloc_uninit::<f32>(out1)?;
21084 let cfg = LaunchConfig {
21085 grid_dim: (nb(out0) + nb(out1), 1, 1),
21086 block_dim: (32, ROWS_PER_BLOCK, 1),
21087 shared_mem_bytes: 0,
21088 };
21089 let (inf, oi0, oi1, mi) = (in_f as i32, out0 as i32, out1 as i32, 1i32);
21090 let __s_b = self.gpu.stream();
21091 let mut b = __s_b.launch_builder(&f);
21092 b.arg(w0)
21093 .arg(w1)
21094 .arg(aq)
21095 .arg(ad)
21096 .arg(&mut y0)
21097 .arg(&mut y1)
21098 .arg(&inf)
21099 .arg(&oi0)
21100 .arg(&oi1)
21101 .arg(&mi)
21102 .arg(&s0)
21103 .arg(&s1);
21104 unsafe {
21105 b.launch(cfg)?;
21106 }
21107 Ok((y0, y1))
21108 }
21109
21110 #[allow(clippy::too_many_arguments)]
21115 pub fn matvec_bf16_f32acc_x4_rows_arm_raw(
21116 &self,
21117 w: &CudaSlice<u8>,
21118 x: &CudaSlice<f32>,
21119 y: &mut CudaSlice<f32>,
21120 in_f: usize,
21121 out_f: usize,
21122 t: usize,
21123 use_arm: bool,
21124 ) -> Result<(), Box<dyn std::error::Error>> {
21125 let name = if use_arm {
21126 "matvec_bf16_f32acc_x4_rows_pf"
21127 } else {
21128 "matvec_bf16_f32acc_x4_rows"
21129 };
21130 let f = self.func(name);
21131 let cfg = LaunchConfig {
21132 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
21133 block_dim: (mmv_block(), 1, 1),
21134 shared_mem_bytes: 0,
21135 };
21136 let (ini, outi) = (in_f as i32, out_f as i32);
21137 let __s_b = self.gpu.stream();
21138 let mut b = __s_b.launch_builder(&f);
21139 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21140 unsafe {
21141 b.launch(cfg)?;
21142 }
21143 Ok(())
21144 }
21145
21146 pub fn gemv_v2_ksplit(&self, in_f: usize, out_f: usize, t: usize) -> usize {
21152 let blocks = out_f.div_ceil(GEMV_V2_ROWS) * t.max(1);
21153 let want = 2 * self.sm_count().max(1) as usize;
21154 if blocks >= want || blocks == 0 {
21155 return 1;
21156 }
21157 let per_step = mmv_block() as usize * 8;
21158 let max_ks = (in_f / per_step).max(1);
21159 want.div_ceil(blocks).clamp(1, max_ks)
21160 }
21161
21162 #[allow(clippy::too_many_arguments)]
21167 pub fn matvec_bf16_v2_raw(
21168 &self,
21169 w: &CudaSlice<u8>,
21170 x: &CudaSlice<f32>,
21171 y: &mut CudaSlice<f32>,
21172 in_f: usize,
21173 out_f: usize,
21174 t: usize,
21175 ksplit: usize,
21176 ) -> Result<(), Box<dyn std::error::Error>> {
21177 if t == 0 || !in_f.is_multiple_of(8) || x.len() < t * in_f || y.len() < t * out_f {
21178 return Err("matvec_bf16_v2 geometry".into());
21179 }
21180 let nb = mmv_block();
21181 let smem = (GEMV_V2_ROWS as u32) * nb * 4;
21182 let rows = out_f.div_ceil(GEMV_V2_ROWS) as u32;
21183 let (ini, outi) = (in_f as i32, out_f as i32);
21184 if ksplit <= 1 {
21185 let f = self.func("matvec_bf16_v2");
21186 let cfg = LaunchConfig {
21187 grid_dim: (rows, t as u32, 1),
21188 block_dim: (nb, 1, 1),
21189 shared_mem_bytes: smem,
21190 };
21191 let __s_b = self.gpu.stream();
21192 let mut b = __s_b.launch_builder(&f);
21193 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21194 unsafe {
21195 b.launch(cfg)?;
21196 }
21197 return Ok(());
21198 }
21199 let chunk = in_f.div_ceil(8).div_ceil(ksplit) * 8;
21204 let eff = in_f.div_ceil(chunk).max(1);
21205 let mut part = self.alloc_uninit::<f32>(eff * t * out_f)?;
21206 let ks = eff as i32;
21207 {
21208 let f = self.func("matvec_bf16_v2_sk");
21209 let cfg = LaunchConfig {
21210 grid_dim: (rows, t as u32, eff as u32),
21211 block_dim: (nb, 1, 1),
21212 shared_mem_bytes: smem,
21213 };
21214 let __s_b = self.gpu.stream();
21215 let mut b = __s_b.launch_builder(&f);
21216 b.arg(w).arg(x).arg(&mut part).arg(&ini).arg(&outi).arg(&ks);
21217 unsafe {
21218 b.launch(cfg)?;
21219 }
21220 }
21221 let f = self.func("matvec_bf16_v2_sk_combine");
21222 let n = (out_f * t) as u32;
21223 let cfg = LaunchConfig {
21224 grid_dim: (n.div_ceil(256), 1, 1),
21225 block_dim: (256, 1, 1),
21226 shared_mem_bytes: 0,
21227 };
21228 let ti = t as i32;
21229 let __s_b = self.gpu.stream();
21230 let mut b = __s_b.launch_builder(&f);
21231 b.arg(&part).arg(&mut *y).arg(&outi).arg(&ti).arg(&ks);
21232 unsafe {
21233 b.launch(cfg)?;
21234 }
21235 Ok(())
21236 }
21237
21238 pub fn gemv_v3_fits(&self) -> bool {
21242 gemv_v3_fits()
21243 }
21244
21245 pub fn matvec_bf16_v3_raw(
21253 &self,
21254 w: &CudaSlice<u8>,
21255 x: &CudaSlice<f32>,
21256 y: &mut CudaSlice<f32>,
21257 in_f: usize,
21258 out_f: usize,
21259 t: usize,
21260 ) -> Result<(), Box<dyn std::error::Error>> {
21261 if t == 0 || !in_f.is_multiple_of(8) || x.len() < t * in_f || y.len() < t * out_f {
21262 return Err("matvec_bf16_v3 geometry".into());
21263 }
21264 let nb = mmv_block();
21265 let smem = gemv_v3_smem_bytes(nb as usize);
21266 if smem > 48 * 1024 {
21267 return Err("matvec_bf16_v3 smem over the 48 KB default cap".into());
21268 }
21269 let f = self.func("matvec_bf16_v3");
21270 let cfg = LaunchConfig {
21271 grid_dim: (out_f.div_ceil(GEMV_V2_ROWS) as u32, t as u32, 1),
21272 block_dim: (nb, 1, 1),
21273 shared_mem_bytes: smem as u32,
21274 };
21275 let (ini, outi) = (in_f as i32, out_f as i32);
21276 let __s_b = self.gpu.stream();
21277 let mut b = __s_b.launch_builder(&f);
21278 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21279 unsafe {
21280 b.launch(cfg)?;
21281 }
21282 Ok(())
21283 }
21284
21285 #[allow(clippy::too_many_arguments)]
21292 pub fn qmatvec_q8_0_rp_v2_raw(
21293 &self,
21294 w: &CudaSlice<u8>,
21295 aq: &CudaSlice<i8>,
21296 ad: &CudaSlice<f32>,
21297 y: &mut CudaSlice<f32>,
21298 in_f: usize,
21299 out_f: usize,
21300 t: usize,
21301 ) -> Result<(), Box<dyn std::error::Error>> {
21302 self.qmatvec_q8_0_rp_v2_raw_arm(w, aq, ad, y, in_f, out_f, t, q8_row_ilp_on())
21303 }
21304
21305 #[allow(clippy::too_many_arguments)]
21308 pub fn qmatvec_q8_0_rp_v2_raw_arm(
21309 &self,
21310 w: &CudaSlice<u8>,
21311 aq: &CudaSlice<i8>,
21312 ad: &CudaSlice<f32>,
21313 y: &mut CudaSlice<f32>,
21314 in_f: usize,
21315 out_f: usize,
21316 t: usize,
21317 ilp: bool,
21318 ) -> Result<(), Box<dyn std::error::Error>> {
21319 if t == 0 || !in_f.is_multiple_of(32) || y.len() < t * out_f {
21320 return Err("qmatvec_q8_0_rp_v2 geometry".into());
21321 }
21322 let smem = Self::q8_v2_smem_bytes(in_f);
21323 if smem > 48 * 1024 {
21324 return Err("qmatvec_q8_0_rp_v2 smem over the 48 KB default cap".into());
21325 }
21326 if ilp {
21327 q8_row_ilp_note("qmatvec_q8_0_mmvq_rp_v2");
21328 }
21329 let f = self.func(if ilp {
21330 "qmatvec_q8_0_mmvq_rp_v2_ilp"
21331 } else {
21332 "qmatvec_q8_0_mmvq_rp_v2"
21333 });
21334 let cfg = LaunchConfig {
21335 grid_dim: ((out_f as u32).div_ceil(Q8_V2_ROWS), t as u32, 1),
21336 block_dim: (32, Q8_V2_ROWS, 1),
21337 shared_mem_bytes: smem as u32,
21338 };
21339 let (ini, outi, mi, rb) = (in_f as i32, out_f as i32, t as i32, 0i64);
21340 let __s_b = self.gpu.stream();
21341 let mut b = __s_b.launch_builder(&f);
21342 b.arg(w)
21343 .arg(aq)
21344 .arg(ad)
21345 .arg(&mut *y)
21346 .arg(&ini)
21347 .arg(&outi)
21348 .arg(&mi)
21349 .arg(&rb);
21350 unsafe {
21351 b.launch(cfg)?;
21352 }
21353 Ok(())
21354 }
21355
21356 #[allow(clippy::too_many_arguments)]
21362 pub fn qmatvec_q8_0_rows_tw_v2_raw(
21363 &self,
21364 w: &CudaSlice<u8>,
21365 aq: &CudaSlice<i8>,
21366 ad: &CudaSlice<f32>,
21367 y: &mut CudaSlice<f32>,
21368 in_f: usize,
21369 out_f: usize,
21370 t: usize,
21371 ) -> Result<(), Box<dyn std::error::Error>> {
21372 if t == 0 || t > 8 || !in_f.is_multiple_of(32) || y.len() < t * out_f {
21373 return Err("qmatvec_q8_0_rows_tw_v2 geometry".into());
21374 }
21375 let f = self.func("qmatvec_q8_0_rows_tw_v2");
21376 let cfg = LaunchConfig {
21377 grid_dim: ((out_f as u32).div_ceil(Q8_V2_ROWS), 1, 1),
21378 block_dim: (32, Q8_V2_ROWS, 1),
21379 shared_mem_bytes: 0,
21380 };
21381 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
21382 let __s_b = self.gpu.stream();
21383 let mut b = __s_b.launch_builder(&f);
21384 b.arg(w)
21385 .arg(aq)
21386 .arg(ad)
21387 .arg(&mut *y)
21388 .arg(&ini)
21389 .arg(&outi)
21390 .arg(&ti);
21391 unsafe {
21392 b.launch(cfg)?;
21393 }
21394 Ok(())
21395 }
21396
21397 #[allow(clippy::too_many_arguments)]
21402 pub fn qmatvec_q8_0_rows_tw_arm_raw(
21403 &self,
21404 w: &CudaSlice<u8>,
21405 aq: &CudaSlice<i8>,
21406 ad: &CudaSlice<f32>,
21407 y: &mut CudaSlice<f32>,
21408 in_f: usize,
21409 out_f: usize,
21410 t: usize,
21411 use_v2: bool,
21412 ) -> Result<(), Box<dyn std::error::Error>> {
21413 if use_v2 {
21414 return self.qmatvec_q8_0_rows_tw_v2_raw(w, aq, ad, y, in_f, out_f, t);
21415 }
21416 let f = self.func("qmatvec_q8_0_rows_tw");
21417 let cfg = LaunchConfig {
21418 grid_dim: ((out_f as u32).div_ceil(4), 1, 1),
21419 block_dim: (32, 4, 1),
21420 shared_mem_bytes: 0,
21421 };
21422 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
21423 let __s_b = self.gpu.stream();
21424 let mut b = __s_b.launch_builder(&f);
21425 b.arg(w)
21426 .arg(aq)
21427 .arg(ad)
21428 .arg(&mut *y)
21429 .arg(&ini)
21430 .arg(&outi)
21431 .arg(&ti);
21432 unsafe {
21433 b.launch(cfg)?;
21434 }
21435 Ok(())
21436 }
21437
21438 #[allow(clippy::too_many_arguments)]
21453 #[allow(clippy::map_entry)] pub fn kda_proj_fused6_q8rp_raw(
21455 &self,
21456 bq: &CudaSlice<u8>,
21457 bk: &CudaSlice<u8>,
21458 bv: &CudaSlice<u8>,
21459 wfa: &CudaSlice<f32>,
21460 wga: &CudaSlice<f32>,
21461 wb: &CudaSlice<f32>,
21462 x: &CudaSlice<f32>,
21463 outs: &mut [CudaSlice<f32>; 6],
21464 in_f: usize,
21465 dims: [usize; 6],
21466 t: usize,
21467 ) -> Result<(), Box<dyn std::error::Error>> {
21468 self.kda_proj_fused6_q8rp_raw_arm(
21469 bq,
21470 bk,
21471 bv,
21472 wfa,
21473 wga,
21474 wb,
21475 x,
21476 outs,
21477 in_f,
21478 dims,
21479 t,
21480 q8_row_ilp_on(),
21481 )
21482 }
21483
21484 #[allow(clippy::too_many_arguments)]
21487 pub fn kda_proj_fused6_q8rp_raw_arm(
21488 &self,
21489 bq: &CudaSlice<u8>,
21490 bk: &CudaSlice<u8>,
21491 bv: &CudaSlice<u8>,
21492 wfa: &CudaSlice<f32>,
21493 wga: &CudaSlice<f32>,
21494 wb: &CudaSlice<f32>,
21495 x: &CudaSlice<f32>,
21496 outs: &mut [CudaSlice<f32>; 6],
21497 in_f: usize,
21498 dims: [usize; 6],
21499 t: usize,
21500 ilp: bool,
21501 ) -> Result<(), Box<dyn std::error::Error>> {
21502 self.kda_proj_fused6_q8rp_raw_pre(
21503 bq, bk, bv, wfa, wga, wb, x, outs, in_f, dims, t, ilp, None,
21504 )
21505 }
21506
21507 #[allow(clippy::too_many_arguments)]
21514 pub fn kda_proj_fused6_q8rp_raw_pre(
21515 &self,
21516 bq: &CudaSlice<u8>,
21517 bk: &CudaSlice<u8>,
21518 bv: &CudaSlice<u8>,
21519 wfa: &CudaSlice<f32>,
21520 wga: &CudaSlice<f32>,
21521 wb: &CudaSlice<f32>,
21522 x: &CudaSlice<f32>,
21523 outs: &mut [CudaSlice<f32>; 6],
21524 in_f: usize,
21525 dims: [usize; 6],
21526 t: usize,
21527 ilp: bool,
21528 pre_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
21529 ) -> Result<(), Box<dyn std::error::Error>> {
21530 use cudarc::driver::DevicePtr;
21531 if t == 0 || t > 32 || !in_f.is_multiple_of(32) || x.len() < t * in_f {
21532 return Err("kda_proj_fused6_q8rp geometry".into());
21533 }
21534 let smem = Self::q8_v2_smem_bytes(in_f);
21535 if smem > 48 * 1024 {
21536 return Err("kda_proj_fused6_q8rp smem over the 48 KB default cap".into());
21537 }
21538 for (i, (o, want)) in outs.iter().zip(dims).enumerate() {
21539 if o.len() < t * want {
21540 return Err(format!("kda_proj_fused6_q8rp: output {i} too small").into());
21541 }
21542 }
21543 let nblk = in_f / 32;
21544 let keys: Vec<(u64, u32, u32)> = {
21546 let s = self.gpu.stream();
21547 [(bq, dims[0]), (bk, dims[1]), (bv, dims[2])]
21548 .iter()
21549 .map(|(d, out)| {
21550 let (p, _g) = d.device_ptr(&s);
21551 (p, in_f as u32, *out as u32)
21552 })
21553 .collect()
21554 };
21555 {
21556 let mut mirrors = self
21557 .w8_mirrors
21558 .lock()
21559 .map_err(|_| "w8 mirror map is poisoned")?;
21560 for ((d, out), key) in [(bq, dims[0]), (bk, dims[1]), (bv, dims[2])]
21561 .iter()
21562 .zip(&keys)
21563 {
21564 if !mirrors.contains_key(key) {
21565 let mut interleaved = self.alloc_u8_uninit(out * Self::q8_0_row_bytes(in_f))?;
21566 self.encode_q8_0_from_bf16(d, &mut interleaved, in_f, *out)?;
21567 let planar = self.build_q8_rp4_raw(&interleaved, in_f, *out)?;
21568 mirrors.insert(*key, planar);
21569 }
21570 }
21571 }
21572 let akey = in_f * 64 + t.min(32);
21575 if pre_q8.is_none() {
21576 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21577 if let std::collections::hash_map::Entry::Vacant(slot) = act.entry(akey) {
21578 let aq = self.alloc_i8_uninit(32 * in_f)?;
21579 let ad = self.alloc_uninit::<f32>(32 * nblk)?;
21580 slot.insert((aq, ad));
21581 }
21582 let (aq, ad) = act.get_mut(&akey).expect("just inserted");
21583 self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
21584 }
21585 let mirrors = self
21586 .w8_mirrors
21587 .lock()
21588 .map_err(|_| "w8 mirror map is poisoned")?;
21589 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21590 let (aq, ad): (&CudaSlice<i8>, &CudaSlice<f32>) = match pre_q8 {
21591 Some((aq, ad)) => {
21592 if aq.len() < t * in_f || ad.len() < t * nblk {
21593 return Err("kda_proj_fused6_q8rp: pre-quantized activation too small".into());
21594 }
21595 (aq, ad)
21596 }
21597 None => {
21598 let (aq, ad) = act.get(&akey).expect("built above");
21599 (aq, ad)
21600 }
21601 };
21602 let m0 = mirrors.get(&keys[0]).expect("built above");
21603 let m1 = mirrors.get(&keys[1]).expect("built above");
21604 let m2 = mirrors.get(&keys[2]).expect("built above");
21605 let blocks: usize = dims.iter().map(|d| d.div_ceil(Q8_V2_ROWS as usize)).sum();
21606 if ilp {
21607 q8_row_ilp_note("qmatvec_kda6_q8f32_rp_v2");
21608 }
21609 let f = self.func(if ilp {
21610 "qmatvec_kda6_q8f32_rp_v2_ilp"
21611 } else {
21612 "qmatvec_kda6_q8f32_rp_v2"
21613 });
21614 let cfg = LaunchConfig {
21615 grid_dim: (blocks as u32, t as u32, 1),
21616 block_dim: (32, Q8_V2_ROWS, 1),
21617 shared_mem_bytes: smem as u32,
21618 };
21619 let inf = in_f as i32;
21620 let d = dims.map(|v| v as i32);
21621 let mi = t as i32;
21622 let [o0, o1, o2, o3, o4, o5] = outs;
21623 let stream = self.gpu.stream();
21624 let mut b = stream.launch_builder(&f);
21625 b.arg(m0)
21626 .arg(m1)
21627 .arg(m2)
21628 .arg(wfa)
21629 .arg(wga)
21630 .arg(wb)
21631 .arg(aq)
21632 .arg(ad)
21633 .arg(x)
21634 .arg(&mut *o0)
21635 .arg(&mut *o1)
21636 .arg(&mut *o2)
21637 .arg(&mut *o3)
21638 .arg(&mut *o4)
21639 .arg(&mut *o5)
21640 .arg(&inf)
21641 .arg(&d[0])
21642 .arg(&d[1])
21643 .arg(&d[2])
21644 .arg(&d[3])
21645 .arg(&d[4])
21646 .arg(&d[5])
21647 .arg(&mi);
21648 unsafe {
21649 b.launch(cfg)?;
21650 }
21651 Ok(())
21652 }
21653
21654 #[allow(clippy::too_many_arguments)]
21659 pub fn moe_gate_up_preclamp8_q8_v2(
21660 &self,
21661 gp: WPtr8,
21662 up: WPtr8,
21663 aq: &CudaSlice<i8>,
21664 ad: &CudaSlice<f32>,
21665 gs: F32x8,
21666 us: F32x8,
21667 limit: f32,
21668 in_f: usize,
21669 n_ff: usize,
21670 n_used: usize,
21671 qt_g: i32,
21672 qt_u: i32,
21673 rb_g: usize,
21674 rb_u: usize,
21675 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21676 debug_assert!(
21677 limit > 1e-6,
21678 "moe_gate_up_preclamp8_q8_v2 needs a live limit; use moe_gate_up_silu8_q8"
21679 );
21680 const ROWS: u32 = 8;
21681 let f = self.func("moe_gate_up_preclamp8_q8_v2");
21682 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
21683 let cfg = LaunchConfig {
21684 grid_dim: ((n_ff as u32).div_ceil(ROWS), n_used as u32, 1),
21685 block_dim: (32, ROWS, 1),
21686 shared_mem_bytes: 0,
21687 };
21688 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
21689 let __s_b = self.gpu.stream();
21690 let mut b = __s_b.launch_builder(&f);
21691 b.arg(&gp)
21692 .arg(&up)
21693 .arg(aq)
21694 .arg(ad)
21695 .arg(&gs)
21696 .arg(&us)
21697 .arg(&limit)
21698 .arg(&mut act)
21699 .arg(&inf)
21700 .arg(&nff)
21701 .arg(&qt_g)
21702 .arg(&qt_u)
21703 .arg(&rbg)
21704 .arg(&rbu);
21705 unsafe {
21706 b.launch(cfg)?;
21707 }
21708 Ok(act)
21709 }
21710
21711 #[allow(clippy::too_many_arguments)]
21717 pub fn moe_down8_fma_q8_v2(
21718 &self,
21719 dp: WPtr8,
21720 w: F32x8,
21721 aq2: &CudaSlice<i8>,
21722 ad2: &CudaSlice<f32>,
21723 dst: &mut cudarc::driver::CudaViewMut<f32>,
21724 in_f: usize,
21725 out_f: usize,
21726 n_used: usize,
21727 qt: i32,
21728 rb: usize,
21729 ) -> Result<(), Box<dyn std::error::Error>> {
21730 let f = self.func("moe_down8_fma_q8_v2");
21731 let cfg = LaunchConfig {
21732 grid_dim: (out_f as u32, 1, 1),
21733 block_dim: (32, 8, 1),
21734 shared_mem_bytes: 0,
21735 };
21736 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
21737 let __s_b = self.gpu.stream();
21738 let mut b = __s_b.launch_builder(&f);
21739 b.arg(&dp)
21740 .arg(&w)
21741 .arg(aq2)
21742 .arg(ad2)
21743 .arg(dst)
21744 .arg(&inf)
21745 .arg(&outf)
21746 .arg(&nu)
21747 .arg(&qt)
21748 .arg(&rbi);
21749 unsafe {
21750 b.launch(cfg)?;
21751 }
21752 Ok(())
21753 }
21754
21755 pub fn batched_supports(&self, qtype: i32) -> bool {
21759 matches!(
21760 qtype,
21761 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
21762 )
21763 }
21764
21765 pub fn iq_fast_enabled() -> bool {
21773 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21774 *ON.get_or_init(|| {
21775 std::env::var("MEMRA_IQ_FAST")
21776 .map(|v| v != "0")
21777 .unwrap_or(true)
21778 })
21779 }
21780
21781 pub fn b8_enabled() -> bool {
21784 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21785 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
21786 }
21787
21788 pub fn batched_mcols(m: usize) -> usize {
21790 if m == 2 {
21791 2
21792 } else if m <= 4 {
21793 4
21794 } else if m <= 8 {
21795 8
21796 } else {
21797 16
21798 }
21799 }
21800
21801 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
21806 Some(match (qtype, mcols) {
21807 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
21808 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
21809 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
21810 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
21816 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
21817 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
21818 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
21819 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
21822 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
21823 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
21824 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
21825 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
21828 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
21829 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
21830 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
21831 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
21832 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
21833 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
21834 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
21835 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
21839 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
21840 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
21841 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
21842 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
21846 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
21847 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
21848 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
21849 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
21850 _ => return None,
21851 })
21852 }
21853
21854 pub fn sm_count(&self) -> i32 {
21888 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
21889 *SMS.get_or_init(|| {
21890 use cudarc::driver::sys::CUdevice_attribute_enum as A;
21891 self.gpu
21892 .ctx
21893 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
21894 .unwrap_or(82)
21895 })
21896 }
21897
21898 #[allow(clippy::too_many_arguments)]
21899 #[allow(clippy::if_same_then_else)] pub fn batched_variant(
21902 &self,
21903 _m: usize,
21904 in_f: usize,
21905 out_f: usize,
21906 qtype: i32,
21907 mcols: usize,
21908 rp: bool,
21909 ) -> &'static str {
21910 if qtype == QT_Q8_0 {
21915 return if rp { "rp" } else { "base" };
21916 }
21917 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21922 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
21923 let sc_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(256) && (in_f / 64 <= 272);
21924 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
21925 let sms = *SMS.get_or_init(|| {
21926 use cudarc::driver::sys::CUdevice_attribute_enum as A;
21927 self.gpu
21928 .ctx
21929 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
21930 .unwrap_or(82)
21931 });
21932 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
21952 let variant: &'static str = if qtype == QT_Q4_0 {
21953 let v = if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
21958 "r2"
21959 } else {
21960 "base"
21961 };
21962 if rp {
21964 match v {
21965 "r2" => "r2_rp",
21966 _ => "rp",
21967 }
21968 } else {
21969 v
21970 }
21971 } else if qtype != QT_NVFP4 && !kq_r2 {
21972 "base"
21973 } else if kq_r2 && rp {
21974 "rp"
21978 } else if kq_r2 {
21979 #[allow(clippy::manual_div_ceil)]
21982 let blocks = (out_f + 7) / 8;
21984 let waves = blocks as f64 / (7 * sms as usize) as f64;
21985 let filled = blocks >= 4 * sms as usize;
21986 let use_r2 = if qtype == QT_Q4_K {
21987 filled
21988 } else {
21989 waves >= 2.0
21990 };
21991 if use_r2 { "r2" } else { "base" }
21992 } else if mcols == 8 {
21993 if rp {
22004 if sc_ok { "rpsc" } else { "rpr2w8" }
22005 } else {
22006 "r2w8"
22007 }
22008 } else if mcols >= 4 {
22009 #[allow(clippy::manual_div_ceil)]
22013 let blocks = (out_f + 7) / 8;
22015 let r7 = 7 * sms as usize;
22016 let r8 = 8 * sms as usize;
22017 let waves = blocks as f64 / r7 as f64;
22018 let filled = blocks >= 4 * sms as usize;
22019 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
22023 if rp { "rpr2w8" } else { "r2w8" }
22027 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
22028 if rp { "rpr2" } else { "r2" }
22031 } else {
22032 if rp { "rp" } else { "pf" }
22036 }
22037 } else if in_f >= 6144 {
22038 if rp { "rpr2" } else { "r2" }
22042 } else if rp {
22043 #[allow(clippy::manual_div_ceil)]
22048 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
22050 if sc_ok && (0.9..=1.1).contains(&waves) {
22051 "rpsc"
22052 } else {
22053 "rp"
22054 }
22055 } else {
22056 "base"
22057 };
22058 variant
22059 }
22060
22061 #[allow(clippy::too_many_arguments)]
22062 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_mmvq_batched(
22065 &self,
22066 bytes: &CudaSlice<u8>,
22067 aq: &CudaSlice<i8>,
22068 ad: &CudaSlice<f32>,
22069 m: usize,
22070 in_f: usize,
22071 out_f: usize,
22072 qtype: i32,
22073 row_bytes: usize,
22074 mcols: usize,
22075 scale: f32,
22076 rp: bool,
22077 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22078 const ROWS_PER_BLOCK: u32 = 4;
22079 let forced: Option<&'static str> = {
22084 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
22085 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
22086 .as_deref()
22087 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
22088 };
22089 let variant = match forced {
22090 Some(v) if !rp || v.contains("rp") => v,
22091 _ => self.batched_variant(m, in_f, out_f, qtype, mcols, rp),
22092 };
22093 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
22094 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
22095 })?;
22096 let variant = if mcols == 16 {
22100 if rp { "rp" } else { "base" }
22101 } else {
22102 variant
22103 };
22104 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22111 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
22112 if b567
22113 && qtype == QT_NVFP4
22114 && rp
22115 && mcols == 8
22116 && (5..=7).contains(&m)
22117 && matches!(variant, "rpsc" | "rpr2w8")
22118 {
22119 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
22120 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
22122 let cfg = LaunchConfig {
22123 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
22124 block_dim: (32, ROWS_PER_BLOCK, 1),
22125 shared_mem_bytes: 0,
22126 };
22127 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22128 let __s_b = self.gpu.stream();
22129 let mut b = __s_b.launch_builder(&f);
22130 b.arg(bytes)
22131 .arg(aq)
22132 .arg(ad)
22133 .arg(&mut y)
22134 .arg(&inf)
22135 .arg(&outf)
22136 .arg(&mi)
22137 .arg(&rb);
22138 unsafe {
22139 b.launch(cfg)?;
22140 }
22141 if scale != 1.0 {
22142 self.scale_inplace(&mut y, scale, m * out_f)?;
22143 }
22144 return Ok(y);
22145 }
22146 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
22147 "base" => (base_name.into(), ROWS_PER_BLOCK),
22148 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
22149 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
22150 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
22154 debug_assert!(
22155 !rp || name.contains("_rp"),
22156 "rp weight dispatched to a GGUF-layout kernel"
22157 );
22158 let f = self.func(&name);
22159 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
22160 let smem = if name.contains("_r2sm_rp") {
22162 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
22163 } else {
22164 0
22165 };
22166 let cfg = LaunchConfig {
22167 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
22168 block_dim: (32, ROWS_PER_BLOCK, 1),
22169 shared_mem_bytes: smem,
22170 };
22171 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22172 let __s_b = self.gpu.stream();
22173 let mut b = __s_b.launch_builder(&f);
22174 b.arg(bytes)
22175 .arg(aq)
22176 .arg(ad)
22177 .arg(&mut y)
22178 .arg(&inf)
22179 .arg(&outf)
22180 .arg(&mi)
22181 .arg(&rb);
22182 unsafe {
22183 b.launch(cfg)?;
22184 }
22185 if scale != 1.0 {
22186 self.scale_inplace(&mut y, scale, m * out_f)?;
22187 }
22188 Ok(y)
22189 }
22190
22191 #[allow(clippy::too_many_arguments)] pub fn qmatvec_batched_raw(
22196 &self,
22197 bytes: &CudaSlice<u8>,
22198 x: &CudaSlice<f32>,
22199 m: usize,
22200 in_f: usize,
22201 out_f: usize,
22202 qtype: i32,
22203 row_bytes: usize,
22204 mcols: usize,
22205 rp: bool,
22206 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22207 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
22208 self.qmatvec_mmvq_batched(
22209 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
22210 )
22211 }
22212
22213 #[allow(clippy::too_many_arguments)] pub fn qmatvec_nvfp4_batched_raw(
22216 &self,
22217 bytes: &CudaSlice<u8>,
22218 x: &CudaSlice<f32>,
22219 m: usize,
22220 in_f: usize,
22221 out_f: usize,
22222 row_bytes: usize,
22223 mcols: usize,
22224 rp: bool,
22225 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22226 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
22227 }
22228
22229 fn try_fp4_gemm(
22233 &self,
22234 w: &crate::model::GpuTensor,
22235 x: &CudaSlice<f32>,
22236 m: usize,
22237 in_f: usize,
22238 out_f: usize,
22239 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22240 use crate::model::GpuTensor;
22241 if cfg!(memra_portable_cuda) {
22242 return Ok(None);
22243 }
22244 if std::env::var("MEMRA_FP4").is_ok() {
22253 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
22254 assert!(
22255 konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a"),
22256 "MEMRA_FP4 forces the native mxf4 block-scale GEMM (qmatvec_gemm_nvfp4_fp4), \
22257 which only the sm_120a fatbin contains — this is an sm_{} build. Unset \
22258 MEMRA_FP4; the W4A8 int8 path is the correct default for NVFP4 weights.",
22259 env!("MEMRA_BUILT_CUDA_ARCH")
22260 );
22261 }
22262 if std::env::var("MEMRA_FP4").is_err() {
22263 return Ok(None);
22264 }
22265 #[cfg(memra_cutlass)]
22274 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
22275 if let GpuTensor::Quant {
22276 bytes,
22277 qtype,
22278 scale,
22279 row_bytes,
22280 cutlass,
22281 ..
22282 } = w
22283 {
22284 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
22285 if let Some(cw) = cutlass {
22286 let y = self.cutlass_fp4_gemm(
22288 &cw.b_packed,
22289 &cw.sfb_swizzled,
22290 x,
22291 *scale,
22292 m,
22293 out_f,
22294 in_f,
22295 )?;
22296 return Ok(Some(y));
22297 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
22298 let (b_packed, sfb_sw) =
22303 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
22304 let y =
22305 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
22306 return Ok(Some(y));
22307 }
22308 }
22309 }
22310 }
22311 if let GpuTensor::Quant {
22312 bytes,
22313 qtype,
22314 row_bytes,
22315 scale,
22316 rp,
22317 ..
22318 } = w
22319 {
22320 if *qtype == QT_NVFP4 && in_f.is_multiple_of(64) && !*rp {
22323 let y =
22324 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
22325 return Ok(Some(y));
22326 }
22327 }
22328 Ok(None)
22329 }
22330
22331 #[allow(clippy::too_many_arguments)] pub fn rms_norm_f16out(
22336 &self,
22337 x: &CudaSlice<f32>,
22338 w: &CudaSlice<f32>,
22339 dst: &mut CudaSlice<f32>,
22340 dst16: &mut CudaSlice<u8>,
22341 ncols: usize,
22342 nrows: usize,
22343 eps: f32,
22344 ) -> Result<(), Box<dyn std::error::Error>> {
22345 let f = self.func("rms_norm_f16out_f32");
22346 let cfg = LaunchConfig {
22347 grid_dim: (nrows as u32, 1, 1),
22348 block_dim: (rms_block(), 1, 1),
22349 shared_mem_bytes: 0,
22350 };
22351 let (nc, e) = (ncols as i32, eps);
22352 let __s_b = self.gpu.stream();
22353 let mut b = __s_b.launch_builder(&f);
22354 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
22355 unsafe {
22356 b.launch(cfg)?;
22357 }
22358 Ok(())
22359 }
22360
22361 #[allow(clippy::too_many_arguments)]
22364 pub fn add_rms_norm_f16out(
22365 &self,
22366 a: &CudaSlice<f32>,
22367 b: &CudaSlice<f32>,
22368 w: &CudaSlice<f32>,
22369 res: &mut CudaSlice<f32>,
22370 dst: &mut CudaSlice<f32>,
22371 dst16: &mut CudaSlice<u8>,
22372 ncols: usize,
22373 nrows: usize,
22374 eps: f32,
22375 ) -> Result<(), Box<dyn std::error::Error>> {
22376 let f = self.func("add_rms_norm_f16out_f32");
22377 let cfg = LaunchConfig {
22378 grid_dim: (nrows as u32, 1, 1),
22379 block_dim: (rms_block(), 1, 1),
22380 shared_mem_bytes: 0,
22381 };
22382 let (nc, e) = (ncols as i32, eps);
22383 let __s_lb = self.gpu.stream();
22384 let mut lb = __s_lb.launch_builder(&f);
22385 lb.arg(a)
22386 .arg(b)
22387 .arg(w)
22388 .arg(res)
22389 .arg(dst)
22390 .arg(dst16)
22391 .arg(&nc)
22392 .arg(&e);
22393 unsafe {
22394 lb.launch(cfg)?;
22395 }
22396 Ok(())
22397 }
22398
22399 pub fn matmul_group_xh(
22402 &self,
22403 ws: &[&crate::model::GpuTensor],
22404 x: &CudaSlice<f32>,
22405 xh: &CudaSlice<u8>,
22406 m: usize,
22407 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22408 let mut out = Vec::with_capacity(ws.len());
22409 let in_f = ws[0].in_features();
22410 for w in ws {
22411 if w.in_features() == in_f
22412 && m >= 16
22413 && !self.verify_exact_on()
22414 && let Some(y) = self.try_f16_gemm_pre(w, xh, m)?
22415 {
22416 out.push(y);
22417 continue;
22418 }
22419 out.push(self.matmul(w, x, m)?);
22420 }
22421 Ok(out)
22422 }
22423
22424 pub fn gdn_pad_mask(
22427 &self,
22428 beta: &mut CudaSlice<f32>,
22429 g_log: &mut CudaSlice<f32>,
22430 len_d: &CudaSlice<i32>,
22431 h: usize,
22432 t: usize,
22433 ) -> Result<(), Box<dyn std::error::Error>> {
22434 let f = self.func("gdn_pad_mask_f32");
22435 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
22436 let (hi, ti) = (h as i32, t as i32);
22437 let __s_b = self.gpu.stream();
22438 let mut b = __s_b.launch_builder(&f);
22439 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
22440 unsafe {
22441 b.launch(cfg)?;
22442 }
22443 Ok(())
22444 }
22445
22446 pub fn row_gather_dev(
22449 &self,
22450 src: &CudaSlice<f32>,
22451 dst: &mut CudaSlice<f32>,
22452 len_d: &CudaSlice<i32>,
22453 ncols: usize,
22454 ) -> Result<(), Box<dyn std::error::Error>> {
22455 let f = self.func("row_gather_dev_f32");
22456 let cfg = LaunchConfig::for_num_elems(ncols as u32);
22457 let nc = ncols as i32;
22458 let __s_b = self.gpu.stream();
22459 let mut b = __s_b.launch_builder(&f);
22460 b.arg(src).arg(dst).arg(len_d).arg(&nc);
22461 unsafe {
22462 b.launch(cfg)?;
22463 }
22464 Ok(())
22465 }
22466
22467 pub fn matmul_group(
22474 &self,
22475 ws: &[&crate::model::GpuTensor],
22476 x: &CudaSlice<f32>,
22477 m: usize,
22478 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22479 use crate::model::GpuTensor;
22480 let mut out = Vec::with_capacity(ws.len());
22481 let any_mirror = ws
22482 .iter()
22483 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
22484 if m >= 16 && any_mirror && !self.verify_exact_on() {
22485 let in_f = ws[0].in_features();
22486 let xh = self.f16_act(x, m * in_f)?;
22487 for w in ws {
22488 if w.in_features() == in_f
22489 && let Some(y) = self.try_f16_gemm_pre(w, &xh, m)?
22490 {
22491 out.push(y);
22492 continue;
22493 }
22494 out.push(self.matmul(w, x, m)?);
22495 }
22496 return Ok(out);
22497 }
22498 for w in ws {
22499 out.push(self.matmul(w, x, m)?);
22500 }
22501 Ok(out)
22502 }
22503
22504 pub fn matmul_group_multi(
22511 &self,
22512 ws: &[&crate::model::GpuTensor],
22513 xs: &[&CudaSlice<f32>],
22514 ms: &[usize],
22515 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
22516 assert_eq!(xs.len(), ms.len());
22517 let in_f = ws[0].in_features();
22518 let total: usize = ms.iter().sum();
22519 let mut xcat = self.uninit(total * in_f)?;
22520 let mut off = 0usize;
22521 for (x, &m) in xs.iter().zip(ms) {
22522 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
22523 off += m;
22524 }
22525 let ys = self.matmul_group(ws, &xcat, total)?;
22526 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
22527 for (w, y) in ws.iter().zip(ys) {
22528 let out_f = w.out_features();
22529 let mut off = 0usize;
22530 for (s, &m) in ms.iter().enumerate() {
22531 let mut ys_s = self.uninit(m * out_f)?;
22532 let src = y.slice(off * out_f..(off + m) * out_f);
22533 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
22534 out[s].push(ys_s);
22535 off += m;
22536 }
22537 }
22538 Ok(out)
22539 }
22540
22541 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
22551 use crate::model::GpuTensor;
22552 if !legacy_quant_gemm_allowed(
22553 cfg!(memra_portable_cuda),
22554 cfg!(memra_hopper_mma),
22555 std::env::var_os("MEMRA_NO_GEMM").is_some(),
22556 ) {
22557 return false;
22558 }
22559 match w {
22560 GpuTensor::Quant { qtype, .. } => {
22561 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
22562 || (*qtype == QT_NVFP4 && w.in_features().is_multiple_of(64))
22563 }
22564 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
22565 }
22566 }
22567
22568 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_gemm(
22576 &self,
22577 w: &crate::model::GpuTensor,
22578 aq: &CudaSlice<i8>,
22579 ad: &CudaSlice<f32>,
22580 m: usize,
22581 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22582 use crate::model::GpuTensor;
22583 let in_f = w.in_features();
22584 let out_f = w.out_features();
22585 let (bytes, qtype, row_bytes, scale, rp) = match w {
22586 GpuTensor::Quant {
22587 bytes,
22588 qtype,
22589 row_bytes,
22590 scale,
22591 rp,
22592 ..
22593 } => (bytes, *qtype, *row_bytes, *scale, *rp),
22594 _ => unreachable!("gemm_supports guaranteed Quant"),
22595 };
22596 if cfg!(memra_hopper_mma)
22602 && qtype == QT_Q8_0
22603 && out_f.is_multiple_of(64)
22604 && wgmma_gemm_enabled()
22605 && let GpuTensor::Quant { rp4: Some(m4), .. } = w
22606 {
22607 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
22608 if scale != 1.0 {
22609 self.scale_inplace(&mut y, scale, m * out_f)?;
22610 }
22611 return Ok(y);
22612 }
22613 let name = match qtype {
22614 QT_Q8_0 => "qmatvec_gemm_q8_0",
22615 QT_Q4_K => "qmatvec_gemm_q4_K",
22616 QT_Q4_0 => {
22617 if rp {
22618 "qmatvec_gemm_q4_0_rp"
22619 } else {
22620 "qmatvec_gemm_q4_0"
22621 }
22622 }
22623 QT_Q5_K => "qmatvec_gemm_q5_K",
22624 QT_Q6_K => "qmatvec_gemm_q6_K",
22625 QT_NVFP4 => {
22626 if rp {
22627 "qmatvec_gemm_nvfp4_rp"
22628 } else {
22629 "qmatvec_gemm_nvfp4"
22630 }
22631 }
22632 _ => unreachable!(),
22633 };
22634 let f = self.func(name);
22635 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);
22640 let k1_tile = if is_k1 {
22642 k1_launch_override().unwrap_or((128, 128, 8))
22643 } else {
22644 (128, 128, 8)
22645 };
22646 let (bm, bn): (u32, u32) = if is_k1 {
22647 (k1_tile.0, k1_tile.1)
22648 } else {
22649 (64, 256)
22650 };
22651 let warps: u32 = if is_k1 {
22652 k1_tile.2
22653 } else {
22654 match qtype {
22655 QT_NVFP4 => 8,
22656 _ => 4,
22657 }
22658 };
22659 let cfg = LaunchConfig {
22660 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
22661 block_dim: (32, warps, 1),
22662 shared_mem_bytes: 0,
22663 };
22664 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22665 let __s_b = self.gpu.stream();
22666 let mut b = __s_b.launch_builder(&f);
22667 b.arg(bytes)
22668 .arg(aq)
22669 .arg(ad)
22670 .arg(&mut y)
22671 .arg(&inf)
22672 .arg(&outf)
22673 .arg(&mi)
22674 .arg(&rb);
22675 unsafe {
22676 b.launch(cfg)?;
22677 }
22678 if scale != 1.0 {
22679 self.scale_inplace(&mut y, scale, m * out_f)?;
22680 }
22681 Ok(y)
22682 }
22683
22684 #[allow(clippy::too_many_arguments)]
22689 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_gemm_raw(
22692 &self,
22693 bytes: &CudaSlice<u8>,
22694 x: &CudaSlice<f32>,
22695 m: usize,
22696 in_f: usize,
22697 out_f: usize,
22698 qtype: i32,
22699 row_bytes: usize,
22700 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22701 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
22702 let name = match qtype {
22703 QT_Q8_0 => "qmatvec_gemm_q8_0",
22704 QT_Q4_K => "qmatvec_gemm_q4_K",
22705 QT_Q4_0 => "qmatvec_gemm_q4_0",
22706 QT_Q5_K => "qmatvec_gemm_q5_K",
22707 QT_Q6_K => "qmatvec_gemm_q6_K",
22708 QT_NVFP4 => "qmatvec_gemm_nvfp4",
22709 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
22710 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
22711 };
22712 let f = self.func(name);
22713 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);
22717 let k1_tile = if is_k1 {
22719 k1_launch_override().unwrap_or((128, 128, 8))
22720 } else {
22721 (128, 128, 8)
22722 };
22723 let (bm, bn): (u32, u32) = if is_k1 {
22724 (k1_tile.0, k1_tile.1)
22725 } else {
22726 (64, 256)
22727 };
22728 let warps: u32 = if is_k1 {
22729 k1_tile.2
22730 } else {
22731 match qtype {
22732 QT_NVFP4 | QT_NVFP4_RP => 8,
22733 _ => 4,
22734 }
22735 };
22736 let cfg = LaunchConfig {
22737 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
22738 block_dim: (32, warps, 1),
22739 shared_mem_bytes: 0,
22740 };
22741 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22742 let __s_b = self.gpu.stream();
22743 let mut b = __s_b.launch_builder(&f);
22744 b.arg(bytes)
22745 .arg(&aq)
22746 .arg(&ad)
22747 .arg(&mut y)
22748 .arg(&inf)
22749 .arg(&outf)
22750 .arg(&mi)
22751 .arg(&rb);
22752 unsafe {
22753 b.launch(cfg)?;
22754 }
22755 Ok(y)
22756 }
22757
22758 pub fn qmatvec_gemm_q8_0_wgmma_raw(
22765 &self,
22766 rp4: &CudaSlice<u8>,
22767 aq: &CudaSlice<i8>,
22768 ad: &CudaSlice<f32>,
22769 m: usize,
22770 in_f: usize,
22771 out_f: usize,
22772 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22773 assert!(
22774 out_f.is_multiple_of(64) && in_f.is_multiple_of(32),
22775 "wgmma GEMM needs out_f%64==0, in_f%32==0"
22776 );
22777 let f = self.func("qmatvec_gemm_q8_0_wgmma");
22778 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
22780 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
22781 block_dim: (128, 1, 1),
22782 shared_mem_bytes: 0,
22783 };
22784 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
22785 let __s_b = self.gpu.stream();
22786 let mut b = __s_b.launch_builder(&f);
22787 b.arg(rp4)
22788 .arg(aq)
22789 .arg(ad)
22790 .arg(&mut y)
22791 .arg(&inf)
22792 .arg(&outf)
22793 .arg(&mi);
22794 unsafe {
22795 b.launch(cfg)?;
22796 }
22797 Ok(y)
22798 }
22799
22800 pub fn scale_inplace(
22802 &self,
22803 y: &mut CudaSlice<f32>,
22804 s: f32,
22805 n: usize,
22806 ) -> Result<(), Box<dyn std::error::Error>> {
22807 let f = self.func("scale_f32");
22808 let cfg = LaunchConfig::for_num_elems(n as u32);
22809 let (sf, ni) = (s, n as i32);
22810 let __s_b = self.gpu.stream();
22811 let mut b = __s_b.launch_builder(&f);
22812 b.arg(y).arg(&sf).arg(&ni);
22813 unsafe {
22814 b.launch(cfg)?;
22815 }
22816 Ok(())
22817 }
22818
22819 pub fn bf16_to_f32(
22824 &self,
22825 data: &cudarc::driver::CudaView<'_, u8>,
22826 n: usize,
22827 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22828 let mut out = self.alloc_uninit::<f32>(n)?;
22829 let f = self.func("bf16_to_f32");
22830 let cfg = LaunchConfig::for_num_elems(n as u32);
22831 let ni = n as i32;
22832 let __s_b = self.gpu.stream();
22833 let mut b = __s_b.launch_builder(&f);
22834 b.arg(data).arg(&mut out).arg(&ni);
22835 unsafe {
22836 b.launch(cfg)?;
22837 }
22838 Ok(out)
22839 }
22840
22841 #[allow(clippy::too_many_arguments)] fn linear_bf16_chunked(
22849 &self,
22850 x: &CudaSlice<f32>,
22851 data: &CudaSlice<u8>,
22852 m: usize,
22853 in_f: usize,
22854 out_f: usize,
22855 exact: bool,
22856 canonical_chunk_rows: Option<usize>,
22857 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22858 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22862 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22863 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22864 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
22865 let started = timing.then(std::time::Instant::now);
22866 let result =
22867 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
22868 if let Some(started) = started {
22869 use std::sync::atomic::Ordering;
22870 self.stream().synchronize()?;
22871 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
22872 + started.elapsed().as_nanos() as u64;
22873 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
22874 + (in_f * out_f * 2) as u64;
22875 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
22876 if calls.is_multiple_of(1024) {
22877 eprintln!(
22878 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
22879 weight_gb={:.2}",
22880 ns as f64 / 1.0e6,
22881 ns as f64 / calls as f64 / 1.0e3,
22882 wb as f64 / 1.0e9,
22883 );
22884 }
22885 }
22886 result
22887 }
22888
22889 pub(crate) fn bf16_mmv_on() -> bool {
22894 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22895 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
22896 }
22897
22898 fn matvec_bf16(
22901 &self,
22902 data: &CudaSlice<u8>,
22903 x: &CudaSlice<f32>,
22904 in_f: usize,
22905 out_f: usize,
22906 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22907 if data.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) {
22908 return Err(format!(
22909 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
22910 data.len(),
22911 x.len()
22912 )
22913 .into());
22914 }
22915 let mut y = self.alloc_uninit::<f32>(out_f)?;
22916 let f = self.func("matvec_bf16_f32acc");
22917 let cfg = LaunchConfig {
22918 grid_dim: (out_f as u32, 1, 1),
22919 block_dim: (mmv_block(), 1, 1),
22920 shared_mem_bytes: 0,
22921 };
22922 let ini = in_f as i32;
22923 let __s_bld = self.gpu.stream();
22924 let mut bld = __s_bld.launch_builder(&f);
22925 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
22926 unsafe {
22927 bld.launch(cfg)?;
22928 }
22929 Ok(y)
22930 }
22931
22932 #[allow(clippy::too_many_arguments)]
22936 #[allow(clippy::too_many_arguments)]
22941 #[allow(clippy::too_many_arguments)]
22946 pub fn qk_norm_rope_append_inc_dcw_rows(
22947 &self,
22948 q_raw_t: &CudaSlice<f32>,
22949 k_raw_t: &CudaSlice<f32>,
22950 v_raw_t: &CudaSlice<f32>,
22951 qw: &CudaSlice<f32>,
22952 kw: &CudaSlice<f32>,
22953 q_out_t: &mut CudaSlice<f32>,
22954 k_out_t: &mut CudaSlice<f32>,
22955 tab: &CudaSlice<u64>,
22956 pos_t: &CudaSlice<i32>,
22957 same_session: bool,
22958 t: usize,
22959 kv_dim_k: usize,
22960 kv_dim_v: usize,
22961 k_tok_bytes: usize,
22962 v_tok_bytes: usize,
22963 head_dim: usize,
22964 n_dims: usize,
22965 nh_q: usize,
22966 nh_k: usize,
22967 eps: f32,
22968 freq_base: f32,
22969 freq_scale: f32,
22970 ff: Option<&CudaSlice<f32>>,
22971 ) -> Result<(), Box<dyn std::error::Error>> {
22972 if head_dim != 128
22973 || kv_dim_v != kv_dim_k
22974 || kv_dim_k != nh_k * head_dim
22975 || t == 0
22976 || t > 32
22977 || tab.len() < t * 6
22978 || pos_t.len() < t
22979 || q_raw_t.len() < t * nh_q * head_dim
22980 || k_raw_t.len() < t * nh_k * head_dim
22981 || v_raw_t.len() < t * kv_dim_v
22982 || q_out_t.len() < t * nh_q * head_dim
22983 || k_out_t.len() < t * nh_k * head_dim
22984 {
22985 return Err(format!(
22986 "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
22987 nh_q={nh_q} nh_k={nh_k}"
22988 )
22989 .into());
22990 }
22991 let f = self.func("qk_norm_rope_append_inc_dcw_rows");
22992 let same_t: i32 = if same_session { t as i32 } else { 0 };
22993 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
22994 let cfg = LaunchConfig {
22995 grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
22996 block_dim: (128, 1, 1),
22997 shared_mem_bytes: 0,
22998 };
22999 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
23000 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23001 let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
23002 let null: u64 = 0;
23003 let __s_b = self.gpu.stream();
23004 let mut b = __s_b.launch_builder(&f);
23005 b.arg(q_raw_t)
23006 .arg(k_raw_t)
23007 .arg(v_raw_t)
23008 .arg(qw)
23009 .arg(kw)
23010 .arg(q_out_t)
23011 .arg(k_out_t)
23012 .arg(tab)
23013 .arg(pos_t)
23014 .arg(&same_t)
23015 .arg(&kvk)
23016 .arg(&kvv)
23017 .arg(&ktb)
23018 .arg(&vtb)
23019 .arg(&hd)
23020 .arg(&nd)
23021 .arg(&nq)
23022 .arg(&nk)
23023 .arg(&eps)
23024 .arg(&theta_scale)
23025 .arg(&freq_scale);
23026 match ff {
23027 Some(freqs) => {
23028 b.arg(freqs);
23029 }
23030 None => {
23031 b.arg(&null);
23032 }
23033 }
23034 unsafe {
23035 b.launch(cfg)?;
23036 }
23037 Ok(())
23038 }
23039
23040 #[allow(clippy::too_many_arguments)] pub fn qk_norm_rope_append_inc_dcw(
23042 &self,
23043 q_raw: &CudaSlice<f32>,
23044 k_raw: &CudaSlice<f32>,
23045 v_raw: &CudaSlice<f32>,
23046 qw: &CudaSlice<f32>,
23047 kw: &CudaSlice<f32>,
23048 q_out: &mut CudaSlice<f32>,
23049 k_out: &mut CudaSlice<f32>,
23050 pos: &CudaSlice<i32>,
23051 k_plane: &mut CudaSlice<u8>,
23052 v_plane: &mut CudaSlice<u8>,
23053 len_dev: &CudaSlice<i32>,
23056 base_dev: Option<&CudaSlice<i32>>,
23057 done_ctr: &mut CudaSlice<u32>,
23058 kv_dim_k: usize,
23059 kv_dim_v: usize,
23060 k_tok_bytes: usize,
23061 v_tok_bytes: usize,
23062 head_dim: usize,
23063 n_dims: usize,
23064 nh_q: usize,
23065 nh_k: usize,
23066 eps: f32,
23067 freq_base: f32,
23068 freq_scale: f32,
23069 ff: Option<&CudaSlice<f32>>,
23070 ) -> Result<(), Box<dyn std::error::Error>> {
23071 if head_dim != 128
23072 || kv_dim_v != kv_dim_k
23073 || kv_dim_k != nh_k * head_dim
23074 || q_raw.len() < nh_q * head_dim
23075 || k_raw.len() < nh_k * head_dim
23076 || v_raw.len() < kv_dim_v
23077 || q_out.len() < nh_q * head_dim
23078 || k_out.len() < nh_k * head_dim
23079 || pos.is_empty()
23080 || done_ctr.is_empty()
23081 {
23082 return Err(format!(
23083 "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}"
23084 )
23085 .into());
23086 }
23087 let f = self.func("qk_norm_rope_append_inc_dcw");
23088 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
23089 let cfg = LaunchConfig {
23090 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
23091 block_dim: (128, 1, 1),
23092 shared_mem_bytes: 0,
23093 };
23094 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
23095 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23096 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
23097 let null: u64 = 0;
23098 let __s_b = self.gpu.stream();
23099 let mut b = __s_b.launch_builder(&f);
23100 b.arg(q_raw)
23101 .arg(k_raw)
23102 .arg(v_raw)
23103 .arg(qw)
23104 .arg(kw)
23105 .arg(q_out)
23106 .arg(k_out)
23107 .arg(pos)
23108 .arg(&mut *k_plane)
23109 .arg(&mut *v_plane)
23110 .arg(len_dev);
23111 match base_dev {
23112 Some(base) => {
23113 b.arg(base);
23114 }
23115 None => {
23116 b.arg(&null);
23117 }
23118 }
23119 b.arg(&mut *done_ctr)
23120 .arg(&kvk)
23121 .arg(&kvv)
23122 .arg(&ktb)
23123 .arg(&vtb)
23124 .arg(&hd)
23125 .arg(&nd)
23126 .arg(&nq)
23127 .arg(&eps)
23128 .arg(&theta_scale)
23129 .arg(&freq_scale);
23130 match ff {
23131 Some(freqs) => {
23132 b.arg(freqs);
23133 }
23134 None => {
23135 b.arg(&null);
23136 }
23137 }
23138 unsafe {
23139 b.launch(cfg)?;
23140 }
23141 Ok(())
23142 }
23143
23144 #[allow(clippy::too_many_arguments)] pub fn qk_norm_rope_into(
23146 &self,
23147 q_raw: &CudaSlice<f32>,
23148 k_raw: &CudaSlice<f32>,
23149 qw: &CudaSlice<f32>,
23150 kw: &CudaSlice<f32>,
23151 q_out: &mut CudaSlice<f32>,
23152 k_out: &mut CudaSlice<f32>,
23153 pos: &CudaSlice<i32>,
23154 head_dim: usize,
23155 n_dims: usize,
23156 nh_q: usize,
23157 nh_k: usize,
23158 eps: f32,
23159 freq_base: f32,
23160 freq_scale: f32,
23161 ff: Option<&CudaSlice<f32>>,
23162 ) -> Result<(), Box<dyn std::error::Error>> {
23163 if head_dim > 512
23164 || q_raw.len() < nh_q * head_dim
23165 || k_raw.len() < nh_k * head_dim
23166 || q_out.len() < nh_q * head_dim
23167 || k_out.len() < nh_k * head_dim
23168 || qw.len() < head_dim
23169 || kw.len() < head_dim
23170 || pos.is_empty()
23171 {
23172 return Err(format!(
23173 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
23174 )
23175 .into());
23176 }
23177 let f = self.func("qk_norm_rope_f32");
23178 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
23179 let cfg = LaunchConfig {
23180 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
23181 block_dim: (128, 1, 1),
23182 shared_mem_bytes: 0,
23183 };
23184 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
23185 let __s_b = self.gpu.stream();
23186 let mut b = __s_b.launch_builder(&f);
23187 b.arg(q_raw)
23188 .arg(k_raw)
23189 .arg(qw)
23190 .arg(kw)
23191 .arg(q_out)
23192 .arg(k_out)
23193 .arg(pos)
23194 .arg(&hd)
23195 .arg(&nd)
23196 .arg(&nq)
23197 .arg(&eps)
23198 .arg(&theta_scale)
23199 .arg(&freq_scale);
23200 match ff {
23201 Some(ffv) => {
23202 b.arg(ffv);
23203 unsafe {
23204 b.launch(cfg)?;
23205 }
23206 }
23207 None => {
23208 let null: u64 = 0;
23209 b.arg(&null);
23210 unsafe {
23211 b.launch(cfg)?;
23212 }
23213 }
23214 }
23215 Ok(())
23216 }
23217
23218 #[allow(clippy::too_many_arguments)]
23221 pub fn matvec_f32_b4_into(
23222 &self,
23223 w: [&CudaSlice<f32>; 4],
23224 x: &CudaSlice<f32>,
23225 y: &mut CudaSlice<f32>,
23226 block_cols: usize,
23227 out_f: usize,
23228 ) -> Result<(), Box<dyn std::error::Error>> {
23229 if !block_cols.is_multiple_of(4)
23230 || x.len() < 4 * block_cols
23231 || y.len() < out_f
23232 || w.iter().any(|w| w.len() != out_f * block_cols)
23233 {
23234 return Err(format!(
23235 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
23236 x.len()
23237 )
23238 .into());
23239 }
23240 let f = self.func("matvec_f32_b4");
23241 let cfg = LaunchConfig {
23242 grid_dim: (out_f as u32, 1, 1),
23243 block_dim: (128, 1, 1),
23244 shared_mem_bytes: 0,
23245 };
23246 let (bc, of) = (block_cols as i32, out_f as i32);
23247 let __s_b = self.gpu.stream();
23248 let mut b = __s_b.launch_builder(&f);
23249 b.arg(w[0])
23250 .arg(w[1])
23251 .arg(w[2])
23252 .arg(w[3])
23253 .arg(x)
23254 .arg(y)
23255 .arg(&bc)
23256 .arg(&of);
23257 unsafe {
23258 b.launch(cfg)?;
23259 }
23260 Ok(())
23261 }
23262
23263 pub fn axpy_rows_seq_into(
23266 &self,
23267 x: &CudaSlice<f32>,
23268 w: &CudaSlice<f32>,
23269 y: &mut CudaSlice<f32>,
23270 width: usize,
23271 n_rows: usize,
23272 ) -> Result<(), Box<dyn std::error::Error>> {
23273 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
23274 return Err(format!(
23275 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
23276 x.len(),
23277 w.len(),
23278 y.len()
23279 )
23280 .into());
23281 }
23282 let f = self.func("axpy_rows_seq_f32");
23283 let cfg = LaunchConfig::for_num_elems(width as u32);
23284 let (wi, nr) = (width as i32, n_rows as i32);
23285 let __s_b = self.gpu.stream();
23286 let mut b = __s_b.launch_builder(&f);
23287 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
23288 unsafe {
23289 b.launch(cfg)?;
23290 }
23291 Ok(())
23292 }
23293
23294 pub fn axpy_rows_seq_tokens_into(
23297 &self,
23298 x: &CudaSlice<f32>,
23299 w: &CudaSlice<f32>,
23300 y: &mut CudaSlice<f32>,
23301 width: usize,
23302 slots: usize,
23303 tokens: usize,
23304 ) -> Result<(), Box<dyn std::error::Error>> {
23305 let rows = slots
23306 .checked_mul(tokens)
23307 .ok_or("axpy_rows_seq_tokens row count overflow")?;
23308 if x.len() < rows * width || w.len() < rows || y.len() < tokens * width {
23309 return Err(format!(
23310 "axpy_rows_seq_tokens geometry x={} w={} y={} width={width} \
23311 slots={slots} tokens={tokens}",
23312 x.len(),
23313 w.len(),
23314 y.len()
23315 )
23316 .into());
23317 }
23318 let f = self.func("axpy_rows_seq_tokens_f32");
23319 let block = 256u32;
23320 let cfg = LaunchConfig {
23321 grid_dim: ((width as u32).div_ceil(block), tokens as u32, 1),
23322 block_dim: (block, 1, 1),
23323 shared_mem_bytes: 0,
23324 };
23325 let (wi, sl, tk) = (width as i32, slots as i32, tokens as i32);
23326 let __s_b = self.gpu.stream();
23327 let mut b = __s_b.launch_builder(&f);
23328 b.arg(x).arg(w).arg(y).arg(&wi).arg(&sl).arg(&tk);
23329 unsafe {
23330 b.launch(cfg)?;
23331 }
23332 Ok(())
23333 }
23334
23335 #[allow(clippy::too_many_arguments)]
23339 pub fn axpy_rows_seq_md_off_into(
23340 &self,
23341 x: &CudaSlice<f32>,
23342 w_route: &CudaSlice<f32>,
23343 md: &CudaSlice<f32>,
23344 sel: &CudaSlice<i32>,
23345 y: &mut CudaSlice<f32>,
23346 width: usize,
23347 n_rows: usize,
23348 row0: usize,
23349 ) -> Result<(), Box<dyn std::error::Error>> {
23350 if x.len() < (row0 + n_rows) * width
23351 || w_route.len() < row0 + n_rows
23352 || sel.len() < row0 + n_rows
23353 || y.len() < width
23354 {
23355 return Err(format!(
23356 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
23357 rows={n_rows} row0={row0}",
23358 x.len(),
23359 w_route.len(),
23360 sel.len(),
23361 y.len()
23362 )
23363 .into());
23364 }
23365 let f = self.func("axpy_rows_seq_md_off_f32");
23366 let cfg = LaunchConfig::for_num_elems(width as u32);
23367 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
23368 let __s_b = self.gpu.stream();
23369 let mut b = __s_b.launch_builder(&f);
23370 b.arg(x)
23371 .arg(w_route)
23372 .arg(md)
23373 .arg(sel)
23374 .arg(y)
23375 .arg(&wi)
23376 .arg(&nr)
23377 .arg(&r0);
23378 unsafe {
23379 b.launch(cfg)?;
23380 }
23381 Ok(())
23382 }
23383
23384 #[allow(clippy::too_many_arguments)]
23387 pub fn axpy_rows_seq_md_into(
23388 &self,
23389 x: &CudaSlice<f32>,
23390 w_route: &CudaSlice<f32>,
23391 md: &CudaSlice<f32>,
23392 sel: &CudaSlice<i32>,
23393 y: &mut CudaSlice<f32>,
23394 width: usize,
23395 n_rows: usize,
23396 ) -> Result<(), Box<dyn std::error::Error>> {
23397 if x.len() < n_rows * width
23398 || w_route.len() < n_rows
23399 || sel.len() < n_rows
23400 || y.len() < width
23401 {
23402 return Err(format!(
23403 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
23404 x.len(),
23405 w_route.len(),
23406 sel.len(),
23407 y.len()
23408 )
23409 .into());
23410 }
23411 let f = self.func("axpy_rows_seq_md_f32");
23412 let cfg = LaunchConfig::for_num_elems(width as u32);
23413 let (wi, nr) = (width as i32, n_rows as i32);
23414 let __s_b = self.gpu.stream();
23415 let mut b = __s_b.launch_builder(&f);
23416 b.arg(x)
23417 .arg(w_route)
23418 .arg(md)
23419 .arg(sel)
23420 .arg(y)
23421 .arg(&wi)
23422 .arg(&nr);
23423 unsafe {
23424 b.launch(cfg)?;
23425 }
23426 Ok(())
23427 }
23428
23429 #[allow(clippy::too_many_arguments)]
23431 #[allow(clippy::too_many_arguments)]
23435 pub fn matvec_bf16_qkvg_tcol_into(
23436 &self,
23437 wq: &CudaSlice<u8>,
23438 wk: &CudaSlice<u8>,
23439 wv: &CudaSlice<u8>,
23440 wg: &CudaSlice<u8>,
23441 x_t: &CudaSlice<f32>,
23442 yq: &mut CudaSlice<f32>,
23443 yk: &mut CudaSlice<f32>,
23444 yv: &mut CudaSlice<f32>,
23445 yg: &mut CudaSlice<f32>,
23446 in_f: usize,
23447 out_q: usize,
23448 out_kv: usize,
23449 out_g: usize,
23450 t: usize,
23451 ) -> Result<(), Box<dyn std::error::Error>> {
23452 if t == 0
23453 || t > 8
23454 || !in_f.is_multiple_of(8)
23455 || x_t.len() < t * in_f
23456 || yq.len() < t * out_q
23457 || yk.len() < t * out_kv
23458 || yv.len() < t * out_kv
23459 || (out_g > 0 && yg.len() < t * out_g)
23460 {
23461 return Err("matvec_bf16_qkvg_tcol geometry".into());
23462 }
23463 let grid = out_q + 2 * out_kv + out_g;
23464 let cfg = LaunchConfig {
23465 grid_dim: (grid as u32, 1, 1),
23466 block_dim: (mmv_block(), 1, 1),
23467 shared_mem_bytes: 0,
23468 };
23469 let (ini, oq, okv, og, ti) = (
23470 in_f as i32,
23471 out_q as i32,
23472 out_kv as i32,
23473 out_g as i32,
23474 t as i32,
23475 );
23476 let __s_b = self.gpu.stream();
23477 let f = self.func("matvec_bf16_qkvg_tcol");
23483 let mut b = __s_b.launch_builder(&f);
23484 b.arg(wq)
23485 .arg(wk)
23486 .arg(wv)
23487 .arg(wg)
23488 .arg(x_t)
23489 .arg(yq)
23490 .arg(yk)
23491 .arg(yv)
23492 .arg(yg)
23493 .arg(&ini)
23494 .arg(&oq)
23495 .arg(&okv)
23496 .arg(&og)
23497 .arg(&ti);
23498 unsafe {
23499 b.launch(cfg)?;
23500 }
23501 Ok(())
23502 }
23503
23504 #[allow(clippy::too_many_arguments)] pub fn matvec_bf16_qkvg_into(
23506 &self,
23507 wq: &CudaSlice<u8>,
23508 wk: &CudaSlice<u8>,
23509 wv: &CudaSlice<u8>,
23510 wg: &CudaSlice<u8>,
23511 x: &CudaSlice<f32>,
23512 yq: &mut CudaSlice<f32>,
23513 yk: &mut CudaSlice<f32>,
23514 yv: &mut CudaSlice<f32>,
23515 yg: &mut CudaSlice<f32>,
23516 in_f: usize,
23517 out_q: usize,
23518 out_kv: usize,
23519 out_g: usize,
23520 ) -> Result<(), Box<dyn std::error::Error>> {
23521 if !in_f.is_multiple_of(8)
23522 || wq.len() != out_q * in_f * 2
23523 || wk.len() != out_kv * in_f * 2
23524 || wv.len() != out_kv * in_f * 2
23525 || wg.len() < out_g * in_f * 2
23526 || x.len() < in_f
23527 || yq.len() < out_q
23528 || yk.len() < out_kv
23529 || yv.len() < out_kv
23530 || (out_g > 0 && yg.len() < out_g)
23531 {
23532 return Err(format!(
23533 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
23534 )
23535 .into());
23536 }
23537 let f = self.func("matvec_bf16_qkvg");
23538 let cfg = LaunchConfig {
23539 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
23540 block_dim: (mmv_block(), 1, 1),
23541 shared_mem_bytes: 0,
23542 };
23543 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
23544 let __s_b = self.gpu.stream();
23545 let mut b = __s_b.launch_builder(&f);
23546 b.arg(wq)
23547 .arg(wk)
23548 .arg(wv)
23549 .arg(wg)
23550 .arg(x)
23551 .arg(yq)
23552 .arg(yk)
23553 .arg(yv)
23554 .arg(yg)
23555 .arg(&inf)
23556 .arg(&oq)
23557 .arg(&okv)
23558 .arg(&og);
23559 unsafe {
23560 b.launch(cfg)?;
23561 }
23562 Ok(())
23563 }
23564
23565 pub fn matvec_bf16_b4_into(
23567 &self,
23568 w: [&CudaSlice<u8>; 4],
23569 x: &CudaSlice<f32>,
23570 y: &mut CudaSlice<f32>,
23571 block_cols: usize,
23572 out_f: usize,
23573 ) -> Result<(), Box<dyn std::error::Error>> {
23574 if !block_cols.is_multiple_of(8)
23575 || x.len() < 4 * block_cols
23576 || y.len() < out_f
23577 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
23578 {
23579 return Err(format!(
23580 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
23581 x.len()
23582 )
23583 .into());
23584 }
23585 let f = self.func("matvec_bf16_b4");
23586 let cfg = LaunchConfig {
23587 grid_dim: (out_f as u32, 1, 1),
23588 block_dim: (mmv_block(), 1, 1),
23589 shared_mem_bytes: 0,
23590 };
23591 let (bc, of) = (block_cols as i32, out_f as i32);
23592 let __s_b = self.gpu.stream();
23593 let mut b = __s_b.launch_builder(&f);
23594 b.arg(w[0])
23595 .arg(w[1])
23596 .arg(w[2])
23597 .arg(w[3])
23598 .arg(x)
23599 .arg(y)
23600 .arg(&bc)
23601 .arg(&of);
23602 unsafe {
23603 b.launch(cfg)?;
23604 }
23605 Ok(())
23606 }
23607
23608 pub fn matvec_bf16_b4_tcol_into(
23613 &self,
23614 w: [&CudaSlice<u8>; 4],
23615 x_t: &CudaSlice<f32>,
23616 y_t: &mut CudaSlice<f32>,
23617 block_cols: usize,
23618 out_f: usize,
23619 t: usize,
23620 ) -> Result<(), Box<dyn std::error::Error>> {
23621 if !block_cols.is_multiple_of(8)
23622 || t == 0
23623 || t > 8
23624 || x_t.len() < t * 4 * block_cols
23625 || y_t.len() < t * out_f
23626 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
23627 {
23628 return Err(format!(
23629 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
23630 x_t.len()
23631 )
23632 .into());
23633 }
23634 let cfg = LaunchConfig {
23638 grid_dim: (out_f as u32, 1, 1),
23639 block_dim: (mmv_block(), 1, 1),
23640 shared_mem_bytes: 0,
23641 };
23642 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
23643 let __s_b = self.gpu.stream();
23644 let f = self.func("matvec_bf16_b4_tcol");
23645 let mut b = __s_b.launch_builder(&f);
23646 b.arg(w[0])
23647 .arg(w[1])
23648 .arg(w[2])
23649 .arg(w[3])
23650 .arg(x_t)
23651 .arg(y_t)
23652 .arg(&bc)
23653 .arg(&of)
23654 .arg(&ti);
23655 unsafe {
23656 b.launch(cfg)?;
23657 }
23658 Ok(())
23659 }
23660
23661 pub fn q8_0_row_bytes(in_f: usize) -> usize {
23664 in_f / 32 * 34
23665 }
23666
23667 pub fn q8_v2_smem_bytes(in_f: usize) -> usize {
23671 in_f + (in_f / 32) * 4
23672 }
23673
23674 pub fn encode_q8_0_from_bf16(
23678 &self,
23679 w_bf16: &CudaSlice<u8>,
23680 out: &mut CudaSlice<u8>,
23681 in_f: usize,
23682 out_f: usize,
23683 ) -> Result<(), Box<dyn std::error::Error>> {
23684 if !in_f.is_multiple_of(32)
23685 || w_bf16.len() < in_f * out_f * 2
23686 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
23687 {
23688 return Err(format!(
23689 "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
23690 w_bf16.len(),
23691 out.len()
23692 )
23693 .into());
23694 }
23695 let f = self.func("encode_q8_0_rows_from_bf16");
23696 const PAIRS_PER_BLOCK: u32 = 4;
23699 let pairs = (out_f * (in_f / 32)) as u64;
23700 let cfg = LaunchConfig {
23701 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
23702 block_dim: (32, PAIRS_PER_BLOCK, 1),
23703 shared_mem_bytes: 0,
23704 };
23705 let (ini, outi) = (in_f as i32, out_f as i32);
23706 let __s_b = self.gpu.stream();
23707 let mut b = __s_b.launch_builder(&f);
23708 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
23709 unsafe {
23710 b.launch(cfg)?;
23711 }
23712 Ok(())
23713 }
23714
23715 pub fn encode_q8_0_from_bf16_view(
23719 &self,
23720 w_bf16: &cudarc::driver::CudaView<'_, u8>,
23721 out: &mut CudaSlice<u8>,
23722 in_f: usize,
23723 out_f: usize,
23724 ) -> Result<(), Box<dyn std::error::Error>> {
23725 if !in_f.is_multiple_of(32)
23726 || w_bf16.len() < in_f * out_f * 2
23727 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
23728 {
23729 return Err(format!(
23730 "encode_q8_0_from_bf16_view geometry in={in_f} out={out_f} src={} dst={}",
23731 w_bf16.len(),
23732 out.len()
23733 )
23734 .into());
23735 }
23736 let f = self.func("encode_q8_0_rows_from_bf16");
23737 const PAIRS_PER_BLOCK: u32 = 4;
23738 let pairs = (out_f * (in_f / 32)) as u64;
23739 let cfg = LaunchConfig {
23740 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
23741 block_dim: (32, PAIRS_PER_BLOCK, 1),
23742 shared_mem_bytes: 0,
23743 };
23744 let (ini, outi) = (in_f as i32, out_f as i32);
23745 let __s_b = self.gpu.stream();
23746 let mut b = __s_b.launch_builder(&f);
23747 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
23748 unsafe {
23749 b.launch(cfg)?;
23750 }
23751 Ok(())
23752 }
23753
23754 #[allow(clippy::too_many_arguments)]
23759 pub fn qmatvec_q8_0_qkv_rp_into(
23760 &self,
23761 wq: &CudaSlice<u8>,
23762 wk: &CudaSlice<u8>,
23763 wv: &CudaSlice<u8>,
23764 aq: &CudaSlice<i8>,
23765 ad: &CudaSlice<f32>,
23766 yq: &mut CudaSlice<f32>,
23767 yk: &mut CudaSlice<f32>,
23768 yv: &mut CudaSlice<f32>,
23769 in_f: usize,
23770 out_q: usize,
23771 out_kv: usize,
23772 ) -> Result<(), Box<dyn std::error::Error>> {
23773 const ROWS_PER_BLOCK: u32 = 4; let rows = out_q + 2 * out_kv;
23775 let nblk = in_f / 32;
23776 if !in_f.is_multiple_of(32)
23777 || aq.len() < in_f
23778 || ad.len() < nblk
23779 || yq.len() < out_q
23780 || yk.len() < out_kv
23781 || yv.len() < out_kv
23782 || wq.len() < out_q * nblk * 34
23783 || wk.len() < out_kv * nblk * 34
23784 || wv.len() < out_kv * nblk * 34
23785 {
23786 return Err(
23787 format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
23788 );
23789 }
23790 let f = self.func("qmatvec_q8_0_qkv_rp");
23791 let cfg = LaunchConfig {
23792 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
23793 block_dim: (32, ROWS_PER_BLOCK, 1),
23794 shared_mem_bytes: 0,
23795 };
23796 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
23797 let __s_b = self.gpu.stream();
23798 let mut b = __s_b.launch_builder(&f);
23799 b.arg(wq)
23800 .arg(wk)
23801 .arg(wv)
23802 .arg(aq)
23803 .arg(ad)
23804 .arg(yq)
23805 .arg(yk)
23806 .arg(yv)
23807 .arg(&ini)
23808 .arg(&oq)
23809 .arg(&okv);
23810 unsafe {
23811 b.launch(cfg)?;
23812 }
23813 Ok(())
23814 }
23815
23816 #[allow(clippy::too_many_arguments)]
23820 pub fn qmatvec_q8_0_b4_rp_into(
23821 &self,
23822 w: [&CudaSlice<u8>; 4],
23823 aq: &CudaSlice<i8>,
23824 ad: &CudaSlice<f32>,
23825 y: &mut CudaSlice<f32>,
23826 block_cols: usize,
23827 out_f: usize,
23828 ) -> Result<(), Box<dyn std::error::Error>> {
23829 const ROWS_PER_BLOCK: u32 = 4; let nblk = block_cols / 32;
23831 if !block_cols.is_multiple_of(32)
23832 || aq.len() < 4 * block_cols
23833 || ad.len() < 4 * nblk
23834 || y.len() < out_f
23835 || w.iter().any(|p| p.len() < out_f * nblk * 34)
23836 {
23837 return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
23838 }
23839 let f = self.func("qmatvec_q8_0_b4_rp");
23840 let cfg = LaunchConfig {
23841 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
23842 block_dim: (32, ROWS_PER_BLOCK, 1),
23843 shared_mem_bytes: 0,
23844 };
23845 let (bc, of) = (block_cols as i32, out_f as i32);
23846 let __s_b = self.gpu.stream();
23847 let mut b = __s_b.launch_builder(&f);
23848 b.arg(w[0])
23849 .arg(w[1])
23850 .arg(w[2])
23851 .arg(w[3])
23852 .arg(aq)
23853 .arg(ad)
23854 .arg(y)
23855 .arg(&bc)
23856 .arg(&of);
23857 unsafe {
23858 b.launch(cfg)?;
23859 }
23860 Ok(())
23861 }
23862
23863 #[allow(clippy::map_entry)] fn matvec_bf16_via_q8_mirror_t(
23867 &self,
23868 data: &CudaSlice<u8>,
23869 x: &CudaSlice<f32>,
23870 y: &mut CudaSlice<f32>,
23871 in_f: usize,
23872 out_f: usize,
23873 t: usize,
23874 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
23875 use cudarc::driver::DevicePtr;
23876 let key = {
23877 let s = self.gpu.stream();
23878 let (p, _g) = data.device_ptr(&s);
23879 (p, in_f as u32, out_f as u32)
23880 };
23881 {
23882 let mut mirrors = self
23883 .w8_mirrors
23884 .lock()
23885 .map_err(|_| "w8 mirror map is poisoned")?;
23886 if !mirrors.contains_key(&key) {
23887 if crate::glm5_graph_capture_open() {
23894 return Err(format!(
23895 "MEMRA_GLM5_W8: the q8_0 mirror for a {in_f}x{out_f} weight would be built inside an \
23896 open CUDA graph capture (its quantize kernels recorded, never executed; memra#131). \
23897 Warm the walk before capturing."
23898 )
23899 .into());
23900 }
23901 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
23902 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
23903 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
23904 mirrors.insert(key, planar);
23905 }
23906 }
23907 let nblk = in_f / 32;
23908 let akey = in_f * 64 + t.min(32);
23910 {
23911 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
23912 if let std::collections::hash_map::Entry::Vacant(slot) = act.entry(akey) {
23913 let aq = self.alloc_i8_uninit(32 * in_f)?;
23914 let ad = self.alloc_uninit::<f32>(32 * nblk)?;
23915 slot.insert((aq, ad));
23916 }
23917 let (aq, ad) = act.get_mut(&akey).expect("just inserted");
23918 self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
23919 }
23920 let mirrors = self
23921 .w8_mirrors
23922 .lock()
23923 .map_err(|_| "w8 mirror map is poisoned")?;
23924 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
23925 let mirror = mirrors.get(&key).expect("built above");
23926 let (aq, ad) = act.get(&akey).expect("built above");
23927 const ROWS_PER_BLOCK: u32 = 4;
23928 let (ini, of) = (in_f as i32, out_f as i32);
23929 if b200_gemv_v2_on() && q8t_wonce_on() && t <= 8 {
23936 if GEMV_V2_Q8_ROWS_TW_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0
23937 {
23938 eprintln!(
23939 "[b200-gemv-v2] engaged arm=q8_rows_tw_v2 t={t} in_f={in_f} out_f={out_f} \
23940 (W8 verify walk; MEMRA_B200_GEMV_V2=1)"
23941 );
23942 }
23943 self.qmatvec_q8_0_rows_tw_v2_raw(mirror, aq, ad, y, in_f, out_f, t)?;
23944 return Ok(Some(()));
23945 }
23946 if q8t_wonce_on() && t <= 32 {
23947 let f = self.func(if t <= 8 {
23948 "qmatvec_q8_0_rows_tw"
23949 } else {
23950 "qmatvec_q8_0_rows_tw32"
23951 });
23952 let cfg = LaunchConfig {
23953 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
23954 block_dim: (32, ROWS_PER_BLOCK, 1),
23955 shared_mem_bytes: 0,
23956 };
23957 let ti = t as i32;
23958 let __s_b = self.gpu.stream();
23959 let mut b = __s_b.launch_builder(&f);
23960 b.arg(mirror)
23961 .arg(aq)
23962 .arg(ad)
23963 .arg(&mut *y)
23964 .arg(&ini)
23965 .arg(&of)
23966 .arg(&ti);
23967 unsafe {
23968 b.launch(cfg)?;
23969 }
23970 return Ok(Some(()));
23971 }
23972 let f = self.func("qmatvec_q8_0_rows_t");
23973 let cfg = LaunchConfig {
23974 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
23975 block_dim: (32, ROWS_PER_BLOCK, 1),
23976 shared_mem_bytes: 0,
23977 };
23978 let __s_b = self.gpu.stream();
23979 let mut b = __s_b.launch_builder(&f);
23980 b.arg(mirror)
23981 .arg(aq)
23982 .arg(ad)
23983 .arg(&mut *y)
23984 .arg(&ini)
23985 .arg(&of);
23986 unsafe {
23987 b.launch(cfg)?;
23988 }
23989 Ok(Some(()))
23990 }
23991
23992 #[allow(clippy::map_entry)] fn matvec_bf16_via_q8_mirror(
23996 &self,
23997 data: &CudaSlice<u8>,
23998 x: &CudaSlice<f32>,
23999 y: &mut CudaSlice<f32>,
24000 in_f: usize,
24001 out_f: usize,
24002 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
24003 use cudarc::driver::DevicePtr;
24004 let key = {
24005 let s = self.gpu.stream();
24006 let (p, _g) = data.device_ptr(&s);
24007 (p, in_f as u32, out_f as u32)
24008 };
24009 {
24010 let mut mirrors = self
24011 .w8_mirrors
24012 .lock()
24013 .map_err(|_| "w8 mirror map is poisoned")?;
24014 if !mirrors.contains_key(&key) {
24015 if crate::glm5_graph_capture_open() {
24022 return Err(format!(
24023 "MEMRA_GLM5_W8: the q8_0 mirror for a {in_f}x{out_f} weight would be built inside an \
24024 open CUDA graph capture (its quantize kernels recorded, never executed; memra#131). \
24025 Warm the walk before capturing."
24026 )
24027 .into());
24028 }
24029 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
24030 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
24031 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
24032 mirrors.insert(key, planar);
24033 if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
24039 eprintln!(
24040 "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
24041 mirrors.len()
24042 );
24043 }
24044 }
24045 }
24046 let nblk = in_f / 32;
24047 {
24048 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24049 if !act.contains_key(&in_f) {
24050 let aq = self.alloc_uninit::<i8>(in_f)?;
24051 let ad = self.alloc_uninit::<f32>(nblk)?;
24052 act.insert(in_f, (aq, ad));
24053 }
24054 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
24055 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
24056 }
24057 let mirrors = self
24058 .w8_mirrors
24059 .lock()
24060 .map_err(|_| "w8 mirror map is poisoned")?;
24061 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24062 let mirror = mirrors.get(&key).expect("built above");
24063 let (aq, ad) = act.get(&in_f).expect("built above");
24064 if b200_gemv_v2_on() && in_f.is_multiple_of(32) && Self::q8_v2_smem_bytes(in_f) <= 48 * 1024
24074 {
24075 if GEMV_V2_Q8_RP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
24076 eprintln!(
24077 "[b200-gemv-v2] engaged arm=q8_rp_v2 t=1 in_f={in_f} out_f={out_f} \
24078 (W8 posture; MEMRA_B200_GEMV_V2=1)"
24079 );
24080 }
24081 self.qmatvec_q8_0_rp_v2_raw(mirror, aq, ad, y, in_f, out_f, 1)?;
24082 return Ok(Some(()));
24083 }
24084 self.qmatvec_mmvq_into(
24085 mirror,
24086 aq,
24087 ad,
24088 1,
24089 in_f,
24090 out_f,
24091 QT_Q8_0,
24092 Self::q8_0_row_bytes(in_f),
24093 1.0,
24094 true,
24095 y,
24096 )?;
24097 Ok(Some(()))
24098 }
24099
24100 #[allow(clippy::too_many_arguments)]
24105 pub fn qmatvec_q8_0_qkv_rp_t_into(
24106 &self,
24107 wq: &CudaSlice<u8>,
24108 wk: &CudaSlice<u8>,
24109 wv: &CudaSlice<u8>,
24110 aq: &CudaSlice<i8>,
24111 ad: &CudaSlice<f32>,
24112 yq: &mut CudaSlice<f32>,
24113 yk: &mut CudaSlice<f32>,
24114 yv: &mut CudaSlice<f32>,
24115 in_f: usize,
24116 out_q: usize,
24117 out_kv: usize,
24118 t: usize,
24119 ) -> Result<(), Box<dyn std::error::Error>> {
24120 const ROWS_PER_BLOCK: u32 = 4;
24121 let rows = out_q + 2 * out_kv;
24122 let nblk = in_f / 32;
24123 if !in_f.is_multiple_of(32)
24124 || t == 0
24125 || aq.len() < t * in_f
24126 || ad.len() < t * nblk
24127 || yq.len() < t * out_q
24128 || yk.len() < t * out_kv
24129 || yv.len() < t * out_kv
24130 {
24131 return Err(format!("q8_0 qkv rp_t geometry in={in_f} t={t}").into());
24132 }
24133 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
24134 if q8t_wonce_on() && t <= 32 {
24138 let f = self.func(if t <= 8 {
24139 "qmatvec_q8_0_qkv_rp_tw"
24140 } else {
24141 "qmatvec_q8_0_qkv_rp_tw32"
24142 });
24143 let cfg = LaunchConfig {
24144 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24145 block_dim: (32, ROWS_PER_BLOCK, 1),
24146 shared_mem_bytes: 0,
24147 };
24148 let ti = t as i32;
24149 let __s_b = self.gpu.stream();
24150 let mut b = __s_b.launch_builder(&f);
24151 b.arg(wq)
24152 .arg(wk)
24153 .arg(wv)
24154 .arg(aq)
24155 .arg(ad)
24156 .arg(yq)
24157 .arg(yk)
24158 .arg(yv)
24159 .arg(&ini)
24160 .arg(&oq)
24161 .arg(&okv)
24162 .arg(&ti);
24163 unsafe {
24164 b.launch(cfg)?;
24165 }
24166 return Ok(());
24167 }
24168 let f = self.func("qmatvec_q8_0_qkv_rp_t");
24169 let cfg = LaunchConfig {
24170 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
24171 block_dim: (32, ROWS_PER_BLOCK, 1),
24172 shared_mem_bytes: 0,
24173 };
24174 let __s_b = self.gpu.stream();
24175 let mut b = __s_b.launch_builder(&f);
24176 b.arg(wq)
24177 .arg(wk)
24178 .arg(wv)
24179 .arg(aq)
24180 .arg(ad)
24181 .arg(yq)
24182 .arg(yk)
24183 .arg(yv)
24184 .arg(&ini)
24185 .arg(&oq)
24186 .arg(&okv);
24187 unsafe {
24188 b.launch(cfg)?;
24189 }
24190 Ok(())
24191 }
24192
24193 #[allow(clippy::too_many_arguments)]
24196 pub fn qmatvec_q8_0_b4_rp_t_into(
24197 &self,
24198 w: [&CudaSlice<u8>; 4],
24199 aq: &CudaSlice<i8>,
24200 ad: &CudaSlice<f32>,
24201 y: &mut CudaSlice<f32>,
24202 block_cols: usize,
24203 out_f: usize,
24204 t: usize,
24205 ) -> Result<(), Box<dyn std::error::Error>> {
24206 const ROWS_PER_BLOCK: u32 = 4;
24207 let nblk = block_cols / 32;
24208 if !block_cols.is_multiple_of(32)
24209 || t == 0
24210 || aq.len() < t * 4 * block_cols
24211 || ad.len() < t * 4 * nblk
24212 || y.len() < t * out_f
24213 {
24214 return Err(format!("q8_0 b4 rp_t geometry cols={block_cols} t={t}").into());
24215 }
24216 let (bc, of) = (block_cols as i32, out_f as i32);
24217 if q8t_wonce_on() && t <= 32 {
24220 let f = self.func(if t <= 8 {
24221 "qmatvec_q8_0_b4_rp_tw"
24222 } else {
24223 "qmatvec_q8_0_b4_rp_tw32"
24224 });
24225 let cfg = LaunchConfig {
24226 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24227 block_dim: (32, ROWS_PER_BLOCK, 1),
24228 shared_mem_bytes: 0,
24229 };
24230 let ti = t as i32;
24231 let __s_b = self.gpu.stream();
24232 let mut b = __s_b.launch_builder(&f);
24233 b.arg(w[0])
24234 .arg(w[1])
24235 .arg(w[2])
24236 .arg(w[3])
24237 .arg(aq)
24238 .arg(ad)
24239 .arg(y)
24240 .arg(&bc)
24241 .arg(&of)
24242 .arg(&ti);
24243 unsafe {
24244 b.launch(cfg)?;
24245 }
24246 return Ok(());
24247 }
24248 let f = self.func("qmatvec_q8_0_b4_rp_t");
24249 let cfg = LaunchConfig {
24250 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
24251 block_dim: (32, ROWS_PER_BLOCK, 1),
24252 shared_mem_bytes: 0,
24253 };
24254 let __s_b = self.gpu.stream();
24255 let mut b = __s_b.launch_builder(&f);
24256 b.arg(w[0])
24257 .arg(w[1])
24258 .arg(w[2])
24259 .arg(w[3])
24260 .arg(aq)
24261 .arg(ad)
24262 .arg(y)
24263 .arg(&bc)
24264 .arg(&of);
24265 unsafe {
24266 b.launch(cfg)?;
24267 }
24268 Ok(())
24269 }
24270
24271 #[allow(clippy::map_entry)] fn matvec_bf16_view_via_q8_mirror(
24282 &self,
24283 data: &cudarc::driver::CudaView<'_, u8>,
24284 x: &CudaSlice<f32>,
24285 y: &mut CudaSlice<f32>,
24286 in_f: usize,
24287 out_f: usize,
24288 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
24289 use cudarc::driver::DevicePtr;
24290 let key = {
24291 let s = self.gpu.stream();
24292 let (p, _g) = data.device_ptr(&s);
24293 (p, in_f as u32, out_f as u32)
24294 };
24295 {
24296 let mut mirrors = self
24297 .w8_mirrors
24298 .lock()
24299 .map_err(|_| "w8 mirror map is poisoned")?;
24300 if !mirrors.contains_key(&key) {
24301 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
24302 self.encode_q8_0_from_bf16_view(data, &mut interleaved, in_f, out_f)?;
24303 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
24304 mirrors.insert(key, planar);
24305 eprintln!("[w8-view] mirror built in_f={in_f} out_f={out_f}");
24309 }
24310 }
24311 let nblk = in_f / 32;
24312 {
24313 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24314 if !act.contains_key(&in_f) {
24315 let aq = self.alloc_uninit::<i8>(in_f)?;
24316 let ad = self.alloc_uninit::<f32>(nblk)?;
24317 act.insert(in_f, (aq, ad));
24318 }
24319 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
24320 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
24321 }
24322 let mirrors = self
24323 .w8_mirrors
24324 .lock()
24325 .map_err(|_| "w8 mirror map is poisoned")?;
24326 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24327 let mirror = mirrors.get(&key).expect("built above");
24328 let (aq, ad) = act.get(&in_f).expect("built above");
24329 self.qmatvec_mmvq_into(
24330 mirror,
24331 aq,
24332 ad,
24333 1,
24334 in_f,
24335 out_f,
24336 QT_Q8_0,
24337 Self::q8_0_row_bytes(in_f),
24338 1.0,
24339 true,
24340 y,
24341 )?;
24342 Ok(Some(()))
24343 }
24344
24345 pub fn matvec_bf16_into(
24346 &self,
24347 data: &CudaSlice<u8>,
24348 x: &CudaSlice<f32>,
24349 y: &mut CudaSlice<f32>,
24350 in_f: usize,
24351 out_f: usize,
24352 ) -> Result<(), Box<dyn std::error::Error>> {
24353 if data.len() != in_f * out_f * 2
24354 || x.len() < in_f
24355 || !in_f.is_multiple_of(8)
24356 || y.len() < out_f
24357 {
24358 return Err(format!(
24359 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24360 data.len(),
24361 x.len(),
24362 y.len()
24363 )
24364 .into());
24365 }
24366 if step_tp_w8_on()
24373 && w8_hybrid_on()
24374 && in_f.is_multiple_of(32)
24375 && out_f >= 64
24376 && let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)?
24377 {
24378 return Ok(());
24379 }
24380 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24384 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
24385 && in_f <= 2048;
24386 if x4 {
24387 let f = self.func("matvec_bf16_f32acc_x4");
24388 let cfg = LaunchConfig {
24389 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
24390 block_dim: (mmv_block(), 1, 1),
24391 shared_mem_bytes: 0,
24392 };
24393 let (ini, outi) = (in_f as i32, out_f as i32);
24394 let __s_b = self.gpu.stream();
24395 let mut b = __s_b.launch_builder(&f);
24396 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
24397 unsafe {
24398 b.launch(cfg)?;
24399 }
24400 return Ok(());
24401 }
24402 let f = self.func("matvec_bf16_f32acc");
24403 let cfg = LaunchConfig {
24404 grid_dim: (out_f as u32, 1, 1),
24405 block_dim: (mmv_block(), 1, 1),
24406 shared_mem_bytes: 0,
24407 };
24408 let ini = in_f as i32;
24409 let __s_b = self.gpu.stream();
24410 let mut b = __s_b.launch_builder(&f);
24411 b.arg(data).arg(x).arg(y).arg(&ini);
24412 unsafe {
24413 b.launch(cfg)?;
24414 }
24415 Ok(())
24416 }
24417
24418 pub fn matvec_bf16_views_into(
24423 &self,
24424 data: &CudaSlice<u8>,
24425 x: &cudarc::driver::CudaView<'_, f32>,
24426 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
24427 in_f: usize,
24428 out_f: usize,
24429 ) -> Result<(), Box<dyn std::error::Error>> {
24430 if data.len() != in_f * out_f * 2
24431 || x.len() < in_f
24432 || !in_f.is_multiple_of(8)
24433 || y.len() < out_f
24434 {
24435 return Err(format!(
24436 "matvec_bf16_views_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24437 data.len(),
24438 x.len(),
24439 y.len()
24440 )
24441 .into());
24442 }
24443 let f = self.func("matvec_bf16_f32acc");
24444 let cfg = LaunchConfig {
24445 grid_dim: (out_f as u32, 1, 1),
24446 block_dim: (mmv_block(), 1, 1),
24447 shared_mem_bytes: 0,
24448 };
24449 let ini = in_f as i32;
24450 let __s_b = self.gpu.stream();
24451 let mut b = __s_b.launch_builder(&f);
24452 b.arg(data).arg(x).arg(y).arg(&ini);
24453 unsafe {
24454 b.launch(cfg)?;
24455 }
24456 Ok(())
24457 }
24458
24459 pub fn matvec_bf16_view_into(
24462 &self,
24463 data: &cudarc::driver::CudaView<'_, u8>,
24464 x: &CudaSlice<f32>,
24465 y: &mut CudaSlice<f32>,
24466 in_f: usize,
24467 out_f: usize,
24468 ) -> Result<(), Box<dyn std::error::Error>> {
24469 if data.len() != in_f * out_f * 2
24470 || x.len() < in_f
24471 || !in_f.is_multiple_of(8)
24472 || y.len() < out_f
24473 {
24474 return Err(format!(
24475 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24476 data.len(),
24477 x.len(),
24478 y.len()
24479 )
24480 .into());
24481 }
24482 if w8_view_on()
24483 && step_tp_w8_on()
24484 && w8_hybrid_on()
24485 && in_f.is_multiple_of(32)
24486 && out_f >= 64
24487 && let Some(()) = self.matvec_bf16_view_via_q8_mirror(data, x, y, in_f, out_f)?
24488 {
24489 return Ok(());
24490 }
24491 let f = self.func("matvec_bf16_f32acc");
24492 let cfg = LaunchConfig {
24493 grid_dim: (out_f as u32, 1, 1),
24494 block_dim: (mmv_block(), 1, 1),
24495 shared_mem_bytes: 0,
24496 };
24497 let ini = in_f as i32;
24498 let __s_b = self.gpu.stream();
24499 let mut b = __s_b.launch_builder(&f);
24500 b.arg(data).arg(x).arg(y).arg(&ini);
24501 unsafe {
24502 b.launch(cfg)?;
24503 }
24504 Ok(())
24505 }
24506
24507 pub fn matvec_bf16_raw_out(
24510 &self,
24511 w: &CudaSlice<u8>,
24512 x: &CudaSlice<f32>,
24513 y_raw: u64,
24514 in_f: usize,
24515 out_f: usize,
24516 ) -> Result<(), Box<dyn std::error::Error>> {
24517 if w.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) || y_raw == 0 {
24518 return Err("matvec_bf16_raw_out geometry".into());
24519 }
24520 let f = self.func("matvec_bf16_f32acc");
24521 let cfg = LaunchConfig {
24522 grid_dim: (out_f as u32, 1, 1),
24523 block_dim: (mmv_block(), 1, 1),
24524 shared_mem_bytes: 0,
24525 };
24526 let ini = in_f as i32;
24527 let __s_b = self.gpu.stream();
24528 let mut b = __s_b.launch_builder(&f);
24529 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
24530 unsafe {
24531 b.launch(cfg)?;
24532 }
24533 Ok(())
24534 }
24535
24536 pub fn add3_raw(
24540 &self,
24541 a: &CudaSlice<f32>,
24542 b: &CudaSlice<f32>,
24543 sh_raw: u64,
24544 scale_raw: u64,
24545 dst: &mut CudaSlice<f32>,
24546 n: usize,
24547 ) -> Result<(), Box<dyn std::error::Error>> {
24548 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
24549 return Err("add3_raw geometry".into());
24550 }
24551 let f = self.func("add3_f32");
24552 let cfg = LaunchConfig {
24553 grid_dim: ((n as u32).div_ceil(256), 1, 1),
24554 block_dim: (256, 1, 1),
24555 shared_mem_bytes: 0,
24556 };
24557 let ni = n as i32;
24558 let __s_b = self.gpu.stream();
24559 let mut bld = __s_b.launch_builder(&f);
24560 bld.arg(a)
24561 .arg(b)
24562 .arg(&sh_raw)
24563 .arg(&scale_raw)
24564 .arg(dst)
24565 .arg(&ni);
24566 unsafe {
24567 bld.launch(cfg)?;
24568 }
24569 Ok(())
24570 }
24571
24572 pub fn matvec_bf16_down_addscale_into(
24575 &self,
24576 w: &CudaSlice<u8>,
24577 x: &CudaSlice<f32>,
24578 scale: &CudaSlice<f32>,
24579 dst: &mut CudaSlice<f32>,
24580 in_f: usize,
24581 out_f: usize,
24582 ) -> Result<(), Box<dyn std::error::Error>> {
24583 if w.len() != in_f * out_f * 2
24584 || x.len() < in_f
24585 || !in_f.is_multiple_of(8)
24586 || dst.len() < out_f
24587 || scale.is_empty()
24588 {
24589 return Err("matvec_bf16_down_addscale geometry".into());
24590 }
24591 let f = self.func("matvec_bf16_down_addscale");
24592 let cfg = LaunchConfig {
24593 grid_dim: (out_f as u32, 1, 1),
24594 block_dim: (mmv_block(), 1, 1),
24595 shared_mem_bytes: 0,
24596 };
24597 let ini = in_f as i32;
24598 let __s_b = self.gpu.stream();
24599 let mut b = __s_b.launch_builder(&f);
24600 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
24601 unsafe {
24602 b.launch(cfg)?;
24603 }
24604 Ok(())
24605 }
24606
24607 #[allow(clippy::too_many_arguments)]
24611 pub fn matvec_bf16_dual_silu_rows_into(
24612 &self,
24613 wg: &CudaSlice<u8>,
24614 wu: &CudaSlice<u8>,
24615 x: &CudaSlice<f32>,
24616 act: &mut CudaSlice<f32>,
24617 in_f: usize,
24618 out_f: usize,
24619 limit: Option<f32>,
24620 t: usize,
24621 ) -> Result<(), Box<dyn std::error::Error>> {
24622 if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
24623 return Err("matvec_bf16_dual_silu_rows geometry".into());
24624 }
24625 let f = self.func("matvec_bf16_dual_silu_rows");
24626 let cfg = LaunchConfig {
24627 grid_dim: (out_f as u32, t as u32, 1),
24628 block_dim: (mmv_block(), 1, 1),
24629 shared_mem_bytes: 0,
24630 };
24631 let (ini, outi) = (in_f as i32, out_f as i32);
24632 let lim = limit.unwrap_or(0.0);
24633 let __s_b = self.gpu.stream();
24634 let mut b = __s_b.launch_builder(&f);
24635 b.arg(wg)
24636 .arg(wu)
24637 .arg(x)
24638 .arg(&mut *act)
24639 .arg(&ini)
24640 .arg(&outi)
24641 .arg(&lim);
24642 unsafe {
24643 b.launch(cfg)?;
24644 }
24645 Ok(())
24646 }
24647
24648 pub fn matvec_bf16_rows_into(
24650 &self,
24651 w: &CudaSlice<u8>,
24652 x: &CudaSlice<f32>,
24653 y: &mut CudaSlice<f32>,
24654 in_f: usize,
24655 out_f: usize,
24656 t: usize,
24657 ) -> Result<(), Box<dyn std::error::Error>> {
24658 if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || !in_f.is_multiple_of(8)
24659 {
24660 return Err("matvec_bf16_rows geometry".into());
24661 }
24662 if (2..=32).contains(&t)
24667 && step_tp_w8_on()
24668 && w8_hybrid_on()
24669 && in_f.is_multiple_of(32)
24670 && out_f >= 64
24671 && let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)?
24672 {
24673 return Ok(());
24674 }
24675 if t == 1
24681 && step_tp_w8_on()
24682 && w8_hybrid_on()
24683 && in_f.is_multiple_of(32)
24684 && out_f >= 64
24685 && let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)?
24686 {
24687 return Ok(());
24688 }
24689 if (2..=32).contains(&t)
24695 && glm5_w8_on()
24696 && in_f.is_multiple_of(32)
24697 && out_f >= 64
24698 && let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)?
24699 {
24700 if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
24701 eprintln!(
24702 "[glm5-w8] engaged t={t} in_f={in_f} out_f={out_f} (q8_0 mirror, \
24703 MEMRA_GLM5_W8=1)"
24704 );
24705 }
24706 return Ok(());
24707 }
24708 if t == 1
24712 && glm5_w8_on()
24713 && in_f.is_multiple_of(32)
24714 && out_f >= 64
24715 && let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)?
24716 {
24717 if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
24718 eprintln!(
24719 "[glm5-w8] engaged t=1 in_f={in_f} out_f={out_f} (q8_0 mirror, \
24720 MEMRA_GLM5_W8=1)"
24721 );
24722 }
24723 return Ok(());
24724 }
24725 if (2..=16).contains(&t) && bf16_tcols_wide_on() {
24733 if BF16_TCOLS_WIDE_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
24734 eprintln!(
24735 "[bf16-tcols-wide] engaged: t={t} in_f={in_f} out_f={out_f} rides the \
24736 weight-once tcols class (MEMRA_BF16_TCOLS_WIDE=1)"
24737 );
24738 }
24739 if t <= 8 {
24740 return self.matvec_bf16_tcols_into(w, x, y, in_f, out_f, t);
24741 }
24742 return self.matvec_bf16_tcols16_into(w, x, y, in_f, out_f, t);
24743 }
24744 if b200_gemv_v2_on() {
24753 let ksplit = self.gemv_v2_ksplit(in_f, out_f, t);
24754 let v3 = b200_gemv_v2_level() >= 2 && ksplit == 1 && gemv_v3_fits();
24758 if GEMV_V2_BF16_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
24759 let arm = if v3 { "v3 (cp.async staged)" } else { "v2" };
24760 eprintln!(
24761 "[b200-gemv-v2] engaged arm={arm} t={t} in_f={in_f} out_f={out_f} \
24762 ksplit={ksplit} (MEMRA_B200_GEMV_V2={})",
24763 b200_gemv_v2_level()
24764 );
24765 }
24766 if v3 {
24767 return self.matvec_bf16_v3_raw(w, x, y, in_f, out_f, t);
24768 }
24769 return self.matvec_bf16_v2_raw(w, x, y, in_f, out_f, t, ksplit);
24770 }
24771 if b200_bf16_gemv_lt_on() && self.bf16_gemv_lt_into(w, x, y, in_f, out_f, t)? {
24779 return Ok(());
24780 }
24781 let kname = if b200_matvec_arm_on() {
24788 "matvec_bf16_f32acc_x4_rows_pf"
24789 } else {
24790 "matvec_bf16_f32acc_x4_rows"
24791 };
24792 let f = self.func(kname);
24793 let cfg = LaunchConfig {
24794 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
24795 block_dim: (mmv_block(), 1, 1),
24796 shared_mem_bytes: 0,
24797 };
24798 let (ini, outi) = (in_f as i32, out_f as i32);
24799 let __s_b = self.gpu.stream();
24800 let mut b = __s_b.launch_builder(&f);
24801 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
24802 unsafe {
24803 b.launch(cfg)?;
24804 }
24805 Ok(())
24806 }
24807
24808 #[allow(clippy::too_many_arguments)]
24814 pub fn matvec_bf16_col_range_into(
24815 &self,
24816 w: &CudaSlice<u8>,
24817 x: &CudaSlice<f32>,
24818 y: &mut CudaSlice<f32>,
24819 in_f: usize,
24820 out_f: usize,
24821 k_start: usize,
24822 k_len: usize,
24823 ) -> Result<(), Box<dyn std::error::Error>> {
24824 let weight_bytes = in_f
24825 .checked_mul(out_f)
24826 .and_then(|elements| elements.checked_mul(2))
24827 .ok_or("BF16 column-range matvec geometry")?;
24828 let k_end = k_start
24829 .checked_add(k_len)
24830 .ok_or("BF16 column-range matvec geometry")?;
24831 if w.len() < weight_bytes
24832 || x.len() < k_len
24833 || y.len() < out_f
24834 || in_f > i32::MAX as usize
24835 || out_f > i32::MAX as usize
24836 || k_start > i32::MAX as usize
24837 || k_len > i32::MAX as usize
24838 || in_f == 0
24839 || out_f == 0
24840 || k_len == 0
24841 || !in_f.is_multiple_of(8)
24842 || !k_start.is_multiple_of(8)
24843 || !k_len.is_multiple_of(8)
24844 || k_end > in_f
24845 {
24846 return Err("BF16 column-range matvec geometry".into());
24847 }
24848 let function = self.func("matvec_bf16_f32acc_x4_range");
24849 let config = LaunchConfig {
24850 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
24851 block_dim: (mmv_block(), 1, 1),
24852 shared_mem_bytes: 0,
24853 };
24854 let (in_f, out_f, k_start, k_len) =
24855 (in_f as i32, out_f as i32, k_start as i32, k_len as i32);
24856 let stream = self.gpu.stream();
24857 let mut launch = stream.launch_builder(&function);
24858 launch
24859 .arg(w)
24860 .arg(x)
24861 .arg(y)
24862 .arg(&in_f)
24863 .arg(&out_f)
24864 .arg(&k_start)
24865 .arg(&k_len);
24866 unsafe {
24867 launch.launch(config)?;
24868 }
24869 Ok(())
24870 }
24871
24872 pub fn matvec_bf16_tcols_into(
24880 &self,
24881 w: &CudaSlice<u8>,
24882 x: &CudaSlice<f32>,
24883 y: &mut CudaSlice<f32>,
24884 in_f: usize,
24885 out_f: usize,
24886 t: usize,
24887 ) -> Result<(), Box<dyn std::error::Error>> {
24888 if x.len() < t * in_f
24889 || y.len() < t * out_f
24890 || !(2..=8).contains(&t)
24891 || !in_f.is_multiple_of(8)
24892 {
24893 return Err("matvec_bf16_tcols geometry".into());
24894 }
24895 let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
24906 if rf
24907 && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
24908 == 0
24909 {
24910 eprintln!(
24911 "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
24912 shared across the t token columns + intra-warp shuffles at the identical \
24913 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
24914 );
24915 }
24916 let x1 = bf16_tcols_x1_on();
24917 let (fname, grid_x) = match (x1, rf) {
24918 (true, true) => ("matvec_bf16_f32acc_x1_tcols_rf", out_f as u32),
24919 (true, false) => ("matvec_bf16_f32acc_x1_tcols", out_f as u32),
24920 (false, true) => ("matvec_bf16_f32acc_x4_tcols_rf", out_f.div_ceil(4) as u32),
24921 (false, false) => ("matvec_bf16_f32acc_x4_tcols", out_f.div_ceil(4) as u32),
24922 };
24923 if x1 && BF16_TCOLS_X1_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
24924 eprintln!(
24925 "[bf16-tcols-x1] engaged: one-row-per-block tcols grid \
24926 (MEMRA_BF16_TCOLS_X1=1)"
24927 );
24928 }
24929 let f = self.func(fname);
24930 let cfg = LaunchConfig {
24931 grid_dim: (grid_x, 1, 1),
24932 block_dim: (mmv_block(), 1, 1),
24933 shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
24934 };
24935 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
24936 let __s_b = self.gpu.stream();
24937 let mut b = __s_b.launch_builder(&f);
24938 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
24939 unsafe {
24940 b.launch(cfg)?;
24941 }
24942 Ok(())
24943 }
24944
24945 pub fn matvec_bf16_tcols16_into(
24951 &self,
24952 w: &CudaSlice<u8>,
24953 x: &CudaSlice<f32>,
24954 y: &mut CudaSlice<f32>,
24955 in_f: usize,
24956 out_f: usize,
24957 t: usize,
24958 ) -> Result<(), Box<dyn std::error::Error>> {
24959 if x.len() < t * in_f
24960 || y.len() < t * out_f
24961 || !(9..=16).contains(&t)
24962 || !in_f.is_multiple_of(8)
24963 {
24964 return Err("matvec_bf16_tcols16 geometry".into());
24965 }
24966 let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
24970 if rf
24971 && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
24972 == 0
24973 {
24974 eprintln!(
24975 "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
24976 shared across the t token columns + intra-warp shuffles at the identical \
24977 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
24978 );
24979 }
24980 let f = self.func(if rf {
24981 "matvec_bf16_f32acc_x4_tcols16_rf"
24982 } else {
24983 "matvec_bf16_f32acc_x4_tcols16"
24984 });
24985 let cfg = LaunchConfig {
24986 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
24987 block_dim: (mmv_block(), 1, 1),
24988 shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
24989 };
24990 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
24991 let __s_b = self.gpu.stream();
24992 let mut b = __s_b.launch_builder(&f);
24993 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
24994 unsafe {
24995 b.launch(cfg)?;
24996 }
24997 Ok(())
24998 }
24999
25000 #[allow(clippy::too_many_arguments)]
25006 pub fn matvec_bf16_tcols_gate_kernel_into(
25007 &self,
25008 kernel: &str,
25009 w: &CudaSlice<u8>,
25010 x: &CudaSlice<f32>,
25011 y: &mut CudaSlice<f32>,
25012 in_f: usize,
25013 out_f: usize,
25014 t: usize,
25015 ) -> Result<(), Box<dyn std::error::Error>> {
25016 let (grid_x, t_max) = match kernel {
25017 "matvec_bf16_f32acc_x1_tcols_rf" | "matvec_bf16_f32acc_x1_tcols_rf_redshift" => {
25018 (out_f as u32, 8usize)
25019 }
25020 "matvec_bf16_f32acc_x4_tcols_rf" => (out_f.div_ceil(4) as u32, 8usize),
25021 "matvec_bf16_f32acc_x4_tcols16_rf" => (out_f.div_ceil(4) as u32, 16usize),
25022 _ => return Err("matvec_bf16_tcols_gate_kernel_into: unknown kernel".into()),
25023 };
25024 if x.len() < t * in_f
25025 || y.len() < t * out_f
25026 || !(1..=t_max).contains(&t)
25027 || !in_f.is_multiple_of(8)
25028 || !mmv_block().is_power_of_two()
25029 {
25030 return Err("matvec_bf16_tcols_gate_kernel geometry".into());
25031 }
25032 let f = self.func(kernel);
25033 let cfg = LaunchConfig {
25034 grid_dim: (grid_x, 1, 1),
25035 block_dim: (mmv_block(), 1, 1),
25036 shared_mem_bytes: (t as u32) * mmv_block() * 4,
25037 };
25038 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
25039 let __s_b = self.gpu.stream();
25040 let mut b = __s_b.launch_builder(&f);
25041 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
25042 unsafe {
25043 b.launch(cfg)?;
25044 }
25045 Ok(())
25046 }
25047
25048 pub fn matmul_rows_exact(
25056 &self,
25057 w: &crate::model::GpuTensor,
25058 x: &CudaSlice<f32>,
25059 m: usize,
25060 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25061 use crate::model::GpuTensor;
25062 if let GpuTensor::FloatBf16 { data, .. } = w
25067 && (2..=32).contains(&m)
25068 && Self::bf16_mmv_on()
25069 && w.in_features().is_multiple_of(32)
25070 && w.out_features() >= 64
25071 && glm5_w8_on()
25072 {
25073 let (in_f, out_f) = (w.in_features(), w.out_features());
25074 let mut y = self.vws_uninit(m * out_f)?;
25075 if let Some(()) = self.matvec_bf16_via_q8_mirror_t(data, x, &mut y, in_f, out_f, m)? {
25076 if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25077 eprintln!(
25078 "[glm5-w8] engaged rows-exact m={m} in_f={in_f} out_f={out_f} \
25079 (q8_0 mirror, MEMRA_GLM5_W8=1)"
25080 );
25081 }
25082 return Ok(y);
25083 }
25084 }
25085 if let GpuTensor::FloatBf16 { data, .. } = w
25086 && (2..=8).contains(&m)
25087 && Self::bf16_mmv_on()
25088 && w.in_features().is_multiple_of(8)
25089 && !(step_tp_w8_on() && w8_hybrid_on())
25090 && !glm5_w8_on()
25091 {
25092 let (in_f, out_f) = (w.in_features(), w.out_features());
25093 let mut y = self.vws_uninit(m * out_f)?;
25096 self.matvec_bf16_tcols_into(data, x, &mut y, in_f, out_f, m)?;
25097 return Ok(y);
25098 }
25099 self.matmul_decode_exact(w, x, m)
25100 }
25101
25102 #[allow(clippy::too_many_arguments)] pub fn matvec_bf16_dual_silu_into(
25104 &self,
25105 wg: &CudaSlice<u8>,
25106 wu: &CudaSlice<u8>,
25107 x: &CudaSlice<f32>,
25108 act: &mut CudaSlice<f32>,
25109 in_f: usize,
25110 out_f: usize,
25111 limit: Option<f32>,
25112 ) -> Result<(), Box<dyn std::error::Error>> {
25113 if wg.len() != in_f * out_f * 2
25114 || wu.len() != in_f * out_f * 2
25115 || x.len() < in_f
25116 || !in_f.is_multiple_of(8)
25117 || act.len() < out_f
25118 {
25119 return Err("matvec_bf16_dual_silu geometry".into());
25120 }
25121 let f = self.func("matvec_bf16_dual_silu");
25122 let cfg = LaunchConfig {
25123 grid_dim: (out_f as u32, 1, 1),
25124 block_dim: (mmv_block(), 1, 1),
25125 shared_mem_bytes: 0,
25126 };
25127 let (ini, outi) = (in_f as i32, out_f as i32);
25128 let lim = limit.unwrap_or(0.0);
25129 let __s_b = self.gpu.stream();
25130 let mut b = __s_b.launch_builder(&f);
25131 b.arg(wg)
25132 .arg(wu)
25133 .arg(x)
25134 .arg(act)
25135 .arg(&ini)
25136 .arg(&outi)
25137 .arg(&lim);
25138 unsafe {
25139 b.launch(cfg)?;
25140 }
25141 Ok(())
25142 }
25143
25144 #[allow(clippy::too_many_arguments)]
25147 pub fn matvec_bf16_dual_view_into(
25148 &self,
25149 wg: &cudarc::driver::CudaView<'_, u8>,
25150 wu: &cudarc::driver::CudaView<'_, u8>,
25151 x: &CudaSlice<f32>,
25152 yg: &mut CudaSlice<f32>,
25153 yu: &mut CudaSlice<f32>,
25154 in_f: usize,
25155 out_f: usize,
25156 ) -> Result<(), Box<dyn std::error::Error>> {
25157 if wg.len() != in_f * out_f * 2
25158 || wu.len() != in_f * out_f * 2
25159 || x.len() < in_f
25160 || !in_f.is_multiple_of(8)
25161 || yg.len() < out_f
25162 || yu.len() < out_f
25163 {
25164 return Err(format!(
25165 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
25166 wg.len(),
25167 wu.len(),
25168 x.len()
25169 )
25170 .into());
25171 }
25172 let f = self.func("matvec_bf16_dual");
25173 let cfg = LaunchConfig {
25174 grid_dim: ((2 * out_f) as u32, 1, 1),
25175 block_dim: (mmv_block(), 1, 1),
25176 shared_mem_bytes: 0,
25177 };
25178 let (ini, outi) = (in_f as i32, out_f as i32);
25179 let __s_b = self.gpu.stream();
25180 let mut b = __s_b.launch_builder(&f);
25181 b.arg(wg)
25182 .arg(wu)
25183 .arg(x)
25184 .arg(yg)
25185 .arg(yu)
25186 .arg(&ini)
25187 .arg(&outi);
25188 unsafe {
25189 b.launch(cfg)?;
25190 }
25191 Ok(())
25192 }
25193
25194 #[allow(clippy::too_many_arguments)]
25196 pub fn matvec_bf16_dual_into(
25197 &self,
25198 wg: &CudaSlice<u8>,
25199 wu: &CudaSlice<u8>,
25200 x: &CudaSlice<f32>,
25201 yg: &mut CudaSlice<f32>,
25202 yu: &mut CudaSlice<f32>,
25203 in_f: usize,
25204 out_f: usize,
25205 ) -> Result<(), Box<dyn std::error::Error>> {
25206 if wg.len() != in_f * out_f * 2
25207 || wu.len() != in_f * out_f * 2
25208 || x.len() < in_f
25209 || !in_f.is_multiple_of(8)
25210 || yg.len() < out_f
25211 || yu.len() < out_f
25212 {
25213 return Err(format!(
25214 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
25215 wg.len(),
25216 wu.len(),
25217 x.len()
25218 )
25219 .into());
25220 }
25221 let f = self.func("matvec_bf16_dual");
25222 let cfg = LaunchConfig {
25223 grid_dim: ((2 * out_f) as u32, 1, 1),
25224 block_dim: (mmv_block(), 1, 1),
25225 shared_mem_bytes: 0,
25226 };
25227 let (ini, outi) = (in_f as i32, out_f as i32);
25228 let __s_b = self.gpu.stream();
25229 let mut b = __s_b.launch_builder(&f);
25230 b.arg(wg)
25231 .arg(wu)
25232 .arg(x)
25233 .arg(yg)
25234 .arg(yu)
25235 .arg(&ini)
25236 .arg(&outi);
25237 unsafe {
25238 b.launch(cfg)?;
25239 }
25240 Ok(())
25241 }
25242
25243 #[allow(dead_code)] pub(crate) fn matvec_bf16_dual(
25247 &self,
25248 wg: &CudaSlice<u8>,
25249 wu: &CudaSlice<u8>,
25250 x: &CudaSlice<f32>,
25251 in_f: usize,
25252 out_f: usize,
25253 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25254 if wg.len() != in_f * out_f * 2
25255 || wu.len() != in_f * out_f * 2
25256 || x.len() < in_f
25257 || !in_f.is_multiple_of(8)
25258 {
25259 return Err(format!(
25260 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
25261 wg.len(),
25262 wu.len(),
25263 x.len()
25264 )
25265 .into());
25266 }
25267 let mut yg = self.alloc_uninit::<f32>(out_f)?;
25268 let mut yu = self.alloc_uninit::<f32>(out_f)?;
25269 let f = self.func("matvec_bf16_dual");
25270 let cfg = LaunchConfig {
25271 grid_dim: ((2 * out_f) as u32, 1, 1),
25272 block_dim: (mmv_block(), 1, 1),
25273 shared_mem_bytes: 0,
25274 };
25275 let (ini, outi) = (in_f as i32, out_f as i32);
25276 let __s_b = self.gpu.stream();
25277 let mut b = __s_b.launch_builder(&f);
25278 b.arg(wg)
25279 .arg(wu)
25280 .arg(x)
25281 .arg(&mut yg)
25282 .arg(&mut yu)
25283 .arg(&ini)
25284 .arg(&outi);
25285 unsafe {
25286 b.launch(cfg)?;
25287 }
25288 Ok((yg, yu))
25289 }
25290
25291 #[allow(clippy::too_many_arguments)]
25292 #[allow(clippy::manual_is_multiple_of)] fn linear_bf16_chunked_inner(
25294 &self,
25295 x: &CudaSlice<f32>,
25296 data: &CudaSlice<u8>,
25297 m: usize,
25298 in_f: usize,
25299 out_f: usize,
25300 exact: bool,
25301 canonical_chunk_rows: Option<usize>,
25302 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25303 const CHUNK_BYTES: usize = 256 << 20;
25304 if m == 1
25307 && !exact
25308 && canonical_chunk_rows.is_none()
25309 && in_f.is_multiple_of(8)
25310 && Self::bf16_mmv_on()
25311 {
25312 return self.matvec_bf16(data, x, in_f, out_f);
25313 }
25314 if m >= 16
25319 && !exact
25320 && canonical_chunk_rows.is_none()
25321 && data.len() == in_f * out_f * 2
25322 && crate::f16_ffi::pp_bf16_enabled()
25323 {
25324 if let Some(y) = self.bf16_tc_gemm(data, x, m, in_f, out_f)? {
25327 return Ok(y);
25328 }
25329 }
25330 let row_bytes = in_f
25331 .checked_mul(std::mem::size_of::<f32>())
25332 .ok_or("BF16 chunk row byte count overflow")?;
25333 if row_bytes == 0 || out_f == 0 {
25334 return Err("BF16 chunk dimensions must be nonzero".into());
25335 }
25336 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
25337 let chunk_rows = match canonical_chunk_rows {
25338 Some(0) => {
25339 return Err("canonical BF16 chunk rows must be nonzero".into());
25340 }
25341 Some(rows) if rows > max_chunk_rows => {
25342 return Err(format!(
25343 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
25344 )
25345 .into());
25346 }
25347 Some(rows) if out_f % rows != 0 => {
25348 return Err(format!(
25349 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
25350 )
25351 .into());
25352 }
25353 Some(rows) => rows,
25354 None => max_chunk_rows,
25355 };
25356 if chunk_rows >= out_f {
25357 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
25358 return if exact {
25359 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
25360 } else {
25361 self.linear(x, &wf32, m, in_f, out_f)
25362 };
25363 }
25364 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
25365 let mut r0 = 0usize;
25366 while r0 < out_f {
25367 let rows = chunk_rows.min(out_f - r0);
25368 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
25369 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
25370 let yc = if exact {
25371 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
25372 } else {
25373 self.linear(x, &wf32, m, in_f, rows)?
25374 };
25375 for mi in 0..m {
25377 let src = yc.slice(mi * rows..(mi + 1) * rows);
25378 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
25379 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
25380 }
25381 r0 += rows;
25382 }
25383 Ok(y)
25384 }
25385
25386 pub fn linear_bf16_resident(
25390 &self,
25391 x: &CudaSlice<f32>,
25392 data: &CudaSlice<u8>,
25393 m: usize,
25394 in_f: usize,
25395 out_f: usize,
25396 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25397 if data.len() != in_f * out_f * 2 {
25398 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
25399 }
25400 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
25401 }
25402
25403 pub fn linear_bf16_resident_canonical_rows(
25409 &self,
25410 x: &CudaSlice<f32>,
25411 data: &CudaSlice<u8>,
25412 m: usize,
25413 in_f: usize,
25414 out_f: usize,
25415 canonical_chunk_rows: usize,
25416 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25417 if data.len() != in_f * out_f * 2 {
25418 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
25419 }
25420 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
25421 }
25422
25423 pub fn linear_f32_resident_canonical_rows(
25428 &self,
25429 x: &CudaSlice<f32>,
25430 data: &CudaSlice<f32>,
25431 m: usize,
25432 in_f: usize,
25433 out_f: usize,
25434 canonical_chunk_rows: usize,
25435 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25436 self.linear_f32_resident_canonical_rows_inner(
25437 x,
25438 data,
25439 m,
25440 in_f,
25441 out_f,
25442 canonical_chunk_rows,
25443 false,
25444 )
25445 }
25446
25447 pub fn linear_f32_resident_canonical_rows_strided(
25453 &self,
25454 x: &CudaSlice<f32>,
25455 data: &CudaSlice<f32>,
25456 m: usize,
25457 in_f: usize,
25458 out_f: usize,
25459 canonical_chunk_rows: usize,
25460 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25461 self.linear_f32_resident_canonical_rows_inner(
25462 x,
25463 data,
25464 m,
25465 in_f,
25466 out_f,
25467 canonical_chunk_rows,
25468 true,
25469 )
25470 }
25471
25472 #[allow(clippy::too_many_arguments)]
25473 #[allow(clippy::manual_is_multiple_of)] fn linear_f32_resident_canonical_rows_inner(
25476 &self,
25477 x: &CudaSlice<f32>,
25478 data: &CudaSlice<f32>,
25479 m: usize,
25480 in_f: usize,
25481 out_f: usize,
25482 canonical_chunk_rows: usize,
25483 strided_output: bool,
25484 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25485 if data.len() != in_f * out_f {
25486 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
25487 }
25488 if canonical_chunk_rows == 0
25489 || canonical_chunk_rows > out_f
25490 || out_f % canonical_chunk_rows != 0
25491 {
25492 return Err(format!(
25493 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
25494 )
25495 .into());
25496 }
25497 if canonical_chunk_rows == out_f {
25498 return self.linear(x, data, m, in_f, out_f);
25499 }
25500
25501 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
25502 let input = x.slice(0..x.len());
25503 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
25504 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
25505 if m == 1 {
25506 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
25507 self.linear_device_into(
25508 &input,
25509 &weights,
25510 &mut destination,
25511 1,
25512 in_f,
25513 canonical_chunk_rows,
25514 )?;
25515 continue;
25516 }
25517 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
25518 if strided_output {
25519 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
25520 } else {
25521 for token in 0..m {
25522 let source = chunk
25523 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
25524 let mut destination =
25525 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
25526 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
25527 }
25528 }
25529 }
25530 Ok(y)
25531 }
25532
25533 #[allow(clippy::manual_is_multiple_of)] pub fn linear_f32_resident_canonical_rows_t1_into(
25539 &self,
25540 x: &CudaSlice<f32>,
25541 data: &CudaSlice<f32>,
25542 y: &mut CudaSlice<f32>,
25543 in_f: usize,
25544 out_f: usize,
25545 canonical_chunk_rows: usize,
25546 ) -> Result<(), Box<dyn std::error::Error>> {
25547 if data.len() != in_f * out_f {
25548 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
25549 }
25550 if y.len() != out_f || x.len() != in_f {
25551 return Err(format!(
25552 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
25553 x.len(),
25554 y.len()
25555 )
25556 .into());
25557 }
25558 if canonical_chunk_rows == 0
25559 || canonical_chunk_rows > out_f
25560 || out_f % canonical_chunk_rows != 0
25561 {
25562 return Err(format!(
25563 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
25564 )
25565 .into());
25566 }
25567 let input = x.slice(0..x.len());
25568 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
25569 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
25570 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
25571 self.linear_device_into(
25572 &input,
25573 &weights,
25574 &mut destination,
25575 1,
25576 in_f,
25577 canonical_chunk_rows,
25578 )?;
25579 }
25580 Ok(())
25581 }
25582
25583 pub fn linear_t1_into(
25586 &self,
25587 x: &cudarc::driver::CudaView<'_, f32>,
25588 w: &cudarc::driver::CudaView<'_, f32>,
25589 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
25590 in_f: usize,
25591 out_f: usize,
25592 ) -> Result<(), Box<dyn std::error::Error>> {
25593 self.linear_device_into(x, w, y, 1, in_f, out_f)
25594 }
25595
25596 pub fn linear_decode_exact(
25603 &self,
25604 x: &CudaSlice<f32>,
25605 w: &CudaSlice<f32>,
25606 m_tokens: usize,
25607 in_f: usize,
25608 out_f: usize,
25609 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25610 if m_tokens == 1 {
25611 return self.linear(x, w, 1, in_f, out_f);
25612 }
25613 if crate::f32_gemv_kernel_on() {
25617 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
25618 if self.gemv_f32_rows_into(x, w, &mut y, m_tokens, in_f, out_f)? {
25619 return Ok(y);
25620 }
25621 }
25622 let xv = self.view(x, m_tokens * in_f);
25623 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
25624 for t in 0..m_tokens {
25625 let row = xv.slice(t * in_f..(t + 1) * in_f);
25626 let mut xr = self.alloc_uninit::<f32>(in_f)?;
25627 self.copy_view_into(&mut xr, 0, &row, in_f)?;
25628 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
25629 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
25630 }
25631 Ok(y)
25632 }
25633
25634 pub fn linear(
25635 &self,
25636 x: &CudaSlice<f32>,
25637 w: &CudaSlice<f32>,
25638 m_tokens: usize,
25639 in_f: usize,
25640 out_f: usize,
25641 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25642 self.linear_device(x, w, m_tokens, in_f, out_f)
25643 }
25644
25645 fn linear_device<I>(
25646 &self,
25647 x: &I,
25648 w: &I,
25649 m_tokens: usize,
25650 in_f: usize,
25651 out_f: usize,
25652 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
25653 where
25654 I: cudarc::driver::DevicePtr<f32>,
25655 {
25656 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)?;
25658 Ok(c)
25659 }
25660
25661 fn linear_device_into<I, O>(
25662 &self,
25663 x: &I,
25664 w: &I,
25665 c: &mut O,
25666 m_tokens: usize,
25667 in_f: usize,
25668 out_f: usize,
25669 ) -> Result<(), Box<dyn std::error::Error>>
25670 where
25671 I: cudarc::driver::DevicePtr<f32>,
25672 O: cudarc::driver::DevicePtrMut<f32>,
25673 {
25674 if crate::f32_gemv_kernel_on() && self.gemv_f32_rows_into(x, w, c, m_tokens, in_f, out_f)? {
25675 return Ok(());
25676 }
25677 use cudarc::cublaslt::{Matmul, MatmulConfig};
25678 let cfg = MatmulConfig {
25679 transa: true,
25680 transb: false,
25681 transc: false,
25682 m: out_f as u64,
25683 n: m_tokens as u64,
25684 k: in_f as u64,
25685 alpha: 1.0,
25686 lda: in_f as i64,
25687 ldb: in_f as i64,
25688 beta: 0.0,
25689 ldc: out_f as i64,
25690 stride_a: None,
25691 stride_b: None,
25692 stride_c: None,
25693 stride_bias: None,
25694 batch_size: None,
25695 };
25696 let blas = self.gpu.blas();
25697 unsafe {
25698 blas.matmul(cfg, w, x, c, None, None)?;
25699 }
25700 Ok(())
25701 }
25702
25703 #[allow(clippy::too_many_arguments)] pub fn sdpa_naive(
25713 &self,
25714 q: &CudaSlice<f32>,
25715 k: &CudaSlice<f32>,
25716 v: &CudaSlice<f32>,
25717 o: &mut CudaSlice<f32>,
25718 head_dim: usize,
25719 n_head: usize,
25720 n_head_kv: usize,
25721 t: usize,
25722 t_kv: usize,
25723 scale: f32,
25724 causal: bool,
25725 ) -> Result<(), Box<dyn std::error::Error>> {
25726 if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
25727 return self.sdpa_naive_gmem(
25728 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
25729 );
25730 }
25731 let f = self.func("sdpa_naive_f32");
25732 let cfg = LaunchConfig {
25733 grid_dim: (n_head as u32, t as u32, 1),
25734 block_dim: (128, 1, 1),
25735 shared_mem_bytes: (t_kv * 4) as u32,
25736 };
25737 let (hd, nh, nhkv, ti, tkvi, cz) = (
25738 head_dim as i32,
25739 n_head as i32,
25740 n_head_kv as i32,
25741 t as i32,
25742 t_kv as i32,
25743 causal as i32,
25744 );
25745 let __s_b = self.gpu.stream();
25746 let mut b = __s_b.launch_builder(&f);
25747 b.arg(q)
25748 .arg(k)
25749 .arg(v)
25750 .arg(o)
25751 .arg(&hd)
25752 .arg(&nh)
25753 .arg(&nhkv)
25754 .arg(&ti)
25755 .arg(&tkvi)
25756 .arg(&scale)
25757 .arg(&cz);
25758 unsafe {
25759 b.launch(cfg)?;
25760 }
25761 Ok(())
25762 }
25763
25764 #[allow(clippy::too_many_arguments)]
25773 pub fn sdpa_naive_gmem(
25774 &self,
25775 q: &CudaSlice<f32>,
25776 k: &CudaSlice<f32>,
25777 v: &CudaSlice<f32>,
25778 o: &mut CudaSlice<f32>,
25779 head_dim: usize,
25780 n_head: usize,
25781 n_head_kv: usize,
25782 t: usize,
25783 t_kv: usize,
25784 scale: f32,
25785 causal: bool,
25786 ) -> Result<(), Box<dyn std::error::Error>> {
25787 let ws_len = n_head
25788 .checked_mul(t)
25789 .and_then(|x| x.checked_mul(t_kv))
25790 .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
25791 let ws_bytes = ws_len
25792 .checked_mul(std::mem::size_of::<f32>())
25793 .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
25794 if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
25795 return Err(format!(
25796 "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
25797 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
25798 needs a tiled/flash kernel, not the naive oracle"
25799 )
25800 .into());
25801 }
25802 let mut scores = self.uninit(ws_len)?;
25803 let f = self.func("sdpa_naive_gmem_f32");
25804 let cfg = LaunchConfig {
25805 grid_dim: (n_head as u32, t as u32, 1),
25806 block_dim: (128, 1, 1),
25807 shared_mem_bytes: 0,
25808 };
25809 let (hd, nh, nhkv, ti, tkvi, cz) = (
25810 head_dim as i32,
25811 n_head as i32,
25812 n_head_kv as i32,
25813 t as i32,
25814 t_kv as i32,
25815 causal as i32,
25816 );
25817 let __s_b = self.gpu.stream();
25818 let mut b = __s_b.launch_builder(&f);
25819 b.arg(q)
25820 .arg(k)
25821 .arg(v)
25822 .arg(o)
25823 .arg(&mut scores)
25824 .arg(&hd)
25825 .arg(&nh)
25826 .arg(&nhkv)
25827 .arg(&ti)
25828 .arg(&tkvi)
25829 .arg(&scale)
25830 .arg(&cz);
25831 unsafe {
25832 b.launch(cfg)?;
25833 }
25834 Ok(())
25835 }
25836
25837 #[allow(clippy::too_many_arguments)]
25842 pub fn sdpa_naive_island(
25843 &self,
25844 q: &CudaSlice<f32>,
25845 k: &CudaSlice<f32>,
25846 v: &CudaSlice<f32>,
25847 o: &mut CudaSlice<f32>,
25848 span_id: &CudaSlice<i32>,
25849 head_dim: usize,
25850 n_head: usize,
25851 n_head_kv: usize,
25852 t: usize,
25853 t_kv: usize,
25854 scale: f32,
25855 window: usize,
25856 ) -> Result<(), Box<dyn std::error::Error>> {
25857 let f = self.func("sdpa_naive_island_f32");
25858 let cfg = LaunchConfig {
25859 grid_dim: (n_head as u32, t as u32, 1),
25860 block_dim: (128, 1, 1),
25861 shared_mem_bytes: (t_kv * 4) as u32,
25862 };
25863 let (hd, nh, nhkv, ti, tkvi, wi) = (
25864 head_dim as i32,
25865 n_head as i32,
25866 n_head_kv as i32,
25867 t as i32,
25868 t_kv as i32,
25869 window as i32,
25870 );
25871 let __s_b = self.gpu.stream();
25872 let mut b = __s_b.launch_builder(&f);
25873 b.arg(q)
25874 .arg(k)
25875 .arg(v)
25876 .arg(o)
25877 .arg(span_id)
25878 .arg(&hd)
25879 .arg(&nh)
25880 .arg(&nhkv)
25881 .arg(&ti)
25882 .arg(&tkvi)
25883 .arg(&scale)
25884 .arg(&wi);
25885 unsafe {
25886 b.launch(cfg)?;
25887 }
25888 Ok(())
25889 }
25890
25891 #[allow(clippy::too_many_arguments)]
25893 pub fn sdpa_naive_w(
25894 &self,
25895 q: &CudaSlice<f32>,
25896 k: &CudaSlice<f32>,
25897 v: &CudaSlice<f32>,
25898 o: &mut CudaSlice<f32>,
25899 head_dim: usize,
25900 n_head: usize,
25901 n_head_kv: usize,
25902 t: usize,
25903 t_kv: usize,
25904 scale: f32,
25905 causal: bool,
25906 window: usize,
25907 ) -> Result<(), Box<dyn std::error::Error>> {
25908 let f = self.func("sdpa_naive_w_f32");
25909 let cfg = LaunchConfig {
25910 grid_dim: (n_head as u32, t as u32, 1),
25911 block_dim: (128, 1, 1),
25912 shared_mem_bytes: (t_kv * 4) as u32,
25913 };
25914 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
25915 head_dim as i32,
25916 n_head as i32,
25917 n_head_kv as i32,
25918 t as i32,
25919 t_kv as i32,
25920 causal as i32,
25921 window as i32,
25922 );
25923 let __s_b = self.gpu.stream();
25924 let mut b = __s_b.launch_builder(&f);
25925 b.arg(q)
25926 .arg(k)
25927 .arg(v)
25928 .arg(o)
25929 .arg(&hd)
25930 .arg(&nh)
25931 .arg(&nhkv)
25932 .arg(&ti)
25933 .arg(&tkvi)
25934 .arg(&scale)
25935 .arg(&cz)
25936 .arg(&wi);
25937 unsafe {
25938 b.launch(cfg)?;
25939 }
25940 Ok(())
25941 }
25942
25943 #[allow(clippy::too_many_arguments)]
25959 pub fn sdpa_naive_w_lo(
25960 &self,
25961 q: &CudaSlice<f32>,
25962 k: &CudaSlice<f32>,
25963 v: &CudaSlice<f32>,
25964 o: &mut CudaSlice<f32>,
25965 head_dim: usize,
25966 n_head: usize,
25967 n_head_kv: usize,
25968 t: usize,
25969 t_kv: usize,
25970 scale: f32,
25971 causal: bool,
25972 window: usize,
25973 kv_floor: usize,
25974 ) -> Result<(), Box<dyn std::error::Error>> {
25975 if kv_floor > t_kv - t {
25976 return Err(format!(
25977 "sdpa_naive_w_lo: kv_floor {kv_floor} would clip the query rows themselves \
25978 (t_kv {t_kv} - t {t})"
25979 )
25980 .into());
25981 }
25982 let kv_lo = if window > 0 {
25983 (t_kv - t + 1).saturating_sub(window)
25984 } else {
25985 0
25986 }
25987 .max(kv_floor);
25988 let smem = (t_kv - kv_lo) * 4;
25989 if smem > 48 * 1024 {
25990 return Err(format!(
25991 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
25992 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
25993 a window this wide needs the multi-pass long-ctx kernel"
25994 )
25995 .into());
25996 }
25997 let f = self.func("sdpa_naive_w_lo_f32");
25998 let cfg = LaunchConfig {
25999 grid_dim: (n_head as u32, t as u32, 1),
26000 block_dim: (128, 1, 1),
26001 shared_mem_bytes: smem as u32,
26002 };
26003 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
26004 head_dim as i32,
26005 n_head as i32,
26006 n_head_kv as i32,
26007 t as i32,
26008 t_kv as i32,
26009 causal as i32,
26010 window as i32,
26011 kv_lo as i32,
26012 );
26013 let __s_b = self.gpu.stream();
26014 let mut b = __s_b.launch_builder(&f);
26015 b.arg(q)
26016 .arg(k)
26017 .arg(v)
26018 .arg(o)
26019 .arg(&hd)
26020 .arg(&nh)
26021 .arg(&nhkv)
26022 .arg(&ti)
26023 .arg(&tkvi)
26024 .arg(&scale)
26025 .arg(&cz)
26026 .arg(&wi)
26027 .arg(&lo);
26028 unsafe {
26029 b.launch(cfg)?;
26030 }
26031 Ok(())
26032 }
26033
26034 #[allow(clippy::too_many_arguments)] pub fn sdpa_naive_view(
26037 &self,
26038 q: &CudaSlice<f32>,
26039 k: &cudarc::driver::CudaView<f32>,
26040 v: &cudarc::driver::CudaView<f32>,
26041 o: &mut CudaSlice<f32>,
26042 head_dim: usize,
26043 n_head: usize,
26044 n_head_kv: usize,
26045 t: usize,
26046 t_kv: usize,
26047 scale: f32,
26048 causal: bool,
26049 ) -> Result<(), Box<dyn std::error::Error>> {
26050 let f = self.func("sdpa_naive_f32");
26051 let cfg = LaunchConfig {
26052 grid_dim: (n_head as u32, t as u32, 1),
26053 block_dim: (128, 1, 1),
26054 shared_mem_bytes: (t_kv * 4) as u32,
26055 };
26056 let (hd, nh, nhkv, ti, tkvi, cz) = (
26057 head_dim as i32,
26058 n_head as i32,
26059 n_head_kv as i32,
26060 t as i32,
26061 t_kv as i32,
26062 causal as i32,
26063 );
26064 let __s_b = self.gpu.stream();
26065 let mut b = __s_b.launch_builder(&f);
26066 b.arg(q)
26067 .arg(k)
26068 .arg(v)
26069 .arg(o)
26070 .arg(&hd)
26071 .arg(&nh)
26072 .arg(&nhkv)
26073 .arg(&ti)
26074 .arg(&tkvi)
26075 .arg(&scale)
26076 .arg(&cz);
26077 unsafe {
26078 b.launch(cfg)?;
26079 }
26080 Ok(())
26081 }
26082
26083 #[allow(clippy::too_many_arguments)]
26091 pub fn fa_dequant_kv_view_f32(
26092 &self,
26093 k: &cudarc::driver::CudaView<u8>,
26094 v: &cudarc::driver::CudaView<u8>,
26095 kf: &mut CudaSlice<f32>,
26096 vf: &mut CudaSlice<f32>,
26097 kv_dim_k: usize,
26098 kv_dim_v: usize,
26099 t_kv: usize,
26100 k_tok_bytes: usize,
26101 v_tok_bytes: usize,
26102 g: bool,
26103 ) -> Result<(), Box<dyn std::error::Error>> {
26104 let f = if g {
26105 self.func_g("fa_dequant_kv_ws_f32")
26106 } else {
26107 self.func("fa_dequant_kv_ws_f32")
26108 };
26109 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
26110 #[allow(clippy::manual_div_ceil)]
26111 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26113 let cfg = LaunchConfig {
26114 grid_dim: (nblk.max(1), 1, 1),
26115 block_dim: (256, 1, 1),
26116 shared_mem_bytes: 0,
26117 };
26118 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
26119 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26120 let __s_b = self.gpu.stream();
26121 let mut b = __s_b.launch_builder(&f);
26122 b.arg(k)
26123 .arg(v)
26124 .arg(&mut *kf)
26125 .arg(&mut *vf)
26126 .arg(&kdk)
26127 .arg(&kdv)
26128 .arg(&tkvi)
26129 .arg(&ktb)
26130 .arg(&vtb);
26131 unsafe {
26132 b.launch(cfg)?;
26133 }
26134 Ok(())
26135 }
26136
26137 #[allow(clippy::too_many_arguments)]
26138 pub fn sdpa_naive_quantized_view(
26139 &self,
26140 q: &CudaSlice<f32>,
26141 k: &cudarc::driver::CudaView<u8>,
26142 v: &cudarc::driver::CudaView<u8>,
26143 o: &mut CudaSlice<f32>,
26144 head_dim: usize,
26145 n_head: usize,
26146 n_head_kv: usize,
26147 t: usize,
26148 t_kv: usize,
26149 scale: f32,
26150 causal: bool,
26151 k_tok_bytes: usize,
26152 v_tok_bytes: usize,
26153 ) -> Result<(), Box<dyn std::error::Error>> {
26154 let kv_dim = n_head_kv * head_dim;
26155 let mut kf = self.uninit(t_kv * kv_dim)?;
26156 let mut vf = self.uninit(t_kv * kv_dim)?;
26157 let f = self.func("fa_dequant_kv_ws_f32");
26158 let total = (2 * t_kv * kv_dim) as u64;
26159 #[allow(clippy::manual_div_ceil)]
26160 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26162 let cfg = LaunchConfig {
26163 grid_dim: (nblk.max(1), 1, 1),
26164 block_dim: (256, 1, 1),
26165 shared_mem_bytes: 0,
26166 };
26167 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
26168 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
26169 let __s_b = self.gpu.stream();
26170 let mut b = __s_b.launch_builder(&f);
26171 b.arg(k)
26172 .arg(v)
26173 .arg(&mut kf)
26174 .arg(&mut vf)
26175 .arg(&kv_dim_i)
26176 .arg(&kv_dim_i)
26177 .arg(&t_kv_i)
26178 .arg(&k_tok_bytes_i)
26179 .arg(&v_tok_bytes_i);
26180 unsafe { b.launch(cfg)? };
26181 self.sdpa_naive(
26182 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26183 )
26184 }
26185
26186 #[allow(clippy::too_many_arguments)]
26198 pub fn sdpa_naive_w_quantized_view(
26199 &self,
26200 q: &CudaSlice<f32>,
26201 k: &cudarc::driver::CudaView<u8>,
26202 v: &cudarc::driver::CudaView<u8>,
26203 o: &mut CudaSlice<f32>,
26204 head_dim: usize,
26205 n_head: usize,
26206 n_head_kv: usize,
26207 t: usize,
26208 t_kv: usize,
26209 scale: f32,
26210 causal: bool,
26211 window: usize,
26212 k_tok_bytes: usize,
26213 v_tok_bytes: usize,
26214 ) -> Result<(), Box<dyn std::error::Error>> {
26215 let kv_dim = n_head_kv * head_dim;
26216 let mut kf = self.uninit(t_kv * kv_dim)?;
26217 let mut vf = self.uninit(t_kv * kv_dim)?;
26218 let f = self.func("fa_dequant_kv_ws_f32");
26219 let total = (2 * t_kv * kv_dim) as u64;
26220 #[allow(clippy::manual_div_ceil)]
26221 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26223 let cfg = LaunchConfig {
26224 grid_dim: (nblk.max(1), 1, 1),
26225 block_dim: (256, 1, 1),
26226 shared_mem_bytes: 0,
26227 };
26228 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
26229 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
26230 let __s_b = self.gpu.stream();
26231 let mut b = __s_b.launch_builder(&f);
26232 b.arg(k)
26233 .arg(v)
26234 .arg(&mut kf)
26235 .arg(&mut vf)
26236 .arg(&kv_dim_i)
26237 .arg(&kv_dim_i)
26238 .arg(&t_kv_i)
26239 .arg(&k_tok_bytes_i)
26240 .arg(&v_tok_bytes_i);
26241 unsafe { b.launch(cfg)? };
26242 self.sdpa_naive_w(
26243 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
26244 )
26245 }
26246
26247 #[allow(clippy::too_many_arguments)]
26251 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill(
26254 &self,
26255 q: &CudaSlice<f32>,
26256 k: &CudaSlice<f32>,
26257 v: &CudaSlice<f32>,
26258 o: &mut CudaSlice<f32>,
26259 head_dim: usize,
26260 n_head: usize,
26261 n_head_kv: usize,
26262 t: usize,
26263 t_kv: usize,
26264 scale: f32,
26265 causal: bool,
26266 ) -> Result<(), Box<dyn std::error::Error>> {
26267 if portable_mma_gated() {
26268 return self.sdpa_naive(
26269 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26270 );
26271 }
26272 let fa3_on = head_dim == 256
26280 && causal
26281 && t == t_kv
26282 && match std::env::var("MEMRA_FA3").as_deref() {
26283 Ok("0") => false,
26284 Ok("1") => {
26288 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
26289 true
26290 }
26291 _ => cfg!(memra_hopper_mma),
26292 };
26293 if fa3_on {
26294 let n = t * n_head * head_dim;
26295 let nkv = t * n_head_kv * head_dim;
26296 let mut q16 = self.alloc_u8_uninit(n * 2)?;
26297 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
26298 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
26299 self.f32_to_bf16_into(q, &mut q16, n)?;
26300 self.f32_to_bf16_into(k, &mut k16, nkv)?;
26301 self.f32_to_bf16_into(v, &mut v16, nkv)?;
26302 let rc = {
26303 use cudarc::driver::{DevicePtr, DevicePtrMut};
26304 let stream = self.gpu.stream();
26305 let (qp, _g1) = q16.device_ptr(&stream);
26306 let (kp, _g2) = k16.device_ptr(&stream);
26307 let (vp, _g3) = v16.device_ptr(&stream);
26308 let (op, _g4) = o.device_ptr_mut(&stream);
26309 unsafe {
26310 memra_fa3_prefill(
26311 qp as *const core::ffi::c_void,
26312 kp as *const core::ffi::c_void,
26313 vp as *const core::ffi::c_void,
26314 op as *mut f32,
26315 t as i32,
26316 n_head as i32,
26317 n_head_kv as i32,
26318 head_dim as i32,
26319 scale,
26320 stream.cu_stream() as *mut core::ffi::c_void,
26321 )
26322 }
26323 };
26324 if rc != 0 {
26325 return Err(format!("memra_fa3_prefill rc={rc}").into());
26326 }
26327 return Ok(());
26328 }
26329 const BK: usize = 32;
26337 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
26340 let (block_q, warps, w2_sfx): (usize, u32, &str) =
26341 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
26342 let hd_sfx = fa_hd_suffix(head_dim)?;
26346 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
26347 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
26352 let (kb16, vb16) = if bf16kv {
26353 let n = t_kv * n_head_kv * head_dim;
26354 let mut kb = self.alloc_u8_uninit(n * 2)?;
26355 let mut vb = self.alloc_u8_uninit(n * 2)?;
26356 let fcv = self.func("f32_to_bf16_bulk");
26357 let ni = n as i64;
26358 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26359 let __s_b = self.gpu.stream();
26360 let mut b = __s_b.launch_builder(&fcv);
26361 b.arg(k).arg(&mut kb).arg(&ni);
26362 unsafe {
26363 b.launch(cfgc)?;
26364 }
26365 let __s_b = self.gpu.stream();
26366 let mut b = __s_b.launch_builder(&fcv);
26367 b.arg(v).arg(&mut vb).arg(&ni);
26368 unsafe {
26369 b.launch(cfgc)?;
26370 }
26371 (Some(kb), Some(vb))
26372 } else {
26373 (None, None)
26374 };
26375 let f = self.func(&if bf16kv {
26376 format!("fa_prefill_bf16kv_pp{hd_sfx}")
26377 } else {
26378 format!(
26379 "fa_prefill_f32{}{}{hd_sfx}",
26380 if floor { "" } else { "_pp" },
26381 if floor { "" } else { w2_sfx }
26382 )
26383 });
26384 let kv_stages = if bf16kv { 2 } else { 1 };
26387 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
26388 + 4 * (block_q * BK + 2 * block_q)) as u32;
26389 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26390 f.set_attribute(
26391 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26392 shmem as i32,
26393 )?;
26394 let cfg = LaunchConfig {
26395 grid_dim: (
26396 (t as u32 + block_q as u32 - 1) / block_q as u32,
26397 n_head as u32,
26398 1,
26399 ),
26400 block_dim: (32, warps, 1),
26401 shared_mem_bytes: shmem,
26402 };
26403 let (hd, nh, nhkv, ti, tkvi, cz) = (
26404 head_dim as i32,
26405 n_head as i32,
26406 n_head_kv as i32,
26407 t as i32,
26408 t_kv as i32,
26409 causal as i32,
26410 );
26411 let __s_b = self.gpu.stream();
26412 let mut b = __s_b.launch_builder(&f);
26413 b.arg(q);
26414 match (&kb16, &vb16) {
26415 (Some(kb), Some(vb)) => {
26416 b.arg(kb).arg(vb);
26417 }
26418 _ => {
26419 b.arg(k).arg(v);
26420 }
26421 }
26422 b.arg(o)
26423 .arg(&hd)
26424 .arg(&nh)
26425 .arg(&nhkv)
26426 .arg(&ti)
26427 .arg(&tkvi)
26428 .arg(&scale)
26429 .arg(&cz);
26430 unsafe {
26431 b.launch(cfg)?;
26432 }
26433 Ok(())
26434 }
26435
26436 #[allow(clippy::too_many_arguments)]
26440 pub fn fa_prefill_w(
26441 &self,
26442 q: &CudaSlice<f32>,
26443 k: &CudaSlice<f32>,
26444 v: &CudaSlice<f32>,
26445 o: &mut CudaSlice<f32>,
26446 head_dim: usize,
26447 n_head: usize,
26448 n_head_kv: usize,
26449 t: usize,
26450 t_kv: usize,
26451 scale: f32,
26452 causal: bool,
26453 window: usize,
26454 ) -> Result<(), Box<dyn std::error::Error>> {
26455 if portable_mma_gated() {
26458 return self.sdpa_naive_w(
26459 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
26460 );
26461 }
26462 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26466 let faw_f32 =
26467 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
26468 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
26469 self.fa_prefill_w_arm(
26470 q,
26471 k,
26472 v,
26473 o,
26474 head_dim,
26475 n_head,
26476 n_head_kv,
26477 t,
26478 t_kv,
26479 scale,
26480 causal,
26481 window,
26482 floor || faw_f32,
26483 floor,
26484 )
26485 }
26486
26487 #[allow(clippy::too_many_arguments)]
26490 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_w_pre(
26492 &self,
26493 qb: &CudaSlice<u8>,
26494 kb: &CudaSlice<u8>,
26495 vb: &CudaSlice<u8>,
26496 o: &mut CudaSlice<f32>,
26497 head_dim: usize,
26498 n_head: usize,
26499 n_head_kv: usize,
26500 t: usize,
26501 t_kv: usize,
26502 scale: f32,
26503 causal: bool,
26504 window: usize,
26505 v_f16: bool,
26506 ) -> Result<(), Box<dyn std::error::Error>> {
26507 const BLOCK_Q: usize = 64;
26508 const BK: usize = 32;
26509 debug_assert_eq!(head_dim, 256);
26510 let hp = fa_f16pv_on()
26511 && faw_hp_on()
26512 && n_head.is_multiple_of(2)
26513 && (n_head / n_head_kv).is_multiple_of(2);
26514 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
26515 if hp {
26516 const BLOCK_QH: usize = 32;
26517 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
26520 let vh: &CudaSlice<u8> = if v_f16 {
26521 vb
26522 } else {
26523 let n = t_kv * n_head_kv * head_dim;
26524 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
26525 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
26526 }
26527 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
26528 vguard.as_ref().unwrap()
26529 };
26530 let f = self.func("fa_prefill_w_bf16_p1h2");
26531 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
26532 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26533 f.set_attribute(
26534 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26535 shmem as i32,
26536 )?;
26537 let cfg = LaunchConfig {
26538 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
26539 block_dim: (32, 4, 1),
26540 shared_mem_bytes: shmem,
26541 };
26542 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26543 head_dim as i32,
26544 n_head as i32,
26545 n_head_kv as i32,
26546 t as i32,
26547 t_kv as i32,
26548 causal as i32,
26549 window as i32,
26550 );
26551 let __s_b = self.gpu.stream();
26552 let mut b = __s_b.launch_builder(&f);
26553 b.arg(qb)
26554 .arg(kb)
26555 .arg(vh)
26556 .arg(o)
26557 .arg(&hd)
26558 .arg(&nh)
26559 .arg(&nhkv)
26560 .arg(&ti)
26561 .arg(&tkvi)
26562 .arg(&scale)
26563 .arg(&cz)
26564 .arg(&wi);
26565 unsafe {
26566 b.launch(cfg)?;
26567 }
26568 return Ok(());
26569 }
26570 let f = self.func("fa_prefill_w_bf16_p1");
26571 let shmem =
26572 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
26573 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26574 f.set_attribute(
26575 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26576 shmem as i32,
26577 )?;
26578 let cfg = LaunchConfig {
26579 grid_dim: (
26580 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
26581 n_head as u32,
26582 1,
26583 ),
26584 block_dim: (32, 4, 1),
26585 shared_mem_bytes: shmem,
26586 };
26587 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26588 head_dim as i32,
26589 n_head as i32,
26590 n_head_kv as i32,
26591 t as i32,
26592 t_kv as i32,
26593 causal as i32,
26594 window as i32,
26595 );
26596 let __s_b = self.gpu.stream();
26597 let mut b = __s_b.launch_builder(&f);
26598 b.arg(qb)
26599 .arg(kb)
26600 .arg(vb)
26601 .arg(o)
26602 .arg(&hd)
26603 .arg(&nh)
26604 .arg(&nhkv)
26605 .arg(&ti)
26606 .arg(&tkvi)
26607 .arg(&scale)
26608 .arg(&cz)
26609 .arg(&wi);
26610 unsafe {
26611 b.launch(cfg)?;
26612 }
26613 Ok(())
26614 }
26615
26616 #[allow(clippy::too_many_arguments)]
26618 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_w_arm(
26620 &self,
26621 q: &CudaSlice<f32>,
26622 k: &CudaSlice<f32>,
26623 v: &CudaSlice<f32>,
26624 o: &mut CudaSlice<f32>,
26625 head_dim: usize,
26626 n_head: usize,
26627 n_head_kv: usize,
26628 t: usize,
26629 t_kv: usize,
26630 scale: f32,
26631 causal: bool,
26632 window: usize,
26633 f32_stage: bool,
26634 floor: bool,
26635 ) -> Result<(), Box<dyn std::error::Error>> {
26636 const BLOCK_Q: usize = 64;
26637 const BK: usize = 32;
26638 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
26639 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26643 let p1 = !floor
26644 && !f32_stage
26645 && *P1_ON.get_or_init(|| {
26646 std::env::var("MEMRA_FAW_P1")
26647 .map(|v| v != "0")
26648 .unwrap_or(true)
26649 });
26650 let hp = p1
26651 && fa_f16pv_on()
26652 && faw_hp_on()
26653 && n_head.is_multiple_of(2)
26654 && (n_head / n_head_kv).is_multiple_of(2);
26655 if hp {
26656 const BLOCK_QH: usize = 32;
26657 let f = self.func("fa_prefill_w_bf16_p1h2");
26658 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
26659 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26660 f.set_attribute(
26661 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26662 shmem as i32,
26663 )?;
26664 let cfg = LaunchConfig {
26665 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
26666 block_dim: (32, 4, 1),
26667 shared_mem_bytes: shmem,
26668 };
26669 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26670 head_dim as i32,
26671 n_head as i32,
26672 n_head_kv as i32,
26673 t as i32,
26674 t_kv as i32,
26675 causal as i32,
26676 window as i32,
26677 );
26678 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
26679 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
26680 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
26681 let __s_b = self.gpu.stream();
26682 let mut b = __s_b.launch_builder(&f);
26683 b.arg(&qb)
26684 .arg(&kb)
26685 .arg(&vh)
26686 .arg(o)
26687 .arg(&hd)
26688 .arg(&nh)
26689 .arg(&nhkv)
26690 .arg(&ti)
26691 .arg(&tkvi)
26692 .arg(&scale)
26693 .arg(&cz)
26694 .arg(&wi);
26695 unsafe {
26696 b.launch(cfg)?;
26697 }
26698 return Ok(());
26699 }
26700 if p1 {
26701 let f = self.func("fa_prefill_w_bf16_p1");
26702 let shmem =
26703 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
26704 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26705 f.set_attribute(
26706 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26707 shmem as i32,
26708 )?;
26709 let cfg = LaunchConfig {
26710 grid_dim: (
26711 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
26712 n_head as u32,
26713 1,
26714 ),
26715 block_dim: (32, 4, 1),
26716 shared_mem_bytes: shmem,
26717 };
26718 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26719 head_dim as i32,
26720 n_head as i32,
26721 n_head_kv as i32,
26722 t as i32,
26723 t_kv as i32,
26724 causal as i32,
26725 window as i32,
26726 );
26727 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
26728 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
26729 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
26730 let __s_b = self.gpu.stream();
26731 let mut b = __s_b.launch_builder(&f);
26732 b.arg(&qb)
26733 .arg(&kb)
26734 .arg(&vb)
26735 .arg(o)
26736 .arg(&hd)
26737 .arg(&nh)
26738 .arg(&nhkv)
26739 .arg(&ti)
26740 .arg(&tkvi)
26741 .arg(&scale)
26742 .arg(&cz)
26743 .arg(&wi);
26744 unsafe {
26745 b.launch(cfg)?;
26746 }
26747 return Ok(());
26748 }
26749 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26752 let g4 = !floor
26753 && !f32_stage
26754 && n_head_kv == 1
26755 && n_head.is_multiple_of(4)
26756 && *G4_ON.get_or_init(|| {
26757 std::env::var("MEMRA_FAW_G4")
26758 .map(|v| v != "0")
26759 .unwrap_or(true)
26760 });
26761 if g4 {
26762 const SP_M: usize = 16;
26763 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26766 let o2 = *O2_ON.get_or_init(|| {
26767 std::env::var("MEMRA_FAW_O2")
26768 .map(|v| v != "0")
26769 .unwrap_or(true)
26770 });
26771 let f = self.func(if o2 {
26772 "fa_prefill_w_bf16_g4o2"
26773 } else {
26774 "fa_prefill_w_bf16_g4"
26775 });
26776 let shmem = if o2 {
26777 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
26778 } else {
26779 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
26780 as u32
26781 };
26782 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26783 f.set_attribute(
26784 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26785 shmem as i32,
26786 )?;
26787 let cfg = LaunchConfig {
26788 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
26789 block_dim: (32, 4, 1),
26790 shared_mem_bytes: shmem,
26791 };
26792 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26793 head_dim as i32,
26794 n_head as i32,
26795 n_head_kv as i32,
26796 t as i32,
26797 t_kv as i32,
26798 causal as i32,
26799 window as i32,
26800 );
26801 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
26802 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
26803 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
26804 let __s_b = self.gpu.stream();
26805 let mut b = __s_b.launch_builder(&f);
26806 b.arg(&qb)
26807 .arg(&kb)
26808 .arg(&vb)
26809 .arg(o)
26810 .arg(&hd)
26811 .arg(&nh)
26812 .arg(&nhkv)
26813 .arg(&ti)
26814 .arg(&tkvi)
26815 .arg(&scale)
26816 .arg(&cz)
26817 .arg(&wi);
26818 unsafe {
26819 b.launch(cfg)?;
26820 }
26821 return Ok(());
26822 }
26823 let f = self.func(if floor {
26824 "fa_prefill_w_f32"
26825 } else if f32_stage {
26826 "fa_prefill_w_f32_pp"
26827 } else {
26828 "fa_prefill_w_bf16_pp"
26829 });
26830 let shmem =
26831 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
26832 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26833 f.set_attribute(
26834 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26835 shmem as i32,
26836 )?;
26837 let cfg = LaunchConfig {
26838 grid_dim: (
26839 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
26840 n_head as u32,
26841 1,
26842 ),
26843 block_dim: (32, 4, 1),
26844 shared_mem_bytes: shmem,
26845 };
26846 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26847 head_dim as i32,
26848 n_head as i32,
26849 n_head_kv as i32,
26850 t as i32,
26851 t_kv as i32,
26852 causal as i32,
26853 window as i32,
26854 );
26855 if f32_stage {
26856 let __s_b = self.gpu.stream();
26857 let mut b = __s_b.launch_builder(&f);
26858 b.arg(q)
26859 .arg(k)
26860 .arg(v)
26861 .arg(o)
26862 .arg(&hd)
26863 .arg(&nh)
26864 .arg(&nhkv)
26865 .arg(&ti)
26866 .arg(&tkvi)
26867 .arg(&scale)
26868 .arg(&cz)
26869 .arg(&wi);
26870 unsafe {
26871 b.launch(cfg)?;
26872 }
26873 } else {
26874 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
26875 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
26876 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
26877 let __s_b = self.gpu.stream();
26878 let mut b = __s_b.launch_builder(&f);
26879 b.arg(&qb)
26880 .arg(&kb)
26881 .arg(&vb)
26882 .arg(o)
26883 .arg(&hd)
26884 .arg(&nh)
26885 .arg(&nhkv)
26886 .arg(&ti)
26887 .arg(&tkvi)
26888 .arg(&scale)
26889 .arg(&cz)
26890 .arg(&wi);
26891 unsafe {
26892 b.launch(cfg)?;
26893 }
26894 }
26895 Ok(())
26896 }
26897
26898 #[allow(clippy::too_many_arguments)]
26902 pub fn fa_prefill_hd512(
26903 &self,
26904 q: &CudaSlice<f32>,
26905 k: &CudaSlice<f32>,
26906 v: &CudaSlice<f32>,
26907 o: &mut CudaSlice<f32>,
26908 head_dim: usize,
26909 n_head: usize,
26910 n_head_kv: usize,
26911 t: usize,
26912 t_kv: usize,
26913 scale: f32,
26914 causal: bool,
26915 ) -> Result<(), Box<dyn std::error::Error>> {
26916 if portable_mma_gated() {
26918 return self.sdpa_naive(
26919 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26920 );
26921 }
26922 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26928 let f32_stage =
26929 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
26930 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26934 let sp = !f32_stage
26935 && *SP_ON.get_or_init(|| {
26936 std::env::var("MEMRA_FA512_SP")
26937 .map(|v| v != "0")
26938 .unwrap_or(true)
26939 });
26940 self.fa_prefill_hd512_arm(
26941 q,
26942 k,
26943 v,
26944 o,
26945 head_dim,
26946 n_head,
26947 n_head_kv,
26948 t,
26949 t_kv,
26950 scale,
26951 causal,
26952 f32_stage,
26953 sp,
26954 sp && fa_f16pv_on(),
26955 )
26956 }
26957
26958 #[allow(clippy::too_many_arguments)]
26960 pub fn fa_prefill_hd512_pre(
26961 &self,
26962 qb: &CudaSlice<u8>,
26963 kb: &CudaSlice<u8>,
26964 vb: &CudaSlice<u8>,
26965 o: &mut CudaSlice<f32>,
26966 head_dim: usize,
26967 n_head: usize,
26968 n_head_kv: usize,
26969 t: usize,
26970 t_kv: usize,
26971 scale: f32,
26972 causal: bool,
26973 v_f16: bool,
26974 ) -> Result<(), Box<dyn std::error::Error>> {
26975 debug_assert_eq!(head_dim, 512);
26976 const SP_M: usize = 16;
26977 const BKS: usize = 32;
26978 let f16pv = fa_f16pv_on();
26982 let nw = 2;
26983 let hp = f16pv
26984 && fa512_hp_on()
26985 && n_head.is_multiple_of(2)
26986 && (n_head / n_head_kv).is_multiple_of(2);
26987 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
26988 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
26989 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
26990 let n = t_kv * n_head_kv * head_dim;
26992 let need = n * 2;
26993 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
26994 *vguard = Some(self.alloc_uninit::<u8>(need)?);
26995 }
26996 let dst = vguard.as_mut().unwrap();
26997 self.bf16_to_f16_into(vb, n, dst)?;
26998 vguard.as_ref().unwrap()
26999 } else {
27000 vb
27001 };
27002 let f = self.func(if hp {
27003 "fa_prefill_bf16_hd512_sp16h2"
27004 } else {
27005 match (f16pv, nw) {
27006 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
27007 (true, _) => "fa_prefill_bf16_hd512_sp16",
27008 _ => "fa_prefill_bf16_hd512_sp",
27009 }
27010 });
27011 let (nwarp, npart) = if hp {
27012 (4usize, 4usize)
27013 } else if nw > 2 {
27014 (nw, nw)
27015 } else {
27016 (2, 1)
27017 };
27018 let shmem = if hp {
27020 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
27021 as u32
27022 } else {
27023 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
27024 + 4 * (npart * SP_M * BKS + SP_M)) as u32
27025 };
27026 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27027 f.set_attribute(
27028 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27029 shmem as i32,
27030 )?;
27031 let grid_y = if hp {
27032 (n_head / 2) as u32
27033 } else {
27034 n_head as u32
27035 };
27036 let cfg = LaunchConfig {
27037 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
27038 block_dim: (32, nwarp as u32, 1),
27039 shared_mem_bytes: shmem,
27040 };
27041 let (hd, nh, nhkv, ti, tkvi, cz) = (
27042 head_dim as i32,
27043 n_head as i32,
27044 n_head_kv as i32,
27045 t as i32,
27046 t_kv as i32,
27047 causal as i32,
27048 );
27049 let __s_b = self.gpu.stream();
27050 let mut b = __s_b.launch_builder(&f);
27051 b.arg(qb)
27052 .arg(kb)
27053 .arg(vref)
27054 .arg(o)
27055 .arg(&hd)
27056 .arg(&nh)
27057 .arg(&nhkv)
27058 .arg(&ti)
27059 .arg(&tkvi)
27060 .arg(&scale)
27061 .arg(&cz);
27062 unsafe {
27063 b.launch(cfg)?;
27064 }
27065 Ok(())
27066 }
27067
27068 #[allow(clippy::too_many_arguments)]
27075 pub fn mla_attn_gathered_tc(
27076 &self,
27077 q_lat_bf: &CudaSlice<u8>, cache_bf: &CudaSlice<u8>, idx: &CudaSlice<i32>, o_lat: &mut CudaSlice<f32>, n_head: usize,
27082 kv_rank: usize,
27083 t_q: usize,
27084 width: usize,
27085 scale: f32,
27086 ) -> Result<(), Box<dyn std::error::Error>> {
27087 if kv_rank != 512 {
27088 return Err(format!(
27089 "mla_attn_gathered_tc is stamped at kv_rank 512 (the glm5_next latent width); \
27090 got {kv_rank} — the caller's door must fall back to the f32 gathered kernel"
27091 )
27092 .into());
27093 }
27094 if t_q == 0 || n_head == 0 {
27095 return Ok(());
27096 }
27097 const SP_M: usize = 16;
27098 const BKS: usize = 32;
27099 const HD: usize = 512;
27100 let f = self.func("fa_mla_gathered_bf16");
27101 let shmem =
27103 (2 * (SP_M * HD + BKS * HD + SP_M * BKS) + 4 * (SP_M * BKS + SP_M) + 4 * BKS) as u32;
27104 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27105 f.set_attribute(
27106 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27107 shmem as i32,
27108 )?;
27109 let cfg = LaunchConfig {
27110 grid_dim: (t_q as u32, (n_head as u32).div_ceil(SP_M as u32), 1),
27111 block_dim: (32, 2, 1),
27112 shared_mem_bytes: shmem,
27113 };
27114 let (nh, tq, w) = (n_head as i32, t_q as i32, width as i32);
27115 let __s_b = self.gpu.stream();
27116 let mut b = __s_b.launch_builder(&f);
27117 b.arg(q_lat_bf)
27118 .arg(cache_bf)
27119 .arg(idx)
27120 .arg(o_lat)
27121 .arg(&nh)
27122 .arg(&tq)
27123 .arg(&w)
27124 .arg(&scale);
27125 unsafe {
27126 b.launch(cfg)?;
27127 }
27128 Ok(())
27129 }
27130
27131 #[allow(clippy::too_many_arguments)]
27134 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_hd512_arm(
27136 &self,
27137 q: &CudaSlice<f32>,
27138 k: &CudaSlice<f32>,
27139 v: &CudaSlice<f32>,
27140 o: &mut CudaSlice<f32>,
27141 head_dim: usize,
27142 n_head: usize,
27143 n_head_kv: usize,
27144 t: usize,
27145 t_kv: usize,
27146 scale: f32,
27147 causal: bool,
27148 f32_stage: bool,
27149 sp: bool,
27150 f16pv: bool,
27151 ) -> Result<(), Box<dyn std::error::Error>> {
27152 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
27153 if sp && !f32_stage {
27154 const SP_M: usize = 16;
27158 const BKS: usize = 32;
27159 let nw = 2;
27160 let hp = f16pv
27161 && fa512_hp_on()
27162 && n_head.is_multiple_of(2)
27163 && (n_head / n_head_kv).is_multiple_of(2);
27164 let f = self.func(if hp {
27165 "fa_prefill_bf16_hd512_sp16h2"
27166 } else {
27167 match (f16pv, nw) {
27168 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
27169 (true, _) => "fa_prefill_bf16_hd512_sp16",
27170 _ => "fa_prefill_bf16_hd512_sp",
27171 }
27172 });
27173 let (nwarp, npart) = if hp {
27174 (4usize, 4usize)
27175 } else if nw > 2 {
27176 (nw, nw)
27177 } else {
27178 (2, 1)
27179 };
27180 let shmem = if hp {
27181 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
27182 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
27183 } else {
27184 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
27185 + 4 * (npart * SP_M * BKS + SP_M)) as u32
27186 };
27187 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27188 f.set_attribute(
27189 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27190 shmem as i32,
27191 )?;
27192 let grid_y = if hp {
27193 (n_head / 2) as u32
27194 } else {
27195 n_head as u32
27196 };
27197 let cfg = LaunchConfig {
27198 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
27199 block_dim: (32, nwarp as u32, 1),
27200 shared_mem_bytes: shmem,
27201 };
27202 let (hd, nh, nhkv, ti, tkvi, cz) = (
27203 head_dim as i32,
27204 n_head as i32,
27205 n_head_kv as i32,
27206 t as i32,
27207 t_kv as i32,
27208 causal as i32,
27209 );
27210 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27211 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27212 let vb = if f16pv {
27213 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
27214 } else {
27215 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
27216 };
27217 let __s_b = self.gpu.stream();
27218 let mut b = __s_b.launch_builder(&f);
27219 b.arg(&qb)
27220 .arg(&kb)
27221 .arg(&vb)
27222 .arg(o)
27223 .arg(&hd)
27224 .arg(&nh)
27225 .arg(&nhkv)
27226 .arg(&ti)
27227 .arg(&tkvi)
27228 .arg(&scale)
27229 .arg(&cz);
27230 unsafe {
27231 b.launch(cfg)?;
27232 }
27233 return Ok(());
27234 }
27235 const BLOCK_Q: usize = 32;
27236 const BK: usize = 32;
27237 const HALF: usize = 256;
27238 let f = self.func(if f32_stage {
27239 "fa_prefill_f32_hd512"
27240 } else {
27241 "fa_prefill_bf16_hd512"
27242 });
27243 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
27245 + 4 * BLOCK_Q) as u32;
27246 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27247 f.set_attribute(
27248 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27249 shmem as i32,
27250 )?;
27251 let cfg = LaunchConfig {
27252 grid_dim: (
27253 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27254 n_head as u32,
27255 2,
27256 ),
27257 block_dim: (32, 2, 1),
27258 shared_mem_bytes: shmem,
27259 };
27260 let (hd, nh, nhkv, ti, tkvi, cz) = (
27261 head_dim as i32,
27262 n_head as i32,
27263 n_head_kv as i32,
27264 t as i32,
27265 t_kv as i32,
27266 causal as i32,
27267 );
27268 if f32_stage {
27269 let __s_b = self.gpu.stream();
27270 let mut b = __s_b.launch_builder(&f);
27271 b.arg(q)
27272 .arg(k)
27273 .arg(v)
27274 .arg(o)
27275 .arg(&hd)
27276 .arg(&nh)
27277 .arg(&nhkv)
27278 .arg(&ti)
27279 .arg(&tkvi)
27280 .arg(&scale)
27281 .arg(&cz);
27282 unsafe {
27283 b.launch(cfg)?;
27284 }
27285 } else {
27286 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27287 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27288 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27289 let __s_b = self.gpu.stream();
27290 let mut b = __s_b.launch_builder(&f);
27291 b.arg(&qb)
27292 .arg(&kb)
27293 .arg(&vb)
27294 .arg(o)
27295 .arg(&hd)
27296 .arg(&nh)
27297 .arg(&nhkv)
27298 .arg(&ti)
27299 .arg(&tkvi)
27300 .arg(&scale)
27301 .arg(&cz);
27302 unsafe {
27303 b.launch(cfg)?;
27304 }
27305 }
27306 Ok(())
27307 }
27308
27309 #[allow(clippy::too_many_arguments)]
27313 pub fn rope_neox2_bf16e(
27314 &self,
27315 q: &mut CudaSlice<f32>,
27316 k: &mut CudaSlice<f32>,
27317 qb: &mut CudaSlice<u8>,
27318 kb: &mut CudaSlice<u8>,
27319 pos: &CudaSlice<i32>,
27320 head_dim: usize,
27321 n_dims: usize,
27322 nh_q: usize,
27323 nh_k: usize,
27324 n_tokens: usize,
27325 base: f32,
27326 freq_scale: f32,
27327 ff: Option<&CudaSlice<f32>>,
27328 ) -> Result<(), Box<dyn std::error::Error>> {
27329 let f = self.func("rope_neox2_bf16e_f32");
27330 let rows = ((nh_q + nh_k) * n_tokens) as u32;
27331 let cfg = LaunchConfig {
27332 grid_dim: (rows, 1, 1),
27333 block_dim: ((head_dim / 2) as u32, 1, 1),
27334 shared_mem_bytes: 0,
27335 };
27336 let theta_scale = base.powf(-2.0 / n_dims as f32);
27337 let (hd, nd, nhq, nhk, nt) = (
27338 head_dim as i32,
27339 n_dims as i32,
27340 nh_q as i32,
27341 nh_k as i32,
27342 n_tokens as i32,
27343 );
27344 let __s_b = self.gpu.stream();
27345 let mut b = __s_b.launch_builder(&f);
27346 match ff {
27347 Some(t) => {
27348 b.arg(&mut *q)
27349 .arg(&mut *k)
27350 .arg(&mut *qb)
27351 .arg(&mut *kb)
27352 .arg(pos)
27353 .arg(&hd)
27354 .arg(&nd)
27355 .arg(&nhq)
27356 .arg(&nhk)
27357 .arg(&nt)
27358 .arg(&theta_scale)
27359 .arg(&freq_scale)
27360 .arg(t);
27361 unsafe {
27362 b.launch(cfg)?;
27363 }
27364 }
27365 None => {
27366 let null: u64 = 0;
27367 b.arg(&mut *q)
27368 .arg(&mut *k)
27369 .arg(&mut *qb)
27370 .arg(&mut *kb)
27371 .arg(pos)
27372 .arg(&hd)
27373 .arg(&nd)
27374 .arg(&nhq)
27375 .arg(&nhk)
27376 .arg(&nt)
27377 .arg(&theta_scale)
27378 .arg(&freq_scale)
27379 .arg(&null);
27380 unsafe {
27381 b.launch(cfg)?;
27382 }
27383 }
27384 }
27385 Ok(())
27386 }
27387
27388 pub fn f32_to_bf16(
27391 &self,
27392 x: &CudaSlice<f32>,
27393 n: usize,
27394 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27395 assert!(
27396 n.is_multiple_of(4),
27397 "f32_to_bf16 requires n % 4 == 0, got {n}"
27398 );
27399 let mut y = self.alloc_uninit::<u8>(n * 2)?;
27400 let f = self.func("f32_to_bf16_flat");
27401 let n_i = n as i64;
27402 let cfg = LaunchConfig {
27403 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
27404 block_dim: (256, 1, 1),
27405 shared_mem_bytes: 0,
27406 };
27407 let __s_b = self.gpu.stream();
27408 let mut b = __s_b.launch_builder(&f);
27409 b.arg(x).arg(&mut y).arg(&n_i);
27410 unsafe {
27411 b.launch(cfg)?;
27412 }
27413 Ok(y)
27414 }
27415
27416 pub fn f32_to_f16(
27417 &self,
27418 x: &CudaSlice<f32>,
27419 n: usize,
27420 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27421 assert!(
27422 n.is_multiple_of(4),
27423 "f32_to_f16 requires n % 4 == 0, got {n}"
27424 );
27425 let mut y = self.alloc_uninit::<u8>(n * 2)?;
27426 let f = self.func("f32_to_f16_flat");
27427 let n_i = n as i64;
27428 let cfg = LaunchConfig {
27429 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
27430 block_dim: (256, 1, 1),
27431 shared_mem_bytes: 0,
27432 };
27433 let __s_b = self.gpu.stream();
27434 let mut b = __s_b.launch_builder(&f);
27435 b.arg(x).arg(&mut y).arg(&n_i);
27436 unsafe {
27437 b.launch(cfg)?;
27438 }
27439 Ok(y)
27440 }
27441
27442 pub fn bf16_to_f16(
27444 &self,
27445 xb: &CudaSlice<u8>,
27446 n: usize,
27447 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27448 let mut y = self.alloc_uninit::<u8>(n * 2)?;
27449 self.bf16_to_f16_into(xb, n, &mut y)?;
27450 Ok(y)
27451 }
27452
27453 pub fn bf16_to_f16_into(
27455 &self,
27456 xb: &CudaSlice<u8>,
27457 n: usize,
27458 y: &mut CudaSlice<u8>,
27459 ) -> Result<(), Box<dyn std::error::Error>> {
27460 assert!(
27461 n.is_multiple_of(2),
27462 "bf16_to_f16 requires n % 2 == 0, got {n}"
27463 );
27464 assert!(y.len() >= n * 2);
27465 let f = self.func("bf16_to_f16_flat");
27466 let n2 = (n / 2) as i64;
27467 let cfg = LaunchConfig {
27468 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
27469 block_dim: (256, 1, 1),
27470 shared_mem_bytes: 0,
27471 };
27472 let __s_b = self.gpu.stream();
27473 let mut b = __s_b.launch_builder(&f);
27474 b.arg(xb).arg(y).arg(&n2);
27475 unsafe {
27476 b.launch(cfg)?;
27477 }
27478 Ok(())
27479 }
27480
27481 #[allow(clippy::too_many_arguments)]
27486 pub fn fa_prefill_vl8(
27487 &self,
27488 seqs: &[FaSeqVl],
27489 head_dim: usize,
27490 n_head: usize,
27491 n_head_kv: usize,
27492 scale: f32,
27493 ) -> Result<(), Box<dyn std::error::Error>> {
27494 const BK: usize = 32;
27495 let b = seqs.len();
27496 assert!((1..=8).contains(&b));
27497 let mut packed = [FaSeqVl::default(); 8];
27498 packed[..b].copy_from_slice(seqs);
27499 let v = FaVl8(packed);
27500 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
27501 let ept = (n_head_kv * head_dim) as i32;
27502 {
27503 let f = self.func("fa_mirror_vl");
27504 let max_n = (max_t as i64) * ept as i64;
27505 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
27506 for which in 0..2i32 {
27507 let cfg = LaunchConfig {
27508 grid_dim: (blocks, 1, b as u32),
27509 block_dim: (256, 1, 1),
27510 shared_mem_bytes: 0,
27511 };
27512 let __s_lb = self.gpu.stream();
27513 let mut lb = __s_lb.launch_builder(&f);
27514 lb.arg(&v).arg(&ept).arg(&which);
27515 unsafe {
27516 lb.launch(cfg)?;
27517 }
27518 }
27519 }
27520 let hd_sfx = fa_hd_suffix(head_dim)?;
27521 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
27522 let block_q = 64usize;
27523 let kv_stages = 2usize;
27524 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
27525 + 4 * (block_q * BK + 2 * block_q)) as u32;
27526 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27527 f.set_attribute(
27528 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27529 shmem as i32,
27530 )?;
27531 let cfg = LaunchConfig {
27532 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
27533 block_dim: (32, 4, 1),
27534 shared_mem_bytes: shmem,
27535 };
27536 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27537 let __s_lb = self.gpu.stream();
27538 let mut lb = __s_lb.launch_builder(&f);
27539 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
27540 unsafe {
27541 lb.launch(cfg)?;
27542 }
27543 Ok(())
27544 }
27545
27546 #[allow(clippy::too_many_arguments)]
27550 pub fn attn_pre_vl8(
27551 &self,
27552 seqs: &[AttnPreVl],
27553 wq: &CudaSlice<f32>,
27554 wk: &CudaSlice<f32>,
27555 head_dim: usize,
27556 rope_dims: usize,
27557 n_head: usize,
27558 n_head_kv: usize,
27559 eps: f32,
27560 freq_base: f32,
27561 freq_scale: f32,
27562 kv_dim_k: usize,
27563 kv_dim_v: usize,
27564 k_tok_bytes: usize,
27565 v_tok_bytes: usize,
27566 ) -> Result<(), Box<dyn std::error::Error>> {
27567 let b = seqs.len();
27568 assert!((1..=8).contains(&b));
27569 let mut packed = [AttnPreVl::default(); 8];
27570 packed[..b].copy_from_slice(seqs);
27571 let v = AttnPreVl8(packed);
27572 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
27573 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27574 {
27575 let f = self.func("q_gate_split_vl");
27576 let n = max_t * (n_head * head_dim) as u32;
27577 let cfg = LaunchConfig {
27578 grid_dim: (n.div_ceil(256), 1, b as u32),
27579 block_dim: (256, 1, 1),
27580 shared_mem_bytes: 0,
27581 };
27582 let __s_lb = self.gpu.stream();
27583 let mut lb = __s_lb.launch_builder(&f);
27584 lb.arg(&v).arg(&hd).arg(&nh);
27585 unsafe {
27586 lb.launch(cfg)?;
27587 }
27588 }
27589 {
27590 let f = self.func("attn_rms_vl");
27591 let cfg = LaunchConfig {
27592 grid_dim: (max_t * n_head as u32, 2, b as u32),
27593 block_dim: (rms_block(), 1, 1),
27594 shared_mem_bytes: 0,
27595 };
27596 let __s_lb = self.gpu.stream();
27597 let mut lb = __s_lb.launch_builder(&f);
27598 lb.arg(&v)
27599 .arg(wq)
27600 .arg(wk)
27601 .arg(&hd)
27602 .arg(&nh)
27603 .arg(&nhkv)
27604 .arg(&eps);
27605 unsafe {
27606 lb.launch(cfg)?;
27607 }
27608 }
27609 {
27610 let f = self.func("attn_rope_vl");
27611 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
27612 let nd = rope_dims as i32;
27613 let cfg = LaunchConfig {
27614 grid_dim: (max_t * n_head as u32, 2, b as u32),
27615 block_dim: ((head_dim / 2) as u32, 1, 1),
27616 shared_mem_bytes: 0,
27617 };
27618 let __s_lb = self.gpu.stream();
27619 let mut lb = __s_lb.launch_builder(&f);
27620 lb.arg(&v)
27621 .arg(&hd)
27622 .arg(&nd)
27623 .arg(&nh)
27624 .arg(&nhkv)
27625 .arg(&theta_scale)
27626 .arg(&freq_scale);
27627 unsafe {
27628 lb.launch(cfg)?;
27629 }
27630 }
27631 {
27632 let f = self.func("append_kv_vl");
27633 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
27634 let cfg = LaunchConfig {
27635 grid_dim: (nblk, max_t, b as u32),
27636 block_dim: (32, 1, 1),
27637 shared_mem_bytes: 0,
27638 };
27639 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
27640 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27641 let __s_lb = self.gpu.stream();
27642 let mut lb = __s_lb.launch_builder(&f);
27643 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
27644 unsafe {
27645 lb.launch(cfg)?;
27646 }
27647 }
27648 Ok(())
27649 }
27650
27651 #[allow(clippy::too_many_arguments)]
27656 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_view(
27659 &self,
27660 q: &CudaSlice<f32>,
27661 k: &cudarc::driver::CudaView<u8>,
27662 v: &cudarc::driver::CudaView<u8>,
27663 o: &mut CudaSlice<f32>,
27664 head_dim: usize,
27665 n_head: usize,
27666 n_head_kv: usize,
27667 t: usize,
27668 t_kv: usize,
27669 scale: f32,
27670 causal: bool,
27671 k_tok_bytes: usize,
27672 v_tok_bytes: usize,
27673 g: bool,
27674 ) -> Result<(), Box<dyn std::error::Error>> {
27675 if portable_mma_gated() {
27676 return self.sdpa_naive_quantized_view(
27677 q,
27678 k,
27679 v,
27680 o,
27681 head_dim,
27682 n_head,
27683 n_head_kv,
27684 t,
27685 t_kv,
27686 scale,
27687 causal,
27688 k_tok_bytes,
27689 v_tok_bytes,
27690 );
27691 }
27692 const BLOCK_Q: usize = 64;
27693 const BK: usize = 32;
27694 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
27697 let f = if g {
27698 self.func_g(&name)
27699 } else {
27700 self.func(&name)
27701 };
27702 let shmem =
27703 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
27704 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27705 f.set_attribute(
27706 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27707 shmem as i32,
27708 )?;
27709 let cfg = LaunchConfig {
27710 grid_dim: (
27711 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27712 n_head as u32,
27713 1,
27714 ),
27715 block_dim: (32, 4, 1),
27716 shared_mem_bytes: shmem,
27717 };
27718 let (hd, nh, nhkv, ti, tkvi, cz) = (
27719 head_dim as i32,
27720 n_head as i32,
27721 n_head_kv as i32,
27722 t as i32,
27723 t_kv as i32,
27724 causal as i32,
27725 );
27726 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27727 let __s_b = self.gpu.stream();
27728 let mut b = __s_b.launch_builder(&f);
27729 b.arg(q)
27730 .arg(k)
27731 .arg(v)
27732 .arg(o)
27733 .arg(&hd)
27734 .arg(&nh)
27735 .arg(&nhkv)
27736 .arg(&ti)
27737 .arg(&tkvi)
27738 .arg(&scale)
27739 .arg(&cz)
27740 .arg(&ktb)
27741 .arg(&vtb);
27742 unsafe {
27743 b.launch(cfg)?;
27744 }
27745 Ok(())
27746 }
27747
27748 #[allow(clippy::too_many_arguments)]
27758 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_view_ws(
27760 &self,
27761 q: &CudaSlice<f32>,
27762 k: &cudarc::driver::CudaView<u8>,
27763 v: &cudarc::driver::CudaView<u8>,
27764 o: &mut CudaSlice<f32>,
27765 head_dim: usize,
27766 n_head: usize,
27767 n_head_kv: usize,
27768 t: usize,
27769 t_kv: usize,
27770 scale: f32,
27771 causal: bool,
27772 k_tok_bytes: usize,
27773 v_tok_bytes: usize,
27774 g: bool,
27775 ) -> Result<(), Box<dyn std::error::Error>> {
27776 if portable_mma_gated() {
27777 return self.sdpa_naive_quantized_view(
27778 q,
27779 k,
27780 v,
27781 o,
27782 head_dim,
27783 n_head,
27784 n_head_kv,
27785 t,
27786 t_kv,
27787 scale,
27788 causal,
27789 k_tok_bytes,
27790 v_tok_bytes,
27791 );
27792 }
27793 const BLOCK_Q: usize = 64;
27794 const BK: usize = 32;
27795 let kv_dim_k = n_head_kv * head_dim;
27796 let kv_dim_v = n_head_kv * head_dim;
27797 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
27799 let mut guard = self.prime_deqw_ws.lock().unwrap();
27801 let need_grow = match guard.as_ref() {
27802 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
27803 None => true,
27804 };
27805 if need_grow {
27806 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
27807 let (ck, cv) = guard
27808 .as_ref()
27809 .map(|(a, b)| (a.len(), b.len()))
27810 .unwrap_or((0, 0));
27811 *guard = Some((
27812 self.alloc_u8(grow(ck, k_ws_bytes))?,
27813 self.alloc_u8(grow(cv, v_ws_bytes))?,
27814 ));
27815 }
27816 let (kw, vw) = guard.as_mut().unwrap();
27817 {
27819 let f = if g {
27821 self.func_g("fa_dequant_kv_ws_bf16")
27822 } else {
27823 self.func("fa_dequant_kv_ws_bf16")
27824 };
27825 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
27826 #[allow(clippy::manual_div_ceil)]
27827 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
27829 let cfg = LaunchConfig {
27830 grid_dim: (nblk.max(1), 1, 1),
27831 block_dim: (256, 1, 1),
27832 shared_mem_bytes: 0,
27833 };
27834 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
27835 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
27836 let __s_b = self.gpu.stream();
27837 let mut b = __s_b.launch_builder(&f);
27838 b.arg(k)
27839 .arg(v)
27840 .arg(&mut *kw)
27841 .arg(&mut *vw)
27842 .arg(&kdk)
27843 .arg(&kdv)
27844 .arg(&tkvi)
27845 .arg(&ktb)
27846 .arg(&vtb);
27847 unsafe {
27848 b.launch(cfg)?;
27849 }
27850 }
27851 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
27859 .map(|v| v != "0")
27860 .unwrap_or(true);
27861 {
27862 let hd_sfx = fa_hd_suffix(head_dim)?;
27863 let f = self.func(&format!(
27864 "fa_prefill_qw{}{hd_sfx}",
27865 if db { "_db" } else { "" }
27866 ));
27867 let shmem = if db {
27868 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
27870 } else {
27871 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
27872 };
27873 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27874 f.set_attribute(
27875 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27876 shmem as i32,
27877 )?;
27878 let cfg = LaunchConfig {
27879 grid_dim: (
27880 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27881 n_head as u32,
27882 1,
27883 ),
27884 block_dim: (32, 4, 1),
27885 shared_mem_bytes: shmem,
27886 };
27887 let (hd, nh, nhkv, ti, tkvi, cz) = (
27888 head_dim as i32,
27889 n_head as i32,
27890 n_head_kv as i32,
27891 t as i32,
27892 t_kv as i32,
27893 causal as i32,
27894 );
27895 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
27896 let __s_b = self.gpu.stream();
27897 let mut b = __s_b.launch_builder(&f);
27898 b.arg(q)
27899 .arg(&*kw)
27900 .arg(&*vw)
27901 .arg(o)
27902 .arg(&hd)
27903 .arg(&nh)
27904 .arg(&nhkv)
27905 .arg(&ti)
27906 .arg(&tkvi)
27907 .arg(&scale)
27908 .arg(&cz)
27909 .arg(&kdk)
27910 .arg(&kdv);
27911 unsafe {
27912 b.launch(cfg)?;
27913 }
27914 }
27915 Ok(())
27916 }
27917
27918 #[allow(clippy::too_many_arguments)]
27934 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_view_ws_w_hd128(
27936 &self,
27937 q: &CudaSlice<f32>,
27938 k: &cudarc::driver::CudaView<u8>,
27939 v: &cudarc::driver::CudaView<u8>,
27940 o: &mut CudaSlice<f32>,
27941 head_dim: usize,
27942 n_head: usize,
27943 n_head_kv: usize,
27944 t: usize,
27945 t_kv: usize,
27946 scale: f32,
27947 causal: bool,
27948 window: usize,
27949 k_tok_bytes: usize,
27950 v_tok_bytes: usize,
27951 ) -> Result<(), Box<dyn std::error::Error>> {
27952 assert_eq!(
27953 head_dim, 128,
27954 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
27955 );
27956 if portable_mma_gated() {
27957 return self.sdpa_naive_w_quantized_view(
27958 q,
27959 k,
27960 v,
27961 o,
27962 head_dim,
27963 n_head,
27964 n_head_kv,
27965 t,
27966 t_kv,
27967 scale,
27968 causal,
27969 window,
27970 k_tok_bytes,
27971 v_tok_bytes,
27972 );
27973 }
27974 const BLOCK_Q: usize = 64;
27975 const BK: usize = 32;
27976 let kv_dim_k = n_head_kv * head_dim;
27977 let kv_dim_v = n_head_kv * head_dim;
27978 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
27980 let mut guard = self.prime_deqw_ws.lock().unwrap();
27981 let need_grow = match guard.as_ref() {
27982 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
27983 None => true,
27984 };
27985 if need_grow {
27986 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
27987 let (ck, cv) = guard
27988 .as_ref()
27989 .map(|(a, b)| (a.len(), b.len()))
27990 .unwrap_or((0, 0));
27991 *guard = Some((
27992 self.alloc_u8(grow(ck, k_ws_bytes))?,
27993 self.alloc_u8(grow(cv, v_ws_bytes))?,
27994 ));
27995 }
27996 let (kw, vw) = guard.as_mut().unwrap();
27997 {
28000 let f = self.func("fa_dequant_kv_ws_bf16");
28001 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
28002 #[allow(clippy::manual_div_ceil)]
28003 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
28005 let cfg = LaunchConfig {
28006 grid_dim: (nblk.max(1), 1, 1),
28007 block_dim: (256, 1, 1),
28008 shared_mem_bytes: 0,
28009 };
28010 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
28011 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28012 let __s_b = self.gpu.stream();
28013 let mut b = __s_b.launch_builder(&f);
28014 b.arg(k)
28015 .arg(v)
28016 .arg(&mut *kw)
28017 .arg(&mut *vw)
28018 .arg(&kdk)
28019 .arg(&kdv)
28020 .arg(&tkvi)
28021 .arg(&ktb)
28022 .arg(&vtb);
28023 unsafe {
28024 b.launch(cfg)?;
28025 }
28026 }
28027 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
28029 .map(|v| v != "0")
28030 .unwrap_or(true);
28031 {
28032 let f = self.func(if db {
28033 "fa_prefill_qw_db_w_hd128"
28034 } else {
28035 "fa_prefill_qw_w_hd128"
28036 });
28037 let shmem = if db {
28038 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
28039 } else {
28040 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
28041 };
28042 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28043 f.set_attribute(
28044 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28045 shmem as i32,
28046 )?;
28047 let cfg = LaunchConfig {
28048 grid_dim: (
28049 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
28050 n_head as u32,
28051 1,
28052 ),
28053 block_dim: (32, 4, 1),
28054 shared_mem_bytes: shmem,
28055 };
28056 let (hd, nh, nhkv, ti, tkvi, cz) = (
28057 head_dim as i32,
28058 n_head as i32,
28059 n_head_kv as i32,
28060 t as i32,
28061 t_kv as i32,
28062 causal as i32,
28063 );
28064 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
28065 let __s_b = self.gpu.stream();
28066 let mut b = __s_b.launch_builder(&f);
28067 b.arg(q)
28068 .arg(&*kw)
28069 .arg(&*vw)
28070 .arg(o)
28071 .arg(&hd)
28072 .arg(&nh)
28073 .arg(&nhkv)
28074 .arg(&ti)
28075 .arg(&tkvi)
28076 .arg(&scale)
28077 .arg(&cz)
28078 .arg(&kdk)
28079 .arg(&kdv)
28080 .arg(&wnd);
28081 unsafe {
28082 b.launch(cfg)?;
28083 }
28084 }
28085 Ok(())
28086 }
28087
28088 #[allow(clippy::too_many_arguments)] pub fn fa_decode(
28093 &self,
28094 q: &CudaSlice<f32>,
28095 k: &cudarc::driver::CudaView<u8>,
28096 v: &cudarc::driver::CudaView<u8>,
28097 o: &mut CudaSlice<f32>,
28098 head_dim: usize,
28099 n_head: usize,
28100 n_head_kv: usize,
28101 t_kv: usize,
28102 scale: f32,
28103 k_tok_bytes: usize,
28104 v_tok_bytes: usize,
28105 ) -> Result<(), Box<dyn std::error::Error>> {
28106 self.fa_decode_kvmod(
28107 q,
28108 k,
28109 v,
28110 o,
28111 head_dim,
28112 n_head,
28113 n_head_kv,
28114 t_kv,
28115 scale,
28116 k_tok_bytes,
28117 v_tok_bytes,
28118 false,
28119 )
28120 }
28121
28122 #[allow(clippy::too_many_arguments)]
28126 #[allow(clippy::too_many_arguments)]
28130 #[allow(clippy::too_many_arguments)]
28131 fn fa_decode_scalar_unified(
28132 &self,
28133 q: &cudarc::driver::CudaView<f32>,
28134 k: &cudarc::driver::CudaView<u8>,
28135 v: &cudarc::driver::CudaView<u8>,
28136 o: &mut cudarc::driver::CudaViewMut<f32>,
28137 head_dim: usize,
28138 n_head: usize,
28139 n_head_kv: usize,
28140 t_kv_host: usize,
28141 t_kv_dev: Option<&CudaSlice<i32>>,
28142 scale: f32,
28143 n_splits: usize,
28144 split_keys: usize,
28145 k_tok_bytes: usize,
28146 v_tok_bytes: usize,
28147 g: bool,
28148 part_o: &mut CudaSlice<f32>,
28149 part_m: &mut CudaSlice<f32>,
28150 part_l: &mut CudaSlice<f32>,
28151 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
28152 ) -> Result<(), Box<dyn std::error::Error>> {
28153 let f = if g {
28154 self.func_g("fa_decode_f32")
28155 } else {
28156 self.fa_func("fa_decode_f32", head_dim)
28157 };
28158 let cfg = LaunchConfig {
28159 grid_dim: (n_head as u32, n_splits as u32, 1),
28160 block_dim: (head_dim as u32, 1, 1),
28161 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
28162 };
28163 let (hd, nh, nhkv, nsp) = (
28164 head_dim as i32,
28165 n_head as i32,
28166 n_head_kv as i32,
28167 n_splits as i32,
28168 );
28169 let (ktb, vtb, tkvi, ski) = (
28170 k_tok_bytes as i64,
28171 v_tok_bytes as i64,
28172 t_kv_host as i32,
28173 split_keys as i32,
28174 );
28175 let __s_b = self.gpu.stream();
28176 let mut b = __s_b.launch_builder(&f);
28177 match t_kv_dev {
28178 Some(d) => {
28179 b.arg(q)
28180 .arg(k)
28181 .arg(v)
28182 .arg(&mut *part_o)
28183 .arg(&mut *part_m)
28184 .arg(&mut *part_l)
28185 .arg(&hd)
28186 .arg(&nh)
28187 .arg(&nhkv)
28188 .arg(&tkvi)
28189 .arg(d)
28190 .arg(&scale)
28191 .arg(&nsp)
28192 .arg(&ski)
28193 .arg(&ktb)
28194 .arg(&vtb);
28195 unsafe {
28196 b.launch(cfg)?;
28197 }
28198 }
28199 None => {
28200 let null: u64 = 0;
28201 b.arg(q)
28202 .arg(k)
28203 .arg(v)
28204 .arg(&mut *part_o)
28205 .arg(&mut *part_m)
28206 .arg(&mut *part_l)
28207 .arg(&hd)
28208 .arg(&nh)
28209 .arg(&nhkv)
28210 .arg(&tkvi)
28211 .arg(&null)
28212 .arg(&scale)
28213 .arg(&nsp)
28214 .arg(&ski)
28215 .arg(&ktb)
28216 .arg(&vtb);
28217 unsafe {
28218 b.launch(cfg)?;
28219 }
28220 }
28221 }
28222 let cfg2 = LaunchConfig {
28223 grid_dim: (n_head as u32, 1, 1),
28224 block_dim: (head_dim as u32, 1, 1),
28225 shared_mem_bytes: 0,
28226 };
28227 if let Some((oq, od)) = q8_out {
28228 let fc = if g {
28230 self.func_g("fa_decode_combine_q8_1")
28231 } else {
28232 self.fa_func("fa_decode_combine_q8_1", head_dim)
28233 };
28234 let __s_b2 = self.gpu.stream();
28235 let mut b2 = __s_b2.launch_builder(&fc);
28236 b2.arg(&*part_o)
28237 .arg(&*part_m)
28238 .arg(&*part_l)
28239 .arg(oq)
28240 .arg(od)
28241 .arg(&hd)
28242 .arg(&nh)
28243 .arg(&nsp);
28244 unsafe {
28245 b2.launch(cfg2)?;
28246 }
28247 return Ok(());
28248 }
28249 let fc = if g {
28250 self.func_g("fa_decode_combine_f32")
28251 } else {
28252 self.fa_func("fa_decode_combine_f32", head_dim)
28253 };
28254 let __s_b2 = self.gpu.stream();
28255 let mut b2 = __s_b2.launch_builder(&fc);
28256 b2.arg(&*part_o)
28257 .arg(&*part_m)
28258 .arg(&*part_l)
28259 .arg(o)
28260 .arg(&hd)
28261 .arg(&nh)
28262 .arg(&nsp);
28263 unsafe {
28264 b2.launch(cfg2)?;
28265 }
28266 Ok(())
28267 }
28268
28269 #[allow(clippy::too_many_arguments)] pub fn fa_decode_kvmod(
28271 &self,
28272 q: &CudaSlice<f32>,
28273 k: &cudarc::driver::CudaView<u8>,
28274 v: &cudarc::driver::CudaView<u8>,
28275 o: &mut CudaSlice<f32>,
28276 head_dim: usize,
28277 n_head: usize,
28278 n_head_kv: usize,
28279 t_kv: usize,
28280 scale: f32,
28281 k_tok_bytes: usize,
28282 v_tok_bytes: usize,
28283 g: bool,
28284 ) -> Result<(), Box<dyn std::error::Error>> {
28285 let q_view = q.as_view();
28286 let mut o_view = o.as_view_mut();
28287 self.fa_decode_kvmod_view(
28288 &q_view,
28289 k,
28290 v,
28291 &mut o_view,
28292 head_dim,
28293 n_head,
28294 n_head_kv,
28295 t_kv,
28296 scale,
28297 k_tok_bytes,
28298 v_tok_bytes,
28299 g,
28300 )
28301 }
28302
28303 #[allow(clippy::too_many_arguments)]
28308 #[allow(clippy::manual_div_ceil)] pub fn fa_decode_kvmod_view(
28310 &self,
28311 q: &cudarc::driver::CudaView<f32>,
28312 k: &cudarc::driver::CudaView<u8>,
28313 v: &cudarc::driver::CudaView<u8>,
28314 o: &mut cudarc::driver::CudaViewMut<f32>,
28315 head_dim: usize,
28316 n_head: usize,
28317 n_head_kv: usize,
28318 t_kv: usize,
28319 scale: f32,
28320 k_tok_bytes: usize,
28321 v_tok_bytes: usize,
28322 g: bool,
28323 ) -> Result<(), Box<dyn std::error::Error>> {
28324 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
28345 if g && head_dim == 256 && !fa_v4_at(t_kv) {
28349 fa_vec = false;
28350 }
28351 let sp = fa_split_keys(t_kv, n_head_kv);
28352 let n_splits = if fa_vec {
28353 ((t_kv + sp - 1) / sp).max(1)
28354 } else {
28355 ((t_kv + 255) / 256).max(1)
28356 };
28357 let o_len = n_head * n_splits * head_dim;
28358 let ml_len = n_head * n_splits;
28359 let mut part_guard = self.fa_part_pool.lock().unwrap();
28360 if part_guard
28361 .as_ref()
28362 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
28363 .unwrap_or(true)
28364 {
28365 let old = part_guard.take();
28376 let (co, cm) = old
28377 .as_ref()
28378 .map(|pp| (pp.0.len(), pp.1.len()))
28379 .unwrap_or((0, 0));
28380 if let Some(old) = old {
28381 self.fa_part_retired.lock().unwrap().push(old);
28382 }
28383 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
28384 eprintln!(
28385 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
28386 co, o_len, cm, ml_len
28387 );
28388 }
28389 *part_guard =
28390 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
28391 }
28392 let pg = part_guard.as_mut().unwrap();
28393 self.gpu
28394 .stream()
28395 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
28396 self.gpu
28397 .stream()
28398 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
28399 self.gpu
28400 .stream()
28401 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
28402 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28403 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
28404 let (hd, nh, nhkv, tkvi, nsp) = (
28405 head_dim as i32,
28406 n_head as i32,
28407 n_head_kv as i32,
28408 t_kv as i32,
28409 n_splits as i32,
28410 );
28411 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28412 let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
28416 let fa512_min = fa512_min_tkv();
28421 let deep = fa_vec && head_dim == 256 && fa_v4_at(t_kv) && !g;
28424 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
28425 let gqa = (n_head / n_head_kv).max(1) as u32;
28428 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
28429 (
28430 fv,
28431 LaunchConfig {
28432 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28433 block_dim: (32, gqa, 1),
28434 shared_mem_bytes: 0,
28435 },
28436 )
28437 } else if fa_vec && head_dim <= 256 {
28438 let gqa = (n_head / n_head_kv).max(1) as u32;
28439 if fa_v4_at(t_kv) && head_dim == 256 {
28440 let v4name = if deep {
28444 "fa_decode_vec_q_v4_deep"
28445 } else {
28446 "fa_decode_vec_q_v4"
28447 };
28448 let fv = if g {
28449 self.func_g(v4name)
28450 } else {
28451 self.func(v4name)
28452 };
28453 let shmem = (if deep { 12160 } else { 11520 }
28456 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
28457 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28458 fv.set_attribute(
28459 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28460 shmem as i32,
28461 )?;
28462 (
28463 fv,
28464 LaunchConfig {
28465 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28466 block_dim: (32, gqa, 1),
28467 shared_mem_bytes: shmem,
28468 },
28469 )
28470 } else if fa_v3_active(head_dim) {
28471 let fv = if g {
28474 self.func_g("fa_decode_vec_q_v3")
28475 } else {
28476 self.func("fa_decode_vec_q_v3")
28477 };
28478 let shmem = (32 * head_dim * 2) as u32; (
28480 fv,
28481 LaunchConfig {
28482 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28483 block_dim: (32, gqa, 1),
28484 shared_mem_bytes: shmem,
28485 },
28486 )
28487 } else {
28488 let fv = if g {
28492 self.func_g("fa_decode_vec_q_v2")
28493 } else {
28494 self.func("fa_decode_vec_q_v2")
28495 };
28496 let shmem = (2 * 32 * head_dim * 2) as u32; (
28498 fv,
28499 LaunchConfig {
28500 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28501 block_dim: (32, gqa, 1),
28502 shared_mem_bytes: shmem,
28503 },
28504 )
28505 }
28506 } else {
28507 return self.fa_decode_scalar_unified(
28510 q,
28511 k,
28512 v,
28513 o,
28514 head_dim,
28515 n_head,
28516 n_head_kv,
28517 t_kv,
28518 None,
28519 scale,
28520 n_splits,
28521 if fa_vec { sp } else { 256 },
28522 k_tok_bytes,
28523 v_tok_bytes,
28524 g,
28525 part_o,
28526 part_m,
28527 part_l,
28528 None,
28529 );
28530 };
28531 let __s_b = self.gpu.stream();
28532 let mut b = __s_b.launch_builder(&f);
28533 b.arg(q)
28534 .arg(k)
28535 .arg(v)
28536 .arg(&mut *part_o)
28537 .arg(&mut *part_m)
28538 .arg(&mut *part_l)
28539 .arg(&hd)
28540 .arg(&nh)
28541 .arg(&nhkv)
28542 .arg(&tkvi)
28543 .arg(&scale)
28544 .arg(&nsp)
28545 .arg(&ktb)
28546 .arg(&vtb);
28547 unsafe {
28548 b.launch(cfg)?;
28549 }
28550 let (fc, cfg2) = (
28553 if g {
28554 self.func_g("fa_decode_combine_f32")
28555 } else {
28556 self.fa_func("fa_decode_combine_f32", head_dim)
28557 },
28558 LaunchConfig {
28559 grid_dim: (n_head as u32, 1, 1),
28560 block_dim: (head_dim as u32, 1, 1),
28561 shared_mem_bytes: 0,
28562 },
28563 );
28564 let __s_b2 = self.gpu.stream();
28565 let mut b2 = __s_b2.launch_builder(&fc);
28566 b2.arg(&*part_o)
28567 .arg(&*part_m)
28568 .arg(&*part_l)
28569 .arg(o)
28570 .arg(&hd)
28571 .arg(&nh)
28572 .arg(&nsp);
28573 unsafe {
28574 b2.launch(cfg2)?;
28575 }
28576 Ok(())
28577 }
28578
28579 #[allow(clippy::too_many_arguments)]
28590 pub fn fa_decode_batch_seqs_v4(
28591 &self,
28592 q: &CudaSlice<f32>,
28593 kv_ptrs: &cudarc::driver::CudaView<u64>,
28594 pos_seq: &CudaSlice<i32>,
28595 o: &mut CudaSlice<f32>,
28596 head_dim: usize,
28597 n_head: usize,
28598 n_head_kv: usize,
28599 b_n: usize,
28600 t_kv_max: usize,
28601 scale: f32,
28602 split_keys: usize,
28603 k_tok_bytes: usize,
28604 v_tok_bytes: usize,
28605 ) -> Result<(), Box<dyn std::error::Error>> {
28606 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
28607 #[allow(clippy::manual_div_ceil)]
28608 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
28610 let o_len = b_n * n_head * n_splits_max * head_dim;
28611 let ml_len = b_n * n_head * n_splits_max;
28612 let mut part_guard = self.fa_part_pool.lock().unwrap();
28613 if part_guard
28614 .as_ref()
28615 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
28616 .unwrap_or(true)
28617 {
28618 let old = part_guard.take();
28629 let (co, cm) = old
28630 .as_ref()
28631 .map(|pp| (pp.0.len(), pp.1.len()))
28632 .unwrap_or((0, 0));
28633 if let Some(old) = old {
28634 self.fa_part_retired.lock().unwrap().push(old);
28635 }
28636 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
28637 eprintln!(
28638 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
28639 co, o_len, cm, ml_len
28640 );
28641 }
28642 *part_guard =
28643 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
28644 }
28645 let pg = part_guard.as_mut().unwrap();
28646 self.gpu
28647 .stream()
28648 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
28649 self.gpu
28650 .stream()
28651 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
28652 self.gpu
28653 .stream()
28654 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
28655 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28656 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
28657 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
28658 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28659 let gqa = (n_head / n_head_kv).max(1) as u32;
28660 let f = self.func("fa_decode_vec_q_seqs_v4");
28661 let shmem = (11520 + 32 * head_dim * 2) as u32;
28663 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28664 f.set_attribute(
28665 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28666 shmem as i32,
28667 )?;
28668 let cfg = LaunchConfig {
28669 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
28670 block_dim: (32, gqa, 1),
28671 shared_mem_bytes: shmem,
28672 };
28673 {
28674 let __s_b = self.gpu.stream();
28675 let mut b = __s_b.launch_builder(&f);
28676 b.arg(q)
28677 .arg(kv_ptrs)
28678 .arg(pos_seq)
28679 .arg(&mut *part_o)
28680 .arg(&mut *part_m)
28681 .arg(&mut *part_l)
28682 .arg(&hd)
28683 .arg(&nh)
28684 .arg(&nhkv)
28685 .arg(&scale)
28686 .arg(&nspm)
28687 .arg(&spk)
28688 .arg(&ktb)
28689 .arg(&vtb);
28690 unsafe {
28691 b.launch(cfg)?;
28692 }
28693 }
28694 let fc = self.func("fa_decode_combine_seqs");
28695 let cfg2 = LaunchConfig {
28696 grid_dim: (n_head as u32, b_n as u32, 1),
28697 block_dim: (head_dim as u32, 1, 1),
28698 shared_mem_bytes: 0,
28699 };
28700 let __s_b2 = self.gpu.stream();
28701 let mut b2 = __s_b2.launch_builder(&fc);
28702 b2.arg(&*part_o)
28703 .arg(&*part_m)
28704 .arg(&*part_l)
28705 .arg(o)
28706 .arg(&hd)
28707 .arg(&nh)
28708 .arg(pos_seq)
28709 .arg(&nspm)
28710 .arg(&spk);
28711 unsafe {
28712 b2.launch(cfg2)?;
28713 }
28714 Ok(())
28715 }
28716
28717 #[allow(clippy::too_many_arguments)]
28724 pub fn append_kv_quantized_seqs(
28725 &self,
28726 k_rows: &CudaSlice<f32>,
28727 v_rows: &CudaSlice<f32>,
28728 kv_ptrs: &cudarc::driver::CudaView<u64>,
28729 pos_seq: &CudaSlice<i32>,
28730 b_n: usize,
28731 kv_dim_k: usize,
28732 kv_dim_v: usize,
28733 k_tok_bytes: usize,
28734 v_tok_bytes: usize,
28735 ) -> Result<(), Box<dyn std::error::Error>> {
28736 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
28737 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
28738 let cfg = LaunchConfig {
28739 grid_dim: (nblk, b_n as u32, 1),
28740 block_dim: (32, 1, 1),
28741 shared_mem_bytes: 0,
28742 };
28743 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
28744 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28745 let __s_b = self.gpu.stream();
28746 let mut b = __s_b.launch_builder(&f);
28747 b.arg(k_rows)
28748 .arg(v_rows)
28749 .arg(kv_ptrs)
28750 .arg(pos_seq)
28751 .arg(&kdk)
28752 .arg(&kdv)
28753 .arg(&ktb)
28754 .arg(&vtb);
28755 unsafe {
28756 b.launch(cfg)?;
28757 }
28758 Ok(())
28759 }
28760
28761 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
28767 std::env::var("MEMRA_NO_FA_VEC").is_err()
28768 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
28769 && base_len + 1 >= fa_vec_min_tkv()
28770 && head_dim <= 256
28771 && head_dim.is_multiple_of(32)
28772 }
28773
28774 #[allow(clippy::too_many_arguments)]
28783 pub fn fa_decode_rows(
28784 &self,
28785 q: &CudaSlice<f32>,
28786 k: &cudarc::driver::CudaView<u8>,
28787 v: &cudarc::driver::CudaView<u8>,
28788 o: &mut CudaSlice<f32>,
28789 head_dim: usize,
28790 n_head: usize,
28791 n_head_kv: usize,
28792 base_len: usize,
28793 t: usize,
28794 scale: f32,
28795 k_tok_bytes: usize,
28796 v_tok_bytes: usize,
28797 base_dev: Option<(&CudaSlice<i32>, i32)>,
28801 kv_shared: bool,
28804 g: bool,
28808 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
28811 ) -> Result<(), Box<dyn std::error::Error>> {
28812 debug_assert!(
28813 base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim.is_multiple_of(32)
28814 );
28815 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
28822 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
28823 let v = *SP512.get_or_init(|| {
28826 std::env::var("MEMRA_FA_SP512")
28827 .ok()
28828 .and_then(|x| x.parse().ok())
28829 .unwrap_or(0)
28830 });
28831 sp = if v >= 8 {
28832 v
28833 } else {
28834 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
28835 };
28836 }
28837 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
28838 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28839 let gqa = (n_head / n_head_kv).max(1) as u32;
28840 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
28851 groups.push((0, t, sp));
28852 } else {
28853 let mut r0 = 0usize;
28854 while r0 < t {
28855 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
28856 let mut r1 = r0 + 1;
28857 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
28858 r1 += 1;
28859 }
28860 groups.push((r0, r1 - r0, sp_g));
28861 r0 = r1;
28862 }
28863 }
28864 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
28868 let v3 = fa_v3_active(head_dim);
28869 let _ = kv_shared;
28874 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
28877 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
28891 let tb512 = head_dim == 512
28893 && sp <= 32
28894 && n_head / n_head_kv.max(1) <= 16
28895 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
28896 let fname = if tb512 {
28897 "fa_decode_vec_q_rows_v4_512_tb"
28898 } else if i2 {
28899 "fa_decode_vec_q_rows_dpl16_i2"
28900 } else if head_dim == 512 {
28901 "fa_decode_vec_q_rows_dpl16"
28902 }
28903 else if v4 {
28905 "fa_decode_vec_q_rows_v4"
28906 } else if v3 {
28907 "fa_decode_vec_q_rows_v3"
28908 } else {
28909 "fa_decode_vec_q_rows_v2"
28910 };
28911 let f = if head_dim == 512 {
28912 self.fa_func(fname, head_dim)
28913 } else if g {
28914 self.func_g(fname)
28922 } else {
28923 self.func(fname)
28924 };
28925 let shmem = if tb512 {
28926 let gk = Self::gkv_on();
28928 let sh =
28929 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
28930 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28931 f.set_attribute(
28932 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28933 sh as i32,
28934 )?;
28935 sh
28936 } else {
28937 let sh = (if v4 {
28939 11520 + 32 * head_dim * if g { 1 } else { 2 }
28940 } else if v3 {
28941 32 * head_dim * 2
28942 } else {
28943 2 * 32 * head_dim * 2
28944 }) as u32;
28945 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28946 f.set_attribute(
28947 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28948 sh as i32,
28949 )?;
28950 sh
28951 };
28952 for &(r0, t_g, sp_g) in &groups {
28956 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
28957 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
28958 let base_i = (base_len + r0) as i32;
28959 let o_len = t_g * n_head * n_splits_g * head_dim;
28960 let ml_len = t_g * n_head * n_splits_g;
28961 let mut part_guard = self.fa_part_pool.lock().unwrap();
28962 if part_guard
28963 .as_ref()
28964 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
28965 .unwrap_or(true)
28966 {
28967 let old = part_guard.take();
28978 let (co, cm) = old
28979 .as_ref()
28980 .map(|pp| (pp.0.len(), pp.1.len()))
28981 .unwrap_or((0, 0));
28982 if let Some(old) = old {
28983 self.fa_part_retired.lock().unwrap().push(old);
28984 }
28985 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
28986 eprintln!(
28987 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
28988 co, o_len, cm, ml_len
28989 );
28990 }
28991 *part_guard =
28992 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
28993 }
28994 let pg = part_guard.as_mut().unwrap();
28995 self.gpu
28996 .stream()
28997 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
28998 self.gpu
28999 .stream()
29000 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29001 self.gpu
29002 .stream()
29003 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29004 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29005 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
29006 let qv = self.view(q, t * n_head * head_dim);
29007 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
29008 let cfg = LaunchConfig {
29009 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
29010 block_dim: (32, gqa, 1),
29011 shared_mem_bytes: shmem,
29012 };
29013 {
29014 let __s_b = self.gpu.stream();
29015 let mut b = __s_b.launch_builder(&f);
29016 if tb512 {
29017 let (bd, plus) =
29019 base_dev.expect("hd512 rows twin requires a device base counter");
29020 let plus_g = plus + r0 as i32;
29021 let nr = t_g as i32;
29022 if Self::pdl_on() && Self::pdl_wb_on() {
29023 use cudarc::driver::{DevicePtr, DevicePtrMut};
29025 let s = &self.gpu.stream();
29026 let (pq, _b0) = q_g.device_ptr(s);
29027 let (pk, _b1) = k.device_ptr(s);
29028 let (pv, _b2) = v.device_ptr(s);
29029 let (po, _b3) = part_o.device_ptr_mut(s);
29030 let (pm, _b4) = part_m.device_ptr_mut(s);
29031 let (pl, _b5) = part_l.device_ptr_mut(s);
29032 let (pb, _b6) = bd.device_ptr(s);
29033 let mut ps = [
29034 &pq as *const _ as *mut std::ffi::c_void,
29035 &pk as *const _ as *mut _,
29036 &pv as *const _ as *mut _,
29037 &po as *const _ as *mut _,
29038 &pm as *const _ as *mut _,
29039 &pl as *const _ as *mut _,
29040 &hd as *const _ as *mut _,
29041 &nh as *const _ as *mut _,
29042 &nhkv as *const _ as *mut _,
29043 &pb as *const _ as *mut _,
29044 &plus_g as *const _ as *mut _,
29045 &scale as *const _ as *mut _,
29046 &nspm as *const _ as *mut _,
29047 &spk as *const _ as *mut _,
29048 &ktb as *const _ as *mut _,
29049 &vtb as *const _ as *mut _,
29050 &nr as *const _ as *mut _,
29051 ];
29052 unsafe {
29053 self.launch_pdl_flash(
29054 Self::gkv_on(),
29055 "fa_decode_vec_q_rows_v4_512_tb",
29056 (n_head_kv as u32, n_splits_g as u32, 1),
29057 (32, gqa, 1),
29058 shmem,
29059 &mut ps,
29060 )?;
29061 }
29062 } else {
29063 let cfg_tb = LaunchConfig {
29064 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
29065 block_dim: (32, gqa, 1),
29066 shared_mem_bytes: shmem,
29067 };
29068 b.arg(&q_g)
29069 .arg(k)
29070 .arg(v)
29071 .arg(&mut *part_o)
29072 .arg(&mut *part_m)
29073 .arg(&mut *part_l)
29074 .arg(&hd)
29075 .arg(&nh)
29076 .arg(&nhkv)
29077 .arg(bd)
29078 .arg(&plus_g)
29079 .arg(&scale)
29080 .arg(&nspm)
29081 .arg(&spk)
29082 .arg(&ktb)
29083 .arg(&vtb)
29084 .arg(&nr);
29085 unsafe {
29086 b.launch(cfg_tb)?;
29087 }
29088 }
29089 } else if head_dim == 512 {
29090 let (bd, plus) =
29091 base_dev.expect("hd512 rows twin requires a device base counter");
29092 let plus_g = plus + r0 as i32;
29093 b.arg(&q_g)
29094 .arg(k)
29095 .arg(v)
29096 .arg(&mut *part_o)
29097 .arg(&mut *part_m)
29098 .arg(&mut *part_l)
29099 .arg(&hd)
29100 .arg(&nh)
29101 .arg(&nhkv)
29102 .arg(bd)
29103 .arg(&plus_g)
29104 .arg(&scale)
29105 .arg(&nspm)
29106 .arg(&spk)
29107 .arg(&ktb)
29108 .arg(&vtb);
29109 unsafe {
29110 b.launch(cfg)?;
29111 }
29112 } else {
29113 b.arg(&q_g)
29114 .arg(k)
29115 .arg(v)
29116 .arg(&mut *part_o)
29117 .arg(&mut *part_m)
29118 .arg(&mut *part_l)
29119 .arg(&hd)
29120 .arg(&nh)
29121 .arg(&nhkv)
29122 .arg(&base_i)
29123 .arg(&scale)
29124 .arg(&nspm)
29125 .arg(&spk)
29126 .arg(&ktb)
29127 .arg(&vtb);
29128 unsafe {
29129 b.launch(cfg)?;
29130 }
29131 }
29132 }
29133 let cfg2 = LaunchConfig {
29134 grid_dim: (n_head as u32, t_g as u32, 1),
29135 block_dim: (head_dim as u32, 1, 1),
29136 shared_mem_bytes: 0,
29137 };
29138 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
29139 if head_dim == 512 {
29140 let (bd, plus) = base_dev.unwrap();
29143 let plus_g = plus + r0 as i32;
29144 if let Some((oq, od)) = q8_out.as_mut() {
29145 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
29147 if Self::pdl_on() && Self::pdl_wb_on() {
29148 use cudarc::driver::{DevicePtr, DevicePtrMut};
29150 let s = &self.gpu.stream();
29151 let (po, _g0) = part_o.device_ptr(s);
29152 let (pm, _g1) = part_m.device_ptr(s);
29153 let (pl, _g2) = part_l.device_ptr(s);
29154 let (pq, _g3) = oq.device_ptr_mut(s);
29155 let (pd, _g4) = od.device_ptr_mut(s);
29156 let (pb, _g5) = bd.device_ptr(s);
29157 let mut ps = [
29158 &po as *const _ as *mut std::ffi::c_void,
29159 &pm as *const _ as *mut _,
29160 &pl as *const _ as *mut _,
29161 &pq as *const _ as *mut _,
29162 &pd as *const _ as *mut _,
29163 &hd as *const _ as *mut _,
29164 &nh as *const _ as *mut _,
29165 &pb as *const _ as *mut _,
29166 &plus_g as *const _ as *mut _,
29167 &nspm as *const _ as *mut _,
29168 &spk as *const _ as *mut _,
29169 ];
29170 unsafe {
29171 self.launch_pdl_flash(
29172 Self::gkv_on(),
29173 "fa_decode_combine_rows_dc_q8_1",
29174 cfg2.grid_dim,
29175 cfg2.block_dim,
29176 0,
29177 &mut ps,
29178 )?;
29179 }
29180 continue;
29181 }
29182 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
29183 let __s_b2 = self.gpu.stream();
29184 let mut b2 = __s_b2.launch_builder(&fc);
29185 b2.arg(&*part_o)
29186 .arg(&*part_m)
29187 .arg(&*part_l)
29188 .arg(&mut **oq)
29189 .arg(&mut **od)
29190 .arg(&hd)
29191 .arg(&nh)
29192 .arg(bd)
29193 .arg(&plus_g)
29194 .arg(&nspm)
29195 .arg(&spk);
29196 unsafe {
29197 b2.launch(cfg2)?;
29198 }
29199 continue;
29200 }
29201 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
29202 let __s_b2 = self.gpu.stream();
29203 let mut b2 = __s_b2.launch_builder(&fc);
29204 b2.arg(&*part_o)
29205 .arg(&*part_m)
29206 .arg(&*part_l)
29207 .arg(&mut o_g)
29208 .arg(&hd)
29209 .arg(&nh)
29210 .arg(bd)
29211 .arg(&plus_g)
29212 .arg(&nspm)
29213 .arg(&spk);
29214 unsafe {
29215 b2.launch(cfg2)?;
29216 }
29217 } else {
29218 assert!(
29221 q8_out.is_none(),
29222 "rows q8 emit requires the hd512 dc combine"
29223 );
29224 let fc = self.func("fa_decode_combine_rows");
29225 let __s_b2 = self.gpu.stream();
29226 let mut b2 = __s_b2.launch_builder(&fc);
29227 b2.arg(&*part_o)
29228 .arg(&*part_m)
29229 .arg(&*part_l)
29230 .arg(&mut o_g)
29231 .arg(&hd)
29232 .arg(&nh)
29233 .arg(&base_i)
29234 .arg(&nspm)
29235 .arg(&spk);
29236 unsafe {
29237 b2.launch(cfg2)?;
29238 }
29239 }
29240 }
29241 Ok(())
29242 }
29243
29244 #[allow(clippy::too_many_arguments)]
29248 pub fn fa_decode_rows_w(
29249 &self,
29250 q: &CudaSlice<f32>,
29251 k: &cudarc::driver::CudaView<u8>,
29252 v: &cudarc::driver::CudaView<u8>,
29253 o: &mut CudaSlice<f32>,
29254 head_dim: usize,
29255 n_head: usize,
29256 n_head_kv: usize,
29257 base_dev: &CudaSlice<i32>,
29258 base_plus: i32,
29259 t: usize,
29260 scale: f32,
29261 window: usize,
29262 k_tok_bytes: usize,
29263 v_tok_bytes: usize,
29264 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
29265 ) -> Result<(), Box<dyn std::error::Error>> {
29266 debug_assert!(head_dim == 256);
29271 let sp = {
29279 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29280 let v = *SPW.get_or_init(|| {
29281 std::env::var("MEMRA_FA_SPW")
29282 .ok()
29283 .and_then(|x| x.parse().ok())
29284 .unwrap_or(0)
29285 });
29286 if v >= 8 {
29287 v
29288 } else {
29289 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
29290 }
29291 };
29292 #[allow(clippy::manual_div_ceil)]
29293 let n_splits_max = (window + sp - 1) / sp;
29295 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29296 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
29297 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29298 let gqa = (n_head / n_head_kv).max(1) as u32;
29299 let o_len = t * n_head * n_splits_max * head_dim;
29300 let ml_len = t * n_head * n_splits_max;
29301 let mut part_guard = self.fa_part_pool.lock().unwrap();
29302 if part_guard
29303 .as_ref()
29304 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29305 .unwrap_or(true)
29306 {
29307 let old = part_guard.take();
29318 let (co, cm) = old
29319 .as_ref()
29320 .map(|pp| (pp.0.len(), pp.1.len()))
29321 .unwrap_or((0, 0));
29322 if let Some(old) = old {
29323 self.fa_part_retired.lock().unwrap().push(old);
29324 }
29325 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29326 eprintln!(
29327 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29328 co, o_len, cm, ml_len
29329 );
29330 }
29331 *part_guard =
29332 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29333 }
29334 let pg = part_guard.as_mut().unwrap();
29335 self.gpu
29336 .stream()
29337 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29338 self.gpu
29339 .stream()
29340 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29341 self.gpu
29342 .stream()
29343 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29344 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29345 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29351 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
29352 std::env::var("MEMRA_FA_SMEM_TKV")
29353 .ok()
29354 .and_then(|v| v.parse().ok())
29355 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
29356 });
29357 use cudarc::driver::sys::CUfunction_attribute_enum as A;
29363 let wg = Self::wkv_on();
29368 let sp2 =
29371 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
29372 if sp2 {
29373 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
29374 if Self::pdl_on() && Self::pdl_wb_on() {
29375 use cudarc::driver::{DevicePtr, DevicePtrMut};
29377 let s = &self.gpu.stream();
29378 let (pq, _b0) = q.device_ptr(s);
29379 let (pk, _b1) = k.device_ptr(s);
29380 let (pv, _b2) = v.device_ptr(s);
29381 let (po, _b3) = part_o.device_ptr_mut(s);
29382 let (pm, _b4) = part_m.device_ptr_mut(s);
29383 let (pl, _b5) = part_l.device_ptr_mut(s);
29384 let (pb, _b6) = base_dev.device_ptr(s);
29385 let mut ps = [
29386 &pq as *const _ as *mut std::ffi::c_void,
29387 &pk as *const _ as *mut _,
29388 &pv as *const _ as *mut _,
29389 &po as *const _ as *mut _,
29390 &pm as *const _ as *mut _,
29391 &pl as *const _ as *mut _,
29392 &hd as *const _ as *mut _,
29393 &nh as *const _ as *mut _,
29394 &nhkv as *const _ as *mut _,
29395 &pb as *const _ as *mut _,
29396 &base_plus as *const _ as *mut _,
29397 &scale as *const _ as *mut _,
29398 &nspm as *const _ as *mut _,
29399 &spk as *const _ as *mut _,
29400 &ktb as *const _ as *mut _,
29401 &vtb as *const _ as *mut _,
29402 &wini as *const _ as *mut _,
29403 ];
29404 unsafe {
29405 self.launch_pdl_flash(
29406 wg,
29407 "fa_decode_vec_q_rows_v4_w_sp",
29408 (n_head_kv as u32, n_splits_max as u32, t as u32),
29409 (32, gqa + 1, 1),
29410 sh,
29411 &mut ps,
29412 )?;
29413 }
29414 } else {
29415 let f = if wg {
29416 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
29417 } else {
29418 self.func("fa_decode_vec_q_rows_v4_w_sp")
29419 };
29420 f.set_attribute(
29421 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29422 sh as i32,
29423 )?;
29424 let cfg = LaunchConfig {
29425 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
29426 block_dim: (32, gqa + 1, 1),
29427 shared_mem_bytes: sh,
29428 };
29429 let __s_b = self.gpu.stream();
29430 let mut b = __s_b.launch_builder(&f);
29431 b.arg(q)
29432 .arg(k)
29433 .arg(v)
29434 .arg(&mut *part_o)
29435 .arg(&mut *part_m)
29436 .arg(&mut *part_l)
29437 .arg(&hd)
29438 .arg(&nh)
29439 .arg(&nhkv)
29440 .arg(base_dev)
29441 .arg(&base_plus)
29442 .arg(&scale)
29443 .arg(&nspm)
29444 .arg(&spk)
29445 .arg(&ktb)
29446 .arg(&vtb)
29447 .arg(&wini);
29448 unsafe {
29449 b.launch(cfg)?;
29450 }
29451 }
29452 } else {
29453 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
29454 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
29456 use cudarc::driver::{DevicePtr, DevicePtrMut};
29457 let s = &self.gpu.stream();
29458 let (pq, _b0) = q.device_ptr(s);
29459 let (pk, _b1) = k.device_ptr(s);
29460 let (pv, _b2) = v.device_ptr(s);
29461 let (po, _b3) = part_o.device_ptr_mut(s);
29462 let (pm, _b4) = part_m.device_ptr_mut(s);
29463 let (pl, _b5) = part_l.device_ptr_mut(s);
29464 let (pb, _b6) = base_dev.device_ptr(s);
29465 let mut ps = [
29466 &pq as *const _ as *mut std::ffi::c_void,
29467 &pk as *const _ as *mut _,
29468 &pv as *const _ as *mut _,
29469 &po as *const _ as *mut _,
29470 &pm as *const _ as *mut _,
29471 &pl as *const _ as *mut _,
29472 &hd as *const _ as *mut _,
29473 &nh as *const _ as *mut _,
29474 &nhkv as *const _ as *mut _,
29475 &pb as *const _ as *mut _,
29476 &base_plus as *const _ as *mut _,
29477 &scale as *const _ as *mut _,
29478 &nspm as *const _ as *mut _,
29479 &spk as *const _ as *mut _,
29480 &ktb as *const _ as *mut _,
29481 &vtb as *const _ as *mut _,
29482 &wini as *const _ as *mut _,
29483 ];
29484 unsafe {
29485 self.launch_pdl_flash(
29486 wg,
29487 "fa_decode_vec_q_rows_v4_w",
29488 (n_head_kv as u32, n_splits_max as u32, t as u32),
29489 (32, gqa, 1),
29490 sh,
29491 &mut ps,
29492 )?;
29493 }
29494 } else {
29495 let pick = |name: &str| {
29496 if wg {
29497 self.func_g(name)
29498 } else {
29499 self.func(name)
29500 }
29501 };
29502 let (f, sh) = if fa_v4_at(window) {
29503 let f = pick("fa_decode_vec_q_rows_v4_w");
29504 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
29505 } else if smem_tkv > 0 && window >= smem_tkv {
29506 (
29509 pick("fa_decode_vec_q_rows_smem_w"),
29510 (2 * 32 * head_dim * 2) as u32,
29511 )
29512 } else {
29513 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
29514 };
29515 f.set_attribute(
29516 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29517 sh as i32,
29518 )?;
29519 let cfg = LaunchConfig {
29520 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
29521 block_dim: (32, gqa, 1),
29522 shared_mem_bytes: sh,
29523 };
29524 let __s_b = self.gpu.stream();
29525 let mut b = __s_b.launch_builder(&f);
29526 b.arg(q)
29527 .arg(k)
29528 .arg(v)
29529 .arg(&mut *part_o)
29530 .arg(&mut *part_m)
29531 .arg(&mut *part_l)
29532 .arg(&hd)
29533 .arg(&nh)
29534 .arg(&nhkv)
29535 .arg(base_dev)
29536 .arg(&base_plus)
29537 .arg(&scale)
29538 .arg(&nspm)
29539 .arg(&spk)
29540 .arg(&ktb)
29541 .arg(&vtb)
29542 .arg(&wini);
29543 unsafe {
29544 b.launch(cfg)?;
29545 }
29546 }
29547 }
29548 let cfg2 = LaunchConfig {
29549 grid_dim: (n_head as u32, t as u32, 1),
29550 block_dim: (head_dim as u32, 1, 1),
29551 shared_mem_bytes: 0,
29552 };
29553 if let Some((oq, od)) = q8_out {
29554 if Self::pdl_on() && Self::pdl_wb_on() {
29557 use cudarc::driver::{DevicePtr, DevicePtrMut};
29559 let s = &self.gpu.stream();
29560 let (po, _g0) = part_o.device_ptr(s);
29561 let (pm, _g1) = part_m.device_ptr(s);
29562 let (pl, _g2) = part_l.device_ptr(s);
29563 let (pq, _g3) = oq.device_ptr_mut(s);
29564 let (pd, _g4) = od.device_ptr_mut(s);
29565 let mut ps = [
29566 &po as *const _ as *mut std::ffi::c_void,
29567 &pm as *const _ as *mut _,
29568 &pl as *const _ as *mut _,
29569 &pq as *const _ as *mut _,
29570 &pd as *const _ as *mut _,
29571 &hd as *const _ as *mut _,
29572 &nh as *const _ as *mut _,
29573 &nspm as *const _ as *mut _,
29574 &spk as *const _ as *mut _,
29575 &wini as *const _ as *mut _,
29576 ];
29577 unsafe {
29578 self.launch_pdl_flash(
29579 wg,
29580 "fa_decode_combine_rows_w_q8_1",
29581 cfg2.grid_dim,
29582 cfg2.block_dim,
29583 0,
29584 &mut ps,
29585 )?;
29586 }
29587 return Ok(());
29588 }
29589 let fc = if wg {
29590 self.func_g("fa_decode_combine_rows_w_q8_1")
29591 } else {
29592 self.func("fa_decode_combine_rows_w_q8_1")
29593 };
29594 let __s_b2 = self.gpu.stream();
29595 let mut b2 = __s_b2.launch_builder(&fc);
29596 b2.arg(&*part_o)
29597 .arg(&*part_m)
29598 .arg(&*part_l)
29599 .arg(oq)
29600 .arg(od)
29601 .arg(&hd)
29602 .arg(&nh)
29603 .arg(&nspm)
29604 .arg(&spk)
29605 .arg(&wini);
29606 unsafe {
29607 b2.launch(cfg2)?;
29608 }
29609 return Ok(());
29610 }
29611 let fc = if wg {
29612 self.func_g("fa_decode_combine_rows_w")
29613 } else {
29614 self.func("fa_decode_combine_rows_w")
29615 };
29616 let __s_b2 = self.gpu.stream();
29617 let mut b2 = __s_b2.launch_builder(&fc);
29618 b2.arg(&*part_o)
29619 .arg(&*part_m)
29620 .arg(&*part_l)
29621 .arg(o)
29622 .arg(&hd)
29623 .arg(&nh)
29624 .arg(&nspm)
29625 .arg(&spk)
29626 .arg(&wini);
29627 unsafe {
29628 b2.launch(cfg2)?;
29629 }
29630 Ok(())
29631 }
29632
29633 #[allow(clippy::too_many_arguments)]
29639 pub fn fa_decode_rows_dc(
29640 &self,
29641 q: &CudaSlice<f32>,
29642 k: &cudarc::driver::CudaView<u8>,
29643 v: &cudarc::driver::CudaView<u8>,
29644 o: &mut CudaSlice<f32>,
29645 head_dim: usize,
29646 n_head: usize,
29647 n_head_kv: usize,
29648 base_dev: &CudaSlice<i32>,
29649 t_kv_upper: usize,
29650 t: usize,
29651 scale: f32,
29652 k_tok_bytes: usize,
29653 v_tok_bytes: usize,
29654 base_plus: i32,
29655 g: bool,
29656 ) -> Result<(), Box<dyn std::error::Error>> {
29657 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
29658 assert!(
29659 v4 || fa_v3_active(head_dim),
29660 "stream fa rows requires the v3 or v4 lane"
29661 );
29662 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
29663 if v4 {
29664 let sp = fa_split_keys(t_kv_upper, n_head_kv);
29665 #[allow(clippy::manual_div_ceil)]
29666 let n_splits_max = (t_kv_upper + sp - 1) / sp;
29668 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29669 let (nspm, spk) = (n_splits_max as i32, sp as i32);
29670 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29671 let gqa = (n_head / n_head_kv).max(1) as u32;
29672 let o_len = t * n_head * n_splits_max * head_dim;
29673 let ml_len = t * n_head * n_splits_max;
29674 let mut part_guard = self.fa_part_pool.lock().unwrap();
29675 if part_guard
29676 .as_ref()
29677 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29678 .unwrap_or(true)
29679 {
29680 let old = part_guard.take();
29691 let (co, cm) = old
29692 .as_ref()
29693 .map(|pp| (pp.0.len(), pp.1.len()))
29694 .unwrap_or((0, 0));
29695 if let Some(old) = old {
29696 self.fa_part_retired.lock().unwrap().push(old);
29697 }
29698 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29699 eprintln!(
29700 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29701 co, o_len, cm, ml_len
29702 );
29703 }
29704 *part_guard =
29705 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29706 }
29707 let pg = part_guard.as_mut().unwrap();
29708 self.gpu
29709 .stream()
29710 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29711 self.gpu
29712 .stream()
29713 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29714 self.gpu
29715 .stream()
29716 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29717 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29718 let f = if g {
29719 self.func_g("fa_decode_vec_q_rows_v4_dc")
29720 } else {
29721 self.func("fa_decode_vec_q_rows_v4_dc")
29722 };
29723 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
29724 use cudarc::driver::sys::CUfunction_attribute_enum as A;
29725 f.set_attribute(
29726 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29727 sh as i32,
29728 )?;
29729 let cfg = LaunchConfig {
29730 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
29731 block_dim: (32, gqa, 1),
29732 shared_mem_bytes: sh,
29733 };
29734 let __s_b = self.gpu.stream();
29735 let mut b = __s_b.launch_builder(&f);
29736 b.arg(q)
29737 .arg(k)
29738 .arg(v)
29739 .arg(&mut *part_o)
29740 .arg(&mut *part_m)
29741 .arg(&mut *part_l)
29742 .arg(&hd)
29743 .arg(&nh)
29744 .arg(&nhkv)
29745 .arg(base_dev)
29746 .arg(&base_plus)
29747 .arg(&scale)
29748 .arg(&nspm)
29749 .arg(&spk)
29750 .arg(&ktb)
29751 .arg(&vtb);
29752 unsafe {
29753 b.launch(cfg)?;
29754 }
29755 let fc = self.func("fa_decode_combine_rows_dc");
29756 let cfg2 = LaunchConfig {
29757 grid_dim: (n_head as u32, t as u32, 1),
29758 block_dim: (head_dim as u32, 1, 1),
29759 shared_mem_bytes: 0,
29760 };
29761 let __s_b2 = self.gpu.stream();
29762 let mut b2 = __s_b2.launch_builder(&fc);
29763 b2.arg(&*part_o)
29764 .arg(&*part_m)
29765 .arg(&*part_l)
29766 .arg(o)
29767 .arg(&hd)
29768 .arg(&nh)
29769 .arg(base_dev)
29770 .arg(&base_plus)
29771 .arg(&nspm)
29772 .arg(&spk);
29773 unsafe {
29774 b2.launch(cfg2)?;
29775 }
29776 return Ok(());
29777 }
29778 let sp = fa_split_keys(t_kv_upper, n_head_kv);
29779 #[allow(clippy::manual_div_ceil)]
29780 let n_splits_max = (t_kv_upper + sp - 1) / sp;
29782 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29783 let (nspm, spk) = (n_splits_max as i32, sp as i32);
29784 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29785 let gqa = (n_head / n_head_kv).max(1) as u32;
29786 let o_len = t * n_head * n_splits_max * head_dim;
29787 let ml_len = t * n_head * n_splits_max;
29788 let mut part_guard = self.fa_part_pool.lock().unwrap();
29789 if part_guard
29790 .as_ref()
29791 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29792 .unwrap_or(true)
29793 {
29794 let old = part_guard.take();
29805 let (co, cm) = old
29806 .as_ref()
29807 .map(|pp| (pp.0.len(), pp.1.len()))
29808 .unwrap_or((0, 0));
29809 if let Some(old) = old {
29810 self.fa_part_retired.lock().unwrap().push(old);
29811 }
29812 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29813 eprintln!(
29814 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29815 co, o_len, cm, ml_len
29816 );
29817 }
29818 *part_guard =
29819 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29820 }
29821 let pg = part_guard.as_mut().unwrap();
29822 self.gpu
29823 .stream()
29824 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29825 self.gpu
29826 .stream()
29827 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29828 self.gpu
29829 .stream()
29830 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29831 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29832 let f = self.func("fa_decode_vec_q_rows_v3_dc");
29833 let sh = (32 * head_dim * 2) as u32;
29834 use cudarc::driver::sys::CUfunction_attribute_enum as A;
29835 f.set_attribute(
29836 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29837 sh as i32,
29838 )?;
29839 let cfg = LaunchConfig {
29840 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
29841 block_dim: (32, gqa, 1),
29842 shared_mem_bytes: sh,
29843 };
29844 let __s_b = self.gpu.stream();
29845 let mut b = __s_b.launch_builder(&f);
29846 b.arg(q)
29847 .arg(k)
29848 .arg(v)
29849 .arg(&mut *part_o)
29850 .arg(&mut *part_m)
29851 .arg(&mut *part_l)
29852 .arg(&hd)
29853 .arg(&nh)
29854 .arg(&nhkv)
29855 .arg(base_dev)
29856 .arg(&scale)
29857 .arg(&nspm)
29858 .arg(&spk)
29859 .arg(&ktb)
29860 .arg(&vtb);
29861 unsafe {
29862 b.launch(cfg)?;
29863 }
29864 let fc = self.func("fa_decode_combine_rows_dc");
29865 let cfg2 = LaunchConfig {
29866 grid_dim: (n_head as u32, t as u32, 1),
29867 block_dim: (head_dim as u32, 1, 1),
29868 shared_mem_bytes: 0,
29869 };
29870 let plus0 = 0i32;
29871 let __s_b2 = self.gpu.stream();
29872 let mut b2 = __s_b2.launch_builder(&fc);
29873 b2.arg(&*part_o)
29874 .arg(&*part_m)
29875 .arg(&*part_l)
29876 .arg(o)
29877 .arg(&hd)
29878 .arg(&nh)
29879 .arg(base_dev)
29880 .arg(&plus0)
29881 .arg(&nspm)
29882 .arg(&spk);
29883 unsafe {
29884 b2.launch(cfg2)?;
29885 }
29886 Ok(())
29887 }
29888
29889 #[allow(clippy::too_many_arguments)] pub fn fa_decode_dc(
29901 &self,
29902 q: &CudaSlice<f32>,
29903 k: &cudarc::driver::CudaView<u8>,
29904 v: &cudarc::driver::CudaView<u8>,
29905 o: &mut CudaSlice<f32>,
29906 head_dim: usize,
29907 n_head: usize,
29908 n_head_kv: usize,
29909 t_kv_dev: &CudaSlice<i32>,
29910 bucket_max: usize,
29911 scale: f32,
29912 k_tok_bytes: usize,
29913 v_tok_bytes: usize,
29914 g: bool,
29915 ) -> Result<(), Box<dyn std::error::Error>> {
29916 self.fa_decode_dc_q8(
29917 q,
29918 k,
29919 v,
29920 o,
29921 head_dim,
29922 n_head,
29923 n_head_kv,
29924 t_kv_dev,
29925 bucket_max,
29926 scale,
29927 k_tok_bytes,
29928 v_tok_bytes,
29929 g,
29930 None,
29931 )
29932 }
29933
29934 #[allow(clippy::too_many_arguments)]
29937 #[allow(clippy::manual_div_ceil)] pub fn fa_decode_dc_q8(
29939 &self,
29940 q: &CudaSlice<f32>,
29941 k: &cudarc::driver::CudaView<u8>,
29942 v: &cudarc::driver::CudaView<u8>,
29943 o: &mut CudaSlice<f32>,
29944 head_dim: usize,
29945 n_head: usize,
29946 n_head_kv: usize,
29947 t_kv_dev: &CudaSlice<i32>,
29948 bucket_max: usize,
29949 scale: f32,
29950 k_tok_bytes: usize,
29951 v_tok_bytes: usize,
29952 g: bool,
29953 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
29954 ) -> Result<(), Box<dyn std::error::Error>> {
29955 let mut fa_vec =
29963 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
29964 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
29965 fa_vec = false;
29966 } let sp = fa_split_keys(bucket_max, n_head_kv);
29968 let n_splits = if fa_vec {
29969 ((bucket_max + sp - 1) / sp).max(1)
29970 } else {
29971 ((bucket_max + 255) / 256).max(1)
29972 };
29973 let o_len = n_head * n_splits * head_dim;
29974 let ml_len = n_head * n_splits;
29975 let mut part_guard = self.fa_part_pool.lock().unwrap();
29976 if part_guard
29977 .as_ref()
29978 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29979 .unwrap_or(true)
29980 {
29981 let old = part_guard.take();
29992 let (co, cm) = old
29993 .as_ref()
29994 .map(|pp| (pp.0.len(), pp.1.len()))
29995 .unwrap_or((0, 0));
29996 if let Some(old) = old {
29997 self.fa_part_retired.lock().unwrap().push(old);
29998 }
29999 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
30000 eprintln!(
30001 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
30002 co, o_len, cm, ml_len
30003 );
30004 }
30005 *part_guard =
30006 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30007 }
30008 let pg = part_guard.as_mut().unwrap();
30009 self.gpu
30010 .stream()
30011 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
30012 self.gpu
30013 .stream()
30014 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
30015 self.gpu
30016 .stream()
30017 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
30018 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30019 let (hd, nh, nhkv, nsp) = (
30020 head_dim as i32,
30021 n_head as i32,
30022 n_head_kv as i32,
30023 n_splits as i32,
30024 );
30025 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30026 let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
30027 let deep = fa_vec && head_dim == 256 && fa_v4_at(bucket_max) && !g;
30030 let (f, cfg) = if fa_vec
30031 && head_dim == 512
30032 && bucket_max >= {
30033 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
30034 *FA512_MIN_DC.get_or_init(|| {
30035 std::env::var("MEMRA_FA512_MIN")
30036 .ok()
30037 .and_then(|v| v.parse().ok())
30038 .unwrap_or(512)
30039 })
30040 } {
30041 let gqa = (n_head / n_head_kv).max(1) as u32;
30043 (
30044 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
30045 LaunchConfig {
30046 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30047 block_dim: (32, gqa, 1),
30048 shared_mem_bytes: 0,
30049 },
30050 )
30051 } else if fa_vec && head_dim == 512 {
30052 let q_view = q.as_view();
30055 let mut o_view = o.as_view_mut();
30056 return self.fa_decode_scalar_unified(
30057 &q_view,
30058 k,
30059 v,
30060 &mut o_view,
30061 head_dim,
30062 n_head,
30063 n_head_kv,
30064 0,
30065 Some(t_kv_dev),
30066 scale,
30067 n_splits,
30068 sp,
30069 k_tok_bytes,
30070 v_tok_bytes,
30071 g,
30072 &mut *part_o,
30073 &mut *part_m,
30074 &mut *part_l,
30075 q8_out,
30076 );
30077 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
30078 let gqa = (n_head / n_head_kv).max(1) as u32;
30081 let fv = if g {
30082 self.func_g("fa_decode_vec_q_v4_dc")
30083 } else if deep {
30084 self.func("fa_decode_vec_q_v4_deep_dc")
30085 } else {
30086 self.func("fa_decode_vec_q_v4_dc")
30087 };
30088 let shmem =
30089 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
30090 use cudarc::driver::sys::CUfunction_attribute_enum as A;
30091 fv.set_attribute(
30092 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30093 shmem as i32,
30094 )?;
30095 (
30096 fv,
30097 LaunchConfig {
30098 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30099 block_dim: (32, gqa, 1),
30100 shared_mem_bytes: shmem,
30101 },
30102 )
30103 } else if fa_vec && fa_v3_active(head_dim) {
30104 let gqa = (n_head / n_head_kv).max(1) as u32;
30107 let fv = if g {
30108 self.func_g("fa_decode_vec_q_v3_dc")
30109 } else {
30110 self.func("fa_decode_vec_q_v3_dc")
30111 };
30112 let shmem = (32 * head_dim * 2) as u32; (
30114 fv,
30115 LaunchConfig {
30116 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30117 block_dim: (32, gqa, 1),
30118 shared_mem_bytes: shmem,
30119 },
30120 )
30121 } else if fa_vec {
30122 let gqa = (n_head / n_head_kv).max(1) as u32;
30125 let fv = if g {
30126 self.func_g("fa_decode_vec_q_v2_dc")
30127 } else {
30128 self.func("fa_decode_vec_q_v2_dc")
30129 };
30130 let shmem = (2 * 32 * head_dim * 2) as u32; (
30132 fv,
30133 LaunchConfig {
30134 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30135 block_dim: (32, gqa, 1),
30136 shared_mem_bytes: shmem,
30137 },
30138 )
30139 } else {
30140 let q_view = q.as_view();
30141 let mut o_view = o.as_view_mut();
30142 return self.fa_decode_scalar_unified(
30143 &q_view,
30144 k,
30145 v,
30146 &mut o_view,
30147 head_dim,
30148 n_head,
30149 n_head_kv,
30150 0,
30151 Some(t_kv_dev),
30152 scale,
30153 n_splits,
30154 if fa_vec { sp } else { 256 },
30155 k_tok_bytes,
30156 v_tok_bytes,
30157 g,
30158 &mut *part_o,
30159 &mut *part_m,
30160 &mut *part_l,
30161 q8_out,
30162 );
30163 };
30164 let ski = sp as i32; let __s_b = self.gpu.stream();
30166 let mut b = __s_b.launch_builder(&f);
30167 b.arg(q)
30168 .arg(k)
30169 .arg(v)
30170 .arg(&mut *part_o)
30171 .arg(&mut *part_m)
30172 .arg(&mut *part_l)
30173 .arg(&hd)
30174 .arg(&nh)
30175 .arg(&nhkv)
30176 .arg(t_kv_dev)
30177 .arg(&scale)
30178 .arg(&nsp)
30179 .arg(&ski)
30180 .arg(&ktb)
30181 .arg(&vtb);
30182 unsafe {
30183 b.launch(cfg)?;
30184 }
30185 let cfg2 = LaunchConfig {
30186 grid_dim: (n_head as u32, 1, 1),
30187 block_dim: (head_dim as u32, 1, 1),
30188 shared_mem_bytes: 0,
30189 };
30190 if let Some((oq, od)) = q8_out {
30191 let fc = if g {
30192 self.func_g("fa_decode_combine_q8_1")
30193 } else {
30194 self.fa_func("fa_decode_combine_q8_1", head_dim)
30195 };
30196 let __s_b2 = self.gpu.stream();
30197 let mut b2 = __s_b2.launch_builder(&fc);
30198 b2.arg(&*part_o)
30199 .arg(&*part_m)
30200 .arg(&*part_l)
30201 .arg(oq)
30202 .arg(od)
30203 .arg(&hd)
30204 .arg(&nh)
30205 .arg(&nsp);
30206 unsafe {
30207 b2.launch(cfg2)?;
30208 }
30209 return Ok(());
30210 }
30211 let fc = if g {
30212 self.func_g("fa_decode_combine_f32")
30213 } else {
30214 self.fa_func("fa_decode_combine_f32", head_dim)
30215 };
30216 let __s_b2 = self.gpu.stream();
30217 let mut b2 = __s_b2.launch_builder(&fc);
30218 b2.arg(&*part_o)
30219 .arg(&*part_m)
30220 .arg(&*part_l)
30221 .arg(o)
30222 .arg(&hd)
30223 .arg(&nh)
30224 .arg(&nsp);
30225 unsafe {
30226 b2.launch(cfg2)?;
30227 }
30228 Ok(())
30229 }
30230
30231 #[allow(clippy::too_many_arguments)]
30235 pub fn append_kv_quantized_dcw(
30236 &self,
30237 k_row: &CudaSlice<f32>,
30238 v_row: &CudaSlice<f32>,
30239 kc: &mut CudaSlice<u8>,
30240 vc: &mut CudaSlice<u8>,
30241 len_dev: &CudaSlice<i32>,
30242 base_dev: Option<&CudaSlice<i32>>,
30243 kv_dim_k: usize,
30244 kv_dim_v: usize,
30245 k_tok_bytes: usize,
30246 v_tok_bytes: usize,
30247 ) -> Result<(), Box<dyn std::error::Error>> {
30248 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
30249 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
30250 let cfg = LaunchConfig {
30251 grid_dim: (nblk, 1, 1),
30252 block_dim: (32, 1, 1),
30253 shared_mem_bytes: 0,
30254 };
30255 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
30256 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30257 let null: u64 = 0;
30258 let __s_b = self.gpu.stream();
30259 let mut b = __s_b.launch_builder(&f);
30260 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
30261 match base_dev {
30262 Some(base) => {
30263 b.arg(base);
30264 }
30265 None => {
30266 b.arg(&null);
30267 }
30268 }
30269 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
30270 unsafe {
30271 b.launch(cfg)?;
30272 }
30273 Ok(())
30274 }
30275
30276 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
30278 let f = self.func("inc_i32");
30279 let cfg = LaunchConfig {
30280 grid_dim: (1, 1, 1),
30281 block_dim: (1, 1, 1),
30282 shared_mem_bytes: 0,
30283 };
30284 let __s_b = self.gpu.stream();
30285 let mut b = __s_b.launch_builder(&f);
30286 b.arg(counter);
30287 unsafe {
30288 b.launch(cfg)?;
30289 }
30290 Ok(())
30291 }
30292
30293 #[allow(clippy::too_many_arguments)]
30302 #[allow(clippy::type_complexity)] fn fa_part_alloc(
30327 &self,
30328 o_len: usize,
30329 ml_len: usize,
30330 co: usize,
30331 cm: usize,
30332 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
30333 static GROWS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
30334 let n = GROWS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
30335 if n < 64 {
30336 eprintln!(
30337 "[fa-pool] grow #{n} dev={} o_len {co} -> {o_len} ml_len {cm} -> {ml_len} (retired kept, zero={})",
30338 self.ctx().ordinal(),
30339 fa_part_zero_on()
30340 );
30341 }
30342 let mut po = self.alloc_uninit::<f32>(o_len)?;
30343 let mut pm = self.alloc_uninit::<f32>(ml_len)?;
30344 let mut pl = self.alloc_uninit::<f32>(ml_len)?;
30345 if fa_part_zero_on() {
30346 self.gpu.stream().memset_zeros(&mut po)?;
30347 self.gpu.stream().memset_zeros(&mut pm)?;
30348 self.gpu.stream().memset_zeros(&mut pl)?;
30349 }
30350 Ok((po, pm, pl))
30351 }
30352
30353 fn fa_part_pool_grow(
30354 &self,
30355 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
30356 o_len: usize,
30357 ml_len: usize,
30358 ) -> Result<(), Box<dyn std::error::Error>> {
30359 if part_guard
30360 .as_ref()
30361 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30362 .unwrap_or(true)
30363 {
30364 let old = part_guard.take();
30365 let (co, cm) = old
30366 .as_ref()
30367 .map(|pp| (pp.0.len(), pp.1.len()))
30368 .unwrap_or((0, 0));
30369 if let Some(old) = old {
30370 self.fa_part_retired.lock().unwrap().push(old);
30371 }
30372 *part_guard =
30385 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30386 }
30387 Ok(())
30388 }
30389
30390 pub fn fa_dcw_pool_ensure(
30393 &self,
30394 head_dim: usize,
30395 n_head: usize,
30396 n_head_kv: usize,
30397 bucket_max: usize,
30398 ) -> Result<(), Box<dyn std::error::Error>> {
30399 let sp = fa_split_keys(bucket_max, n_head_kv);
30400 #[allow(clippy::manual_div_ceil)]
30401 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
30403 let o_len = n_head * n_splits * head_dim;
30404 let ml_len = n_head * n_splits;
30405 let mut part_guard = self.fa_part_pool.lock().unwrap();
30406 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
30407 }
30408
30409 #[allow(clippy::too_many_arguments)]
30418 pub fn fa_decode_dcw_rows(
30419 &self,
30420 q_rows: &CudaSlice<f32>,
30421 tab: &CudaSlice<u64>,
30422 o_rows: &mut CudaSlice<f32>,
30423 t: usize,
30424 head_dim: usize,
30425 n_head: usize,
30426 n_head_kv: usize,
30427 window: usize,
30428 max_ns: usize,
30429 scale: f32,
30430 k_tok_bytes: usize,
30431 v_tok_bytes: usize,
30432 gate_rows: &CudaSlice<f32>,
30433 ) -> Result<(), Box<dyn std::error::Error>> {
30434 if std::env::var("MEMRA_NO_FA_VEC").is_ok()
30435 || head_dim > 256
30436 || !head_dim.is_multiple_of(32)
30437 {
30438 return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
30439 }
30440 if fa_sm_count() < 128
30441 || std::env::var("MEMRA_FA_SPLIT").is_ok()
30442 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
30443 || std::env::var("MEMRA_FA_SP16").is_ok()
30444 {
30445 return Err(
30446 "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
30447 (or a <128-SM rig) keep the per-row path"
30448 .into(),
30449 );
30450 }
30451 if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
30452 return Err("fa_decode_dcw_rows geometry".into());
30453 }
30454 let o_len = t * n_head * max_ns * head_dim;
30455 let ml_len = t * n_head * max_ns;
30456 let mut part_guard = self.fa_part_pool.lock().unwrap();
30457 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
30458 let pg = part_guard.as_mut().unwrap();
30459 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30460 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
30461 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30462 let (win, mns) = (window as i32, max_ns as i32);
30463 let gqa = (n_head / n_head_kv).max(1) as u32;
30464 let smem = (32 * head_dim * 2) as u32;
30465 let f = self.func("fa_decode_vec_q_v3_dcw_rows");
30466 let cfg = LaunchConfig {
30467 grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
30468 block_dim: (32, gqa, 1),
30469 shared_mem_bytes: smem,
30470 };
30471 {
30472 let __s_b = self.gpu.stream();
30473 let mut b = __s_b.launch_builder(&f);
30474 b.arg(q_rows)
30475 .arg(tab)
30476 .arg(&mut *part_o)
30477 .arg(&mut *part_m)
30478 .arg(&mut *part_l)
30479 .arg(&hd)
30480 .arg(&nh)
30481 .arg(&nhkv)
30482 .arg(&win)
30483 .arg(&scale)
30484 .arg(&mns)
30485 .arg(&ktb)
30486 .arg(&vtb);
30487 unsafe {
30488 b.launch(cfg)?;
30489 }
30490 }
30491 let fc = {
30495 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
30496 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
30497 self.func("fa_decode_combine_gate_f32_s")
30498 } else {
30499 self.func("fa_decode_combine_gate_f32")
30500 }
30501 };
30502 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
30503 let nht = (t * n_head) as i32;
30504 let cfg2 = LaunchConfig {
30505 grid_dim: ((t * n_head) as u32, 1, 1),
30506 block_dim: (head_dim as u32, 1, 1),
30507 shared_mem_bytes: if combine_shared {
30508 (2 * max_ns * 4) as u32
30509 } else {
30510 0
30511 },
30512 };
30513 let __s_b2 = self.gpu.stream();
30514 let mut b2 = __s_b2.launch_builder(&fc);
30515 b2.arg(&*part_o)
30516 .arg(&*part_m)
30517 .arg(&*part_l)
30518 .arg(gate_rows)
30519 .arg(o_rows)
30520 .arg(&hd)
30521 .arg(&nht)
30522 .arg(&mns);
30523 unsafe {
30524 b2.launch(cfg2)?;
30525 }
30526 Ok(())
30527 }
30528
30529 #[allow(clippy::too_many_arguments)] pub fn fa_decode_dcw(
30531 &self,
30532 q: &CudaSlice<f32>,
30533 k_ring: &cudarc::driver::CudaView<u8>,
30534 v_ring: &cudarc::driver::CudaView<u8>,
30535 o: &mut CudaSlice<f32>,
30536 head_dim: usize,
30537 n_head: usize,
30538 n_head_kv: usize,
30539 len_dev: &CudaSlice<i32>,
30540 base_dev: Option<&CudaSlice<i32>>,
30541 window: usize,
30542 bucket_max: usize,
30543 scale: f32,
30544 k_tok_bytes: usize,
30545 v_tok_bytes: usize,
30546 fused_gate: Option<&CudaSlice<f32>>,
30550 ) -> Result<(), Box<dyn std::error::Error>> {
30551 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
30552 if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) {
30553 return Err("fa_decode_dcw supports the v3-vec class only (bucket >= vec floor, head_dim <= 256, head_dim % 32 == 0)"
30554 .into());
30555 }
30556 let sp = fa_split_keys(bucket_max, n_head_kv);
30557 #[allow(clippy::manual_div_ceil)]
30558 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
30560 let o_len = n_head * n_splits * head_dim;
30561 let ml_len = n_head * n_splits;
30562 let mut part_guard = self.fa_part_pool.lock().unwrap();
30563 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
30564 let pg = part_guard.as_mut().unwrap();
30565 {
30570 self.gpu
30571 .stream()
30572 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
30573 self.gpu
30574 .stream()
30575 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
30576 self.gpu
30577 .stream()
30578 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
30579 }
30580 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30581 let (hd, nh, nhkv, nsp) = (
30582 head_dim as i32,
30583 n_head as i32,
30584 n_head_kv as i32,
30585 n_splits as i32,
30586 );
30587 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30588 let (ski, win) = (sp as i32, window as i32);
30589 let gqa = (n_head / n_head_kv).max(1) as u32;
30590 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
30594 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
30599 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
30600 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
30601 std::sync::Mutex::new(None);
30602 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
30606 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
30607 && (n_head / n_head_kv).is_multiple_of(2)
30608 && (n_head / n_head_kv) >= 2;
30609 let f = if fprof {
30610 self.func("fa_decode_vec_q_v3_dcw_prof")
30611 } else if hs2 {
30612 self.func("fa_decode_vec_q_v3_dcw_hs2")
30613 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
30614 self.func("fa_decode_vec_q_v3_dcw_u8")
30615 } else {
30616 self.func("fa_decode_vec_q_v3_dcw")
30617 };
30618 let cfg = LaunchConfig {
30619 grid_dim: if hs2 {
30620 ((2 * n_head_kv) as u32, n_splits as u32, 1)
30621 } else {
30622 (n_head_kv as u32, n_splits as u32, 1)
30623 },
30624 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
30625 shared_mem_bytes: smem,
30626 };
30627 let null: u64 = 0;
30628 let __s_b = self.gpu.stream();
30629 let mut b = __s_b.launch_builder(&f);
30630 b.arg(q)
30631 .arg(k_ring)
30632 .arg(v_ring)
30633 .arg(&mut *part_o)
30634 .arg(&mut *part_m)
30635 .arg(&mut *part_l)
30636 .arg(&hd)
30637 .arg(&nh)
30638 .arg(&nhkv)
30639 .arg(len_dev);
30640 match base_dev {
30641 Some(base) => {
30642 b.arg(base);
30643 }
30644 None => {
30645 b.arg(&null);
30646 }
30647 }
30648 b.arg(&win)
30649 .arg(&scale)
30650 .arg(&nsp)
30651 .arg(&ski)
30652 .arg(&ktb)
30653 .arg(&vtb);
30654 if fprof {
30655 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
30656 if guard
30657 .as_ref()
30658 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
30659 {
30660 *guard = Some((self.ctx().ordinal(), self.htod_u64(&[0u64; 8])?));
30661 }
30662 let (_, buf) = guard.as_mut().expect("armed above");
30663 b.arg(&*buf);
30664 unsafe {
30665 b.launch(cfg)?;
30666 }
30667 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
30668 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
30669 if n.is_multiple_of(430) {
30670 self.stream().synchronize()?;
30671 let h = self.dtoh_u64(buf)?;
30672 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
30673 let tot: u64 = h[..6].iter().sum();
30674 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
30675 for (i, name) in phases.iter().enumerate() {
30676 let pct = if tot > 0 {
30677 h[i] as f64 / tot as f64 * 100.0
30678 } else {
30679 0.0
30680 };
30681 line.push_str(&format!(" {name}={pct:.1}%"));
30682 }
30683 if h[6] > 0 {
30684 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
30685 }
30686 eprintln!("{line}");
30687 }
30688 } else {
30689 unsafe {
30690 b.launch(cfg)?;
30691 }
30692 }
30693 let mut combine_shared = false;
30694 let fc = if fused_gate.is_some() {
30695 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
30698 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
30699 combine_shared = true;
30700 self.func("fa_decode_combine_gate_f32_s")
30701 } else {
30702 self.func("fa_decode_combine_gate_f32")
30703 }
30704 } else {
30705 self.fa_func("fa_decode_combine_f32", head_dim)
30706 };
30707 let cfg2 = LaunchConfig {
30708 grid_dim: (n_head as u32, 1, 1),
30709 block_dim: (head_dim as u32, 1, 1),
30710 shared_mem_bytes: if combine_shared {
30711 (2 * n_splits * 4) as u32
30712 } else {
30713 0
30714 },
30715 };
30716 let __s_b2 = self.gpu.stream();
30717 let mut b2 = __s_b2.launch_builder(&fc);
30718 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
30719 if let Some(gate_row) = fused_gate {
30720 b2.arg(gate_row);
30721 }
30722 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
30723 unsafe {
30724 b2.launch(cfg2)?;
30725 }
30726 Ok(())
30727 }
30728
30729 #[allow(clippy::manual_div_ceil)] pub fn fa_geom_eager(
30736 &self,
30737 t_kv: usize,
30738 head_dim: usize,
30739 n_head_kv: usize,
30740 g: bool,
30741 ) -> (bool, usize) {
30742 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
30746 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
30752 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim.is_multiple_of(32));
30753 if g && head_dim == 256 && !fa_v4_at(t_kv) {
30759 fa_vec = false;
30760 }
30761 let sp = fa_split_keys(t_kv, n_head_kv);
30762 let n_splits = if fa_vec {
30763 ((t_kv + sp - 1) / sp).max(1)
30764 } else {
30765 ((t_kv + 255) / 256).max(1)
30766 };
30767 (fa_vec, n_splits)
30768 }
30769
30770 pub fn fa_bucket_key(
30776 &self,
30777 t_kv: usize,
30778 head_dim: usize,
30779 n_head_kv: usize,
30780 g: bool,
30781 ) -> (bool, usize) {
30782 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
30783 }
30784
30785 #[allow(clippy::type_complexity)] pub fn capture_graph_retained<F>(
30798 &self,
30799 step: F,
30800 ) -> Result<
30801 (
30802 cudarc::driver::CudaGraph,
30803 Vec<Box<dyn std::any::Any + Send>>,
30804 ),
30805 Box<dyn std::error::Error>,
30806 >
30807 where
30808 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
30809 {
30810 use cudarc::driver::sys::CUgraphInstantiate_flags;
30811 self.capture_graph_retained_flags(
30812 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
30813 step,
30814 )
30815 }
30816
30817 #[allow(clippy::type_complexity)] pub fn capture_graph_retained_flags<F>(
30823 &self,
30824 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
30825 mut step: F,
30826 ) -> Result<
30827 (
30828 cudarc::driver::CudaGraph,
30829 Vec<Box<dyn std::any::Any + Send>>,
30830 ),
30831 Box<dyn std::error::Error>,
30832 >
30833 where
30834 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
30835 {
30836 use cudarc::driver::sys::CUstreamCaptureMode;
30837 self.capture_keep.lock().unwrap().clear();
30845 let was_tracking = self.gpu.ctx.is_event_tracking();
30846 if was_tracking {
30847 unsafe {
30848 self.gpu.ctx.disable_event_tracking();
30849 }
30850 }
30851 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
30852 self.capture_keep_on
30853 .store(true, std::sync::atomic::Ordering::Relaxed);
30854 let w = (|| {
30855 step(self)?;
30856 step(self)
30857 })();
30858 self.capture_keep_on
30859 .store(false, std::sync::atomic::Ordering::Relaxed);
30860 w?;
30861 self.gpu.stream().synchronize()?;
30862 self.gpu
30863 .stream()
30864 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
30865 let r = step(self);
30866 let g = self.gpu.stream().end_capture(flags);
30867 r?;
30868 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
30869 graph.upload()?;
30870 Ok(graph)
30871 };
30872 let result = run();
30873 self.capture_keep_on
30874 .store(false, std::sync::atomic::Ordering::Relaxed);
30875 if was_tracking {
30876 unsafe {
30877 self.gpu.ctx.enable_event_tracking();
30878 }
30879 }
30880 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
30881 Ok((result?, keeper))
30882 }
30883
30884 #[allow(clippy::type_complexity)] pub fn capture_graph_retained_nowarm<F>(
30891 &self,
30892 mut step: F,
30893 ) -> Result<
30894 (
30895 cudarc::driver::CudaGraph,
30896 Vec<Box<dyn std::any::Any + Send>>,
30897 ),
30898 Box<dyn std::error::Error>,
30899 >
30900 where
30901 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
30902 {
30903 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
30904 let was_tracking = self.gpu.ctx.is_event_tracking();
30905 if was_tracking {
30906 unsafe {
30907 self.gpu.ctx.disable_event_tracking();
30908 }
30909 }
30910 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
30911 self.gpu.stream().synchronize()?;
30912 self.gpu
30913 .stream()
30914 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
30915 let r = step(self);
30916 let g = self.gpu.stream().end_capture(
30917 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
30918 );
30919 r?;
30920 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
30921 graph.upload()?;
30922 Ok(graph)
30923 };
30924 let result = run();
30925 if was_tracking {
30926 unsafe {
30927 self.gpu.ctx.enable_event_tracking();
30928 }
30929 }
30930 Ok((result?, Vec::new()))
30931 }
30932
30933 pub fn capture_graph<F>(
30934 &self,
30935 mut step: F,
30936 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
30937 where
30938 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
30939 {
30940 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
30941 let was_tracking = self.gpu.ctx.is_event_tracking();
30949 if was_tracking {
30950 unsafe {
30951 self.gpu.ctx.disable_event_tracking();
30952 }
30953 }
30954 let warmups = {
30977 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
30978 *W.get_or_init(|| {
30979 std::env::var("MEMRA_GRAPH_WARMUPS")
30980 .ok()
30981 .and_then(|v| v.parse().ok())
30982 .filter(|n| *n >= 1)
30983 .unwrap_or(1)
30984 })
30985 };
30986 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
30987 for _ in 0..warmups {
30989 step(self)?;
30990 }
30991 self.gpu.stream().synchronize()?;
30992 self.gpu
30994 .stream()
30995 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
30996 let r = step(self);
30999 let g = self.gpu.stream().end_capture(
31000 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
31001 );
31002 r?;
31003 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
31004 graph.upload()?;
31005 Ok(graph)
31006 };
31007 let result = run();
31008 if was_tracking {
31009 unsafe {
31010 self.gpu.ctx.enable_event_tracking();
31011 }
31012 }
31013 result
31014 }
31015
31016 #[allow(clippy::too_many_arguments)] pub fn gdn_scan_s128_view(
31019 &self,
31020 q: &CudaSlice<f32>,
31021 k: &CudaSlice<f32>,
31022 v: &CudaSlice<f32>,
31023 g: &CudaSlice<f32>,
31024 beta: &CudaSlice<f32>,
31025 state_in: &cudarc::driver::CudaView<f32>,
31026 state_out: &mut cudarc::driver::CudaViewMut<f32>,
31027 o: &mut CudaSlice<f32>,
31028 n_head: usize,
31029 t: usize,
31030 scale: f32,
31031 ) -> Result<(), Box<dyn std::error::Error>> {
31032 let f = self.func("gdn_scan_s128");
31033 const S_V: u32 = 128;
31034 const WARP: u32 = 32;
31035 const COLS: u32 = 4;
31036 let cfg = LaunchConfig {
31037 grid_dim: (n_head as u32, 1, S_V / COLS),
31038 block_dim: (WARP, COLS, 1),
31039 shared_mem_bytes: 0,
31040 };
31041 let (h, ti) = (n_head as i32, t as i32);
31042 let __s_b = self.gpu.stream();
31043 let mut b = __s_b.launch_builder(&f);
31044 b.arg(q)
31045 .arg(k)
31046 .arg(v)
31047 .arg(g)
31048 .arg(beta)
31049 .arg(state_in)
31050 .arg(state_out)
31051 .arg(o)
31052 .arg(&h)
31053 .arg(&ti)
31054 .arg(&scale);
31055 unsafe {
31056 b.launch(cfg)?;
31057 }
31058 Ok(())
31059 }
31060
31061 #[allow(clippy::too_many_arguments)]
31063 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_view(
31066 &self,
31067 x: &cudarc::driver::CudaView<f32>,
31068 w: &CudaSlice<f32>,
31069 y: &mut CudaSlice<f32>,
31070 conv_dim: usize,
31071 t: usize,
31072 d_conv: usize,
31073 silu: bool,
31074 ) -> Result<(), Box<dyn std::error::Error>> {
31075 let f = self.func("ssm_conv1d_silu_f32");
31076 let cfg = LaunchConfig {
31078 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
31079 block_dim: (256, 1, 1),
31080 shared_mem_bytes: 0,
31081 };
31082 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
31083 let __s_b = self.gpu.stream();
31084 let mut b = __s_b.launch_builder(&f);
31085 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
31086 unsafe {
31087 b.launch(cfg)?;
31088 }
31089 Ok(())
31090 }
31091
31092 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_tm(
31100 &self,
31101 qkv_tm: &CudaSlice<f32>,
31102 w: &CudaSlice<f32>,
31103 y: &mut CudaSlice<f32>,
31104 conv_dim: usize,
31105 t: usize,
31106 d_conv: usize,
31107 ) -> Result<(), Box<dyn std::error::Error>> {
31108 let f = self.func("ssm_conv1d_tm_f32");
31109 let cfg = LaunchConfig {
31110 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31111 block_dim: (256, 1, 1),
31112 shared_mem_bytes: 0,
31113 };
31114 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31115 let __s_b = self.gpu.stream();
31116 let mut b = __s_b.launch_builder(&f);
31117 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
31118 unsafe {
31119 b.launch(cfg)?;
31120 }
31121 Ok(())
31122 }
31123
31124 #[allow(clippy::too_many_arguments)] pub fn ssm_conv1d_tm_state(
31133 &self,
31134 qkv_tm: &CudaSlice<f32>,
31135 conv_state: &mut CudaSlice<f32>,
31136 w: &CudaSlice<f32>,
31137 y: &mut CudaSlice<f32>,
31138 conv_dim: usize,
31139 t: usize,
31140 d_conv: usize,
31141 ) -> Result<(), Box<dyn std::error::Error>> {
31142 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
31143 }
31144
31145 #[allow(clippy::too_many_arguments)]
31148 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_tm_state_pad(
31150 &self,
31151 qkv_tm: &CudaSlice<f32>,
31152 conv_state: &mut CudaSlice<f32>,
31153 w: &CudaSlice<f32>,
31154 y: &mut CudaSlice<f32>,
31155 conv_dim: usize,
31156 t: usize,
31157 d_conv: usize,
31158 pad_len: Option<&CudaSlice<i32>>,
31159 ) -> Result<(), Box<dyn std::error::Error>> {
31160 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
31161 let ring_old = if t < d_conv - 1 {
31165 Some(self.clone_dtod(conv_state)?)
31166 } else {
31167 None
31168 };
31169 {
31170 let f = self.func("ssm_conv1d_tm_state_f32");
31171 let cfg = LaunchConfig {
31172 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31173 block_dim: (256, 1, 1),
31174 shared_mem_bytes: 0,
31175 };
31176 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31177 let __s_b = self.gpu.stream();
31178 let mut b = __s_b.launch_builder(&f);
31179 b.arg(qkv_tm)
31180 .arg(&*conv_state)
31181 .arg(w)
31182 .arg(y)
31183 .arg(&cd)
31184 .arg(&ti)
31185 .arg(&dc);
31186 unsafe {
31187 b.launch(cfg)?;
31188 }
31189 }
31190 match (ring_old, pad_len) {
31191 (None, Some(len_d)) => {
31192 let f = self.func("ssm_conv_ring_update_dev_f32");
31193 let n = conv_dim * (d_conv - 1);
31194 let cfg = LaunchConfig::for_num_elems(n as u32);
31195 let (cd, dc) = (conv_dim as i32, d_conv as i32);
31196 let __s_b = self.gpu.stream();
31197 let mut b = __s_b.launch_builder(&f);
31198 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
31199 unsafe {
31200 b.launch(cfg)?;
31201 }
31202 }
31203 (None, None) => {
31204 let f = self.func("ssm_conv_ring_update_f32");
31205 let n = conv_dim * (d_conv - 1);
31206 let cfg = LaunchConfig::for_num_elems(n as u32);
31207 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31208 let __s_b = self.gpu.stream();
31209 let mut b = __s_b.launch_builder(&f);
31210 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
31211 unsafe {
31212 b.launch(cfg)?;
31213 }
31214 }
31215 (Some(old), _) => {
31216 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
31217 }
31218 }
31219 Ok(())
31220 }
31221
31222 #[allow(clippy::too_many_arguments)]
31224 pub fn ssm_conv_ring_rebuild(
31230 &self,
31231 qkv_tm: &CudaSlice<f32>,
31232 ring_old: &CudaSlice<f32>,
31233 conv_state: &mut CudaSlice<f32>,
31234 conv_dim: usize,
31235 tc: usize,
31236 d_conv: usize,
31237 ) -> Result<(), Box<dyn std::error::Error>> {
31238 let f = self.func("ssm_conv_ring_rebuild_f32");
31239 let n = conv_dim * (d_conv - 1);
31240 let cfg = LaunchConfig::for_num_elems(n as u32);
31241 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
31242 let __s_b = self.gpu.stream();
31243 let mut b = __s_b.launch_builder(&f);
31244 b.arg(qkv_tm)
31245 .arg(ring_old)
31246 .arg(conv_state)
31247 .arg(&cd)
31248 .arg(&ti)
31249 .arg(&dc);
31250 unsafe {
31251 b.launch(cfg)?;
31252 }
31253 Ok(())
31254 }
31255
31256 #[allow(clippy::too_many_arguments)]
31261 pub fn gdn_prep_decode(
31262 &self,
31263 conv_out: &CudaSlice<f32>,
31264 beta_raw: &CudaSlice<f32>,
31265 alpha: &CudaSlice<f32>,
31266 dt_bias: &CudaSlice<f32>,
31267 a: &CudaSlice<f32>,
31268 q_l2: &mut CudaSlice<f32>,
31269 k_l2: &mut CudaSlice<f32>,
31270 v_g: &mut CudaSlice<f32>,
31271 beta: &mut CudaSlice<f32>,
31272 g_log: &mut CudaSlice<f32>,
31273 d_state: usize,
31274 num_v: usize,
31275 num_k: usize,
31276 key_dim: usize,
31277 eps: f32,
31278 ) -> Result<(), Box<dyn std::error::Error>> {
31279 let f = self.func("gdn_prep_decode_f32");
31280 let cfg = LaunchConfig {
31281 grid_dim: (num_v as u32, 1, 1),
31282 block_dim: (32, 4, 1),
31283 shared_mem_bytes: 0,
31284 };
31285 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
31286 let __s_b = self.gpu.stream();
31287 let mut b = __s_b.launch_builder(&f);
31288 b.arg(conv_out)
31289 .arg(beta_raw)
31290 .arg(alpha)
31291 .arg(dt_bias)
31292 .arg(a)
31293 .arg(q_l2)
31294 .arg(k_l2)
31295 .arg(v_g)
31296 .arg(beta)
31297 .arg(g_log)
31298 .arg(&ds)
31299 .arg(&nv)
31300 .arg(&nk)
31301 .arg(&kd)
31302 .arg(&eps);
31303 unsafe {
31304 b.launch(cfg)?;
31305 }
31306 Ok(())
31307 }
31308
31309 #[allow(clippy::too_many_arguments)]
31313 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_gdn(
31315 &self,
31316 qkv_tm: &CudaSlice<f32>,
31317 w: &CudaSlice<f32>,
31318 q_g: &mut CudaSlice<f32>,
31319 k_g: &mut CudaSlice<f32>,
31320 v_g: &mut CudaSlice<f32>,
31321 conv_dim: usize,
31322 t: usize,
31323 d_conv: usize,
31324 d_state: usize,
31325 num_v: usize,
31326 num_k: usize,
31327 key_dim: usize,
31328 ) -> Result<(), Box<dyn std::error::Error>> {
31329 let f = self.func("ssm_conv1d_gdn_f32");
31330 let cfg = LaunchConfig {
31331 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31332 block_dim: (256, 1, 1),
31333 shared_mem_bytes: 0,
31334 };
31335 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31336 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
31337 let __s_b = self.gpu.stream();
31338 let mut b = __s_b.launch_builder(&f);
31339 b.arg(qkv_tm)
31340 .arg(w)
31341 .arg(q_g)
31342 .arg(k_g)
31343 .arg(v_g)
31344 .arg(&cd)
31345 .arg(&ti)
31346 .arg(&dc)
31347 .arg(&ds)
31348 .arg(&nv)
31349 .arg(&nk)
31350 .arg(&kd);
31351 unsafe {
31352 b.launch(cfg)?;
31353 }
31354 Ok(())
31355 }
31356
31357 #[allow(clippy::too_many_arguments)]
31358 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d(
31361 &self,
31362 x: &CudaSlice<f32>,
31363 w: &CudaSlice<f32>,
31364 y: &mut CudaSlice<f32>,
31365 conv_dim: usize,
31366 t: usize,
31367 d_conv: usize,
31368 silu: bool,
31369 ) -> Result<(), Box<dyn std::error::Error>> {
31370 let f = self.func("ssm_conv1d_silu_f32");
31371 let cfg = LaunchConfig {
31372 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
31373 block_dim: (256, 1, 1),
31374 shared_mem_bytes: 0,
31375 };
31376 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
31377 let __s_b = self.gpu.stream();
31378 let mut b = __s_b.launch_builder(&f);
31379 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
31380 unsafe {
31381 b.launch(cfg)?;
31382 }
31383 Ok(())
31384 }
31385
31386 #[allow(clippy::too_many_arguments)] pub fn gdn_scan_s128(
31390 &self,
31391 q: &CudaSlice<f32>,
31392 k: &CudaSlice<f32>,
31393 v: &CudaSlice<f32>,
31394 g: &CudaSlice<f32>,
31395 beta: &CudaSlice<f32>,
31396 state_in: &CudaSlice<f32>,
31397 state_out: &mut CudaSlice<f32>,
31398 o: &mut CudaSlice<f32>,
31399 n_head: usize,
31400 t: usize,
31401 scale: f32,
31402 ) -> Result<(), Box<dyn std::error::Error>> {
31403 let f = self.func("gdn_scan_s128");
31404 const S_V: u32 = 128;
31405 const WARP: u32 = 32;
31406 const COLS_PER_BLOCK: u32 = 4;
31407 let cfg = LaunchConfig {
31408 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
31409 block_dim: (WARP, COLS_PER_BLOCK, 1),
31410 shared_mem_bytes: 0,
31411 };
31412 let (h, ti) = (n_head as i32, t as i32);
31413 let __s_b = self.gpu.stream();
31414 let mut b = __s_b.launch_builder(&f);
31415 b.arg(q)
31416 .arg(k)
31417 .arg(v)
31418 .arg(g)
31419 .arg(beta)
31420 .arg(state_in)
31421 .arg(state_out)
31422 .arg(o)
31423 .arg(&h)
31424 .arg(&ti)
31425 .arg(&scale);
31426 unsafe {
31427 b.launch(cfg)?;
31428 }
31429 Ok(())
31430 }
31431
31432 #[allow(clippy::too_many_arguments)]
31437 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_fused_decode_b(
31439 &self,
31440 qkv_cols: &CudaSlice<f32>,
31441 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
31442 w: &CudaSlice<f32>,
31443 conv_outs: &mut CudaSlice<f32>,
31444 conv_dim: usize,
31445 d_conv: usize,
31446 b_n: usize,
31447 ) -> Result<(), Box<dyn std::error::Error>> {
31448 let f = self.func("ssm_conv1d_fused_decode_b_f32");
31449 let cfg = LaunchConfig {
31450 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
31451 block_dim: (256, 1, 1),
31452 shared_mem_bytes: 0,
31453 };
31454 let (cd, dc) = (conv_dim as i32, d_conv as i32);
31455 let __s_b = self.gpu.stream();
31456 let mut b = __s_b.launch_builder(&f);
31457 b.arg(qkv_cols)
31458 .arg(conv_state_ptrs)
31459 .arg(w)
31460 .arg(conv_outs)
31461 .arg(&cd)
31462 .arg(&dc);
31463 unsafe {
31464 b.launch(cfg)?;
31465 }
31466 Ok(())
31467 }
31468
31469 #[allow(clippy::too_many_arguments)]
31470 pub fn gdn_prep_decode_b(
31471 &self,
31472 conv_outs: &CudaSlice<f32>,
31473 beta_raws: &CudaSlice<f32>,
31474 alphas: &CudaSlice<f32>,
31475 dt_bias: &CudaSlice<f32>,
31476 a: &CudaSlice<f32>,
31477 q_l2: &mut CudaSlice<f32>,
31478 k_l2: &mut CudaSlice<f32>,
31479 v_g: &mut CudaSlice<f32>,
31480 beta: &mut CudaSlice<f32>,
31481 g_log: &mut CudaSlice<f32>,
31482 d_state: usize,
31483 num_v: usize,
31484 num_k: usize,
31485 key_dim: usize,
31486 eps: f32,
31487 conv_dim: usize,
31488 b_n: usize,
31489 ) -> Result<(), Box<dyn std::error::Error>> {
31490 let f = self.func("gdn_prep_decode_b_f32");
31491 let cfg = LaunchConfig {
31492 grid_dim: (num_v as u32, 1, b_n as u32),
31493 block_dim: (32, 4, 1),
31494 shared_mem_bytes: 0,
31495 };
31496 let (ds, nv, nk, kd, cd) = (
31497 d_state as i32,
31498 num_v as i32,
31499 num_k as i32,
31500 key_dim as i32,
31501 conv_dim as i32,
31502 );
31503 let __s_b = self.gpu.stream();
31504 let mut b = __s_b.launch_builder(&f);
31505 b.arg(conv_outs)
31506 .arg(beta_raws)
31507 .arg(alphas)
31508 .arg(dt_bias)
31509 .arg(a)
31510 .arg(q_l2)
31511 .arg(k_l2)
31512 .arg(v_g)
31513 .arg(beta)
31514 .arg(g_log)
31515 .arg(&ds)
31516 .arg(&nv)
31517 .arg(&nk)
31518 .arg(&kd)
31519 .arg(&eps)
31520 .arg(&cd);
31521 unsafe {
31522 b.launch(cfg)?;
31523 }
31524 Ok(())
31525 }
31526
31527 #[allow(clippy::too_many_arguments)]
31528 pub fn gdn_scan_s128_batched(
31529 &self,
31530 q: &CudaSlice<f32>,
31531 k: &CudaSlice<f32>,
31532 v: &CudaSlice<f32>,
31533 g: &CudaSlice<f32>,
31534 beta: &CudaSlice<f32>,
31535 state_in_ptrs: &cudarc::driver::CudaView<u64>,
31536 state_out_ptrs: &cudarc::driver::CudaView<u64>,
31537 o: &mut CudaSlice<f32>,
31538 n_head: usize,
31539 b_n: usize,
31540 scale: f32,
31541 ) -> Result<(), Box<dyn std::error::Error>> {
31542 let f = self.func("gdn_scan_s128_b");
31543 const S_V: u32 = 128;
31544 const WARP: u32 = 32;
31545 const COLS_PER_BLOCK: u32 = 4;
31546 let cfg = LaunchConfig {
31547 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
31548 block_dim: (WARP, COLS_PER_BLOCK, 1),
31549 shared_mem_bytes: 0,
31550 };
31551 let h = n_head as i32;
31552 let __s_b = self.gpu.stream();
31553 let mut b = __s_b.launch_builder(&f);
31554 b.arg(q)
31555 .arg(k)
31556 .arg(v)
31557 .arg(g)
31558 .arg(beta)
31559 .arg(state_in_ptrs)
31560 .arg(state_out_ptrs)
31561 .arg(o)
31562 .arg(&h)
31563 .arg(&scale);
31564 unsafe {
31565 b.launch(cfg)?;
31566 }
31567 Ok(())
31568 }
31569
31570 #[allow(clippy::too_many_arguments)]
31576 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_fused_decode_b_view(
31578 &self,
31579 qkv_cols: &cudarc::driver::CudaView<f32>,
31580 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
31581 w: &CudaSlice<f32>,
31582 conv_outs: &mut CudaSlice<f32>,
31583 conv_dim: usize,
31584 d_conv: usize,
31585 b_n: usize,
31586 ) -> Result<(), Box<dyn std::error::Error>> {
31587 let f = self.func("ssm_conv1d_fused_decode_b_f32");
31588 let cfg = LaunchConfig {
31589 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
31590 block_dim: (256, 1, 1),
31591 shared_mem_bytes: 0,
31592 };
31593 let (cd, dc) = (conv_dim as i32, d_conv as i32);
31594 let __s_b = self.gpu.stream();
31595 let mut b = __s_b.launch_builder(&f);
31596 b.arg(qkv_cols)
31597 .arg(conv_state_ptrs)
31598 .arg(w)
31599 .arg(conv_outs)
31600 .arg(&cd)
31601 .arg(&dc);
31602 unsafe {
31603 b.launch(cfg)?;
31604 }
31605 Ok(())
31606 }
31607
31608 #[allow(clippy::too_many_arguments)]
31609 pub fn gdn_prep_decode_b_view(
31610 &self,
31611 conv_outs: &CudaSlice<f32>,
31612 beta_raws: &cudarc::driver::CudaView<f32>,
31613 alphas: &cudarc::driver::CudaView<f32>,
31614 dt_bias: &CudaSlice<f32>,
31615 a: &CudaSlice<f32>,
31616 q_l2: &mut CudaSlice<f32>,
31617 k_l2: &mut CudaSlice<f32>,
31618 v_g: &mut CudaSlice<f32>,
31619 beta: &mut CudaSlice<f32>,
31620 g_log: &mut CudaSlice<f32>,
31621 d_state: usize,
31622 num_v: usize,
31623 num_k: usize,
31624 key_dim: usize,
31625 eps: f32,
31626 conv_dim: usize,
31627 b_n: usize,
31628 ) -> Result<(), Box<dyn std::error::Error>> {
31629 let f = self.func("gdn_prep_decode_b_f32");
31630 let cfg = LaunchConfig {
31631 grid_dim: (num_v as u32, 1, b_n as u32),
31632 block_dim: (32, 4, 1),
31633 shared_mem_bytes: 0,
31634 };
31635 let (ds, nv, nk, kd, cd) = (
31636 d_state as i32,
31637 num_v as i32,
31638 num_k as i32,
31639 key_dim as i32,
31640 conv_dim as i32,
31641 );
31642 let __s_b = self.gpu.stream();
31643 let mut b = __s_b.launch_builder(&f);
31644 b.arg(conv_outs)
31645 .arg(beta_raws)
31646 .arg(alphas)
31647 .arg(dt_bias)
31648 .arg(a)
31649 .arg(q_l2)
31650 .arg(k_l2)
31651 .arg(v_g)
31652 .arg(beta)
31653 .arg(g_log)
31654 .arg(&ds)
31655 .arg(&nv)
31656 .arg(&nk)
31657 .arg(&kd)
31658 .arg(&eps)
31659 .arg(&cd);
31660 unsafe {
31661 b.launch(cfg)?;
31662 }
31663 Ok(())
31664 }
31665
31666 #[allow(clippy::too_many_arguments)]
31667 pub fn gdn_scan_s128_batched_view(
31668 &self,
31669 q: &CudaSlice<f32>,
31670 k: &CudaSlice<f32>,
31671 v: &CudaSlice<f32>,
31672 g: &CudaSlice<f32>,
31673 beta: &CudaSlice<f32>,
31674 state_in_ptrs: &cudarc::driver::CudaView<u64>,
31675 state_out_ptrs: &cudarc::driver::CudaView<u64>,
31676 o: &mut cudarc::driver::CudaViewMut<f32>,
31677 n_head: usize,
31678 b_n: usize,
31679 scale: f32,
31680 ) -> Result<(), Box<dyn std::error::Error>> {
31681 let f = self.func("gdn_scan_s128_b");
31682 const S_V: u32 = 128;
31683 const WARP: u32 = 32;
31684 const COLS_PER_BLOCK: u32 = 4;
31685 let cfg = LaunchConfig {
31686 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
31687 block_dim: (WARP, COLS_PER_BLOCK, 1),
31688 shared_mem_bytes: 0,
31689 };
31690 let h = n_head as i32;
31691 let __s_b = self.gpu.stream();
31692 let mut b = __s_b.launch_builder(&f);
31693 b.arg(q)
31694 .arg(k)
31695 .arg(v)
31696 .arg(g)
31697 .arg(beta)
31698 .arg(state_in_ptrs)
31699 .arg(state_out_ptrs)
31700 .arg(o)
31701 .arg(&h)
31702 .arg(&scale);
31703 unsafe {
31704 b.launch(cfg)?;
31705 }
31706 Ok(())
31707 }
31708
31709 pub fn gdn_chunked_enabled() -> bool {
31718 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31719 *E.get_or_init(|| {
31720 std::env::var("MEMRA_GDN_CHUNKED")
31721 .map(|v| v != "0")
31722 .unwrap_or(true)
31723 })
31724 }
31725
31726 pub fn gdn_chunk_size() -> usize {
31731 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
31732 *C.get_or_init(|| {
31733 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
31734 .ok()
31735 .and_then(|v| v.parse().ok())
31736 .unwrap_or(32);
31737 c.clamp(32, 128) / 32 * 32
31738 })
31739 }
31740
31741 #[allow(clippy::too_many_arguments)]
31746 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
31749 #[allow(clippy::too_many_arguments)]
31750 pub fn gdn_chunk_k123(
31751 &self,
31752 q: &CudaSlice<f32>,
31753 k: &CudaSlice<f32>,
31754 v: &CudaSlice<f32>,
31755 g: &CudaSlice<f32>,
31756 beta: &CudaSlice<f32>,
31757 wb16: Option<&mut CudaSlice<u8>>,
31758 n_head: usize,
31759 t: usize,
31760 c: usize,
31761 hk: usize,
31762 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
31763 ) -> Result<
31764 (
31765 CudaSlice<f32>,
31766 CudaSlice<f32>,
31767 CudaSlice<f32>,
31768 CudaSlice<f32>,
31769 ),
31770 Box<dyn std::error::Error>,
31771 > {
31772 const D: usize = 128;
31773 let h = n_head;
31774 #[allow(clippy::manual_div_ceil)]
31775 let nc = (t + c - 1) / c;
31777 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
31778 let mut gcum = self.uninit(t * h)?;
31779 let mut a = self.uninit(nc * h * c * c)?;
31780 let mut p = self.uninit(nc * h * c * c)?;
31781 let mut u = self.uninit(nc * h * c * D)?;
31782 let mut w = self.uninit(nc * h * c * D)?;
31783 {
31784 let f = self.func("gdn_chunk_cumgate_f32");
31786 let cfg = LaunchConfig {
31787 grid_dim: (nc as u32, h as u32, 1),
31788 block_dim: (32, 1, 1),
31789 shared_mem_bytes: 0,
31790 };
31791 let __s_b = self.gpu.stream();
31792 let mut b = __s_b.launch_builder(&f);
31793 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
31794 unsafe {
31795 b.launch(cfg)?;
31796 }
31797 }
31798 if let Some((qb, kb, pb)) = k2w {
31799 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
31802 let f = self.func("gdn_k2_wgmma");
31803 let cfg = LaunchConfig {
31804 grid_dim: (nc as u32, h as u32, 1),
31805 block_dim: (128, 1, 1),
31806 shared_mem_bytes: 0,
31807 };
31808 let hki = hk as i32;
31809 let __s_b = self.gpu.stream();
31810 let mut b = __s_b.launch_builder(&f);
31811 b.arg(qb)
31812 .arg(kb)
31813 .arg(&gcum)
31814 .arg(beta)
31815 .arg(&mut a)
31816 .arg(&mut *pb)
31817 .arg(&hi)
31818 .arg(&ti)
31819 .arg(&ci)
31820 .arg(&hki);
31821 unsafe {
31822 b.launch(cfg)?;
31823 }
31824 } else if c <= 64 && !portable_mma_gated() {
31825 let f = self.func("gdn_chunk_attn_f32");
31827 f.set_attribute(
31828 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
31829 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
31830 )?;
31831 #[allow(clippy::manual_div_ceil)]
31832 let jt = ((c + 31) / 32) as u32;
31834 let cfg = LaunchConfig {
31835 grid_dim: (nc as u32, h as u32, jt),
31836 block_dim: (256, 1, 1),
31837 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
31838 };
31839 let hki = hk as i32;
31840 let __s_b = self.gpu.stream();
31841 let mut b = __s_b.launch_builder(&f);
31842 b.arg(q)
31843 .arg(k)
31844 .arg(&gcum)
31845 .arg(beta)
31846 .arg(&mut a)
31847 .arg(&mut p)
31848 .arg(&hi)
31849 .arg(&ti)
31850 .arg(&ci)
31851 .arg(&hki);
31852 unsafe {
31853 b.launch(cfg)?;
31854 }
31855 } else {
31856 assert!(
31858 hk == h,
31859 "generic K2 is broadcast-only (de-broadcast rides C==32)"
31860 );
31861 let f = self.func("gdn_chunk_attn_g_f32");
31862 let cfg = LaunchConfig {
31863 grid_dim: (nc as u32, h as u32, 1),
31864 block_dim: (32, 8, 1),
31865 shared_mem_bytes: 0,
31866 };
31867 let __s_b = self.gpu.stream();
31868 let mut b = __s_b.launch_builder(&f);
31869 b.arg(q)
31870 .arg(k)
31871 .arg(&gcum)
31872 .arg(beta)
31873 .arg(&mut a)
31874 .arg(&mut p)
31875 .arg(&hi)
31876 .arg(&ti)
31877 .arg(&ci);
31878 unsafe {
31879 b.launch(cfg)?;
31880 }
31881 }
31882 {
31883 let cfg = LaunchConfig {
31885 grid_dim: (nc as u32, h as u32, 1),
31886 block_dim: (256, 1, 1),
31887 shared_mem_bytes: 0,
31888 };
31889 match c {
31890 32 | 64 => {
31891 let f = self.func(if c == 32 {
31892 "gdn_chunk_solve32_f32"
31893 } else {
31894 "gdn_chunk_solve64_f32"
31895 });
31896 let wb: u64 = match wb16 {
31898 Some(d) => self.addr_u8(d),
31899 None => 0,
31900 };
31901 let hki = hk as i32;
31902 let __s_b = self.gpu.stream();
31903 let mut b = __s_b.launch_builder(&f);
31904 b.arg(v)
31905 .arg(k)
31906 .arg(&a)
31907 .arg(&gcum)
31908 .arg(&mut u)
31909 .arg(&mut w)
31910 .arg(&wb)
31911 .arg(&hi)
31912 .arg(&ti)
31913 .arg(&hki);
31914 unsafe {
31915 b.launch(cfg)?;
31916 }
31917 }
31918 _ => {
31919 assert!(hk == h, "generic K3 is broadcast-only");
31920 let f = self.func("gdn_chunk_solve_f32");
31921 let __s_b = self.gpu.stream();
31922 let mut b = __s_b.launch_builder(&f);
31923 b.arg(v)
31924 .arg(k)
31925 .arg(&a)
31926 .arg(&gcum)
31927 .arg(&mut u)
31928 .arg(&mut w)
31929 .arg(&hi)
31930 .arg(&ti)
31931 .arg(&ci);
31932 unsafe {
31933 b.launch(cfg)?;
31934 }
31935 }
31936 }
31937 }
31938 Ok((gcum, p, u, w))
31939 }
31940
31941 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
31950 !portable_mma_gated()
31951 && c == 32
31952 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
31953 Ok("1") => true,
31954 Ok("0") => false,
31955 _ => gdn_mma_default_on(),
31956 }
31957 }
31958
31959 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
31966 cfg!(memra_hopper_mma)
31967 && self.gdn_mma_enabled(c)
31968 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
31969 }
31970
31971 #[allow(clippy::too_many_arguments)]
31976 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_gdn_state_pad(
31978 &self,
31979 qkv_tm: &cudarc::driver::CudaView<f32>,
31980 conv_state: &mut CudaSlice<f32>,
31981 w: &CudaSlice<f32>,
31982 q_g: &mut CudaSlice<f32>,
31983 k_g: &mut CudaSlice<f32>,
31984 v_g: &mut CudaSlice<f32>,
31985 conv_dim: usize,
31986 t: usize,
31987 d_conv: usize,
31988 d_state: usize,
31989 num_v: usize,
31990 num_k: usize,
31991 key_dim: usize,
31992 hk: usize,
31993 pad_len: Option<&CudaSlice<i32>>,
31994 ) -> Result<(), Box<dyn std::error::Error>> {
31995 assert!(
31996 t >= d_conv - 1,
31997 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
31998 );
31999 {
32000 let f = self.func("ssm_conv1d_gdn_state_f32");
32001 let cfg = LaunchConfig {
32002 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
32003 block_dim: (256, 1, 1),
32004 shared_mem_bytes: 0,
32005 };
32006 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
32007 let (ds, nv, nk, kd, hki) = (
32008 d_state as i32,
32009 num_v as i32,
32010 num_k as i32,
32011 key_dim as i32,
32012 hk as i32,
32013 );
32014 let __s_b = self.gpu.stream();
32015 let mut b = __s_b.launch_builder(&f);
32016 b.arg(qkv_tm)
32017 .arg(&*conv_state)
32018 .arg(w)
32019 .arg(q_g)
32020 .arg(k_g)
32021 .arg(v_g)
32022 .arg(&cd)
32023 .arg(&ti)
32024 .arg(&dc)
32025 .arg(&ds)
32026 .arg(&nv)
32027 .arg(&nk)
32028 .arg(&kd)
32029 .arg(&hki);
32030 unsafe {
32031 b.launch(cfg)?;
32032 }
32033 }
32034 match pad_len {
32035 Some(len_d) => {
32036 let f = self.func("ssm_conv_ring_update_dev_f32");
32037 let n = conv_dim * (d_conv - 1);
32038 let cfg = LaunchConfig::for_num_elems(n as u32);
32039 let (cd, dc) = (conv_dim as i32, d_conv as i32);
32040 let __s_b = self.gpu.stream();
32041 let mut b = __s_b.launch_builder(&f);
32042 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
32043 unsafe {
32044 b.launch(cfg)?;
32045 }
32046 }
32047 None => {
32048 let f = self.func("ssm_conv_ring_update_f32");
32049 let n = conv_dim * (d_conv - 1);
32050 let cfg = LaunchConfig::for_num_elems(n as u32);
32051 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
32052 let __s_b = self.gpu.stream();
32053 let mut b = __s_b.launch_builder(&f);
32054 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
32055 unsafe {
32056 b.launch(cfg)?;
32057 }
32058 }
32059 }
32060 Ok(())
32061 }
32062
32063 pub fn gdn_chunk_alloc(
32067 &self,
32068 n_head: usize,
32069 t: usize,
32070 c: usize,
32071 hk: usize,
32072 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
32073 const D: usize = 128;
32074 assert!(
32075 c == 32,
32076 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
32077 );
32078 let h = n_head;
32079 #[allow(clippy::manual_div_ceil)]
32080 let nc = (t + c - 1) / c;
32082 Ok(GdnChunkBufs {
32083 gcum: self.uninit(t * h)?,
32084 a: self.uninit(nc * h * c * c)?,
32085 p: self.uninit(nc * h * c * c)?,
32086 u: self.uninit(nc * h * c * D)?,
32087 w: self.uninit(nc * h * c * D)?,
32088 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
32089 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
32090 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
32091 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
32092 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
32093 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
32094 o: self.uninit(D * h * t)?,
32095 t,
32096 nc,
32097 })
32098 }
32099
32100 pub fn f32_to_bf16_v(
32102 &self,
32103 x: &cudarc::driver::CudaView<f32>,
32104 dst: &mut CudaSlice<u8>,
32105 n: usize,
32106 ) -> Result<(), Box<dyn std::error::Error>> {
32107 let f = self.func("f32_to_bf16_bulk");
32108 let ni = n as i64;
32109 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
32110 let __s_b = self.gpu.stream();
32111 let mut b = __s_b.launch_builder(&f);
32112 b.arg(x).arg(dst).arg(&ni);
32113 unsafe {
32114 b.launch(cfg)?;
32115 }
32116 Ok(())
32117 }
32118
32119 pub fn f32_to_bf16_into(
32121 &self,
32122 x: &CudaSlice<f32>,
32123 dst: &mut CudaSlice<u8>,
32124 n: usize,
32125 ) -> Result<(), Box<dyn std::error::Error>> {
32126 let f = self.func("f32_to_bf16_bulk");
32127 let ni = n as i64;
32128 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
32129 let __s_b = self.gpu.stream();
32130 let mut b = __s_b.launch_builder(&f);
32131 b.arg(x).arg(dst).arg(&ni);
32132 unsafe {
32133 b.launch(cfg)?;
32134 }
32135 Ok(())
32136 }
32137
32138 pub fn gdn_chunk_k123_vl8(
32141 &self,
32142 seqs: &[GdnSeqVl],
32143 n_head: usize,
32144 hk: usize,
32145 wq: Option<&GdnWVl8>,
32146 ) -> Result<(), Box<dyn std::error::Error>> {
32147 let b = seqs.len();
32148 assert!((1..=8).contains(&b), "gdn_chunk_k123_vl8: 1..=8 sequences");
32149 let mut packed = [GdnSeqVl::default(); 8];
32150 packed[..b].copy_from_slice(seqs);
32151 let v = GdnVl8(packed);
32152 let (hi, ci) = (n_head as i32, 32i32);
32153 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
32154 {
32155 let f = self.func("gdn_chunk_cumgate_vl");
32156 let cfg = LaunchConfig {
32157 grid_dim: (max_nc, n_head as u32, b as u32),
32158 block_dim: (32, 1, 1),
32159 shared_mem_bytes: 0,
32160 };
32161 let __s_lb = self.gpu.stream();
32162 let mut lb = __s_lb.launch_builder(&f);
32163 lb.arg(&v).arg(&hi).arg(&ci);
32164 unsafe {
32165 lb.launch(cfg)?;
32166 }
32167 }
32168 let hki = hk as i32;
32169 if let Some(w) = wq {
32170 let f = self.func("gdn_k2_wgmma_vl");
32172 let cfg = LaunchConfig {
32173 grid_dim: (max_nc, n_head as u32, b as u32),
32174 block_dim: (128, 1, 1),
32175 shared_mem_bytes: 0,
32176 };
32177 let __s_lb = self.gpu.stream();
32178 let mut lb = __s_lb.launch_builder(&f);
32179 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
32180 unsafe {
32181 lb.launch(cfg)?;
32182 }
32183 } else {
32184 let f = self.func("gdn_chunk_attn_vl");
32185 f.set_attribute(
32186 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
32187 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
32188 )?;
32189 let cfg = LaunchConfig {
32190 grid_dim: (max_nc, n_head as u32, b as u32),
32191 block_dim: (256, 1, 1),
32192 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
32193 };
32194 let __s_lb = self.gpu.stream();
32195 let mut lb = __s_lb.launch_builder(&f);
32196 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
32197 unsafe {
32198 lb.launch(cfg)?;
32199 }
32200 }
32201 {
32202 let f = self.func("gdn_chunk_solve32_vl");
32203 let cfg = LaunchConfig {
32204 grid_dim: (max_nc, n_head as u32, b as u32),
32205 block_dim: (256, 1, 1),
32206 shared_mem_bytes: 0,
32207 };
32208 let __s_lb = self.gpu.stream();
32209 let mut lb = __s_lb.launch_builder(&f);
32210 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
32211 unsafe {
32212 lb.launch(cfg)?;
32213 }
32214 }
32215 Ok(())
32216 }
32217
32218 #[allow(clippy::too_many_arguments)]
32222 pub fn gdn_prep_vl8(
32223 &self,
32224 seqs: &[GdnPrepVl],
32225 conv_w: &CudaSlice<f32>,
32226 dt_bias: &CudaSlice<f32>,
32227 a: &CudaSlice<f32>,
32228 conv_dim: usize,
32229 d_conv: usize,
32230 d_state: usize,
32231 num_v: usize,
32232 num_k: usize,
32233 key_dim: usize,
32234 hk: usize,
32235 eps: f32,
32236 ) -> Result<(), Box<dyn std::error::Error>> {
32237 let b = seqs.len();
32238 assert!((1..=8).contains(&b));
32239 let mut packed = [GdnPrepVl::default(); 8];
32240 packed[..b].copy_from_slice(seqs);
32241 let v = GdnPrepVl8(packed);
32242 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
32243 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
32244 {
32245 let f = self.func("ssm_conv1d_gdn_state_vl");
32246 let cfg = LaunchConfig {
32247 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
32248 block_dim: (256, 1, 1),
32249 shared_mem_bytes: 0,
32250 };
32251 let (dsi, nvi, nki, kdi, hki) = (
32252 d_state as i32,
32253 num_v as i32,
32254 num_k as i32,
32255 key_dim as i32,
32256 hk as i32,
32257 );
32258 let __s_lb = self.gpu.stream();
32259 let mut lb = __s_lb.launch_builder(&f);
32260 lb.arg(&v)
32261 .arg(conv_w)
32262 .arg(&cdi)
32263 .arg(&dci)
32264 .arg(&dsi)
32265 .arg(&nvi)
32266 .arg(&nki)
32267 .arg(&kdi)
32268 .arg(&hki);
32269 unsafe {
32270 lb.launch(cfg)?;
32271 }
32272 }
32273 {
32274 let f = self.func("ssm_conv_ring_update_vl");
32275 let n = (conv_dim * (d_conv - 1)) as u32;
32276 let cfg = LaunchConfig {
32277 grid_dim: (n.div_ceil(256), 1, b as u32),
32278 block_dim: (256, 1, 1),
32279 shared_mem_bytes: 0,
32280 };
32281 let __s_lb = self.gpu.stream();
32282 let mut lb = __s_lb.launch_builder(&f);
32283 lb.arg(&v).arg(&cdi).arg(&dci);
32284 unsafe {
32285 lb.launch(cfg)?;
32286 }
32287 }
32288 if Self::l2_v2_on(d_state) {
32289 let f = self.func("gdn_l2_v2_vl");
32290 let cfg = LaunchConfig {
32291 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
32292 block_dim: (256, 1, 1),
32293 shared_mem_bytes: 0,
32294 };
32295 let (dsi, nvi) = (d_state as i32, hk as i32);
32296 let __s_lb = self.gpu.stream();
32297 let mut lb = __s_lb.launch_builder(&f);
32298 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
32299 unsafe {
32300 lb.launch(cfg)?;
32301 }
32302 } else {
32303 let f = self.func("gdn_l2_vl");
32304 let cfg = LaunchConfig {
32305 grid_dim: (max_t * hk as u32, 2, b as u32),
32306 block_dim: (256, 1, 1),
32307 shared_mem_bytes: 0,
32308 };
32309 let (dsi, nvi) = (d_state as i32, hk as i32);
32310 let __s_lb = self.gpu.stream();
32311 let mut lb = __s_lb.launch_builder(&f);
32312 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
32313 unsafe {
32314 lb.launch(cfg)?;
32315 }
32316 }
32317 {
32318 let f = self.func("gdn_gate_prep_vl");
32319 let n = max_t * num_v as u32;
32320 let cfg = LaunchConfig {
32321 grid_dim: (n.div_ceil(256), 1, b as u32),
32322 block_dim: (256, 1, 1),
32323 shared_mem_bytes: 0,
32324 };
32325 let nvi = num_v as i32;
32326 let __s_lb = self.gpu.stream();
32327 let mut lb = __s_lb.launch_builder(&f);
32328 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
32329 unsafe {
32330 lb.launch(cfg)?;
32331 }
32332 }
32333 Ok(())
32334 }
32335
32336 pub fn gdn_mirror_vl8(
32338 &self,
32339 seqs: &[GdnSeqVl],
32340 n_head: usize,
32341 which: i32,
32342 hk: usize,
32343 ) -> Result<(), Box<dyn std::error::Error>> {
32344 let b = seqs.len();
32345 assert!((1..=8).contains(&b));
32346 let mut packed = [GdnSeqVl::default(); 8];
32347 packed[..b].copy_from_slice(seqs);
32348 let v = GdnVl8(packed);
32349 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
32350 let max_n = seqs
32351 .iter()
32352 .map(|s| {
32353 if which == 0 {
32354 s.t as i64 * ept as i64
32355 } else {
32356 s.nc as i64 * ept as i64 * 32
32357 }
32358 })
32359 .max()
32360 .unwrap();
32361 let f = self.func("gdn_mirror_vl");
32362 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
32363 let cfg = LaunchConfig {
32364 grid_dim: (blocks, 1, b as u32),
32365 block_dim: (256, 1, 1),
32366 shared_mem_bytes: 0,
32367 };
32368 let __s_lb = self.gpu.stream();
32369 let mut lb = __s_lb.launch_builder(&f);
32370 lb.arg(&v).arg(&ept).arg(&which);
32371 unsafe {
32372 lb.launch(cfg)?;
32373 }
32374 Ok(())
32375 }
32376
32377 pub fn gdn_tail_vl8(
32379 &self,
32380 seqs: &[GdnPrepVl],
32381 norm_w: &CudaSlice<f32>,
32382 d_state: usize,
32383 num_v: usize,
32384 eps: f32,
32385 ) -> Result<(), Box<dyn std::error::Error>> {
32386 let b = seqs.len();
32387 assert!((1..=8).contains(&b));
32388 let mut packed = [GdnPrepVl::default(); 8];
32389 packed[..b].copy_from_slice(seqs);
32390 let v = GdnPrepVl8(packed);
32391 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
32392 let f = self.func("gated_rmsnorm_f16out_vl");
32393 let cfg = LaunchConfig {
32395 grid_dim: (max_t * num_v as u32, 1, b as u32),
32396 block_dim: (128, 1, 1),
32397 shared_mem_bytes: 0,
32398 };
32399 let (dsi, nvi) = (d_state as i32, num_v as i32);
32400 let __s_lb = self.gpu.stream();
32401 let mut lb = __s_lb.launch_builder(&f);
32402 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
32403 unsafe {
32404 lb.launch(cfg)?;
32405 }
32406 Ok(())
32407 }
32408
32409 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
32412 use cudarc::driver::DevicePtr;
32413 let s = self.gpu.stream();
32414 let (p, _g) = x.device_ptr(&s);
32415 p
32416 }
32417 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
32418 use cudarc::driver::DevicePtrMut;
32419 let s = self.gpu.stream();
32420 let (p, _g) = x.device_ptr_mut(&s);
32421 p
32422 }
32423 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
32424 use cudarc::driver::DevicePtr;
32425 let s = self.gpu.stream();
32426 let (p, _g) = x.device_ptr(&s);
32427 p
32428 }
32429 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
32430 use cudarc::driver::DevicePtr;
32431 let s = self.gpu.stream();
32432 let (p, _g) = x.device_ptr(&s);
32433 p
32434 }
32435
32436 pub fn gdn_chunk_vl8(
32440 &self,
32441 seqs: &[GdnSeqVl],
32442 n_head: usize,
32443 scale: f32,
32444 hk: usize,
32445 wq: Option<&GdnWVl8>,
32446 ) -> Result<(), Box<dyn std::error::Error>> {
32447 const NSPLIT: u32 = 4;
32448 let b = seqs.len();
32449 assert!((1..=8).contains(&b), "gdn_chunk_vl8: 1..=8 sequences");
32450 let mut packed = [GdnSeqVl::default(); 8];
32451 packed[..b].copy_from_slice(seqs);
32452 let v = GdnVl8(packed);
32453 let (hi, ci) = (n_head as i32, 32i32);
32454 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
32455 let hki = hk as i32;
32456 if let Some(w) = wq {
32457 let f = self.func("gdn_k45_wgmma_vl");
32459 let cfg = LaunchConfig {
32460 grid_dim: (n_head as u32, NSPLIT, b as u32),
32461 block_dim: (256, 1, 1),
32462 shared_mem_bytes: 0,
32463 };
32464 let __s_lb = self.gpu.stream();
32465 let mut lb = __s_lb.launch_builder(&f);
32466 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
32467 unsafe {
32468 lb.launch(cfg)?;
32469 }
32470 let _ = max_nc;
32471 return Ok(());
32472 }
32473 {
32474 let f = self.func("gdn_chunk_state_mma_vl");
32475 let cfg = LaunchConfig {
32476 grid_dim: (n_head as u32, NSPLIT, b as u32),
32477 block_dim: (256, 1, 1),
32478 shared_mem_bytes: 0,
32479 };
32480 let __s_lb = self.gpu.stream();
32481 let mut lb = __s_lb.launch_builder(&f);
32482 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
32483 unsafe {
32484 lb.launch(cfg)?;
32485 }
32486 }
32487 {
32488 let f = self.func("gdn_chunk_output_mma_vl");
32489 let cfg = LaunchConfig {
32490 grid_dim: (max_nc, n_head as u32, b as u32),
32491 block_dim: (256, 1, 1),
32492 shared_mem_bytes: 0,
32493 };
32494 let __s_lb = self.gpu.stream();
32495 let mut lb = __s_lb.launch_builder(&f);
32496 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
32497 unsafe {
32498 lb.launch(cfg)?;
32499 }
32500 }
32501 Ok(())
32502 }
32503 #[allow(clippy::too_many_arguments)] pub fn gdn_scan_chunked(
32505 &self,
32506 q: &CudaSlice<f32>,
32507 k: &CudaSlice<f32>,
32508 v: &CudaSlice<f32>,
32509 g: &CudaSlice<f32>,
32510 beta: &CudaSlice<f32>,
32511 kb16_pre: Option<&CudaSlice<u8>>,
32512 qb16_pre: Option<&CudaSlice<u8>>,
32513 state_in: &CudaSlice<f32>,
32514 state_out: &mut CudaSlice<f32>,
32515 o: &mut CudaSlice<f32>,
32516 n_head: usize,
32517 t: usize,
32518 scale: f32,
32519 c: usize,
32520 hk: usize,
32521 ) -> Result<(), Box<dyn std::error::Error>> {
32522 const D: usize = 128;
32523 const NSPLIT: u32 = 4;
32524 assert!(
32525 (1..=128).contains(&c),
32526 "gdn_scan_chunked: C must be in 1..=128"
32527 );
32528 let h = n_head;
32529 #[allow(clippy::manual_div_ceil)]
32530 let nc = (t + c - 1) / c;
32532 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
32533 let gdn_mma_pre = !portable_mma_gated()
32538 && c == 32
32539 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
32540 Ok("1") => true,
32541 Ok("0") => false,
32542 _ => gdn_mma_default_on(),
32543 };
32544 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
32545 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
32546 } else {
32547 None
32548 };
32549 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
32554 && gdn_mma_pre
32555 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
32556 let nk = t * hk * D;
32557 let mut kb16_local: Option<CudaSlice<u8>> = None;
32558 if gdn_mma_pre && kb16_pre.is_none() {
32559 let mut kb = self.alloc_u8_uninit(nk * 2)?;
32560 let f = self.func("f32_to_bf16_bulk");
32561 let n2 = nk as i64;
32562 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
32563 let __s_b = self.gpu.stream();
32564 let mut b = __s_b.launch_builder(&f);
32565 b.arg(k).arg(&mut kb).arg(&n2);
32566 unsafe {
32567 b.launch(cfg2)?;
32568 }
32569 kb16_local = Some(kb);
32570 }
32571 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
32572 if let Some(kb) = kb16_pre {
32573 assert!(kb.len() >= nk * 2, "kb16_pre too small");
32574 }
32575 let mut qb16: Option<CudaSlice<u8>> = None;
32576 let mut pb16: Option<CudaSlice<u8>> = None;
32577 if gdn_wgmma_pre {
32578 if qb16_pre.is_none() {
32581 let mut qb = self.alloc_u8_uninit(nk * 2)?;
32582 let f = self.func("f32_to_bf16_bulk");
32583 let n2 = nk as i64;
32584 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
32585 let __s_b = self.gpu.stream();
32586 let mut b = __s_b.launch_builder(&f);
32587 b.arg(q).arg(&mut qb).arg(&n2);
32588 unsafe {
32589 b.launch(cfg2)?;
32590 }
32591 qb16 = Some(qb);
32592 } else if let Some(qb) = qb16_pre {
32593 assert!(qb.len() >= nk * 2, "qb16_pre too small");
32594 }
32595 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
32596 }
32597 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
32598 let k2w = if gdn_wgmma_pre {
32599 Some((
32600 *qb16_ref0.as_ref().unwrap(),
32601 *kb16_ref0.as_ref().unwrap(),
32602 pb16.as_mut().unwrap(),
32603 ))
32604 } else {
32605 None
32606 };
32607 let (gcum, p, u, w) =
32608 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
32609 let _ = &w;
32610 let mut y = self.uninit(nc * h * c * D)?;
32611 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
32625 && c == 32
32626 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
32627 Ok("1") => true,
32628 Ok("0") => false,
32629 _ => gdn_mma_default_on(),
32630 };
32631 if gdn_mma {
32632 let wb16 = wb16_pre
32633 .take()
32634 .expect("mma path pre-allocates wb16 (K3 store fold)");
32635 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
32636 if gdn_wgmma_pre {
32648 let qb16 = qb16_ref0.unwrap();
32650 let pb16 = pb16.as_ref().unwrap();
32651 {
32652 let f = self.func("gdn_k45_wgmma");
32653 let cfg = LaunchConfig {
32654 grid_dim: (h as u32, 4, 1),
32655 block_dim: (256, 1, 1),
32656 shared_mem_bytes: 0,
32657 };
32658 let hki = hk as i32;
32659 let __s_b = self.gpu.stream();
32660 let mut b = __s_b.launch_builder(&f);
32661 b.arg(kb16_ref)
32662 .arg(&gcum)
32663 .arg(beta)
32664 .arg(&u)
32665 .arg(&wb16)
32666 .arg(qb16)
32667 .arg(pb16)
32668 .arg(o)
32669 .arg(&scale)
32670 .arg(state_in)
32671 .arg(&mut *state_out)
32672 .arg(&hi)
32673 .arg(&ti)
32674 .arg(&ci)
32675 .arg(&hki);
32676 unsafe {
32677 b.launch(cfg)?;
32678 }
32679 }
32680 return Ok(());
32681 }
32682 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
32686 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
32687 {
32688 let f = self.func("gdn_chunk_state_mma");
32689 let cfg = LaunchConfig {
32690 grid_dim: (h as u32, NSPLIT, 1),
32691 block_dim: (256, 1, 1),
32692 shared_mem_bytes: 0,
32693 };
32694 let hki = hk as i32;
32695 let __s_b = self.gpu.stream();
32696 let mut b = __s_b.launch_builder(&f);
32697 b.arg(kb16_ref)
32698 .arg(&gcum)
32699 .arg(beta)
32700 .arg(&u)
32701 .arg(&wb16)
32702 .arg(&mut y16)
32703 .arg(&mut ssnap16)
32704 .arg(state_in)
32705 .arg(&mut *state_out)
32706 .arg(&hi)
32707 .arg(&ti)
32708 .arg(&ci)
32709 .arg(&hki);
32710 unsafe {
32711 b.launch(cfg)?;
32712 }
32713 }
32714 {
32715 let f = self.func("gdn_chunk_output_mma");
32717 #[allow(clippy::manual_div_ceil)]
32718 let jt = ((c + 31) / 32) as u32;
32720 let cfg = LaunchConfig {
32721 grid_dim: (nc as u32, h as u32, jt),
32722 block_dim: (256, 1, 1),
32723 shared_mem_bytes: 0,
32724 };
32725 let hki = hk as i32;
32726 let __s_b = self.gpu.stream();
32727 let mut b = __s_b.launch_builder(&f);
32728 b.arg(q)
32729 .arg(&gcum)
32730 .arg(&p)
32731 .arg(&y16)
32732 .arg(&ssnap16)
32733 .arg(o)
32734 .arg(&hi)
32735 .arg(&ti)
32736 .arg(&ci)
32737 .arg(&scale)
32738 .arg(&hki);
32739 unsafe {
32740 b.launch(cfg)?;
32741 }
32742 }
32743 return Ok(());
32744 }
32745 {
32746 let f = self.func("gdn_chunk_state_f32");
32748 let cfg = LaunchConfig {
32749 grid_dim: (h as u32, NSPLIT, 1),
32750 block_dim: (256, 1, 1),
32751 shared_mem_bytes: 0,
32752 };
32753 let __s_b = self.gpu.stream();
32754 let mut b = __s_b.launch_builder(&f);
32755 b.arg(k)
32756 .arg(&gcum)
32757 .arg(beta)
32758 .arg(&u)
32759 .arg(&w)
32760 .arg(&mut y)
32761 .arg(&mut ssnap)
32762 .arg(state_in)
32763 .arg(&mut *state_out)
32764 .arg(&hi)
32765 .arg(&ti)
32766 .arg(&ci);
32767 unsafe {
32768 b.launch(cfg)?;
32769 }
32770 }
32771 {
32772 let f = self.func("gdn_chunk_output_f32");
32774 #[allow(clippy::manual_div_ceil)]
32775 let jt = ((c + 31) / 32) as u32;
32777 let cfg = LaunchConfig {
32778 grid_dim: (nc as u32, h as u32, jt),
32779 block_dim: (256, 1, 1),
32780 shared_mem_bytes: 0,
32781 };
32782 let __s_b = self.gpu.stream();
32783 let mut b = __s_b.launch_builder(&f);
32784 b.arg(q)
32785 .arg(&gcum)
32786 .arg(&p)
32787 .arg(&y)
32788 .arg(&ssnap)
32789 .arg(o)
32790 .arg(&hi)
32791 .arg(&ti)
32792 .arg(&ci)
32793 .arg(&scale);
32794 unsafe {
32795 b.launch(cfg)?;
32796 }
32797 }
32798 Ok(())
32799 }
32800
32801 #[allow(clippy::too_many_arguments)]
32810 #[allow(clippy::too_many_arguments)]
32811 pub fn gdn_scan_prefill(
32812 &self,
32813 q: &CudaSlice<f32>,
32814 k: &CudaSlice<f32>,
32815 v: &CudaSlice<f32>,
32816 g: &CudaSlice<f32>,
32817 beta: &CudaSlice<f32>,
32818 kb16_pre: Option<&CudaSlice<u8>>,
32819 qb16_pre: Option<&CudaSlice<u8>>,
32820 state_in: &CudaSlice<f32>,
32821 state_out: &mut CudaSlice<f32>,
32822 o: &mut CudaSlice<f32>,
32823 n_head: usize,
32824 t: usize,
32825 scale: f32,
32826 hk: usize,
32827 ) -> Result<(), Box<dyn std::error::Error>> {
32828 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
32829 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
32830 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
32831 }
32832 if Self::gdn_chunked_enabled() && t >= 16 {
32833 self.gdn_scan_chunked(
32834 q,
32835 k,
32836 v,
32837 g,
32838 beta,
32839 kb16_pre,
32840 qb16_pre,
32841 state_in,
32842 state_out,
32843 o,
32844 n_head,
32845 t,
32846 scale,
32847 Self::gdn_chunk_size(),
32848 hk,
32849 )
32850 } else {
32851 assert!(
32852 hk == n_head,
32853 "s128 scan is broadcast-only (prep guarantees by predicate)"
32854 );
32855 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
32856 }
32857 }
32858
32859 #[allow(clippy::too_many_arguments)]
32861 fn gdn_scan_diff(
32862 &self,
32863 q: &CudaSlice<f32>,
32864 k: &CudaSlice<f32>,
32865 v: &CudaSlice<f32>,
32866 g: &CudaSlice<f32>,
32867 beta: &CudaSlice<f32>,
32868 state_in: &CudaSlice<f32>,
32869 state_out: &mut CudaSlice<f32>,
32870 o: &mut CudaSlice<f32>,
32871 n_head: usize,
32872 t: usize,
32873 scale: f32,
32874 ) -> Result<(), Box<dyn std::error::Error>> {
32875 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
32876 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
32877 let mut o_c = self.uninit(o.len())?;
32878 let mut st_c = self.uninit(state_out.len())?;
32879 self.gdn_scan_chunked(
32880 q,
32881 k,
32882 v,
32883 g,
32884 beta,
32885 None,
32886 None,
32887 state_in,
32888 &mut st_c,
32889 &mut o_c,
32890 n_head,
32891 t,
32892 scale,
32893 Self::gdn_chunk_size(),
32894 n_head,
32895 )?;
32896 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
32897 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
32898 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
32899 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
32900 let mut max_abs = 0f32;
32901 let mut max_rel = 0f32;
32902 let mut sum_rel = 0f64;
32903 for (x, y) in a.iter().zip(b) {
32904 let ad = (x - y).abs();
32905 let rel = ad / x.abs().max(y.abs()).max(1e-3);
32906 if ad > max_abs {
32907 max_abs = ad;
32908 }
32909 if rel > max_rel {
32910 max_rel = rel;
32911 }
32912 sum_rel += rel as f64;
32913 }
32914 (max_abs, max_rel, sum_rel / a.len() as f64)
32915 };
32916 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
32917 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
32918 println!(
32919 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
32920 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
32921 Self::gdn_chunk_size()
32922 );
32923 Ok(())
32924 }
32925
32926 pub fn gdn_glog(
32928 &self,
32929 alpha: &CudaSlice<f32>,
32930 dt_bias: &CudaSlice<f32>,
32931 a: &CudaSlice<f32>,
32932 g_log: &mut CudaSlice<f32>,
32933 n_head: usize,
32934 t: usize,
32935 ) -> Result<(), Box<dyn std::error::Error>> {
32936 let f = self.func("gdn_glog_f32");
32937 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
32938 let (h, ti) = (n_head as i32, t as i32);
32939 let __s_b = self.gpu.stream();
32940 let mut b = __s_b.launch_builder(&f);
32941 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
32942 unsafe {
32943 b.launch(cfg)?;
32944 }
32945 Ok(())
32946 }
32947
32948 pub fn sigmoid_v(
32951 &self,
32952 x: &cudarc::driver::CudaView<f32>,
32953 y: &mut CudaSlice<f32>,
32954 n: usize,
32955 ) -> Result<(), Box<dyn std::error::Error>> {
32956 let f = self.func("sigmoid_f32");
32957 let cfg = LaunchConfig::for_num_elems(n as u32);
32958 let ni = n as i32;
32959 let __s_b = self.gpu.stream();
32960 let mut b = __s_b.launch_builder(&f);
32961 b.arg(x).arg(y).arg(&ni);
32962 unsafe {
32963 b.launch(cfg)?;
32964 }
32965 Ok(())
32966 }
32967
32968 pub fn gdn_glog_v(
32969 &self,
32970 alpha: &cudarc::driver::CudaView<f32>,
32971 dt_bias: &CudaSlice<f32>,
32972 a: &CudaSlice<f32>,
32973 g_log: &mut CudaSlice<f32>,
32974 n_head: usize,
32975 t: usize,
32976 ) -> Result<(), Box<dyn std::error::Error>> {
32977 let f = self.func("gdn_glog_f32");
32978 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
32979 let (h, ti) = (n_head as i32, t as i32);
32980 let __s_b = self.gpu.stream();
32981 let mut b = __s_b.launch_builder(&f);
32982 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
32983 unsafe {
32984 b.launch(cfg)?;
32985 }
32986 Ok(())
32987 }
32988
32989 pub fn sigmoid(
32990 &self,
32991 x: &CudaSlice<f32>,
32992 y: &mut CudaSlice<f32>,
32993 n: usize,
32994 ) -> Result<(), Box<dyn std::error::Error>> {
32995 let f = self.func("sigmoid_f32");
32996 let cfg = LaunchConfig::for_num_elems(n as u32);
32997 let ni = n as i32;
32998 let __s_b = self.gpu.stream();
32999 let mut b = __s_b.launch_builder(&f);
33000 b.arg(x).arg(y).arg(&ni);
33001 unsafe {
33002 b.launch(cfg)?;
33003 }
33004 Ok(())
33005 }
33006
33007 pub fn sig_mul_f16out(
33010 &self,
33011 a: &CudaSlice<f32>,
33012 g: &CudaSlice<f32>,
33013 dst: &mut CudaSlice<f32>,
33014 dst16: &mut CudaSlice<u8>,
33015 n: usize,
33016 ) -> Result<(), Box<dyn std::error::Error>> {
33017 let f = self.func("sig_mul_f16out_f32");
33018 let cfg = LaunchConfig::for_num_elems(n as u32);
33019 let ni = n as i32;
33020 let __s_b = self.gpu.stream();
33021 let mut b = __s_b.launch_builder(&f);
33022 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
33023 unsafe {
33024 b.launch(cfg)?;
33025 }
33026 Ok(())
33027 }
33028
33029 #[allow(clippy::too_many_arguments)]
33038 pub fn attn_head_gate(
33039 &self,
33040 a: &CudaSlice<f32>,
33041 g: &CudaSlice<f32>,
33042 dst: &mut CudaSlice<f32>,
33043 dst16: Option<&mut CudaSlice<u8>>,
33044 head_dim: usize,
33045 n_head: usize,
33046 t: usize,
33047 ) -> Result<(), Box<dyn std::error::Error>> {
33048 let f = self.func("attn_head_gate_f32");
33049 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
33050 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
33051 let d16: u64 = match dst16 {
33053 Some(d) => self.addr_u8(d),
33054 None => 0,
33055 };
33056 let __s_b = self.gpu.stream();
33057 let mut b = __s_b.launch_builder(&f);
33058 b.arg(a)
33059 .arg(g)
33060 .arg(dst)
33061 .arg(&d16)
33062 .arg(&hd)
33063 .arg(&nh)
33064 .arg(&ti);
33065 unsafe {
33066 b.launch(cfg)?;
33067 }
33068 Ok(())
33069 }
33070
33071 #[allow(clippy::too_many_arguments)]
33080 pub fn swiglu_clamped_mul_scaled(
33081 &self,
33082 gate: &CudaSlice<f32>,
33083 up: &CudaSlice<f32>,
33084 gs: f32,
33085 us: f32,
33086 limit: f32,
33087 dst: &mut CudaSlice<f32>,
33088 n: usize,
33089 ) -> Result<(), Box<dyn std::error::Error>> {
33090 debug_assert!(
33091 limit > 1e-6,
33092 "swiglu_clamped needs a live limit; use silu_mul_scaled"
33093 );
33094 let f = self.func("swiglu_clamped_mul_scaled_f32");
33095 let cfg = LaunchConfig::for_num_elems(n as u32);
33096 let ni = n as i32;
33097 let __s_b = self.gpu.stream();
33098 let mut b = __s_b.launch_builder(&f);
33099 b.arg(gate)
33100 .arg(up)
33101 .arg(&gs)
33102 .arg(&us)
33103 .arg(&limit)
33104 .arg(dst)
33105 .arg(&ni);
33106 unsafe {
33107 b.launch(cfg)?;
33108 }
33109 Ok(())
33110 }
33111
33112 #[allow(clippy::too_many_arguments)]
33121 pub fn swiglu_preclamped_mul_scaled(
33122 &self,
33123 gate: &CudaSlice<f32>,
33124 up: &CudaSlice<f32>,
33125 gs: f32,
33126 us: f32,
33127 limit: f32,
33128 dst: &mut CudaSlice<f32>,
33129 n: usize,
33130 ) -> Result<(), Box<dyn std::error::Error>> {
33131 debug_assert!(
33132 limit > 1e-6,
33133 "swiglu_preclamped needs a live limit; use silu_mul_scaled"
33134 );
33135 let f = self.func("swiglu_preclamped_mul_scaled_f32");
33136 let cfg = LaunchConfig::for_num_elems(n as u32);
33137 let ni = n as i32;
33138 let __s_b = self.gpu.stream();
33139 let mut b = __s_b.launch_builder(&f);
33140 b.arg(gate)
33141 .arg(up)
33142 .arg(&gs)
33143 .arg(&us)
33144 .arg(&limit)
33145 .arg(dst)
33146 .arg(&ni);
33147 unsafe {
33148 b.launch(cfg)?;
33149 }
33150 Ok(())
33151 }
33152
33153 #[allow(clippy::too_many_arguments)] pub fn gated_rmsnorm(
33156 &self,
33157 o: &CudaSlice<f32>,
33158 w: &CudaSlice<f32>,
33159 z: &CudaSlice<f32>,
33160 dst: &mut CudaSlice<f32>,
33161 ncols: usize,
33162 nrows: usize,
33163 eps: f32,
33164 ) -> Result<(), Box<dyn std::error::Error>> {
33165 let f = self.func("gated_rmsnorm_f32");
33166 let cfg = LaunchConfig {
33167 grid_dim: (nrows as u32, 1, 1),
33168 block_dim: (128, 1, 1),
33169 shared_mem_bytes: 0,
33170 };
33171 let (nc, e) = (ncols as i32, eps);
33172 let __s_b = self.gpu.stream();
33173 let mut b = __s_b.launch_builder(&f);
33174 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
33175 unsafe {
33176 b.launch(cfg)?;
33177 }
33178 Ok(())
33179 }
33180
33181 #[allow(clippy::too_many_arguments)] pub fn gated_rmsnorm_f16out(
33185 &self,
33186 o: &CudaSlice<f32>,
33187 w: &CudaSlice<f32>,
33188 z: &CudaSlice<f32>,
33189 dst: &mut CudaSlice<f32>,
33190 dst16: &mut CudaSlice<u8>,
33191 ncols: usize,
33192 nrows: usize,
33193 eps: f32,
33194 ) -> Result<(), Box<dyn std::error::Error>> {
33195 let f = self.func("gated_rmsnorm_f16out_f32");
33196 let cfg = LaunchConfig {
33198 grid_dim: (nrows as u32, 1, 1),
33199 block_dim: (128, 1, 1),
33200 shared_mem_bytes: 0,
33201 };
33202 let (nc, e) = (ncols as i32, eps);
33203 let __s_b = self.gpu.stream();
33204 let mut b = __s_b.launch_builder(&f);
33205 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
33206 unsafe {
33207 b.launch(cfg)?;
33208 }
33209 Ok(())
33210 }
33211
33212 #[allow(clippy::too_many_arguments)]
33216 pub fn add_rms_norm_zq8(
33217 &self,
33218 a: &CudaSlice<f32>,
33219 b_in: &CudaSlice<f32>,
33220 w: &CudaSlice<f32>,
33221 res: &mut CudaSlice<f32>,
33222 z: &mut CudaSlice<f32>,
33223 ncols: usize,
33224 nrows: usize,
33225 eps: f32,
33226 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
33227 assert!(ncols.is_multiple_of(32));
33228 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
33229 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
33230 let f = self.func("add_rms_norm_zq8");
33231 let cfg = LaunchConfig {
33232 grid_dim: (nrows as u32, 1, 1),
33233 block_dim: (1024, 1, 1),
33234 shared_mem_bytes: 0,
33235 };
33236 let (nc, ep) = (ncols as i32, eps);
33237 let __s_b = self.gpu.stream();
33238 let mut b = __s_b.launch_builder(&f);
33239 b.arg(a)
33240 .arg(b_in)
33241 .arg(w)
33242 .arg(res)
33243 .arg(z)
33244 .arg(&mut q)
33245 .arg(&mut d)
33246 .arg(&nc)
33247 .arg(&ep);
33248 unsafe {
33249 b.launch(cfg)?;
33250 }
33251 Ok((q, d))
33252 }
33253
33254 #[allow(clippy::too_many_arguments)] pub fn gated_rmsnorm_zv(
33260 &self,
33261 o: &CudaSlice<f32>,
33262 w: &CudaSlice<f32>,
33263 z: &cudarc::driver::CudaView<f32>,
33264 dst: &mut CudaSlice<f32>,
33265 ncols: usize,
33266 nrows: usize,
33267 eps: f32,
33268 ) -> Result<(), Box<dyn std::error::Error>> {
33269 let f = self.func("gated_rmsnorm_f32");
33270 let cfg = LaunchConfig {
33271 grid_dim: (nrows as u32, 1, 1),
33272 block_dim: (128, 1, 1),
33273 shared_mem_bytes: 0,
33274 };
33275 let (nc, e) = (ncols as i32, eps);
33276 let __s_b = self.gpu.stream();
33277 let mut b = __s_b.launch_builder(&f);
33278 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
33279 unsafe {
33280 b.launch(cfg)?;
33281 }
33282 Ok(())
33283 }
33284
33285 #[allow(clippy::too_many_arguments)] pub fn gated_rmsnorm_f16out_zv(
33287 &self,
33288 o: &CudaSlice<f32>,
33289 w: &CudaSlice<f32>,
33290 z: &cudarc::driver::CudaView<f32>,
33291 dst: &mut CudaSlice<f32>,
33292 dst16: &mut CudaSlice<u8>,
33293 ncols: usize,
33294 nrows: usize,
33295 eps: f32,
33296 ) -> Result<(), Box<dyn std::error::Error>> {
33297 let f = self.func("gated_rmsnorm_f16out_f32");
33298 let cfg = LaunchConfig {
33300 grid_dim: (nrows as u32, 1, 1),
33301 block_dim: (128, 1, 1),
33302 shared_mem_bytes: 0,
33303 };
33304 let (nc, e) = (ncols as i32, eps);
33305 let __s_b = self.gpu.stream();
33306 let mut b = __s_b.launch_builder(&f);
33307 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
33308 unsafe {
33309 b.launch(cfg)?;
33310 }
33311 Ok(())
33312 }
33313
33314 pub fn gated_rmsnorm_q8_1(
33315 &self,
33316 o: &CudaSlice<f32>,
33317 w: &CudaSlice<f32>,
33318 z: &CudaSlice<f32>,
33319 ncols: usize,
33320 nrows: usize,
33321 eps: f32,
33322 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
33323 assert!(ncols.is_multiple_of(32));
33324 let f = self.func("gated_rmsnorm_q8_1");
33325 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
33326 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
33327 let cfg = LaunchConfig {
33328 grid_dim: (nrows as u32, 1, 1),
33329 block_dim: (128, 1, 1),
33330 shared_mem_bytes: 0,
33331 };
33332 let (nc, ep) = (ncols as i32, eps);
33333 let __s_b = self.gpu.stream();
33334 let mut b = __s_b.launch_builder(&f);
33335 b.arg(o)
33336 .arg(w)
33337 .arg(z)
33338 .arg(&mut out_q)
33339 .arg(&mut out_d)
33340 .arg(&nc)
33341 .arg(&ep);
33342 unsafe {
33343 b.launch(cfg)?;
33344 }
33345 Ok((out_q, out_d))
33346 }
33347
33348 pub fn transpose(
33350 &self,
33351 inp: &CudaSlice<f32>,
33352 rows: usize,
33353 cols: usize,
33354 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
33355 let f = self.func("transpose_f32");
33356 let mut out = self.zeros(rows * cols)?;
33357 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
33358 let (r, c) = (rows as i32, cols as i32);
33359 let __s_b = self.gpu.stream();
33360 let mut b = __s_b.launch_builder(&f);
33361 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
33362 unsafe {
33363 b.launch(cfg)?;
33364 }
33365 Ok(out)
33366 }
33367
33368 pub fn repeat_heads(
33370 &self,
33371 inp: &CudaSlice<f32>,
33372 out: &mut CudaSlice<f32>,
33373 head_dim: usize,
33374 n_in: usize,
33375 n_out: usize,
33376 t: usize,
33377 ) -> Result<(), Box<dyn std::error::Error>> {
33378 let f = self.func("repeat_heads_f32");
33379 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
33380 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
33381 let __s_b = self.gpu.stream();
33382 let mut b = __s_b.launch_builder(&f);
33383 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
33384 unsafe {
33385 b.launch(cfg)?;
33386 }
33387 Ok(())
33388 }
33389
33390 pub fn q_gate_split(
33397 &self,
33398 qf: &CudaSlice<f32>,
33399 q_out: &mut CudaSlice<f32>,
33400 gate_out: &mut CudaSlice<f32>,
33401 head_dim: usize,
33402 n_head: usize,
33403 t: usize,
33404 ) -> Result<(), Box<dyn std::error::Error>> {
33405 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
33406 let out_need = head_dim * n_head * t;
33407 if q_out.len() < out_need || gate_out.len() < out_need {
33408 return Err(format!(
33409 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
33410 q_out.len(),
33411 gate_out.len()
33412 )
33413 .into());
33414 }
33415 let f = self.func("q_gate_split_f32");
33416 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
33417 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
33418 let __s_b = self.gpu.stream();
33419 let mut b = __s_b.launch_builder(&f);
33420 b.arg(qf)
33421 .arg(q_out)
33422 .arg(gate_out)
33423 .arg(&hd)
33424 .arg(&nh)
33425 .arg(&ti);
33426 unsafe {
33427 b.launch(cfg)?;
33428 }
33429 Ok(())
33430 }
33431
33432 #[allow(clippy::too_many_arguments)] pub fn qkv_to_gdn_repack(
33437 &self,
33438 conv_out: &CudaSlice<f32>,
33439 q_g: &mut CudaSlice<f32>,
33440 k_g: &mut CudaSlice<f32>,
33441 v_g: &mut CudaSlice<f32>,
33442 d_state: usize,
33443 num_v: usize,
33444 num_k: usize,
33445 key_dim: usize,
33446 t: usize,
33447 ) -> Result<(), Box<dyn std::error::Error>> {
33448 let f = self.func("qkv_to_gdn_repack_f32");
33449 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
33450 let (ds, nv, nk, kd, ti) = (
33451 d_state as i32,
33452 num_v as i32,
33453 num_k as i32,
33454 key_dim as i32,
33455 t as i32,
33456 );
33457 let __s_b = self.gpu.stream();
33458 let mut b = __s_b.launch_builder(&f);
33459 b.arg(conv_out)
33460 .arg(q_g)
33461 .arg(k_g)
33462 .arg(v_g)
33463 .arg(&ds)
33464 .arg(&nv)
33465 .arg(&nk)
33466 .arg(&kd)
33467 .arg(&ti);
33468 unsafe {
33469 b.launch(cfg)?;
33470 }
33471 Ok(())
33472 }
33473
33474 pub fn conv_left_pad(
33477 &self,
33478 src: &CudaSlice<f32>,
33479 dst: &mut CudaSlice<f32>,
33480 conv_dim: usize,
33481 t: usize,
33482 pad: usize,
33483 ) -> Result<(), Box<dyn std::error::Error>> {
33484 let f = self.func("conv_left_pad_f32");
33485 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
33486 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
33487 let __s_b = self.gpu.stream();
33488 let mut b = __s_b.launch_builder(&f);
33489 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
33490 unsafe {
33491 b.launch(cfg)?;
33492 }
33493 Ok(())
33494 }
33495
33496 pub fn conv_assemble_and_roll(
33500 &self,
33501 qkv_col: &CudaSlice<f32>,
33502 conv_state: &mut CudaSlice<f32>,
33503 conv_in: &mut CudaSlice<f32>,
33504 conv_dim: usize,
33505 pad: usize,
33506 ) -> Result<(), Box<dyn std::error::Error>> {
33507 let f = self.func("conv_assemble_and_roll_f32");
33508 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
33509 let (cd, p) = (conv_dim as i32, pad as i32);
33510 let __s_b = self.gpu.stream();
33511 let mut b = __s_b.launch_builder(&f);
33512 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
33513 unsafe {
33514 b.launch(cfg)?;
33515 }
33516 Ok(())
33517 }
33518
33519 pub fn ssm_conv1d_fused_decode(
33525 &self,
33526 qkv_col: &CudaSlice<f32>,
33527 conv_state: &mut CudaSlice<f32>,
33528 w: &CudaSlice<f32>,
33529 conv_out: &mut CudaSlice<f32>,
33530 conv_dim: usize,
33531 d_conv: usize,
33532 ) -> Result<(), Box<dyn std::error::Error>> {
33533 let f = self.func("ssm_conv1d_fused_decode_f32");
33534 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
33535 let (cd, dc) = (conv_dim as i32, d_conv as i32);
33536 let __s_b = self.gpu.stream();
33537 let mut b = __s_b.launch_builder(&f);
33538 b.arg(qkv_col)
33539 .arg(conv_state)
33540 .arg(w)
33541 .arg(conv_out)
33542 .arg(&cd)
33543 .arg(&dc);
33544 unsafe {
33545 b.launch(cfg)?;
33546 }
33547 Ok(())
33548 }
33549
33550 pub fn slice_range(
33553 &self,
33554 src: &CudaSlice<f32>,
33555 start: usize,
33556 len: usize,
33557 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
33558 let host = self.gpu.stream().clone_dtoh(src)?;
33559 self.gpu.stream().synchronize()?;
33560 self.htod(&host[start..start + len])
33561 }
33562}
33563
33564pub fn f32_gemv_kernel_on() -> bool {
33573 std::env::var("MEMRA_F32_GEMV_KERNEL").as_deref() == Ok("1")
33574}
33575
33576pub fn moe_down_ilp2_on() -> bool {
33585 std::env::var("MEMRA_MOE_DOWN_ILP2").as_deref() == Ok("1")
33586}
33587
33588pub fn mla_wo_zq8_on() -> bool {
33597 std::env::var("MEMRA_MLA_WO_ZQ8").as_deref() == Ok("1")
33598}
33599
33600pub static MLA_WO_ZQ8_DISPATCHES: std::sync::atomic::AtomicU64 =
33602 std::sync::atomic::AtomicU64::new(0);
33603
33604pub fn moe_gateup_ilp2_on() -> bool {
33613 std::env::var("MEMRA_MOE_GATEUP_ILP2").as_deref() == Ok("1")
33614}
33615
33616pub static MOE_GATEUP_ILP2_DISPATCHES: std::sync::atomic::AtomicU64 =
33618 std::sync::atomic::AtomicU64::new(0);
33619
33620pub static MOE_DOWN_ILP2_DISPATCHES: std::sync::atomic::AtomicU64 =
33622 std::sync::atomic::AtomicU64::new(0);
33623
33624pub static F32_GEMV_KERNEL_DISPATCHES: std::sync::atomic::AtomicU64 =
33626 std::sync::atomic::AtomicU64::new(0);
33627
33628#[cfg(test)]
33629mod target_dispatch_tests {
33630 use super::legacy_quant_gemm_allowed;
33631
33632 #[test]
33633 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
33634 assert!(legacy_quant_gemm_allowed(false, false, false));
33636 assert!(!legacy_quant_gemm_allowed(false, false, true));
33637 assert!(!legacy_quant_gemm_allowed(true, false, false));
33639 assert!(!legacy_quant_gemm_allowed(true, false, true));
33640 assert!(legacy_quant_gemm_allowed(true, true, false));
33642 assert!(!legacy_quant_gemm_allowed(true, true, true));
33643 }
33644
33645 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
33646 #[test]
33647 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
33648 assert!(!legacy_quant_gemm_allowed(
33649 cfg!(memra_portable_cuda),
33650 cfg!(memra_hopper_mma),
33651 false
33652 ));
33653 }
33654
33655 #[cfg(memra_hopper_mma)]
33656 #[test]
33657 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
33658 assert!(legacy_quant_gemm_allowed(
33659 cfg!(memra_portable_cuda),
33660 cfg!(memra_hopper_mma),
33661 false
33662 ));
33663 assert!(super::portable_mma_gated() == false);
33664 }
33665}
33666
33667impl memra_kv::KvDev for Engine {
33670 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
33671 Engine::zeros(self, n)
33672 }
33673 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
33674 Engine::uninit(self, n)
33675 }
33676 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
33677 Engine::alloc_u8(self, n)
33678 }
33679 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
33680 Engine::htod_i32(self, v)
33681 }
33682 fn clone_dtod(
33683 &self,
33684 src: &CudaSlice<f32>,
33685 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
33686 Engine::clone_dtod(self, src)
33687 }
33688 fn copy_into(
33689 &self,
33690 dst: &mut CudaSlice<f32>,
33691 off: usize,
33692 src: &CudaSlice<f32>,
33693 len: usize,
33694 ) -> Result<(), Box<dyn std::error::Error>> {
33695 Engine::copy_into(self, dst, off, src, len)
33696 }
33697 fn copy_range_into(
33698 &self,
33699 dst: &mut CudaSlice<f32>,
33700 dst_off: usize,
33701 src: &CudaSlice<f32>,
33702 src_off: usize,
33703 len: usize,
33704 ) -> Result<(), Box<dyn std::error::Error>> {
33705 Engine::copy_range_into(self, dst, dst_off, src, src_off, len)
33706 }
33707 fn set_i32_one(
33708 &self,
33709 d: &mut CudaSlice<i32>,
33710 v: i32,
33711 ) -> Result<(), Box<dyn std::error::Error>> {
33712 Engine::set_i32_one(self, d, v)
33713 }
33714}
33715
33716#[cfg(test)]
33717mod fused_gate_bounds_tests {
33718 use super::*;
33719
33720 #[test]
33733 #[ignore = "requires a CUDA GPU"]
33734 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
33735 let e = Engine::new(0).unwrap();
33736 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
33737 let fused = 2 * head_dim * n_head * t;
33738 let out_n = head_dim * n_head * t;
33739
33740 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
33742 let mut q = e.uninit(out_n).unwrap();
33743 let mut gate = e.uninit(out_n).unwrap();
33744 let err = e
33745 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
33746 .expect_err("half-width wq must be refused, not read past")
33747 .to_string();
33748 assert!(err.contains("NO fused gate"), "{err}");
33749 assert!(err.contains(&format!("{fused}")), "{err}");
33750
33751 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
33754 let wide = e.htod(&host).unwrap();
33755 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
33756 .expect("full-width wq splits");
33757 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
33758 for tok in 0..t {
33759 for hh in 0..n_head {
33760 for d in 0..head_dim {
33761 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
33762 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
33763 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
33764 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
33765 }
33766 }
33767 }
33768
33769 let mut small = e.uninit(out_n - 1).unwrap();
33771 assert!(
33772 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
33773 .is_err()
33774 );
33775 }
33776}
33777
33778#[cfg(test)]
33782mod fused_rope_width_tests {
33783 use super::Engine;
33784
33785 #[test]
33788 fn full_width_is_accepted() {
33789 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
33790 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
33791 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
33792 }
33793
33794 #[test]
33807 fn gemma4_official_artifact_widths_pass() {
33808 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
33809 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
33810 }
33811
33812 #[test]
33815 fn partial_rotary_is_refused_with_the_geometry_named() {
33816 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
33818 .expect_err("partial rotary must refuse");
33819 let msg = err.to_string();
33820 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
33821 assert!(msg.contains("n_rot 64"), "{msg}");
33822 assert!(msg.contains("head_dim 256"), "{msg}");
33823 assert!(
33824 msg.contains("64..256"),
33825 "names the band it would corrupt: {msg}"
33826 );
33827 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
33829 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
33831 }
33832}