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_vrows_t1_dev_forced() -> bool {
375 glm5_vrows_t1_dev_forced_from(
376 std::env::var("MEMRA_GLM5_VROWS_T1_DEV").ok().as_deref(),
377 env!("MEMRA_BUILT_CUDA_ARCH"),
378 glm5_graph_no_capture() && glm5_decode_graph_on(),
379 )
380}
381
382pub fn glm5_vrows_t1_dev_forced_from(
392 v: Option<&str>,
393 built_arch: &str,
394 no_capture_bisect: bool,
395) -> bool {
396 match v.map(str::trim) {
397 Some("1") => true,
398 Some("0") => no_capture_bisect,
399 _ => built_arch == "100a" || no_capture_bisect,
400 }
401}
402
403#[cfg(test)]
404mod glm5_vrows_t1_dev_default_tests {
405 use super::glm5_vrows_t1_dev_forced_from;
406
407 #[test]
408 fn arch_keyed_default_with_override_and_bisect_knob() {
409 assert!(glm5_vrows_t1_dev_forced_from(None, "100a", false));
410 assert!(!glm5_vrows_t1_dev_forced_from(None, "120a", false));
411 assert!(glm5_vrows_t1_dev_forced_from(None, "120a", true));
412 assert!(glm5_vrows_t1_dev_forced_from(Some("1"), "120a", false));
413 assert!(!glm5_vrows_t1_dev_forced_from(Some("0"), "100a", false));
414 assert!(glm5_vrows_t1_dev_forced_from(Some("0"), "100a", true));
415 }
416}
417
418pub fn glm5_graph_trace_on() -> bool {
426 std::env::var("MEMRA_GLM5_GRAPH_TRACE").as_deref() == Ok("1")
427}
428
429type Glm5TraceKey = (&'static str, String, u16);
445
446fn glm5_trace_slots() -> &'static Mutex<std::collections::BTreeSet<Glm5TraceKey>> {
447 static S: std::sync::OnceLock<Mutex<std::collections::BTreeSet<Glm5TraceKey>>> =
448 std::sync::OnceLock::new();
449 S.get_or_init(Default::default)
450}
451
452pub(crate) fn glm5_trace_take_slot(kind: &'static str, arm: &str, il: u16) -> bool {
456 const GLM5_TRACE_MAX_LAYERS: usize = 8;
457 let mut s = glm5_trace_slots().lock().unwrap();
458 if s.iter().filter(|(k, a, _)| *k == kind && a == arm).count() >= GLM5_TRACE_MAX_LAYERS {
459 return false;
460 }
461 s.insert((kind, arm.to_string(), il))
462}
463
464pub fn glm5_trace_reset() {
467 glm5_trace_slots().lock().unwrap().clear();
468}
469
470pub static GLM5_DECODE_GRAPH_REPLAYS: std::sync::atomic::AtomicU64 =
473 std::sync::atomic::AtomicU64::new(0);
474pub static GLM5_DECODE_GRAPH_CAPTURES: std::sync::atomic::AtomicU64 =
475 std::sync::atomic::AtomicU64::new(0);
476pub static GLM5_DECODE_GRAPH_LAYERS: std::sync::atomic::AtomicU64 =
477 std::sync::atomic::AtomicU64::new(0);
478pub static GLM5_DECODE_GRAPH_RECAPTURES: std::sync::atomic::AtomicU64 =
480 std::sync::atomic::AtomicU64::new(0);
481
482pub(crate) static GLM5_GRAPH_CAPTURE_OPEN: std::sync::atomic::AtomicBool =
488 std::sync::atomic::AtomicBool::new(false);
489
490pub(crate) fn glm5_graph_capture_open() -> bool {
491 GLM5_GRAPH_CAPTURE_OPEN.load(std::sync::atomic::Ordering::Relaxed)
492}
493
494pub fn router_prefill_exact_on() -> bool {
504 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
505 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
506}
507
508pub fn router_kernel_on() -> bool {
509 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
510 *ON.get_or_init(|| {
511 let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
512 if !on {
513 eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
514 }
515 on
516 })
517}
518
519pub const ROUTER_BATCH_MIN_T: usize = 8;
534pub fn router_batch_on() -> bool {
535 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
536 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
537}
538mod cpu_experts;
539#[cfg(memra_cutlass)]
540pub mod cutlass_ffi;
541pub mod dsv4_ffi;
542pub mod dsv4_gpu;
543pub mod f16_ffi;
544pub mod fp8_ffi;
545pub mod mmq_ffi;
546pub mod moe_cache;
547pub mod prime_graph;
548pub mod spill;
549mod spill_pread;
550
551const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
558const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
559const KDA_FATBIN: &[u8] = include_bytes!(env!("MEMRA_KDA_FATBIN"));
561const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
562const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
563const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
564const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
565const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
567
568fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
574 assert!(
575 !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
576 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
577 );
578 match std::env::var("MEMRA_GEMM_FATBIN") {
579 Ok(path) => std::borrow::Cow::Owned(
580 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
581 ),
582 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
583 }
584}
585
586pub(crate) const fn portable_mma_gated() -> bool {
593 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
594}
595
596#[track_caller]
613pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
614 assert!(
615 !portable_mma_gated(),
616 "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
617 portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
618 );
619}
620
621pub(crate) const fn gdn_mma_default_on() -> bool {
630 cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
631}
632
633const fn konst_eq(a: &str, b: &str) -> bool {
635 let (a, b) = (a.as_bytes(), b.as_bytes());
636 if a.len() != b.len() {
637 return false;
638 }
639 let mut i = 0;
640 while i < a.len() {
641 if a[i] != b[i] {
642 return false;
643 }
644 i += 1;
645 }
646 true
647}
648
649const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
654 (!portable_cuda || hopper_mma) && !no_gemm
655}
656
657const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
665const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
666const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
667const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
668const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
669
670pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
673
674fn flash_fatbin_bytes() -> &'static [u8] {
676 match kv_cache_formats() {
677 ("q8_0", "q5_1") => FLASH_FATBIN,
678 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
679 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
680 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
681 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
682 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
683 other => unreachable!("kv_cache_formats returned {other:?}"),
684 }
685}
686
687fn k1_launch_override() -> Option<(u32, u32, u32)> {
694 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
695 *K1.get_or_init(|| {
696 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
697 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
698 match p.as_slice() {
699 [bm, bn, w] => Some((*bm, *bn, *w)),
700 _ => None,
701 }
702 })
703}
704
705pub(crate) fn wgmma_gemm_enabled() -> bool {
712 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
713 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
714}
715
716pub const FA_VEC_MIN_TKV: usize = 96;
731pub fn fa_vec_min_tkv() -> usize {
735 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
736 *V.get_or_init(|| {
737 std::env::var("MEMRA_FA_VEC_MIN")
738 .ok()
739 .and_then(|v| v.parse().ok())
740 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
741 })
742}
743
744pub fn fa_f16pv_on() -> bool {
755 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
756 *ON.get_or_init(|| {
757 std::env::var("MEMRA_FA_F16PV")
758 .map(|v| v != "0")
759 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
760 })
761}
762
763pub fn fa512_hp_on() -> bool {
767 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
768 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
769}
770
771pub fn faw_hp_on() -> bool {
775 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
776 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
777}
778
779pub fn fa512_wide_warps() -> usize {
783 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
784 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
785 Ok("1") => 4,
786 _ => 2,
787 })
788}
789
790pub fn fa512_min_tkv() -> usize {
793 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
794 *FA512_MIN.get_or_init(|| {
795 std::env::var("MEMRA_FA512_MIN")
796 .ok()
797 .and_then(|v| v.parse().ok())
798 .unwrap_or(512)
799 })
800}
801pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
805 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
806pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
810pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
816 std::sync::atomic::AtomicBool::new(false);
817pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
824 std::sync::atomic::AtomicBool::new(true);
825pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
826 std::sync::atomic::AtomicUsize::new(16);
827pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
832pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
834pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
839pub use memra_kv::KV_FP8_FORCE;
842pub(crate) fn mmv_block() -> u32 {
847 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
848 *V.get_or_init(|| {
849 std::env::var("MEMRA_MMV_BLOCK")
850 .ok()
851 .and_then(|v| v.parse().ok())
852 .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
853 .unwrap_or(128)
854 })
855}
856
857pub(crate) fn b200_matvec_arm_on() -> bool {
868 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
869 *V.get_or_init(|| {
870 env!("MEMRA_BUILT_CUDA_ARCH") == "100a"
871 && std::env::var("MEMRA_B200_MATVEC_ARM").as_deref() == Ok("1")
872 })
873}
874
875pub(crate) fn b200_gemv_v2_level() -> u8 {
907 static V: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
908 *V.get_or_init(|| {
909 if env!("MEMRA_BUILT_CUDA_ARCH") != "100a" {
910 return 0;
911 }
912 match std::env::var("MEMRA_B200_GEMV_V2").as_deref() {
913 Ok("1") => 1,
914 Ok("2") => 2,
915 _ => 0,
916 }
917 })
918}
919
920pub(crate) fn b200_gemv_v2_on() -> bool {
921 b200_gemv_v2_level() >= 1
922}
923
924pub fn e4m3_row_ilp_level() -> u32 {
948 static LVL: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
949 *LVL.get_or_init(|| u32::from(std::env::var("MEMRA_E4M3_ROW_ILP").as_deref() == Ok("1")))
950}
951
952pub fn e4m3_row_ilp_on() -> bool {
953 e4m3_row_ilp_level() > 0
954}
955
956pub(crate) fn q8_row_ilp_on() -> bool {
963 q8_row_ilp_on_from(
964 std::env::var("MEMRA_Q8_ROW_ILP").ok().as_deref(),
965 env!("MEMRA_BUILT_CUDA_ARCH"),
966 )
967}
968
969pub fn q8_row_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
973 match v.map(str::trim) {
974 Some("1") => true,
975 Some("0") => false,
976 _ => built_arch == "100a",
977 }
978}
979
980pub static Q8_ROW_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
982 std::sync::atomic::AtomicU64::new(0);
983
984pub fn q8_row_ilp_dispatches() -> u64 {
986 Q8_ROW_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
987}
988
989fn q8_row_ilp_note(site: &str) {
990 if Q8_ROW_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
991 eprintln!(
992 "[q8-row-ilp] engaged at {site}: q8_0 row walk with four blocks' loads per lane \
993 ahead of the dp4a chains (MEMRA_Q8_ROW_ILP=1)"
994 );
995 }
996}
997
998pub(crate) fn nvfp4_row_ilp_on() -> bool {
1005 nvfp4_row_ilp_on_from(
1006 std::env::var("MEMRA_NVFP4_ROW_ILP").ok().as_deref(),
1007 env!("MEMRA_BUILT_CUDA_ARCH"),
1008 )
1009}
1010
1011pub fn nvfp4_row_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
1015 match v.map(str::trim) {
1016 Some("1") => true,
1017 Some("0") => false,
1018 _ => built_arch == "100a",
1019 }
1020}
1021
1022pub static NVFP4_ROW_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
1024 std::sync::atomic::AtomicU64::new(0);
1025
1026pub fn nvfp4_row_ilp_dispatches() -> u64 {
1028 NVFP4_ROW_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
1029}
1030
1031pub(crate) fn b200_mr1_fill() -> u32 {
1039 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1040 *V.get_or_init(|| {
1041 std::env::var("MEMRA_B200_MR1_FILL")
1042 .ok()
1043 .and_then(|v| v.trim().parse::<u32>().ok())
1044 .filter(|&n| n >= 1)
1045 .unwrap_or(16)
1046 })
1047}
1048
1049pub(crate) const GEMV_V3_STAGES: usize = 2;
1053
1054pub(crate) const Q8_V2_ROWS: u32 = 8;
1057pub(crate) fn gemv_v3_smem_bytes(nb: usize) -> usize {
1058 GEMV_V3_STAGES * GEMV_V2_ROWS * (nb * 8) * 2 + GEMV_V2_ROWS * nb * 4
1059}
1060
1061pub(crate) fn gemv_v3_fits() -> bool {
1066 gemv_v3_smem_bytes(mmv_block() as usize) <= 48 * 1024
1067}
1068
1069pub(crate) const GEMV_V2_ROWS: usize = 8;
1072
1073pub(crate) static GEMV_V2_Q8_ROWS_TW_DISPATCHES: std::sync::atomic::AtomicU64 =
1084 std::sync::atomic::AtomicU64::new(0);
1085pub(crate) static GEMV_V2_Q8_RP_DISPATCHES: std::sync::atomic::AtomicU64 =
1087 std::sync::atomic::AtomicU64::new(0);
1088pub(crate) static GEMV_V2_BF16_DISPATCHES: std::sync::atomic::AtomicU64 =
1090 std::sync::atomic::AtomicU64::new(0);
1091
1092pub(crate) fn b200_bf16_gemv_lt_on() -> bool {
1115 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1116 *V.get_or_init(|| {
1117 env!("MEMRA_BUILT_CUDA_ARCH") == "100a"
1118 && std::env::var("MEMRA_B200_BF16_GEMV_LT").as_deref() == Ok("1")
1119 })
1120}
1121
1122static STEP37_SERVING_DEFAULTS: std::sync::atomic::AtomicBool =
1148 std::sync::atomic::AtomicBool::new(false);
1149
1150pub fn arm_step37_serving_defaults() {
1151 STEP37_SERVING_DEFAULTS.store(true, std::sync::atomic::Ordering::Relaxed);
1152 crate::cache::set_swa_ring_default(true);
1153 eprintln!(
1154 "[step37-defaults] serving doors armed ON for the SlidingGatedMoe program \
1155 (per-flag =0 kills, =1 forces; owner flip 2026-08-27)"
1156 );
1157}
1158
1159pub(crate) fn step37_defaults_armed() -> bool {
1160 STEP37_SERVING_DEFAULTS.load(std::sync::atomic::Ordering::Relaxed)
1161}
1162
1163pub(crate) fn step37_door(cell: &'static std::sync::OnceLock<Option<bool>>, name: &str) -> bool {
1167 match *cell.get_or_init(|| match std::env::var(name).ok().as_deref() {
1168 Some("1") => Some(true),
1169 Some("0") => Some(false),
1170 _ => None,
1171 }) {
1172 Some(forced) => forced,
1173 None => step37_defaults_armed(),
1174 }
1175}
1176
1177pub(crate) fn w8_hybrid_on() -> bool {
1178 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1179 step37_door(&ENV, "MEMRA_W8_HYBRID")
1180}
1181
1182pub(crate) fn step_tp_w8_on() -> bool {
1183 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1184 step37_door(&ENV, "MEMRA_STEP_TP_W8")
1185}
1186
1187pub(crate) fn glm5_w8_on() -> bool {
1201 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1202 *ON.get_or_init(|| std::env::var("MEMRA_GLM5_W8").as_deref() == Ok("1"))
1203}
1204
1205pub(crate) static GLM5_W8_DISPATCHES: std::sync::atomic::AtomicU64 =
1208 std::sync::atomic::AtomicU64::new(0);
1209
1210pub(crate) fn w8_view_on() -> bool {
1215 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1216 *ON.get_or_init(|| std::env::var("MEMRA_W8_VIEW").as_deref() == Ok("1"))
1217}
1218
1219pub(crate) fn step_gemm_prime_on() -> bool {
1233 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1234 step37_door(&ENV, "MEMRA_STEP_GEMM_PRIME")
1235}
1236
1237pub(crate) fn q8t_wonce_on() -> bool {
1238 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
1239 step37_door(&ENV, "MEMRA_Q8T_WONCE")
1240}
1241
1242pub(crate) fn sig_expf_dev_on() -> bool {
1247 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1248 *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
1249}
1250
1251pub(crate) fn topk_fast_on() -> bool {
1252 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1253 *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
1254}
1255
1256fn sigmoid_topk_kernel(sig_expf: bool, fast: bool, n_used: usize) -> &'static str {
1261 match (sig_expf, fast && n_used <= 8) {
1262 (true, true) => "moe_router_sigmoid_topk_f32_dexp_fast",
1263 (true, false) => "moe_router_sigmoid_topk_f32_dexp",
1264 (false, true) => "moe_router_sigmoid_topk_f32_fast",
1265 (false, false) => "moe_router_sigmoid_topk_f32",
1266 }
1267}
1268
1269#[cfg(test)]
1270mod sigmoid_topk_dispatch_tests {
1271 #[test]
1272 fn fast_kernel_refuses_wide_topk_and_composes_with_dexp() {
1273 use super::sigmoid_topk_kernel;
1274
1275 assert_eq!(
1276 sigmoid_topk_kernel(false, true, 8),
1277 "moe_router_sigmoid_topk_f32_fast"
1278 );
1279 assert_eq!(
1280 sigmoid_topk_kernel(true, true, 8),
1281 "moe_router_sigmoid_topk_f32_dexp_fast"
1282 );
1283 assert_eq!(
1284 sigmoid_topk_kernel(false, true, 9),
1285 "moe_router_sigmoid_topk_f32"
1286 );
1287 assert_eq!(
1288 sigmoid_topk_kernel(true, true, 9),
1289 "moe_router_sigmoid_topk_f32_dexp"
1290 );
1291 }
1292}
1293
1294pub(crate) fn rms_block() -> u32 {
1295 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1296 *V.get_or_init(|| {
1297 std::env::var("MEMRA_RMS_BLOCK")
1298 .ok()
1299 .and_then(|v| v.parse().ok())
1300 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1301 })
1302}
1303
1304pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
1305 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1306 if let Some(forced) = *S.get_or_init(|| {
1307 std::env::var("MEMRA_FA_SPLIT")
1308 .ok()
1309 .and_then(|v| v.parse().ok())
1310 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
1311 }) {
1312 return forced;
1313 }
1314 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
1332 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
1333 {
1334 return if t_kv <= 8192 {
1335 16
1336 } else if t_kv <= 16384 {
1337 64
1338 } else {
1339 128
1340 };
1341 }
1342 let big_rig = fa_sm_count() >= 128;
1343 if big_rig {
1344 let _ = n_head_kv;
1345 if t_kv <= 2048 {
1346 static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
1355 if let Some(sp) = *SHORT.get_or_init(|| {
1356 std::env::var("MEMRA_FA_SP_SHORT")
1357 .ok()
1358 .and_then(|v| v.parse().ok())
1359 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
1360 }) {
1361 return sp;
1362 }
1363 16
1364 } else if t_kv <= 16384 {
1365 64
1366 } else {
1367 128
1368 }
1369 } else if n_head_kv <= 4 {
1370 if t_kv <= 512 {
1391 8
1392 } else if t_kv <= 16384 {
1393 64
1394 } else {
1395 128
1396 }
1397 } else {
1398 if t_kv <= 8192 {
1399 32
1400 } else if t_kv <= 16384 {
1401 64
1402 } else {
1403 128
1404 }
1405 }
1406}
1407
1408pub(crate) fn fa_sm_count() -> i32 {
1411 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
1412 *N.get_or_init(|| {
1413 cudarc::driver::result::init().ok();
1414 cudarc::driver::result::device::get(0)
1415 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
1416 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
1417 .unwrap_or(82)
1418 })
1419}
1420
1421#[allow(clippy::type_complexity)] fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
1426 match head_dim {
1427 256 => Ok(""),
1428 128 => Ok("_hd128"),
1429 d => Err(format!(
1430 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
1431 callers must gate to sdpa_naive"
1432 )
1433 .into()),
1434 }
1435}
1436
1437pub const QT_Q8_0: i32 = 0;
1439pub const QT_Q4_K: i32 = 1;
1440pub const QT_Q6_K: i32 = 2;
1441pub const QT_Q5_K: i32 = 3;
1442pub const QT_Q3_K: i32 = 4;
1443pub const QT_IQ4_XS: i32 = 5;
1444pub const QT_IQ3_S: i32 = 6;
1445pub const QT_NVFP4: i32 = 7;
1446pub const QT_NVFP4_V2: i32 = 107;
1449pub const QT_F8_E4M3: i32 = 10;
1455pub const QT_NVFP4_RP: i32 = 9;
1458pub const QT_F32: i32 = 8;
1460pub const QT_BF16: i32 = 11;
1461pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
1466pub const QT_F8_E4M3_BLK: i32 = 14;
1482
1483pub struct Engine {
1485 pub gpu: memra_runtime::Gpu,
1486 module: Arc<CudaModule>,
1487 hybrid: Arc<CudaModule>,
1488 kda: Arc<CudaModule>,
1490 qmatvec: Arc<CudaModule>,
1491 flash: Arc<CudaModule>,
1492 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
1496 gemm: Arc<CudaModule>,
1497 router: Arc<CudaModule>,
1498 sample: Arc<CudaModule>,
1500 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
1504 w8_mirrors: Mutex<std::collections::HashMap<(u64, u32, u32), CudaSlice<u8>>>,
1515 #[allow(clippy::type_complexity)]
1518 w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
1520 moe_cache_layout: Mutex<Option<Vec<usize>>>,
1524 capture_keep_on: std::sync::atomic::AtomicBool,
1530 verify_exact: std::sync::atomic::AtomicBool,
1535 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
1536 pub copy_stream: Arc<CudaStream>,
1538 #[cfg(memra_cutlass)]
1545 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
1546 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
1550 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
1553 #[allow(clippy::type_complexity)]
1557 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
1559 #[allow(clippy::type_complexity)]
1563 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
1565 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
1567 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
1568 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
1573 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
1578 router_stage: Mutex<Option<PinnedStage>>,
1582 hyper_decode_ws: Mutex<Option<crate::hyper::HyperDecodeWs>>,
1588 mla_seg_ws: Mutex<Option<crate::hybrid_forward::MlaSegWs>>,
1591 verify_ws: Mutex<VerifyWs>,
1603 vrows_macro_dev: Mutex<std::collections::HashMap<(u16, u8), CudaSlice<f32>>>,
1611 shexp_ones: Mutex<Option<CudaSlice<f32>>>,
1619}
1620
1621#[derive(Default)]
1626pub struct VerifyWs {
1627 f32_pool: std::collections::HashMap<usize, Vec<CudaSlice<f32>>>,
1628 i8_pool: std::collections::HashMap<usize, Vec<CudaSlice<i8>>>,
1629 u64_pool: std::collections::HashMap<usize, Vec<CudaSlice<u64>>>,
1630 held_bytes: usize,
1631}
1632
1633const VWS_PER_CLASS_CAP: usize = 16;
1636const VWS_HELD_BYTES_CAP: usize = 256 << 20;
1639
1640impl VerifyWs {
1641 fn take<T>(
1642 pool: &mut std::collections::HashMap<usize, Vec<CudaSlice<T>>>,
1643 held: &mut usize,
1644 n: usize,
1645 ) -> Option<CudaSlice<T>> {
1646 let s = pool.get_mut(&n)?.pop()?;
1647 *held -= n * std::mem::size_of::<T>();
1648 Some(s)
1649 }
1650 fn put<T>(
1651 pool: &mut std::collections::HashMap<usize, Vec<CudaSlice<T>>>,
1652 held: &mut usize,
1653 s: CudaSlice<T>,
1654 ) {
1655 let n = s.len();
1656 let bytes = n * std::mem::size_of::<T>();
1657 if *held + bytes > VWS_HELD_BYTES_CAP {
1658 return; }
1660 let v = pool.entry(n).or_default();
1661 if v.len() >= VWS_PER_CLASS_CAP {
1662 return;
1663 }
1664 v.push(s);
1665 *held += bytes;
1666 }
1667}
1668
1669pub static SCRATCH_ALLOC_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1675
1676fn fa_v2_on() -> bool {
1686 std::env::var("MEMRA_FA_V2")
1692 .map(|v| v != "0")
1693 .unwrap_or(true)
1694}
1695
1696pub(crate) fn fa_part_zero_on() -> bool {
1706 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1707 *ON.get_or_init(|| std::env::var("MEMRA_FA_PART_ZERO").as_deref() == Ok("1"))
1708}
1709
1710pub(crate) fn fa_v3_on() -> bool {
1711 std::env::var("MEMRA_FA_V3")
1715 .map(|v| v != "0")
1716 .unwrap_or(true)
1717}
1718
1719fn fa_v4_mode() -> &'static str {
1724 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
1725 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
1726}
1727fn fa_v4_on() -> bool {
1728 fa_v4_mode() != "0"
1729} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
1738 std::sync::atomic::AtomicUsize::new(1024);
1739pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
1740 std::sync::atomic::AtomicUsize::new(usize::MAX);
1741pub fn fa_v4_at_pub(t_kv: usize) -> bool {
1742 fa_v4_at(t_kv)
1743}
1744fn fa_v4_at(t_kv: usize) -> bool {
1745 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
1746 let mx = *M.get_or_init(|| {
1747 std::env::var("MEMRA_FA_V4_MAX")
1748 .ok()
1749 .and_then(|v| v.parse().ok())
1750 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1751 });
1752 fa_v4_on() && t_kv < mx
1753}
1754pub const FA_DEEP_MIN_DEFAULT: usize = 0;
1768fn fa_deep_at(t_kv: usize) -> bool {
1769 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
1770 return false;
1771 }
1772 let min = std::env::var("MEMRA_FA_DEEP_MIN")
1773 .ok()
1774 .and_then(|v| v.parse().ok())
1775 .unwrap_or(FA_DEEP_MIN_DEFAULT);
1776 t_kv >= min
1777}
1778pub fn fa_deep_at_pub(t_kv: usize) -> bool {
1780 fa_deep_at(t_kv)
1781}
1782
1783fn fa_v3_active(head_dim: usize) -> bool {
1784 fa_v3_on()
1787 && head_dim.is_multiple_of(128)
1788 && kv_cache_formats() == ("q8_0", "q5_1")
1789 && !Engine::kv_fp8_on()
1790}
1791
1792pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1800 std::env::var("MEMRA_NO_FA_VEC").is_err()
1801 && t_kv >= fa_vec_min_tkv()
1802 && head_dim == 256
1803 && fa_v4_at(t_kv)
1804 && !matches!(fa_v4_mode(), "noB3" | "stage")
1805 && !Engine::kv_fp8_on()
1806}
1807pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1809 fa_split_keys(t_kv, n_head_kv)
1810}
1811
1812struct PinnedStage {
1817 ptr: *mut u8,
1818 cap: usize,
1819}
1820unsafe impl Send for PinnedStage {}
1821impl PinnedStage {
1822 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1823 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1824 Ok(PinnedStage { ptr, cap })
1825 }
1826}
1827impl Drop for PinnedStage {
1828 fn drop(&mut self) {
1829 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1830 }
1831}
1832
1833pub struct PinnedHostBuf {
1840 ptr: *mut u8,
1841 len: usize,
1842}
1843unsafe impl Send for PinnedHostBuf {}
1846impl PinnedHostBuf {
1847 pub fn new(len: usize) -> Result<Self, Box<dyn std::error::Error>> {
1850 let ptr = unsafe { cudarc::driver::result::malloc_host(len.max(1), 0)? } as *mut u8;
1851 Ok(PinnedHostBuf { ptr, len })
1852 }
1853 pub fn len(&self) -> usize {
1854 self.len
1855 }
1856 pub fn is_empty(&self) -> bool {
1857 self.len == 0
1858 }
1859 pub fn as_slice(&self) -> &[u8] {
1860 unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
1861 }
1862 pub fn as_mut_slice(&mut self) -> &mut [u8] {
1863 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
1864 }
1865}
1866impl Drop for PinnedHostBuf {
1867 fn drop(&mut self) {
1868 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1869 }
1870}
1871
1872pub const ARGMAX_NB: usize = 256;
1875
1876pub(crate) use memra_fa3_vl as fa3_vl_raw;
1878
1879unsafe extern "C" {
1880 fn memra_fa3_prefill(
1882 q16: *const core::ffi::c_void,
1883 k16: *const core::ffi::c_void,
1884 v16: *const core::ffi::c_void,
1885 o: *mut f32,
1886 t: i32,
1887 h: i32,
1888 hkv: i32,
1889 d: i32,
1890 scale: f32,
1891 stream: *mut core::ffi::c_void,
1892 ) -> i32;
1893 pub(crate) fn memra_fa3_vl(
1895 q16s: *const *const core::ffi::c_void,
1896 k16s: *const *const core::ffi::c_void,
1897 v16s: *const *const core::ffi::c_void,
1898 os: *const *mut f32,
1899 ts: *const i32,
1900 b: i32,
1901 h: i32,
1902 hkv: i32,
1903 d: i32,
1904 scale: f32,
1905 stream: *mut core::ffi::c_void,
1906 ) -> i32;
1907}
1908
1909#[repr(C)]
1914#[derive(Clone, Copy)]
1915pub struct WPtr8(pub [u64; 8]);
1916unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1917
1918#[repr(C)]
1923#[derive(Clone, Copy, Default)]
1924pub struct GdnSeqVl {
1925 pub kb16: u64,
1926 pub gcum: u64,
1927 pub beta: u64,
1928 pub u: u64,
1929 pub wb16: u64,
1930 pub y: u64,
1931 pub ssnap: u64,
1932 pub state_in: u64,
1933 pub state_out: u64,
1934 pub q: u64,
1935 pub p: u64,
1936 pub o: u64,
1937 pub k: u64,
1938 pub v: u64,
1939 pub g: u64,
1940 pub a: u64,
1941 pub w: u64,
1942 pub t: i32,
1943 pub nc: i32,
1944}
1945unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1946#[repr(C)]
1947#[derive(Clone, Copy)]
1948pub struct GdnVl8(pub [GdnSeqVl; 8]);
1949unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1950
1951#[repr(C)]
1954#[derive(Clone, Copy, Default)]
1955pub struct GdnWVl {
1956 pub qb16: u64,
1957 pub pb16: u64,
1958}
1959unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1960#[repr(C)]
1961#[derive(Clone, Copy)]
1962pub struct GdnWVl8(pub [GdnWVl; 8]);
1963unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1964
1965#[repr(C)]
1967#[derive(Clone, Copy, Default)]
1968pub struct GdnPrepVl {
1969 pub qkv: u64,
1970 pub conv_state: u64,
1971 pub conv_out: u64,
1972 pub q_g: u64,
1973 pub k_g: u64,
1974 pub v_g: u64,
1975 pub q_l2: u64,
1976 pub k_l2: u64,
1977 pub beta_raw: u64,
1978 pub alpha: u64,
1979 pub beta: u64,
1980 pub g_log: u64,
1981 pub o: u64,
1982 pub z: u64,
1983 pub gn: u64,
1984 pub gn16: u64,
1985 pub kb16: u64,
1986 pub qb16: u64,
1987 pub t: i32,
1988 pub pad: i32,
1989}
1990unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1991#[repr(C)]
1992#[derive(Clone, Copy)]
1993pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1994unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1995
1996#[repr(C)]
1998#[derive(Clone, Copy, Default)]
1999pub struct FaSeqVl {
2000 pub q: u64,
2001 pub k16: u64,
2002 pub v16: u64,
2003 pub o: u64,
2004 pub kf: u64,
2005 pub vf: u64,
2006 pub t: i32,
2007 pub pad: i32,
2008}
2009unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
2010#[repr(C)]
2011#[derive(Clone, Copy)]
2012pub struct FaVl8(pub [FaSeqVl; 8]);
2013unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
2014
2015#[repr(C)]
2017#[derive(Clone, Copy, Default)]
2018pub struct AttnPreVl {
2019 pub qf: u64,
2020 pub kf: u64,
2021 pub vf: u64,
2022 pub q: u64,
2023 pub gate: u64,
2024 pub qn: u64,
2025 pub kn: u64,
2026 pub kc: u64,
2027 pub vc: u64,
2028 pub t: i32,
2029 pub pad: i32,
2030}
2031unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
2032#[repr(C)]
2033#[derive(Clone, Copy)]
2034pub struct AttnPreVl8(pub [AttnPreVl; 8]);
2035unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
2036
2037pub struct GdnChunkBufs {
2040 pub gcum: CudaSlice<f32>,
2041 pub a: CudaSlice<f32>,
2042 pub p: CudaSlice<f32>,
2043 pub u: CudaSlice<f32>,
2044 pub w: CudaSlice<f32>,
2045 pub kb16: CudaSlice<u8>,
2046 pub wb16: CudaSlice<u8>,
2047 pub y16: CudaSlice<u8>,
2048 pub ssnap16: CudaSlice<u8>,
2049 pub qb16: CudaSlice<u8>,
2050 pub pb16: CudaSlice<u8>,
2051 pub o: CudaSlice<f32>,
2052 pub t: usize,
2053 pub nc: usize,
2054}
2055
2056#[repr(C)]
2058#[derive(Clone, Copy)]
2059pub struct F32x8(pub [f32; 8]);
2060unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
2061
2062pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2066
2067pub static MOE_FUSED_EPI_DISPATCHES: std::sync::atomic::AtomicU64 =
2076 std::sync::atomic::AtomicU64::new(0);
2077
2078pub fn moe_fused_epilogue_dispatches() -> u64 {
2082 MOE_FUSED_EPI_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2083}
2084
2085pub static MOE_VROWS_DISPATCHES: std::sync::atomic::AtomicU64 =
2092 std::sync::atomic::AtomicU64::new(0);
2093
2094pub fn moe_vrows_dispatches() -> u64 {
2096 MOE_VROWS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2097}
2098
2099fn bf16_tcols_wide_on() -> bool {
2108 std::env::var("MEMRA_BF16_TCOLS_WIDE").as_deref() != Ok("0")
2109}
2110
2111pub static BF16_TCOLS_WIDE_DISPATCHES: std::sync::atomic::AtomicU64 =
2115 std::sync::atomic::AtomicU64::new(0);
2116
2117pub fn bf16_tcols_wide_dispatches() -> u64 {
2119 BF16_TCOLS_WIDE_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2120}
2121
2122fn bf16_tcols_x1_on() -> bool {
2130 std::env::var("MEMRA_BF16_TCOLS_X1").as_deref() != Ok("0")
2131}
2132
2133pub static BF16_TCOLS_X1_DISPATCHES: std::sync::atomic::AtomicU64 =
2135 std::sync::atomic::AtomicU64::new(0);
2136
2137pub fn bf16_tcols_x1_dispatches() -> u64 {
2139 BF16_TCOLS_X1_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2140}
2141
2142fn bf16_tcols_red_fused_on() -> bool {
2155 std::env::var("MEMRA_BF16_TCOLS_RED_FUSED").as_deref() == Ok("1")
2156}
2157
2158pub static BF16_TCOLS_RED_FUSED_DISPATCHES: std::sync::atomic::AtomicU64 =
2161 std::sync::atomic::AtomicU64::new(0);
2162
2163pub fn bf16_tcols_red_fused_dispatches() -> u64 {
2165 BF16_TCOLS_RED_FUSED_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2166}
2167
2168pub(crate) fn moe_vrows_pack_on() -> bool {
2175 std::env::var("MEMRA_MOE_VROWS_PACK").as_deref() == Ok("1")
2176}
2177
2178pub static MOE_VROWS_PACK_DISPATCHES: std::sync::atomic::AtomicU64 =
2180 std::sync::atomic::AtomicU64::new(0);
2181
2182pub fn moe_vrows_pack_dispatches() -> u64 {
2184 MOE_VROWS_PACK_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2185}
2186
2187pub(crate) fn moe_vrows_ilp_on() -> bool {
2194 moe_vrows_ilp_on_from(
2195 std::env::var("MEMRA_MOE_VROWS_ILP").ok().as_deref(),
2196 env!("MEMRA_BUILT_CUDA_ARCH"),
2197 )
2198}
2199
2200pub fn moe_vrows_ilp_on_from(v: Option<&str>, built_arch: &str) -> bool {
2205 match v.map(str::trim) {
2206 Some("1") => true,
2207 Some("0") => false,
2208 _ => built_arch == "100a",
2209 }
2210}
2211
2212pub static MOE_VROWS_ILP_DISPATCHES: std::sync::atomic::AtomicU64 =
2214 std::sync::atomic::AtomicU64::new(0);
2215
2216pub fn moe_vrows_ilp_dispatches() -> u64 {
2218 MOE_VROWS_ILP_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2219}
2220
2221fn moe_vrows_ilp_refuse(site: &str, qt: i32) {
2224 static SAID: std::sync::Once = std::sync::Once::new();
2225 SAID.call_once(|| {
2226 eprintln!(
2227 "[moe-vrows-ilp] REFUSED at {site}: MEMRA_MOE_VROWS_ILP is on but the expert qtype is \
2228 {qt}, neither QT_NVFP4 ({QT_NVFP4}) nor QT_NVFP4_V2 ({QT_NVFP4_V2}); the shipped \
2229 kernel runs (the ILP twins hoist the NVFP4 group loads, interleaved or slot-major, \
2230 and no other layout)"
2231 );
2232 });
2233}
2234
2235pub(crate) fn moe_vrows_dev_tables_on() -> bool {
2245 std::env::var("MEMRA_MOE_VROWS_DEV_TABLES").as_deref() == Ok("1")
2246}
2247
2248pub static MOE_VROWS_DEV_TABLES_DISPATCHES: std::sync::atomic::AtomicU64 =
2250 std::sync::atomic::AtomicU64::new(0);
2251
2252pub fn moe_vrows_dev_tables_dispatches() -> u64 {
2254 MOE_VROWS_DEV_TABLES_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2255}
2256
2257pub static MOE_VROWS_ROUTER_SYNCS_AVOIDED: std::sync::atomic::AtomicU64 =
2261 std::sync::atomic::AtomicU64::new(0);
2262
2263pub fn moe_vrows_router_syncs_avoided() -> u64 {
2265 MOE_VROWS_ROUTER_SYNCS_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2266}
2267
2268fn moe_vrows_dedup_stat_on() -> bool {
2278 std::env::var("MEMRA_MOE_VROWS_DEDUP_STAT").as_deref() == Ok("1")
2279}
2280
2281pub static MOE_VROWS_PAIR_VISITS: std::sync::atomic::AtomicU64 =
2283 std::sync::atomic::AtomicU64::new(0);
2284
2285pub static MOE_VROWS_PAIR_DISTINCT: std::sync::atomic::AtomicU64 =
2289 std::sync::atomic::AtomicU64::new(0);
2290
2291pub(crate) fn vrows_overlap_counts(sel_all: &[u32]) -> (u64, u64) {
2298 let mut seen = std::collections::HashSet::with_capacity(sel_all.len());
2299 for &ex in sel_all {
2300 seen.insert(ex);
2301 }
2302 (sel_all.len() as u64, seen.len() as u64)
2303}
2304
2305static MOE_VROWS_DEDUP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2307
2308fn moe_vrows_dedup_report() {
2315 let n = MOE_VROWS_DEDUP_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2316 if n != 0 && !n.is_multiple_of(42) {
2317 return;
2318 }
2319 let (visits, distinct) = moe_vrows_pair_overlap();
2320 if visits == 0 {
2321 return;
2322 }
2323 let repeat = 100.0 * (1.0 - distinct as f64 / visits as f64);
2324 eprintln!(
2325 "[moe-vrows-dedup] layer-calls={} visits={visits} distinct={distinct} \
2326 repeat={repeat:.2}% = the cross-row expert-slab dedup ceiling on the vrows pair \
2327 (MEMRA_MOE_VROWS_DEDUP_STAT=1)",
2328 n + 1
2329 );
2330}
2331
2332pub fn vrows_overlap_counts_for_test(sel_all: &[u32]) -> (u64, u64) {
2335 vrows_overlap_counts(sel_all)
2336}
2337
2338pub fn moe_vrows_pair_overlap() -> (u64, u64) {
2340 (
2341 MOE_VROWS_PAIR_VISITS.load(std::sync::atomic::Ordering::Relaxed),
2342 MOE_VROWS_PAIR_DISTINCT.load(std::sync::atomic::Ordering::Relaxed),
2343 )
2344}
2345
2346pub(crate) fn moe_vrows_dedup_order_on() -> bool {
2363 std::env::var("MEMRA_MOE_VROWS_DEDUP_ORDER").as_deref() == Ok("1")
2364}
2365
2366fn moe_vrows_down_tmaj_on() -> bool {
2375 std::env::var("MEMRA_MOE_VROWS_DOWN_TMAJ").as_deref() == Ok("1")
2376}
2377
2378pub static MOE_VROWS_DEDUP_ORDER_DISPATCHES: std::sync::atomic::AtomicU64 =
2380 std::sync::atomic::AtomicU64::new(0);
2381
2382pub fn moe_vrows_dedup_order_dispatches() -> u64 {
2384 MOE_VROWS_DEDUP_ORDER_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2385}
2386
2387pub static MOE_VROWS_DOWN_TMAJ_DISPATCHES: std::sync::atomic::AtomicU64 =
2389 std::sync::atomic::AtomicU64::new(0);
2390
2391pub fn moe_vrows_down_tmaj_dispatches() -> u64 {
2393 MOE_VROWS_DOWN_TMAJ_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2394}
2395
2396pub static MOE_VROWS_SLAB_READS_AVOIDED: std::sync::atomic::AtomicU64 =
2407 std::sync::atomic::AtomicU64::new(0);
2408
2409pub fn moe_vrows_slab_reads_avoided() -> u64 {
2411 MOE_VROWS_SLAB_READS_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2412}
2413
2414pub(crate) fn vrows_expert_major_order(sel_all: &[u32]) -> Vec<u64> {
2419 let mut ord: Vec<u64> = (0..sel_all.len() as u64).collect();
2420 ord.sort_by_key(|&p| sel_all[p as usize]);
2423 ord
2424}
2425
2426pub fn vrows_expert_major_order_for_test(sel_all: &[u32]) -> Vec<u64> {
2429 vrows_expert_major_order(sel_all)
2430}
2431
2432pub(crate) fn alias_door_from(
2453 general: (&'static str, Option<&str>),
2454 alias: (&'static str, Option<&str>),
2455) -> Result<(bool, &'static str), String> {
2456 match (general.1, alias.1) {
2457 (Some(g), Some(a)) if g != a => Err(format!(
2458 "{}={g:?} and {}={a:?} disagree — the alias and the general flag name ONE door \
2459 (unset one); refused rather than silently picking a precedence winner, and the \
2460 door falls closed to the shipped program",
2461 general.0, alias.0
2462 )),
2463 (Some(g), _) => Ok((g == "1", general.0)),
2464 (None, Some(a)) => Ok((a == "1", alias.0)),
2465 (None, None) => Ok((false, general.0)),
2466 }
2467}
2468
2469fn alias_door(
2486 general: &'static str,
2487 alias: &'static str,
2488 latch: &'static std::sync::atomic::AtomicBool,
2489) -> (bool, &'static str) {
2490 let g = std::env::var(general).ok();
2491 let a = std::env::var(alias).ok();
2492 match alias_door_from((general, g.as_deref()), (alias, a.as_deref())) {
2493 Ok(resolved) => resolved,
2494 Err(msg) => {
2495 if !latch.swap(true, std::sync::atomic::Ordering::Relaxed) {
2496 eprintln!("[flag-alias] {msg}");
2497 }
2498 (false, general)
2499 }
2500 }
2501}
2502
2503pub fn htod_diet_on() -> bool {
2523 htod_diet_armed().0
2524}
2525
2526static HTOD_DIET_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2528 std::sync::atomic::AtomicBool::new(false);
2529
2530pub(crate) fn htod_diet_armed() -> (bool, &'static str) {
2532 alias_door(
2533 "MEMRA_HTOD_DIET",
2534 "MEMRA_GLM5_HTOD_DIET",
2535 &HTOD_DIET_ALIAS_WARNED,
2536 )
2537}
2538
2539pub static HTOD_DIET_AVOIDED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2542
2543pub fn htod_diet_avoided() -> u64 {
2545 HTOD_DIET_AVOIDED.load(std::sync::atomic::Ordering::Relaxed)
2546}
2547
2548pub fn ep_diet_on() -> bool {
2569 ep_diet_armed().0
2570}
2571
2572static EP_DIET_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2574 std::sync::atomic::AtomicBool::new(false);
2575
2576pub(crate) fn ep_diet_armed() -> (bool, &'static str) {
2579 alias_door("MEMRA_EP_DIET", "MEMRA_GLM5_EP_DIET", &EP_DIET_ALIAS_WARNED)
2580}
2581
2582pub fn ep_grouped_prime_on() -> bool {
2600 ep_grouped_prime_armed().0
2601}
2602
2603static EP_GROUPED_PRIME_ALIAS_WARNED: std::sync::atomic::AtomicBool =
2605 std::sync::atomic::AtomicBool::new(false);
2606
2607pub(crate) fn ep_grouped_prime_armed() -> (bool, &'static str) {
2609 alias_door(
2610 "MEMRA_EP_GROUPED_PRIME",
2611 "MEMRA_GLM5_EP_GROUPED_PRIME",
2612 &EP_GROUPED_PRIME_ALIAS_WARNED,
2613 )
2614}
2615
2616fn topk_shards_on() -> bool {
2625 std::env::var("MEMRA_TOPK_SHARDS").as_deref() != Ok("0")
2626}
2627
2628pub static TOPK_SHARDS_DISPATCHES: std::sync::atomic::AtomicU64 =
2630 std::sync::atomic::AtomicU64::new(0);
2631
2632pub fn topk_shards_dispatches() -> u64 {
2634 TOPK_SHARDS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2635}
2636
2637pub fn alloc_trace_on() -> bool {
2644 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2645 *V.get_or_init(|| std::env::var("MEMRA_ALLOC_TRACE").as_deref() == Ok("1"))
2646}
2647
2648#[track_caller]
2649pub(crate) fn alloc_trace_hit(bytes: usize) {
2650 if alloc_trace_on() {
2651 let loc = std::panic::Location::caller();
2652 eprintln!(
2653 "[alloc-trace] {bytes} bytes from {}:{}",
2654 loc.file(),
2655 loc.line()
2656 );
2657 }
2658}
2659
2660pub fn dtoh_trace_on() -> bool {
2667 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2668 *V.get_or_init(|| std::env::var("MEMRA_DTOH_TRACE").as_deref() == Ok("1"))
2669}
2670
2671#[track_caller]
2672pub(crate) fn dtoh_trace_hit(bytes: usize) {
2673 if dtoh_trace_on() {
2674 let loc = std::panic::Location::caller();
2675 eprintln!(
2676 "[dtoh-trace] {bytes} bytes from {}:{}",
2677 loc.file(),
2678 loc.line()
2679 );
2680 }
2681}
2682
2683pub fn moe_expert_rp_on() -> bool {
2696 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2697 *V.get_or_init(|| std::env::var("MEMRA_MOE_EXPERT_RP").as_deref() == Ok("1"))
2698}
2699
2700pub fn rp_qt(rp: bool, qt: i32) -> i32 {
2704 if rp && qt == QT_NVFP4 {
2705 QT_NVFP4_V2
2706 } else {
2707 qt
2708 }
2709}
2710
2711pub fn moe_rp_refuse(rp: bool, path: &str) -> Result<(), Box<dyn std::error::Error>> {
2714 if rp {
2715 return Err(format!(
2716 "{path}: the resident expert slab is split-plane (MEMRA_MOE_EXPERT_RP=1) and this \
2717 path reads experts interleaved; it is not wired for the door (memra#147). Boot \
2718 without MEMRA_MOE_EXPERT_RP for this model or wire the path."
2719 )
2720 .into());
2721 }
2722 Ok(())
2723}
2724
2725pub(crate) fn hyper_batch_solo_on() -> bool {
2752 std::env::var("MEMRA_HYPER_BATCH_SOLO").as_deref() == Ok("1")
2753}
2754
2755pub(crate) fn hc_pre_sink_reg() -> bool {
2799 hc_pre_sink_reg_from(
2800 std::env::var("MEMRA_HC_PRE_SINK_REG").ok().as_deref(),
2801 env!("MEMRA_BUILT_CUDA_ARCH"),
2802 )
2803}
2804
2805pub fn hc_pre_sink_reg_from(v: Option<&str>, built_arch: &str) -> bool {
2809 match v.map(str::trim) {
2810 Some("1") => true,
2811 Some("0") => false,
2812 _ => built_arch == "100a",
2813 }
2814}
2815
2816pub fn hc_pre_split_collapse() -> bool {
2833 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2834 *ON.get_or_init(|| std::env::var("MEMRA_HC_PRE_SPLIT_COLLAPSE").as_deref() == Ok("1"))
2835}
2836
2837pub fn hc_pre_block_default(built_arch: &str) -> usize {
2839 if built_arch == "100a" { 512 } else { 128 }
2840}
2841
2842pub(crate) fn hc_pre_block() -> usize {
2843 let default = hc_pre_block_default(env!("MEMRA_BUILT_CUDA_ARCH"));
2847 match std::env::var("MEMRA_HC_PRE_BLOCK") {
2848 Err(_) => default,
2849 Ok(v) if v.is_empty() => default,
2850 Ok(v) => match v.parse::<usize>() {
2851 Ok(n) if (32..=1024).contains(&n) && n.is_power_of_two() => n,
2852 _ => {
2853 static SAID: std::sync::Once = std::sync::Once::new();
2854 SAID.call_once(|| {
2855 eprintln!(
2856 "[hc-pre-block] MEMRA_HC_PRE_BLOCK={v:?} is not a power of two in \
2857 [32, 1024]; using the build default {default}"
2858 );
2859 });
2860 default
2861 }
2862 },
2863 }
2864}
2865
2866fn verify_ws_on() -> bool {
2867 verify_ws_on_from(
2868 std::env::var("MEMRA_VERIFY_WS").ok().as_deref(),
2869 std::env::var("MEMRA_GLM5_VERIFY_WS").ok().as_deref(),
2870 )
2871}
2872
2873fn verify_ws_on_from(general: Option<&str>, glm5_alias: Option<&str>) -> bool {
2876 general != Some("0") && glm5_alias != Some("0")
2877}
2878
2879pub static VERIFY_WS_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2883
2884pub fn verify_ws_hits() -> u64 {
2886 VERIFY_WS_HITS.load(std::sync::atomic::Ordering::Relaxed)
2887}
2888
2889pub static MLA_TC_PREFILL_DISPATCHES: std::sync::atomic::AtomicU64 =
2896 std::sync::atomic::AtomicU64::new(0);
2897
2898pub fn mla_tc_prefill_dispatches() -> u64 {
2902 MLA_TC_PREFILL_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2903}
2904
2905pub static MOE_GROUPED_PREFILL_DISPATCHES: std::sync::atomic::AtomicU64 =
2911 std::sync::atomic::AtomicU64::new(0);
2912
2913pub fn moe_grouped_prefill_dispatches() -> u64 {
2916 MOE_GROUPED_PREFILL_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
2917}
2918
2919#[must_use = "dropping immediately ends the exact scope"]
2924pub struct ExactScope<'a> {
2925 flag: &'a std::sync::atomic::AtomicBool,
2926 prev: bool,
2927}
2928
2929impl<'a> ExactScope<'a> {
2930 pub(crate) fn set(flag: &'a std::sync::atomic::AtomicBool, on: bool) -> Self {
2931 let prev = flag.load(std::sync::atomic::Ordering::Relaxed);
2932 flag.store(on, std::sync::atomic::Ordering::Relaxed);
2933 ExactScope { flag, prev }
2934 }
2935}
2936
2937impl Drop for ExactScope<'_> {
2938 fn drop(&mut self) {
2939 self.flag
2940 .store(self.prev, std::sync::atomic::Ordering::Relaxed);
2941 }
2942}
2943
2944#[cfg(test)]
2945mod verify_ws_flag_tests {
2946 use super::verify_ws_on_from;
2947
2948 #[test]
2949 fn off_wins_across_general_and_alias() {
2950 assert!(verify_ws_on_from(None, None));
2952 assert!(!verify_ws_on_from(Some("0"), None));
2955 assert!(!verify_ws_on_from(None, Some("0")));
2956 assert!(!verify_ws_on_from(Some("1"), Some("0")));
2957 assert!(!verify_ws_on_from(Some("0"), Some("1")));
2958 assert!(verify_ws_on_from(Some("1"), None));
2960 assert!(verify_ws_on_from(None, Some("1")));
2961 }
2962}
2963
2964#[cfg(test)]
2965mod moe_vrows_ilp_default_tests {
2966 use super::moe_vrows_ilp_on_from;
2967
2968 #[test]
2969 fn nvfp4_row_ilp_arch_keyed_default() {
2970 use super::nvfp4_row_ilp_on_from;
2971 assert!(nvfp4_row_ilp_on_from(None, "100a"));
2972 assert!(!nvfp4_row_ilp_on_from(None, "120a"));
2973 assert!(nvfp4_row_ilp_on_from(Some("1"), "120a"));
2974 assert!(!nvfp4_row_ilp_on_from(Some("0"), "100a"));
2975 }
2976
2977 #[test]
2978 fn q8_row_ilp_arch_keyed_default() {
2979 use super::q8_row_ilp_on_from;
2980 assert!(q8_row_ilp_on_from(None, "100a"));
2981 assert!(!q8_row_ilp_on_from(None, "120a"));
2982 assert!(q8_row_ilp_on_from(Some("1"), "120a"));
2983 assert!(!q8_row_ilp_on_from(Some("0"), "100a"));
2984 }
2985
2986 #[test]
2987 fn arch_keyed_default_with_explicit_override() {
2988 assert!(moe_vrows_ilp_on_from(None, "100a"));
2989 assert!(!moe_vrows_ilp_on_from(None, "120a"));
2990 assert!(!moe_vrows_ilp_on_from(None, "90a"));
2991 assert!(moe_vrows_ilp_on_from(Some("1"), "120a"));
2992 assert!(!moe_vrows_ilp_on_from(Some("0"), "100a"));
2993 assert!(!moe_vrows_ilp_on_from(Some(" 0 "), "100a"));
2994 }
2995}
2996
2997#[cfg(test)]
2998mod hc_pre_default_tests {
2999 use super::{hc_pre_block_default, hc_pre_sink_reg_from};
3000
3001 #[test]
3002 fn arch_keyed_defaults_with_explicit_override() {
3003 assert_eq!(hc_pre_block_default("100a"), 512);
3004 assert_eq!(hc_pre_block_default("120a"), 128);
3005 assert!(hc_pre_sink_reg_from(None, "100a"));
3006 assert!(!hc_pre_sink_reg_from(None, "120a"));
3007 assert!(hc_pre_sink_reg_from(Some("1"), "120a"));
3008 assert!(!hc_pre_sink_reg_from(Some("0"), "100a"));
3009 assert_eq!(
3010 crate::hyper::hc_fused_pre_arm_from(None, "100a"),
3011 crate::hyper::HcFusedPreArm::V2
3012 );
3013 assert_eq!(
3014 crate::hyper::hc_fused_pre_arm_from(None, "120a"),
3015 crate::hyper::HcFusedPreArm::Off
3016 );
3017 assert_eq!(
3018 crate::hyper::hc_fused_pre_arm_from(Some("0"), "100a"),
3019 crate::hyper::HcFusedPreArm::Off
3020 );
3021 assert_eq!(
3022 crate::hyper::hc_fused_pre_arm_from(Some("1"), "100a"),
3023 crate::hyper::HcFusedPreArm::V1
3024 );
3025 }
3026}
3027
3028#[cfg(test)]
3029mod glm5_decode_graph_default_tests {
3030 use super::glm5_decode_graph_on_from;
3031
3032 #[test]
3035 fn unset_and_one_arm_only_zero_disarms() {
3036 assert!(glm5_decode_graph_on_from(None));
3037 assert!(glm5_decode_graph_on_from(Some("1")));
3038 assert!(glm5_decode_graph_on_from(Some("")));
3039 assert!(!glm5_decode_graph_on_from(Some("0")));
3040 assert!(!glm5_decode_graph_on_from(Some(" 0 ")));
3041 }
3042}
3043
3044#[cfg(test)]
3045mod alias_door_tests {
3046 use super::alias_door_from;
3047
3048 const G: &str = "MEMRA_EP_DIET";
3049 const A: &str = "MEMRA_GLM5_EP_DIET";
3050
3051 fn r(g: Option<&str>, a: Option<&str>) -> Result<(bool, &'static str), String> {
3052 alias_door_from((G, g), (A, a))
3053 }
3054
3055 #[test]
3056 fn default_off_and_either_name_arms() {
3057 assert_eq!(r(None, None).unwrap(), (false, G));
3059 assert_eq!(r(Some("1"), None).unwrap(), (true, G));
3061 assert_eq!(r(None, Some("1")).unwrap(), (true, A));
3062 assert_eq!(r(Some("0"), None).unwrap(), (false, G));
3064 assert_eq!(r(None, Some("0")).unwrap(), (false, A));
3065 assert_eq!(r(None, Some("on")).unwrap(), (false, A));
3067 assert_eq!(r(Some(""), None).unwrap(), (false, G));
3068 }
3069
3070 #[test]
3071 fn agreeing_pair_resolves_to_the_general_name() {
3072 assert_eq!(r(Some("1"), Some("1")).unwrap(), (true, G));
3073 assert_eq!(r(Some("0"), Some("0")).unwrap(), (false, G));
3074 }
3075
3076 #[test]
3077 fn disagreeing_pair_refuses_and_names_both() {
3078 for (g, a) in [("1", "0"), ("0", "1")] {
3079 let err = r(Some(g), Some(a)).expect_err("a disagreeing pair must refuse");
3080 assert!(
3081 err.contains(G),
3082 "the refusal must name the general flag: {err}"
3083 );
3084 assert!(err.contains(A), "the refusal must name the alias: {err}");
3085 assert!(err.contains("falls closed"), "{err}");
3088 }
3089 }
3090}
3091
3092#[cfg(test)]
3093mod exact_scope_tests {
3094 use std::sync::atomic::{AtomicBool, Ordering};
3095
3096 #[test]
3097 fn error_path_restores_verify_exact() {
3098 let flag = AtomicBool::new(false);
3103 let failing = |flag: &AtomicBool| -> Result<(), &'static str> {
3104 let _scope = super::ExactScope::set(flag, true);
3105 assert!(flag.load(Ordering::Relaxed), "scope arms the flag");
3106 Err("draft forward failed")? };
3108 assert!(failing(&flag).is_err());
3109 assert!(
3110 !flag.load(Ordering::Relaxed),
3111 "error propagation must restore the pre-scope value"
3112 );
3113 let flag = AtomicBool::new(true);
3115 {
3116 let _scope = super::ExactScope::set(&flag, true);
3117 }
3118 assert!(flag.load(Ordering::Relaxed));
3119 let flag = AtomicBool::new(false);
3121 let scope = super::ExactScope::set(&flag, true);
3122 drop(scope);
3123 assert!(!flag.load(Ordering::Relaxed));
3124 }
3125}
3126
3127impl Engine {
3128 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
3129 let gpu = memra_runtime::Gpu::new(ordinal)?;
3130 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
3134 use cudarc::driver::sys::CUdevice_attribute_enum as A;
3135 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
3136 .and_then(|d| unsafe {
3137 Ok((
3138 cudarc::driver::result::device::get_attribute(
3139 d,
3140 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
3141 )?,
3142 cudarc::driver::result::device::get_attribute(
3143 d,
3144 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
3145 )?,
3146 ))
3147 })
3148 .unwrap_or((0, 0));
3149 let built = env!("MEMRA_BUILT_CUDA_ARCH");
3150 let ok = matches!(
3151 (built, maj, min),
3152 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
3153 );
3154 if !ok {
3155 return Err(format!(
3156 "memra was built for sm_{built} but device {ordinal} reports compute \
3157 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
3158 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
3159 )
3160 .into());
3161 }
3162 }
3163 unsafe {
3168 use cudarc::driver::sys;
3169 let dev: sys::CUdevice = ordinal as sys::CUdevice;
3170 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3171 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
3172 let mut thresh: u64 = u64::MAX;
3173 let _ = sys::cuMemPoolSetAttribute(
3174 pool,
3175 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
3176 &mut thresh as *mut u64 as *mut core::ffi::c_void,
3177 );
3178 }
3179 }
3180 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
3181 let hybrid = gpu
3182 .ctx
3183 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
3184 let kda = gpu.ctx.load_module(Ptx::from_binary(KDA_FATBIN.to_vec()))?;
3185 let qmatvec = gpu
3186 .ctx
3187 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
3188 let flash = gpu
3189 .ctx
3190 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
3191 let gemm = gpu
3192 .ctx
3193 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
3194 let router = gpu
3195 .ctx
3196 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
3197 let sample = gpu
3198 .ctx
3199 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
3200 let copy_stream = gpu.ctx.new_stream()?;
3201 if std::env::var("MEMRA_EVT")
3217 .map(|v| v == "1")
3218 .unwrap_or(false)
3219 {
3220 } else {
3222 unsafe {
3223 gpu.ctx.disable_event_tracking();
3224 }
3225 }
3226 Ok(Self {
3227 gpu,
3228 module,
3229 hybrid,
3230 kda,
3231 qmatvec,
3232 flash,
3233 flash_g: std::sync::OnceLock::new(),
3234 gemm,
3235 router,
3236 sample,
3237 moe_cache: Mutex::new(None),
3238 w8_mirrors: Mutex::new(std::collections::HashMap::new()),
3239 w8_act: Mutex::new(std::collections::HashMap::new()),
3240 moe_cache_layout: Mutex::new(None),
3241 copy_stream,
3242 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
3243 verify_exact: std::sync::atomic::AtomicBool::new(false),
3244 capture_keep: Mutex::new(Vec::new()),
3245 argmax_partials: Mutex::new(None),
3246 prime_deqw_ws: Mutex::new(None),
3247 router_stage: Mutex::new(None),
3248 hyper_decode_ws: Mutex::new(None),
3249 mla_seg_ws: Mutex::new(None),
3250 verify_ws: Mutex::new(VerifyWs::default()),
3251 vrows_macro_dev: Mutex::new(std::collections::HashMap::new()),
3252 shexp_ones: Mutex::new(None),
3253 fp8_scratch: Mutex::new(None),
3254 fa_vf16_scratch: Mutex::new(None),
3255 fa_part_pool: Mutex::new(None),
3256 fa_part_retired: Mutex::new(Vec::new()),
3257 fn_cache: Mutex::new(Default::default()),
3258 f16_scratch: Mutex::new(None),
3259 #[cfg(memra_cutlass)]
3260 cutlass_scratch: Mutex::new(None),
3261 })
3262 }
3263
3264 pub fn ctx(&self) -> &Arc<CudaContext> {
3265 &self.gpu.ctx
3266 }
3267
3268 pub fn pool_cached_bytes(&self) -> usize {
3286 let (reserved, used) = self.pool_reserved_used();
3287 reserved.saturating_sub(used)
3288 }
3289
3290 pub fn device_graph_mem_reserved(&self) -> usize {
3300 use cudarc::driver::sys as cus;
3301 let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
3302 return 0;
3303 };
3304 let mut bytes: u64 = 0;
3305 let rc = unsafe {
3306 cus::cuDeviceGetGraphMemAttribute(
3307 dev,
3308 cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
3309 &mut bytes as *mut u64 as *mut std::ffi::c_void,
3310 )
3311 };
3312 if rc == cus::cudaError_enum::CUDA_SUCCESS {
3313 bytes as usize
3314 } else {
3315 0
3316 }
3317 }
3318
3319 pub fn pool_trim_to_zero(&self) -> usize {
3333 use cudarc::driver::sys;
3334 let (before, _) = self.pool_reserved_used();
3335 unsafe {
3336 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3337 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3338 != sys::CUresult::CUDA_SUCCESS
3339 {
3340 return 0;
3341 }
3342 let _ = sys::cuMemPoolTrimTo(pool, 0);
3343 }
3344 let (after, _) = self.pool_reserved_used();
3345 before.saturating_sub(after)
3346 }
3347
3348 pub fn pool_reserved_used(&self) -> (usize, usize) {
3349 use cudarc::driver::sys;
3350 unsafe {
3351 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3352 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3353 != sys::CUresult::CUDA_SUCCESS
3354 {
3355 return (0, 0);
3356 }
3357 let (mut reserved, mut used) = (0u64, 0u64);
3358 if sys::cuMemPoolGetAttribute(
3359 pool,
3360 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
3361 &mut reserved as *mut u64 as *mut core::ffi::c_void,
3362 ) != sys::CUresult::CUDA_SUCCESS
3363 {
3364 return (0, 0);
3365 }
3366 if sys::cuMemPoolGetAttribute(
3367 pool,
3368 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
3369 &mut used as *mut u64 as *mut core::ffi::c_void,
3370 ) != sys::CUresult::CUDA_SUCCESS
3371 {
3372 return (0, 0);
3373 }
3374 (reserved as usize, used as usize)
3375 }
3376 }
3377
3378 pub fn pool_high_water_reset(&self) -> (usize, usize) {
3387 use cudarc::driver::sys;
3388 unsafe {
3389 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
3390 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
3391 != sys::CUresult::CUDA_SUCCESS
3392 {
3393 return (0, 0);
3394 }
3395 let (mut reserved, mut used) = (0u64, 0u64);
3396 if sys::cuMemPoolGetAttribute(
3397 pool,
3398 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
3399 &mut reserved as *mut u64 as *mut core::ffi::c_void,
3400 ) != sys::CUresult::CUDA_SUCCESS
3401 {
3402 return (0, 0);
3403 }
3404 if sys::cuMemPoolGetAttribute(
3405 pool,
3406 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
3407 &mut used as *mut u64 as *mut core::ffi::c_void,
3408 ) != sys::CUresult::CUDA_SUCCESS
3409 {
3410 return (0, 0);
3411 }
3412 let mut zero: u64 = 0;
3415 let _ = sys::cuMemPoolSetAttribute(
3416 pool,
3417 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
3418 &mut zero as *mut u64 as *mut core::ffi::c_void,
3419 );
3420 let mut zero2: u64 = 0;
3421 let _ = sys::cuMemPoolSetAttribute(
3422 pool,
3423 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_HIGH,
3424 &mut zero2 as *mut u64 as *mut core::ffi::c_void,
3425 );
3426 (reserved as usize, used as usize)
3427 }
3428 }
3429
3430 pub fn stream(&self) -> Arc<CudaStream> {
3433 self.gpu.stream()
3434 }
3435 pub fn gkv_on() -> bool {
3438 memra_kv::gkv_on()
3439 }
3440
3441 pub fn wkv_on() -> bool {
3453 memra_kv::wkv_on()
3454 }
3455
3456 pub fn kv_fp8_on() -> bool {
3462 memra_kv::kv_fp8_on()
3463 }
3464
3465 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
3468 if head_dim == 512 && Self::gkv_on() {
3469 self.func_g(name)
3470 } else {
3471 self.func(name)
3472 }
3473 }
3474
3475 fn func_g(&self, name: &str) -> CudaFunction {
3479 let m = self.flash_g.get_or_init(|| {
3480 self.gpu
3481 .ctx
3482 .load_module(cudarc::nvrtc::Ptx::from_binary(
3483 FLASH_FATBIN_KF8VF8.to_vec(),
3484 ))
3485 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
3486 });
3487 let key = format!("g:{name}");
3488 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
3489 return f.clone();
3490 }
3491 let f = match m.load_function(name) {
3492 Ok(f) => f,
3493 Err(_) => self.func(name),
3494 };
3495 self.fn_cache.lock().unwrap().insert(key, f.clone());
3496 f
3497 }
3498
3499 fn func(&self, name: &str) -> CudaFunction {
3500 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
3503 return f.clone();
3504 }
3505 let f = self
3506 .module
3507 .load_function(name)
3508 .or_else(|_| self.hybrid.load_function(name))
3509 .or_else(|_| self.kda.load_function(name))
3510 .or_else(|_| self.qmatvec.load_function(name))
3511 .or_else(|_| self.flash.load_function(name))
3512 .or_else(|_| self.gemm.load_function(name))
3513 .or_else(|_| self.router.load_function(name))
3514 .or_else(|_| self.sample.load_function(name))
3515 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
3516 self.fn_cache
3517 .lock()
3518 .unwrap()
3519 .insert(name.to_string(), f.clone());
3520 f
3521 }
3522
3523 pub fn scatter_trim_logits(
3526 &self,
3527 src: &CudaSlice<f32>,
3528 d2t: &CudaSlice<u32>,
3529 dst: &mut CudaSlice<f32>,
3530 d_vocab: usize,
3531 n_vocab: usize,
3532 ) -> Result<(), Box<dyn std::error::Error>> {
3533 let f1 = self.func("scatter_trim_logits_f32");
3534 let f2 = self.func("scatter_trim_logits_pass2_f32");
3535 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
3536 let cfg1 = LaunchConfig {
3537 grid_dim: (256, 1, 1),
3538 block_dim: (256, 1, 1),
3539 shared_mem_bytes: 0,
3540 };
3541 let __s_b1 = self.gpu.stream();
3542 let mut b1 = __s_b1.launch_builder(&f1);
3543 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
3544 unsafe {
3545 b1.launch(cfg1)?;
3546 }
3547 let cfg2 = LaunchConfig {
3548 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
3549 block_dim: (256, 1, 1),
3550 shared_mem_bytes: 0,
3551 };
3552 let __s_b2 = self.gpu.stream();
3553 let mut b2 = __s_b2.launch_builder(&f2);
3554 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
3555 unsafe {
3556 b2.launch(cfg2)?;
3557 }
3558 Ok(())
3559 }
3560
3561 #[allow(clippy::too_many_arguments)]
3567 pub fn filter_stats(
3568 &self,
3569 x: &CudaSlice<f32>,
3570 row_stride: usize,
3571 rows: &CudaSlice<i32>,
3572 out_th: &mut CudaSlice<f32>,
3573 out_z: &mut CudaSlice<f32>,
3574 out_max: &mut CudaSlice<f32>,
3575 n: usize,
3576 nrow: usize,
3577 temp: f32,
3578 top_k: i32,
3579 top_p: f32,
3580 min_p: f32,
3581 ) -> Result<(), Box<dyn std::error::Error>> {
3582 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3609 let coop_on =
3610 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
3611 if coop_on && self.sm_count() >= 16 {
3612 let cap = self.sm_count() as usize / 16;
3613 let mut done = 0usize;
3614 while done < nrow {
3615 let chunk = cap.min(nrow - done);
3616 self.filter_stats_coop_chunk(
3617 x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
3618 top_p, min_p,
3619 )?;
3620 done += chunk;
3621 }
3622 return Ok(());
3623 }
3624 self.filter_stats_plain_program(
3625 x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
3626 )
3627 }
3628
3629 #[allow(clippy::too_many_arguments)]
3634 pub fn filter_stats_coop_chunk(
3635 &self,
3636 x: &CudaSlice<f32>,
3637 row_stride: usize,
3638 rows: &CudaSlice<i32>,
3639 row0: usize,
3640 out_th: &mut CudaSlice<f32>,
3641 out_z: &mut CudaSlice<f32>,
3642 out_max: &mut CudaSlice<f32>,
3643 n: usize,
3644 chunk: usize,
3645 temp: f32,
3646 top_k: i32,
3647 top_p: f32,
3648 min_p: f32,
3649 ) -> Result<(), Box<dyn std::error::Error>> {
3650 let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
3651 let f = self.func("filter_stats_coop_f32");
3652 let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
3653 let cfg = LaunchConfig {
3654 grid_dim: (16, chunk as u32, 1),
3655 block_dim: (512, 1, 1),
3656 shared_mem_bytes: 0,
3657 };
3658 let rows_v = rows.slice(row0..row0 + chunk);
3659 let mut th_v = out_th.slice_mut(row0..row0 + chunk);
3660 let mut z_v = out_z.slice_mut(row0..row0 + chunk);
3661 let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
3662 let __s_b = self.gpu.stream();
3663 let mut b = __s_b.launch_builder(&f);
3664 b.arg(x)
3665 .arg(&rs)
3666 .arg(&rows_v)
3667 .arg(&mut th_v)
3668 .arg(&mut z_v)
3669 .arg(&mut mx_v)
3670 .arg(&mut ws)
3671 .arg(&ni)
3672 .arg(&nr)
3673 .arg(&temp)
3674 .arg(&top_k)
3675 .arg(&top_p)
3676 .arg(&min_p);
3677 unsafe {
3678 b.launch_cooperative(cfg)?;
3679 }
3680 Ok(())
3681 }
3682
3683 #[allow(clippy::too_many_arguments)]
3687 pub fn filter_stats_plain_program(
3688 &self,
3689 x: &CudaSlice<f32>,
3690 row_stride: usize,
3691 rows: &CudaSlice<i32>,
3692 out_th: &mut CudaSlice<f32>,
3693 out_z: &mut CudaSlice<f32>,
3694 out_max: &mut CudaSlice<f32>,
3695 n: usize,
3696 nrow: usize,
3697 temp: f32,
3698 top_k: i32,
3699 top_p: f32,
3700 min_p: f32,
3701 ) -> Result<(), Box<dyn std::error::Error>> {
3702 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
3703 let f = self.func("filter_stats_f32");
3704 let cfg = LaunchConfig {
3705 grid_dim: (nrow as u32, 1, 1),
3706 block_dim: (1024, 1, 1),
3707 shared_mem_bytes: 0,
3708 };
3709 let __s_b = self.gpu.stream();
3710 let mut b = __s_b.launch_builder(&f);
3711 b.arg(x)
3712 .arg(&rs)
3713 .arg(rows)
3714 .arg(&mut *out_th)
3715 .arg(&mut *out_z)
3716 .arg(&mut *out_max)
3717 .arg(&ni)
3718 .arg(&nr)
3719 .arg(&temp)
3720 .arg(&top_k)
3721 .arg(&top_p)
3722 .arg(&min_p);
3723 unsafe {
3724 b.launch(cfg)?;
3725 }
3726 Ok(())
3727 }
3728
3729 #[allow(clippy::too_many_arguments)]
3731 pub fn softmax_gather_filtered(
3732 &self,
3733 x: &CudaSlice<f32>,
3734 row_stride: usize,
3735 ids: &CudaSlice<u32>,
3736 rows: &CudaSlice<i32>,
3737 th: &CudaSlice<f32>,
3738 z: &CudaSlice<f32>,
3739 out: &mut CudaSlice<f32>,
3740 n: usize,
3741 npair: usize,
3742 temp: f32,
3743 ) -> Result<(), Box<dyn std::error::Error>> {
3744 let f = self.func("softmax_gather_filtered_f32");
3745 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
3746 let cfg = LaunchConfig {
3747 grid_dim: (npair as u32, 1, 1),
3748 block_dim: (256, 1, 1),
3749 shared_mem_bytes: 0,
3750 };
3751 let __s_b = self.gpu.stream();
3752 let mut b = __s_b.launch_builder(&f);
3753 b.arg(x)
3754 .arg(&rs)
3755 .arg(ids)
3756 .arg(rows)
3757 .arg(th)
3758 .arg(z)
3759 .arg(&mut *out)
3760 .arg(&ni)
3761 .arg(&np)
3762 .arg(&temp);
3763 unsafe {
3764 b.launch(cfg)?;
3765 }
3766 Ok(())
3767 }
3768
3769 #[allow(clippy::too_many_arguments)]
3771 pub fn residual_sample_filtered(
3772 &self,
3773 p: &CudaSlice<f32>,
3774 q: Option<&CudaSlice<f32>>,
3775 n: usize,
3776 temp: f32,
3777 seed: u64,
3778 stream_pos: u32,
3779 p_stats: (f32, f32, f32),
3780 q_stats: (f32, f32, f32),
3781 out_tok: &mut CudaSlice<u32>,
3782 ) -> Result<(), Box<dyn std::error::Error>> {
3783 let f = self.func("residual_sample_filtered_f32");
3784 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3785 let has_q: i32 = q.is_some() as i32;
3786 let qbuf = q.unwrap_or(p);
3787 let (pm, pth, pz) = p_stats;
3788 let (qm, qth, qz) = q_stats;
3789 let cfg = LaunchConfig {
3790 grid_dim: (1, 1, 1),
3791 block_dim: (1024, 1, 1),
3792 shared_mem_bytes: 0,
3793 };
3794 let __s_b = self.gpu.stream();
3795 let mut b = __s_b.launch_builder(&f);
3796 b.arg(p)
3797 .arg(qbuf)
3798 .arg(&has_q)
3799 .arg(&ni)
3800 .arg(&temp)
3801 .arg(&slo)
3802 .arg(&shi)
3803 .arg(&stream_pos)
3804 .arg(&pm)
3805 .arg(&pth)
3806 .arg(&pz)
3807 .arg(&qm)
3808 .arg(&qth)
3809 .arg(&qz)
3810 .arg(&mut *out_tok);
3811 unsafe {
3812 b.launch(cfg)?;
3813 }
3814 Ok(())
3815 }
3816
3817 #[allow(clippy::too_many_arguments)]
3823 pub fn residual_sample_sparse_q(
3824 &self,
3825 p: &CudaSlice<f32>,
3826 cand_ids: &CudaSlice<u32>,
3827 q_probs: &CudaSlice<f32>,
3828 n_cand: usize,
3829 n: usize,
3830 temp: f32,
3831 seed: u64,
3832 stream_pos: u32,
3833 p_stats: (f32, f32, f32),
3834 out_tok: &mut CudaSlice<u32>,
3835 ) -> Result<(), Box<dyn std::error::Error>> {
3836 assert!(
3837 (1..=32).contains(&n_cand),
3838 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
3839 );
3840 let f = self.func("residual_sample_sparse_q_f32");
3841 let (ni, nc) = (n as i32, n_cand as i32);
3842 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3843 let (pm, pth, pz) = p_stats;
3844 let cfg = LaunchConfig {
3845 grid_dim: (1, 1, 1),
3846 block_dim: (1024, 1, 1),
3847 shared_mem_bytes: 0,
3848 };
3849 let __s_b = self.gpu.stream();
3850 let mut b = __s_b.launch_builder(&f);
3851 b.arg(p)
3852 .arg(cand_ids)
3853 .arg(q_probs)
3854 .arg(&nc)
3855 .arg(&ni)
3856 .arg(&temp)
3857 .arg(&slo)
3858 .arg(&shi)
3859 .arg(&stream_pos)
3860 .arg(&pm)
3861 .arg(&pth)
3862 .arg(&pz)
3863 .arg(&mut *out_tok);
3864 unsafe {
3865 b.launch(cfg)?;
3866 }
3867 Ok(())
3868 }
3869
3870 #[allow(clippy::too_many_arguments)]
3872 pub fn gumbel_perturb_filtered(
3873 &self,
3874 x: &CudaSlice<f32>,
3875 y: &mut CudaSlice<f32>,
3876 n: usize,
3877 seed: u64,
3878 stream_pos: u32,
3879 temp: f32,
3880 row_max: f32,
3881 th: f32,
3882 ) -> Result<(), Box<dyn std::error::Error>> {
3883 let f = self.func("gumbel_perturb_filtered_f32");
3884 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3885 let cfg = LaunchConfig {
3886 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3887 block_dim: (256, 1, 1),
3888 shared_mem_bytes: 0,
3889 };
3890 let __s_b = self.gpu.stream();
3891 let mut b = __s_b.launch_builder(&f);
3892 b.arg(x)
3893 .arg(&mut *y)
3894 .arg(&ni)
3895 .arg(&slo)
3896 .arg(&shi)
3897 .arg(&stream_pos)
3898 .arg(&temp)
3899 .arg(&row_max)
3900 .arg(&th);
3901 unsafe {
3902 b.launch(cfg)?;
3903 }
3904 Ok(())
3905 }
3906
3907 #[allow(clippy::too_many_arguments)]
3911 pub fn penalize_logits(
3912 &self,
3913 x: &mut CudaSlice<f32>,
3914 hist: &CudaSlice<u32>,
3915 n_hist: usize,
3916 rep: f32,
3917 freq: f32,
3918 present: f32,
3919 n: usize,
3920 ) -> Result<(), Box<dyn std::error::Error>> {
3921 if n_hist == 0 {
3922 return Ok(());
3923 }
3924 let f = self.func("penalize_logits_f32");
3925 let (nh, ni) = (n_hist as i32, n as i32);
3926 let cfg = LaunchConfig {
3927 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
3928 block_dim: (128, 1, 1),
3929 shared_mem_bytes: 0,
3930 };
3931 let __s_b = self.gpu.stream();
3932 let mut b = __s_b.launch_builder(&f);
3933 b.arg(&mut *x)
3934 .arg(hist)
3935 .arg(&nh)
3936 .arg(&rep)
3937 .arg(&freq)
3938 .arg(&present)
3939 .arg(&ni);
3940 unsafe {
3941 b.launch(cfg)?;
3942 }
3943 Ok(())
3944 }
3945
3946 #[allow(clippy::too_many_arguments)]
3948 pub fn penalize_logits_rows(
3949 &self,
3950 x: &mut CudaSlice<f32>,
3951 hist: &CudaSlice<u32>,
3952 n_hist: usize,
3953 rep: f32,
3954 freq: f32,
3955 present: f32,
3956 n: usize,
3957 nrow: usize,
3958 ) -> Result<(), Box<dyn std::error::Error>> {
3959 if n_hist == 0 || nrow == 0 {
3960 return Ok(());
3961 }
3962 let f = self.func("penalize_logits_rows_f32");
3963 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
3964 let cfg = LaunchConfig {
3965 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
3966 block_dim: (128, 1, 1),
3967 shared_mem_bytes: 0,
3968 };
3969 let __s_b = self.gpu.stream();
3970 let mut b = __s_b.launch_builder(&f);
3971 b.arg(&mut *x)
3972 .arg(hist)
3973 .arg(&nh)
3974 .arg(&rep)
3975 .arg(&freq)
3976 .arg(&present)
3977 .arg(&ni)
3978 .arg(&nr);
3979 unsafe {
3980 b.launch(cfg)?;
3981 }
3982 Ok(())
3983 }
3984
3985 #[allow(clippy::too_many_arguments)]
3991 pub fn penalize_logits_sparse_rows(
3992 &self,
3993 x: &mut CudaSlice<f32>,
3994 ids: &[u32],
3995 counts: &[u32],
3996 offsets: &[i32],
3997 rows: &[i32],
3998 reps: &[f32],
3999 freqs: &[f32],
4000 presents: &[f32],
4001 n: usize,
4002 ) -> Result<(), Box<dyn std::error::Error>> {
4003 let nrow = rows.len();
4004 if nrow == 0 {
4005 return Ok(());
4006 }
4007 let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
4008 let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
4009 let entry_count =
4010 i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
4011 if ids.len() != counts.len()
4012 || offsets.len() != nrow + 1
4013 || reps.len() != nrow
4014 || freqs.len() != nrow
4015 || presents.len() != nrow
4016 || offsets.first().copied() != Some(0)
4017 || offsets.last().copied() != Some(entry_count)
4018 {
4019 return Err("sparse penalty row metadata shape mismatch".into());
4020 }
4021 if counts.contains(&0) {
4022 return Err("sparse penalty counts must be positive".into());
4023 }
4024 let mut max_len = 0usize;
4025 for pair in offsets.windows(2) {
4026 if pair[0] < 0 || pair[1] < pair[0] {
4027 return Err("sparse penalty offsets must be monotonic".into());
4028 }
4029 max_len = max_len.max((pair[1] - pair[0]) as usize);
4030 }
4031 if max_len == 0 {
4032 return Ok(());
4033 }
4034
4035 let mut seen = std::collections::HashSet::with_capacity(ids.len());
4036 for (r, &row) in rows.iter().enumerate() {
4037 if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
4038 return Err("sparse penalty row index exceeds logits shape".into());
4039 }
4040 let begin = offsets[r] as usize;
4041 let end = offsets[r + 1] as usize;
4042 for &id in &ids[begin..end] {
4043 if id as usize >= n {
4044 return Err("sparse penalty token id exceeds logits row".into());
4045 }
4046 if !seen.insert((row, id)) {
4047 return Err("sparse penalty entries must be unique per logits row".into());
4048 }
4049 }
4050 }
4051
4052 unsafe {
4054 self.penalize_logits_sparse_rows_unchecked(
4055 x, ids, counts, offsets, rows, reps, freqs, presents, n,
4056 )
4057 }
4058 }
4059
4060 #[allow(clippy::too_many_arguments)]
4069 pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
4070 &self,
4071 x: &mut CudaSlice<f32>,
4072 ids: &[u32],
4073 counts: &[u32],
4074 offsets: &[i32],
4075 rows: &[i32],
4076 reps: &[f32],
4077 freqs: &[f32],
4078 presents: &[f32],
4079 n: usize,
4080 ) -> Result<(), Box<dyn std::error::Error>> {
4081 let nrow = rows.len();
4082 if nrow == 0 {
4083 return Ok(());
4084 }
4085 let max_len = offsets
4086 .windows(2)
4087 .map(|pair| (pair[1] - pair[0]) as usize)
4088 .max()
4089 .unwrap_or(0);
4090 if max_len == 0 {
4091 return Ok(());
4092 }
4093 let ids_d = self.htod_u32_v(ids)?;
4094 let counts_d = self.htod_u32_v(counts)?;
4095 let offsets_d = self.htod_i32(offsets)?;
4096 let rows_d = self.htod_i32(rows)?;
4097 let reps_d = self.htod(reps)?;
4098 let freqs_d = self.htod(freqs)?;
4099 let presents_d = self.htod(presents)?;
4100 let f = self.func("penalize_logits_sparse_rows_f32");
4101 let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
4102 let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
4103 let cfg = LaunchConfig {
4104 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
4105 block_dim: (128, 1, 1),
4106 shared_mem_bytes: 0,
4107 };
4108 let __s_b = self.gpu.stream();
4109 let mut b = __s_b.launch_builder(&f);
4110 b.arg(&mut *x)
4111 .arg(&ids_d)
4112 .arg(&counts_d)
4113 .arg(&offsets_d)
4114 .arg(&rows_d)
4115 .arg(&reps_d)
4116 .arg(&freqs_d)
4117 .arg(&presents_d)
4118 .arg(&ni)
4119 .arg(&nr);
4120 unsafe {
4121 b.launch(cfg)?;
4122 }
4123 Ok(())
4124 }
4125
4126 #[allow(clippy::too_many_arguments)]
4134 pub fn penalize_logits_rows_inc(
4135 &self,
4136 x: &mut CudaSlice<f32>,
4137 hist: &CudaSlice<u32>,
4138 n_hist0: usize,
4139 rep: f32,
4140 freq: f32,
4141 present: f32,
4142 n: usize,
4143 nrow: usize,
4144 win: usize,
4145 ) -> Result<(), Box<dyn std::error::Error>> {
4146 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
4147 return Ok(());
4148 }
4149 debug_assert!(
4150 hist.len() >= n_hist0 + nrow - 1,
4151 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
4152 );
4153 let f = self.func("penalize_logits_rows_inc_f32");
4154 let max_len = win.min(n_hist0 + nrow - 1).max(1);
4155 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
4156 let cfg = LaunchConfig {
4157 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
4158 block_dim: (128, 1, 1),
4159 shared_mem_bytes: 0,
4160 };
4161 let __s_b = self.gpu.stream();
4162 let mut b = __s_b.launch_builder(&f);
4163 b.arg(&mut *x)
4164 .arg(hist)
4165 .arg(&nh)
4166 .arg(&rep)
4167 .arg(&freq)
4168 .arg(&present)
4169 .arg(&ni)
4170 .arg(&nr)
4171 .arg(&wi);
4172 unsafe {
4173 b.launch(cfg)?;
4174 }
4175 Ok(())
4176 }
4177
4178 pub fn wpf_level() -> u32 {
4186 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
4187 *ON.get_or_init(|| {
4188 std::env::var("MEMRA_WPF")
4189 .ok()
4190 .and_then(|v| v.parse().ok())
4191 .unwrap_or(1)
4192 })
4193 }
4194
4195 pub fn set_verify_exact(&self, on: bool) {
4211 self.verify_exact
4212 .store(on, std::sync::atomic::Ordering::Relaxed);
4213 }
4214 pub(crate) fn verify_exact_on(&self) -> bool {
4215 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
4216 }
4217
4218 pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
4224 ExactScope::set(&self.verify_exact, on)
4225 }
4226
4227 pub fn qkv_append_on() -> bool {
4230 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4231 *ON.get_or_init(|| {
4232 std::env::var("MEMRA_QKV_APPEND")
4233 .map(|v| v != "0")
4234 .unwrap_or(true)
4235 })
4236 }
4237
4238 pub fn pdl_wb_on() -> bool {
4241 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4242 *ON.get_or_init(|| {
4243 std::env::var("MEMRA_PDL_WB")
4244 .map(|v| v != "0")
4245 .unwrap_or(true)
4246 })
4247 }
4248
4249 pub fn norm_ilp_on() -> bool {
4257 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4258 *ON.get_or_init(|| {
4259 std::env::var("MEMRA_NORM_ILP")
4260 .map(|v| v != "0")
4261 .unwrap_or(true)
4262 })
4263 }
4264
4265 pub fn norm_ilp_zq8_on() -> bool {
4270 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4271 *ON.get_or_init(|| {
4272 Self::norm_ilp_on()
4273 && std::env::var("MEMRA_NORM_ILP_ZQ8")
4274 .map(|v| v != "0")
4275 .unwrap_or(true)
4276 })
4277 }
4278
4279 pub fn tk_ffn_dual_on() -> bool {
4287 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4288 *ON.get_or_init(|| {
4289 std::env::var("MEMRA_TK_FFN_DUAL")
4290 .map(|v| v != "0")
4291 .unwrap_or(true)
4292 })
4293 }
4294
4295 pub fn pdl_mmvq_on() -> bool {
4299 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4300 *ON.get_or_init(|| {
4301 std::env::var("MEMRA_PDL_MMVQ")
4302 .map(|v| v != "0")
4303 .unwrap_or(true)
4304 })
4305 }
4306
4307 pub fn pdl_on() -> bool {
4308 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4309 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
4310 }
4311
4312 pub fn pdl_nvfp4q8_on() -> bool {
4318 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4319 *ON.get_or_init(|| {
4320 std::env::var("MEMRA_PDL_NVFP4")
4321 .map(|v| v != "0")
4322 .unwrap_or(true)
4323 })
4324 }
4325
4326 fn q40_mr1_on() -> bool {
4332 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
4333 match *Q40MR.get_or_init(|| {
4334 std::env::var("MEMRA_Q40_MR")
4335 .ok()
4336 .and_then(|v| v.parse().ok())
4337 }) {
4338 Some(v) => v == 1,
4339 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4340 }
4341 }
4342
4343 fn pdl_func_flash(
4348 &self,
4349 g: bool,
4350 name: &'static str,
4351 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4352 use cudarc::driver::sys as cu;
4353 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
4360 std::sync::Mutex::new(None);
4361 #[allow(clippy::type_complexity)] static FNS: std::sync::Mutex<
4363 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
4364 > = std::sync::Mutex::new(None);
4365 let ctx_key = self.ctx().cu_ctx() as usize;
4366 if let Some(&f) = FNS
4367 .lock()
4368 .unwrap()
4369 .get_or_insert_with(Default::default)
4370 .get(&(ctx_key, g, name))
4371 {
4372 return Ok(f as cu::CUfunction);
4373 }
4374 let module = {
4375 let mut mods = MODS.lock().unwrap();
4376 let map = mods.get_or_insert_with(Default::default);
4377 match map.get(&(ctx_key, g)) {
4378 Some(&m) => m,
4379 None => {
4380 let m = self.pdl_load_module_in_ctx(if g {
4381 FLASH_FATBIN_KF8VF8
4382 } else {
4383 FLASH_FATBIN
4384 })?;
4385 map.insert((ctx_key, g), m);
4386 m
4387 }
4388 }
4389 };
4390 let cname = std::ffi::CString::new(name)?;
4391 let mut f: cu::CUfunction = std::ptr::null_mut();
4392 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
4393 if r != cu::CUresult::CUDA_SUCCESS {
4394 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
4395 }
4396 FNS.lock()
4397 .unwrap()
4398 .get_or_insert_with(Default::default)
4399 .insert((ctx_key, g, name), f as usize);
4400 Ok(f)
4401 }
4402
4403 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
4408 use cudarc::driver::sys as cu;
4409 let mut prev: cu::CUcontext = std::ptr::null_mut();
4410 unsafe {
4411 cu::cuCtxGetCurrent(&mut prev).result()?;
4412 }
4413 self.ctx().bind_to_thread()?;
4414 let mut m: cu::CUmodule = std::ptr::null_mut();
4415 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
4416 let restore = if prev.is_null() {
4417 cu::CUresult::CUDA_SUCCESS
4418 } else {
4419 unsafe { cu::cuCtxSetCurrent(prev) }
4420 };
4421 if r != cu::CUresult::CUDA_SUCCESS {
4422 return Err(format!("pdl module load: {r:?}").into());
4423 }
4424 if restore != cu::CUresult::CUDA_SUCCESS {
4425 return Err(format!("pdl module load: ctx restore {restore:?}").into());
4426 }
4427 Ok(m as usize)
4428 }
4429
4430 pub fn raw_kernel_function(
4433 &self,
4434 name: &'static str,
4435 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4436 self.pdl_func(name)
4437 }
4438
4439 fn pdl_func(
4440 &self,
4441 name: &'static str,
4442 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
4443 use cudarc::driver::sys as cu;
4444 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
4447 std::sync::Mutex::new(None);
4448 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
4451 std::sync::Mutex::new(None);
4452 static FNS: std::sync::Mutex<
4453 Option<std::collections::HashMap<(usize, &'static str), usize>>,
4454 > = std::sync::Mutex::new(None);
4455 let ctx_key = self.ctx().cu_ctx() as usize;
4456 if let Some(&f) = FNS
4457 .lock()
4458 .unwrap()
4459 .get_or_insert_with(Default::default)
4460 .get(&(ctx_key, name))
4461 {
4462 return Ok(f as cu::CUfunction);
4463 }
4464 let module = {
4465 let mut mods = MODULES.lock().unwrap();
4466 let map = mods.get_or_insert_with(Default::default);
4467 match map.get(&ctx_key) {
4468 Some(&m) => m,
4469 None => {
4470 let m = self.pdl_load_module_in_ctx(FATBIN)?;
4471 map.insert(ctx_key, m);
4472 m
4473 }
4474 }
4475 };
4476 let cname = std::ffi::CString::new(name)?;
4477 let mut f: cu::CUfunction = std::ptr::null_mut();
4478 let mut r =
4479 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
4480 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
4481 let qmodule = {
4482 let mut mods = QMODULES.lock().unwrap();
4483 let map = mods.get_or_insert_with(Default::default);
4484 match map.get(&ctx_key) {
4485 Some(&m) => m,
4486 None => {
4487 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
4488 map.insert(ctx_key, m);
4489 m
4490 }
4491 }
4492 };
4493 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
4494 }
4495 if r != cu::CUresult::CUDA_SUCCESS {
4496 return Err(format!("pdl_func {name}: {r:?}").into());
4497 }
4498 FNS.lock()
4499 .unwrap()
4500 .get_or_insert_with(Default::default)
4501 .insert((ctx_key, name), f as usize);
4502 Ok(f)
4503 }
4504
4505 unsafe fn launch_pdl_flash(
4517 &self,
4518 g: bool,
4519 name: &'static str,
4520 grid: (u32, u32, u32),
4521 block: (u32, u32, u32),
4522 smem: u32,
4523 params: &mut [*mut std::ffi::c_void],
4524 ) -> Result<(), Box<dyn std::error::Error>> {
4525 use cudarc::driver::sys as cu;
4526 let f = self.pdl_func_flash(g, name)?;
4527 if smem > 0 {
4528 let r =
4530 unsafe {
4531 cu::cuFuncSetAttribute(f,
4532 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
4533 smem as i32)
4534 };
4535 if r != cu::CUresult::CUDA_SUCCESS {
4536 return Err(format!("pdl smem attr {name}: {r:?}").into());
4537 }
4538 }
4539 let mut attr = cu::CUlaunchAttribute {
4540 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
4541 pad: [0; 4],
4542 value: cu::CUlaunchAttributeValue {
4543 programmaticStreamSerializationAllowed: 1,
4544 },
4545 };
4546 let cfg = cu::CUlaunchConfig {
4547 gridDimX: grid.0,
4548 gridDimY: grid.1,
4549 gridDimZ: grid.2,
4550 blockDimX: block.0,
4551 blockDimY: block.1,
4552 blockDimZ: block.2,
4553 sharedMemBytes: smem,
4554 hStream: self.gpu.stream().cu_stream(),
4555 attrs: &mut attr,
4556 numAttrs: 1,
4557 };
4558 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
4559 if r != cu::CUresult::CUDA_SUCCESS {
4560 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
4561 }
4562 Ok(())
4563 }
4564
4565 unsafe fn launch_pdl(
4566 &self,
4567 name: &'static str,
4568 grid: (u32, u32, u32),
4569 block: (u32, u32, u32),
4570 params: &mut [*mut std::ffi::c_void],
4571 ) -> Result<(), Box<dyn std::error::Error>> {
4572 use cudarc::driver::sys as cu;
4573 let f = self.pdl_func(name)?;
4574 let mut attr = cu::CUlaunchAttribute {
4575 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
4576 pad: [0; 4],
4577 value: cu::CUlaunchAttributeValue {
4578 programmaticStreamSerializationAllowed: 1,
4579 },
4580 };
4581 let cfg = cu::CUlaunchConfig {
4582 gridDimX: grid.0,
4583 gridDimY: grid.1,
4584 gridDimZ: grid.2,
4585 blockDimX: block.0,
4586 blockDimY: block.1,
4587 blockDimZ: block.2,
4588 sharedMemBytes: 0,
4589 hStream: self.gpu.stream().cu_stream(),
4590 attrs: &mut attr,
4591 numAttrs: 1,
4592 };
4593 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
4594 if r != cu::CUresult::CUDA_SUCCESS {
4595 return Err(format!("launch_pdl {name}: {r:?}").into());
4596 }
4597 Ok(())
4598 }
4599
4600 pub fn prefetch_weight_l2(
4603 &self,
4604 w: &crate::model::GpuTensor,
4605 ) -> Result<(), Box<dyn std::error::Error>> {
4606 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
4607 let p = rp4.as_ref().unwrap_or(bytes);
4608 self.prefetch_l2(p, p.len())?;
4609 }
4610 Ok(())
4611 }
4612
4613 pub fn gather_row_bf16(
4616 &self,
4617 table: &CudaSlice<u8>,
4618 tok: &CudaSlice<u32>,
4619 idx: usize,
4620 dst: &mut CudaSlice<f32>,
4621 ncols: usize,
4622 ) -> Result<(), Box<dyn std::error::Error>> {
4623 let f = self.func("gather_row_bf16_f32");
4624 let cfg = LaunchConfig {
4625 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
4626 block_dim: (256, 1, 1),
4627 shared_mem_bytes: 0,
4628 };
4629 let (nc, ix) = (ncols as i32, idx as i32);
4630 let __s_b = self.gpu.stream();
4631 let mut b = __s_b.launch_builder(&f);
4632 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
4633 unsafe {
4634 b.launch(cfg)?;
4635 }
4636 Ok(())
4637 }
4638
4639 #[allow(clippy::too_many_arguments)]
4645 pub fn dflash2_dynconv(
4646 &self,
4647 x: &CudaSlice<f32>,
4648 dyn_: &CudaSlice<f32>,
4649 base: &CudaSlice<f32>,
4650 out: &mut CudaSlice<f32>,
4651 rows: usize,
4652 hidden: usize,
4653 group_size: usize,
4654 ksize: usize,
4655 half: usize,
4656 ) -> Result<(), Box<dyn std::error::Error>> {
4657 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
4658 let f = self.func("dflash2_dynconv_f32");
4659 let n = rows * hidden;
4660 let cfg = LaunchConfig {
4661 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4662 block_dim: (256, 1, 1),
4663 shared_mem_bytes: 0,
4664 };
4665 let (ri, hi, gi, ki, hf) = (
4666 rows as i32,
4667 hidden as i32,
4668 group_size as i32,
4669 ksize as i32,
4670 half as i32,
4671 );
4672 let __s_b = self.gpu.stream();
4673 let mut b = __s_b.launch_builder(&f);
4674 b.arg(x)
4675 .arg(dyn_)
4676 .arg(base)
4677 .arg(out)
4678 .arg(&ri)
4679 .arg(&hi)
4680 .arg(&gi)
4681 .arg(&ki)
4682 .arg(&hf);
4683 unsafe {
4684 b.launch(cfg)?;
4685 }
4686 Ok(())
4687 }
4688
4689 pub fn topk_rows(
4693 &self,
4694 logits: &CudaSlice<f32>,
4695 n_rows: usize,
4696 n_cols: usize,
4697 k: usize,
4698 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
4699 assert!((1..=32).contains(&k), "topk_rows supports 1..=32, got {k}");
4700 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
4701 if topk_shards_on() && n_cols >= 16 * 1024 && k <= n_cols / 16 {
4708 if TOPK_SHARDS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
4709 eprintln!(
4710 "[topk-shards] engaged: rows={n_rows} cols={n_cols} k={k} shards=16 \
4711 (MEMRA_TOPK_SHARDS=1)"
4712 );
4713 }
4714 return self.topk_rows_sharded(logits, n_rows, n_cols, k, 16);
4715 }
4716 let f = self.func("topk_rows_f32");
4717 let nth = 256usize;
4718 let mut vals = self.uninit(n_rows * k)?;
4719 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
4720 let cfg = LaunchConfig {
4721 grid_dim: (n_rows as u32, 1, 1),
4722 block_dim: (nth as u32, 1, 1),
4723 shared_mem_bytes: (nth * k * 8) as u32,
4724 };
4725 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
4726 let __s_b = self.gpu.stream();
4727 let mut b = __s_b.launch_builder(&f);
4728 b.arg(logits)
4729 .arg(&nr)
4730 .arg(&nc)
4731 .arg(&ki)
4732 .arg(&mut vals)
4733 .arg(&mut idxs);
4734 unsafe {
4735 b.launch(cfg)?;
4736 }
4737 Ok((vals, idxs))
4738 }
4739
4740 fn topk_rows_sharded(
4746 &self,
4747 logits: &CudaSlice<f32>,
4748 n_rows: usize,
4749 n_cols: usize,
4750 k: usize,
4751 n_shards: usize,
4752 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
4753 assert!((1..=64).contains(&n_shards), "shard merge head cap is 64");
4754 let nth = 256usize;
4755 let mut pvals = self.uninit(n_rows * n_shards * k)?;
4756 let mut pidxs = self.alloc_uninit::<u32>(n_rows * n_shards * k)?;
4757 let f1 = self.func("topk_rows_shard_f32");
4758 let cfg1 = LaunchConfig {
4759 grid_dim: (n_rows as u32, n_shards as u32, 1),
4760 block_dim: (nth as u32, 1, 1),
4761 shared_mem_bytes: (nth * k * 8) as u32,
4762 };
4763 let (nr, nc, ki, ns) = (n_rows as i32, n_cols as i32, k as i32, n_shards as i32);
4764 {
4765 let __s_b = self.gpu.stream();
4766 let mut b = __s_b.launch_builder(&f1);
4767 b.arg(logits)
4768 .arg(&nr)
4769 .arg(&nc)
4770 .arg(&ki)
4771 .arg(&ns)
4772 .arg(&mut pvals)
4773 .arg(&mut pidxs);
4774 unsafe {
4775 b.launch(cfg1)?;
4776 }
4777 }
4778 let mut vals = self.uninit(n_rows * k)?;
4779 let mut idxs = self.alloc_uninit::<u32>(n_rows * k)?;
4780 let f2 = self.func("topk_rows_shard_merge_f32");
4781 let cfg2 = LaunchConfig {
4782 grid_dim: (n_rows as u32, 1, 1),
4783 block_dim: (32, 1, 1),
4784 shared_mem_bytes: 0,
4785 };
4786 let __s_b = self.gpu.stream();
4787 let mut b = __s_b.launch_builder(&f2);
4788 b.arg(&pvals)
4789 .arg(&pidxs)
4790 .arg(&nr)
4791 .arg(&ns)
4792 .arg(&ki)
4793 .arg(&mut vals)
4794 .arg(&mut idxs);
4795 unsafe {
4796 b.launch(cfg2)?;
4797 }
4798 Ok((vals, idxs))
4799 }
4800
4801 pub fn add_row_inplace(
4803 &self,
4804 logits: &mut CudaSlice<f32>,
4805 bias: &CudaSlice<f32>,
4806 n: usize,
4807 row_off: usize,
4808 ) -> Result<(), Box<dyn std::error::Error>> {
4809 let f = self.func("add_row_inplace_f32");
4810 let cfg = LaunchConfig {
4811 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4812 block_dim: (256, 1, 1),
4813 shared_mem_bytes: 0,
4814 };
4815 let (ni, off) = (n as i32, row_off as i64);
4816 let __s_b = self.gpu.stream();
4817 let mut b = __s_b.launch_builder(&f);
4818 b.arg(logits).arg(bias).arg(&ni).arg(&off);
4819 unsafe {
4820 b.launch(cfg)?;
4821 }
4822 Ok(())
4823 }
4824
4825 pub fn prefetch_l2(
4827 &self,
4828 p: &CudaSlice<u8>,
4829 n: usize,
4830 ) -> Result<(), Box<dyn std::error::Error>> {
4831 let f = self.func("prefetch_l2_bytes");
4832 let lines = n.div_ceil(128);
4833 let ni = n as i64;
4834 let cfg = LaunchConfig {
4835 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
4836 block_dim: (256, 1, 1),
4837 shared_mem_bytes: 0,
4838 };
4839 let __s_b = self.gpu.stream();
4840 let mut b = __s_b.launch_builder(&f);
4841 b.arg(p).arg(&ni);
4842 unsafe {
4843 b.launch(cfg)?;
4844 }
4845 Ok(())
4846 }
4847
4848 pub fn router_gemv(
4851 &self,
4852 w: &CudaSlice<f32>,
4853 x: &CudaSlice<f32>,
4854 n_embd: usize,
4855 n_experts: usize,
4856 t: usize,
4857 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4858 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
4864 Ok("0") => false,
4865 Ok(_) => true,
4866 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4867 };
4868 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
4878 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
4879 }
4880
4881 #[allow(clippy::too_many_arguments)] pub fn router_gemv_form(
4885 &self,
4886 w: &CudaSlice<f32>,
4887 x: &CudaSlice<f32>,
4888 n_embd: usize,
4889 n_experts: usize,
4890 t: usize,
4891 w8: bool,
4892 batch: bool,
4893 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4894 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
4895 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
4896 let f = if batch {
4897 self.func("router_gemv_f32_w8_batch")
4898 } else if w8 {
4899 self.func("router_gemv_f32_w8")
4900 } else {
4901 self.func("router_gemv_f32")
4902 };
4903 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
4904 let cfg = if batch {
4905 LaunchConfig {
4906 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
4907 block_dim: (32, 8, 1),
4908 shared_mem_bytes: 0,
4909 }
4910 } else {
4911 LaunchConfig {
4912 grid_dim: (n_experts as u32, t as u32, 1),
4913 block_dim: (32, if w8 { 8 } else { 1 }, 1),
4914 shared_mem_bytes: 0,
4915 }
4916 };
4917 let __s_b = self.gpu.stream();
4918 let mut b = __s_b.launch_builder(&f);
4919 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
4920 unsafe {
4921 b.launch(cfg)?;
4922 }
4923 Ok(y)
4924 }
4925
4926 pub fn router_gemv_into(
4929 &self,
4930 w: &CudaSlice<f32>,
4931 x: &CudaSlice<f32>,
4932 y: &mut CudaSlice<f32>,
4933 n_embd: usize,
4934 n_experts: usize,
4935 t: usize,
4936 ) -> Result<(), Box<dyn std::error::Error>> {
4937 if y.len() < t * n_experts {
4938 return Err("router_gemv_into output too small".into());
4939 }
4940 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
4941 Ok("0") => false,
4942 Ok(_) => true,
4943 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
4944 };
4945 let f = if w8 {
4946 self.func("router_gemv_f32_w8")
4947 } else {
4948 self.func("router_gemv_f32")
4949 };
4950 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
4951 let cfg = LaunchConfig {
4952 grid_dim: (n_experts as u32, t as u32, 1),
4953 block_dim: (32, if w8 { 8 } else { 1 }, 1),
4954 shared_mem_bytes: 0,
4955 };
4956 let __s_b = self.gpu.stream();
4957 let mut b = __s_b.launch_builder(&f);
4958 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
4959 unsafe {
4960 b.launch(cfg)?;
4961 }
4962 Ok(())
4963 }
4964
4965 pub fn rows_permute(
4967 &self,
4968 src: &CudaSlice<f32>,
4969 idx: &CudaSlice<i32>,
4970 nrows: usize,
4971 ncols: usize,
4972 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4973 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
4974 let f = self.func("rows_permute_f32");
4975 let (nc, nr) = (ncols as i32, nrows as i32);
4976 let cfg = LaunchConfig {
4977 grid_dim: (nrows as u32, 1, 1),
4978 block_dim: (256, 1, 1),
4979 shared_mem_bytes: 0,
4980 };
4981 let __s_b = self.gpu.stream();
4982 let mut b = __s_b.launch_builder(&f);
4983 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
4984 unsafe {
4985 b.launch(cfg)?;
4986 }
4987 Ok(dst)
4988 }
4989
4990 pub fn sigmoid_dot_rows(
4995 &self,
4996 x: &CudaSlice<f32>,
4997 w: &CudaSlice<f32>,
4998 n_embd: usize,
4999 t: usize,
5000 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5001 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5004 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
5005 let gs = self.linear(x, w, t, n_embd, 1)?;
5006 let mut g = self.uninit(t)?;
5007 self.sigmoid(&gs, &mut g, t)?;
5008 return Ok(g);
5009 }
5010 let mut g = self.alloc_uninit::<f32>(t)?;
5016 let f = self.func("sigmoid_dot_rows_f32");
5017 let (ne, ti) = (n_embd as i32, t as i32);
5018 let cfg = LaunchConfig {
5019 grid_dim: (t as u32, 1, 1),
5020 block_dim: (32, 8, 1),
5021 shared_mem_bytes: 0,
5022 };
5023 let __s_b = self.gpu.stream();
5024 let mut b = __s_b.launch_builder(&f);
5025 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
5026 unsafe {
5027 b.launch(cfg)?;
5028 }
5029 Ok(g)
5030 }
5031
5032 pub fn sigmoid_dot_rows_into(
5034 &self,
5035 x: &CudaSlice<f32>,
5036 w: &CudaSlice<f32>,
5037 g: &mut CudaSlice<f32>,
5038 n_embd: usize,
5039 t: usize,
5040 ) -> Result<(), Box<dyn std::error::Error>> {
5041 if g.len() < t {
5042 return Err("sigmoid_dot_rows_into output too small".into());
5043 }
5044 let f = self.func("sigmoid_dot_rows_f32");
5045 let (ne, ti) = (n_embd as i32, t as i32);
5046 let cfg = LaunchConfig {
5047 grid_dim: (t as u32, 1, 1),
5048 block_dim: (32, 8, 1),
5049 shared_mem_bytes: 0,
5050 };
5051 let __s_b = self.gpu.stream();
5052 let mut b = __s_b.launch_builder(&f);
5053 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
5054 unsafe {
5055 b.launch(cfg)?;
5056 }
5057 Ok(())
5058 }
5059
5060 pub fn spec_rollback_stream(
5062 &self,
5063 len_ptrs: &CudaSlice<u64>,
5064 pos_start: &CudaSlice<i32>,
5065 acc: &CudaSlice<u32>,
5066 base: usize,
5067 n_rows: usize,
5068 ) -> Result<(), Box<dyn std::error::Error>> {
5069 let f = self.func("spec_rollback_stream");
5070 let (b, nr) = (base as i32, n_rows as i32);
5071 let cfg = LaunchConfig {
5072 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
5073 block_dim: (64, 1, 1),
5074 shared_mem_bytes: 0,
5075 };
5076 let __s_bl = self.gpu.stream();
5077 let mut bl = __s_bl.launch_builder(&f);
5078 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
5079 unsafe {
5080 bl.launch(cfg)?;
5081 }
5082 Ok(())
5083 }
5084
5085 pub fn plain_tok_ring(
5087 &self,
5088 vam: &CudaSlice<u32>,
5089 pos_start: &CudaSlice<i32>,
5090 base: usize,
5091 ring: &mut CudaSlice<u32>,
5092 ) -> Result<(), Box<dyn std::error::Error>> {
5093 let f = self.func("plain_tok_ring");
5094 let (b, cap) = (base as i32, ring.len() as i32);
5095 let cfg = LaunchConfig {
5096 grid_dim: (1, 1, 1),
5097 block_dim: (32, 1, 1),
5098 shared_mem_bytes: 0,
5099 };
5100 let __s_bl = self.gpu.stream();
5101 let mut bl = __s_bl.launch_builder(&f);
5102 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
5103 unsafe {
5104 bl.launch(cfg)?;
5105 }
5106 Ok(())
5107 }
5108
5109 pub fn spec_ring_commit(
5111 &self,
5112 vtok: &CudaSlice<u32>,
5113 acc: &CudaSlice<u32>,
5114 brk: &CudaSlice<u32>,
5115 ring: &mut CudaSlice<u32>,
5116 pend: &mut CudaSlice<u32>,
5117 ) -> Result<(), Box<dyn std::error::Error>> {
5118 let f = self.func("spec_ring_commit");
5119 let cfg = LaunchConfig {
5120 grid_dim: (1, 1, 1),
5121 block_dim: (32, 1, 1),
5122 shared_mem_bytes: 0,
5123 };
5124 let __s_b = self.gpu.stream();
5125 let mut b = __s_b.launch_builder(&f);
5126 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
5127 unsafe {
5128 b.launch(cfg)?;
5129 }
5130 Ok(())
5131 }
5132 pub fn i32_copy_add(
5133 &self,
5134 src: &CudaSlice<i32>,
5135 dst: &mut CudaSlice<i32>,
5136 delta: i32,
5137 ) -> Result<(), Box<dyn std::error::Error>> {
5138 let f = self.func("i32_copy_add");
5139 let cfg = LaunchConfig {
5140 grid_dim: (1, 1, 1),
5141 block_dim: (32, 1, 1),
5142 shared_mem_bytes: 0,
5143 };
5144 let __s_b = self.gpu.stream();
5145 let mut b = __s_b.launch_builder(&f);
5146 b.arg(src).arg(dst).arg(&delta);
5147 unsafe {
5148 b.launch(cfg)?;
5149 }
5150 Ok(())
5151 }
5152 pub fn u32_copy(
5153 &self,
5154 src: &CudaSlice<u32>,
5155 dst: &mut CudaSlice<u32>,
5156 ) -> Result<(), Box<dyn std::error::Error>> {
5157 let f = self.func("u32_copy");
5158 let cfg = LaunchConfig {
5159 grid_dim: (1, 1, 1),
5160 block_dim: (32, 1, 1),
5161 shared_mem_bytes: 0,
5162 };
5163 let __s_b = self.gpu.stream();
5164 let mut b = __s_b.launch_builder(&f);
5165 b.arg(src).arg(dst);
5166 unsafe {
5167 b.launch(cfg)?;
5168 }
5169 Ok(())
5170 }
5171
5172 pub fn spec_adapt_k(
5176 &self,
5177 acc: &CudaSlice<u32>,
5178 brk: &mut CudaSlice<u32>,
5179 floor: usize,
5180 cap: usize,
5181 ) -> Result<(), Box<dyn std::error::Error>> {
5182 let f = self.func("spec_adapt_k");
5183 let (fl, cp) = (floor as i32, cap as i32);
5184 let cfg = LaunchConfig {
5185 grid_dim: (1, 1, 1),
5186 block_dim: (32, 1, 1),
5187 shared_mem_bytes: 0,
5188 };
5189 let __s_b = self.gpu.stream();
5190 let mut b = __s_b.launch_builder(&f);
5191 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
5192 unsafe {
5193 b.launch(cfg)?;
5194 }
5195 Ok(())
5196 }
5197
5198 pub fn spec_accept_greedy_dc(
5200 &self,
5201 preds: &CudaSlice<u32>,
5202 vtok: &CudaSlice<u32>,
5203 last_pred: &CudaSlice<u32>,
5204 brk: &CudaSlice<u32>,
5205 out: &mut CudaSlice<u32>,
5206 ) -> Result<(), Box<dyn std::error::Error>> {
5207 let f = self.func("spec_accept_greedy_dc");
5208 let cfg = LaunchConfig {
5209 grid_dim: (1, 1, 1),
5210 block_dim: (32, 1, 1),
5211 shared_mem_bytes: 0,
5212 };
5213 let __s_b = self.gpu.stream();
5214 let mut b = __s_b.launch_builder(&f);
5215 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
5216 unsafe {
5217 b.launch(cfg)?;
5218 }
5219 Ok(())
5220 }
5221
5222 pub fn pos_iota(
5224 &self,
5225 pos0: &CudaSlice<i32>,
5226 out: &mut CudaSlice<i32>,
5227 t: usize,
5228 ) -> Result<(), Box<dyn std::error::Error>> {
5229 let f = self.func("pos_iota_i32");
5230 let ti = t as i32;
5231 let cfg = LaunchConfig {
5232 grid_dim: (1, 1, 1),
5233 block_dim: (t.max(1) as u32, 1, 1),
5234 shared_mem_bytes: 0,
5235 };
5236 let __s_b = self.gpu.stream();
5237 let mut b = __s_b.launch_builder(&f);
5238 b.arg(pos0).arg(out).arg(&ti);
5239 unsafe {
5240 b.launch(cfg)?;
5241 }
5242 Ok(())
5243 }
5244 #[allow(clippy::too_many_arguments)]
5245 pub fn append_kv_quantized_rows_dc(
5246 &self,
5247 k_rows: &CudaSlice<f32>,
5248 v_rows: &CudaSlice<f32>,
5249 kc: &mut CudaSlice<u8>,
5250 vc: &mut CudaSlice<u8>,
5251 t0_dev: &CudaSlice<i32>,
5252 t: usize,
5253 kv_dim_k: usize,
5254 kv_dim_v: usize,
5255 k_tok_bytes: usize,
5256 v_tok_bytes: usize,
5257 g: bool,
5258 ) -> Result<(), Box<dyn std::error::Error>> {
5259 let f = if g {
5260 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
5261 } else {
5262 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
5263 };
5264 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
5265 let cfg = LaunchConfig {
5266 grid_dim: (nblk, t as u32, 1),
5267 block_dim: (32, 1, 1),
5268 shared_mem_bytes: 0,
5269 };
5270 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
5271 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5272 let __s_b = self.gpu.stream();
5273 let mut b = __s_b.launch_builder(&f);
5274 b.arg(k_rows)
5275 .arg(v_rows)
5276 .arg(kc)
5277 .arg(vc)
5278 .arg(t0_dev)
5279 .arg(&kdk)
5280 .arg(&kdv)
5281 .arg(&ktb)
5282 .arg(&vtb);
5283 unsafe {
5284 b.launch(cfg)?;
5285 }
5286 Ok(())
5287 }
5288
5289 #[allow(clippy::too_many_arguments)]
5292 pub fn append_kv_quantized_row_dc_inc(
5293 &self,
5294 k_row: &CudaSlice<f32>,
5295 v_row: &CudaSlice<f32>,
5296 kc: &mut CudaSlice<u8>,
5297 vc: &mut CudaSlice<u8>,
5298 t0_dev: &mut CudaSlice<i32>,
5299 kv_dim_k: usize,
5300 kv_dim_v: usize,
5301 k_tok_bytes: usize,
5302 v_tok_bytes: usize,
5303 g: bool,
5304 ) -> Result<(), Box<dyn std::error::Error>> {
5305 let f = if g {
5306 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
5307 } else {
5308 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
5309 };
5310 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
5311 let cfg = LaunchConfig {
5312 grid_dim: (1, 1, 1),
5313 block_dim: (nthreads, 1, 1),
5314 shared_mem_bytes: 0,
5315 };
5316 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
5317 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5318 let __s_b = self.gpu.stream();
5319 let mut b = __s_b.launch_builder(&f);
5320 b.arg(k_row)
5321 .arg(v_row)
5322 .arg(kc)
5323 .arg(vc)
5324 .arg(t0_dev)
5325 .arg(&kdk)
5326 .arg(&kdv)
5327 .arg(&ktb)
5328 .arg(&vtb);
5329 unsafe {
5330 b.launch(cfg)?;
5331 }
5332 Ok(())
5333 }
5334
5335 pub fn pack_tok_p(
5337 &self,
5338 tok: &CudaSlice<u32>,
5339 p: &CudaSlice<f32>,
5340 out: &mut CudaSlice<u32>,
5341 slot: usize,
5342 ) -> Result<(), Box<dyn std::error::Error>> {
5343 let f = self.func("pack_tok_p");
5344 let sl = slot as i32;
5345 let cfg = LaunchConfig {
5346 grid_dim: (1, 1, 1),
5347 block_dim: (32, 1, 1),
5348 shared_mem_bytes: 0,
5349 };
5350 let __s_b = self.gpu.stream();
5351 let mut b = __s_b.launch_builder(&f);
5352 b.arg(tok).arg(p).arg(out).arg(&sl);
5353 unsafe {
5354 b.launch(cfg)?;
5355 }
5356 Ok(())
5357 }
5358 pub fn tok_map_u32(
5359 &self,
5360 tok: &mut CudaSlice<u32>,
5361 map: &CudaSlice<u32>,
5362 ) -> Result<(), Box<dyn std::error::Error>> {
5363 let f = self.func("tok_map_u32");
5364 let cfg = LaunchConfig {
5365 grid_dim: (1, 1, 1),
5366 block_dim: (32, 1, 1),
5367 shared_mem_bytes: 0,
5368 };
5369 let __s_b = self.gpu.stream();
5370 let mut b = __s_b.launch_builder(&f);
5371 b.arg(tok).arg(map);
5372 unsafe {
5373 b.launch(cfg)?;
5374 }
5375 Ok(())
5376 }
5377
5378 #[allow(clippy::too_many_arguments)]
5380 pub fn spec_assemble_verify(
5381 &self,
5382 tokp: &CudaSlice<u32>,
5383 pend: &CudaSlice<u32>,
5384 d2t: Option<&CudaSlice<u32>>,
5385 vtok: &mut CudaSlice<u32>,
5386 brk: &mut CudaSlice<u32>,
5387 p_min: f32,
5388 k: usize,
5389 pmin0: bool,
5390 ) -> Result<(), Box<dyn std::error::Error>> {
5391 let f = self.func("spec_assemble_verify");
5392 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
5393 let cfg = LaunchConfig {
5394 grid_dim: (1, 1, 1),
5395 block_dim: (32, 1, 1),
5396 shared_mem_bytes: 0,
5397 };
5398 let __s_b = self.gpu.stream();
5399 let mut b = __s_b.launch_builder(&f);
5400 match d2t {
5401 Some(m) => {
5402 b.arg(tokp)
5403 .arg(pend)
5404 .arg(m)
5405 .arg(vtok)
5406 .arg(brk)
5407 .arg(&p_min)
5408 .arg(&ki)
5409 .arg(&pm);
5410 unsafe {
5411 b.launch(cfg)?;
5412 }
5413 }
5414 None => {
5415 let null: u64 = 0;
5416 b.arg(tokp)
5417 .arg(pend)
5418 .arg(&null)
5419 .arg(vtok)
5420 .arg(brk)
5421 .arg(&p_min)
5422 .arg(&ki)
5423 .arg(&pm);
5424 unsafe {
5425 b.launch(cfg)?;
5426 }
5427 }
5428 }
5429 Ok(())
5430 }
5431
5432 #[allow(clippy::too_many_arguments)]
5434 pub fn ssm_conv_ring_rebuild_dc(
5435 &self,
5436 qkv_tm: &CudaSlice<f32>,
5437 ring_old: &CudaSlice<f32>,
5438 conv_state: &mut CudaSlice<f32>,
5439 conv_dim: usize,
5440 acc: &CudaSlice<u32>,
5441 base: usize,
5442 t_v: usize,
5443 d_conv: usize,
5444 ) -> Result<(), Box<dyn std::error::Error>> {
5445 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
5446 let n = conv_dim * (d_conv - 1);
5447 let cfg = LaunchConfig::for_num_elems(n as u32);
5448 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
5449 let __s_b = self.gpu.stream();
5450 let mut b = __s_b.launch_builder(&f);
5451 b.arg(qkv_tm)
5452 .arg(ring_old)
5453 .arg(conv_state)
5454 .arg(&cd)
5455 .arg(acc)
5456 .arg(&b0)
5457 .arg(&tv)
5458 .arg(&dc);
5459 unsafe {
5460 b.launch(cfg)?;
5461 }
5462 Ok(())
5463 }
5464 #[allow(clippy::too_many_arguments)]
5465 pub fn gdn_scan_s128_dc(
5466 &self,
5467 q: &CudaSlice<f32>,
5468 k: &CudaSlice<f32>,
5469 v: &CudaSlice<f32>,
5470 g: &CudaSlice<f32>,
5471 beta: &CudaSlice<f32>,
5472 state_in: &CudaSlice<f32>,
5473 state_out: &mut CudaSlice<f32>,
5474 o: &mut CudaSlice<f32>,
5475 n_head: usize,
5476 acc: &CudaSlice<u32>,
5477 base: usize,
5478 t_v: usize,
5479 scale: f32,
5480 ) -> Result<(), Box<dyn std::error::Error>> {
5481 let f = self.func("gdn_scan_s128_dc");
5482 const S_V: u32 = 128;
5483 const WARP: u32 = 32;
5484 const COLS_PER_BLOCK: u32 = 4;
5485 let cfg = LaunchConfig {
5486 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
5487 block_dim: (WARP, COLS_PER_BLOCK, 1),
5488 shared_mem_bytes: 0,
5489 };
5490 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
5491 let __s_b = self.gpu.stream();
5492 let mut b = __s_b.launch_builder(&f);
5493 b.arg(q)
5494 .arg(k)
5495 .arg(v)
5496 .arg(g)
5497 .arg(beta)
5498 .arg(state_in)
5499 .arg(state_out)
5500 .arg(o)
5501 .arg(&h)
5502 .arg(acc)
5503 .arg(&b0)
5504 .arg(&tv)
5505 .arg(&scale);
5506 unsafe {
5507 b.launch(cfg)?;
5508 }
5509 Ok(())
5510 }
5511
5512 pub fn spec_rollback_kv(
5514 &self,
5515 len_ptrs: &CudaSlice<u64>,
5516 saved: &CudaSlice<i32>,
5517 acc: &CudaSlice<u32>,
5518 base: usize,
5519 n_layer: usize,
5520 ) -> Result<(), Box<dyn std::error::Error>> {
5521 let f = self.func("spec_rollback_kv");
5522 let (b, nl) = (base as i32, n_layer as i32);
5523 let cfg = LaunchConfig {
5524 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
5525 block_dim: (64, 1, 1),
5526 shared_mem_bytes: 0,
5527 };
5528 let __s_bl = self.gpu.stream();
5529 let mut bl = __s_bl.launch_builder(&f);
5530 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
5531 unsafe {
5532 bl.launch(cfg)?;
5533 }
5534 Ok(())
5535 }
5536
5537 pub fn spec_fork_valid(
5539 &self,
5540 acc: &CudaSlice<u32>,
5541 optimistic_pending: u32,
5542 valid: &mut CudaSlice<u32>,
5543 ) -> Result<(), Box<dyn std::error::Error>> {
5544 let f = self.func("spec_fork_valid");
5545 let cfg = LaunchConfig {
5546 grid_dim: (1, 1, 1),
5547 block_dim: (1, 1, 1),
5548 shared_mem_bytes: 0,
5549 };
5550 let __s_bl = self.gpu.stream();
5551 let mut bl = __s_bl.launch_builder(&f);
5552 bl.arg(acc).arg(&optimistic_pending).arg(valid);
5553 unsafe {
5554 bl.launch(cfg)?;
5555 }
5556 Ok(())
5557 }
5558
5559 pub fn spec_fork_reconcile_kv(
5561 &self,
5562 len_ptrs: &CudaSlice<u64>,
5563 saved: &CudaSlice<i32>,
5564 acc: &CudaSlice<u32>,
5565 valid: &CudaSlice<u32>,
5566 base: usize,
5567 n_layer: usize,
5568 ) -> Result<(), Box<dyn std::error::Error>> {
5569 let f = self.func("spec_fork_reconcile_kv");
5570 let (b, nl) = (base as i32, n_layer as i32);
5571 let cfg = LaunchConfig {
5572 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
5573 block_dim: (64, 1, 1),
5574 shared_mem_bytes: 0,
5575 };
5576 let __s_bl = self.gpu.stream();
5577 let mut bl = __s_bl.launch_builder(&f);
5578 bl.arg(len_ptrs)
5579 .arg(saved)
5580 .arg(acc)
5581 .arg(valid)
5582 .arg(&b)
5583 .arg(&nl);
5584 unsafe {
5585 bl.launch(cfg)?;
5586 }
5587 Ok(())
5588 }
5589
5590 pub fn spec_fork_restore_f32(
5592 &self,
5593 snapshot: &CudaSlice<f32>,
5594 state: &mut CudaSlice<f32>,
5595 valid: &CudaSlice<u32>,
5596 ) -> Result<(), Box<dyn std::error::Error>> {
5597 assert_eq!(
5598 snapshot.len(),
5599 state.len(),
5600 "fork recurrent snapshot shape mismatch"
5601 );
5602 let f = self.func("spec_fork_restore_f32");
5603 let n = state.len() as i32;
5604 #[allow(clippy::manual_clamp)]
5605 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
5607 let cfg = LaunchConfig {
5608 grid_dim: (blocks, 1, 1),
5609 block_dim: (256, 1, 1),
5610 shared_mem_bytes: 0,
5611 };
5612 let __s_bl = self.gpu.stream();
5613 let mut bl = __s_bl.launch_builder(&f);
5614 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
5615 unsafe {
5616 bl.launch(cfg)?;
5617 }
5618 Ok(())
5619 }
5620
5621 pub fn spec_seed_gather(
5624 &self,
5625 vx: &CudaSlice<f32>,
5626 fill_prev: &CudaSlice<f32>,
5627 acc: &CudaSlice<u32>,
5628 h_seed: &mut CudaSlice<f32>,
5629 base: usize,
5630 n_embd: usize,
5631 ) -> Result<(), Box<dyn std::error::Error>> {
5632 let f = self.func("spec_seed_gather");
5633 let (b, ne) = (base as i32, n_embd as i32);
5634 let cfg = LaunchConfig {
5635 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
5636 block_dim: (256, 1, 1),
5637 shared_mem_bytes: 0,
5638 };
5639 let __s_bl = self.gpu.stream();
5640 let mut bl = __s_bl.launch_builder(&f);
5641 bl.arg(vx)
5642 .arg(fill_prev)
5643 .arg(acc)
5644 .arg(h_seed)
5645 .arg(&b)
5646 .arg(&ne);
5647 unsafe {
5648 bl.launch(cfg)?;
5649 }
5650 Ok(())
5651 }
5652
5653 pub fn spec_accept_greedy(
5655 &self,
5656 preds: &CudaSlice<u32>,
5657 draft: &CudaSlice<u32>,
5658 last_pred: u32,
5659 base: usize,
5660 k_round: usize,
5661 out: &mut CudaSlice<u32>,
5662 ) -> Result<(), Box<dyn std::error::Error>> {
5663 let f = self.func("spec_accept_greedy");
5664 let (b, k) = (base as i32, k_round as i32);
5665 let cfg = LaunchConfig {
5666 grid_dim: (1, 1, 1),
5667 block_dim: (32, 1, 1),
5668 shared_mem_bytes: 0,
5669 };
5670 let __s_bl = self.gpu.stream();
5671 let mut bl = __s_bl.launch_builder(&f);
5672 bl.arg(preds)
5673 .arg(draft)
5674 .arg(&last_pred)
5675 .arg(&b)
5676 .arg(&k)
5677 .arg(out);
5678 unsafe {
5679 bl.launch(cfg)?;
5680 }
5681 Ok(())
5682 }
5683
5684 pub fn gumbel_perturb(
5691 &self,
5692 x: &CudaSlice<f32>,
5693 y: &mut CudaSlice<f32>,
5694 n: usize,
5695 seed: u64,
5696 stream_pos: u32,
5697 temp: f32,
5698 ) -> Result<(), Box<dyn std::error::Error>> {
5699 let f = self.func("gumbel_perturb_f32");
5700 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5701 let cfg = LaunchConfig {
5702 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5703 block_dim: (256, 1, 1),
5704 shared_mem_bytes: 0,
5705 };
5706 let __s_b = self.gpu.stream();
5707 let mut b = __s_b.launch_builder(&f);
5708 b.arg(x)
5709 .arg(&mut *y)
5710 .arg(&ni)
5711 .arg(&slo)
5712 .arg(&shi)
5713 .arg(&stream_pos)
5714 .arg(&temp);
5715 unsafe {
5716 b.launch(cfg)?;
5717 }
5718 Ok(())
5719 }
5720
5721 pub fn mask_logits_col(
5729 &self,
5730 logits: &mut CudaSlice<f32>,
5731 mask: &CudaSlice<u32>,
5732 col: usize,
5733 n: usize,
5734 mask_words: usize,
5735 ) -> Result<(), Box<dyn std::error::Error>> {
5736 let f = self.func("mask_logits_f32");
5737 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
5738 let cfg = LaunchConfig {
5739 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
5740 block_dim: (256, 1, 1),
5741 shared_mem_bytes: 0,
5742 };
5743 let __s_b = self.gpu.stream();
5744 let mut b = __s_b.launch_builder(&f);
5745 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
5746 unsafe {
5747 b.launch(cfg)?;
5748 }
5749 Ok(())
5750 }
5751
5752 #[allow(clippy::too_many_arguments)] pub fn gumbel_perturb_col(
5760 &self,
5761 x: &CudaSlice<f32>,
5762 col: usize,
5763 y: &mut CudaSlice<f32>,
5764 n: usize,
5765 seed: u64,
5766 stream_pos: u32,
5767 temp: f32,
5768 ) -> Result<(), Box<dyn std::error::Error>> {
5769 let f = self.func("gumbel_perturb_f32");
5770 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5771 let col_view = x.slice(col * n..(col + 1) * n);
5772 let cfg = LaunchConfig {
5773 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5774 block_dim: (256, 1, 1),
5775 shared_mem_bytes: 0,
5776 };
5777 let __s_b = self.gpu.stream();
5778 let mut b = __s_b.launch_builder(&f);
5779 b.arg(&col_view)
5780 .arg(&mut *y)
5781 .arg(&ni)
5782 .arg(&slo)
5783 .arg(&shi)
5784 .arg(&stream_pos)
5785 .arg(&temp);
5786 unsafe {
5787 b.launch(cfg)?;
5788 }
5789 Ok(())
5790 }
5791
5792 #[allow(clippy::too_many_arguments)]
5798 pub fn gumbel_perturb_filtered_col(
5799 &self,
5800 x: &CudaSlice<f32>,
5801 col: usize,
5802 y: &mut CudaSlice<f32>,
5803 n: usize,
5804 seed: u64,
5805 stream_pos: u32,
5806 temp: f32,
5807 stat_max: &CudaSlice<f32>,
5808 stat_th: &CudaSlice<f32>,
5809 stat_idx: usize,
5810 ) -> Result<(), Box<dyn std::error::Error>> {
5811 let f = self.func("gumbel_perturb_filtered_col_f32");
5812 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5813 let (ci, si) = (col as i32, stat_idx as i32);
5814 let cfg = LaunchConfig {
5815 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5816 block_dim: (256, 1, 1),
5817 shared_mem_bytes: 0,
5818 };
5819 let __s_b = self.gpu.stream();
5820 let mut b = __s_b.launch_builder(&f);
5821 b.arg(x)
5822 .arg(&ci)
5823 .arg(&mut *y)
5824 .arg(&ni)
5825 .arg(&slo)
5826 .arg(&shi)
5827 .arg(&stream_pos)
5828 .arg(&temp)
5829 .arg(stat_max)
5830 .arg(stat_th)
5831 .arg(&si);
5832 unsafe {
5833 b.launch(cfg)?;
5834 }
5835 Ok(())
5836 }
5837
5838 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
5843 let f = self.func("memra_sctr_inc");
5844 let cfg = LaunchConfig {
5845 grid_dim: (1, 1, 1),
5846 block_dim: (1, 1, 1),
5847 shared_mem_bytes: 0,
5848 };
5849 let __s_b = self.gpu.stream();
5850 let mut b = __s_b.launch_builder(&f);
5851 b.arg(&mut *ctr);
5852 unsafe {
5853 b.launch(cfg)?;
5854 }
5855 Ok(())
5856 }
5857
5858 pub fn gumbel_perturb_ctr(
5863 &self,
5864 x: &CudaSlice<f32>,
5865 y: &mut CudaSlice<f32>,
5866 n: usize,
5867 seed: u64,
5868 ctr: &CudaSlice<u32>,
5869 temp: f32,
5870 ) -> Result<(), Box<dyn std::error::Error>> {
5871 let f = self.func("gumbel_perturb_ctr_f32");
5872 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5873 let cfg = LaunchConfig {
5874 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5875 block_dim: (256, 1, 1),
5876 shared_mem_bytes: 0,
5877 };
5878 let __s_b = self.gpu.stream();
5879 let mut b = __s_b.launch_builder(&f);
5880 b.arg(x)
5881 .arg(&mut *y)
5882 .arg(&ni)
5883 .arg(&slo)
5884 .arg(&shi)
5885 .arg(ctr)
5886 .arg(&temp);
5887 unsafe {
5888 b.launch(cfg)?;
5889 }
5890 Ok(())
5891 }
5892
5893 #[allow(clippy::too_many_arguments)]
5901 pub fn gumbel_perturb_filtered_ctr(
5902 &self,
5903 x: &CudaSlice<f32>,
5904 y: &mut CudaSlice<f32>,
5905 n: usize,
5906 seed: u64,
5907 ctr: &CudaSlice<u32>,
5908 temp: f32,
5909 stat_max: &CudaSlice<f32>,
5910 stat_th: &CudaSlice<f32>,
5911 ) -> Result<(), Box<dyn std::error::Error>> {
5912 let f = self.func("gumbel_perturb_filtered_ctr_f32");
5913 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5914 let cfg = LaunchConfig {
5915 grid_dim: (n.div_ceil(256) as u32, 1, 1),
5916 block_dim: (256, 1, 1),
5917 shared_mem_bytes: 0,
5918 };
5919 let __s_b = self.gpu.stream();
5920 let mut b = __s_b.launch_builder(&f);
5921 b.arg(x)
5922 .arg(&mut *y)
5923 .arg(&ni)
5924 .arg(&slo)
5925 .arg(&shi)
5926 .arg(ctr)
5927 .arg(&temp)
5928 .arg(stat_max)
5929 .arg(stat_th);
5930 unsafe {
5931 b.launch(cfg)?;
5932 }
5933 Ok(())
5934 }
5935
5936 #[allow(clippy::too_many_arguments)] pub fn softmax_gather(
5941 &self,
5942 x: &CudaSlice<f32>,
5943 row_stride: usize,
5944 ids: &CudaSlice<u32>,
5945 rows: &CudaSlice<i32>,
5946 out: &mut CudaSlice<f32>,
5947 n: usize,
5948 npair: usize,
5949 temp: f32,
5950 ) -> Result<(), Box<dyn std::error::Error>> {
5951 let f = self.func("softmax_gather_f32");
5952 let (ni, rs) = (n as i32, row_stride as i64);
5953 let np = npair as i32;
5954 let cfg = LaunchConfig {
5955 grid_dim: (npair as u32, 1, 1),
5956 block_dim: (256, 1, 1),
5957 shared_mem_bytes: 0,
5958 };
5959 let __s_b = self.gpu.stream();
5960 let mut b = __s_b.launch_builder(&f);
5961 b.arg(x)
5962 .arg(&rs)
5963 .arg(ids)
5964 .arg(rows)
5965 .arg(&mut *out)
5966 .arg(&ni)
5967 .arg(&np)
5968 .arg(&temp);
5969 unsafe {
5970 b.launch(cfg)?;
5971 }
5972 Ok(())
5973 }
5974
5975 #[allow(clippy::too_many_arguments)] pub fn residual_sample(
5980 &self,
5981 p: &CudaSlice<f32>,
5982 q: Option<&CudaSlice<f32>>,
5983 n: usize,
5984 temp: f32,
5985 seed: u64,
5986 stream_pos: u32,
5987 out_tok: &mut CudaSlice<u32>,
5988 ) -> Result<(), Box<dyn std::error::Error>> {
5989 let f = self.func("residual_sample_f32");
5990 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
5991 let nth = 1024u32;
5992 let cfg = LaunchConfig {
5993 grid_dim: (1, 1, 1),
5994 block_dim: (nth, 1, 1),
5995 shared_mem_bytes: 0,
5996 };
5997 let has_q: i32 = q.is_some() as i32;
5998 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
6000 let mut b = __s_b.launch_builder(&f);
6001 b.arg(p)
6002 .arg(qbuf)
6003 .arg(&has_q)
6004 .arg(&ni)
6005 .arg(&temp)
6006 .arg(&slo)
6007 .arg(&shi)
6008 .arg(&stream_pos)
6009 .arg(&mut *out_tok);
6010 unsafe {
6011 b.launch(cfg)?;
6012 }
6013 Ok(())
6014 }
6015
6016 pub fn with_moe_cache<R>(
6021 &self,
6022 max_block_bytes: usize,
6023 f: impl FnOnce(
6024 &mut crate::moe_cache::MoeSlotCache,
6025 &Engine,
6026 ) -> Result<R, Box<dyn std::error::Error>>,
6027 ) -> Result<R, Box<dyn std::error::Error>> {
6028 let mut guard = self.moe_cache.lock().unwrap();
6029 if guard.is_none() {
6030 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
6031 }
6032 let cache = guard.as_mut().unwrap();
6033 f(cache, self)
6034 }
6035
6036 pub fn freeze_moe_cache(&self) {
6039 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
6040 cache.freeze();
6041 }
6042 }
6043
6044 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
6047 self.moe_cache
6048 .lock()
6049 .unwrap()
6050 .as_ref()
6051 .map(crate::moe_cache::MoeSlotCache::export_residency)
6052 }
6053
6054 pub(crate) fn moe_cache_frozen(&self) -> bool {
6055 self.moe_cache
6056 .lock()
6057 .unwrap()
6058 .as_ref()
6059 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
6060 }
6061
6062 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
6069 crate::cpu_experts::configured()
6070 && self.moe_cache_frozen()
6071 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
6072 }
6073
6074 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
6076 assert!(
6077 self.moe_cache.lock().unwrap().is_none(),
6078 "MoE cache layout configured after cache construction"
6079 );
6080 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
6081 }
6082
6083 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
6084 self.moe_cache_layout.lock().unwrap().clone()
6085 }
6086
6087 pub fn moe_cache_enabled() -> bool {
6089 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
6090 }
6091
6092 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
6095 let guard = self.moe_cache.lock().unwrap();
6096 guard
6097 .as_ref()
6098 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
6099 }
6100
6101 #[allow(clippy::type_complexity)] pub fn cpu_expert_stats(
6106 &self,
6107 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
6108 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
6109 }
6110
6111 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
6114 crate::cpu_experts::predictor_stats()
6115 }
6116
6117 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
6118 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
6119 }
6120
6121 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
6124 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
6125 }
6126
6127 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
6130 let guard = self.moe_cache.lock().unwrap();
6131 guard
6132 .as_ref()
6133 .and_then(|cache| cache.pread_stats())
6134 .map(|stats| {
6135 (
6136 stats.reads,
6137 stats.bytes,
6138 stats.read_errors,
6139 stats.short_reads,
6140 stats.fallbacks,
6141 stats.buffer_waits,
6142 stats.ring_full,
6143 )
6144 })
6145 }
6146
6147 pub fn spill_config_fallbacks(&self) -> u64 {
6149 crate::spill_pread::config_fallbacks()
6150 }
6151
6152 pub fn moe_cache_reset_counters(&self) {
6154 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
6155 c.reset_counters();
6156 }
6157 }
6158
6159 #[track_caller]
6160 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6161 crate::alloc_trace_hit(v.len());
6162 Ok(self.gpu.stream().clone_htod(v)?)
6163 }
6164
6165 pub fn htod_bytes_padded(
6169 &self,
6170 v: &[u8],
6171 pad: usize,
6172 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6173 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
6174 {
6175 let mut view = d.slice_mut(0..v.len());
6176 self.gpu.stream().memcpy_htod(v, &mut view)?;
6177 }
6178 Ok(d)
6179 }
6180
6181 pub fn copy_into(
6183 &self,
6184 dst: &mut CudaSlice<f32>,
6185 off: usize,
6186 src: &CudaSlice<f32>,
6187 len: usize,
6188 ) -> Result<(), Box<dyn std::error::Error>> {
6189 let mut view = dst.slice_mut(off..off + len);
6190 self.gpu
6191 .stream()
6192 .memcpy_dtod(&src.slice(0..len), &mut view)?;
6193 Ok(())
6194 }
6195
6196 pub fn copy_range_into(
6200 &self,
6201 dst: &mut CudaSlice<f32>,
6202 dst_off: usize,
6203 src: &CudaSlice<f32>,
6204 src_off: usize,
6205 len: usize,
6206 ) -> Result<(), Box<dyn std::error::Error>> {
6207 let mut view = dst.slice_mut(dst_off..dst_off + len);
6208 self.gpu
6209 .stream()
6210 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut view)?;
6211 Ok(())
6212 }
6213
6214 pub fn copy_u8_into(
6217 &self,
6218 dst: &mut CudaSlice<u8>,
6219 off: usize,
6220 src: &CudaSlice<u8>,
6221 len: usize,
6222 ) -> Result<(), Box<dyn std::error::Error>> {
6223 let cap = dst.len();
6227 let mut view = dst.try_slice_mut(off..off + len).ok_or_else(|| {
6228 format!(
6229 "copy_u8_into dst range [{off},{}) exceeds capacity {cap}",
6230 off + len,
6231 )
6232 })?;
6233 self.gpu
6234 .stream()
6235 .memcpy_dtod(&src.slice(0..len), &mut view)?;
6236 Ok(())
6237 }
6238
6239 pub fn copy_u8_range_into(
6241 &self,
6242 dst: &mut CudaSlice<u8>,
6243 dst_off: usize,
6244 src: &CudaSlice<u8>,
6245 src_off: usize,
6246 len: usize,
6247 ) -> Result<(), Box<dyn std::error::Error>> {
6248 let cap = dst.len();
6251 let mut dst_view = dst.try_slice_mut(dst_off..dst_off + len).ok_or_else(|| {
6252 format!(
6253 "copy_u8_range_into dst range [{dst_off},{}) exceeds capacity {cap}",
6254 dst_off + len,
6255 )
6256 })?;
6257 self.gpu
6258 .stream()
6259 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
6260 Ok(())
6261 }
6262
6263 #[track_caller]
6273 pub fn prepare_kv_append(
6274 &self,
6275 kv: &mut crate::cache::KvLayer,
6276 retain_from: usize,
6277 append_rows: usize,
6278 ) -> Result<usize, Box<dyn std::error::Error>> {
6279 let caller = std::panic::Location::caller();
6280 let base_before = kv.ring.as_ref().map(|r| r.base());
6281 let Some(plan) = kv
6282 .ring
6283 .as_ref()
6284 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
6285 .transpose()
6286 .map_err(|err| -> Box<dyn std::error::Error> {
6287 format!(
6288 "{err} [append len={} retain_from={retain_from} append_rows={append_rows} base={base_before:?} called from {caller}]",
6289 kv.len
6290 )
6291 .into()
6292 })?
6293 else {
6294 return Ok(kv.len);
6295 };
6296 match plan {
6297 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
6298 crate::cache::KvRingAppend::Rebase {
6299 src_row,
6300 keep_rows,
6301 new_base,
6302 write_row,
6303 } => {
6304 if keep_rows > 0 {
6305 let k_len = keep_rows * kv.k_tok_bytes;
6306 let v_len = keep_rows * kv.v_tok_bytes;
6307 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
6308 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
6309 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
6310 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
6311 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
6312 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
6313 }
6314 if std::env::var("MEMRA_KV_REBASE_TRACE").as_deref() == Ok("1") {
6317 eprintln!(
6318 "[kv-rebase] new_base={new_base} keep_rows={keep_rows} len={} \
6319 retain_from={retain_from} called from {caller}",
6320 kv.len
6321 );
6322 }
6323 kv.ring.as_mut().unwrap().apply_rebase(new_base);
6324 if let Some(base_d) = kv.base_d.as_mut() {
6328 self.set_i32_one(base_d, new_base as i32)?;
6329 }
6330 Ok(write_row)
6331 }
6332 }
6333 }
6334
6335 pub fn htod_u8_into(
6338 &self,
6339 dst: &mut CudaSlice<u8>,
6340 off: usize,
6341 src: &[u8],
6342 ) -> Result<(), Box<dyn std::error::Error>> {
6343 let mut view = dst.slice_mut(off..off + src.len());
6344 self.gpu.stream().memcpy_htod(src, &mut view)?;
6345 Ok(())
6346 }
6347
6348 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
6349 b.slice(0..len)
6350 }
6351
6352 pub fn view_u8_range<'a>(
6355 &self,
6356 b: &'a CudaSlice<u8>,
6357 start: usize,
6358 end: usize,
6359 ) -> cudarc::driver::CudaView<'a, u8> {
6360 b.slice(start..end)
6361 }
6362 pub fn view_u8<'a>(
6363 &self,
6364 b: &'a CudaSlice<u8>,
6365 len: usize,
6366 ) -> cudarc::driver::CudaView<'a, u8> {
6367 b.slice(0..len)
6368 }
6369
6370 #[allow(clippy::too_many_arguments)] pub fn append_kv_quantized(
6375 &self,
6376 k_row: &CudaSlice<f32>,
6377 v_row: &CudaSlice<f32>,
6378 kc: &mut CudaSlice<u8>,
6379 vc: &mut CudaSlice<u8>,
6380 t: usize,
6381 kv_dim_k: usize,
6382 kv_dim_v: usize,
6383 k_tok_bytes: usize,
6384 v_tok_bytes: usize,
6385 g: bool,
6386 ) -> Result<(), Box<dyn std::error::Error>> {
6387 let f = if g {
6388 self.func_g("append_quantize_kv_q8_0_q5_1")
6389 } else {
6390 self.func("append_quantize_kv_q8_0_q5_1")
6391 };
6392 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6393 let cfg = LaunchConfig {
6394 grid_dim: (nblk, 1, 1),
6395 block_dim: (32, 1, 1),
6396 shared_mem_bytes: 0,
6397 };
6398 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
6399 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6400 let __s_b = self.gpu.stream();
6401 let mut b = __s_b.launch_builder(&f);
6402 b.arg(k_row)
6403 .arg(v_row)
6404 .arg(kc)
6405 .arg(vc)
6406 .arg(&ti)
6407 .arg(&kdk)
6408 .arg(&kdv)
6409 .arg(&ktb)
6410 .arg(&vtb);
6411 unsafe {
6412 b.launch(cfg)?;
6413 }
6414 Ok(())
6415 }
6416
6417 #[allow(clippy::too_many_arguments)] pub fn append_kv_quantized_dc(
6422 &self,
6423 k_row: &CudaSlice<f32>,
6424 v_row: &CudaSlice<f32>,
6425 kc: &mut CudaSlice<u8>,
6426 vc: &mut CudaSlice<u8>,
6427 t_dev: &CudaSlice<i32>,
6428 kv_dim_k: usize,
6429 kv_dim_v: usize,
6430 k_tok_bytes: usize,
6431 v_tok_bytes: usize,
6432 g: bool,
6433 ) -> Result<(), Box<dyn std::error::Error>> {
6434 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6435 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
6436 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6437 if Self::pdl_on() && Self::pdl_wb_on() {
6439 use cudarc::driver::{DevicePtr, DevicePtrMut};
6440 let s = &self.gpu.stream();
6441 let (pk, _g0) = k_row.device_ptr(s);
6442 let (pv, _g1) = v_row.device_ptr(s);
6443 let (pkc, _g2) = kc.device_ptr_mut(s);
6444 let (pvc, _g3) = vc.device_ptr_mut(s);
6445 let (pt, _g4) = t_dev.device_ptr(s);
6446 let mut ps = [
6447 &pk as *const _ as *mut std::ffi::c_void,
6448 &pv as *const _ as *mut _,
6449 &pkc as *const _ as *mut _,
6450 &pvc as *const _ as *mut _,
6451 &pt as *const _ as *mut _,
6452 &kdk as *const _ as *mut _,
6453 &kdv as *const _ as *mut _,
6454 &ktb as *const _ as *mut _,
6455 &vtb as *const _ as *mut _,
6456 ];
6457 unsafe {
6458 self.launch_pdl_flash(
6459 g,
6460 "append_quantize_kv_q8_0_q5_1_dc",
6461 (nblk, 1, 1),
6462 (32, 1, 1),
6463 0,
6464 &mut ps,
6465 )?;
6466 }
6467 return Ok(());
6468 }
6469 let f = if g {
6470 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
6471 } else {
6472 self.func("append_quantize_kv_q8_0_q5_1_dc")
6473 };
6474 let cfg = LaunchConfig {
6475 grid_dim: (nblk, 1, 1),
6476 block_dim: (32, 1, 1),
6477 shared_mem_bytes: 0,
6478 };
6479 let __s_b = self.gpu.stream();
6480 let mut b = __s_b.launch_builder(&f);
6481 b.arg(k_row)
6482 .arg(v_row)
6483 .arg(kc)
6484 .arg(vc)
6485 .arg(t_dev)
6486 .arg(&kdk)
6487 .arg(&kdv)
6488 .arg(&ktb)
6489 .arg(&vtb);
6490 unsafe {
6491 b.launch(cfg)?;
6492 }
6493 Ok(())
6494 }
6495
6496 #[allow(clippy::too_many_arguments)]
6503 pub fn append_kv_quantized_rows(
6504 &self,
6505 k_rows: &CudaSlice<f32>,
6506 v_rows: &CudaSlice<f32>,
6507 kc: &mut CudaSlice<u8>,
6508 vc: &mut CudaSlice<u8>,
6509 t0: usize,
6510 t: usize,
6511 kv_dim_k: usize,
6512 kv_dim_v: usize,
6513 k_tok_bytes: usize,
6514 v_tok_bytes: usize,
6515 g: bool,
6516 ) -> Result<(), Box<dyn std::error::Error>> {
6517 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
6518 for i in 0..t {
6519 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
6520 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
6521 self.append_kv_quantized_view(
6522 &k_row,
6523 &v_row,
6524 kc,
6525 vc,
6526 t0 + i,
6527 kv_dim_k,
6528 kv_dim_v,
6529 k_tok_bytes,
6530 v_tok_bytes,
6531 g,
6532 )?;
6533 }
6534 return Ok(());
6535 }
6536 let f = if g {
6537 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
6538 } else {
6539 self.func("append_quantize_kv_q8_0_q5_1_rows")
6540 };
6541 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6542 let cfg = LaunchConfig {
6543 grid_dim: (nblk, t as u32, 1),
6544 block_dim: (32, 1, 1),
6545 shared_mem_bytes: 0,
6546 };
6547 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
6548 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6549 let __s_b = self.gpu.stream();
6550 let mut b = __s_b.launch_builder(&f);
6551 b.arg(k_rows)
6552 .arg(v_rows)
6553 .arg(kc)
6554 .arg(vc)
6555 .arg(&t0i)
6556 .arg(&kdk)
6557 .arg(&kdv)
6558 .arg(&ktb)
6559 .arg(&vtb);
6560 unsafe {
6561 b.launch(cfg)?;
6562 }
6563 Ok(())
6564 }
6565
6566 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
6570 let f = self.func("inc_i32");
6571 let cfg = LaunchConfig {
6572 grid_dim: (1, 1, 1),
6573 block_dim: (1, 1, 1),
6574 shared_mem_bytes: 0,
6575 };
6576 let __s_b = self.gpu.stream();
6577 let mut b = __s_b.launch_builder(&f);
6578 b.arg(p);
6579 unsafe {
6580 b.launch(cfg)?;
6581 }
6582 Ok(())
6583 }
6584
6585 #[allow(clippy::too_many_arguments)] pub fn append_kv_quantized_view(
6589 &self,
6590 k_row: &cudarc::driver::CudaView<f32>,
6591 v_row: &cudarc::driver::CudaView<f32>,
6592 kc: &mut CudaSlice<u8>,
6593 vc: &mut CudaSlice<u8>,
6594 t: usize,
6595 kv_dim_k: usize,
6596 kv_dim_v: usize,
6597 k_tok_bytes: usize,
6598 v_tok_bytes: usize,
6599 g: bool,
6600 ) -> Result<(), Box<dyn std::error::Error>> {
6601 let stream = self.gpu.stream();
6602 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
6603 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
6604 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
6605 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
6606 let f = if g {
6607 self.func_g("append_quantize_kv_q8_0_q5_1")
6608 } else {
6609 self.func("append_quantize_kv_q8_0_q5_1")
6610 };
6611 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
6612 let cfg = LaunchConfig {
6613 grid_dim: (nblk, 1, 1),
6614 block_dim: (32, 1, 1),
6615 shared_mem_bytes: 0,
6616 };
6617 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
6618 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
6619 let mut b = stream.launch_builder(&f);
6620 b.arg(k_row)
6621 .arg(v_row)
6622 .arg(kc)
6623 .arg(vc)
6624 .arg(&ti)
6625 .arg(&kdk)
6626 .arg(&kdv)
6627 .arg(&ktb)
6628 .arg(&vtb);
6629 unsafe {
6630 b.launch(cfg)?;
6631 }
6632 Ok(())
6633 }
6634
6635 pub fn copy_view_into(
6638 &self,
6639 dst: &mut CudaSlice<f32>,
6640 off: usize,
6641 src: &cudarc::driver::CudaView<f32>,
6642 len: usize,
6643 ) -> Result<(), Box<dyn std::error::Error>> {
6644 let mut view = dst.slice_mut(off..off + len);
6645 self.gpu
6646 .stream()
6647 .memcpy_dtod(&src.slice(0..len), &mut view)?;
6648 Ok(())
6649 }
6650
6651 pub fn clone_dtod(
6672 &self,
6673 src: &CudaSlice<f32>,
6674 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6675 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
6676 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
6677 Ok(dst)
6678 }
6679
6680 pub fn dtod_copy_view(
6683 &self,
6684 src: &cudarc::driver::CudaView<f32>,
6685 dst: &mut CudaSlice<f32>,
6686 ) -> Result<(), Box<dyn std::error::Error>> {
6687 self.gpu.stream().memcpy_dtod(src, dst)?;
6688 Ok(())
6689 }
6690
6691 pub fn dtod_copy_view_i8(
6693 &self,
6694 src: &cudarc::driver::CudaView<i8>,
6695 dst: &mut CudaSlice<i8>,
6696 ) -> Result<(), Box<dyn std::error::Error>> {
6697 self.gpu.stream().memcpy_dtod(src, dst)?;
6698 Ok(())
6699 }
6700
6701 pub fn dtod_copy_into(
6703 &self,
6704 src: &CudaSlice<f32>,
6705 dst: &mut CudaSlice<f32>,
6706 offset: usize,
6707 ) -> Result<(), Box<dyn std::error::Error>> {
6708 let n = src.len();
6709 let mut dv = dst.slice_mut(offset..offset + n);
6710 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
6711 Ok(())
6712 }
6713
6714 pub fn copy_batch_uniform_f32(
6720 &self,
6721 table: &CudaSlice<u64>,
6722 n: usize,
6723 words: usize,
6724 ) -> Result<(), Box<dyn std::error::Error>> {
6725 if n == 0 || words == 0 {
6726 return Ok(());
6727 }
6728 debug_assert!(
6729 table.len() >= 2 * n,
6730 "pointer table must hold n srcs + n dsts"
6731 );
6732 let f = self.func("copy_batch_uniform_f32");
6733 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
6736 let (ni, wi) = (n as i32, words as i32);
6737 let cfg = LaunchConfig {
6738 grid_dim: (chunks, n as u32, 1),
6739 block_dim: (256, 1, 1),
6740 shared_mem_bytes: 0,
6741 };
6742 let __s = self.gpu.stream();
6743 let mut b = __s.launch_builder(&f);
6744 b.arg(table).arg(&ni).arg(&wi);
6745 unsafe {
6746 b.launch(cfg)?;
6747 }
6748 Ok(())
6749 }
6750
6751 #[allow(clippy::too_many_arguments)] pub fn copy_batch_uniform_kv_u8_set_len(
6756 &self,
6757 table: &CudaSlice<u64>,
6758 n: usize,
6759 rows: usize,
6760 k_row_bytes: usize,
6761 v_row_bytes: usize,
6762 k_src_stride: usize,
6763 v_src_stride: usize,
6764 logical_len: usize,
6765 ) -> Result<(), Box<dyn std::error::Error>> {
6766 if n == 0 || rows == 0 || (k_row_bytes == 0 && v_row_bytes == 0) {
6767 return Ok(());
6768 }
6769 if table.len() < 5 * n {
6770 return Err(format!(
6771 "TP KV repair table has {} words, expected at least {}",
6772 table.len(),
6773 5 * n
6774 )
6775 .into());
6776 }
6777 let ni = i32::try_from(n).map_err(|_| "TP KV repair layer count exceeds i32")?;
6778 let rows = i32::try_from(rows).map_err(|_| "TP KV repair rows exceed i32")?;
6779 let kb = i32::try_from(k_row_bytes).map_err(|_| "TP KV repair K bytes exceed i32")?;
6780 let vb = i32::try_from(v_row_bytes).map_err(|_| "TP KV repair V bytes exceed i32")?;
6781 let ks = i32::try_from(k_src_stride).map_err(|_| "TP KV repair K stride exceeds i32")?;
6782 let vs = i32::try_from(v_src_stride).map_err(|_| "TP KV repair V stride exceeds i32")?;
6783 let len = i32::try_from(logical_len).map_err(|_| "TP KV repair length exceeds i32")?;
6784 let f = self.func("copy_batch_uniform_kv_u8_set_len");
6785 let cfg = LaunchConfig {
6786 grid_dim: (n as u32, 1, 1),
6787 block_dim: (256, 1, 1),
6788 shared_mem_bytes: 0,
6789 };
6790 let stream = self.gpu.stream();
6791 let mut builder = stream.launch_builder(&f);
6792 builder
6793 .arg(table)
6794 .arg(&ni)
6795 .arg(&rows)
6796 .arg(&kb)
6797 .arg(&vb)
6798 .arg(&ks)
6799 .arg(&vs)
6800 .arg(&len);
6801 unsafe {
6802 builder.launch(cfg)?;
6803 }
6804 Ok(())
6805 }
6806
6807 pub fn htod_u64_into(
6810 &self,
6811 v: &[u64],
6812 dst: &mut CudaSlice<u64>,
6813 ) -> Result<(), Box<dyn std::error::Error>> {
6814 let mut view = dst.slice_mut(0..v.len());
6815 self.gpu.stream().memcpy_htod(v, &mut view)?;
6816 Ok(())
6817 }
6818
6819 pub fn htod_f32_into(
6822 &self,
6823 v: &[f32],
6824 dst: &mut CudaSlice<f32>,
6825 ) -> Result<(), Box<dyn std::error::Error>> {
6826 let mut view = dst.slice_mut(0..v.len());
6827 self.gpu.stream().memcpy_htod(v, &mut view)?;
6828 Ok(())
6829 }
6830
6831 pub fn htod_f32_into_at(
6835 &self,
6836 v: &[f32],
6837 dst: &mut CudaSlice<f32>,
6838 off: usize,
6839 ) -> Result<(), Box<dyn std::error::Error>> {
6840 if off + v.len() > dst.len() {
6841 return Err(format!(
6842 "htod_f32_into_at range {}..{} exceeds dst {}",
6843 off,
6844 off + v.len(),
6845 dst.len()
6846 )
6847 .into());
6848 }
6849 let mut view = dst.slice_mut(off..off + v.len());
6850 self.gpu.stream().memcpy_htod(v, &mut view)?;
6851 Ok(())
6852 }
6853
6854 pub fn copy_indirect_src_f32(
6859 &self,
6860 src_entry: &cudarc::driver::CudaView<u64>,
6861 dst: &mut CudaSlice<f32>,
6862 dst_off: usize,
6863 words: usize,
6864 ) -> Result<(), Box<dyn std::error::Error>> {
6865 let f = self.func("copy_indirect_src_f32");
6866 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
6867 let wi = words as i32;
6868 let cfg = LaunchConfig {
6869 grid_dim: (chunks, 1, 1),
6870 block_dim: (256, 1, 1),
6871 shared_mem_bytes: 0,
6872 };
6873 let mut dv = dst.slice_mut(dst_off..dst_off + words);
6874 let __s = self.gpu.stream();
6875 let mut b = __s.launch_builder(&f);
6876 b.arg(src_entry).arg(&mut dv).arg(&wi);
6877 unsafe {
6878 b.launch(cfg)?;
6879 }
6880 Ok(())
6881 }
6882
6883 #[track_caller]
6885 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
6886 self.alloc_uninit::<i8>(n)
6887 }
6888
6889 #[allow(clippy::too_many_arguments)] pub fn qmatvec(
6892 &self,
6893 w: &CudaSlice<u8>,
6894 x: &CudaSlice<f32>,
6895 m: usize,
6896 in_f: usize,
6897 out_f: usize,
6898 qtype: i32,
6899 row_bytes: usize,
6900 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6901 let f = self.func("qmatvec_f32");
6902 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
6904 grid_dim: (out_f as u32, m as u32, 1),
6905 block_dim: (256, 1, 1),
6906 shared_mem_bytes: 0,
6907 };
6908 let (inf, outf, mi, qt, rb) =
6909 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
6910 let __s_b = self.gpu.stream();
6911 let mut b = __s_b.launch_builder(&f);
6912 b.arg(w)
6913 .arg(x)
6914 .arg(&mut y)
6915 .arg(&inf)
6916 .arg(&outf)
6917 .arg(&mi)
6918 .arg(&qt)
6919 .arg(&rb);
6920 unsafe {
6921 b.launch(cfg)?;
6922 }
6923 Ok(y)
6924 }
6925
6926 #[track_caller]
6928 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6929 crate::alloc_trace_hit(n);
6930 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
6931 self.keep_if_capturing(&s);
6932 Ok(s)
6933 }
6934
6935 #[track_caller]
6939 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6940 crate::alloc_trace_hit(n);
6941 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
6942 self.keep_if_capturing(&s);
6943 Ok(s)
6944 }
6945
6946 pub fn memset_zeros_view(
6949 &self,
6950 dst: &mut cudarc::driver::CudaViewMut<f32>,
6951 ) -> Result<(), Box<dyn std::error::Error>> {
6952 self.gpu.stream().memset_zeros(dst)?;
6953 Ok(())
6954 }
6955
6956 pub fn stage_expert(
6962 &self,
6963 host_bytes: &[u8],
6964 scratch: &mut CudaSlice<u8>,
6965 off: usize,
6966 ) -> Result<(), Box<dyn std::error::Error>> {
6967 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
6970 }
6971
6972 pub fn nvfp4_expert_split_repack(
6976 &self,
6977 src: &CudaSlice<u8>,
6978 n_expert: usize,
6979 rows: usize,
6980 nsb64: usize,
6981 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6982 let need = n_expert * rows * nsb64 * 36;
6983 if src.len() < need {
6984 return Err(format!(
6985 "nvfp4_expert_split_repack: slab holds {} bytes, {n_expert} x {rows} x {nsb64} x 36 = {need} needed",
6986 src.len()
6987 )
6988 .into());
6989 }
6990 let f = self.func("nvfp4_expert_split_repack");
6991 let mut dst = self.alloc_u8_uninit(src.len())?; let nblk = (n_expert * rows * nsb64) as u32;
6993 let cfg = LaunchConfig {
6994 grid_dim: (nblk.div_ceil(256), 1, 1),
6995 block_dim: (256, 1, 1),
6996 shared_mem_bytes: 0,
6997 };
6998 let (ne, nr, ns) = (n_expert as i32, rows as i32, nsb64 as i32);
6999 let __s_b = self.gpu.stream();
7000 let mut b = __s_b.launch_builder(&f);
7001 b.arg(src).arg(&mut dst).arg(&ne).arg(&nr).arg(&ns);
7002 unsafe {
7003 b.launch(cfg)?;
7004 }
7005 Ok(dst)
7006 }
7007
7008 pub fn moe_router_topk(
7014 &self,
7015 logits: &CudaSlice<f32>,
7016 t: usize,
7017 n_expert: usize,
7018 n_used: usize,
7019 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7020 let f = self.func("moe_router_topk_f32");
7021 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 {
7024 grid_dim: (t as u32, 1, 1),
7025 block_dim: (n_expert as u32, 1, 1),
7026 shared_mem_bytes: 0,
7027 };
7028 let (ne, nu) = (n_expert as i32, n_used as i32);
7029 let __s_b = self.gpu.stream();
7030 let mut b = __s_b.launch_builder(&f);
7031 b.arg(logits)
7032 .arg(&mut sel_idx)
7033 .arg(&mut sel_w)
7034 .arg(&ne)
7035 .arg(&nu);
7036 unsafe {
7037 b.launch(cfg)?;
7038 }
7039 Ok((sel_idx, sel_w))
7040 }
7041
7042 pub fn moe_router_topk_scaled(
7045 &self,
7046 logits: &CudaSlice<f32>,
7047 t: usize,
7048 n_expert: usize,
7049 n_used: usize,
7050 ex_scale: &CudaSlice<f32>,
7051 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7052 let f = self.func("moe_router_topk_scaled_f32");
7057 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
7058 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
7059 let cfg = LaunchConfig {
7060 grid_dim: (t as u32, 1, 1),
7061 block_dim: (n_expert as u32, 1, 1),
7062 shared_mem_bytes: 0,
7063 };
7064 let (ne, nu) = (n_expert as i32, n_used as i32);
7065 let __s_b = self.gpu.stream();
7066 let mut b = __s_b.launch_builder(&f);
7067 b.arg(logits)
7068 .arg(&mut sel_idx)
7069 .arg(&mut sel_w)
7070 .arg(&ne)
7071 .arg(&nu)
7072 .arg(ex_scale);
7073 unsafe {
7074 b.launch(cfg)?;
7075 }
7076 Ok((sel_idx, sel_w))
7077 }
7078
7079 pub fn moe_router_topk_host(
7087 &self,
7088 logits: &CudaSlice<f32>,
7089 t: usize,
7090 n_expert: usize,
7091 n_used: usize,
7092 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
7093 let f = self.func("moe_router_topk_f32");
7094 let n = t * n_used;
7095 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
7096 let mut sel_w = self.alloc_uninit::<f32>(n)?;
7097 let cfg = LaunchConfig {
7098 grid_dim: (t as u32, 1, 1),
7099 block_dim: (n_expert as u32, 1, 1),
7100 shared_mem_bytes: 0,
7101 };
7102 let (ne, nu) = (n_expert as i32, n_used as i32);
7103 let __s_b = self.gpu.stream();
7104 let mut b = __s_b.launch_builder(&f);
7105 b.arg(logits)
7106 .arg(&mut sel_idx)
7107 .arg(&mut sel_w)
7108 .arg(&ne)
7109 .arg(&nu);
7110 unsafe {
7111 b.launch(cfg)?;
7112 }
7113 let bytes = n * 8;
7115 let mut guard = self.router_stage.lock().unwrap();
7116 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
7117 *guard = Some(PinnedStage::new(bytes.max(4096))?);
7118 }
7119 let stage = guard.as_mut().unwrap();
7120 let (si, sw) = unsafe {
7121 (
7122 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
7123 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
7124 )
7125 };
7126 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()))
7130 }
7131
7132 #[allow(clippy::too_many_arguments)]
7136 pub fn moe_router_sigmoid_topk(
7137 &self,
7138 logits: &CudaSlice<f32>,
7139 t: usize,
7140 n_expert: usize,
7141 n_used: usize,
7142 active_count: usize,
7143 correction_bias: &CudaSlice<f32>,
7144 active: &CudaSlice<u8>,
7145 scaling_factor: f32,
7146 route_norm: bool,
7147 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7148 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
7149 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
7150 return Err(format!(
7151 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
7152 )
7153 .into());
7154 }
7155 if logits.len() < t * n_expert
7156 || correction_bias.len() != n_expert
7157 || active.len() != n_expert
7158 {
7159 return Err(format!(
7160 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
7161 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
7162 ).into());
7163 }
7164 let f = self.func(crate::sigmoid_topk_kernel(
7165 crate::sig_expf_dev_on(),
7166 crate::topk_fast_on(),
7167 n_used,
7168 ));
7169 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
7170 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
7171 let threads = n_expert.div_ceil(32) * 32;
7172 let cfg = LaunchConfig {
7173 grid_dim: (t as u32, 1, 1),
7174 block_dim: (threads as u32, 1, 1),
7175 shared_mem_bytes: 0,
7176 };
7177 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
7178 let __s_b = self.gpu.stream();
7179 let mut b = __s_b.launch_builder(&f);
7180 b.arg(logits)
7181 .arg(correction_bias)
7182 .arg(active)
7183 .arg(&mut sel_idx)
7184 .arg(&mut sel_w)
7185 .arg(&ne)
7186 .arg(&nu)
7187 .arg(&scaling_factor)
7188 .arg(&rn);
7189 unsafe {
7190 b.launch(cfg)?;
7191 }
7192 Ok((sel_idx, sel_w))
7193 }
7194
7195 #[allow(clippy::too_many_arguments)]
7198 pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
7202 if ptr == 0 {
7203 return Err("ring_flag_raw: unarmed flag".into());
7204 }
7205 let f = self.func("memra_ring_flag");
7206 let cfg = LaunchConfig {
7207 grid_dim: (1, 1, 1),
7208 block_dim: (32, 1, 1),
7209 shared_mem_bytes: 0,
7210 };
7211 let __s_b = self.gpu.stream();
7212 let mut b = __s_b.launch_builder(&f);
7213 b.arg(&ptr).arg(&value);
7214 unsafe {
7215 b.launch(cfg)?;
7216 }
7217 Ok(())
7218 }
7219
7220 pub fn moe_sel_w_mirror(
7223 &self,
7224 sel_src: &CudaSlice<i32>,
7225 w_src: &CudaSlice<f32>,
7226 sel_dst: &mut CudaSlice<i32>,
7227 w_dst: &mut CudaSlice<f32>,
7228 n: usize,
7229 ) -> Result<(), Box<dyn std::error::Error>> {
7230 if n == 0
7231 || n > i32::MAX as usize
7232 || sel_src.len() < n
7233 || w_src.len() < n
7234 || sel_dst.len() < n
7235 || w_dst.len() < n
7236 {
7237 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
7238 }
7239 let f = self.func("moe_sel_w_mirror");
7240 let threads = if n <= 32 { 32 } else { 128 };
7241 let cfg = LaunchConfig {
7242 grid_dim: ((n as u32).div_ceil(threads), 1, 1),
7243 block_dim: (threads, 1, 1),
7244 shared_mem_bytes: 0,
7245 };
7246 let ni = n as i32;
7247 let __s_b = self.gpu.stream();
7248 let mut b = __s_b.launch_builder(&f);
7249 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
7250 unsafe {
7251 b.launch(cfg)?;
7252 }
7253 Ok(())
7254 }
7255
7256 #[allow(clippy::too_many_arguments)]
7260 pub fn nvfp4_ep_stage_inputs(
7261 &self,
7262 input_src: &CudaSlice<f32>,
7263 sel_src: &CudaSlice<i32>,
7264 w_src: &CudaSlice<f32>,
7265 input_bf16_dst: &mut CudaSlice<u8>,
7266 sel_dst: &mut CudaSlice<i32>,
7267 w_dst: &mut CudaSlice<f32>,
7268 input_values: usize,
7269 pairs: usize,
7270 copy_weights: bool,
7271 ) -> Result<(), Box<dyn std::error::Error>> {
7272 if input_values == 0
7273 || pairs == 0
7274 || input_src.len() < input_values
7275 || sel_src.len() < pairs
7276 || w_src.len() < pairs
7277 || input_bf16_dst.len() < 2 * input_values
7278 || sel_dst.len() < pairs
7279 || w_dst.len() < pairs
7280 {
7281 return Err(format!(
7282 "W4A16 EP stage geometry input={} sel={} weights={} input_bf16={} \
7283 sel_dst={} weights_dst={} active={input_values} pairs={pairs}",
7284 input_src.len(),
7285 sel_src.len(),
7286 w_src.len(),
7287 input_bf16_dst.len(),
7288 sel_dst.len(),
7289 w_dst.len(),
7290 )
7291 .into());
7292 }
7293 let f = self.func("nvfp4_ep_stage_inputs");
7294 let n = input_values.max(pairs);
7295 let cfg = LaunchConfig::for_num_elems(n as u32);
7296 let (input_values, pairs, copy_weights) =
7297 (input_values as i32, pairs as i32, i32::from(copy_weights));
7298 let __s_b = self.gpu.stream();
7299 let mut b = __s_b.launch_builder(&f);
7300 b.arg(input_src)
7301 .arg(sel_src)
7302 .arg(w_src)
7303 .arg(input_bf16_dst)
7304 .arg(sel_dst)
7305 .arg(w_dst)
7306 .arg(&input_values)
7307 .arg(&pairs)
7308 .arg(©_weights);
7309 unsafe {
7310 b.launch(cfg)?;
7311 }
7312 Ok(())
7313 }
7314
7315 #[allow(clippy::too_many_arguments)]
7318 pub fn nvfp4_ep_stage_inputs_raw(
7319 &self,
7320 input_src: u64,
7321 sel_src: u64,
7322 w_src: u64,
7323 input_bf16_dst: &mut CudaSlice<u8>,
7324 sel_dst: &mut CudaSlice<i32>,
7325 w_dst: &mut CudaSlice<f32>,
7326 input_values: usize,
7327 pairs: usize,
7328 copy_weights: bool,
7329 ) -> Result<(), Box<dyn std::error::Error>> {
7330 if input_src == 0
7331 || sel_src == 0
7332 || w_src == 0
7333 || input_values == 0
7334 || pairs == 0
7335 || input_bf16_dst.len() < 2 * input_values
7336 || sel_dst.len() < pairs
7337 || w_dst.len() < pairs
7338 {
7339 return Err(format!(
7340 "W4A16 EP raw stage geometry input={input_src:#x} sel={sel_src:#x} \
7341 weights={w_src:#x} input_bf16={} sel_dst={} weights_dst={} \
7342 active={input_values} pairs={pairs}",
7343 input_bf16_dst.len(),
7344 sel_dst.len(),
7345 w_dst.len(),
7346 )
7347 .into());
7348 }
7349 let f = self.func("nvfp4_ep_stage_inputs");
7350 let n = input_values.max(pairs);
7351 let cfg = LaunchConfig::for_num_elems(n as u32);
7352 let (input_values, pairs, copy_weights) =
7353 (input_values as i32, pairs as i32, i32::from(copy_weights));
7354 let __s_b = self.gpu.stream();
7355 let mut b = __s_b.launch_builder(&f);
7356 b.arg(&input_src)
7357 .arg(&sel_src)
7358 .arg(&w_src)
7359 .arg(input_bf16_dst)
7360 .arg(sel_dst)
7361 .arg(w_dst)
7362 .arg(&input_values)
7363 .arg(&pairs)
7364 .arg(©_weights);
7365 unsafe {
7366 b.launch(cfg)?;
7367 }
7368 Ok(())
7369 }
7370
7371 #[allow(clippy::too_many_arguments)] pub fn moe_router_sigmoid_topk_into(
7373 &self,
7374 logits: &CudaSlice<f32>,
7375 t: usize,
7376 n_expert: usize,
7377 n_used: usize,
7378 active_count: usize,
7379 correction_bias: &CudaSlice<f32>,
7380 active: &CudaSlice<u8>,
7381 scaling_factor: f32,
7382 route_norm: bool,
7383 sel_idx: &mut CudaSlice<i32>,
7384 sel_w: &mut CudaSlice<f32>,
7385 ) -> Result<(), Box<dyn std::error::Error>> {
7386 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
7387 if n_expert == 0
7388 || n_expert > 1024
7389 || n_used == 0
7390 || n_used > 32 || n_used > n_expert
7392 || logits.len() < t * n_expert
7393 || correction_bias.len() != n_expert
7394 || active.len() != n_expert
7395 || sel_idx.len() < t * n_used
7396 || sel_w.len() < t * n_used
7397 {
7398 return Err("sigmoid router _into geometry mismatch".into());
7399 }
7400 let f = self.func(crate::sigmoid_topk_kernel(
7401 crate::sig_expf_dev_on(),
7402 crate::topk_fast_on(),
7403 n_used,
7404 ));
7405 let threads = n_expert.div_ceil(32) * 32;
7406 let cfg = LaunchConfig {
7407 grid_dim: (t as u32, 1, 1),
7408 block_dim: (threads as u32, 1, 1),
7409 shared_mem_bytes: 0,
7410 };
7411 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
7412 let __s_b = self.gpu.stream();
7413 let mut b = __s_b.launch_builder(&f);
7414 b.arg(logits)
7415 .arg(correction_bias)
7416 .arg(active)
7417 .arg(&mut *sel_idx)
7418 .arg(&mut *sel_w)
7419 .arg(&ne)
7420 .arg(&nu)
7421 .arg(&scaling_factor)
7422 .arg(&rn);
7423 unsafe {
7424 b.launch(cfg)?;
7425 }
7426 Ok(())
7427 }
7428
7429 #[allow(clippy::too_many_arguments)]
7432 pub fn moe_router_sigmoid_topk_host(
7433 &self,
7434 logits: &CudaSlice<f32>,
7435 t: usize,
7436 n_expert: usize,
7437 n_used: usize,
7438 active_count: usize,
7439 correction_bias: &CudaSlice<f32>,
7440 active: &CudaSlice<u8>,
7441 scaling_factor: f32,
7442 route_norm: bool,
7443 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
7444 if glm5_graph_capture_open() {
7458 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"
7459 .into());
7460 }
7461 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
7462 logits,
7463 t,
7464 n_expert,
7465 n_used,
7466 active_count,
7467 correction_bias,
7468 active,
7469 scaling_factor,
7470 route_norm,
7471 )?;
7472 let n = t * n_used;
7473 let bytes = n * 8;
7474 let mut guard = self.router_stage.lock().unwrap();
7475 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
7476 *guard = Some(PinnedStage::new(bytes.max(4096))?);
7477 }
7478 let stage = guard.as_mut().unwrap();
7479 let (si, sw) = unsafe {
7480 (
7481 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
7482 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
7483 )
7484 };
7485 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
7486 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
7487 self.gpu.stream().synchronize()?;
7488 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
7489 }
7490
7491 pub fn stage_expert_async(
7495 &self,
7496 host_bytes: &[u8],
7497 scratch: &mut CudaSlice<u8>,
7498 off: usize,
7499 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
7500 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
7501 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
7502 Ok(self.copy_stream.record_event(None)?)
7503 }
7504
7505 pub fn compute_wait(
7507 &self,
7508 ev: &cudarc::driver::CudaEvent,
7509 ) -> Result<(), Box<dyn std::error::Error>> {
7510 self.gpu.stream().wait(ev)?;
7511 Ok(())
7512 }
7513
7514 #[allow(clippy::too_many_arguments)] pub fn qmatvec_view(
7520 &self,
7521 w: &CudaSlice<u8>,
7522 range: std::ops::Range<usize>,
7523 x: &cudarc::driver::CudaView<f32>,
7524 m: usize,
7525 in_f: usize,
7526 out_f: usize,
7527 qtype: i32,
7528 row_bytes: usize,
7529 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7530 self.qmatvec_view_inner(w, range, x, m, in_f, out_f, qtype, row_bytes)
7531 }
7532
7533 #[allow(clippy::too_many_arguments)]
7537 pub fn qmatvec_view_bf16_activation(
7538 &self,
7539 w: &CudaSlice<u8>,
7540 range: std::ops::Range<usize>,
7541 x: &cudarc::driver::CudaView<f32>,
7542 m: usize,
7543 in_f: usize,
7544 out_f: usize,
7545 qtype: i32,
7546 row_bytes: usize,
7547 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7548 let n = m * in_f;
7549 if x.len() != n {
7550 return Err(format!(
7551 "W4A16 BF16 activation input length {} != {m}x{in_f}",
7552 x.len()
7553 )
7554 .into());
7555 }
7556 let mut x_bf16 = self.alloc_u8_uninit(n * 2)?;
7557 self.f32_to_bf16_v(x, &mut x_bf16, n)?;
7558 let x_f32 = self.bf16_to_f32(&x_bf16.slice(0..n * 2), n)?;
7559 self.qmatvec_view_inner(
7560 w,
7561 range,
7562 &x_f32.slice(0..n),
7563 m,
7564 in_f,
7565 out_f,
7566 qtype,
7567 row_bytes,
7568 )
7569 }
7570
7571 #[allow(clippy::too_many_arguments)]
7572 fn qmatvec_view_inner(
7573 &self,
7574 w: &CudaSlice<u8>,
7575 range: std::ops::Range<usize>,
7576 x: &cudarc::driver::CudaView<f32>,
7577 m: usize,
7578 in_f: usize,
7579 out_f: usize,
7580 qtype: i32,
7581 row_bytes: usize,
7582 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7583 let f = self.func("qmatvec_f32");
7584 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7587 grid_dim: (out_f as u32, m as u32, 1),
7588 block_dim: (256, 1, 1),
7589 shared_mem_bytes: 0,
7590 };
7591 let (inf, outf, mi, qt, rb) =
7592 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
7593 let __s_b = self.gpu.stream();
7594 let mut b = __s_b.launch_builder(&f);
7595 b.arg(&wv)
7596 .arg(x)
7597 .arg(&mut y)
7598 .arg(&inf)
7599 .arg(&outf)
7600 .arg(&mi)
7601 .arg(&qt)
7602 .arg(&rb);
7603 unsafe {
7604 b.launch(cfg)?;
7605 }
7606 Ok(y)
7607 }
7608
7609 #[allow(clippy::too_many_arguments)]
7616 pub fn moe_gate_up_silu8_q8(
7620 &self,
7621 gp: WPtr8,
7622 up: WPtr8,
7623 aq: &CudaSlice<i8>,
7624 ad: &CudaSlice<f32>,
7625 in_f: usize,
7626 n_ff: usize,
7627 n_used: usize,
7628 qt_g: i32,
7629 qt_u: i32,
7630 rb_g: usize,
7631 rb_u: usize,
7632 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7633 let f = self.func("moe_gate_up_silu8_q8");
7634 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7635 let cfg = LaunchConfig {
7636 grid_dim: (n_ff as u32, n_used as u32, 1),
7637 block_dim: (32, 1, 1),
7638 shared_mem_bytes: 0,
7639 };
7640 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7641 let __s_b = self.gpu.stream();
7642 let mut b = __s_b.launch_builder(&f);
7643 b.arg(&gp)
7644 .arg(&up)
7645 .arg(aq)
7646 .arg(ad)
7647 .arg(&mut act)
7648 .arg(&inf)
7649 .arg(&nff)
7650 .arg(&qt_g)
7651 .arg(&qt_u)
7652 .arg(&rbg)
7653 .arg(&rbu);
7654 unsafe {
7655 b.launch(cfg)?;
7656 }
7657 Ok(act)
7658 }
7659
7660 #[allow(clippy::too_many_arguments)]
7673 pub fn moe_gate_up_preclamp8_q8(
7674 &self,
7675 gp: WPtr8,
7676 up: WPtr8,
7677 aq: &CudaSlice<i8>,
7678 ad: &CudaSlice<f32>,
7679 gs: F32x8,
7680 us: F32x8,
7681 limit: f32,
7682 in_f: usize,
7683 n_ff: usize,
7684 n_used: usize,
7685 qt_g: i32,
7686 qt_u: i32,
7687 rb_g: usize,
7688 rb_u: usize,
7689 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7690 debug_assert!(
7691 limit > 1e-6,
7692 "moe_gate_up_preclamp8_q8 needs a live limit; use moe_gate_up_silu8_q8"
7693 );
7694 let f = self.func("moe_gate_up_preclamp8_q8");
7695 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7696 let cfg = LaunchConfig {
7697 grid_dim: (n_ff as u32, n_used as u32, 1),
7698 block_dim: (32, 1, 1),
7699 shared_mem_bytes: 0,
7700 };
7701 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7702 let __s_b = self.gpu.stream();
7703 let mut b = __s_b.launch_builder(&f);
7704 b.arg(&gp)
7705 .arg(&up)
7706 .arg(aq)
7707 .arg(ad)
7708 .arg(&gs)
7709 .arg(&us)
7710 .arg(&limit)
7711 .arg(&mut act)
7712 .arg(&inf)
7713 .arg(&nff)
7714 .arg(&qt_g)
7715 .arg(&qt_u)
7716 .arg(&rbg)
7717 .arg(&rbu);
7718 unsafe {
7719 b.launch(cfg)?;
7720 }
7721 Ok(act)
7722 }
7723
7724 #[allow(clippy::too_many_arguments)]
7725 pub fn moe_down8_fma_q8(
7726 &self,
7727 dp: WPtr8,
7728 w: F32x8,
7729 aq2: &CudaSlice<i8>,
7730 ad2: &CudaSlice<f32>,
7731 dst: &mut cudarc::driver::CudaViewMut<f32>,
7732 in_f: usize,
7733 out_f: usize,
7734 n_used: usize,
7735 qt: i32,
7736 rb: usize,
7737 ) -> Result<(), Box<dyn std::error::Error>> {
7738 let f = self.func("moe_down8_fma_q8");
7739 let cfg = LaunchConfig {
7740 grid_dim: (out_f as u32, 1, 1),
7741 block_dim: (32, 1, 1),
7742 shared_mem_bytes: 0,
7743 };
7744 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
7745 let __s_b = self.gpu.stream();
7746 let mut b = __s_b.launch_builder(&f);
7747 b.arg(&dp)
7748 .arg(&w)
7749 .arg(aq2)
7750 .arg(ad2)
7751 .arg(dst)
7752 .arg(&inf)
7753 .arg(&outf)
7754 .arg(&nu)
7755 .arg(&qt)
7756 .arg(&rbi);
7757 unsafe {
7758 b.launch(cfg)?;
7759 }
7760 Ok(())
7761 }
7762
7763 #[allow(clippy::too_many_arguments)]
7768 pub fn moe_gate_up_preclamp8_q8_w4(
7769 &self,
7770 gp: WPtr8,
7771 up: WPtr8,
7772 aq: &CudaSlice<i8>,
7773 ad: &CudaSlice<f32>,
7774 gs: F32x8,
7775 us: F32x8,
7776 limit: f32,
7777 in_f: usize,
7778 n_ff: usize,
7779 n_used: usize,
7780 qt_g: i32,
7781 qt_u: i32,
7782 rb_g: usize,
7783 rb_u: usize,
7784 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7785 debug_assert!(
7786 limit > 1e-6,
7787 "moe_gate_up_preclamp8_q8_w4 needs a live limit; use moe_gate_up_silu8_q8"
7788 );
7789 const ROWS: u32 = 4; let f = self.func("moe_gate_up_preclamp8_q8_w4");
7791 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7792 let cfg = LaunchConfig {
7793 grid_dim: ((n_ff as u32).div_ceil(ROWS), n_used as u32, 1),
7794 block_dim: (32, ROWS, 1),
7795 shared_mem_bytes: 0,
7796 };
7797 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
7798 let __s_b = self.gpu.stream();
7799 let mut b = __s_b.launch_builder(&f);
7800 b.arg(&gp)
7801 .arg(&up)
7802 .arg(aq)
7803 .arg(ad)
7804 .arg(&gs)
7805 .arg(&us)
7806 .arg(&limit)
7807 .arg(&mut act)
7808 .arg(&inf)
7809 .arg(&nff)
7810 .arg(&qt_g)
7811 .arg(&qt_u)
7812 .arg(&rbg)
7813 .arg(&rbu);
7814 unsafe {
7815 b.launch(cfg)?;
7816 }
7817 Ok(act)
7818 }
7819
7820 #[allow(clippy::too_many_arguments)]
7823 pub fn moe_down8_fma_q8_w4(
7824 &self,
7825 dp: WPtr8,
7826 w: F32x8,
7827 aq2: &CudaSlice<i8>,
7828 ad2: &CudaSlice<f32>,
7829 dst: &mut cudarc::driver::CudaViewMut<f32>,
7830 in_f: usize,
7831 out_f: usize,
7832 n_used: usize,
7833 qt: i32,
7834 rb: usize,
7835 ) -> Result<(), Box<dyn std::error::Error>> {
7836 const ROWS: u32 = 4; let f = self.func("moe_down8_fma_q8_w4");
7838 let cfg = LaunchConfig {
7839 grid_dim: ((out_f as u32).div_ceil(ROWS), 1, 1),
7840 block_dim: (32, ROWS, 1),
7841 shared_mem_bytes: 0,
7842 };
7843 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
7844 let __s_b = self.gpu.stream();
7845 let mut b = __s_b.launch_builder(&f);
7846 b.arg(&dp)
7847 .arg(&w)
7848 .arg(aq2)
7849 .arg(ad2)
7850 .arg(dst)
7851 .arg(&inf)
7852 .arg(&outf)
7853 .arg(&nu)
7854 .arg(&qt)
7855 .arg(&rbi);
7856 unsafe {
7857 b.launch(cfg)?;
7858 }
7859 Ok(())
7860 }
7861
7862 #[allow(clippy::too_many_arguments)]
7874 pub fn moe_vrows_tables_from_sel(
7876 &self,
7877 sel: &CudaSlice<i32>,
7878 selw: &CudaSlice<f32>,
7879 il: u16,
7880 macros: Option<(&[f32], &[f32], &[f32])>,
7881 (pg, pu, pd): (u64, u64, u64),
7882 (sg, su, sd): (usize, usize, usize),
7883 n_pairs: usize,
7884 ptrs: &mut CudaSlice<u64>,
7885 scl: &mut CudaSlice<f32>,
7886 ) -> Result<(), Box<dyn std::error::Error>> {
7887 debug_assert!(sel.len() >= n_pairs && selw.len() >= n_pairs);
7888 debug_assert!(ptrs.len() >= 3 * n_pairs);
7890 debug_assert_eq!(scl.len(), 3 * n_pairs);
7891 let mut mac = self
7894 .vrows_macro_dev
7895 .lock()
7896 .map_err(|_| "vrows macro mirror map is poisoned")?;
7897 if let Some((hg, hu, hd)) = macros {
7898 for (plane, host) in [(0u8, hg), (1u8, hu), (2u8, hd)] {
7899 if let std::collections::hash_map::Entry::Vacant(slot) = mac.entry((il, plane)) {
7902 slot.insert(self.htod(host)?);
7903 }
7904 }
7905 }
7906 let (mg, mu, md, have) = match macros {
7909 Some(_) => (
7910 mac.get(&(il, 0)).expect("gate macro mirror built above"),
7911 mac.get(&(il, 1)).expect("up macro mirror built above"),
7912 mac.get(&(il, 2)).expect("down macro mirror built above"),
7913 1i32,
7914 ),
7915 None => (selw, selw, selw, 0i32),
7916 };
7917 let f = self.func("moe_vrows_tables_from_sel");
7918 let threads = 128u32;
7919 let cfg = LaunchConfig {
7920 grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
7921 block_dim: (threads, 1, 1),
7922 shared_mem_bytes: 0,
7923 };
7924 let (sgi, sui, sdi) = (sg as i64, su as i64, sd as i64);
7925 let (npi, havei) = (n_pairs as i32, have);
7926 let __s_b = self.gpu.stream();
7927 let mut b = __s_b.launch_builder(&f);
7928 b.arg(sel)
7929 .arg(selw)
7930 .arg(mg)
7931 .arg(mu)
7932 .arg(md)
7933 .arg(&mut *ptrs)
7934 .arg(&mut *scl)
7935 .arg(&pg)
7936 .arg(&pu)
7937 .arg(&pd)
7938 .arg(&sgi)
7939 .arg(&sui)
7940 .arg(&sdi)
7941 .arg(&npi)
7942 .arg(&havei);
7943 unsafe {
7944 b.launch(cfg)?;
7945 }
7946 Ok(())
7947 }
7948
7949 pub fn moe_vrows_order_from_sel(
7961 &self,
7962 sel: &CudaSlice<i32>,
7963 n_pairs: usize,
7964 ptrs: &mut CudaSlice<u64>,
7965 ) -> Result<(), Box<dyn std::error::Error>> {
7966 debug_assert!(sel.len() >= n_pairs);
7967 debug_assert!(
7968 ptrs.len() >= 4 * n_pairs,
7969 "the order plane lives at ptrs[3*n_pairs .. 4*n_pairs)"
7970 );
7971 let f = self.func("moe_vrows_order_from_sel");
7972 let threads = 128u32;
7973 let cfg = LaunchConfig {
7974 grid_dim: ((n_pairs as u32).div_ceil(threads), 1, 1),
7975 block_dim: (threads, 1, 1),
7976 shared_mem_bytes: 0,
7977 };
7978 let np = n_pairs as i32;
7979 let __s_b = self.gpu.stream();
7980 let mut b = __s_b.launch_builder(&f);
7981 b.arg(sel).arg(&mut *ptrs).arg(&np);
7982 unsafe {
7983 b.launch(cfg)?;
7984 }
7985 Ok(())
7986 }
7987
7988 #[allow(clippy::too_many_arguments)]
7994 pub fn moe_gate_up_preclamp8_q8_rows(
7996 &self,
7997 ptrs: &CudaSlice<u64>,
7998 scl: &CudaSlice<f32>,
7999 aq: &CudaSlice<i8>,
8000 ad: &CudaSlice<f32>,
8001 limit: f32,
8002 in_f: usize,
8003 n_ff: usize,
8004 n_used: usize,
8005 n_pairs: usize,
8006 qt_g: i32,
8007 qt_u: i32,
8008 rb_g: usize,
8009 rb_u: usize,
8010 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8011 debug_assert!(
8012 limit > 1e-6,
8013 "moe_gate_up_preclamp8_q8_rows needs a live limit; the kernel collapses every gate \
8014 to silu(0) at limit 0"
8015 );
8016 debug_assert!(ptrs.len() >= 3 * n_pairs);
8017 debug_assert_eq!(scl.len(), 3 * n_pairs);
8018 let packed = moe_vrows_pack_on();
8026 let ordered =
8027 !packed && moe_vrows_dedup_order_on() && ptrs.len() >= 4 * n_pairs && n_ff <= 65535;
8028 let ilp = moe_vrows_ilp_on()
8031 && if (qt_g == QT_NVFP4 && qt_u == QT_NVFP4)
8032 || (qt_g == QT_NVFP4_V2 && qt_u == QT_NVFP4_V2)
8033 {
8034 true
8035 } else {
8036 moe_vrows_ilp_refuse("gate/up", if qt_g == QT_NVFP4 { qt_u } else { qt_g });
8037 false
8038 };
8039 if ilp && MOE_VROWS_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
8040 eprintln!(
8041 "[moe-vrows-ilp] engaged: verify-rows MoE pair with four groups' loads per lane \
8042 hoisted ahead of their math (MEMRA_MOE_VROWS_ILP=1, packed={packed})"
8043 );
8044 }
8045 let (f, cfg) = if packed {
8046 if MOE_VROWS_PACK_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
8047 eprintln!(
8048 "[moe-vrows-pack] engaged: 4-warp blocks on the verify-rows MoE pair \
8049 (MEMRA_MOE_VROWS_PACK=1)"
8050 );
8051 }
8052 (
8053 self.func(if ilp {
8054 "moe_gate_up_preclamp8_q8_rows_w4_ilp"
8055 } else {
8056 "moe_gate_up_preclamp8_q8_rows_w4"
8057 }),
8058 LaunchConfig {
8059 grid_dim: ((n_ff as u32).div_ceil(4), n_pairs as u32, 1),
8060 block_dim: (32, 4, 1),
8061 shared_mem_bytes: 0,
8062 },
8063 )
8064 } else if ordered {
8065 if MOE_VROWS_DEDUP_ORDER_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8066 == 0
8067 {
8068 eprintln!(
8069 "[moe-vrows-dedup-order] engaged: verify-rows gate/up walks the pair union \
8070 EXPERT-MAJOR with the pair index as the fastest grid dimension, so the \
8071 21.96%-measured repeat visits read a shared expert slab's rows in adjacent \
8072 blocks (MEMRA_MOE_VROWS_DEDUP_ORDER=1)"
8073 );
8074 }
8075 (
8076 self.func("moe_gate_up_preclamp8_q8_rows_ord"),
8077 LaunchConfig {
8078 grid_dim: (n_pairs as u32, n_ff as u32, 1),
8079 block_dim: (32, 1, 1),
8080 shared_mem_bytes: 0,
8081 },
8082 )
8083 } else {
8084 (
8085 self.func(if ilp {
8086 "moe_gate_up_preclamp8_q8_rows_ilp"
8087 } else {
8088 "moe_gate_up_preclamp8_q8_rows"
8089 }),
8090 LaunchConfig {
8091 grid_dim: (n_ff as u32, n_pairs as u32, 1),
8092 block_dim: (32, 1, 1),
8093 shared_mem_bytes: 0,
8094 },
8095 )
8096 };
8097 let mut act = self.vws_uninit(n_pairs * n_ff)?;
8099 let (inf, nff, nu, np) = (in_f as i32, n_ff as i32, n_used as i32, n_pairs as i32);
8100 let (rbg, rbu) = (rb_g as i64, rb_u as i64);
8101 let __s_b = self.gpu.stream();
8102 let mut b = __s_b.launch_builder(&f);
8103 b.arg(ptrs)
8104 .arg(scl)
8105 .arg(aq)
8106 .arg(ad)
8107 .arg(&limit)
8108 .arg(&mut act)
8109 .arg(&inf)
8110 .arg(&nff)
8111 .arg(&nu)
8112 .arg(&np)
8113 .arg(&qt_g)
8114 .arg(&qt_u)
8115 .arg(&rbg)
8116 .arg(&rbu);
8117 unsafe {
8118 b.launch(cfg)?;
8119 }
8120 Ok(act)
8121 }
8122
8123 #[allow(clippy::too_many_arguments)]
8127 pub fn moe_down8_fma_q8_rows(
8129 &self,
8130 ptrs: &CudaSlice<u64>,
8131 scl: &CudaSlice<f32>,
8132 aq2: &CudaSlice<i8>,
8133 ad2: &CudaSlice<f32>,
8134 dst: &mut CudaSlice<f32>,
8135 in_f: usize,
8136 out_f: usize,
8137 n_used: usize,
8138 n_pairs: usize,
8139 qt: i32,
8140 rb: usize,
8141 ) -> Result<(), Box<dyn std::error::Error>> {
8142 debug_assert!(ptrs.len() >= 3 * n_pairs);
8143 debug_assert_eq!(scl.len(), 3 * n_pairs);
8144 debug_assert_eq!(n_pairs % n_used, 0, "pairs are dense slot-major");
8145 let t = n_pairs / n_used;
8146 debug_assert!(dst.len() >= t * out_f);
8147 let packed = moe_vrows_pack_on();
8149 let tmaj = !packed && moe_vrows_down_tmaj_on() && out_f <= 65535;
8155 let ilp = moe_vrows_ilp_on()
8157 && if qt == QT_NVFP4 || qt == QT_NVFP4_V2 {
8158 true
8159 } else {
8160 moe_vrows_ilp_refuse("down", qt);
8161 false
8162 };
8163 if ilp {
8164 MOE_VROWS_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8165 }
8166 let (f, cfg) = if packed {
8167 (
8168 self.func(if ilp {
8169 "moe_down8_fma_q8_rows_w4_ilp"
8170 } else {
8171 "moe_down8_fma_q8_rows_w4"
8172 }),
8173 LaunchConfig {
8174 grid_dim: ((out_f as u32).div_ceil(4), t as u32, 1),
8175 block_dim: (32, 4, 1),
8176 shared_mem_bytes: 0,
8177 },
8178 )
8179 } else if tmaj {
8180 if MOE_VROWS_DOWN_TMAJ_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8181 == 0
8182 {
8183 eprintln!(
8184 "[moe-vrows-down-tmaj] engaged: verify-rows down/FMA grid transposed to \
8185 (t, out_f) so the verify rows at one output row are adjacent blocks; the \
8186 slot-ordered FMA chain is unchanged (MEMRA_MOE_VROWS_DOWN_TMAJ=1)"
8187 );
8188 }
8189 (
8190 self.func("moe_down8_fma_q8_rows_tmaj"),
8191 LaunchConfig {
8192 grid_dim: (t as u32, out_f as u32, 1),
8193 block_dim: (32, 1, 1),
8194 shared_mem_bytes: 0,
8195 },
8196 )
8197 } else {
8198 (
8199 self.func(if ilp {
8200 "moe_down8_fma_q8_rows_ilp"
8201 } else {
8202 "moe_down8_fma_q8_rows"
8203 }),
8204 LaunchConfig {
8205 grid_dim: (out_f as u32, t as u32, 1),
8206 block_dim: (32, 1, 1),
8207 shared_mem_bytes: 0,
8208 },
8209 )
8210 };
8211 let (inf, outf, nu, np, rbi) = (
8212 in_f as i32,
8213 out_f as i32,
8214 n_used as i32,
8215 n_pairs as i32,
8216 rb as i64,
8217 );
8218 let __s_b = self.gpu.stream();
8219 let mut b = __s_b.launch_builder(&f);
8220 b.arg(ptrs)
8221 .arg(scl)
8222 .arg(aq2)
8223 .arg(ad2)
8224 .arg(dst)
8225 .arg(&inf)
8226 .arg(&outf)
8227 .arg(&nu)
8228 .arg(&np)
8229 .arg(&qt)
8230 .arg(&rbi);
8231 unsafe {
8232 b.launch(cfg)?;
8233 }
8234 Ok(())
8235 }
8236
8237 #[allow(clippy::too_many_arguments)]
8239 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_expert_q8(
8242 &self,
8243 w: &CudaSlice<u8>,
8244 range: std::ops::Range<usize>,
8245 aq: &CudaSlice<i8>,
8246 ad: &CudaSlice<f32>,
8247 m: usize,
8248 in_f: usize,
8249 out_f: usize,
8250 qtype: i32,
8251 row_bytes: usize,
8252 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8253 let f = self.func("qmatvec_expert_q8");
8254 let wv = w.slice(range);
8255 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
8256 const ROWS: u32 = 4; let cfg = LaunchConfig {
8258 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
8259 block_dim: (32, ROWS, 1),
8260 shared_mem_bytes: 0,
8261 };
8262 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8263 let __s_b = self.gpu.stream();
8264 let mut b = __s_b.launch_builder(&f);
8265 b.arg(&wv)
8266 .arg(aq)
8267 .arg(ad)
8268 .arg(&mut y)
8269 .arg(&inf)
8270 .arg(&outf)
8271 .arg(&mi)
8272 .arg(&qtype)
8273 .arg(&rbi);
8274 unsafe {
8275 b.launch(cfg)?;
8276 }
8277 Ok(y)
8278 }
8279
8280 #[allow(clippy::too_many_arguments)] pub fn moe_gate_up_silu8(
8282 &self,
8283 gp: WPtr8,
8284 up: WPtr8,
8285 x: &cudarc::driver::CudaView<f32>,
8286 in_f: usize,
8287 n_ff: usize,
8288 n_used: usize,
8289 qt_g: i32,
8290 qt_u: i32,
8291 rb_g: usize,
8292 rb_u: usize,
8293 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8294 let f = self.func("moe_gate_up_silu8_f32");
8295 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
8297 grid_dim: (n_ff as u32, n_used as u32, 1),
8298 block_dim: (256, 1, 1),
8299 shared_mem_bytes: 0,
8300 };
8301 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
8302 let __s_b = self.gpu.stream();
8303 let mut b = __s_b.launch_builder(&f);
8304 b.arg(&gp)
8305 .arg(&up)
8306 .arg(x)
8307 .arg(&mut act)
8308 .arg(&inf)
8309 .arg(&nff)
8310 .arg(&qt_g)
8311 .arg(&qt_u)
8312 .arg(&rbg)
8313 .arg(&rbu);
8314 unsafe {
8315 b.launch(cfg)?;
8316 }
8317 Ok(act)
8318 }
8319
8320 #[allow(clippy::too_many_arguments)]
8326 pub fn moe_down8_fma_into(
8327 &self,
8328 dp: WPtr8,
8329 w: F32x8,
8330 act: &CudaSlice<f32>,
8331 dst: &mut cudarc::driver::CudaViewMut<f32>,
8332 in_f: usize,
8333 out_f: usize,
8334 n_used: usize,
8335 qt: i32,
8336 rb: usize,
8337 ) -> Result<(), Box<dyn std::error::Error>> {
8338 let f = self.func("moe_down8_fma_f32");
8339 let cfg = LaunchConfig {
8340 grid_dim: (out_f as u32, 1, 1),
8341 block_dim: (256, 1, 1),
8342 shared_mem_bytes: 0,
8343 };
8344 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
8345 let __s_b = self.gpu.stream();
8346 let mut b = __s_b.launch_builder(&f);
8347 b.arg(&dp)
8348 .arg(&w)
8349 .arg(act)
8350 .arg(dst)
8351 .arg(&inf)
8352 .arg(&outf)
8353 .arg(&nu)
8354 .arg(&qt)
8355 .arg(&rbv);
8356 unsafe {
8357 b.launch(cfg)?;
8358 }
8359 Ok(())
8360 }
8361
8362 #[allow(clippy::too_many_arguments)]
8367 #[allow(clippy::too_many_arguments)]
8382 #[allow(clippy::too_many_arguments)]
8384 #[allow(clippy::manual_div_ceil)] pub fn moe_pairs_matvec_q8(
8386 &self,
8387 table: &CudaSlice<u64>,
8388 proj: i32,
8389 pair_tok: &CudaSlice<i32>,
8390 pair_ex: &CudaSlice<i32>,
8391 aq: &CudaSlice<i8>,
8392 ad: &CudaSlice<f32>,
8393 in_f: usize,
8394 out_f: usize,
8395 n_expert: usize,
8396 n_pairs: usize,
8397 qtype: i32,
8398 row_bytes: usize,
8399 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8400 let f = self.func("moe_pairs_matvec_q8");
8401 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8402 const ROWS: u32 = 4;
8403 let cfg = LaunchConfig {
8404 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
8405 block_dim: (32, ROWS, 1),
8406 shared_mem_bytes: 0,
8407 };
8408 let (inf, outf, ne, np, rbi) = (
8409 in_f as i32,
8410 out_f as i32,
8411 n_expert as i32,
8412 n_pairs as i32,
8413 row_bytes as i64,
8414 );
8415 let __s_b = self.gpu.stream();
8416 let mut b = __s_b.launch_builder(&f);
8417 b.arg(table)
8418 .arg(&proj)
8419 .arg(pair_tok)
8420 .arg(pair_ex)
8421 .arg(aq)
8422 .arg(ad)
8423 .arg(&mut y)
8424 .arg(&inf)
8425 .arg(&outf)
8426 .arg(&ne)
8427 .arg(&np)
8428 .arg(&qtype)
8429 .arg(&rbi);
8430 unsafe {
8431 b.launch(cfg)?;
8432 }
8433 Ok(y)
8434 }
8435
8436 #[allow(clippy::too_many_arguments)]
8438 #[allow(clippy::manual_div_ceil)] pub fn moe_pairs_matvec_q8_em(
8440 &self,
8441 table: &CudaSlice<u64>,
8442 proj: i32,
8443 ex_ids: &CudaSlice<i32>,
8444 ex_off: &CudaSlice<i32>,
8445 ex_pairs: &CudaSlice<i32>,
8446 pair_tok: &CudaSlice<i32>,
8447 aq: &CudaSlice<i8>,
8448 ad: &CudaSlice<f32>,
8449 in_f: usize,
8450 out_f: usize,
8451 n_expert: usize,
8452 n_active: usize,
8453 n_pairs: usize,
8454 qtype: i32,
8455 row_bytes: usize,
8456 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8457 let f = self.func("moe_pairs_matvec_q8_em");
8458 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8459 const ROWS: u32 = 4;
8460 let cfg = LaunchConfig {
8461 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
8462 block_dim: (32, ROWS, 1),
8463 shared_mem_bytes: 0,
8464 };
8465 let (inf, outf, ne, na, rbi) = (
8466 in_f as i32,
8467 out_f as i32,
8468 n_expert as i32,
8469 n_active as i32,
8470 row_bytes as i64,
8471 );
8472 let __s_b = self.gpu.stream();
8473 let mut b = __s_b.launch_builder(&f);
8474 b.arg(table)
8475 .arg(&proj)
8476 .arg(ex_ids)
8477 .arg(ex_off)
8478 .arg(ex_pairs)
8479 .arg(pair_tok)
8480 .arg(aq)
8481 .arg(ad)
8482 .arg(&mut y)
8483 .arg(&inf)
8484 .arg(&outf)
8485 .arg(&ne)
8486 .arg(&na)
8487 .arg(&qtype)
8488 .arg(&rbi);
8489 unsafe {
8490 b.launch(cfg)?;
8491 }
8492 Ok(y)
8493 }
8494
8495 #[allow(clippy::too_many_arguments)]
8498 #[allow(clippy::manual_div_ceil)] pub fn moe_pairs_matvec_q8_dec(
8500 &self,
8501 table: &CudaSlice<u64>,
8502 proj: i32,
8503 ex_ids: &CudaSlice<i32>,
8504 ex_off: &CudaSlice<i32>,
8505 ex_pairs: &CudaSlice<i32>,
8506 pair_tok: &CudaSlice<i32>,
8507 aq: &CudaSlice<i8>,
8508 ad: &CudaSlice<f32>,
8509 in_f: usize,
8510 out_f: usize,
8511 n_expert: usize,
8512 n_active: usize,
8513 n_pairs: usize,
8514 qtype: i32,
8515 row_bytes: usize,
8516 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8517 let f = self.func("moe_pairs_matvec_q8_dec");
8518 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
8519 const ROWS: u32 = 4;
8520 let cfg = LaunchConfig {
8521 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
8522 block_dim: (32, ROWS, 1),
8523 shared_mem_bytes: 0,
8524 };
8525 let (inf, outf, ne, na, rbi) = (
8526 in_f as i32,
8527 out_f as i32,
8528 n_expert as i32,
8529 n_active as i32,
8530 row_bytes as i64,
8531 );
8532 let __s_b = self.gpu.stream();
8533 let mut b = __s_b.launch_builder(&f);
8534 b.arg(table)
8535 .arg(&proj)
8536 .arg(ex_ids)
8537 .arg(ex_off)
8538 .arg(ex_pairs)
8539 .arg(pair_tok)
8540 .arg(aq)
8541 .arg(ad)
8542 .arg(&mut y)
8543 .arg(&inf)
8544 .arg(&outf)
8545 .arg(&ne)
8546 .arg(&na)
8547 .arg(&qtype)
8548 .arg(&rbi);
8549 unsafe {
8550 b.launch(cfg)?;
8551 }
8552 Ok(y)
8553 }
8554
8555 pub fn moe_pairs_gelu_mul(
8556 &self,
8557 gate: &CudaSlice<f32>,
8558 up: &CudaSlice<f32>,
8559 n: usize,
8560 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8561 let f = self.func("moe_pairs_gelu_mul");
8562 let mut act = self.alloc_uninit::<f32>(n)?;
8563 let cfg = LaunchConfig::for_num_elems(n as u32);
8564 let nl = n as i64;
8565 let __s_b = self.gpu.stream();
8566 let mut b = __s_b.launch_builder(&f);
8567 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
8568 unsafe {
8569 b.launch(cfg)?;
8570 }
8571 Ok(act)
8572 }
8573
8574 pub fn moe_pairs_silu_mul(
8575 &self,
8576 gate: &CudaSlice<f32>,
8577 up: &CudaSlice<f32>,
8578 n: usize,
8579 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8580 let f = self.func("moe_pairs_silu_mul");
8581 let mut act = self.alloc_uninit::<f32>(n)?;
8582 let cfg = LaunchConfig::for_num_elems(n as u32);
8583 let nl = n as i64;
8584 let __s_b = self.gpu.stream();
8585 let mut b = __s_b.launch_builder(&f);
8586 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
8587 unsafe {
8588 b.launch(cfg)?;
8589 }
8590 Ok(act)
8591 }
8592
8593 #[allow(clippy::too_many_arguments)]
8594 #[allow(clippy::manual_div_ceil)] pub fn moe_pairs_scatter(
8596 &self,
8597 y_down: &CudaSlice<f32>,
8598 pair_w: &CudaSlice<f32>,
8599 tok_pair_off: &CudaSlice<i32>,
8600 tok_pair_ids: &CudaSlice<i32>,
8601 moe_out: &mut CudaSlice<f32>,
8602 t: usize,
8603 n_embd: usize,
8604 ) -> Result<(), Box<dyn std::error::Error>> {
8605 let f = self.func("moe_pairs_scatter");
8606 let cfg = LaunchConfig {
8607 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
8608 block_dim: (256, 1, 1),
8609 shared_mem_bytes: 0,
8610 };
8611 let ne = n_embd as i32;
8612 let __s_b = self.gpu.stream();
8613 let mut b = __s_b.launch_builder(&f);
8614 b.arg(y_down)
8615 .arg(pair_w)
8616 .arg(tok_pair_off)
8617 .arg(tok_pair_ids)
8618 .arg(moe_out)
8619 .arg(&ne);
8620 unsafe {
8621 b.launch(cfg)?;
8622 }
8623 Ok(())
8624 }
8625
8626 #[allow(clippy::too_many_arguments)]
8630 pub fn moe_gate_up_gelu8_dev_q8(
8631 &self,
8632 table: &CudaSlice<u64>,
8633 sel: &cudarc::driver::CudaView<i32>,
8634 aq: &CudaSlice<i8>,
8635 ad: &CudaSlice<f32>,
8636 in_f: usize,
8637 n_ff: usize,
8638 n_used: usize,
8639 n_expert: usize,
8640 qt_g: i32,
8641 qt_u: i32,
8642 rb_g: usize,
8643 rb_u: usize,
8644 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8645 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
8646 let (inf, nff, ne, rbg, rbu) = (
8647 in_f as i32,
8648 n_ff as i32,
8649 n_expert as i32,
8650 rb_g as i64,
8651 rb_u as i64,
8652 );
8653 let f = self.func("moe_gate_up_gelu8_dev_q8");
8654 let cfg = LaunchConfig {
8655 grid_dim: (n_ff as u32, n_used as u32, 1),
8656 block_dim: (32, 1, 1),
8657 shared_mem_bytes: 0,
8658 };
8659 let __s_b = self.gpu.stream();
8660 let mut b = __s_b.launch_builder(&f);
8661 b.arg(table)
8662 .arg(sel)
8663 .arg(aq)
8664 .arg(ad)
8665 .arg(&mut act)
8666 .arg(&inf)
8667 .arg(&nff)
8668 .arg(&ne)
8669 .arg(&qt_g)
8670 .arg(&qt_u)
8671 .arg(&rbg)
8672 .arg(&rbu);
8673 unsafe {
8674 b.launch(cfg)?;
8675 }
8676 Ok(act)
8677 }
8678
8679 #[allow(clippy::too_many_arguments)]
8681 pub fn moe_gate_up_gelu8_dev_q8_rows(
8682 &self,
8683 table: &CudaSlice<u64>,
8684 sel: &CudaSlice<i32>,
8685 aq: &CudaSlice<i8>,
8686 ad: &CudaSlice<f32>,
8687 t: usize,
8688 in_f: usize,
8689 n_ff: usize,
8690 n_used: usize,
8691 n_expert: usize,
8692 qt_g: i32,
8693 qt_u: i32,
8694 rb_g: usize,
8695 rb_u: usize,
8696 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8697 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
8698 let (inf, nff, ne, rbg, rbu, nu) = (
8699 in_f as i32,
8700 n_ff as i32,
8701 n_expert as i32,
8702 rb_g as i64,
8703 rb_u as i64,
8704 n_used as i32,
8705 );
8706 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
8707 let cfg = LaunchConfig {
8708 grid_dim: (n_ff as u32, n_used as u32, t as u32),
8709 block_dim: (32, 1, 1),
8710 shared_mem_bytes: 0,
8711 };
8712 let __s_b = self.gpu.stream();
8713 let mut b = __s_b.launch_builder(&f);
8714 b.arg(table)
8715 .arg(sel)
8716 .arg(aq)
8717 .arg(ad)
8718 .arg(&mut act)
8719 .arg(&inf)
8720 .arg(&nff)
8721 .arg(&ne)
8722 .arg(&qt_g)
8723 .arg(&qt_u)
8724 .arg(&rbg)
8725 .arg(&rbu)
8726 .arg(&nu);
8727 unsafe {
8728 b.launch(cfg)?;
8729 }
8730 Ok(act)
8731 }
8732
8733 #[allow(clippy::too_many_arguments)]
8735 pub fn moe_gate_up_gelu8_dev_q8_csr(
8736 &self,
8737 table: &CudaSlice<u64>,
8738 sel: &CudaSlice<i32>,
8739 aq: &CudaSlice<i8>,
8740 ad: &CudaSlice<f32>,
8741 n_pairs: usize,
8742 in_f: usize,
8743 n_ff: usize,
8744 n_used: usize,
8745 n_expert: usize,
8746 qt_g: i32,
8747 qt_u: i32,
8748 rb_g: usize,
8749 rb_u: usize,
8750 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8751 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
8752 let (inf, nff, ne, rbg, rbu, nu, npi) = (
8753 in_f as i32,
8754 n_ff as i32,
8755 n_expert as i32,
8756 rb_g as i64,
8757 rb_u as i64,
8758 n_used as i32,
8759 n_pairs as i32,
8760 );
8761 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
8762 let cfg = LaunchConfig {
8763 grid_dim: (n_ff as u32, n_pairs as u32, 1),
8764 block_dim: (32, 1, 1),
8765 shared_mem_bytes: 0,
8766 };
8767 let __s_b = self.gpu.stream();
8768 let mut b = __s_b.launch_builder(&f);
8769 b.arg(table)
8770 .arg(sel)
8771 .arg(aq)
8772 .arg(ad)
8773 .arg(&mut act)
8774 .arg(&inf)
8775 .arg(&nff)
8776 .arg(&ne)
8777 .arg(&qt_g)
8778 .arg(&qt_u)
8779 .arg(&rbg)
8780 .arg(&rbu)
8781 .arg(&nu)
8782 .arg(&npi);
8783 unsafe {
8784 b.launch(cfg)?;
8785 }
8786 Ok(act)
8787 }
8788
8789 #[allow(clippy::too_many_arguments)]
8791 pub fn moe_down8_fma_dev_q8_rows_g(
8792 &self,
8793 table: &CudaSlice<u64>,
8794 sel: &CudaSlice<i32>,
8795 w: &CudaSlice<f32>,
8796 aq2: &CudaSlice<i8>,
8797 ad2: &CudaSlice<f32>,
8798 dst: &mut CudaSlice<f32>,
8799 t: usize,
8800 in_f: usize,
8801 out_f: usize,
8802 n_used: usize,
8803 n_expert: usize,
8804 qt: i32,
8805 rb: usize,
8806 ) -> Result<(), Box<dyn std::error::Error>> {
8807 let (inf, outf, nu, ne, rbi) = (
8808 in_f as i32,
8809 out_f as i32,
8810 n_used as i32,
8811 n_expert as i32,
8812 rb as i64,
8813 );
8814 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
8818 let f = self.func(if step_b1_w8 {
8819 "moe_down8_fma_dev_q8_rows_w8"
8820 } else {
8821 "moe_down8_fma_dev_q8_rows_g"
8822 });
8823 let cfg = LaunchConfig {
8824 grid_dim: (out_f as u32, 1, t as u32),
8825 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
8826 shared_mem_bytes: 0,
8827 };
8828 let __s_b = self.gpu.stream();
8829 let mut b = __s_b.launch_builder(&f);
8830 b.arg(table)
8831 .arg(sel)
8832 .arg(w)
8833 .arg(aq2)
8834 .arg(ad2)
8835 .arg(dst)
8836 .arg(&inf)
8837 .arg(&outf)
8838 .arg(&nu)
8839 .arg(&ne)
8840 .arg(&qt)
8841 .arg(&rbi);
8842 unsafe {
8843 b.launch(cfg)?;
8844 }
8845 Ok(())
8846 }
8847
8848 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
8852 let (out_f, in_f) = (2048usize, 2816usize);
8853 let nblk = in_f / 32;
8854 let mut seed = 0x9E3779B97F4A7C15u64;
8855 let mut rng = move || {
8856 seed = seed
8857 .wrapping_mul(6364136223846793005)
8858 .wrapping_add(1442695040888963407);
8859 (seed >> 33) as u8
8860 };
8861 let mut w = vec![0u8; out_f * nblk * 18];
8862 for b in w.iter_mut() {
8863 *b = rng();
8864 }
8865 for r in 0..out_f {
8866 for g in 0..nblk {
8867 let off = (r * nblk + g) * 18;
8868 w[off] = 0x00;
8869 w[off + 1] = 0x2C; }
8871 }
8872 let qplane = out_f * nblk * 16;
8873 let mut wrp = vec![0u8; w.len()];
8874 for r in 0..out_f {
8875 for g in 0..nblk {
8876 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
8877 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
8878 .copy_from_slice(&src[0..2]);
8879 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
8880 }
8881 }
8882 let w_d = self.htod_bytes(&w)?;
8883 let wrp_d = self.htod_bytes(&wrp)?;
8884 let mut aq = vec![0i8; m * in_f];
8885 for v in aq.iter_mut() {
8886 *v = rng() as i8;
8887 }
8888 let aq_d = self.htod_i8(&aq)?;
8889 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
8890 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
8891 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
8892 const RPB: u32 = 4;
8893 let cfg = LaunchConfig {
8894 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
8895 block_dim: (32, RPB, 1),
8896 shared_mem_bytes: 0,
8897 };
8898 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
8899 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
8900 let fb = self.func("qmatvec_q4_0_mmvq_b4");
8901 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
8902 {
8903 let __s_b = self.gpu.stream();
8904 let mut b = __s_b.launch_builder(&fb);
8905 b.arg(&w_d)
8906 .arg(&aq_d)
8907 .arg(&ad_d)
8908 .arg(&mut y0)
8909 .arg(&inf)
8910 .arg(&outf)
8911 .arg(&mi)
8912 .arg(&rb);
8913 unsafe {
8914 b.launch(cfg)?;
8915 }
8916 let __s_b = self.gpu.stream();
8917 let mut b = __s_b.launch_builder(&fr);
8918 b.arg(&wrp_d)
8919 .arg(&aq_d)
8920 .arg(&ad_d)
8921 .arg(&mut y1)
8922 .arg(&inf)
8923 .arg(&outf)
8924 .arg(&mi)
8925 .arg(&qp);
8926 unsafe {
8927 b.launch(cfg)?;
8928 }
8929 }
8930 self.gpu.stream().synchronize()?;
8931 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
8932 let nd = h0
8933 .iter()
8934 .zip(&h1)
8935 .filter(|(a, b)| a.to_bits() != b.to_bits())
8936 .count();
8937 if nd != 0 {
8938 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
8939 }
8940 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
8941 self.gpu.stream().synchronize()?;
8942 let t0 = std::time::Instant::now();
8943 for _ in 0..500 {
8944 if rp {
8945 let __s_b = self.gpu.stream();
8946 let mut b = __s_b.launch_builder(&fr);
8947 b.arg(&wrp_d)
8948 .arg(&aq_d)
8949 .arg(&ad_d)
8950 .arg(&mut y1)
8951 .arg(&inf)
8952 .arg(&outf)
8953 .arg(&mi)
8954 .arg(&qp);
8955 unsafe {
8956 b.launch(cfg)?;
8957 }
8958 } else {
8959 let __s_b = self.gpu.stream();
8960 let mut b = __s_b.launch_builder(&fb);
8961 b.arg(&w_d)
8962 .arg(&aq_d)
8963 .arg(&ad_d)
8964 .arg(&mut y0)
8965 .arg(&inf)
8966 .arg(&outf)
8967 .arg(&mi)
8968 .arg(&rb);
8969 unsafe {
8970 b.launch(cfg)?;
8971 }
8972 }
8973 }
8974 self.gpu.stream().synchronize()?;
8975 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
8976 };
8977 let _ = time(false)?;
8978 let _ = time(true)?; Ok((time(false)?, time(true)?))
8980 }
8981
8982 pub fn build_q4_rp4(
8987 &self,
8988 t: &mut crate::model::GpuTensor,
8989 ) -> Result<(), Box<dyn std::error::Error>> {
8990 use crate::model::GpuTensor;
8991 let GpuTensor::Quant {
8992 bytes,
8993 qtype,
8994 row_bytes,
8995 ne,
8996 rp4,
8997 ..
8998 } = t
8999 else {
9000 return Ok(());
9001 };
9002 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
9003 return Ok(());
9004 }
9005 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9006 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
9007 return Ok(());
9008 }
9009 let nblk = in_f / 32;
9010 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
9011 let f = self.func("q4_0_split_rp_build");
9012 let n = (out_f * nblk) as i32;
9013 let cfg = LaunchConfig {
9014 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
9015 block_dim: (256, 1, 1),
9016 shared_mem_bytes: 0,
9017 };
9018 let (of, nb) = (out_f as i32, nblk as i32);
9019 let _ = n;
9020 let __s_b = self.gpu.stream();
9021 let mut b = __s_b.launch_builder(&f);
9022 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
9023 unsafe {
9024 b.launch(cfg)?;
9025 }
9026 *rp4 = Some(dst);
9027 Ok(())
9028 }
9029
9030 pub fn build_q8_rp4(
9035 &self,
9036 t: &mut crate::model::GpuTensor,
9037 ) -> Result<(), Box<dyn std::error::Error>> {
9038 use crate::model::GpuTensor;
9039 let GpuTensor::Quant {
9040 bytes,
9041 qtype,
9042 row_bytes,
9043 ne,
9044 rp4,
9045 ..
9046 } = t
9047 else {
9048 return Ok(());
9049 };
9050 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
9051 return Ok(());
9052 }
9053 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9054 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
9055 return Ok(());
9056 }
9057 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
9058 Ok(())
9059 }
9060
9061 pub fn build_q8_rp4_raw(
9064 &self,
9065 bytes: &CudaSlice<u8>,
9066 in_f: usize,
9067 out_f: usize,
9068 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9069 assert!(in_f.is_multiple_of(32));
9070 let nblk = in_f / 32;
9071 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
9072 let f = self.func("q8_0_split_rp_build");
9073 let cfg = LaunchConfig {
9074 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
9075 block_dim: (256, 1, 1),
9076 shared_mem_bytes: 0,
9077 };
9078 let (of, nb) = (out_f as i32, nblk as i32);
9079 let __s_b = self.gpu.stream();
9080 let mut b = __s_b.launch_builder(&f);
9081 b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
9082 unsafe {
9083 b.launch(cfg)?;
9084 }
9085 Ok(dst)
9086 }
9087
9088 pub fn build_q4k_rp4(
9096 &self,
9097 t: &mut crate::model::GpuTensor,
9098 ) -> Result<(), Box<dyn std::error::Error>> {
9099 use crate::model::GpuTensor;
9100 let GpuTensor::Quant {
9101 bytes,
9102 qtype,
9103 row_bytes,
9104 ne,
9105 rp4,
9106 ..
9107 } = t
9108 else {
9109 return Ok(());
9110 };
9111 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
9112 return Ok(());
9113 }
9114 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9115 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
9116 return Ok(());
9117 }
9118 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
9119 Ok(())
9120 }
9121
9122 pub fn build_q6k_rp4(
9123 &self,
9124 t: &mut crate::model::GpuTensor,
9125 ) -> Result<(), Box<dyn std::error::Error>> {
9126 use crate::model::GpuTensor;
9127 let GpuTensor::Quant {
9128 bytes,
9129 qtype,
9130 row_bytes,
9131 ne,
9132 rp4,
9133 ..
9134 } = t
9135 else {
9136 return Ok(());
9137 };
9138 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
9139 return Ok(());
9140 }
9141 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
9142 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
9143 return Ok(());
9144 }
9145 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
9146 Ok(())
9147 }
9148
9149 pub fn build_kq_rp4_raw(
9151 &self,
9152 bytes: &CudaSlice<u8>,
9153 in_f: usize,
9154 out_f: usize,
9155 qtype: i32,
9156 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9157 assert!(in_f.is_multiple_of(256));
9158 let nsbk = in_f / 256;
9159 let (sb_bytes, kname) = match qtype {
9160 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
9161 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
9162 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
9163 };
9164 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
9165 let f = self.func(kname);
9166 let cfg = LaunchConfig {
9167 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
9168 block_dim: (256, 1, 1),
9169 shared_mem_bytes: 0,
9170 };
9171 let (of, nb) = (out_f as i32, nsbk as i32);
9172 let __s_b = self.gpu.stream();
9173 let mut b = __s_b.launch_builder(&f);
9174 b.arg(bytes).arg(&mut dst).arg(&of).arg(&nb);
9175 unsafe {
9176 b.launch(cfg)?;
9177 }
9178 Ok(dst)
9179 }
9180
9181 pub fn kqrp_enabled() -> bool {
9185 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9186 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
9187 Ok("0") => false,
9188 Ok(_) => true,
9189 Err(_) => cfg!(memra_hopper_mma),
9190 })
9191 }
9192
9193 pub fn build_q4_rp_swap(
9199 &self,
9200 t: &mut crate::model::GpuTensor,
9201 ) -> Result<bool, Box<dyn std::error::Error>> {
9202 use crate::model::GpuTensor;
9203 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
9213 return Ok(false);
9214 }
9215 self.build_q4_rp4(t)?;
9216 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
9218 return Ok(false);
9219 };
9220 match rp4.take() {
9221 Some(split) => {
9222 *bytes = split; *rp = true;
9224 Ok(true)
9225 }
9226 None => Ok(false),
9227 }
9228 }
9229
9230 pub fn q4rp_enabled() -> bool {
9232 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9233 *ON.get_or_init(|| {
9234 std::env::var("MEMRA_Q4RP")
9235 .map(|v| v != "0")
9236 .unwrap_or(true)
9237 })
9238 }
9239
9240 #[allow(clippy::manual_div_ceil)] pub fn copy_rows_strided(
9244 &self,
9245 src: &CudaSlice<f32>,
9246 dst: &mut CudaSlice<f32>,
9247 row_elems: usize,
9248 n_rows: usize,
9249 src_stride: usize,
9250 src_off: usize,
9251 ) -> Result<(), Box<dyn std::error::Error>> {
9252 let f = self.func("copy_rows_strided_f32");
9253 let cfg = LaunchConfig {
9254 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
9255 block_dim: (256, 1, 1),
9256 shared_mem_bytes: 0,
9257 };
9258 let (re, nr) = (row_elems as i32, n_rows as i32);
9259 let (st, off) = (src_stride as i64, src_off as i64);
9260 let __s_b = self.gpu.stream();
9261 let mut b = __s_b.launch_builder(&f);
9262 b.arg(src)
9263 .arg(&mut *dst)
9264 .arg(&re)
9265 .arg(&nr)
9266 .arg(&st)
9267 .arg(&off);
9268 unsafe {
9269 b.launch(cfg)?;
9270 }
9271 Ok(())
9272 }
9273
9274 #[allow(clippy::manual_div_ceil)] pub fn place_rows_strided(
9281 &self,
9282 src: &CudaSlice<f32>,
9283 dst: &mut CudaSlice<f32>,
9284 row_elems: usize,
9285 n_rows: usize,
9286 dst_stride: usize,
9287 dst_off: usize,
9288 ) -> Result<(), Box<dyn std::error::Error>> {
9289 if row_elems == 0 || n_rows == 0 {
9290 return Err("strided row placement requires nonzero rows and row width".into());
9291 }
9292 let src_len = n_rows
9293 .checked_mul(row_elems)
9294 .ok_or("strided row placement source size overflow")?;
9295 let dst_len = n_rows
9296 .checked_sub(1)
9297 .and_then(|rows| rows.checked_mul(dst_stride))
9298 .and_then(|base| base.checked_add(dst_off))
9299 .and_then(|base| base.checked_add(row_elems))
9300 .ok_or("strided row placement destination size overflow")?;
9301 let row_end = dst_off
9302 .checked_add(row_elems)
9303 .ok_or("strided row placement row size overflow")?;
9304 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
9305 return Err(format!(
9306 "strided row placement geometry mismatch: src={} need_src={src_len} \
9307 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
9308 dst_stride={dst_stride} dst_off={dst_off}",
9309 src.len(),
9310 dst.len(),
9311 )
9312 .into());
9313 }
9314 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
9315 return Err("strided row placement exceeds CUDA kernel geometry".into());
9316 }
9317 let f = self.func("place_rows_strided_f32");
9318 let cfg = LaunchConfig {
9319 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
9320 block_dim: (256, 1, 1),
9321 shared_mem_bytes: 0,
9322 };
9323 let (re, nr) = (row_elems as i32, n_rows as i32);
9324 let (st, off) = (dst_stride as i64, dst_off as i64);
9325 let __s_b = self.gpu.stream();
9326 let mut b = __s_b.launch_builder(&f);
9327 b.arg(src)
9328 .arg(&mut *dst)
9329 .arg(&re)
9330 .arg(&nr)
9331 .arg(&st)
9332 .arg(&off);
9333 unsafe {
9334 b.launch(cfg)?;
9335 }
9336 Ok(())
9337 }
9338
9339 pub fn u32_set_k(
9341 &self,
9342 dst: &mut CudaSlice<u32>,
9343 v: u32,
9344 idx: usize,
9345 ) -> Result<(), Box<dyn std::error::Error>> {
9346 let f = self.func("u32_set_k");
9347 let cfg = LaunchConfig {
9348 grid_dim: (1, 1, 1),
9349 block_dim: (1, 1, 1),
9350 shared_mem_bytes: 0,
9351 };
9352 let ii = idx as i32;
9353 let __s_b = self.gpu.stream();
9354 let mut b = __s_b.launch_builder(&f);
9355 b.arg(dst).arg(&v).arg(&ii);
9356 unsafe {
9357 b.launch(cfg)?;
9358 }
9359 Ok(())
9360 }
9361
9362 pub fn i32_add_k(
9364 &self,
9365 d: &mut CudaSlice<i32>,
9366 v: i32,
9367 ) -> Result<(), Box<dyn std::error::Error>> {
9368 let f = self.func("i32_add_k");
9369 let cfg = LaunchConfig {
9370 grid_dim: (1, 1, 1),
9371 block_dim: (32, 1, 1),
9372 shared_mem_bytes: 0,
9373 };
9374 let __s_b = self.gpu.stream();
9375 let mut b = __s_b.launch_builder(&f);
9376 b.arg(d).arg(&v);
9377 unsafe {
9378 b.launch(cfg)?;
9379 }
9380 Ok(())
9381 }
9382
9383 pub fn i32_iota_from(
9385 &self,
9386 ctr: &CudaSlice<i32>,
9387 dst: &mut CudaSlice<i32>,
9388 n: usize,
9389 ) -> Result<(), Box<dyn std::error::Error>> {
9390 let f = self.func("i32_iota_from");
9391 let cfg = LaunchConfig::for_num_elems(n as u32);
9392 let ni = n as i32;
9393 let __s_b = self.gpu.stream();
9394 let mut b = __s_b.launch_builder(&f);
9395 b.arg(ctr).arg(dst).arg(&ni);
9396 unsafe {
9397 b.launch(cfg)?;
9398 }
9399 Ok(())
9400 }
9401
9402 pub fn u32_map_k(
9404 &self,
9405 buf: &mut CudaSlice<u32>,
9406 map: &CudaSlice<u32>,
9407 idx: usize,
9408 ) -> Result<(), Box<dyn std::error::Error>> {
9409 let f = self.func("u32_map_k");
9410 let cfg = LaunchConfig {
9411 grid_dim: (1, 1, 1),
9412 block_dim: (1, 1, 1),
9413 shared_mem_bytes: 0,
9414 };
9415 let ii = idx as i32;
9416 let __s_b = self.gpu.stream();
9417 let mut b = __s_b.launch_builder(&f);
9418 b.arg(buf).arg(map).arg(&ii);
9419 unsafe {
9420 b.launch(cfg)?;
9421 }
9422 Ok(())
9423 }
9424
9425 #[allow(clippy::too_many_arguments)]
9427 pub fn u32_pack2(
9428 &self,
9429 a: &CudaSlice<u32>,
9430 off_a: usize,
9431 n1: usize,
9432 b_in: &CudaSlice<u32>,
9433 n2: usize,
9434 out: &mut CudaSlice<u32>,
9435 ) -> Result<(), Box<dyn std::error::Error>> {
9436 let f = self.func("u32_pack2");
9437 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
9438 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
9439 let __s_b = self.gpu.stream();
9440 let mut b = __s_b.launch_builder(&f);
9441 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
9442 unsafe {
9443 b.launch(cfg)?;
9444 }
9445 Ok(())
9446 }
9447
9448 pub fn moe_w_exscale(
9450 &self,
9451 w: &mut CudaSlice<f32>,
9452 sel: &CudaSlice<i32>,
9453 s: &CudaSlice<f32>,
9454 n: usize,
9455 ) -> Result<(), Box<dyn std::error::Error>> {
9456 let f = self.func("moe_w_exscale");
9457 let cfg = LaunchConfig::for_num_elems(n as u32);
9458 let ni = n as i32;
9459 let __s_b = self.gpu.stream();
9460 let mut b = __s_b.launch_builder(&f);
9461 b.arg(w).arg(sel).arg(s).arg(&ni);
9462 unsafe {
9463 b.launch(cfg)?;
9464 }
9465 Ok(())
9466 }
9467
9468 pub fn moe_w_scale_by_expert(
9471 &self,
9472 w: &mut CudaSlice<f32>,
9473 sel: &CudaSlice<i32>,
9474 macros: &CudaSlice<f32>,
9475 n_expert: usize,
9476 n: usize,
9477 ) -> Result<(), Box<dyn std::error::Error>> {
9478 let f = self.func("moe_w_scale_by_expert");
9479 let cfg = LaunchConfig {
9480 grid_dim: (n.div_ceil(64) as u32, 1, 1),
9481 block_dim: (64, 1, 1),
9482 shared_mem_bytes: 0,
9483 };
9484 let (ne, nn) = (n_expert as i32, n as i32);
9485 let __s_b = self.gpu.stream();
9486 let mut b = __s_b.launch_builder(&f);
9487 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
9488 unsafe {
9489 b.launch(cfg)?;
9490 }
9491 Ok(())
9492 }
9493
9494 #[allow(clippy::too_many_arguments)] pub fn moe_gate_up_silu8_dev_q8(
9496 &self,
9497 table: &CudaSlice<u64>,
9498 sel: &cudarc::driver::CudaView<i32>,
9499 aq: &CudaSlice<i8>,
9500 ad: &CudaSlice<f32>,
9501 in_f: usize,
9502 n_ff: usize,
9503 n_used: usize,
9504 n_expert: usize,
9505 qt_g: i32,
9506 qt_u: i32,
9507 rb_g: usize,
9508 rb_u: usize,
9509 macros: &CudaSlice<f32>,
9510 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9511 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
9512 let (mode, wpb) = GU.get_or_init(|| {
9513 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
9514 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
9515 .ok()
9516 .and_then(|v| v.parse().ok())
9517 .unwrap_or(4u32)
9518 .clamp(1, 16);
9519 (mode, wpb)
9520 });
9521 let (mode, wpb) = (mode.as_str(), *wpb);
9522 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
9523 let (inf, nff, ne, rbg, rbu) = (
9524 in_f as i32,
9525 n_ff as i32,
9526 n_expert as i32,
9527 rb_g as i64,
9528 rb_u as i64,
9529 );
9530 let (f, cfg) = match mode {
9531 "1" | "2" | "4" => {
9532 let rpw: u32 = mode.parse().unwrap();
9533 let f = self.func(match rpw {
9534 1 => "moe_gate_up_silu8_dev_q8_r1",
9535 2 => "moe_gate_up_silu8_dev_q8_r2",
9536 _ => "moe_gate_up_silu8_dev_q8_r4",
9537 });
9538 let rows_per_block = (rpw * wpb) as usize;
9539 let gx = n_ff.div_ceil(rows_per_block) as u32;
9540 (
9541 f,
9542 LaunchConfig {
9543 grid_dim: (gx, n_used as u32, 1),
9544 block_dim: (32, wpb, 1),
9545 shared_mem_bytes: 0,
9546 },
9547 )
9548 }
9549 "j8" if n_used <= 32 => (
9550 self.func("moe_gate_up_silu8_dev_q8_j8"),
9551 LaunchConfig {
9552 grid_dim: (n_ff as u32, 1, 1),
9553 block_dim: (32, n_used as u32, 1),
9554 shared_mem_bytes: 0,
9555 },
9556 ),
9557 "vsm2" => {
9559 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
9560 let sh = (rb_g + rb_u) as u32;
9561 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9562 f.set_attribute(
9563 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
9564 sh as i32,
9565 )?;
9566 (
9567 f,
9568 LaunchConfig {
9569 grid_dim: (n_ff as u32, n_used as u32, 1),
9570 block_dim: (32, 1, 1),
9571 shared_mem_bytes: sh,
9572 },
9573 )
9574 }
9575 "vsm" => {
9576 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
9577 let sh = (rb_g + rb_u) as u32;
9578 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9579 f.set_attribute(
9580 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
9581 sh as i32,
9582 )?;
9583 (
9584 f,
9585 LaunchConfig {
9586 grid_dim: (n_ff as u32, n_used as u32, 1),
9587 block_dim: (32, 1, 1),
9588 shared_mem_bytes: sh,
9589 },
9590 )
9591 }
9592 "sg" => (
9593 self.func("moe_gate_up_silu8_dev_q8_sg"),
9594 LaunchConfig {
9595 grid_dim: (n_ff as u32, n_used as u32, 1),
9596 block_dim: (32, 1, 1),
9597 shared_mem_bytes: 0,
9598 },
9599 ),
9600 "j8sg" if n_used <= 32 => (
9601 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
9602 LaunchConfig {
9603 grid_dim: (n_ff as u32, 1, 1),
9604 block_dim: (32, n_used as u32, 1),
9605 shared_mem_bytes: 0,
9606 },
9607 ),
9608 "u64" if in_f == 2048 => (
9609 self.func("moe_gate_up_silu8_dev_q8_u64"),
9610 LaunchConfig {
9611 grid_dim: (n_ff as u32, n_used as u32, 1),
9612 block_dim: (32, 1, 1),
9613 shared_mem_bytes: 0,
9614 },
9615 ),
9616 "gs4" if in_f == 2048 => (
9617 self.func("moe_gate_up_silu8_dev_q8_gs4"),
9618 LaunchConfig {
9619 grid_dim: (n_ff as u32, n_used as u32, 1),
9620 block_dim: (32, 4, 1),
9621 shared_mem_bytes: 0,
9622 },
9623 ),
9624 "v" | "" => (
9626 self.func("moe_gate_up_silu8_dev_q8_v"),
9627 LaunchConfig {
9628 grid_dim: (n_ff as u32, n_used as u32, 1),
9629 block_dim: (32, 1, 1),
9630 shared_mem_bytes: 0,
9631 },
9632 ),
9633 "s2" => (
9634 self.func("moe_gate_up_silu8_dev_q8_s2"),
9635 LaunchConfig {
9636 grid_dim: (n_ff as u32, n_used as u32, 1),
9637 block_dim: (32, 2, 1),
9638 shared_mem_bytes: 0,
9639 },
9640 ),
9641 "s2z" => {
9642 let rz = wpb.min(16); (
9644 self.func("moe_gate_up_silu8_dev_q8_s2z"),
9645 LaunchConfig {
9646 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
9647 block_dim: (32, 2, rz),
9648 shared_mem_bytes: 0,
9649 },
9650 )
9651 }
9652 _ => (
9653 self.func("moe_gate_up_silu8_dev_q8"),
9654 LaunchConfig {
9655 grid_dim: (n_ff as u32, n_used as u32, 1),
9656 block_dim: (32, 1, 1),
9657 shared_mem_bytes: 0,
9658 },
9659 ),
9660 };
9661 let __s_b = self.gpu.stream();
9662 let mut b = __s_b.launch_builder(&f);
9663 b.arg(table)
9664 .arg(sel)
9665 .arg(aq)
9666 .arg(ad)
9667 .arg(&mut act)
9668 .arg(&inf)
9669 .arg(&nff)
9670 .arg(&ne)
9671 .arg(&qt_g)
9672 .arg(&qt_u)
9673 .arg(&rbg)
9674 .arg(&rbu)
9675 .arg(macros);
9676 unsafe {
9677 b.launch(cfg)?;
9678 }
9679 Ok(act)
9680 }
9681
9682 #[allow(clippy::too_many_arguments)]
9683 pub fn moe_down8_fma_dev_q8(
9684 &self,
9685 table: &CudaSlice<u64>,
9686 sel: &cudarc::driver::CudaView<i32>,
9687 w: &cudarc::driver::CudaView<f32>,
9688 aq2: &CudaSlice<i8>,
9689 ad2: &CudaSlice<f32>,
9690 dst: &mut cudarc::driver::CudaViewMut<f32>,
9691 in_f: usize,
9692 out_f: usize,
9693 n_used: usize,
9694 n_expert: usize,
9695 qt: i32,
9696 rb: usize,
9697 ) -> Result<(), Box<dyn std::error::Error>> {
9698 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
9699 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
9700 let (inf, outf, nu, ne, rbi) = (
9701 in_f as i32,
9702 out_f as i32,
9703 n_used as i32,
9704 n_expert as i32,
9705 rb as i64,
9706 );
9707 let (f, cfg) = match mode.as_str() {
9710 m @ ("1" | "2" | "4") if n_used <= 8 => {
9711 let rpw: usize = m.parse().unwrap();
9712 let f = self.func(match rpw {
9713 1 => "moe_down8_fma_dev_q8_w8r1",
9714 2 => "moe_down8_fma_dev_q8_w8r2",
9715 _ => "moe_down8_fma_dev_q8_w8r4",
9716 });
9717 (
9718 f,
9719 LaunchConfig {
9720 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
9721 block_dim: (32, n_used as u32, 1),
9722 shared_mem_bytes: 0,
9723 },
9724 )
9725 }
9726 "h2" if in_f == 512 => (
9727 self.func("moe_down8_fma_dev_q8_h2"),
9728 LaunchConfig {
9729 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9730 block_dim: (32, 1, 1),
9731 shared_mem_bytes: 0,
9732 },
9733 ),
9734 "" if in_f == 704 && n_used <= 8 => (
9737 self.func("moe_down8_fma_dev_q8_w8r2"),
9738 LaunchConfig {
9739 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9740 block_dim: (32, n_used as u32, 1),
9741 shared_mem_bytes: 0,
9742 },
9743 ),
9744 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
9748 self.func("moe_down8_fma_dev_q8_w8h2v"),
9749 LaunchConfig {
9750 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9751 block_dim: (32, n_used as u32, 1),
9752 shared_mem_bytes: 0,
9753 },
9754 ),
9755 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
9756 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
9757 LaunchConfig {
9758 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
9759 block_dim: (32, n_used as u32, 1),
9760 shared_mem_bytes: 0,
9761 },
9762 ),
9763 "w8h2r2" if in_f == 512 && n_used <= 8 => (
9764 self.func("moe_down8_fma_dev_q8_w8h2r2"),
9765 LaunchConfig {
9766 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
9767 block_dim: (32, n_used as u32, 1),
9768 shared_mem_bytes: 0,
9769 },
9770 ),
9771 "w8h2" if in_f == 512 && n_used <= 8 => (
9772 self.func("moe_down8_fma_dev_q8_w8h2"),
9773 LaunchConfig {
9774 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
9775 block_dim: (32, n_used as u32, 1),
9776 shared_mem_bytes: 0,
9777 },
9778 ),
9779 _ => (
9780 self.func("moe_down8_fma_dev_q8"),
9781 LaunchConfig {
9782 grid_dim: (out_f as u32, 1, 1),
9783 block_dim: (32, 1, 1),
9784 shared_mem_bytes: 0,
9785 },
9786 ),
9787 };
9788 let __s_b = self.gpu.stream();
9789 let mut b = __s_b.launch_builder(&f);
9790 b.arg(table)
9791 .arg(sel)
9792 .arg(w)
9793 .arg(aq2)
9794 .arg(ad2)
9795 .arg(dst)
9796 .arg(&inf)
9797 .arg(&outf)
9798 .arg(&nu)
9799 .arg(&ne)
9800 .arg(&qt)
9801 .arg(&rbi);
9802 unsafe {
9803 b.launch(cfg)?;
9804 }
9805 Ok(())
9806 }
9807
9808 #[allow(clippy::too_many_arguments)]
9815 pub fn moe_gate_up_silu8_dev_q8_rows(
9816 &self,
9817 table: &CudaSlice<u64>,
9818 sel: &CudaSlice<i32>,
9819 aq: &CudaSlice<i8>,
9820 ad: &CudaSlice<f32>,
9821 t: usize,
9822 in_f: usize,
9823 n_ff: usize,
9824 n_used: usize,
9825 n_expert: usize,
9826 qt_g: i32,
9827 qt_u: i32,
9828 rb_g: usize,
9829 rb_u: usize,
9830 macros: &CudaSlice<f32>,
9831 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9832 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
9833 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
9834 let cfg = LaunchConfig {
9835 grid_dim: (n_ff as u32, n_used as u32, t as u32),
9836 block_dim: (32, 1, 1),
9837 shared_mem_bytes: 0,
9838 };
9839 let (inf, nff, ne, nu, rbg, rbu) = (
9840 in_f as i32,
9841 n_ff as i32,
9842 n_expert as i32,
9843 n_used as i32,
9844 rb_g as i64,
9845 rb_u as i64,
9846 );
9847 let __s_b = self.gpu.stream();
9848 let mut b = __s_b.launch_builder(&f);
9849 b.arg(table)
9850 .arg(sel)
9851 .arg(aq)
9852 .arg(ad)
9853 .arg(&mut act)
9854 .arg(&inf)
9855 .arg(&nff)
9856 .arg(&ne)
9857 .arg(&qt_g)
9858 .arg(&qt_u)
9859 .arg(&rbg)
9860 .arg(&rbu)
9861 .arg(&nu)
9862 .arg(macros);
9863 unsafe {
9864 b.launch(cfg)?;
9865 }
9866 Ok(act)
9867 }
9868
9869 #[allow(clippy::too_many_arguments)]
9874 pub fn moe_down8_fma_dev_q8_rows(
9875 &self,
9876 table: &CudaSlice<u64>,
9877 sel: &CudaSlice<i32>,
9878 w: &CudaSlice<f32>,
9879 aq2: &CudaSlice<i8>,
9880 ad2: &CudaSlice<f32>,
9881 dst: &mut CudaSlice<f32>,
9882 t: usize,
9883 in_f: usize,
9884 out_f: usize,
9885 n_used: usize,
9886 n_expert: usize,
9887 qt: i32,
9888 rb: usize,
9889 ) -> Result<(), Box<dyn std::error::Error>> {
9890 assert!(
9891 in_f == 512 && n_used <= 8,
9892 "down rows twin is w8h2v shape-gated"
9893 );
9894 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
9895 let cfg = LaunchConfig {
9896 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
9897 block_dim: (32, n_used as u32, 1),
9898 shared_mem_bytes: 0,
9899 };
9900 let (inf, outf, nu, ne, rbi) = (
9901 in_f as i32,
9902 out_f as i32,
9903 n_used as i32,
9904 n_expert as i32,
9905 rb as i64,
9906 );
9907 let __s_b = self.gpu.stream();
9908 let mut b = __s_b.launch_builder(&f);
9909 b.arg(table)
9910 .arg(sel)
9911 .arg(w)
9912 .arg(aq2)
9913 .arg(ad2)
9914 .arg(dst)
9915 .arg(&inf)
9916 .arg(&outf)
9917 .arg(&nu)
9918 .arg(&ne)
9919 .arg(&qt)
9920 .arg(&rbi);
9921 unsafe {
9922 b.launch(cfg)?;
9923 }
9924 Ok(())
9925 }
9926
9927 #[allow(clippy::too_many_arguments)]
9931 pub fn moe_gate_up_silu8_dev_q8_csr(
9932 &self,
9933 table: &CudaSlice<u64>,
9934 sel: &CudaSlice<i32>,
9935 aq: &CudaSlice<i8>,
9936 ad: &CudaSlice<f32>,
9937 n_pairs: usize,
9938 in_f: usize,
9939 n_ff: usize,
9940 n_used: usize,
9941 n_expert: usize,
9942 qt_g: i32,
9943 qt_u: i32,
9944 rb_g: usize,
9945 rb_u: usize,
9946 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9947 let f = if qt_g == crate::QT_NVFP4 {
9950 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
9951 } else {
9952 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
9953 };
9954 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
9955 let cfg = LaunchConfig {
9956 grid_dim: (n_ff as u32, n_pairs as u32, 1),
9957 block_dim: (32, 1, 1),
9958 shared_mem_bytes: 0,
9959 };
9960 let (inf, nff, ne, nu, npi, rbg, rbu) = (
9961 in_f as i32,
9962 n_ff as i32,
9963 n_expert as i32,
9964 n_used as i32,
9965 n_pairs as i32,
9966 rb_g as i64,
9967 rb_u as i64,
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 .arg(&nu)
9984 .arg(&npi);
9985 unsafe {
9986 b.launch(cfg)?;
9987 }
9988 Ok(act)
9989 }
9990
9991 #[allow(clippy::too_many_arguments)]
9995 pub fn moe_down8_fma_dev_q8_variant(
9996 &self,
9997 variant: &str,
9998 table: &CudaSlice<u64>,
9999 sel: &cudarc::driver::CudaView<i32>,
10000 w: &cudarc::driver::CudaView<f32>,
10001 aq2: &CudaSlice<i8>,
10002 ad2: &CudaSlice<f32>,
10003 dst: &mut cudarc::driver::CudaViewMut<f32>,
10004 in_f: usize,
10005 out_f: usize,
10006 n_used: usize,
10007 n_expert: usize,
10008 qt: i32,
10009 rb: usize,
10010 ) -> Result<(), Box<dyn std::error::Error>> {
10011 let (inf, outf, nu, ne, rbi) = (
10012 in_f as i32,
10013 out_f as i32,
10014 n_used as i32,
10015 n_expert as i32,
10016 rb as i64,
10017 );
10018 let (f, cfg) = match variant {
10019 "w8h2" | "w8h2v" => (
10020 self.func(if variant == "w8h2" {
10021 "moe_down8_fma_dev_q8_w8h2"
10022 } else {
10023 "moe_down8_fma_dev_q8_w8h2v"
10024 }),
10025 LaunchConfig {
10026 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
10027 block_dim: (32, n_used as u32, 1),
10028 shared_mem_bytes: 0,
10029 },
10030 ),
10031 "w8h2r2" | "w8h2r2v" => (
10032 self.func(if variant == "w8h2r2" {
10033 "moe_down8_fma_dev_q8_w8h2r2"
10034 } else {
10035 "moe_down8_fma_dev_q8_w8h2r2v"
10036 }),
10037 LaunchConfig {
10038 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
10039 block_dim: (32, n_used as u32, 1),
10040 shared_mem_bytes: 0,
10041 },
10042 ),
10043 _ => (
10044 self.func("moe_down8_fma_dev_q8"),
10045 LaunchConfig {
10046 grid_dim: (out_f as u32, 1, 1),
10047 block_dim: (32, 1, 1),
10048 shared_mem_bytes: 0,
10049 },
10050 ),
10051 };
10052 let __s_b = self.gpu.stream();
10053 let mut b = __s_b.launch_builder(&f);
10054 b.arg(table)
10055 .arg(sel)
10056 .arg(w)
10057 .arg(aq2)
10058 .arg(ad2)
10059 .arg(dst)
10060 .arg(&inf)
10061 .arg(&outf)
10062 .arg(&nu)
10063 .arg(&ne)
10064 .arg(&qt)
10065 .arg(&rbi);
10066 unsafe {
10067 b.launch(cfg)?;
10068 }
10069 Ok(())
10070 }
10071
10072 #[allow(clippy::too_many_arguments)]
10074 pub fn moe_gate_up_silu8_dev_q8_variant(
10075 &self,
10076 variant: &str,
10077 table: &CudaSlice<u64>,
10078 sel: &cudarc::driver::CudaView<i32>,
10079 aq: &CudaSlice<i8>,
10080 ad: &CudaSlice<f32>,
10081 in_f: usize,
10082 n_ff: usize,
10083 n_used: usize,
10084 n_expert: usize,
10085 qt_g: i32,
10086 qt_u: i32,
10087 rb_g: usize,
10088 rb_u: usize,
10089 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10090 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
10091 let (inf, nff, ne, rbg, rbu) = (
10092 in_f as i32,
10093 n_ff as i32,
10094 n_expert as i32,
10095 rb_g as i64,
10096 rb_u as i64,
10097 );
10098 let f = self.func(if variant == "v" {
10099 "moe_gate_up_silu8_dev_q8_v"
10100 } else {
10101 "moe_gate_up_silu8_dev_q8"
10102 });
10103 let cfg = LaunchConfig {
10104 grid_dim: (n_ff as u32, n_used as u32, 1),
10105 block_dim: (32, 1, 1),
10106 shared_mem_bytes: 0,
10107 };
10108 let __s_b = self.gpu.stream();
10109 let mut b = __s_b.launch_builder(&f);
10110 b.arg(table)
10111 .arg(sel)
10112 .arg(aq)
10113 .arg(ad)
10114 .arg(&mut act)
10115 .arg(&inf)
10116 .arg(&nff)
10117 .arg(&ne)
10118 .arg(&qt_g)
10119 .arg(&qt_u)
10120 .arg(&rbg)
10121 .arg(&rbu);
10122 unsafe {
10123 b.launch(cfg)?;
10124 }
10125 Ok(act)
10126 }
10127
10128 #[allow(clippy::too_many_arguments)] pub fn moe_gate_up_silu8_dev(
10130 &self,
10131 table: &CudaSlice<u64>,
10132 sel: &cudarc::driver::CudaView<i32>,
10133 x: &cudarc::driver::CudaView<f32>,
10134 in_f: usize,
10135 n_ff: usize,
10136 n_used: usize,
10137 n_expert: usize,
10138 qt_g: i32,
10139 qt_u: i32,
10140 rb_g: usize,
10141 rb_u: usize,
10142 macros: &CudaSlice<f32>,
10143 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10144 let f = self.func("moe_gate_up_silu8_dev");
10145 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
10147 grid_dim: (n_ff as u32, n_used as u32, 1),
10148 block_dim: (256, 1, 1),
10149 shared_mem_bytes: 0,
10150 };
10151 let (inf, nff, ne, rbg, rbu) = (
10152 in_f as i32,
10153 n_ff as i32,
10154 n_expert as i32,
10155 rb_g as i64,
10156 rb_u as i64,
10157 );
10158 let __s_b = self.gpu.stream();
10159 let mut b = __s_b.launch_builder(&f);
10160 b.arg(table)
10161 .arg(sel)
10162 .arg(x)
10163 .arg(&mut act)
10164 .arg(&inf)
10165 .arg(&nff)
10166 .arg(&ne)
10167 .arg(&qt_g)
10168 .arg(&qt_u)
10169 .arg(&rbg)
10170 .arg(&rbu)
10171 .arg(macros);
10172 unsafe {
10173 b.launch(cfg)?;
10174 }
10175 Ok(act)
10176 }
10177
10178 #[allow(clippy::too_many_arguments)]
10181 pub fn moe_down8_fma_dev(
10182 &self,
10183 table: &CudaSlice<u64>,
10184 sel: &cudarc::driver::CudaView<i32>,
10185 w: &cudarc::driver::CudaView<f32>,
10186 act: &CudaSlice<f32>,
10187 dst: &mut cudarc::driver::CudaViewMut<f32>,
10188 in_f: usize,
10189 out_f: usize,
10190 n_used: usize,
10191 n_expert: usize,
10192 qt: i32,
10193 rb: usize,
10194 ) -> Result<(), Box<dyn std::error::Error>> {
10195 let f = self.func("moe_down8_fma_dev");
10196 let cfg = LaunchConfig {
10197 grid_dim: (out_f as u32, 1, 1),
10198 block_dim: (256, 1, 1),
10199 shared_mem_bytes: 0,
10200 };
10201 let (inf, outf, nu, ne, rbv) = (
10202 in_f as i32,
10203 out_f as i32,
10204 n_used as i32,
10205 n_expert as i32,
10206 rb as i64,
10207 );
10208 let __s_b = self.gpu.stream();
10209 let mut b = __s_b.launch_builder(&f);
10210 b.arg(table)
10211 .arg(sel)
10212 .arg(w)
10213 .arg(act)
10214 .arg(dst)
10215 .arg(&inf)
10216 .arg(&outf)
10217 .arg(&nu)
10218 .arg(&ne)
10219 .arg(&qt)
10220 .arg(&rbv);
10221 unsafe {
10222 b.launch(cfg)?;
10223 }
10224 Ok(())
10225 }
10226
10227 pub fn axpy_into(
10229 &self,
10230 src: &CudaSlice<f32>,
10231 alpha: f32,
10232 dst: &mut cudarc::driver::CudaViewMut<f32>,
10233 n: usize,
10234 ) -> Result<(), Box<dyn std::error::Error>> {
10235 let f = self.func("axpy_f32");
10236 let cfg = LaunchConfig::for_num_elems(n as u32);
10237 let (a, ni) = (alpha, n as i32);
10238 let __s_b = self.gpu.stream();
10239 let mut b = __s_b.launch_builder(&f);
10240 b.arg(src).arg(dst).arg(&a).arg(&ni);
10241 unsafe {
10242 b.launch(cfg)?;
10243 }
10244 Ok(())
10245 }
10246
10247 pub fn axpy_host_into(
10249 &self,
10250 src: &cudarc::driver::CudaView<'_, f32>,
10251 alpha: f32,
10252 dst: &mut cudarc::driver::CudaViewMut<f32>,
10253 n: usize,
10254 ) -> Result<(), Box<dyn std::error::Error>> {
10255 let f = self.func("axpy_host_f32");
10256 let cfg = LaunchConfig::for_num_elems(n as u32);
10257 let (a, ni) = (alpha, n as i32);
10258 let __s_b = self.gpu.stream();
10259 let mut b = __s_b.launch_builder(&f);
10260 b.arg(src).arg(dst).arg(&a).arg(&ni);
10261 unsafe {
10262 b.launch(cfg)?;
10263 }
10264 Ok(())
10265 }
10266
10267 pub fn add_scaled_rows(
10269 &self,
10270 src: &CudaSlice<f32>,
10271 scale: &CudaSlice<f32>,
10272 dst: &mut CudaSlice<f32>,
10273 ncols: usize,
10274 nrows: usize,
10275 ) -> Result<(), Box<dyn std::error::Error>> {
10276 let f = self.func("add_scaled_rows_f32");
10277 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10278 let (nc, nr) = (ncols as i32, nrows as i32);
10279 let __s_b = self.gpu.stream();
10280 let mut b = __s_b.launch_builder(&f);
10281 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
10282 unsafe {
10283 b.launch(cfg)?;
10284 }
10285 Ok(())
10286 }
10287
10288 pub fn add_scaled_rows_ones(
10293 &self,
10294 src: &CudaSlice<f32>,
10295 dst: &mut CudaSlice<f32>,
10296 ncols: usize,
10297 nrows: usize,
10298 ) -> Result<(), Box<dyn std::error::Error>> {
10299 let mut guard = self
10300 .shexp_ones
10301 .lock()
10302 .map_err(|_| "shexp ones buffer is poisoned")?;
10303 if guard.as_ref().map(|b| b.len() < nrows).unwrap_or(true) {
10304 *guard = Some(self.htod(&vec![1.0f32; nrows.max(64)])?);
10307 }
10308 let ones = guard.as_ref().expect("just ensured");
10309 let f = self.func("add_scaled_rows_f32");
10310 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10311 let (nc, nr) = (ncols as i32, nrows as i32);
10312 let __s_b = self.gpu.stream();
10313 let mut b = __s_b.launch_builder(&f);
10314 b.arg(src).arg(ones).arg(&mut *dst).arg(&nc).arg(&nr);
10315 unsafe {
10316 b.launch(cfg)?;
10317 }
10318 Ok(())
10319 }
10320
10321 pub fn i32_mirror_store(
10325 &self,
10326 dst: &mut CudaSlice<i32>,
10327 v: i32,
10328 ) -> Result<(), Box<dyn std::error::Error>> {
10329 if crate::htod_diet_on() {
10330 HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10331 return self.i32_set_k(dst, v);
10332 }
10333 self.gpu.stream().memcpy_htod(&[v], dst)?;
10334 Ok(())
10335 }
10336
10337 pub fn scale_rows(
10340 &self,
10341 y: &mut CudaSlice<f32>,
10342 s: &CudaSlice<f32>,
10343 ncols: usize,
10344 nrows: usize,
10345 ) -> Result<(), Box<dyn std::error::Error>> {
10346 let f = self.func("scale_rows_f32");
10347 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
10348 let (nc, nr) = (ncols as i32, nrows as i32);
10349 let __s_b = self.gpu.stream();
10350 let mut b = __s_b.launch_builder(&f);
10351 b.arg(&mut *y).arg(s).arg(&nc).arg(&nr);
10352 unsafe {
10353 b.launch(cfg)?;
10354 }
10355 Ok(())
10356 }
10357
10358 #[allow(clippy::too_many_arguments)]
10362 pub fn moe_prime_join_scatter(
10363 &self,
10364 y0: &CudaSlice<f32>,
10365 y1: &CudaSlice<f32>,
10366 inv: &CudaSlice<i32>,
10367 w: &CudaSlice<f32>,
10368 out: &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("moe_prime_join_scatter_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(y0)
10379 .arg(y1)
10380 .arg(inv)
10381 .arg(w)
10382 .arg(&mut *out)
10383 .arg(&nc)
10384 .arg(&nu)
10385 .arg(&ti);
10386 unsafe {
10387 b.launch(cfg)?;
10388 }
10389 Ok(())
10390 }
10391
10392 pub fn moe_pairs_weighted_scatter(
10395 &self,
10396 y: &CudaSlice<f32>,
10397 w: &CudaSlice<f32>,
10398 out: &mut CudaSlice<f32>,
10399 ncols: usize,
10400 n_used: usize,
10401 t: usize,
10402 ) -> Result<(), Box<dyn std::error::Error>> {
10403 let f = self.func("moe_pairs_weighted_scatter_f32");
10404 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10405 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10406 let __s_b = self.gpu.stream();
10407 let mut b = __s_b.launch_builder(&f);
10408 b.arg(y).arg(w).arg(&mut *out).arg(&nc).arg(&nu).arg(&ti);
10409 unsafe {
10410 b.launch(cfg)?;
10411 }
10412 Ok(())
10413 }
10414
10415 pub fn gather_rows(
10419 &self,
10420 src: &CudaSlice<f32>,
10421 idx: &CudaSlice<i32>,
10422 dst: &mut CudaSlice<f32>,
10423 ncols: usize,
10424 m_e: usize,
10425 ) -> Result<(), Box<dyn std::error::Error>> {
10426 let f = self.func("gather_rows_f32");
10427 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
10428 let (nc, me) = (ncols as i32, m_e as i32);
10429 let __s_b = self.gpu.stream();
10430 let mut b = __s_b.launch_builder(&f);
10431 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
10432 unsafe {
10433 b.launch(cfg)?;
10434 }
10435 Ok(())
10436 }
10437
10438 #[allow(clippy::too_many_arguments)] pub fn scatter_slot(
10444 &self,
10445 src: &CudaSlice<f32>,
10446 tok_idx: &CudaSlice<i32>,
10447 slot_idx: &CudaSlice<i32>,
10448 weight: &CudaSlice<f32>,
10449 dst: &mut CudaSlice<f32>,
10450 wbuf: &mut CudaSlice<f32>,
10451 ncols: usize,
10452 n_used: usize,
10453 m_e: usize,
10454 ) -> Result<(), Box<dyn std::error::Error>> {
10455 let f = self.func("scatter_add_slot_f32");
10456 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
10457 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
10458 let __s_b = self.gpu.stream();
10459 let mut b = __s_b.launch_builder(&f);
10460 b.arg(src)
10461 .arg(tok_idx)
10462 .arg(slot_idx)
10463 .arg(weight)
10464 .arg(dst)
10465 .arg(wbuf)
10466 .arg(&nc)
10467 .arg(&nu)
10468 .arg(&me);
10469 unsafe {
10470 b.launch(cfg)?;
10471 }
10472 Ok(())
10473 }
10474
10475 pub fn reduce_slots(
10479 &self,
10480 slots: &CudaSlice<f32>,
10481 wbuf: &CudaSlice<f32>,
10482 dst: &mut CudaSlice<f32>,
10483 ncols: usize,
10484 n_used: usize,
10485 t: usize,
10486 ) -> Result<(), Box<dyn std::error::Error>> {
10487 let f = self.func("reduce_slots_f32");
10488 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10489 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10490 let __s_b = self.gpu.stream();
10491 let mut b = __s_b.launch_builder(&f);
10492 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
10493 unsafe {
10494 b.launch(cfg)?;
10495 }
10496 Ok(())
10497 }
10498
10499 pub fn reduce_slots_host(
10504 &self,
10505 slots: &CudaSlice<f32>,
10506 wbuf: &CudaSlice<f32>,
10507 dst: &mut CudaSlice<f32>,
10508 ncols: usize,
10509 n_used: usize,
10510 t: usize,
10511 ) -> Result<(), Box<dyn std::error::Error>> {
10512 let f = self.func("reduce_slots_host_f32");
10513 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
10514 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
10515 let __s_b = self.gpu.stream();
10516 let mut b = __s_b.launch_builder(&f);
10517 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
10518 unsafe {
10519 b.launch(cfg)?;
10520 }
10521 Ok(())
10522 }
10523
10524 pub fn quantize_q8_1_view(
10531 &self,
10532 x: &cudarc::driver::CudaView<f32>,
10533 m: usize,
10534 in_f: usize,
10535 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10536 let f = self.func("quantize_q8_1");
10537 let nblk = in_f / 32;
10538 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
10539 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
10540 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10541 let (inf, mi) = (in_f as i32, m as i32);
10542 let __s_b = self.gpu.stream();
10543 let mut b = __s_b.launch_builder(&f);
10544 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
10545 unsafe {
10546 b.launch(cfg)?;
10547 }
10548 Ok((q, d))
10549 }
10550
10551 pub fn quantize_q8_1(
10552 &self,
10553 x: &CudaSlice<f32>,
10554 m: usize,
10555 in_f: usize,
10556 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10557 let nblk = in_f / 32;
10558 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);
10562 let (inf, mi) = (in_f as i32, m as i32);
10563 if Self::pdl_on() && Self::pdl_wb_on() {
10564 {
10565 use cudarc::driver::{DevicePtr, DevicePtrMut};
10566 let s = &self.gpu.stream();
10567 let (px, _g0) = x.device_ptr(s);
10568 let (pq, _g1) = q.device_ptr_mut(s);
10569 let (pd, _g2) = d.device_ptr_mut(s);
10570 let mut ps = [
10571 &px as *const _ as *mut std::ffi::c_void,
10572 &pq as *const _ as *mut _,
10573 &pd as *const _ as *mut _,
10574 &inf as *const _ as *mut _,
10575 &mi as *const _ as *mut _,
10576 ];
10577 unsafe {
10578 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10579 }
10580 }
10581 return Ok((q, d));
10582 }
10583 let f = self.func("quantize_q8_1");
10584 let __s_b = self.gpu.stream();
10585 let mut b = __s_b.launch_builder(&f);
10586 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
10587 unsafe {
10588 b.launch(cfg)?;
10589 }
10590 Ok((q, d))
10591 }
10592
10593 pub fn quantize_fp4_act(
10597 &self,
10598 x: &CudaSlice<f32>,
10599 m: usize,
10600 in_f: usize,
10601 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
10602 let f = self.func("quantize_fp4_act");
10603 let nb16 = in_f / 16;
10604 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);
10607 let (inf, mi) = (in_f as i32, m as i32);
10608 let __s_b = self.gpu.stream();
10609 let mut b = __s_b.launch_builder(&f);
10610 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
10611 unsafe {
10612 b.launch(cfg)?;
10613 }
10614 Ok((aq4, ad4))
10615 }
10616
10617 #[allow(clippy::too_many_arguments)] pub fn qmatvec_gemm_nvfp4_fp4(
10623 &self,
10624 bytes: &CudaSlice<u8>,
10625 x: &CudaSlice<f32>,
10626 m: usize,
10627 in_f: usize,
10628 out_f: usize,
10629 row_bytes: usize,
10630 scale: f32,
10631 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10632 assert!(
10633 in_f.is_multiple_of(64),
10634 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
10635 );
10636 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
10637 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
10638 if scale != 1.0 {
10639 self.scale_inplace(&mut y, scale, m * out_f)?;
10640 }
10641 Ok(y)
10642 }
10643
10644 #[allow(clippy::too_many_arguments)]
10647 #[allow(clippy::manual_div_ceil)] fn fp4_gemm_launch(
10650 &self,
10651 bytes: &CudaSlice<u8>,
10652 aq4: &CudaSlice<u32>,
10653 ad4: &CudaSlice<u8>,
10654 m: usize,
10655 in_f: usize,
10656 out_f: usize,
10657 row_bytes: usize,
10658 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10659 let f = self.func("qmatvec_gemm_nvfp4_fp4");
10660 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
10662 const BN: u32 = 256;
10663 let cfg = LaunchConfig {
10664 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
10665 block_dim: (32, 4, 1),
10666 shared_mem_bytes: 0,
10667 };
10668 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10669 let __s_b = self.gpu.stream();
10670 let mut b = __s_b.launch_builder(&f);
10671 b.arg(bytes)
10672 .arg(aq4)
10673 .arg(ad4)
10674 .arg(&mut y)
10675 .arg(&inf)
10676 .arg(&outf)
10677 .arg(&mi)
10678 .arg(&rb);
10679 unsafe {
10680 b.launch(cfg)?;
10681 }
10682 Ok(y)
10683 }
10684
10685 pub fn qmatvec_gemm_nvfp4_fp4_raw(
10687 &self,
10688 bytes: &CudaSlice<u8>,
10689 x: &CudaSlice<f32>,
10690 m: usize,
10691 in_f: usize,
10692 out_f: usize,
10693 row_bytes: usize,
10694 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10695 assert!(
10696 in_f.is_multiple_of(64),
10697 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
10698 );
10699 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
10700 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
10701 }
10702
10703 pub fn qmatvec_q8_0_fast(
10705 &self,
10706 w: &CudaSlice<u8>,
10707 x: &CudaSlice<f32>,
10708 m: usize,
10709 in_f: usize,
10710 out_f: usize,
10711 row_bytes: usize,
10712 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10713 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10714 let f = self.func("qmatvec_q8_0_dp4a");
10715 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10717 grid_dim: (out_f as u32, m as u32, 1),
10718 block_dim: (128, 1, 1),
10719 shared_mem_bytes: 0,
10720 };
10721 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10722 let __s_b = self.gpu.stream();
10723 let mut b = __s_b.launch_builder(&f);
10724 b.arg(w)
10725 .arg(&aq)
10726 .arg(&ad)
10727 .arg(&mut y)
10728 .arg(&inf)
10729 .arg(&outf)
10730 .arg(&mi)
10731 .arg(&rb);
10732 unsafe {
10733 b.launch(cfg)?;
10734 }
10735 Ok(y)
10736 }
10737
10738 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
10741 &self,
10742 w: &CudaSlice<u8>,
10743 x: &CudaSlice<f32>,
10744 m: usize,
10745 in_f: usize,
10746 out_f: usize,
10747 row_bytes: usize,
10748 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10749 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10750 let f = self.func("qmatvec_q4_K_dp4a");
10751 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10753 grid_dim: (out_f as u32, m as u32, 1),
10754 block_dim: (128, 1, 1),
10755 shared_mem_bytes: 0,
10756 };
10757 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10758 let __s_b = self.gpu.stream();
10759 let mut b = __s_b.launch_builder(&f);
10760 b.arg(w)
10761 .arg(&aq)
10762 .arg(&ad)
10763 .arg(&mut y)
10764 .arg(&inf)
10765 .arg(&outf)
10766 .arg(&mi)
10767 .arg(&rb);
10768 unsafe {
10769 b.launch(cfg)?;
10770 }
10771 Ok(y)
10772 }
10773
10774 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
10777 &self,
10778 w: &CudaSlice<u8>,
10779 x: &CudaSlice<f32>,
10780 m: usize,
10781 in_f: usize,
10782 out_f: usize,
10783 row_bytes: usize,
10784 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10785 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10786 let f = self.func("qmatvec_q6_K_dp4a");
10787 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10789 grid_dim: (out_f as u32, m as u32, 1),
10790 block_dim: (128, 1, 1),
10791 shared_mem_bytes: 0,
10792 };
10793 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10794 let __s_b = self.gpu.stream();
10795 let mut b = __s_b.launch_builder(&f);
10796 b.arg(w)
10797 .arg(&aq)
10798 .arg(&ad)
10799 .arg(&mut y)
10800 .arg(&inf)
10801 .arg(&outf)
10802 .arg(&mi)
10803 .arg(&rb);
10804 unsafe {
10805 b.launch(cfg)?;
10806 }
10807 Ok(y)
10808 }
10809
10810 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
10813 &self,
10814 w: &CudaSlice<u8>,
10815 x: &CudaSlice<f32>,
10816 m: usize,
10817 in_f: usize,
10818 out_f: usize,
10819 row_bytes: usize,
10820 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10821 self.qmatvec_dp4a_named(
10822 "qmatvec_q5_K_dp4a",
10823 &w.slice(0..w.len()),
10824 x,
10825 m,
10826 in_f,
10827 out_f,
10828 row_bytes,
10829 )
10830 }
10831 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
10834 &self,
10835 w: &CudaSlice<u8>,
10836 x: &CudaSlice<f32>,
10837 m: usize,
10838 in_f: usize,
10839 out_f: usize,
10840 row_bytes: usize,
10841 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10842 self.qmatvec_dp4a_named(
10843 "qmatvec_q3_K_dp4a",
10844 &w.slice(0..w.len()),
10845 x,
10846 m,
10847 in_f,
10848 out_f,
10849 row_bytes,
10850 )
10851 }
10852 pub fn qmatvec_nvfp4_fast_rp(
10854 &self,
10855 w: &CudaSlice<u8>,
10856 x: &CudaSlice<f32>,
10857 m: usize,
10858 in_f: usize,
10859 out_f: usize,
10860 row_bytes: usize,
10861 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10862 assert!(
10863 in_f.is_multiple_of(64),
10864 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10865 );
10866 self.qmatvec_dp4a_named(
10867 "qmatvec_nvfp4_dp4a_rp",
10868 &w.slice(0..w.len()),
10869 x,
10870 m,
10871 in_f,
10872 out_f,
10873 row_bytes,
10874 )
10875 }
10876 pub fn qmatvec_nvfp4_fast(
10878 &self,
10879 w: &cudarc::driver::CudaView<'_, u8>,
10880 x: &CudaSlice<f32>,
10881 m: usize,
10882 in_f: usize,
10883 out_f: usize,
10884 row_bytes: usize,
10885 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10886 assert!(
10889 in_f.is_multiple_of(64),
10890 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10891 );
10892 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
10893 }
10894 pub fn qmatvec_nvfp4_fast_v2(
10899 &self,
10900 w: &cudarc::driver::CudaView<'_, u8>,
10901 x: &CudaSlice<f32>,
10902 m: usize,
10903 in_f: usize,
10904 out_f: usize,
10905 row_bytes: usize,
10906 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10907 assert!(
10908 in_f.is_multiple_of(64),
10909 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10910 );
10911 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
10912 }
10913 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
10916 &self,
10917 w: &CudaSlice<u8>,
10918 x: &CudaSlice<f32>,
10919 m: usize,
10920 in_f: usize,
10921 out_f: usize,
10922 row_bytes: usize,
10923 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10924 self.qmatvec_dp4a_named(
10925 "qmatvec_iq4_XS_dp4a",
10926 &w.slice(0..w.len()),
10927 x,
10928 m,
10929 in_f,
10930 out_f,
10931 row_bytes,
10932 )
10933 }
10934
10935 #[allow(clippy::too_many_arguments)] fn qmatvec_dp4a_named(
10938 &self,
10939 name: &str,
10940 w: &cudarc::driver::CudaView<'_, u8>,
10941 x: &CudaSlice<f32>,
10942 m: usize,
10943 in_f: usize,
10944 out_f: usize,
10945 row_bytes: usize,
10946 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10947 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
10948 let f = self.func(name);
10949 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
10951 grid_dim: (out_f as u32, m as u32, 1),
10952 block_dim: (128, 1, 1),
10953 shared_mem_bytes: 0,
10954 };
10955 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
10956 let __s_b = self.gpu.stream();
10957 let mut b = __s_b.launch_builder(&f);
10958 b.arg(w)
10959 .arg(&aq)
10960 .arg(&ad)
10961 .arg(&mut y)
10962 .arg(&inf)
10963 .arg(&outf)
10964 .arg(&mi)
10965 .arg(&rb);
10966 unsafe {
10967 b.launch(cfg)?;
10968 }
10969 Ok(y)
10970 }
10971
10972 #[allow(clippy::too_many_arguments)]
10978 pub fn qmatvec_nvfp4_fast_prequant_into(
10979 &self,
10980 w: &CudaSlice<u8>,
10981 aq: &CudaSlice<i8>,
10982 ad: &CudaSlice<f32>,
10983 y: &mut CudaSlice<f32>,
10984 m: usize,
10985 in_f: usize,
10986 out_f: usize,
10987 row_bytes: usize,
10988 ) -> Result<(), Box<dyn std::error::Error>> {
10989 assert!(
10990 in_f.is_multiple_of(64),
10991 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
10992 );
10993 if y.len() < m * out_f {
10994 return Err(format!(
10995 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
10996 y.len()
10997 )
10998 .into());
10999 }
11000 let f = self.func("qmatvec_nvfp4_dp4a");
11001 let cfg = LaunchConfig {
11002 grid_dim: (out_f as u32, m as u32, 1),
11003 block_dim: (128, 1, 1),
11004 shared_mem_bytes: 0,
11005 };
11006 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11007 let __s_b = self.gpu.stream();
11008 let mut b = __s_b.launch_builder(&f);
11009 b.arg(w)
11010 .arg(aq)
11011 .arg(ad)
11012 .arg(y)
11013 .arg(&inf)
11014 .arg(&outf)
11015 .arg(&mi)
11016 .arg(&rb);
11017 unsafe {
11018 b.launch(cfg)?;
11019 }
11020 Ok(())
11021 }
11022
11023 #[allow(clippy::too_many_arguments)]
11026 pub fn matvec_f32_qkv_into(
11027 &self,
11028 wq: &CudaSlice<f32>,
11029 wk: &CudaSlice<f32>,
11030 wv: &CudaSlice<f32>,
11031 wg: &CudaSlice<f32>,
11032 x: &CudaSlice<f32>,
11033 yq: &mut CudaSlice<f32>,
11034 yk: &mut CudaSlice<f32>,
11035 yv: &mut CudaSlice<f32>,
11036 yg: &mut CudaSlice<f32>,
11037 in_f: usize,
11038 out_q: usize,
11039 out_kv: usize,
11040 out_g: usize,
11041 ) -> Result<(), Box<dyn std::error::Error>> {
11042 if !in_f.is_multiple_of(4)
11043 || wq.len() != out_q * in_f
11044 || wk.len() != out_kv * in_f
11045 || wv.len() != out_kv * in_f
11046 || wg.len() < out_g * in_f
11047 || x.len() < in_f
11048 || yq.len() < out_q
11049 || yk.len() < out_kv
11050 || yv.len() < out_kv
11051 || (out_g > 0 && yg.len() < out_g)
11052 {
11053 return Err(format!(
11054 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
11055 wq={} wk={} wv={} wg={}",
11056 wq.len(),
11057 wk.len(),
11058 wv.len(),
11059 wg.len()
11060 )
11061 .into());
11062 }
11063 let f = self.func("matvec_f32_qkv");
11064 let cfg = LaunchConfig {
11065 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
11066 block_dim: (128, 1, 1),
11067 shared_mem_bytes: 0,
11068 };
11069 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
11070 let __s_b = self.gpu.stream();
11071 let mut b = __s_b.launch_builder(&f);
11072 b.arg(wq)
11073 .arg(wk)
11074 .arg(wv)
11075 .arg(wg)
11076 .arg(x)
11077 .arg(yq)
11078 .arg(yk)
11079 .arg(yv)
11080 .arg(yg)
11081 .arg(&inf)
11082 .arg(&oq)
11083 .arg(&okv)
11084 .arg(&og);
11085 unsafe {
11086 b.launch(cfg)?;
11087 }
11088 Ok(())
11089 }
11090
11091 #[allow(clippy::too_many_arguments)]
11103 pub fn qmatvec_nvfp4_sel_gu_into(
11104 &self,
11105 gate_bank: &CudaSlice<u8>,
11106 up_bank: &CudaSlice<u8>,
11107 sel: &CudaSlice<i32>,
11108 aq: &CudaSlice<i8>,
11109 ad: &CudaSlice<f32>,
11110 yg: &mut CudaSlice<f32>,
11111 yu: &mut CudaSlice<f32>,
11112 n_sel: usize,
11113 in_f: usize,
11114 out_f: usize,
11115 row_bytes: usize,
11116 expert_stride: usize,
11117 slot_major: bool,
11118 ) -> Result<(), Box<dyn std::error::Error>> {
11119 assert!(
11120 in_f.is_multiple_of(64),
11121 "NVFP4 dp4a requires in_f % 64 == 0"
11122 );
11123 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
11124 return Err("NVFP4 gu sel geometry".into());
11125 }
11126 if !slot_major {
11127 return Err(
11128 "NVFP4 gu sel fusion reads slot-major rows: these banks are block_nvfp4 \
11129 v1 (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
11130 .into(),
11131 );
11132 }
11133 static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
11137 let rpw = *RPW.get_or_init(|| {
11138 std::env::var("MEMRA_NVFP4_SEL_GU_RPW")
11139 .ok()
11140 .and_then(|v| v.parse().ok())
11141 .filter(|r| *r == 2 || *r == 4)
11142 .unwrap_or(1)
11143 });
11144 let rpw = if out_f.is_multiple_of(rpw) { rpw } else { 1 };
11145 static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11150 let wpr =
11151 *WPR.get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_GU_WPR").as_deref() == Ok("1"));
11152 let f = self.func(match (wpr, rpw) {
11153 (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
11154 (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
11155 (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
11156 _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
11157 });
11158 let cfg = LaunchConfig {
11159 grid_dim: if wpr {
11160 (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
11161 } else if rpw == 1 {
11162 ((2 * out_f) as u32, n_sel as u32, 1)
11163 } else {
11164 ((out_f / rpw) as u32, n_sel as u32, 1)
11165 },
11166 block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
11167 shared_mem_bytes: 0,
11168 };
11169 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11170 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11171 let (ars, adrs) = (0i64, 0i64);
11172 let __s_b = self.gpu.stream();
11173 let mut b = __s_b.launch_builder(&f);
11174 b.arg(gate_bank)
11175 .arg(up_bank)
11176 .arg(sel)
11177 .arg(aq)
11178 .arg(ad)
11179 .arg(yg)
11180 .arg(yu)
11181 .arg(&inf)
11182 .arg(&outf)
11183 .arg(&ns)
11184 .arg(&rb)
11185 .arg(&es)
11186 .arg(&ars)
11187 .arg(&adrs);
11188 unsafe {
11189 b.launch(cfg)?;
11190 }
11191 Ok(())
11192 }
11193
11194 #[allow(clippy::too_many_arguments)]
11206 pub fn qmatvec_nvfp4_sel_down8_into(
11207 &self,
11208 bank: &CudaSlice<u8>,
11209 sel: &CudaSlice<i32>,
11210 aq: &CudaSlice<i8>,
11211 ad: &CudaSlice<f32>,
11212 route_w: &CudaSlice<f32>,
11213 md: &CudaSlice<f32>,
11214 dst: &mut CudaSlice<f32>,
11215 n_sel: usize,
11216 in_f: usize,
11217 out_f: usize,
11218 row_bytes: usize,
11219 expert_stride: usize,
11220 act_row_stride: usize,
11221 ad_row_stride: usize,
11222 slot_major: bool,
11223 ) -> Result<(), Box<dyn std::error::Error>> {
11224 if !in_f.is_multiple_of(64)
11225 || n_sel == 0
11226 || n_sel > 8
11227 || (in_f >> 5) > 32
11228 || dst.len() < out_f
11229 || sel.len() < n_sel
11230 || route_w.len() < n_sel
11231 {
11232 return Err(format!(
11233 "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
11234 dst.len()
11235 )
11236 .into());
11237 }
11238 if !slot_major {
11239 return Err(
11240 "NVFP4 sel down8 reads slot-major rows: this shard is block_nvfp4 v1 \
11241 (arm MEMRA_NVFP4_BANK_SM to build slot-major TP banks)"
11242 .into(),
11243 );
11244 }
11245 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
11246 let cfg = LaunchConfig {
11247 grid_dim: (out_f as u32, 1, 1),
11248 block_dim: (32, n_sel as u32, 1),
11249 shared_mem_bytes: 0,
11250 };
11251 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11252 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11253 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
11254 let __s_b = self.gpu.stream();
11255 let mut b = __s_b.launch_builder(&f);
11256 b.arg(bank)
11257 .arg(sel)
11258 .arg(aq)
11259 .arg(ad)
11260 .arg(route_w)
11261 .arg(md)
11262 .arg(dst)
11263 .arg(&inf)
11264 .arg(&outf)
11265 .arg(&ns)
11266 .arg(&rb)
11267 .arg(&es)
11268 .arg(&ars)
11269 .arg(&adrs);
11270 unsafe {
11271 b.launch(cfg)?;
11272 }
11273 Ok(())
11274 }
11275
11276 #[allow(clippy::too_many_arguments)]
11279 pub fn qmatvec_nvfp4_sel_gu_ep_into(
11280 &self,
11281 gate_bank: &CudaSlice<u8>,
11282 up_bank: &CudaSlice<u8>,
11283 sel: &CudaSlice<i32>,
11284 aq: &CudaSlice<i8>,
11285 ad: &CudaSlice<f32>,
11286 yg: &mut CudaSlice<f32>,
11287 yu: &mut CudaSlice<f32>,
11288 n_sel: usize,
11289 in_f: usize,
11290 out_f: usize,
11291 row_bytes: usize,
11292 expert_stride: usize,
11293 owner: usize,
11294 ) -> Result<(), Box<dyn std::error::Error>> {
11295 assert!(
11296 in_f.is_multiple_of(64),
11297 "NVFP4 dp4a requires in_f % 64 == 0"
11298 );
11299 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
11300 return Err("NVFP4 gu ep geometry".into());
11301 }
11302 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
11303 let cfg = LaunchConfig {
11304 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
11305 block_dim: (128, 1, 1),
11306 shared_mem_bytes: 0,
11307 };
11308 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
11309 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11310 let (ars, adrs) = (0i64, 0i64);
11311 let __s_b = self.gpu.stream();
11312 let mut b = __s_b.launch_builder(&f);
11313 b.arg(gate_bank)
11314 .arg(up_bank)
11315 .arg(sel)
11316 .arg(aq)
11317 .arg(ad)
11318 .arg(yg)
11319 .arg(yu)
11320 .arg(&inf)
11321 .arg(&outf)
11322 .arg(&ns)
11323 .arg(&rb)
11324 .arg(&es)
11325 .arg(&ars)
11326 .arg(&adrs)
11327 .arg(&own);
11328 unsafe {
11329 b.launch(cfg)?;
11330 }
11331 Ok(())
11332 }
11333
11334 #[allow(clippy::too_many_arguments)]
11336 pub fn silu_mul_scaled_q8_1_sel_ep_into(
11337 &self,
11338 gate: &CudaSlice<f32>,
11339 up: &CudaSlice<f32>,
11340 gmac: &CudaSlice<f32>,
11341 umac: &CudaSlice<f32>,
11342 sel: &CudaSlice<i32>,
11343 limit: Option<f32>,
11344 out_q: &mut CudaSlice<i8>,
11345 out_d: &mut CudaSlice<f32>,
11346 n_per: usize,
11347 n_sel: usize,
11348 owner: usize,
11349 ) -> Result<(), Box<dyn std::error::Error>> {
11350 if !n_per.is_multiple_of(32)
11351 || out_q.len() < n_sel * n_per
11352 || out_d.len() < n_sel * n_per / 32
11353 {
11354 return Err("NVFP4 silu ep geometry".into());
11355 }
11356 let f = self.func("silu_mul_scaled_q8_1_sel_ep");
11357 let warps = n_sel * n_per / 32;
11358 let cfg = LaunchConfig {
11359 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
11360 block_dim: (128, 1, 1),
11361 shared_mem_bytes: 0,
11362 };
11363 let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
11364 let (lim, has) = match limit {
11365 Some(l) => (l, 1i32),
11366 None => (0.0f32, 0i32),
11367 };
11368 let __s_b = self.gpu.stream();
11369 let mut b = __s_b.launch_builder(&f);
11370 b.arg(gate)
11371 .arg(up)
11372 .arg(gmac)
11373 .arg(umac)
11374 .arg(sel)
11375 .arg(&lim)
11376 .arg(&has)
11377 .arg(out_q)
11378 .arg(out_d)
11379 .arg(&np)
11380 .arg(&ns)
11381 .arg(&own);
11382 unsafe {
11383 b.launch(cfg)?;
11384 }
11385 Ok(())
11386 }
11387
11388 #[allow(clippy::too_many_arguments)]
11390 pub fn qmatvec_nvfp4_sel_down8_ep_into(
11391 &self,
11392 bank: &CudaSlice<u8>,
11393 sel: &CudaSlice<i32>,
11394 aq: &CudaSlice<i8>,
11395 ad: &CudaSlice<f32>,
11396 route_w: &CudaSlice<f32>,
11397 md: &CudaSlice<f32>,
11398 dst: &mut CudaSlice<f32>,
11399 n_sel: usize,
11400 in_f: usize,
11401 out_f: usize,
11402 row_bytes: usize,
11403 expert_stride: usize,
11404 act_row_stride: usize,
11405 ad_row_stride: usize,
11406 owner: usize,
11407 ) -> Result<(), Box<dyn std::error::Error>> {
11408 if !in_f.is_multiple_of(64)
11409 || n_sel == 0
11410 || n_sel > 8
11411 || (in_f >> 5) > 64
11412 || dst.len() < out_f
11413 {
11414 return Err("NVFP4 down8 ep geometry".into());
11415 }
11416 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
11417 let cfg = LaunchConfig {
11418 grid_dim: (out_f as u32, 1, 1),
11419 block_dim: (32, n_sel as u32, 1),
11420 shared_mem_bytes: 0,
11421 };
11422 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
11423 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11424 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
11425 let __s_b = self.gpu.stream();
11426 let mut b = __s_b.launch_builder(&f);
11427 b.arg(bank)
11428 .arg(sel)
11429 .arg(aq)
11430 .arg(ad)
11431 .arg(route_w)
11432 .arg(md)
11433 .arg(dst)
11434 .arg(&inf)
11435 .arg(&outf)
11436 .arg(&ns)
11437 .arg(&rb)
11438 .arg(&es)
11439 .arg(&ars)
11440 .arg(&adrs)
11441 .arg(&own);
11442 unsafe {
11443 b.launch(cfg)?;
11444 }
11445 Ok(())
11446 }
11447
11448 #[allow(clippy::too_many_arguments)] pub fn qmatvec_nvfp4_sel_into(
11463 &self,
11464 bank: &CudaSlice<u8>,
11465 sel: &CudaSlice<i32>,
11466 aq: &CudaSlice<i8>,
11467 ad: &CudaSlice<f32>,
11468 y: &mut CudaSlice<f32>,
11469 n_sel: usize,
11470 in_f: usize,
11471 out_f: usize,
11472 row_bytes: usize,
11473 expert_stride: usize,
11474 act_row_stride: usize,
11475 ad_row_stride: usize,
11476 slot_major: bool,
11477 ) -> Result<(), Box<dyn std::error::Error>> {
11478 assert!(
11479 in_f.is_multiple_of(64),
11480 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
11481 );
11482 if y.len() < n_sel * out_f || sel.len() < n_sel {
11483 return Err(format!(
11484 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
11485 y.len(),
11486 sel.len()
11487 )
11488 .into());
11489 }
11490 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
11497 let mode = *MR.get_or_init(|| {
11498 if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
11499 2
11500 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
11501 1
11502 } else {
11503 0
11504 }
11505 });
11506 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
11507 let mode = if slot_major { 3 } else { mode };
11511 static SM_STREAM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11515 let sm_stream = mode == 3
11516 && *SM_STREAM
11517 .get_or_init(|| std::env::var("MEMRA_NVFP4_SEL_SM_STREAM").as_deref() == Ok("1"))
11518 && row_bytes.is_multiple_of(16)
11519 && in_f <= 4096;
11520 let kname = match (mode, sm_stream) {
11521 (3, true) => "qmatvec_nvfp4_dp4a_sel_v2s",
11522 (3, false) => "qmatvec_nvfp4_dp4a_sel_v2",
11523 (2, _) => "qmatvec_nvfp4_dp4a_sel_stream",
11524 (1, _) => "qmatvec_nvfp4_dp4a_sel_mr4",
11525 _ => "qmatvec_nvfp4_dp4a_sel",
11526 };
11527 {
11533 static SEEN_SEL: std::sync::Mutex<Vec<(&'static str, usize, usize)>> =
11534 std::sync::Mutex::new(Vec::new());
11535 let combo = (kname, in_f, out_f);
11536 let mut seen = SEEN_SEL.lock().unwrap();
11537 if !seen.contains(&combo) {
11538 seen.push(combo);
11539 eprintln!(
11540 "[nvfp4-sel] kernel={kname} slot_major={slot_major} in_f={in_f} \
11541 out_f={out_f} nsb={} row_bytes={row_bytes}",
11542 in_f >> 5
11543 );
11544 }
11545 }
11546 let f = self.func(kname);
11547 let nsb = in_f >> 5;
11552 let fit_block: u32 = if (mode == 0 || mode == 3) && !sm_stream && nsb <= 32 {
11553 32
11554 } else if mode == 1 {
11555 512
11556 } else {
11557 128
11558 };
11559 let cfg = LaunchConfig {
11560 grid_dim: (
11561 if sm_stream {
11562 (out_f as u32).div_ceil(8)
11563 } else {
11564 match mode {
11565 2 => (out_f as u32).div_ceil(16),
11566 1 => (out_f as u32).div_ceil(4),
11567 _ => out_f as u32,
11568 }
11569 },
11570 n_sel as u32,
11571 1,
11572 ),
11573 block_dim: (fit_block, 1, 1),
11574 shared_mem_bytes: 0,
11575 };
11576 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11577 let (rb, es, ars, adrs) = (
11578 row_bytes as i64,
11579 expert_stride as i64,
11580 act_row_stride as i64,
11581 ad_row_stride as i64,
11582 );
11583 let __s_b = self.gpu.stream();
11584 let mut b = __s_b.launch_builder(&f);
11585 b.arg(bank)
11586 .arg(sel)
11587 .arg(aq)
11588 .arg(ad)
11589 .arg(y)
11590 .arg(&inf)
11591 .arg(&outf)
11592 .arg(&ns)
11593 .arg(&rb)
11594 .arg(&es)
11595 .arg(&ars)
11596 .arg(&adrs);
11597 unsafe {
11598 b.launch(cfg)?;
11599 }
11600 Ok(())
11601 }
11602
11603 #[allow(clippy::too_many_arguments)]
11606 pub fn qmatvec_nvfp4_bf16_sel_dual_rows_into(
11607 &self,
11608 gate_bank: &CudaSlice<u8>,
11609 up_bank: &CudaSlice<u8>,
11610 sel: &CudaSlice<i32>,
11611 token_rows: &CudaSlice<i32>,
11612 x_bf16: &CudaSlice<u8>,
11613 gate_out: &mut CudaSlice<f32>,
11614 up_out: &mut CudaSlice<f32>,
11615 n_sel: usize,
11616 in_f: usize,
11617 out_f: usize,
11618 row_bytes: usize,
11619 expert_stride: usize,
11620 tokens: usize,
11621 ) -> Result<(), Box<dyn std::error::Error>> {
11622 if !in_f.is_multiple_of(64)
11623 || sel.len() < n_sel
11624 || token_rows.len() < n_sel
11625 || gate_out.len() < n_sel * out_f
11626 || up_out.len() < n_sel * out_f
11627 || x_bf16.len() < 2 * in_f * tokens
11628 {
11629 return Err(format!(
11630 "W4A16 NVFP4 dual selected rows geometry sel={} token_rows={} x={} gate={} up={} \
11631 n_sel={n_sel} tokens={tokens} in={in_f} out={out_f}",
11632 sel.len(),
11633 token_rows.len(),
11634 x_bf16.len(),
11635 gate_out.len(),
11636 up_out.len(),
11637 )
11638 .into());
11639 }
11640 let adjacent_rows = tokens > 1;
11641 let f = if adjacent_rows {
11642 self.func("qmatvec_nvfp4_bf16_sel_quad_rows")
11643 } else {
11644 self.func("qmatvec_nvfp4_bf16_sel_dual_rows")
11645 };
11646 let cfg = LaunchConfig {
11647 grid_dim: (
11648 if adjacent_rows {
11649 out_f.div_ceil(2) as u32
11650 } else {
11651 (2 * out_f) as u32
11652 },
11653 n_sel as u32,
11654 1,
11655 ),
11656 block_dim: (256, 1, 1),
11657 shared_mem_bytes: 0,
11658 };
11659 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11660 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11661 let __s_b = self.gpu.stream();
11662 let mut b = __s_b.launch_builder(&f);
11663 b.arg(gate_bank)
11664 .arg(up_bank)
11665 .arg(sel)
11666 .arg(token_rows)
11667 .arg(x_bf16)
11668 .arg(gate_out)
11669 .arg(up_out)
11670 .arg(&inf)
11671 .arg(&outf)
11672 .arg(&ns)
11673 .arg(&rb)
11674 .arg(&es);
11675 unsafe {
11676 b.launch(cfg)?;
11677 }
11678 Ok(())
11679 }
11680
11681 #[allow(clippy::too_many_arguments)]
11684 pub fn qmatvec_nvfp4_bf16_ep_dual_slots_into(
11685 &self,
11686 gate_bank: &CudaSlice<u8>,
11687 up_bank: &CudaSlice<u8>,
11688 sel: &CudaSlice<i32>,
11689 x_bf16: &CudaSlice<u8>,
11690 gate_out: &mut CudaSlice<f32>,
11691 up_out: &mut CudaSlice<f32>,
11692 n_pairs: usize,
11693 top_k: usize,
11694 in_f: usize,
11695 out_f: usize,
11696 owner_start: usize,
11697 owner_end: usize,
11698 row_bytes: usize,
11699 expert_stride: usize,
11700 ) -> Result<(), Box<dyn std::error::Error>> {
11701 let tokens = n_pairs.div_ceil(top_k);
11702 if top_k == 0
11703 || owner_start >= owner_end
11704 || !in_f.is_multiple_of(64)
11705 || sel.len() < n_pairs
11706 || gate_out.len() < n_pairs * out_f
11707 || up_out.len() < n_pairs * out_f
11708 || x_bf16.len() < 2 * in_f * tokens
11709 {
11710 return Err(format!(
11711 "W4A16 NVFP4 device EP dual-slot geometry sel={} x={} gate={} up={} \
11712 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11713 owner={owner_start}..{owner_end}",
11714 sel.len(),
11715 x_bf16.len(),
11716 gate_out.len(),
11717 up_out.len(),
11718 )
11719 .into());
11720 }
11721 let pair_parallel = tokens > 1;
11722 let f = if pair_parallel {
11723 self.func("qmatvec_nvfp4_bf16_ep_quad_pairs")
11724 } else {
11725 self.func("qmatvec_nvfp4_bf16_ep_dual_slots")
11726 };
11727 let cfg = LaunchConfig {
11728 grid_dim: (
11729 if pair_parallel {
11730 out_f.div_ceil(2) as u32
11731 } else {
11732 (2 * out_f) as u32
11733 },
11734 if pair_parallel { n_pairs as u32 } else { 1 },
11735 1,
11736 ),
11737 block_dim: (256, 1, 1),
11738 shared_mem_bytes: 0,
11739 };
11740 let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11741 let (os, oe) = (owner_start as i32, owner_end as i32);
11742 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11743 let __s_b = self.gpu.stream();
11744 let mut b = __s_b.launch_builder(&f);
11745 b.arg(gate_bank)
11746 .arg(up_bank)
11747 .arg(sel)
11748 .arg(x_bf16)
11749 .arg(gate_out)
11750 .arg(up_out)
11751 .arg(&inf)
11752 .arg(&outf)
11753 .arg(&np)
11754 .arg(&tk)
11755 .arg(&os)
11756 .arg(&oe)
11757 .arg(&rb)
11758 .arg(&es);
11759 unsafe {
11760 b.launch(cfg)?;
11761 }
11762 Ok(())
11763 }
11764
11765 #[allow(clippy::too_many_arguments)]
11767 pub fn qmatvec_nvfp4_q8_ep_dual_slots_into(
11768 &self,
11769 gate_bank: &CudaSlice<u8>,
11770 up_bank: &CudaSlice<u8>,
11771 sel: &CudaSlice<i32>,
11772 aq: &CudaSlice<i8>,
11773 ad: &CudaSlice<f32>,
11774 gate_out: &mut CudaSlice<f32>,
11775 up_out: &mut CudaSlice<f32>,
11776 n_pairs: usize,
11777 top_k: usize,
11778 in_f: usize,
11779 out_f: usize,
11780 owner_start: usize,
11781 owner_end: usize,
11782 row_bytes: usize,
11783 expert_stride: usize,
11784 ) -> Result<(), Box<dyn std::error::Error>> {
11785 let tokens = n_pairs.div_ceil(top_k);
11786 if top_k == 0
11787 || owner_start >= owner_end
11788 || !in_f.is_multiple_of(64)
11789 || sel.len() < n_pairs
11790 || aq.len() < tokens * in_f
11791 || ad.len() < tokens * (in_f / 32)
11792 || gate_out.len() < n_pairs * out_f
11793 || up_out.len() < n_pairs * out_f
11794 {
11795 return Err(format!(
11796 "W4A8 NVFP4 device EP gate/up geometry sel={} aq={} ad={} gate={} up={} \
11797 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11798 owner={owner_start}..{owner_end}",
11799 sel.len(),
11800 aq.len(),
11801 ad.len(),
11802 gate_out.len(),
11803 up_out.len(),
11804 )
11805 .into());
11806 }
11807 let f = self.func("qmatvec_nvfp4_q8_ep_dual_slots");
11808 let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
11809 let cfg = LaunchConfig {
11810 grid_dim: (out_f as u32, 1, 1),
11811 block_dim: (threads, 1, 1),
11812 shared_mem_bytes: 0,
11813 };
11814 let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11815 let (os, oe) = (owner_start as i32, owner_end as i32);
11816 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11817 let __s_b = self.gpu.stream();
11818 let mut b = __s_b.launch_builder(&f);
11819 b.arg(gate_bank)
11820 .arg(up_bank)
11821 .arg(sel)
11822 .arg(aq)
11823 .arg(ad)
11824 .arg(gate_out)
11825 .arg(up_out)
11826 .arg(&inf)
11827 .arg(&outf)
11828 .arg(&np)
11829 .arg(&tk)
11830 .arg(&os)
11831 .arg(&oe)
11832 .arg(&rb)
11833 .arg(&es);
11834 unsafe {
11835 b.launch(cfg)?;
11836 }
11837 Ok(())
11838 }
11839
11840 #[allow(clippy::too_many_arguments)]
11843 pub fn qmatvec_nvfp4_q8_ep_paired_slots_into(
11844 &self,
11845 gate_bank: &CudaSlice<u8>,
11846 up_bank: &CudaSlice<u8>,
11847 sel: &CudaSlice<i32>,
11848 aq: &CudaSlice<i8>,
11849 ad: &CudaSlice<f32>,
11850 gate_out: &mut CudaSlice<f32>,
11851 up_out: &mut CudaSlice<f32>,
11852 n_pairs: usize,
11853 top_k: usize,
11854 in_f: usize,
11855 out_f: usize,
11856 owner_start: usize,
11857 owner_end: usize,
11858 row_bytes: usize,
11859 expert_stride: usize,
11860 ) -> Result<(), Box<dyn std::error::Error>> {
11861 let tokens = n_pairs.div_ceil(top_k);
11862 if top_k == 0
11863 || owner_start >= owner_end
11864 || !in_f.is_multiple_of(64)
11865 || sel.len() < n_pairs
11866 || aq.len() < tokens * in_f
11867 || ad.len() < tokens * (in_f / 32)
11868 || gate_out.len() < n_pairs * out_f
11869 || up_out.len() < n_pairs * out_f
11870 {
11871 return Err(format!(
11872 "W4A8 NVFP4 paired gate/up geometry sel={} aq={} ad={} gate={} up={} \
11873 pairs={n_pairs} top_k={top_k} in={in_f} out={out_f} \
11874 owner={owner_start}..{owner_end}",
11875 sel.len(),
11876 aq.len(),
11877 ad.len(),
11878 gate_out.len(),
11879 up_out.len(),
11880 )
11881 .into());
11882 }
11883 let f = self.func("qmatvec_nvfp4_q8_ep_paired_slots");
11884 let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
11885 let cfg = LaunchConfig {
11886 grid_dim: (out_f as u32, 1, 1),
11887 block_dim: (threads, 1, 1),
11888 shared_mem_bytes: 0,
11889 };
11890 let (inf, outf, np, tk) = (in_f as i32, out_f as i32, n_pairs as i32, top_k as i32);
11891 let (os, oe) = (owner_start as i32, owner_end as i32);
11892 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11893 let __s_b = self.gpu.stream();
11894 let mut b = __s_b.launch_builder(&f);
11895 b.arg(gate_bank)
11896 .arg(up_bank)
11897 .arg(sel)
11898 .arg(aq)
11899 .arg(ad)
11900 .arg(gate_out)
11901 .arg(up_out)
11902 .arg(&inf)
11903 .arg(&outf)
11904 .arg(&np)
11905 .arg(&tk)
11906 .arg(&os)
11907 .arg(&oe)
11908 .arg(&rb)
11909 .arg(&es);
11910 unsafe {
11911 b.launch(cfg)?;
11912 }
11913 Ok(())
11914 }
11915
11916 #[allow(clippy::too_many_arguments)]
11918 pub fn qmatvec_nvfp4_bf16_sel_down_rows_raw(
11919 &self,
11920 bank: &CudaSlice<u8>,
11921 sel: &CudaSlice<i32>,
11922 global_pairs: &CudaSlice<i32>,
11923 activation_bf16: &CudaSlice<u8>,
11924 macros_down: &CudaSlice<f32>,
11925 dst_raw: u64,
11926 n_sel: usize,
11927 in_f: usize,
11928 out_f: usize,
11929 row_bytes: usize,
11930 expert_stride: usize,
11931 total_pairs: usize,
11932 ) -> Result<(), Box<dyn std::error::Error>> {
11933 if !in_f.is_multiple_of(64)
11934 || sel.len() < n_sel
11935 || global_pairs.len() < n_sel
11936 || activation_bf16.len() < 2 * n_sel * in_f
11937 || dst_raw == 0
11938 {
11939 return Err(format!(
11940 "W4A16 NVFP4 down rows geometry sel={} pairs={} act={} dst_raw={dst_raw:#x} \
11941 n_sel={n_sel} total_pairs={total_pairs} in={in_f} out={out_f}",
11942 sel.len(),
11943 global_pairs.len(),
11944 activation_bf16.len(),
11945 )
11946 .into());
11947 }
11948 let f = self.func("qmatvec_nvfp4_bf16_sel_down_rows");
11949 let cfg = LaunchConfig {
11950 grid_dim: (out_f.div_ceil(2) as u32, n_sel as u32, 1),
11951 block_dim: (256, 1, 1),
11952 shared_mem_bytes: 0,
11953 };
11954 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
11955 let (rb, es) = (row_bytes as i64, expert_stride as i64);
11956 let __s_b = self.gpu.stream();
11957 let mut b = __s_b.launch_builder(&f);
11958 b.arg(bank)
11959 .arg(sel)
11960 .arg(global_pairs)
11961 .arg(activation_bf16)
11962 .arg(macros_down)
11963 .arg(&dst_raw)
11964 .arg(&inf)
11965 .arg(&outf)
11966 .arg(&ns)
11967 .arg(&rb)
11968 .arg(&es);
11969 unsafe {
11970 b.launch(cfg)?;
11971 }
11972 Ok(())
11973 }
11974
11975 #[allow(clippy::too_many_arguments)]
11978 pub fn qmatvec_nvfp4_bf16_ep_down_slots_raw(
11979 &self,
11980 bank: &CudaSlice<u8>,
11981 sel: &CudaSlice<i32>,
11982 activation_bf16: &CudaSlice<u8>,
11983 macros_down: &CudaSlice<f32>,
11984 dst_raw: u64,
11985 n_pairs: usize,
11986 in_f: usize,
11987 out_f: usize,
11988 owner_start: usize,
11989 owner_end: usize,
11990 row_bytes: usize,
11991 expert_stride: usize,
11992 ) -> Result<(), Box<dyn std::error::Error>> {
11993 if owner_start >= owner_end
11994 || !in_f.is_multiple_of(64)
11995 || sel.len() < n_pairs
11996 || activation_bf16.len() < 2 * n_pairs * in_f
11997 || dst_raw == 0
11998 {
11999 return Err(format!(
12000 "W4A16 NVFP4 device EP down-slot geometry sel={} act={} dst_raw={dst_raw:#x} \
12001 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
12002 sel.len(),
12003 activation_bf16.len(),
12004 )
12005 .into());
12006 }
12007 let f = self.func("qmatvec_nvfp4_bf16_ep_down_slots");
12008 let cfg = LaunchConfig {
12009 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
12010 block_dim: (256, 1, 1),
12011 shared_mem_bytes: 0,
12012 };
12013 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12014 let (os, oe) = (owner_start as i32, owner_end as i32);
12015 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12016 let __s_b = self.gpu.stream();
12017 let mut b = __s_b.launch_builder(&f);
12018 b.arg(bank)
12019 .arg(sel)
12020 .arg(activation_bf16)
12021 .arg(macros_down)
12022 .arg(&dst_raw)
12023 .arg(&inf)
12024 .arg(&outf)
12025 .arg(&np)
12026 .arg(&os)
12027 .arg(&oe)
12028 .arg(&rb)
12029 .arg(&es);
12030 unsafe {
12031 b.launch(cfg)?;
12032 }
12033 Ok(())
12034 }
12035
12036 #[allow(clippy::too_many_arguments)]
12038 pub fn qmatvec_nvfp4_bf16_ep_down_pairs_raw(
12039 &self,
12040 bank: &CudaSlice<u8>,
12041 sel: &CudaSlice<i32>,
12042 activation_bf16: &CudaSlice<u8>,
12043 macros_down: &CudaSlice<f32>,
12044 dst_raw: u64,
12045 n_pairs: usize,
12046 in_f: usize,
12047 out_f: usize,
12048 owner_start: usize,
12049 owner_end: usize,
12050 row_bytes: usize,
12051 expert_stride: usize,
12052 ) -> Result<(), Box<dyn std::error::Error>> {
12053 if owner_start >= owner_end
12054 || !in_f.is_multiple_of(64)
12055 || sel.len() < n_pairs
12056 || activation_bf16.len() < 2 * n_pairs * in_f
12057 || dst_raw == 0
12058 {
12059 return Err(format!(
12060 "W4A16 NVFP4 device EP down-pair geometry sel={} act={} dst_raw={dst_raw:#x} \
12061 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
12062 sel.len(),
12063 activation_bf16.len(),
12064 )
12065 .into());
12066 }
12067 let f = self.func("qmatvec_nvfp4_bf16_ep_down_pairs");
12068 let cfg = LaunchConfig {
12069 grid_dim: (out_f.div_ceil(2) as u32, n_pairs as u32, 1),
12070 block_dim: (256, 1, 1),
12071 shared_mem_bytes: 0,
12072 };
12073 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12074 let (os, oe) = (owner_start as i32, owner_end as i32);
12075 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12076 let __s_b = self.gpu.stream();
12077 let mut b = __s_b.launch_builder(&f);
12078 b.arg(bank)
12079 .arg(sel)
12080 .arg(activation_bf16)
12081 .arg(macros_down)
12082 .arg(&dst_raw)
12083 .arg(&inf)
12084 .arg(&outf)
12085 .arg(&np)
12086 .arg(&os)
12087 .arg(&oe)
12088 .arg(&rb)
12089 .arg(&es);
12090 unsafe {
12091 b.launch(cfg)?;
12092 }
12093 Ok(())
12094 }
12095
12096 #[allow(clippy::too_many_arguments)]
12098 pub fn silu_mul_scaled_host_expf_bf16_sel_into(
12099 &self,
12100 gate: &CudaSlice<f32>,
12101 up: &CudaSlice<f32>,
12102 gate_macros: &CudaSlice<f32>,
12103 up_macros: &CudaSlice<f32>,
12104 sel: &CudaSlice<i32>,
12105 limit: Option<f32>,
12106 output_bf16: &mut CudaSlice<u8>,
12107 n_per: usize,
12108 n_sel: usize,
12109 ) -> Result<(), Box<dyn std::error::Error>> {
12110 let n = n_per * n_sel;
12111 if sel.len() < n_sel || gate.len() < n || up.len() < n || output_bf16.len() < 2 * n {
12112 return Err(format!(
12113 "W4A16 selected activation geometry sel={} gate={} up={} out={} \
12114 n_per={n_per} n_sel={n_sel}",
12115 sel.len(),
12116 gate.len(),
12117 up.len(),
12118 output_bf16.len(),
12119 )
12120 .into());
12121 }
12122 let (limit, has_limit) = match limit {
12123 Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
12124 Some(limit) => {
12125 return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
12126 }
12127 None => (0.0f32, 0i32),
12128 };
12129 let f = self.func("silu_mul_scaled_host_expf_bf16_sel");
12130 let cfg = LaunchConfig::for_num_elems(n as u32);
12131 let (np, ns) = (n_per as i32, n_sel as i32);
12132 let __s_b = self.gpu.stream();
12133 let mut b = __s_b.launch_builder(&f);
12134 b.arg(gate)
12135 .arg(up)
12136 .arg(gate_macros)
12137 .arg(up_macros)
12138 .arg(sel)
12139 .arg(&limit)
12140 .arg(&has_limit)
12141 .arg(output_bf16)
12142 .arg(&np)
12143 .arg(&ns);
12144 unsafe {
12145 b.launch(cfg)?;
12146 }
12147 Ok(())
12148 }
12149
12150 #[allow(clippy::too_many_arguments)]
12153 pub fn silu_mul_scaled_host_expf_bf16_ep_slots_into(
12154 &self,
12155 gate: &CudaSlice<f32>,
12156 up: &CudaSlice<f32>,
12157 gate_macros: &CudaSlice<f32>,
12158 up_macros: &CudaSlice<f32>,
12159 sel: &CudaSlice<i32>,
12160 owner_start: usize,
12161 owner_end: usize,
12162 limit: Option<f32>,
12163 output_bf16: &mut CudaSlice<u8>,
12164 n_per: usize,
12165 n_pairs: usize,
12166 ) -> Result<(), Box<dyn std::error::Error>> {
12167 let n = n_per * n_pairs;
12168 if owner_start >= owner_end
12169 || sel.len() < n_pairs
12170 || gate.len() < n
12171 || up.len() < n
12172 || output_bf16.len() < 2 * n
12173 {
12174 return Err(format!(
12175 "W4A16 device EP activation geometry sel={} gate={} up={} out={} \
12176 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
12177 sel.len(),
12178 gate.len(),
12179 up.len(),
12180 output_bf16.len(),
12181 )
12182 .into());
12183 }
12184 let (limit, has_limit) = match limit {
12185 Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
12186 Some(limit) => {
12187 return Err(format!("W4A16 selected activation limit {limit} is invalid").into());
12188 }
12189 None => (0.0f32, 0i32),
12190 };
12191 let f = self.func("silu_mul_scaled_host_expf_bf16_ep_slots");
12192 let cfg = LaunchConfig::for_num_elems(n as u32);
12193 let (np, pairs) = (n_per as i32, n_pairs as i32);
12194 let (os, oe) = (owner_start as i32, owner_end as i32);
12195 let __s_b = self.gpu.stream();
12196 let mut b = __s_b.launch_builder(&f);
12197 b.arg(gate)
12198 .arg(up)
12199 .arg(gate_macros)
12200 .arg(up_macros)
12201 .arg(sel)
12202 .arg(&limit)
12203 .arg(&has_limit)
12204 .arg(output_bf16)
12205 .arg(&np)
12206 .arg(&pairs)
12207 .arg(&os)
12208 .arg(&oe);
12209 unsafe {
12210 b.launch(cfg)?;
12211 }
12212 Ok(())
12213 }
12214
12215 #[allow(clippy::too_many_arguments)]
12217 pub fn silu_mul_scaled_host_expf_q8_ep_slots_into(
12218 &self,
12219 gate: &CudaSlice<f32>,
12220 up: &CudaSlice<f32>,
12221 gate_macros: &CudaSlice<f32>,
12222 up_macros: &CudaSlice<f32>,
12223 sel: &CudaSlice<i32>,
12224 owner_start: usize,
12225 owner_end: usize,
12226 limit: Option<f32>,
12227 output_q8: &mut CudaSlice<i8>,
12228 output_scales: &mut CudaSlice<f32>,
12229 n_per: usize,
12230 n_pairs: usize,
12231 ) -> Result<(), Box<dyn std::error::Error>> {
12232 let n = n_per * n_pairs;
12233 if owner_start >= owner_end
12234 || !n_per.is_multiple_of(32)
12235 || sel.len() < n_pairs
12236 || gate.len() < n
12237 || up.len() < n
12238 || output_q8.len() < n
12239 || output_scales.len() < n / 32
12240 {
12241 return Err(format!(
12242 "W4A8 device EP activation geometry sel={} gate={} up={} q8={} scales={} \
12243 n_per={n_per} pairs={n_pairs} owner={owner_start}..{owner_end}",
12244 sel.len(),
12245 gate.len(),
12246 up.len(),
12247 output_q8.len(),
12248 output_scales.len(),
12249 )
12250 .into());
12251 }
12252 let (limit, has_limit) = match limit {
12253 Some(limit) if limit.is_finite() && limit > 1e-6 => (limit, 1i32),
12254 Some(limit) => {
12255 return Err(format!("W4A8 selected activation limit {limit} is invalid").into());
12256 }
12257 None => (0.0f32, 0i32),
12258 };
12259 let f = self.func("silu_mul_scaled_host_expf_q8_ep_slots");
12260 let warps = n / 32;
12261 let cfg = LaunchConfig {
12262 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
12263 block_dim: (128, 1, 1),
12264 shared_mem_bytes: 0,
12265 };
12266 let (np, pairs) = (n_per as i32, n_pairs as i32);
12267 let (os, oe) = (owner_start as i32, owner_end as i32);
12268 let __s_b = self.gpu.stream();
12269 let mut b = __s_b.launch_builder(&f);
12270 b.arg(gate)
12271 .arg(up)
12272 .arg(gate_macros)
12273 .arg(up_macros)
12274 .arg(sel)
12275 .arg(&limit)
12276 .arg(&has_limit)
12277 .arg(output_q8)
12278 .arg(output_scales)
12279 .arg(&np)
12280 .arg(&pairs)
12281 .arg(&os)
12282 .arg(&oe);
12283 unsafe {
12284 b.launch(cfg)?;
12285 }
12286 Ok(())
12287 }
12288
12289 #[allow(clippy::too_many_arguments)]
12292 pub fn qmatvec_nvfp4_bf16_sel_down_fma_into(
12293 &self,
12294 bank: &CudaSlice<u8>,
12295 sel: &CudaSlice<i32>,
12296 activation_bf16: &CudaSlice<u8>,
12297 route_weights: &CudaSlice<f32>,
12298 macros_down: &CudaSlice<f32>,
12299 dst: &mut cudarc::driver::CudaViewMut<f32>,
12300 n_sel: usize,
12301 in_f: usize,
12302 out_f: usize,
12303 row_bytes: usize,
12304 expert_stride: usize,
12305 ) -> Result<(), Box<dyn std::error::Error>> {
12306 if !in_f.is_multiple_of(64)
12307 || sel.len() < n_sel
12308 || route_weights.len() < n_sel
12309 || activation_bf16.len() < 2 * n_sel * in_f
12310 || dst.len() < out_f
12311 {
12312 return Err(format!(
12313 "W4A16 NVFP4 down selected geometry sel={} act={} weights={} dst={} \
12314 n_sel={n_sel} in={in_f} out={out_f}",
12315 sel.len(),
12316 activation_bf16.len(),
12317 route_weights.len(),
12318 dst.len(),
12319 )
12320 .into());
12321 }
12322 let f = self.func("qmatvec_nvfp4_bf16_sel_down_fma");
12323 let cfg = LaunchConfig {
12324 grid_dim: (out_f as u32, 1, 1),
12325 block_dim: (256, 1, 1),
12326 shared_mem_bytes: 0,
12327 };
12328 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
12329 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12330 let __s_b = self.gpu.stream();
12331 let mut b = __s_b.launch_builder(&f);
12332 b.arg(bank)
12333 .arg(sel)
12334 .arg(activation_bf16)
12335 .arg(route_weights)
12336 .arg(macros_down)
12337 .arg(dst)
12338 .arg(&inf)
12339 .arg(&outf)
12340 .arg(&ns)
12341 .arg(&rb)
12342 .arg(&es);
12343 unsafe {
12344 b.launch(cfg)?;
12345 }
12346 Ok(())
12347 }
12348
12349 #[allow(clippy::too_many_arguments)]
12351 pub fn qmatvec_nvfp4_bf16_ep_down_fma_into(
12352 &self,
12353 bank: &CudaSlice<u8>,
12354 sel: &CudaSlice<i32>,
12355 activation_bf16: &CudaSlice<u8>,
12356 route_weights: &CudaSlice<f32>,
12357 macros_down: &CudaSlice<f32>,
12358 dst: &mut cudarc::driver::CudaViewMut<f32>,
12359 n_pairs: usize,
12360 in_f: usize,
12361 out_f: usize,
12362 owner_start: usize,
12363 owner_end: usize,
12364 row_bytes: usize,
12365 expert_stride: usize,
12366 ) -> Result<(), Box<dyn std::error::Error>> {
12367 if owner_start >= owner_end
12368 || !in_f.is_multiple_of(64)
12369 || sel.len() < n_pairs
12370 || route_weights.len() < n_pairs
12371 || activation_bf16.len() < 2 * n_pairs * in_f
12372 || dst.len() < out_f
12373 {
12374 return Err(format!(
12375 "W4A16 device EP down-FMA geometry sel={} act={} weights={} dst={} \
12376 pairs={n_pairs} in={in_f} out={out_f} owner={owner_start}..{owner_end}",
12377 sel.len(),
12378 activation_bf16.len(),
12379 route_weights.len(),
12380 dst.len(),
12381 )
12382 .into());
12383 }
12384 let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
12385 let cfg = LaunchConfig {
12386 grid_dim: (out_f as u32, 1, 1),
12387 block_dim: (256, 1, 1),
12388 shared_mem_bytes: 0,
12389 };
12390 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12391 let (os, oe) = (owner_start as i32, owner_end as i32);
12392 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12393 let __s_b = self.gpu.stream();
12394 let mut b = __s_b.launch_builder(&f);
12395 b.arg(bank)
12396 .arg(sel)
12397 .arg(activation_bf16)
12398 .arg(route_weights)
12399 .arg(macros_down)
12400 .arg(dst)
12401 .arg(&inf)
12402 .arg(&outf)
12403 .arg(&np)
12404 .arg(&os)
12405 .arg(&oe)
12406 .arg(&rb)
12407 .arg(&es);
12408 unsafe {
12409 b.launch(cfg)?;
12410 }
12411 Ok(())
12412 }
12413
12414 #[allow(clippy::too_many_arguments)]
12417 pub fn qmatvec_nvfp4_bf16_ep_down_fma_raw(
12418 &self,
12419 bank: &CudaSlice<u8>,
12420 sel: &CudaSlice<i32>,
12421 activation_bf16: &CudaSlice<u8>,
12422 route_weights: &CudaSlice<f32>,
12423 macros_down: &CudaSlice<f32>,
12424 dst_raw: u64,
12425 n_pairs: usize,
12426 in_f: usize,
12427 out_f: usize,
12428 owner_start: usize,
12429 owner_end: usize,
12430 row_bytes: usize,
12431 expert_stride: usize,
12432 ) -> Result<(), Box<dyn std::error::Error>> {
12433 if dst_raw == 0
12434 || owner_start >= owner_end
12435 || !in_f.is_multiple_of(64)
12436 || sel.len() < n_pairs
12437 || route_weights.len() < n_pairs
12438 || activation_bf16.len() < 2 * n_pairs * in_f
12439 {
12440 return Err(format!(
12441 "W4A16 device EP raw down-FMA geometry sel={} act={} weights={} \
12442 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12443 owner={owner_start}..{owner_end}",
12444 sel.len(),
12445 activation_bf16.len(),
12446 route_weights.len(),
12447 )
12448 .into());
12449 }
12450 let f = self.func("qmatvec_nvfp4_bf16_ep_down_fma");
12451 let cfg = LaunchConfig {
12452 grid_dim: (out_f as u32, 1, 1),
12453 block_dim: (256, 1, 1),
12454 shared_mem_bytes: 0,
12455 };
12456 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12457 let (os, oe) = (owner_start as i32, owner_end as i32);
12458 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12459 let __s_b = self.gpu.stream();
12460 let mut b = __s_b.launch_builder(&f);
12461 b.arg(bank)
12462 .arg(sel)
12463 .arg(activation_bf16)
12464 .arg(route_weights)
12465 .arg(macros_down)
12466 .arg(&dst_raw)
12467 .arg(&inf)
12468 .arg(&outf)
12469 .arg(&np)
12470 .arg(&os)
12471 .arg(&oe)
12472 .arg(&rb)
12473 .arg(&es);
12474 unsafe {
12475 b.launch(cfg)?;
12476 }
12477 Ok(())
12478 }
12479
12480 #[allow(clippy::too_many_arguments)]
12483 pub fn qmatvec_nvfp4_q8_ep_down_slots_raw(
12484 &self,
12485 bank: &CudaSlice<u8>,
12486 sel: &CudaSlice<i32>,
12487 aq: &CudaSlice<i8>,
12488 ad: &CudaSlice<f32>,
12489 macros_down: &CudaSlice<f32>,
12490 dst_raw: u64,
12491 n_pairs: usize,
12492 in_f: usize,
12493 out_f: usize,
12494 owner_start: usize,
12495 owner_end: usize,
12496 row_bytes: usize,
12497 expert_stride: usize,
12498 ) -> Result<(), Box<dyn std::error::Error>> {
12499 if dst_raw == 0
12500 || owner_start >= owner_end
12501 || !in_f.is_multiple_of(64)
12502 || sel.len() < n_pairs
12503 || aq.len() < n_pairs * in_f
12504 || ad.len() < n_pairs * (in_f / 32)
12505 {
12506 return Err(format!(
12507 "W4A8 device EP raw down-slot geometry sel={} aq={} ad={} \
12508 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12509 owner={owner_start}..{owner_end}",
12510 sel.len(),
12511 aq.len(),
12512 ad.len(),
12513 )
12514 .into());
12515 }
12516 let f = self.func("qmatvec_nvfp4_q8_ep_down_slots");
12517 let threads = ((in_f / 32).div_ceil(32) * 32).clamp(32, 256) as u32;
12518 let cfg = LaunchConfig {
12519 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
12520 block_dim: (threads, 1, 1),
12521 shared_mem_bytes: 0,
12522 };
12523 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12524 let (os, oe) = (owner_start as i32, owner_end as i32);
12525 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12526 let __s_b = self.gpu.stream();
12527 let mut b = __s_b.launch_builder(&f);
12528 b.arg(bank)
12529 .arg(sel)
12530 .arg(aq)
12531 .arg(ad)
12532 .arg(macros_down)
12533 .arg(&dst_raw)
12534 .arg(&inf)
12535 .arg(&outf)
12536 .arg(&np)
12537 .arg(&os)
12538 .arg(&oe)
12539 .arg(&rb)
12540 .arg(&es);
12541 unsafe {
12542 b.launch(cfg)?;
12543 }
12544 Ok(())
12545 }
12546
12547 #[allow(clippy::too_many_arguments)]
12549 pub fn qmatvec_nvfp4_q8_ep_down_fma_raw(
12550 &self,
12551 bank: &CudaSlice<u8>,
12552 sel: &CudaSlice<i32>,
12553 aq: &CudaSlice<i8>,
12554 ad: &CudaSlice<f32>,
12555 route_weights: &CudaSlice<f32>,
12556 macros_down: &CudaSlice<f32>,
12557 dst_raw: u64,
12558 n_pairs: usize,
12559 in_f: usize,
12560 out_f: usize,
12561 owner_start: usize,
12562 owner_end: usize,
12563 row_bytes: usize,
12564 expert_stride: usize,
12565 ) -> Result<(), Box<dyn std::error::Error>> {
12566 if dst_raw == 0
12567 || owner_start >= owner_end
12568 || !in_f.is_multiple_of(64)
12569 || sel.len() < n_pairs
12570 || aq.len() < n_pairs * in_f
12571 || ad.len() < n_pairs * (in_f / 32)
12572 || route_weights.len() < n_pairs
12573 {
12574 return Err(format!(
12575 "W4A8 device EP raw down-FMA geometry sel={} aq={} ad={} weights={} \
12576 dst={dst_raw:#x} pairs={n_pairs} in={in_f} out={out_f} \
12577 owner={owner_start}..{owner_end}",
12578 sel.len(),
12579 aq.len(),
12580 ad.len(),
12581 route_weights.len(),
12582 )
12583 .into());
12584 }
12585 let f = self.func("qmatvec_nvfp4_q8_ep_down_fma");
12586 let cfg = LaunchConfig {
12587 grid_dim: (out_f as u32, 1, 1),
12588 block_dim: (256, 1, 1),
12589 shared_mem_bytes: 0,
12590 };
12591 let (inf, outf, np) = (in_f as i32, out_f as i32, n_pairs as i32);
12592 let (os, oe) = (owner_start as i32, owner_end as i32);
12593 let (rb, es) = (row_bytes as i64, expert_stride as i64);
12594 let __s_b = self.gpu.stream();
12595 let mut b = __s_b.launch_builder(&f);
12596 b.arg(bank)
12597 .arg(sel)
12598 .arg(aq)
12599 .arg(ad)
12600 .arg(route_weights)
12601 .arg(macros_down)
12602 .arg(&dst_raw)
12603 .arg(&inf)
12604 .arg(&outf)
12605 .arg(&np)
12606 .arg(&os)
12607 .arg(&oe)
12608 .arg(&rb)
12609 .arg(&es);
12610 unsafe {
12611 b.launch(cfg)?;
12612 }
12613 Ok(())
12614 }
12615
12616 #[allow(clippy::too_many_arguments)]
12621 pub fn silu_mul_scaled_q8_1_sel_into(
12622 &self,
12623 gate: &CudaSlice<f32>,
12624 up: &CudaSlice<f32>,
12625 gmac: &CudaSlice<f32>,
12626 umac: &CudaSlice<f32>,
12627 sel: &CudaSlice<i32>,
12628 limit: Option<f32>,
12629 out_q: &mut CudaSlice<i8>,
12630 out_d: &mut CudaSlice<f32>,
12631 n_per: usize,
12632 n_sel: usize,
12633 ) -> Result<(), Box<dyn std::error::Error>> {
12634 let n = n_per * n_sel;
12635 if !n_per.is_multiple_of(32) || out_q.len() < n || out_d.len() < n / 32 {
12636 return Err(format!(
12637 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
12638 out_q.len(),
12639 out_d.len()
12640 )
12641 .into());
12642 }
12643 if let Some(limit) = limit {
12644 if limit <= 1e-6 {
12645 return Err(format!(
12646 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
12647 )
12648 .into());
12649 }
12650 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
12651 let cfg = LaunchConfig::for_num_elems(n as u32);
12652 let (np, ns) = (n_per as i32, n_sel as i32);
12653 let __s_b = self.gpu.stream();
12654 let mut b = __s_b.launch_builder(&f);
12655 b.arg(gate)
12656 .arg(up)
12657 .arg(gmac)
12658 .arg(umac)
12659 .arg(sel)
12660 .arg(&limit)
12661 .arg(out_q)
12662 .arg(out_d)
12663 .arg(&np)
12664 .arg(&ns);
12665 unsafe {
12666 b.launch(cfg)?;
12667 }
12668 return Ok(());
12669 }
12670 let f = self.func("silu_mul_scaled_q8_1_sel");
12671 let cfg = LaunchConfig::for_num_elems(n as u32);
12672 let (np, ns) = (n_per as i32, n_sel as i32);
12673 let __s_b = self.gpu.stream();
12674 let mut b = __s_b.launch_builder(&f);
12675 b.arg(gate)
12676 .arg(up)
12677 .arg(gmac)
12678 .arg(umac)
12679 .arg(sel)
12680 .arg(out_q)
12681 .arg(out_d)
12682 .arg(&np)
12683 .arg(&ns);
12684 unsafe {
12685 b.launch(cfg)?;
12686 }
12687 Ok(())
12688 }
12689
12690 #[track_caller]
12691 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12692 crate::alloc_trace_hit(v.len() * 4);
12693 Ok(self.gpu.stream().clone_htod(v)?)
12694 }
12695 #[track_caller]
12696 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
12697 crate::alloc_trace_hit(v.len() * 4);
12698 Ok(self.gpu.stream().clone_htod(v)?)
12699 }
12700 #[track_caller]
12702 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
12703 crate::alloc_trace_hit(v.len());
12704 Ok(self.gpu.stream().clone_htod(v)?)
12705 }
12706 #[track_caller]
12707 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
12708 crate::alloc_trace_hit(v.len() * 8);
12709 Ok(self.gpu.stream().clone_htod(v)?)
12710 }
12711 #[track_caller]
12713 pub fn dtoh_view(
12714 &self,
12715 d: &cudarc::driver::CudaView<f32>,
12716 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
12717 crate::dtoh_trace_hit(d.len() * 4);
12718 let v = self.gpu.stream().clone_dtoh(d)?;
12719 self.gpu.stream().synchronize()?;
12720 Ok(v)
12721 }
12722 #[track_caller]
12723 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
12724 crate::dtoh_trace_hit(d.len() * 4);
12725 let v = self.gpu.stream().clone_dtoh(d)?;
12726 self.gpu.stream().synchronize()?;
12727 Ok(v)
12728 }
12729 #[track_caller]
12733 pub fn dtoh_pair(
12734 &self,
12735 a: &CudaSlice<f32>,
12736 b: &CudaSlice<f32>,
12737 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12738 crate::dtoh_trace_hit((a.len() + b.len()) * 4);
12739 let av = self.gpu.stream().clone_dtoh(a)?;
12740 let bv = self.gpu.stream().clone_dtoh(b)?;
12741 self.gpu.stream().synchronize()?;
12742 Ok((av, bv))
12743 }
12744 #[track_caller]
12747 pub fn dtoh_pair_views(
12748 &self,
12749 a: &cudarc::driver::CudaView<f32>,
12750 b: &cudarc::driver::CudaView<f32>,
12751 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12752 crate::dtoh_trace_hit((a.len() + b.len()) * 4);
12753 let av = self.gpu.stream().clone_dtoh(a)?;
12754 let bv = self.gpu.stream().clone_dtoh(b)?;
12755 self.gpu.stream().synchronize()?;
12756 Ok((av, bv))
12757 }
12758 #[track_caller]
12760 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
12761 crate::dtoh_trace_hit(d.len() * 4);
12762 let v = self.gpu.stream().clone_dtoh(d)?;
12763 self.gpu.stream().synchronize()?;
12764 Ok(v)
12765 }
12766 #[track_caller]
12770 pub fn dtoh_i8(&self, d: &CudaSlice<i8>) -> Result<Vec<i8>, Box<dyn std::error::Error>> {
12771 crate::dtoh_trace_hit(d.len());
12772 let v = self.gpu.stream().clone_dtoh(d)?;
12773 self.gpu.stream().synchronize()?;
12774 Ok(v)
12775 }
12776 #[track_caller]
12777 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
12778 crate::dtoh_trace_hit(d.len());
12779 let v = self.gpu.stream().clone_dtoh(d)?;
12780 self.gpu.stream().synchronize()?;
12781 Ok(v)
12782 }
12783 #[track_caller]
12784 pub fn dtoh_u8_view(
12785 &self,
12786 d: &cudarc::driver::CudaView<u8>,
12787 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
12788 crate::dtoh_trace_hit(d.len());
12789 let v = self.gpu.stream().clone_dtoh(d)?;
12790 self.gpu.stream().synchronize()?;
12791 Ok(v)
12792 }
12793 #[track_caller]
12801 pub fn dtoh_u8_into_pinned(
12802 &self,
12803 d: &CudaSlice<u8>,
12804 dst: &mut PinnedHostBuf,
12805 n: usize,
12806 ) -> Result<(), Box<dyn std::error::Error>> {
12807 crate::dtoh_trace_hit(d.len());
12808 if n > d.len() || n > dst.len() {
12809 return Err(format!(
12810 "dtoh_u8_into_pinned range {n} exceeds src {} or pinned dst {}",
12811 d.len(),
12812 dst.len(),
12813 )
12814 .into());
12815 }
12816 if n == 0 {
12817 return Ok(());
12818 }
12819 let host = &mut dst.as_mut_slice()[..n];
12820 self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
12821 self.gpu.stream().synchronize()?;
12822 Ok(())
12823 }
12824 #[track_caller]
12828 pub fn dtoh_f32_into_pinned(
12829 &self,
12830 d: &CudaSlice<f32>,
12831 dst: &mut PinnedHostBuf,
12832 n: usize,
12833 ) -> Result<(), Box<dyn std::error::Error>> {
12834 crate::dtoh_trace_hit(d.len() * 4);
12835 if n > d.len() || n * std::mem::size_of::<f32>() > dst.len() {
12836 return Err(format!(
12837 "dtoh_f32_into_pinned range {n} floats exceeds src {} or pinned dst {} bytes",
12838 d.len(),
12839 dst.len(),
12840 )
12841 .into());
12842 }
12843 if n == 0 {
12844 return Ok(());
12845 }
12846 let host: &mut [f32] = unsafe {
12848 std::slice::from_raw_parts_mut(dst.as_mut_slice().as_mut_ptr() as *mut f32, n)
12849 };
12850 self.gpu.stream().memcpy_dtoh(&d.slice(0..n), host)?;
12851 self.gpu.stream().synchronize()?;
12852 Ok(())
12853 }
12854
12855 pub fn htod_f32_from_pinned_async(
12863 &self,
12864 src: &PinnedHostBuf,
12865 n: usize,
12866 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12867 if n * std::mem::size_of::<f32>() > src.len() {
12868 return Err(format!(
12869 "htod_f32_from_pinned_async range {n} floats exceeds pinned src {} bytes",
12870 src.len(),
12871 )
12872 .into());
12873 }
12874 let mut d = self.uninit(n)?;
12875 if n == 0 {
12876 return Ok(d);
12877 }
12878 let host: &[f32] =
12880 unsafe { std::slice::from_raw_parts(src.as_slice().as_ptr() as *const f32, n) };
12881 let stream = self.gpu.stream();
12882 let s: &CudaStream = &stream;
12883 {
12884 let (pd, _guard) = d.device_ptr_mut(s);
12885 unsafe { cudarc::driver::result::memcpy_htod_async(pd, host, s.cu_stream())? };
12889 }
12890 Ok(d)
12891 }
12892
12893 #[allow(clippy::too_many_arguments)]
12900 pub fn copy_2d_dtod_async(
12903 &self,
12904 dst: &mut CudaSlice<f32>,
12905 dst_off_floats: usize,
12906 dst_pitch_floats: usize,
12907 src: &CudaSlice<f32>,
12908 src_pitch_floats: usize,
12909 width_floats: usize,
12910 height: usize,
12911 ) -> Result<(), Box<dyn std::error::Error>> {
12912 if height == 0 || width_floats == 0 {
12913 return Ok(());
12914 }
12915 let f = std::mem::size_of::<f32>();
12916 let dst_need = dst_off_floats + (height - 1) * dst_pitch_floats + width_floats;
12917 let src_need = (height - 1) * src_pitch_floats + width_floats;
12918 if width_floats > src_pitch_floats
12919 || width_floats > dst_pitch_floats
12920 || dst_need > dst.len()
12921 || src_need > src.len()
12922 {
12923 return Err(format!(
12924 "copy_2d_dtod_async out of bounds: dst needs {dst_need} of {}, src needs \
12925 {src_need} of {}, width {width_floats} pitches {src_pitch_floats}/{dst_pitch_floats}",
12926 dst.len(),
12927 src.len(),
12928 )
12929 .into());
12930 }
12931 let stream = self.gpu.stream();
12932 let s: &CudaStream = &stream;
12933 let (sp, _g0) = src.device_ptr(s);
12934 let (dp, _g1) = dst.device_ptr_mut(s);
12935 let copy = cudarc::driver::sys::CUDA_MEMCPY2D_st {
12936 srcXInBytes: 0,
12937 srcY: 0,
12938 srcMemoryType: cudarc::driver::sys::CUmemorytype_enum::CU_MEMORYTYPE_DEVICE,
12939 srcHost: std::ptr::null(),
12940 srcDevice: sp,
12941 srcArray: std::ptr::null_mut(),
12942 srcPitch: src_pitch_floats * f,
12943 dstXInBytes: dst_off_floats * f,
12944 dstY: 0,
12945 dstMemoryType: cudarc::driver::sys::CUmemorytype_enum::CU_MEMORYTYPE_DEVICE,
12946 dstHost: std::ptr::null_mut(),
12947 dstDevice: dp,
12948 dstArray: std::ptr::null_mut(),
12949 dstPitch: dst_pitch_floats * f,
12950 WidthInBytes: width_floats * f,
12951 Height: height,
12952 };
12953 unsafe { cudarc::driver::sys::cuMemcpy2DAsync_v2(©, s.cu_stream()).result()? };
12956 Ok(())
12957 }
12958
12959 pub fn copy_peer_from_async(
12966 &self,
12967 dst: &mut CudaSlice<f32>,
12968 src_engine: &Engine,
12969 src: &CudaSlice<f32>,
12970 n: usize,
12971 ) -> Result<(), Box<dyn std::error::Error>> {
12972 if n > src.len() || n > dst.len() {
12973 return Err(format!(
12974 "copy_peer_from_async range {n} exceeds src {} or dst {}",
12975 src.len(),
12976 dst.len(),
12977 )
12978 .into());
12979 }
12980 if n == 0 {
12981 return Ok(());
12982 }
12983 let stream = src_engine.gpu.stream();
12984 let s_src: &CudaStream = &stream;
12985 let (sp, _g0) = src.device_ptr(s_src);
12986 let (dp, _g1) = dst.device_ptr_mut(s_src);
12987 src_engine.ctx().bind_to_thread()?;
12988 let r = unsafe {
12991 cudarc::driver::result::memcpy_peer_async(
12992 self.ctx().cu_ctx(),
12993 dp,
12994 src_engine.ctx().cu_ctx(),
12995 sp,
12996 n * std::mem::size_of::<f32>(),
12997 s_src.cu_stream(),
12998 )
12999 };
13000 self.ctx().bind_to_thread()?;
13001 r?;
13002 Ok(())
13003 }
13004
13005 pub fn free_mem_mb(&self) -> u64 {
13007 self.ctx()
13008 .mem_get_info()
13009 .map(|(f, _)| f as u64 >> 20)
13010 .unwrap_or(0)
13011 }
13012
13013 #[track_caller]
13014 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13015 crate::alloc_trace_hit(n * 4);
13016 SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13017 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
13018 self.keep_if_capturing(&s);
13019 Ok(s)
13020 }
13021
13022 pub(crate) fn hyper_ws_take(&self) -> Option<crate::hyper::HyperDecodeWs> {
13026 self.hyper_decode_ws.lock().unwrap().take()
13027 }
13028
13029 pub(crate) fn hyper_ws_put(&self, ws: crate::hyper::HyperDecodeWs) {
13030 *self.hyper_decode_ws.lock().unwrap() = Some(ws);
13031 }
13032
13033 pub(crate) fn mla_seg_ws_take(
13036 &self,
13037 nh: usize,
13038 dn: usize,
13039 dr: usize,
13040 r: usize,
13041 q_lora: usize,
13042 ) -> Result<crate::hybrid_forward::MlaSegWs, Box<dyn std::error::Error>> {
13043 let held = self.mla_seg_ws.lock().unwrap().take();
13044 match held {
13045 Some(ws) if ws.sig == (nh, dn, dr, r, q_lora) => Ok(ws),
13046 _ => crate::hybrid_forward::MlaSegWs::new(self, nh, dn, dr, r, q_lora),
13047 }
13048 }
13049
13050 pub(crate) fn mla_seg_ws_put(&self, ws: crate::hybrid_forward::MlaSegWs) {
13051 *self.mla_seg_ws.lock().unwrap() = Some(ws);
13052 }
13053
13054 pub(crate) fn mla_seg_ws_on() -> bool {
13059 std::env::var("MEMRA_MLA_SEG_WS").as_deref() == Ok("1")
13060 }
13061
13062 pub(crate) fn vws_uninit(
13071 &self,
13072 n: usize,
13073 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13074 if verify_ws_on() {
13075 let mut ws = self.verify_ws.lock().unwrap();
13076 let ws = &mut *ws;
13077 if let Some(s) = VerifyWs::take(&mut ws.f32_pool, &mut ws.held_bytes, n) {
13078 if VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
13079 eprintln!(
13080 "[glm5-verify-ws] engaged: verify-walk buffers recycling through \
13081 the size-keyed pool (MEMRA_GLM5_VERIFY_WS=1)"
13082 );
13083 }
13084 return Ok(s);
13085 }
13086 }
13087 self.alloc_uninit::<f32>(n)
13088 }
13089
13090 pub(crate) fn vws_uninit_i8(
13092 &self,
13093 n: usize,
13094 ) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
13095 if verify_ws_on() {
13096 let mut ws = self.verify_ws.lock().unwrap();
13097 let ws = &mut *ws;
13098 if let Some(s) = VerifyWs::take(&mut ws.i8_pool, &mut ws.held_bytes, n) {
13099 VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13100 return Ok(s);
13101 }
13102 }
13103 self.alloc_uninit::<i8>(n)
13104 }
13105
13106 pub(crate) fn vws_uninit_u64(
13108 &self,
13109 n: usize,
13110 ) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
13111 if verify_ws_on() {
13112 let mut ws = self.verify_ws.lock().unwrap();
13113 let ws = &mut *ws;
13114 if let Some(s) = VerifyWs::take(&mut ws.u64_pool, &mut ws.held_bytes, n) {
13115 VERIFY_WS_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13116 return Ok(s);
13117 }
13118 }
13119 self.alloc_uninit::<u64>(n)
13120 }
13121
13122 #[allow(clippy::type_complexity)] pub(crate) fn glm5_graph_keep(&self) -> &Mutex<Vec<Box<dyn std::any::Any + Send>>> {
13128 &self.capture_keep
13129 }
13130
13131 pub(crate) fn vws_recycle(&self, s: CudaSlice<f32>) {
13134 if glm5_graph_capture_open() {
13135 self.capture_keep.lock().unwrap().push(Box::new(s));
13136 return;
13137 }
13138 if verify_ws_on() {
13139 let mut ws = self.verify_ws.lock().unwrap();
13140 let ws = &mut *ws;
13141 VerifyWs::put(&mut ws.f32_pool, &mut ws.held_bytes, s);
13142 }
13143 }
13144
13145 pub(crate) fn vws_recycle_i8(&self, s: CudaSlice<i8>) {
13147 if glm5_graph_capture_open() {
13148 self.capture_keep.lock().unwrap().push(Box::new(s));
13149 return;
13150 }
13151 if verify_ws_on() {
13152 let mut ws = self.verify_ws.lock().unwrap();
13153 let ws = &mut *ws;
13154 VerifyWs::put(&mut ws.i8_pool, &mut ws.held_bytes, s);
13155 }
13156 }
13157
13158 pub(crate) fn vws_recycle_u64(&self, s: CudaSlice<u64>) {
13160 if glm5_graph_capture_open() {
13161 self.capture_keep.lock().unwrap().push(Box::new(s));
13162 return;
13163 }
13164 if verify_ws_on() {
13165 let mut ws = self.verify_ws.lock().unwrap();
13166 let ws = &mut *ws;
13167 VerifyWs::put(&mut ws.u64_pool, &mut ws.held_bytes, s);
13168 }
13169 }
13170
13171 pub fn prob_of_token_device(
13180 &self,
13181 logits: &CudaSlice<f32>,
13182 tok: &CudaSlice<u32>,
13183 n_vocab: usize,
13184 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13185 let nb = ARGMAX_NB;
13186 let mut part = self.alloc_uninit::<f32>(nb)?;
13187 let mut p = self.alloc_uninit::<f32>(1)?;
13188 let f1 = self.func("prob_of_token_partial_f32");
13189 let cfg1 = LaunchConfig {
13190 grid_dim: (nb as u32, 1, 1),
13191 block_dim: (256, 1, 1),
13192 shared_mem_bytes: 0,
13193 };
13194 let nv = n_vocab as i32;
13195 let __s_b1 = self.gpu.stream();
13196 let mut b1 = __s_b1.launch_builder(&f1);
13197 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
13198 unsafe {
13199 b1.launch(cfg1)?;
13200 }
13201 let f2 = self.func("prob_of_token_final_f32");
13202 let cfg2 = LaunchConfig {
13203 grid_dim: (1, 1, 1),
13204 block_dim: (256, 1, 1),
13205 shared_mem_bytes: 0,
13206 };
13207 let nbi = nb as i32;
13208 let __s_b2 = self.gpu.stream();
13209 let mut b2 = __s_b2.launch_builder(&f2);
13210 b2.arg(&part).arg(&mut p).arg(&nbi);
13211 unsafe {
13212 b2.launch(cfg2)?;
13213 }
13214 Ok(p)
13215 }
13216
13217 pub fn prob_of_token_device_col(
13224 &self,
13225 logits: &CudaSlice<f32>,
13226 tok_all: &CudaSlice<u32>,
13227 tok_idx: usize,
13228 p_out: &mut CudaSlice<f32>,
13229 p_idx: usize,
13230 n_vocab: usize,
13231 ) -> Result<(), Box<dyn std::error::Error>> {
13232 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
13233 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
13234 let nb = ARGMAX_NB;
13235 let mut part = self.alloc_uninit::<f32>(nb)?;
13236 let f1 = self.func("prob_of_token_partial_f32");
13237 let cfg1 = LaunchConfig {
13238 grid_dim: (nb as u32, 1, 1),
13239 block_dim: (256, 1, 1),
13240 shared_mem_bytes: 0,
13241 };
13242 let nv = n_vocab as i32;
13243 let __s_b1 = self.gpu.stream();
13244 let mut b1 = __s_b1.launch_builder(&f1);
13245 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
13246 unsafe {
13247 b1.launch(cfg1)?;
13248 }
13249 let f2 = self.func("prob_of_token_final_f32");
13250 let cfg2 = LaunchConfig {
13251 grid_dim: (1, 1, 1),
13252 block_dim: (256, 1, 1),
13253 shared_mem_bytes: 0,
13254 };
13255 let nbi = nb as i32;
13256 let __s_b2 = self.gpu.stream();
13257 let mut b2 = __s_b2.launch_builder(&f2);
13258 b2.arg(&part).arg(&mut p_v).arg(&nbi);
13259 unsafe {
13260 b2.launch(cfg2)?;
13261 }
13262 Ok(())
13263 }
13264
13265 pub fn prob_of_token_device_into(
13266 &self,
13267 logits: &CudaSlice<f32>,
13268 tok: &CudaSlice<u32>,
13269 p_out: &mut CudaSlice<f32>,
13270 n_vocab: usize,
13271 ) -> Result<(), Box<dyn std::error::Error>> {
13272 let nb = ARGMAX_NB;
13273 let mut part = self.alloc_uninit::<f32>(nb)?;
13274 let f1 = self.func("prob_of_token_partial_f32");
13275 let cfg1 = LaunchConfig {
13276 grid_dim: (nb as u32, 1, 1),
13277 block_dim: (256, 1, 1),
13278 shared_mem_bytes: 0,
13279 };
13280 let nv = n_vocab as i32;
13281 let __s_b1 = self.gpu.stream();
13282 let mut b1 = __s_b1.launch_builder(&f1);
13283 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
13284 unsafe {
13285 b1.launch(cfg1)?;
13286 }
13287 let f2 = self.func("prob_of_token_final_f32");
13288 let cfg2 = LaunchConfig {
13289 grid_dim: (1, 1, 1),
13290 block_dim: (256, 1, 1),
13291 shared_mem_bytes: 0,
13292 };
13293 let nbi = nb as i32;
13294 let __s_b2 = self.gpu.stream();
13295 let mut b2 = __s_b2.launch_builder(&f2);
13296 b2.arg(&part).arg(p_out).arg(&nbi);
13297 unsafe {
13298 b2.launch(cfg2)?;
13299 }
13300 Ok(())
13301 }
13302
13303 pub fn u32_hist_append(
13306 &self,
13307 tok: &CudaSlice<u32>,
13308 hist: &mut CudaSlice<u32>,
13309 idx: &mut CudaSlice<i32>,
13310 ) -> Result<(), Box<dyn std::error::Error>> {
13311 let f = self.func("u32_hist_append");
13312 let cfg = LaunchConfig {
13313 grid_dim: (1, 1, 1),
13314 block_dim: (32, 1, 1),
13315 shared_mem_bytes: 0,
13316 };
13317 let __s_b = self.gpu.stream();
13318 let mut b = __s_b.launch_builder(&f);
13319 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
13320 unsafe {
13321 b.launch(cfg)?;
13322 }
13323 Ok(())
13324 }
13325
13326 pub fn argmax_token_device(
13327 &self,
13328 logits: &CudaSlice<f32>,
13329 n_vocab: usize,
13330 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13331 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
13332 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
13333 Ok(tok)
13334 }
13335 pub fn argmax_token_device_into(
13342 &self,
13343 logits: &CudaSlice<f32>,
13344 tok: &mut CudaSlice<u32>,
13345 n_vocab: usize,
13346 ) -> Result<(), Box<dyn std::error::Error>> {
13347 let nb = ARGMAX_NB;
13348 let f1 = self.func("argmax_partial_f32");
13349 let f2 = self.func("argmax_final_f32");
13350 let mut guard = self.argmax_partials.lock().unwrap();
13351 if guard.is_none() {
13352 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
13355 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
13356 *guard = Some((pv, pi));
13357 }
13358 let (part_v, part_i) = guard.as_mut().unwrap();
13359 let nv = n_vocab as i32;
13360 let nbi = nb as i32;
13361 let cfg1 = LaunchConfig {
13363 grid_dim: (nb as u32, 1, 1),
13364 block_dim: (256, 1, 1),
13365 shared_mem_bytes: 0,
13366 };
13367 let __s_b1 = self.gpu.stream();
13368 let mut b1 = __s_b1.launch_builder(&f1);
13369 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
13370 unsafe {
13371 b1.launch(cfg1)?;
13372 }
13373 let cfg2 = LaunchConfig {
13375 grid_dim: (1, 1, 1),
13376 block_dim: (256, 1, 1),
13377 shared_mem_bytes: 0,
13378 };
13379 let __s_b2 = self.gpu.stream();
13380 let mut b2 = __s_b2.launch_builder(&f2);
13381 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
13382 unsafe {
13383 b2.launch(cfg2)?;
13384 }
13385 Ok(())
13386 }
13387 pub fn argmax_token_device_col(
13393 &self,
13394 logits: &CudaSlice<f32>,
13395 col: usize,
13396 n_vocab: usize,
13397 toks: &mut CudaSlice<u32>,
13398 out_idx: usize,
13399 ) -> Result<(), Box<dyn std::error::Error>> {
13400 let nb = ARGMAX_NB;
13401 let f1 = self.func("argmax_partial_f32");
13402 let f2 = self.func("argmax_final_f32");
13403 let mut guard = self.argmax_partials.lock().unwrap();
13404 if guard.is_none() {
13405 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
13406 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
13407 *guard = Some((pv, pi));
13408 }
13409 let (part_v, part_i) = guard.as_mut().unwrap();
13410 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
13411 let nv = n_vocab as i32;
13412 let nbi = nb as i32;
13413 let cfg1 = LaunchConfig {
13414 grid_dim: (nb as u32, 1, 1),
13415 block_dim: (256, 1, 1),
13416 shared_mem_bytes: 0,
13417 };
13418 let __s_b1 = self.gpu.stream();
13419 let mut b1 = __s_b1.launch_builder(&f1);
13420 b1.arg(&col_view)
13421 .arg(&mut *part_v)
13422 .arg(&mut *part_i)
13423 .arg(&nv);
13424 unsafe {
13425 b1.launch(cfg1)?;
13426 }
13427 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
13428 let cfg2 = LaunchConfig {
13429 grid_dim: (1, 1, 1),
13430 block_dim: (256, 1, 1),
13431 shared_mem_bytes: 0,
13432 };
13433 let __s_b2 = self.gpu.stream();
13434 let mut b2 = __s_b2.launch_builder(&f2);
13435 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
13436 unsafe {
13437 b2.launch(cfg2)?;
13438 }
13439 Ok(())
13440 }
13441 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13443 Ok(self.gpu.stream().clone_htod(v)?)
13444 }
13445 #[track_caller]
13446 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
13447 crate::dtoh_trace_hit(d.len() * 8);
13448 let v = self.gpu.stream().clone_dtoh(d)?;
13449 self.gpu.stream().synchronize()?;
13450 Ok(v)
13451 }
13452
13453 #[track_caller]
13454 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
13455 crate::dtoh_trace_hit(d.len() * 4);
13456 let v = self.gpu.stream().clone_dtoh(d)?;
13457 self.gpu.stream().synchronize()?;
13458 Ok(v)
13459 }
13460 pub fn htod_u32_into(
13464 &self,
13465 dst: &mut CudaSlice<u32>,
13466 src: &[u32],
13467 ) -> Result<(), Box<dyn std::error::Error>> {
13468 let mut view = dst.slice_mut(0..src.len());
13469 self.gpu.stream().memcpy_htod(src, &mut view)?;
13470 Ok(())
13471 }
13472
13473 pub fn htod_i32_into(
13476 &self,
13477 dst: &mut CudaSlice<i32>,
13478 src: &[i32],
13479 ) -> Result<(), Box<dyn std::error::Error>> {
13480 let mut view = dst.slice_mut(0..src.len());
13481 self.gpu.stream().memcpy_htod(src, &mut view)?;
13482 Ok(())
13483 }
13484
13485 #[track_caller]
13486 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13487 crate::alloc_trace_hit(n * 4);
13488 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
13489 self.keep_if_capturing(&s);
13490 Ok(s)
13491 }
13492 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device_into(
13496 &self,
13497 embd: &CudaSlice<u8>,
13498 token_d: &CudaSlice<u32>,
13499 x_out: &mut CudaSlice<f32>,
13500 n_embd: usize,
13501 qtype: i32,
13502 row_bytes: usize,
13503 ) -> Result<(), Box<dyn std::error::Error>> {
13504 let f = self.func("embed_gather_u32");
13505 let cfg = LaunchConfig {
13506 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
13507 block_dim: (256, 1, 1),
13508 shared_mem_bytes: 0,
13509 };
13510 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
13511 let __s_b = self.gpu.stream();
13512 let mut b = __s_b.launch_builder(&f);
13513 b.arg(embd)
13514 .arg(token_d)
13515 .arg(x_out)
13516 .arg(&ne)
13517 .arg(&qt)
13518 .arg(&rb);
13519 unsafe {
13520 b.launch(cfg)?;
13521 }
13522 Ok(())
13523 }
13524 #[track_caller]
13526 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
13527 crate::dtoh_trace_hit(d.len() * 4);
13528 let v = self.gpu.stream().clone_dtoh(d)?;
13529 self.gpu.stream().synchronize()?;
13530 Ok(v[0])
13531 }
13532 pub fn i32_set_k(
13539 &self,
13540 dst: &mut CudaSlice<i32>,
13541 v: i32,
13542 ) -> Result<(), Box<dyn std::error::Error>> {
13543 let f = self.func("i32_set_k");
13544 let cfg = LaunchConfig {
13545 grid_dim: (1, 1, 1),
13546 block_dim: (1, 1, 1),
13547 shared_mem_bytes: 0,
13548 };
13549 let idx = 0i32;
13550 let __s_b = self.gpu.stream();
13551 let mut b = __s_b.launch_builder(&f);
13552 b.arg(dst).arg(&v).arg(&idx);
13553 unsafe {
13554 b.launch(cfg)?;
13555 }
13556 Ok(())
13557 }
13558
13559 pub fn set_i32_one(
13560 &self,
13561 d: &mut CudaSlice<i32>,
13562 v: i32,
13563 ) -> Result<(), Box<dyn std::error::Error>> {
13564 self.gpu.stream().memcpy_htod(&[v], d)?;
13565 Ok(())
13566 }
13567 pub fn set_u32_one(
13570 &self,
13571 d: &mut CudaSlice<u32>,
13572 v: u32,
13573 ) -> Result<(), Box<dyn std::error::Error>> {
13574 self.gpu.stream().memcpy_htod(&[v], d)?;
13575 Ok(())
13576 }
13577 #[track_caller]
13579 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
13580 crate::dtoh_trace_hit(d.len() * 4);
13581 let v = self.gpu.stream().clone_dtoh(d)?;
13582 self.gpu.stream().synchronize()?;
13583 Ok(v[0])
13584 }
13585 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
13587 Ok(self.gpu.stream().clone_htod(bytes)?)
13588 }
13589 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device(
13594 &self,
13595 embd: &CudaSlice<u8>,
13596 token_d: &CudaSlice<u32>,
13597 n_embd: usize,
13598 qtype: i32,
13599 row_bytes: usize,
13600 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13601 let f = self.func("embed_gather_u32");
13602 let mut x = self.alloc_uninit::<f32>(n_embd)?;
13603 let cfg = LaunchConfig {
13604 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
13605 block_dim: (256, 1, 1),
13606 shared_mem_bytes: 0,
13607 };
13608 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
13609 let __s_b = self.gpu.stream();
13610 let mut b = __s_b.launch_builder(&f);
13611 b.arg(embd)
13612 .arg(token_d)
13613 .arg(&mut x)
13614 .arg(&ne)
13615 .arg(&qt)
13616 .arg(&rb);
13617 unsafe {
13618 b.launch(cfg)?;
13619 }
13620 Ok(x)
13621 }
13622
13623 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device_t(
13628 &self,
13629 embd: &CudaSlice<u8>,
13630 tokens: &[u32],
13631 n_embd: usize,
13632 qtype: i32,
13633 row_bytes: usize,
13634 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13635 let t = tokens.len();
13636 let tok_d = self.gpu.stream().clone_htod(tokens)?;
13637 let f = self.func("embed_gather_u32_t");
13638 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13639 let cfg = LaunchConfig {
13640 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13641 block_dim: (256, 1, 1),
13642 shared_mem_bytes: 0,
13643 };
13644 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13645 let __s_b = self.gpu.stream();
13646 let mut b = __s_b.launch_builder(&f);
13647 b.arg(embd)
13648 .arg(&tok_d)
13649 .arg(&mut x)
13650 .arg(&ne)
13651 .arg(&qt)
13652 .arg(&rb)
13653 .arg(&ti);
13654 unsafe {
13655 b.launch(cfg)?;
13656 }
13657 Ok(x)
13658 }
13659
13660 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device_tv(
13666 &self,
13667 embd: &CudaSlice<u8>,
13668 tok_v: &cudarc::driver::CudaView<u32>,
13669 t: usize,
13670 n_embd: usize,
13671 qtype: i32,
13672 row_bytes: usize,
13673 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13674 let f = self.func("embed_gather_u32_t");
13675 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13676 let cfg = LaunchConfig {
13677 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13678 block_dim: (256, 1, 1),
13679 shared_mem_bytes: 0,
13680 };
13681 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13682 let __s_b = self.gpu.stream();
13683 let mut b = __s_b.launch_builder(&f);
13684 b.arg(embd)
13685 .arg(tok_v)
13686 .arg(&mut x)
13687 .arg(&ne)
13688 .arg(&qt)
13689 .arg(&rb)
13690 .arg(&ti);
13691 unsafe {
13692 b.launch(cfg)?;
13693 }
13694 Ok(x)
13695 }
13696
13697 #[allow(clippy::manual_div_ceil)] pub fn embed_gather_device_td(
13699 &self,
13700 embd: &CudaSlice<u8>,
13701 tok_d: &CudaSlice<u32>,
13702 t: usize,
13703 n_embd: usize,
13704 qtype: i32,
13705 row_bytes: usize,
13706 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13707 let f = self.func("embed_gather_u32_t");
13708 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
13709 let cfg = LaunchConfig {
13710 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
13711 block_dim: (256, 1, 1),
13712 shared_mem_bytes: 0,
13713 };
13714 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
13715 let __s_b = self.gpu.stream();
13716 let mut b = __s_b.launch_builder(&f);
13717 b.arg(embd)
13718 .arg(tok_d)
13719 .arg(&mut x)
13720 .arg(&ne)
13721 .arg(&qt)
13722 .arg(&rb)
13723 .arg(&ti);
13724 unsafe {
13725 b.launch(cfg)?;
13726 }
13727 Ok(x)
13728 }
13729
13730 #[inline]
13736 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
13738 if self
13739 .capture_keep_on
13740 .load(std::sync::atomic::Ordering::Relaxed)
13741 {
13742 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
13743 }
13744 }
13745
13746 #[track_caller]
13747 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
13748 &self,
13749 n: usize,
13750 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
13751 SCRATCH_ALLOC_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13752 crate::alloc_trace_hit(n * std::mem::size_of::<T>());
13753 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
13754 {
13758 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13759 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
13760 use cudarc::driver::DevicePtrMut;
13762 let n_bytes = s.len() * std::mem::size_of::<T>();
13763 let stream = self.gpu.stream();
13764 let (p_, _g) = s.device_ptr_mut(&stream);
13765 unsafe {
13766 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
13767 .result()?;
13768 }
13769 }
13770 }
13771 self.keep_if_capturing(&s);
13772 Ok(s)
13773 }
13774
13775 #[track_caller]
13780 pub fn uninit_q8_pair(
13781 &self,
13782 n: usize,
13783 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13784 Ok((
13785 self.alloc_uninit::<i8>(n)?,
13786 self.alloc_uninit::<f32>(n / 32)?,
13787 ))
13788 }
13789
13790 #[track_caller]
13791 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13792 self.alloc_uninit::<f32>(n)
13793 }
13794
13795 #[track_caller]
13797 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
13798 self.alloc_uninit::<i8>(n)
13799 }
13800
13801 #[track_caller]
13803 pub fn uninit_i32(&self, n: usize) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
13804 self.alloc_uninit::<i32>(n)
13805 }
13806
13807 #[allow(clippy::too_many_arguments)]
13811 pub fn rms_norm3(
13812 &self,
13813 x: &CudaSlice<f32>,
13814 w0: &CudaSlice<f32>,
13815 w1: &CudaSlice<f32>,
13816 w2: &CudaSlice<f32>,
13817 d0: &mut CudaSlice<f32>,
13818 d1: &mut CudaSlice<f32>,
13819 d2: &mut CudaSlice<f32>,
13820 ncols: usize,
13821 nrows: usize,
13822 eps: f32,
13823 ) -> Result<(), Box<dyn std::error::Error>> {
13824 let f = self.func("rms_norm3_f32");
13825 let cfg = LaunchConfig {
13826 grid_dim: (nrows as u32, 1, 1),
13827 block_dim: (rms_block(), 1, 1),
13828 shared_mem_bytes: 0,
13829 };
13830 let (nc, e) = (ncols as i32, eps);
13831 let __s_b = self.gpu.stream();
13832 let mut b = __s_b.launch_builder(&f);
13833 b.arg(x)
13834 .arg(w0)
13835 .arg(w1)
13836 .arg(w2)
13837 .arg(d0)
13838 .arg(d1)
13839 .arg(d2)
13840 .arg(&nc)
13841 .arg(&e);
13842 unsafe {
13843 b.launch(cfg)?;
13844 }
13845 Ok(())
13846 }
13847
13848 #[allow(clippy::too_many_arguments)]
13850 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
13853 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13854 *WARP_ON.get_or_init(|| {
13855 std::env::var("MEMRA_QKVNORM_W")
13856 .map(|v| v != "0")
13857 .unwrap_or(true)
13858 }) && ncols.is_multiple_of(4)
13859 && rows >= 64
13860 }
13861
13862 #[allow(clippy::too_many_arguments)]
13865 pub fn rms_norm_qkv_w4b(
13866 &self,
13867 q: &CudaSlice<f32>,
13868 k: &CudaSlice<f32>,
13869 v: &CudaSlice<f32>,
13870 wq: &CudaSlice<f32>,
13871 wk: &CudaSlice<f32>,
13872 wv: &CudaSlice<f32>,
13873 dq: &mut CudaSlice<f32>,
13874 dk: &mut CudaSlice<f32>,
13875 dv: &mut CudaSlice<f32>,
13876 dvb: &mut CudaSlice<u8>,
13877 ncols: usize,
13878 rq: usize,
13879 rk: usize,
13880 eps: f32,
13881 vf16: bool,
13882 ) -> Result<(), Box<dyn std::error::Error>> {
13883 assert!(ncols.is_multiple_of(4) && rq + 2 * rk >= 64);
13884 let f = self.func("rms_norm_qkv_w4b_f32");
13885 let rows = (rq + 2 * rk) as u32;
13886 let cfg = LaunchConfig {
13887 grid_dim: (rows.div_ceil(8), 1, 1),
13888 block_dim: (256, 1, 1),
13889 shared_mem_bytes: 0,
13890 };
13891 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
13892 let vf = vf16 as i32;
13893 let __s_b = self.gpu.stream();
13894 let mut b = __s_b.launch_builder(&f);
13895 b.arg(q)
13896 .arg(k)
13897 .arg(v)
13898 .arg(wq)
13899 .arg(wk)
13900 .arg(wv)
13901 .arg(dq)
13902 .arg(dk)
13903 .arg(dv)
13904 .arg(&mut *dvb)
13905 .arg(&nc)
13906 .arg(&rqi)
13907 .arg(&rki)
13908 .arg(&rvi)
13909 .arg(&e)
13910 .arg(&vf);
13911 unsafe {
13912 b.launch(cfg)?;
13913 }
13914 Ok(())
13915 }
13916
13917 #[allow(clippy::too_many_arguments)] pub fn rms_norm_qkv(
13919 &self,
13920 q: &CudaSlice<f32>,
13921 k: &CudaSlice<f32>,
13922 v: &CudaSlice<f32>,
13923 wq: &CudaSlice<f32>,
13924 wk: &CudaSlice<f32>,
13925 wv: &CudaSlice<f32>,
13926 dq: &mut CudaSlice<f32>,
13927 dk: &mut CudaSlice<f32>,
13928 dv: &mut CudaSlice<f32>,
13929 ncols: usize,
13930 rq: usize,
13931 rk: usize,
13932 eps: f32,
13933 ) -> Result<(), Box<dyn std::error::Error>> {
13934 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13938 let warp_on = *WARP_ON.get_or_init(|| {
13939 std::env::var("MEMRA_QKVNORM_W")
13940 .map(|v| v != "0")
13941 .unwrap_or(true)
13942 });
13943 if warp_on && ncols.is_multiple_of(4) && rq + 2 * rk >= 64 {
13946 let f = self.func("rms_norm_qkv_w4_f32");
13947 let rows = (rq + 2 * rk) as u32;
13948 let cfg = LaunchConfig {
13949 grid_dim: (rows.div_ceil(8), 1, 1),
13950 block_dim: (256, 1, 1),
13951 shared_mem_bytes: 0,
13952 };
13953 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
13954 let __s_b = self.gpu.stream();
13955 let mut b = __s_b.launch_builder(&f);
13956 b.arg(q)
13957 .arg(k)
13958 .arg(v)
13959 .arg(wq)
13960 .arg(wk)
13961 .arg(wv)
13962 .arg(dq)
13963 .arg(dk)
13964 .arg(dv)
13965 .arg(&nc)
13966 .arg(&rqi)
13967 .arg(&rki)
13968 .arg(&rvi)
13969 .arg(&e);
13970 unsafe {
13971 b.launch(cfg)?;
13972 }
13973 return Ok(());
13974 }
13975 let f = self.func("rms_norm_qkv_f32");
13976 let grid = (rq + 2 * rk) as u32;
13977 let cfg = LaunchConfig {
13978 grid_dim: (grid, 1, 1),
13979 block_dim: (rms_block(), 1, 1),
13980 shared_mem_bytes: 0,
13981 };
13982 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
13983 let __s_b = self.gpu.stream();
13984 let mut b = __s_b.launch_builder(&f);
13985 b.arg(q)
13986 .arg(k)
13987 .arg(v)
13988 .arg(wq)
13989 .arg(wk)
13990 .arg(wv)
13991 .arg(dq)
13992 .arg(dk)
13993 .arg(dv)
13994 .arg(&nc)
13995 .arg(&rqi)
13996 .arg(&rki)
13997 .arg(&e);
13998 unsafe {
13999 b.launch(cfg)?;
14000 }
14001 Ok(())
14002 }
14003
14004 #[allow(clippy::too_many_arguments)]
14006 pub fn rms_norm2x(
14007 &self,
14008 a: &CudaSlice<f32>,
14009 bb: &CudaSlice<f32>,
14010 wa: &CudaSlice<f32>,
14011 wb: &CudaSlice<f32>,
14012 da: &mut CudaSlice<f32>,
14013 db: &mut CudaSlice<f32>,
14014 ncols: usize,
14015 nrows: usize,
14016 eps: f32,
14017 ) -> Result<(), Box<dyn std::error::Error>> {
14018 let f = self.func("rms_norm2x_f32");
14019 let cfg = LaunchConfig {
14020 grid_dim: (2 * nrows as u32, 1, 1),
14021 block_dim: (rms_block(), 1, 1),
14022 shared_mem_bytes: 0,
14023 };
14024 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
14025 let __s_b = self.gpu.stream();
14026 let mut b = __s_b.launch_builder(&f);
14027 b.arg(a)
14028 .arg(bb)
14029 .arg(wa)
14030 .arg(wb)
14031 .arg(da)
14032 .arg(db)
14033 .arg(&nc)
14034 .arg(&nr)
14035 .arg(&e);
14036 unsafe {
14037 b.launch(cfg)?;
14038 }
14039 Ok(())
14040 }
14041
14042 pub fn softcap(
14044 &self,
14045 y: &mut CudaSlice<f32>,
14046 cap: f32,
14047 n: usize,
14048 ) -> Result<(), Box<dyn std::error::Error>> {
14049 let f = self.func("softcap_f32");
14050 let cfg = LaunchConfig::for_num_elems(n as u32);
14051 let ni = n as i32;
14052 let __s_b = self.gpu.stream();
14053 let mut b = __s_b.launch_builder(&f);
14054 b.arg(y).arg(&cap).arg(&ni);
14055 unsafe {
14056 b.launch(cfg)?;
14057 }
14058 Ok(())
14059 }
14060
14061 pub fn mask_ids_rows(
14064 &self,
14065 y: &mut CudaSlice<f32>,
14066 ids: &CudaSlice<i32>,
14067 n_ids: usize,
14068 n_vocab: usize,
14069 t: usize,
14070 ) -> Result<(), Box<dyn std::error::Error>> {
14071 let f = self.func("mask_ids_rows_f32");
14072 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
14073 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
14074 let __s_b = self.gpu.stream();
14075 let mut b = __s_b.launch_builder(&f);
14076 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
14077 unsafe {
14078 b.launch(cfg)?;
14079 }
14080 Ok(())
14081 }
14082
14083 #[allow(clippy::too_many_arguments)]
14085 pub fn add_scale_rms_norm(
14086 &self,
14087 a: &CudaSlice<f32>,
14088 b_in: &CudaSlice<f32>,
14089 c: f32,
14090 w: &CudaSlice<f32>,
14091 res: &mut CudaSlice<f32>,
14092 dst: &mut CudaSlice<f32>,
14093 ncols: usize,
14094 nrows: usize,
14095 eps: f32,
14096 ) -> Result<(), Box<dyn std::error::Error>> {
14097 let f = self.func("add_scale_rms_norm_f32");
14098 let cfg = LaunchConfig {
14099 grid_dim: (nrows as u32, 1, 1),
14100 block_dim: (rms_block(), 1, 1),
14101 shared_mem_bytes: 0,
14102 };
14103 let (nc, e2) = (ncols as i32, eps);
14104 let __s_b = self.gpu.stream();
14105 let mut b = __s_b.launch_builder(&f);
14106 b.arg(a)
14107 .arg(b_in)
14108 .arg(&c)
14109 .arg(w)
14110 .arg(res)
14111 .arg(dst)
14112 .arg(&nc)
14113 .arg(&e2);
14114 unsafe {
14115 b.launch(cfg)?;
14116 }
14117 Ok(())
14118 }
14119
14120 #[allow(clippy::too_many_arguments)]
14123 pub fn add_scale_rms_norm_q8_1(
14124 &self,
14125 a: &CudaSlice<f32>,
14126 b_in: &CudaSlice<f32>,
14127 c: f32,
14128 w: &CudaSlice<f32>,
14129 res: &mut CudaSlice<f32>,
14130 ncols: usize,
14131 nrows: usize,
14132 eps: f32,
14133 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14134 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14135 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14136 let (nc, e2) = (ncols as i32, eps);
14137 if Self::pdl_on() && Self::pdl_wb_on() {
14138 {
14139 use cudarc::driver::{DevicePtr, DevicePtrMut};
14140 let s = &self.gpu.stream();
14141 let (pa, _g0) = a.device_ptr(s);
14142 let (pb, _g1) = b_in.device_ptr(s);
14143 let (pw, _g2) = w.device_ptr(s);
14144 let (pr, _g3) = res.device_ptr_mut(s);
14145 let (pq, _g4) = out_q.device_ptr_mut(s);
14146 let (pd, _g5) = out_d.device_ptr_mut(s);
14147 let mut ps = [
14148 &pa as *const _ as *mut std::ffi::c_void,
14149 &pb as *const _ as *mut _,
14150 &c as *const _ as *mut _,
14151 &pw as *const _ as *mut _,
14152 &pr as *const _ as *mut _,
14153 &pq as *const _ as *mut _,
14154 &pd as *const _ as *mut _,
14155 &nc as *const _ as *mut _,
14156 &e2 as *const _ as *mut _,
14157 ];
14158 unsafe {
14159 self.launch_pdl(
14160 "add_scale_rms_norm_q8_1",
14161 (nrows as u32, 1, 1),
14162 (rms_block(), 1, 1),
14163 &mut ps,
14164 )?;
14165 }
14166 }
14167 return Ok((out_q, out_d));
14168 }
14169 let f = self.func("add_scale_rms_norm_q8_1");
14170 let cfg = LaunchConfig {
14171 grid_dim: (nrows as u32, 1, 1),
14172 block_dim: (rms_block(), 1, 1),
14173 shared_mem_bytes: 0,
14174 };
14175 let __s_b = self.gpu.stream();
14176 let mut b = __s_b.launch_builder(&f);
14177 b.arg(a)
14178 .arg(b_in)
14179 .arg(&c)
14180 .arg(w)
14181 .arg(res)
14182 .arg(&mut out_q)
14183 .arg(&mut out_d)
14184 .arg(&nc)
14185 .arg(&e2);
14186 unsafe {
14187 b.launch(cfg)?;
14188 }
14189 Ok((out_q, out_d))
14190 }
14191
14192 #[allow(clippy::too_many_arguments)]
14194 pub fn add_scale_rms_norm_q8_1_into(
14195 &self,
14196 a: &CudaSlice<f32>,
14197 b_in: &CudaSlice<f32>,
14198 c: f32,
14199 w: &CudaSlice<f32>,
14200 res: &mut CudaSlice<f32>,
14201 ncols: usize,
14202 nrows: usize,
14203 eps: f32,
14204 out_q: &mut CudaSlice<i8>,
14205 out_d: &mut CudaSlice<f32>,
14206 ) -> Result<(), Box<dyn std::error::Error>> {
14207 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
14208 let (nc, e2) = (ncols as i32, eps);
14209 if Self::pdl_on() && Self::pdl_wb_on() {
14210 use cudarc::driver::{DevicePtr, DevicePtrMut};
14211 let s = &self.gpu.stream();
14212 let (pa, _g0) = a.device_ptr(s);
14213 let (pb, _g1) = b_in.device_ptr(s);
14214 let (pw, _g2) = w.device_ptr(s);
14215 let (pr, _g3) = res.device_ptr_mut(s);
14216 let (pq, _g4) = out_q.device_ptr_mut(s);
14217 let (pd, _g5) = out_d.device_ptr_mut(s);
14218 let mut ps = [
14219 &pa as *const _ as *mut std::ffi::c_void,
14220 &pb as *const _ as *mut _,
14221 &c as *const _ as *mut _,
14222 &pw as *const _ as *mut _,
14223 &pr as *const _ as *mut _,
14224 &pq as *const _ as *mut _,
14225 &pd as *const _ as *mut _,
14226 &nc as *const _ as *mut _,
14227 &e2 as *const _ as *mut _,
14228 ];
14229 unsafe {
14230 self.launch_pdl(
14231 "add_scale_rms_norm_q8_1",
14232 (nrows as u32, 1, 1),
14233 (rms_block(), 1, 1),
14234 &mut ps,
14235 )?;
14236 }
14237 return Ok(());
14238 }
14239 let f = self.func("add_scale_rms_norm_q8_1");
14240 let cfg = LaunchConfig {
14241 grid_dim: (nrows as u32, 1, 1),
14242 block_dim: (rms_block(), 1, 1),
14243 shared_mem_bytes: 0,
14244 };
14245 let __s_b = self.gpu.stream();
14246 let mut b = __s_b.launch_builder(&f);
14247 b.arg(a)
14248 .arg(b_in)
14249 .arg(&c)
14250 .arg(w)
14251 .arg(res)
14252 .arg(&mut *out_q)
14253 .arg(&mut *out_d)
14254 .arg(&nc)
14255 .arg(&e2);
14256 unsafe {
14257 b.launch(cfg)?;
14258 }
14259 Ok(())
14260 }
14261
14262 #[allow(clippy::too_many_arguments)]
14265 pub fn rms_pre_add_scale_rms_norm_q8_1(
14266 &self,
14267 a: &CudaSlice<f32>,
14268 wa: &CudaSlice<f32>,
14269 b_in: &CudaSlice<f32>,
14270 c: f32,
14271 w: &CudaSlice<f32>,
14272 res: &mut CudaSlice<f32>,
14273 ncols: usize,
14274 nrows: usize,
14275 eps: f32,
14276 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14277 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14278 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14279 let (nc, e2) = (ncols as i32, eps);
14280 if Self::pdl_on() {
14281 {
14282 use cudarc::driver::{DevicePtr, DevicePtrMut};
14283 let s = &self.gpu.stream();
14284 let (pa, _g0) = a.device_ptr(s);
14285 let (pwa, _g1) = wa.device_ptr(s);
14286 let (pb, _g2) = b_in.device_ptr(s);
14287 let (pw, _g3) = w.device_ptr(s);
14288 let (pr, _g4) = res.device_ptr_mut(s);
14289 let (pq, _g5) = out_q.device_ptr_mut(s);
14290 let (pd, _g6) = out_d.device_ptr_mut(s);
14291 let mut ps = [
14292 &pa as *const _ as *mut std::ffi::c_void,
14293 &pwa as *const _ as *mut _,
14294 &pb as *const _ as *mut _,
14295 &c as *const _ as *mut _,
14296 &pw as *const _ as *mut _,
14297 &pr as *const _ as *mut _,
14298 &pq as *const _ as *mut _,
14299 &pd as *const _ as *mut _,
14300 &nc as *const _ as *mut _,
14301 &e2 as *const _ as *mut _,
14302 ];
14303 unsafe {
14304 self.launch_pdl(
14305 "rms_pre_add_scale_rms_norm_q8_1",
14306 (nrows as u32, 1, 1),
14307 (rms_block(), 1, 1),
14308 &mut ps,
14309 )?;
14310 }
14311 }
14312 return Ok((out_q, out_d));
14313 }
14314 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
14315 let cfg = LaunchConfig {
14316 grid_dim: (nrows as u32, 1, 1),
14317 block_dim: (rms_block(), 1, 1),
14318 shared_mem_bytes: 0,
14319 };
14320 let __s_b = self.gpu.stream();
14321 let mut b = __s_b.launch_builder(&f);
14322 b.arg(a)
14323 .arg(wa)
14324 .arg(b_in)
14325 .arg(&c)
14326 .arg(w)
14327 .arg(res)
14328 .arg(&mut out_q)
14329 .arg(&mut out_d)
14330 .arg(&nc)
14331 .arg(&e2);
14332 unsafe {
14333 b.launch(cfg)?;
14334 }
14335 Ok((out_q, out_d))
14336 }
14337
14338 pub fn gelu_tanh_mul_q8_1(
14341 &self,
14342 gate: &CudaSlice<f32>,
14343 up: &cudarc::driver::CudaView<f32>,
14344 act: &mut CudaSlice<f32>,
14345 ncols: usize,
14346 nrows: usize,
14347 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14348 debug_assert!(ncols.is_multiple_of(128));
14349 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
14350 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14351 let nc = ncols as i32;
14352 if Self::pdl_on() {
14353 {
14354 use cudarc::driver::{DevicePtr, DevicePtrMut};
14355 let s = &self.gpu.stream();
14356 let (pg, _g0) = gate.device_ptr(s);
14357 let (pu, _g1) = up.device_ptr(s);
14358 let (pact, _g2) = act.device_ptr_mut(s);
14359 let (pq, _g3) = out_q.device_ptr_mut(s);
14360 let (pd, _g4) = out_d.device_ptr_mut(s);
14361 let mut ps = [
14362 &pg as *const _ as *mut std::ffi::c_void,
14363 &pu as *const _ as *mut _,
14364 &pact as *const _ as *mut _,
14365 &pq as *const _ as *mut _,
14366 &pd as *const _ as *mut _,
14367 &nc as *const _ as *mut _,
14368 ];
14369 unsafe {
14370 self.launch_pdl(
14371 "gelu_tanh_mul_q8_1",
14372 (nrows as u32, 1, 1),
14373 (rms_block(), 1, 1),
14374 &mut ps,
14375 )?;
14376 }
14377 }
14378 return Ok((out_q, out_d));
14379 }
14380 let f = self.func("gelu_tanh_mul_q8_1");
14381 let cfg = LaunchConfig {
14382 grid_dim: (nrows as u32, 1, 1),
14383 block_dim: (rms_block(), 1, 1),
14384 shared_mem_bytes: 0,
14385 };
14386 let __s_b = self.gpu.stream();
14387 let mut b = __s_b.launch_builder(&f);
14388 b.arg(gate)
14389 .arg(up)
14390 .arg(act)
14391 .arg(&mut out_q)
14392 .arg(&mut out_d)
14393 .arg(&nc);
14394 unsafe {
14395 b.launch(cfg)?;
14396 }
14397 Ok((out_q, out_d))
14398 }
14399
14400 #[allow(clippy::too_many_arguments)]
14402 pub fn gelu_tanh_mul_q8_1_into(
14403 &self,
14404 gate: &CudaSlice<f32>,
14405 up: &cudarc::driver::CudaView<f32>,
14406 act: &mut CudaSlice<f32>,
14407 ncols: usize,
14408 nrows: usize,
14409 out_q: &mut CudaSlice<i8>,
14410 out_d: &mut CudaSlice<f32>,
14411 ) -> Result<(), Box<dyn std::error::Error>> {
14412 debug_assert!(ncols.is_multiple_of(128));
14413 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
14414 let nc = ncols as i32;
14415 if Self::pdl_on() {
14416 use cudarc::driver::{DevicePtr, DevicePtrMut};
14417 let s = &self.gpu.stream();
14418 let (pg, _g0) = gate.device_ptr(s);
14419 let (pu, _g1) = up.device_ptr(s);
14420 let (pact, _g2) = act.device_ptr_mut(s);
14421 let (pq, _g3) = out_q.device_ptr_mut(s);
14422 let (pd, _g4) = out_d.device_ptr_mut(s);
14423 let mut ps = [
14424 &pg as *const _ as *mut std::ffi::c_void,
14425 &pu as *const _ as *mut _,
14426 &pact as *const _ as *mut _,
14427 &pq as *const _ as *mut _,
14428 &pd as *const _ as *mut _,
14429 &nc as *const _ as *mut _,
14430 ];
14431 unsafe {
14432 self.launch_pdl(
14433 "gelu_tanh_mul_q8_1",
14434 (nrows as u32, 1, 1),
14435 (rms_block(), 1, 1),
14436 &mut ps,
14437 )?;
14438 }
14439 return Ok(());
14440 }
14441 let f = self.func("gelu_tanh_mul_q8_1");
14442 let cfg = LaunchConfig {
14443 grid_dim: (nrows as u32, 1, 1),
14444 block_dim: (rms_block(), 1, 1),
14445 shared_mem_bytes: 0,
14446 };
14447 let __s_b = self.gpu.stream();
14448 let mut b = __s_b.launch_builder(&f);
14449 b.arg(gate)
14450 .arg(up)
14451 .arg(&mut *act)
14452 .arg(&mut *out_q)
14453 .arg(&mut *out_d)
14454 .arg(&nc);
14455 unsafe {
14456 b.launch(cfg)?;
14457 }
14458 Ok(())
14459 }
14460
14461 #[allow(clippy::too_many_arguments)]
14463 #[allow(clippy::type_complexity)] pub fn add_rms_norm3_q8z(
14465 &self,
14466 a: &CudaSlice<f32>,
14467 b_in: &CudaSlice<f32>,
14468 w0: &CudaSlice<f32>,
14469 w1: &CudaSlice<f32>,
14470 w2: &CudaSlice<f32>,
14471 res: &mut CudaSlice<f32>,
14472 out1: &mut CudaSlice<f32>,
14473 ncols: usize,
14474 nrows: usize,
14475 eps: f32,
14476 ) -> Result<
14477 (
14478 (CudaSlice<i8>, CudaSlice<f32>),
14479 (CudaSlice<i8>, CudaSlice<f32>),
14480 ),
14481 Box<dyn std::error::Error>,
14482 > {
14483 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
14484 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14485 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
14486 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
14487 let f = self.func("add_rms_norm3_q8z_f32");
14488 let cfg = LaunchConfig {
14489 grid_dim: (nrows as u32, 1, 1),
14490 block_dim: (rms_block(), 1, 1),
14491 shared_mem_bytes: 0,
14492 };
14493 let (nc, e2) = (ncols as i32, eps);
14494 let __s_b = self.gpu.stream();
14495 let mut b = __s_b.launch_builder(&f);
14496 b.arg(a)
14497 .arg(b_in)
14498 .arg(w0)
14499 .arg(w1)
14500 .arg(w2)
14501 .arg(res)
14502 .arg(&mut q0)
14503 .arg(&mut d0)
14504 .arg(out1)
14505 .arg(&mut q2)
14506 .arg(&mut d2)
14507 .arg(&nc)
14508 .arg(&e2);
14509 unsafe {
14510 b.launch(cfg)?;
14511 }
14512 Ok(((q0, d0), (q2, d2)))
14513 }
14514
14515 #[allow(clippy::too_many_arguments)]
14517 pub fn add_rms_norm3(
14518 &self,
14519 a: &CudaSlice<f32>,
14520 b_in: &CudaSlice<f32>,
14521 w0: &CudaSlice<f32>,
14522 w1: &CudaSlice<f32>,
14523 w2: &CudaSlice<f32>,
14524 res: &mut CudaSlice<f32>,
14525 d0: &mut CudaSlice<f32>,
14526 d1: &mut CudaSlice<f32>,
14527 d2: &mut CudaSlice<f32>,
14528 ncols: usize,
14529 nrows: usize,
14530 eps: f32,
14531 ) -> Result<(), Box<dyn std::error::Error>> {
14532 let f = self.func("add_rms_norm3_f32");
14533 let cfg = LaunchConfig {
14534 grid_dim: (nrows as u32, 1, 1),
14535 block_dim: (rms_block(), 1, 1),
14536 shared_mem_bytes: 0,
14537 };
14538 let (nc, e2) = (ncols as i32, eps);
14539 let __s_b = self.gpu.stream();
14540 let mut b = __s_b.launch_builder(&f);
14541 b.arg(a)
14542 .arg(b_in)
14543 .arg(w0)
14544 .arg(w1)
14545 .arg(w2)
14546 .arg(res)
14547 .arg(d0)
14548 .arg(d1)
14549 .arg(d2)
14550 .arg(&nc)
14551 .arg(&e2);
14552 unsafe {
14553 b.launch(cfg)?;
14554 }
14555 Ok(())
14556 }
14557
14558 pub fn add_scale(
14560 &self,
14561 a: &CudaSlice<f32>,
14562 b_in: &CudaSlice<f32>,
14563 c: f32,
14564 dst: &mut CudaSlice<f32>,
14565 n: usize,
14566 ) -> Result<(), Box<dyn std::error::Error>> {
14567 let f = self.func("add_scale_f32");
14568 let cfg = LaunchConfig::for_num_elems(n as u32);
14569 let ni = n as i32;
14570 let __s_b = self.gpu.stream();
14571 let mut b = __s_b.launch_builder(&f);
14572 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
14573 unsafe {
14574 b.launch(cfg)?;
14575 }
14576 Ok(())
14577 }
14578
14579 #[allow(clippy::too_many_arguments)] pub fn layer_norm_bias(
14582 &self,
14583 x: &CudaSlice<f32>,
14584 w: &CudaSlice<f32>,
14585 b: &CudaSlice<f32>,
14586 dst: &mut CudaSlice<f32>,
14587 ncols: usize,
14588 nrows: usize,
14589 eps: f32,
14590 ) -> Result<(), Box<dyn std::error::Error>> {
14591 let f = self.func("layer_norm_bias_f32");
14592 let (nc, e) = (ncols as i32, eps);
14593 let cfg = LaunchConfig {
14594 grid_dim: (nrows as u32, 1, 1),
14595 block_dim: (256, 1, 1),
14596 shared_mem_bytes: 0,
14597 };
14598 let __s_b = self.gpu.stream();
14599 let mut lb = __s_b.launch_builder(&f);
14600 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
14601 unsafe {
14602 lb.launch(cfg)?;
14603 }
14604 Ok(())
14605 }
14606
14607 pub fn gelu_tanh(
14609 &self,
14610 x: &CudaSlice<f32>,
14611 dst: &mut CudaSlice<f32>,
14612 n: usize,
14613 ) -> Result<(), Box<dyn std::error::Error>> {
14614 let f = self.func("gelu_tanh_f32");
14615 let ni = n as i64;
14616 let cfg = LaunchConfig {
14617 grid_dim: (n.div_ceil(256) as u32, 1, 1),
14618 block_dim: (256, 1, 1),
14619 shared_mem_bytes: 0,
14620 };
14621 let __s_b = self.gpu.stream();
14622 let mut lb = __s_b.launch_builder(&f);
14623 lb.arg(x).arg(&mut *dst).arg(&ni);
14624 unsafe {
14625 lb.launch(cfg)?;
14626 }
14627 Ok(())
14628 }
14629
14630 pub fn row_softmax(
14632 &self,
14633 x: &mut CudaSlice<f32>,
14634 ncols: usize,
14635 nrows: usize,
14636 ) -> Result<(), Box<dyn std::error::Error>> {
14637 let f = self.func("row_softmax_f32");
14638 let nc = ncols as i32;
14639 let cfg = LaunchConfig {
14640 grid_dim: (nrows as u32, 1, 1),
14641 block_dim: (256, 1, 1),
14642 shared_mem_bytes: 0,
14643 };
14644 let __s_b = self.gpu.stream();
14645 let mut lb = __s_b.launch_builder(&f);
14646 lb.arg(&mut *x).arg(&nc);
14647 unsafe {
14648 lb.launch(cfg)?;
14649 }
14650 Ok(())
14651 }
14652
14653 pub fn rms_norm(
14654 &self,
14655 x: &CudaSlice<f32>,
14656 w: &CudaSlice<f32>,
14657 dst: &mut CudaSlice<f32>,
14658 ncols: usize,
14659 nrows: usize,
14660 eps: f32,
14661 ) -> Result<(), Box<dyn std::error::Error>> {
14662 let (nc, e) = (ncols as i32, eps);
14663 let kname = if Self::norm_ilp_on() {
14664 "rms_norm_f32_v2"
14665 } else {
14666 "rms_norm_f32"
14667 };
14668 if Self::pdl_on() && Self::pdl_wb_on() {
14669 use cudarc::driver::{DevicePtr, DevicePtrMut};
14670 let s = &self.gpu.stream();
14671 let (px, _g0) = x.device_ptr(s);
14672 let (pw, _g1) = w.device_ptr(s);
14673 let (pd, _g2) = dst.device_ptr_mut(s);
14674 let mut ps = [
14675 &px as *const _ as *mut std::ffi::c_void,
14676 &pw as *const _ as *mut _,
14677 &pd as *const _ as *mut _,
14678 &nc as *const _ as *mut _,
14679 &e as *const _ as *mut _,
14680 ];
14681 unsafe {
14682 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
14683 }
14684 return Ok(());
14685 }
14686 let f = self.func(kname);
14687 let cfg = LaunchConfig {
14688 grid_dim: (nrows as u32, 1, 1),
14689 block_dim: (rms_block(), 1, 1),
14690 shared_mem_bytes: 0,
14691 };
14692 let __s_b = self.gpu.stream();
14693 let mut b = __s_b.launch_builder(&f);
14694 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
14695 unsafe {
14696 b.launch(cfg)?;
14697 }
14698 Ok(())
14699 }
14700
14701 pub fn rms_norm_decode(
14709 &self,
14710 x: &CudaSlice<f32>,
14711 w: &CudaSlice<f32>,
14712 dst: &mut CudaSlice<f32>,
14713 ncols: usize,
14714 nrows: usize,
14715 eps: f32,
14716 ) -> Result<(), Box<dyn std::error::Error>> {
14717 let f = self.func(if Self::norm_ilp_on() {
14718 "rms_norm_f32_v2"
14719 } else {
14720 "rms_norm_f32"
14721 });
14722 let cfg = LaunchConfig {
14723 grid_dim: (nrows as u32, 1, 1),
14724 block_dim: (1024, 1, 1),
14725 shared_mem_bytes: 0,
14726 };
14727 let (nc, e) = (ncols as i32, eps);
14728 let __s_b = self.gpu.stream();
14729 let mut b = __s_b.launch_builder(&f);
14730 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
14731 unsafe {
14732 b.launch(cfg)?;
14733 }
14734 Ok(())
14735 }
14736
14737 pub fn rms_norm_q8_1(
14741 &self,
14742 x: &CudaSlice<f32>,
14743 w: &CudaSlice<f32>,
14744 ncols: usize,
14745 nrows: usize,
14746 eps: f32,
14747 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14748 let nblk = ncols / 32;
14749 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
14750 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
14751 let (nc, e) = (ncols as i32, eps);
14752 if Self::pdl_on() {
14753 {
14754 use cudarc::driver::{DevicePtr, DevicePtrMut};
14755 let s = &self.gpu.stream();
14756 let (px, _g0) = x.device_ptr(s);
14757 let (pw, _g1) = w.device_ptr(s);
14758 let (pq, _g2) = q.device_ptr_mut(s);
14759 let (pd, _g3) = d.device_ptr_mut(s);
14760 let mut ps = [
14761 &px as *const _ as *mut std::ffi::c_void,
14762 &pw as *const _ as *mut _,
14763 &pq as *const _ as *mut _,
14764 &pd as *const _ as *mut _,
14765 &nc as *const _ as *mut _,
14766 &e as *const _ as *mut _,
14767 ];
14768 unsafe {
14769 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
14770 }
14771 }
14772 return Ok((q, d));
14773 }
14774 let f = self.func("rms_norm_q8_1");
14775 let cfg = LaunchConfig {
14778 grid_dim: (nrows as u32, 1, 1),
14779 block_dim: (1024, 1, 1),
14780 shared_mem_bytes: 0,
14781 };
14782 let __s_b = self.gpu.stream();
14783 let mut b = __s_b.launch_builder(&f);
14784 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
14785 unsafe {
14786 b.launch(cfg)?;
14787 }
14788 Ok((q, d))
14789 }
14790
14791 #[allow(clippy::too_many_arguments)] pub fn rms_norm_q8_1_into(
14795 &self,
14796 x: &CudaSlice<f32>,
14797 w: &CudaSlice<f32>,
14798 ncols: usize,
14799 nrows: usize,
14800 eps: f32,
14801 q: &mut CudaSlice<i8>,
14802 d: &mut CudaSlice<f32>,
14803 ) -> Result<(), Box<dyn std::error::Error>> {
14804 let nblk = ncols / 32;
14805 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
14806 let (nc, e) = (ncols as i32, eps);
14807 if Self::pdl_on() {
14808 use cudarc::driver::{DevicePtr, DevicePtrMut};
14809 let s = &self.gpu.stream();
14810 let (px, _g0) = x.device_ptr(s);
14811 let (pw, _g1) = w.device_ptr(s);
14812 let (pq, _g2) = q.device_ptr_mut(s);
14813 let (pd, _g3) = d.device_ptr_mut(s);
14814 let mut ps = [
14815 &px as *const _ as *mut std::ffi::c_void,
14816 &pw as *const _ as *mut _,
14817 &pq as *const _ as *mut _,
14818 &pd as *const _ as *mut _,
14819 &nc as *const _ as *mut _,
14820 &e as *const _ as *mut _,
14821 ];
14822 unsafe {
14823 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
14824 }
14825 return Ok(());
14826 }
14827 let f = self.func("rms_norm_q8_1");
14828 let cfg = LaunchConfig {
14829 grid_dim: (nrows as u32, 1, 1),
14830 block_dim: (1024, 1, 1),
14831 shared_mem_bytes: 0,
14832 };
14833 let __s_b = self.gpu.stream();
14834 let mut b = __s_b.launch_builder(&f);
14835 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
14836 unsafe {
14837 b.launch(cfg)?;
14838 }
14839 Ok(())
14840 }
14841
14842 pub fn rms_norm_zq8_f32(
14855 &self,
14856 x: &CudaSlice<f32>,
14857 w: &CudaSlice<f32>,
14858 z: &mut CudaSlice<f32>,
14859 ncols: usize,
14860 nrows: usize,
14861 eps: f32,
14862 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14863 self.rms_norm_zq8_f32_arm(x, w, z, ncols, nrows, eps, Self::norm_ilp_zq8_on())
14864 }
14865
14866 #[allow(clippy::too_many_arguments)]
14872 pub fn rms_norm_zq8_f32_arm(
14873 &self,
14874 x: &CudaSlice<f32>,
14875 w: &CudaSlice<f32>,
14876 z: &mut CudaSlice<f32>,
14877 ncols: usize,
14878 nrows: usize,
14879 eps: f32,
14880 ilp: bool,
14881 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14882 assert!(ncols.is_multiple_of(32));
14883 let nblk = ncols / 32;
14884 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?; let mut d = self.alloc_uninit::<f32>(nrows * nblk)?; let f = self.func(if ilp {
14887 "rms_norm_zq8_f32_v2"
14888 } else {
14889 "rms_norm_zq8_f32"
14890 });
14891 let cfg = LaunchConfig {
14895 grid_dim: (nrows as u32, 1, 1),
14896 block_dim: (rms_block(), 1, 1),
14897 shared_mem_bytes: 0,
14898 };
14899 let (nc, ep) = (ncols as i32, eps);
14900 let __s_b = self.gpu.stream();
14901 let mut b = __s_b.launch_builder(&f);
14902 b.arg(x)
14903 .arg(w)
14904 .arg(&mut *z)
14905 .arg(&mut q)
14906 .arg(&mut d)
14907 .arg(&nc)
14908 .arg(&ep);
14909 unsafe {
14910 b.launch(cfg)?;
14911 }
14912 Ok((q, d))
14913 }
14914
14915 pub fn quantize_q8_1_into(
14917 &self,
14918 x: &CudaSlice<f32>,
14919 m: usize,
14920 in_f: usize,
14921 q: &mut CudaSlice<i8>,
14922 d: &mut CudaSlice<f32>,
14923 ) -> Result<(), Box<dyn std::error::Error>> {
14924 let nblk = in_f / 32;
14925 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
14926 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
14927 let (inf, mi) = (in_f as i32, m as i32);
14928 if Self::pdl_on() && Self::pdl_wb_on() {
14929 use cudarc::driver::{DevicePtr, DevicePtrMut};
14930 let s = &self.gpu.stream();
14931 let (px, _g0) = x.device_ptr(s);
14932 let (pq, _g1) = q.device_ptr_mut(s);
14933 let (pd, _g2) = d.device_ptr_mut(s);
14934 let mut ps = [
14935 &px as *const _ as *mut std::ffi::c_void,
14936 &pq as *const _ as *mut _,
14937 &pd as *const _ as *mut _,
14938 &inf as *const _ as *mut _,
14939 &mi as *const _ as *mut _,
14940 ];
14941 unsafe {
14942 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
14943 }
14944 return Ok(());
14945 }
14946 let f = self.func("quantize_q8_1");
14947 let __s_b = self.gpu.stream();
14948 let mut b = __s_b.launch_builder(&f);
14949 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
14950 unsafe {
14951 b.launch(cfg)?;
14952 }
14953 Ok(())
14954 }
14955
14956 #[allow(clippy::too_many_arguments)] pub fn add_rms_norm_q8_1(
14961 &self,
14962 a: &CudaSlice<f32>,
14963 b_in: &CudaSlice<f32>,
14964 w: &CudaSlice<f32>,
14965 res: &mut CudaSlice<f32>,
14966 ncols: usize,
14967 nrows: usize,
14968 eps: f32,
14969 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14970 let nblk = ncols / 32;
14971 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
14972 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
14973 let f = self.func("add_rms_norm_q8_1");
14974 let cfg = LaunchConfig {
14976 grid_dim: (nrows as u32, 1, 1),
14977 block_dim: (1024, 1, 1),
14978 shared_mem_bytes: 0,
14979 };
14980 let (nc, e) = (ncols as i32, eps);
14981 let __s_bld = self.gpu.stream();
14982 let mut bld = __s_bld.launch_builder(&f);
14983 bld.arg(a)
14984 .arg(b_in)
14985 .arg(w)
14986 .arg(res)
14987 .arg(&mut q)
14988 .arg(&mut d)
14989 .arg(&nc)
14990 .arg(&e);
14991 unsafe {
14992 bld.launch(cfg)?;
14993 }
14994 Ok((q, d))
14995 }
14996
14997 #[allow(clippy::too_many_arguments)]
15003 pub fn join_add_rms_norm_raw(
15004 &self,
15005 a0_raw: u64,
15006 a1_raw: u64,
15007 x: &CudaSlice<f32>,
15008 w: &CudaSlice<f32>,
15009 res: &mut CudaSlice<f32>,
15010 dst: &mut CudaSlice<f32>,
15011 ncols: usize,
15012 eps: f32,
15013 ) -> Result<(), Box<dyn std::error::Error>> {
15014 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
15015 return Err("join_add_rms_norm geometry".into());
15016 }
15017 let f = self.func("join_add_rms_norm_f32");
15018 let cfg = LaunchConfig {
15019 grid_dim: (1, 1, 1),
15020 block_dim: (rms_block(), 1, 1),
15021 shared_mem_bytes: 0,
15022 };
15023 let (nc, e) = (ncols as i32, eps);
15024 let __s_b = self.gpu.stream();
15025 let mut b = __s_b.launch_builder(&f);
15026 b.arg(&a0_raw)
15027 .arg(&a1_raw)
15028 .arg(x)
15029 .arg(w)
15030 .arg(&mut *res)
15031 .arg(&mut *dst)
15032 .arg(&nc)
15033 .arg(&e);
15034 unsafe {
15035 b.launch(cfg)?;
15036 }
15037 Ok(())
15038 }
15039
15040 #[allow(clippy::too_many_arguments)] pub fn add_rms_norm(
15042 &self,
15043 a: &CudaSlice<f32>,
15044 b: &CudaSlice<f32>,
15045 w: &CudaSlice<f32>,
15046 res: &mut CudaSlice<f32>,
15047 dst: &mut CudaSlice<f32>,
15048 ncols: usize,
15049 nrows: usize,
15050 eps: f32,
15051 ) -> Result<(), Box<dyn std::error::Error>> {
15052 let (nc, e) = (ncols as i32, eps);
15053 let kname = if Self::norm_ilp_on() {
15054 "add_rms_norm_f32_v2"
15055 } else {
15056 "add_rms_norm_f32"
15057 };
15058 if Self::pdl_on() && Self::pdl_wb_on() {
15059 use cudarc::driver::{DevicePtr, DevicePtrMut};
15060 let s = &self.gpu.stream();
15061 let (pa, _g0) = a.device_ptr(s);
15062 let (pb, _g1) = b.device_ptr(s);
15063 let (pw, _g2) = w.device_ptr(s);
15064 let (pr, _g3) = res.device_ptr_mut(s);
15065 let (pd, _g4) = dst.device_ptr_mut(s);
15066 let mut ps = [
15067 &pa as *const _ as *mut std::ffi::c_void,
15068 &pb as *const _ as *mut _,
15069 &pw as *const _ as *mut _,
15070 &pr as *const _ as *mut _,
15071 &pd as *const _ as *mut _,
15072 &nc as *const _ as *mut _,
15073 &e as *const _ as *mut _,
15074 ];
15075 unsafe {
15076 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
15077 }
15078 return Ok(());
15079 }
15080 let f = self.func(kname);
15081 let cfg = LaunchConfig {
15082 grid_dim: (nrows as u32, 1, 1),
15083 block_dim: (rms_block(), 1, 1),
15084 shared_mem_bytes: 0,
15085 };
15086 let __s_b2 = self.gpu.stream();
15087 let mut b2 = __s_b2.launch_builder(&f);
15088 b2.arg(a)
15089 .arg(b)
15090 .arg(w)
15091 .arg(&mut *res)
15092 .arg(&mut *dst)
15093 .arg(&nc)
15094 .arg(&e);
15095 unsafe {
15096 b2.launch(cfg)?;
15097 }
15098 Ok(())
15099 }
15100
15101 #[allow(clippy::too_many_arguments)]
15104 pub fn rms_pre_add_rms_norm(
15105 &self,
15106 a: &CudaSlice<f32>,
15107 wa: &CudaSlice<f32>,
15108 b: &CudaSlice<f32>,
15109 w: &CudaSlice<f32>,
15110 res: &mut CudaSlice<f32>,
15111 dst: &mut CudaSlice<f32>,
15112 ncols: usize,
15113 nrows: usize,
15114 eps: f32,
15115 ) -> Result<(), Box<dyn std::error::Error>> {
15116 let f = self.func("rms_pre_add_rms_norm_f32");
15117 let cfg = LaunchConfig {
15118 grid_dim: (nrows as u32, 1, 1),
15119 block_dim: (rms_block(), 1, 1),
15120 shared_mem_bytes: 0,
15121 };
15122 let (nc, e) = (ncols as i32, eps);
15123 let __s_b2 = self.gpu.stream();
15124 let mut b2 = __s_b2.launch_builder(&f);
15125 b2.arg(a)
15126 .arg(wa)
15127 .arg(b)
15128 .arg(w)
15129 .arg(&mut *res)
15130 .arg(&mut *dst)
15131 .arg(&nc)
15132 .arg(&e);
15133 unsafe {
15134 b2.launch(cfg)?;
15135 }
15136 Ok(())
15137 }
15138
15139 #[allow(clippy::too_many_arguments)]
15141 pub fn rms_pre_add_rms_norm_q8z(
15142 &self,
15143 a: &CudaSlice<f32>,
15144 wa: &CudaSlice<f32>,
15145 b: &CudaSlice<f32>,
15146 w: &CudaSlice<f32>,
15147 res: &mut CudaSlice<f32>,
15148 dst: &mut CudaSlice<f32>,
15149 ncols: usize,
15150 nrows: usize,
15151 eps: f32,
15152 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15153 debug_assert!(ncols.is_multiple_of(128));
15154 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
15155 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
15156 let (nc, e) = (ncols as i32, eps);
15157 if Self::pdl_on() {
15158 {
15159 use cudarc::driver::{DevicePtr, DevicePtrMut};
15160 let s = &self.gpu.stream();
15161 let (pa, _g0) = a.device_ptr(s);
15162 let (pwa, _g1) = wa.device_ptr(s);
15163 let (pb, _g2) = b.device_ptr(s);
15164 let (pw, _g3) = w.device_ptr(s);
15165 let (pr, _g4) = res.device_ptr_mut(s);
15166 let (pdst, _g5) = dst.device_ptr_mut(s);
15167 let (pq, _g6) = out_q.device_ptr_mut(s);
15168 let (pd, _g7) = out_d.device_ptr_mut(s);
15169 let mut ps = [
15170 &pa as *const _ as *mut std::ffi::c_void,
15171 &pwa as *const _ as *mut _,
15172 &pb as *const _ as *mut _,
15173 &pw as *const _ as *mut _,
15174 &pr as *const _ as *mut _,
15175 &pdst as *const _ as *mut _,
15176 &pq as *const _ as *mut _,
15177 &pd as *const _ as *mut _,
15178 &nc as *const _ as *mut _,
15179 &e as *const _ as *mut _,
15180 ];
15181 unsafe {
15182 self.launch_pdl(
15183 "rms_pre_add_rms_norm_q8z_f32",
15184 (nrows as u32, 1, 1),
15185 (rms_block(), 1, 1),
15186 &mut ps,
15187 )?;
15188 }
15189 }
15190 return Ok((out_q, out_d));
15191 }
15192 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
15193 let cfg = LaunchConfig {
15194 grid_dim: (nrows as u32, 1, 1),
15195 block_dim: (rms_block(), 1, 1),
15196 shared_mem_bytes: 0,
15197 };
15198 let __s_b2 = self.gpu.stream();
15199 let mut b2 = __s_b2.launch_builder(&f);
15200 b2.arg(a)
15201 .arg(wa)
15202 .arg(b)
15203 .arg(w)
15204 .arg(&mut *res)
15205 .arg(&mut *dst)
15206 .arg(&mut out_q)
15207 .arg(&mut out_d)
15208 .arg(&nc)
15209 .arg(&e);
15210 unsafe {
15211 b2.launch(cfg)?;
15212 }
15213 Ok((out_q, out_d))
15214 }
15215
15216 #[allow(clippy::too_many_arguments)]
15220 pub fn rms_pre_add_rms_norm_q8z_into(
15221 &self,
15222 a: &CudaSlice<f32>,
15223 wa: &CudaSlice<f32>,
15224 b: &CudaSlice<f32>,
15225 w: &CudaSlice<f32>,
15226 res: &mut CudaSlice<f32>,
15227 dst: &mut CudaSlice<f32>,
15228 ncols: usize,
15229 nrows: usize,
15230 eps: f32,
15231 out_q: &mut CudaSlice<i8>,
15232 out_d: &mut CudaSlice<f32>,
15233 ) -> Result<(), Box<dyn std::error::Error>> {
15234 debug_assert!(ncols.is_multiple_of(128));
15235 let (nc, e) = (ncols as i32, eps);
15236 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
15237 let cfg = LaunchConfig {
15238 grid_dim: (nrows as u32, 1, 1),
15239 block_dim: (rms_block(), 1, 1),
15240 shared_mem_bytes: 0,
15241 };
15242 let __s_b = self.gpu.stream();
15243 let mut b2 = __s_b.launch_builder(&f);
15244 b2.arg(a)
15245 .arg(wa)
15246 .arg(b)
15247 .arg(w)
15248 .arg(&mut *res)
15249 .arg(&mut *dst)
15250 .arg(&mut *out_q)
15251 .arg(&mut *out_d)
15252 .arg(&nc)
15253 .arg(&e);
15254 unsafe {
15255 b2.launch(cfg)?;
15256 }
15257 Ok(())
15258 }
15259
15260 #[allow(clippy::too_many_arguments)]
15263 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
15264 &self,
15265 a: &CudaSlice<f32>,
15266 wa: &CudaSlice<f32>,
15267 b_in: &CudaSlice<f32>,
15268 c: f32,
15269 w: &CudaSlice<f32>,
15270 res: &mut CudaSlice<f32>,
15271 ncols: usize,
15272 nrows: usize,
15273 eps: f32,
15274 out_q: &mut CudaSlice<i8>,
15275 out_d: &mut CudaSlice<f32>,
15276 ) -> Result<(), Box<dyn std::error::Error>> {
15277 debug_assert!(ncols.is_multiple_of(128));
15278 let (nc, e2) = (ncols as i32, eps);
15279 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
15280 let cfg = LaunchConfig {
15281 grid_dim: (nrows as u32, 1, 1),
15282 block_dim: (rms_block(), 1, 1),
15283 shared_mem_bytes: 0,
15284 };
15285 let __s_b = self.gpu.stream();
15286 let mut b2 = __s_b.launch_builder(&f);
15287 b2.arg(a)
15288 .arg(wa)
15289 .arg(b_in)
15290 .arg(&c)
15291 .arg(w)
15292 .arg(&mut *res)
15293 .arg(&mut *out_q)
15294 .arg(&mut *out_d)
15295 .arg(&nc)
15296 .arg(&e2);
15297 unsafe {
15298 b2.launch(cfg)?;
15299 }
15300 Ok(())
15301 }
15302
15303 pub fn g4_pnfold_on() -> bool {
15311 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15312 *ON.get_or_init(|| {
15313 std::env::var("MEMRA_G4_PNFOLD")
15314 .map(|v| v != "0")
15315 .unwrap_or(true)
15316 })
15317 }
15318
15319 pub fn build_q4_out_concat3(
15323 &self,
15324 w0: &crate::model::GpuTensor,
15325 w1: &crate::model::GpuTensor,
15326 w2: &crate::model::GpuTensor,
15327 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
15328 use crate::model::GpuTensor;
15329 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
15330 match w {
15331 GpuTensor::Quant {
15332 qtype,
15333 row_bytes,
15334 rp,
15335 ..
15336 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
15337 _ => None,
15338 }
15339 };
15340 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
15341 else {
15342 return Ok(None);
15343 };
15344 if rb0 != rb1
15345 || rb0 != rb2
15346 || w0.in_features() != w1.in_features()
15347 || w0.in_features() != w2.in_features()
15348 {
15349 return Ok(None);
15350 }
15351 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
15352 match w {
15353 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
15354 _ => unreachable!(),
15355 }
15356 }
15357 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
15358 let total = rb0 * (o0 + o1 + o2);
15359 let mut cat = self.alloc_u8(total)?;
15360 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
15361 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
15362 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
15363 Ok(Some(GpuTensor::Quant {
15364 bytes: cat,
15365 qtype: QT_Q4_0,
15366 row_bytes: rb0,
15367 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
15368 scale: 1.0,
15369 rp: false,
15370 #[cfg(memra_cutlass)]
15371 cutlass: None,
15372 fp8: None,
15373 blk: None,
15374 rp4: None,
15375 f16: None,
15376 }))
15377 }
15378
15379 fn full_width_rope_only(
15397 kernel: &str,
15398 n_rot: usize,
15399 head_dim: usize,
15400 ) -> Result<(), Box<dyn std::error::Error>> {
15401 if n_rot == head_dim {
15402 return Ok(());
15403 }
15404 Err(format!(
15405 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
15406 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
15407 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
15408 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
15409 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
15410 )
15411 .into())
15412 }
15413
15414 #[allow(clippy::too_many_arguments)]
15418 pub fn rms_norm_qkv_rope_cat(
15419 &self,
15420 qkv: &CudaSlice<f32>,
15421 wq: &CudaSlice<f32>,
15422 wk: &CudaSlice<f32>,
15423 wv: &CudaSlice<f32>,
15424 q: &mut CudaSlice<f32>,
15425 k: &mut CudaSlice<f32>,
15426 v: &mut CudaSlice<f32>,
15427 head_dim: usize,
15428 n_rot: usize,
15429 rq: usize,
15430 rk: usize,
15431 pos: &CudaSlice<i32>,
15432 nh_q: usize,
15433 nh_k: usize,
15434 base: f32,
15435 freq_scale: f32,
15436 ff: Option<&CudaSlice<f32>>,
15437 eps: f32,
15438 ) -> Result<(), Box<dyn std::error::Error>> {
15439 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
15440 let rows = rq + rk + rk;
15441 let theta_scale = base.powf(-2.0 / head_dim as f32);
15442 let (nc, rqi, rki, nhq, nhk) = (
15443 head_dim as i32,
15444 rq as i32,
15445 rk as i32,
15446 nh_q as i32,
15447 nh_k as i32,
15448 );
15449 if Self::pdl_on() {
15450 use cudarc::driver::{DevicePtr, DevicePtrMut};
15451 let s = &self.gpu.stream();
15452 let (pqkv, _g0) = qkv.device_ptr(s);
15453 let (pwq, _g1) = wq.device_ptr(s);
15454 let (pwk, _g2) = wk.device_ptr(s);
15455 let (pwv, _g3) = wv.device_ptr(s);
15456 let (pq, _g4) = q.device_ptr_mut(s);
15457 let (pk, _g5) = k.device_ptr_mut(s);
15458 let (pv, _g6) = v.device_ptr_mut(s);
15459 let (ppos, _g7) = pos.device_ptr(s);
15460 let (pff, _g8) = match ff {
15461 Some(t) => {
15462 let (p, g) = t.device_ptr(s);
15463 (p, Some(g))
15464 }
15465 None => (0, None),
15466 };
15467 let mut ps = [
15468 &pqkv as *const _ as *mut std::ffi::c_void,
15469 &pwq as *const _ as *mut _,
15470 &pwk as *const _ as *mut _,
15471 &pwv as *const _ as *mut _,
15472 &pq as *const _ as *mut _,
15473 &pk as *const _ as *mut _,
15474 &pv as *const _ as *mut _,
15475 &nc as *const _ as *mut _,
15476 &rqi as *const _ as *mut _,
15477 &rki as *const _ as *mut _,
15478 &ppos as *const _ as *mut _,
15479 &nhq as *const _ as *mut _,
15480 &nhk as *const _ as *mut _,
15481 &theta_scale as *const _ as *mut _,
15482 &freq_scale as *const _ as *mut _,
15483 &pff as *const _ as *mut _,
15484 &eps as *const _ as *mut _,
15485 ];
15486 unsafe {
15487 self.launch_pdl(
15488 "rms_norm_qkv_rope_cat_f32",
15489 (rows as u32, 1, 1),
15490 (rms_block(), 1, 1),
15491 &mut ps,
15492 )?;
15493 }
15494 return Ok(());
15495 }
15496 let f = self.func("rms_norm_qkv_rope_cat_f32");
15497 let cfg = LaunchConfig {
15498 grid_dim: (rows as u32, 1, 1),
15499 block_dim: (rms_block(), 1, 1),
15500 shared_mem_bytes: 0,
15501 };
15502 let __s_b = self.gpu.stream();
15503 let mut b = __s_b.launch_builder(&f);
15504 match ff {
15505 Some(t) => {
15506 b.arg(qkv)
15507 .arg(wq)
15508 .arg(wk)
15509 .arg(wv)
15510 .arg(&mut *q)
15511 .arg(&mut *k)
15512 .arg(&mut *v)
15513 .arg(&nc)
15514 .arg(&rqi)
15515 .arg(&rki)
15516 .arg(pos)
15517 .arg(&nhq)
15518 .arg(&nhk)
15519 .arg(&theta_scale)
15520 .arg(&freq_scale)
15521 .arg(t)
15522 .arg(&eps);
15523 unsafe {
15524 b.launch(cfg)?;
15525 }
15526 }
15527 None => {
15528 let null: u64 = 0;
15529 b.arg(qkv)
15530 .arg(wq)
15531 .arg(wk)
15532 .arg(wv)
15533 .arg(&mut *q)
15534 .arg(&mut *k)
15535 .arg(&mut *v)
15536 .arg(&nc)
15537 .arg(&rqi)
15538 .arg(&rki)
15539 .arg(pos)
15540 .arg(&nhq)
15541 .arg(&nhk)
15542 .arg(&theta_scale)
15543 .arg(&freq_scale)
15544 .arg(&null)
15545 .arg(&eps);
15546 unsafe {
15547 b.launch(cfg)?;
15548 }
15549 }
15550 }
15551 Ok(())
15552 }
15553
15554 #[allow(clippy::too_many_arguments)]
15558 pub fn rms_norm_qkv_rope(
15559 &self,
15560 q0: &CudaSlice<f32>,
15561 k0: &CudaSlice<f32>,
15562 v0: &CudaSlice<f32>,
15563 wq: &CudaSlice<f32>,
15564 wk: &CudaSlice<f32>,
15565 wv: &CudaSlice<f32>,
15566 q: &mut CudaSlice<f32>,
15567 k: &mut CudaSlice<f32>,
15568 v: &mut CudaSlice<f32>,
15569 head_dim: usize,
15570 n_rot: usize,
15571 rq: usize,
15572 rk: usize,
15573 pos: &CudaSlice<i32>,
15574 nh_q: usize,
15575 nh_k: usize,
15576 base: f32,
15577 freq_scale: f32,
15578 ff: Option<&CudaSlice<f32>>,
15579 eps: f32,
15580 ) -> Result<(), Box<dyn std::error::Error>> {
15581 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
15582 let f = self.func("rms_norm_qkv_rope_f32");
15583 let rows = rq + rk + rk; let cfg = LaunchConfig {
15585 grid_dim: (rows as u32, 1, 1),
15586 block_dim: (rms_block(), 1, 1),
15587 shared_mem_bytes: 0,
15588 };
15589 let theta_scale = base.powf(-2.0 / head_dim as f32);
15590 let (nc, rqi, rki, nhq, nhk) = (
15591 head_dim as i32,
15592 rq as i32,
15593 rk as i32,
15594 nh_q as i32,
15595 nh_k as i32,
15596 );
15597 let __s_b = self.gpu.stream();
15598 let mut b = __s_b.launch_builder(&f);
15599 match ff {
15600 Some(t) => {
15601 b.arg(q0)
15602 .arg(k0)
15603 .arg(v0)
15604 .arg(wq)
15605 .arg(wk)
15606 .arg(wv)
15607 .arg(&mut *q)
15608 .arg(&mut *k)
15609 .arg(&mut *v)
15610 .arg(&nc)
15611 .arg(&rqi)
15612 .arg(&rki)
15613 .arg(pos)
15614 .arg(&nhq)
15615 .arg(&nhk)
15616 .arg(&theta_scale)
15617 .arg(&freq_scale)
15618 .arg(t)
15619 .arg(&eps);
15620 unsafe {
15621 b.launch(cfg)?;
15622 }
15623 }
15624 None => {
15625 let null: u64 = 0;
15626 b.arg(q0)
15627 .arg(k0)
15628 .arg(v0)
15629 .arg(wq)
15630 .arg(wk)
15631 .arg(wv)
15632 .arg(&mut *q)
15633 .arg(&mut *k)
15634 .arg(&mut *v)
15635 .arg(&nc)
15636 .arg(&rqi)
15637 .arg(&rki)
15638 .arg(pos)
15639 .arg(&nhq)
15640 .arg(&nhk)
15641 .arg(&theta_scale)
15642 .arg(&freq_scale)
15643 .arg(&null)
15644 .arg(&eps);
15645 unsafe {
15646 b.launch(cfg)?;
15647 }
15648 }
15649 }
15650 Ok(())
15651 }
15652
15653 #[allow(clippy::too_many_arguments)]
15659 pub fn rms_norm_qkv_rope_append_dc(
15660 &self,
15661 q0: &CudaSlice<f32>,
15662 k0: &CudaSlice<f32>,
15663 v0: &CudaSlice<f32>,
15664 wq: &CudaSlice<f32>,
15665 wk: &CudaSlice<f32>,
15666 wv: &CudaSlice<f32>,
15667 q: &mut CudaSlice<f32>,
15668 k: &mut CudaSlice<f32>,
15669 v: &mut CudaSlice<f32>,
15670 head_dim: usize,
15671 n_rot: usize,
15672 rq: usize,
15673 rk: usize,
15674 pos: &CudaSlice<i32>,
15675 nh_q: usize,
15676 nh_k: usize,
15677 base: f32,
15678 freq_scale: f32,
15679 ff: Option<&CudaSlice<f32>>,
15680 eps: f32,
15681 kc: &mut CudaSlice<u8>,
15682 vc: &mut CudaSlice<u8>,
15683 t_dev: &CudaSlice<i32>,
15684 k_tok_bytes: usize,
15685 v_tok_bytes: usize,
15686 g: bool,
15687 ) -> Result<(), Box<dyn std::error::Error>> {
15688 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
15689 let rows = rq + rk + rk;
15690 let theta_scale = base.powf(-2.0 / head_dim as f32);
15691 let (nc, rqi, rki, nhq, nhk) = (
15692 head_dim as i32,
15693 rq as i32,
15694 rk as i32,
15695 nh_q as i32,
15696 nh_k as i32,
15697 );
15698 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15699 if Self::pdl_on() && Self::pdl_wb_on() {
15700 use cudarc::driver::{DevicePtr, DevicePtrMut};
15701 let s = &self.gpu.stream();
15702 let (p0, _a0) = q0.device_ptr(s);
15703 let (p1, _a1) = k0.device_ptr(s);
15704 let (p2, _a2) = v0.device_ptr(s);
15705 let (pwq, _a3) = wq.device_ptr(s);
15706 let (pwk, _a4) = wk.device_ptr(s);
15707 let (pwv, _a5) = wv.device_ptr(s);
15708 let (pq, _a6) = q.device_ptr_mut(s);
15709 let (pk, _a7) = k.device_ptr_mut(s);
15710 let (pv, _a8) = v.device_ptr_mut(s);
15711 let (pp, _a9) = pos.device_ptr(s);
15712 let pff: u64 = match ff {
15713 Some(t) => {
15714 let (p, _gg) = t.device_ptr(s);
15715 p
15716 }
15717 None => 0,
15718 };
15719 let (pkc, _a10) = kc.device_ptr_mut(s);
15720 let (pvc, _a11) = vc.device_ptr_mut(s);
15721 let (pt, _a12) = t_dev.device_ptr(s);
15722 let mut ps = [
15723 &p0 as *const _ as *mut std::ffi::c_void,
15724 &p1 as *const _ as *mut _,
15725 &p2 as *const _ as *mut _,
15726 &pwq as *const _ as *mut _,
15727 &pwk as *const _ as *mut _,
15728 &pwv as *const _ as *mut _,
15729 &pq as *const _ as *mut _,
15730 &pk as *const _ as *mut _,
15731 &pv as *const _ as *mut _,
15732 &nc as *const _ as *mut _,
15733 &rqi as *const _ as *mut _,
15734 &rki as *const _ as *mut _,
15735 &pp as *const _ as *mut _,
15736 &nhq as *const _ as *mut _,
15737 &nhk as *const _ as *mut _,
15738 &theta_scale as *const _ as *mut _,
15739 &freq_scale as *const _ as *mut _,
15740 &pff as *const _ as *mut _,
15741 &eps as *const _ as *mut _,
15742 &pkc as *const _ as *mut _,
15743 &pvc as *const _ as *mut _,
15744 &pt as *const _ as *mut _,
15745 &ktb as *const _ as *mut _,
15746 &vtb as *const _ as *mut _,
15747 ];
15748 unsafe {
15749 self.launch_pdl_flash(
15750 g,
15751 "rms_norm_qkv_rope_append_dc_f32",
15752 (rows as u32, 1, 1),
15753 (rms_block(), 1, 1),
15754 0,
15755 &mut ps,
15756 )?;
15757 }
15758 return Ok(());
15759 }
15760 let f = if g {
15761 self.func_g("rms_norm_qkv_rope_append_dc_f32")
15762 } else {
15763 self.func("rms_norm_qkv_rope_append_dc_f32")
15764 };
15765 let cfg = LaunchConfig {
15766 grid_dim: (rows as u32, 1, 1),
15767 block_dim: (rms_block(), 1, 1),
15768 shared_mem_bytes: 0,
15769 };
15770 let __s_b = self.gpu.stream();
15771 let mut b = __s_b.launch_builder(&f);
15772 match ff {
15773 Some(t) => {
15774 b.arg(q0)
15775 .arg(k0)
15776 .arg(v0)
15777 .arg(wq)
15778 .arg(wk)
15779 .arg(wv)
15780 .arg(&mut *q)
15781 .arg(&mut *k)
15782 .arg(&mut *v)
15783 .arg(&nc)
15784 .arg(&rqi)
15785 .arg(&rki)
15786 .arg(pos)
15787 .arg(&nhq)
15788 .arg(&nhk)
15789 .arg(&theta_scale)
15790 .arg(&freq_scale)
15791 .arg(t)
15792 .arg(&eps)
15793 .arg(&mut *kc)
15794 .arg(&mut *vc)
15795 .arg(t_dev)
15796 .arg(&ktb)
15797 .arg(&vtb);
15798 unsafe {
15799 b.launch(cfg)?;
15800 }
15801 }
15802 None => {
15803 let null: u64 = 0;
15804 b.arg(q0)
15805 .arg(k0)
15806 .arg(v0)
15807 .arg(wq)
15808 .arg(wk)
15809 .arg(wv)
15810 .arg(&mut *q)
15811 .arg(&mut *k)
15812 .arg(&mut *v)
15813 .arg(&nc)
15814 .arg(&rqi)
15815 .arg(&rki)
15816 .arg(pos)
15817 .arg(&nhq)
15818 .arg(&nhk)
15819 .arg(&theta_scale)
15820 .arg(&freq_scale)
15821 .arg(&null)
15822 .arg(&eps)
15823 .arg(&mut *kc)
15824 .arg(&mut *vc)
15825 .arg(t_dev)
15826 .arg(&ktb)
15827 .arg(&vtb);
15828 unsafe {
15829 b.launch(cfg)?;
15830 }
15831 }
15832 }
15833 Ok(())
15834 }
15835
15836 #[allow(clippy::too_many_arguments)]
15844 pub fn rms_norm_qkv_rope_append(
15845 &self,
15846 q0: &CudaSlice<f32>,
15847 k0: &CudaSlice<f32>,
15848 v0: &CudaSlice<f32>,
15849 wq: &CudaSlice<f32>,
15850 wk: &CudaSlice<f32>,
15851 wv: &CudaSlice<f32>,
15852 q: &mut CudaSlice<f32>,
15853 k: &mut CudaSlice<f32>,
15854 v: &mut CudaSlice<f32>,
15855 head_dim: usize,
15856 n_rot: usize,
15857 rq: usize,
15858 rk: usize,
15859 pos: &CudaSlice<i32>,
15860 nh_q: usize,
15861 nh_k: usize,
15862 base: f32,
15863 freq_scale: f32,
15864 ff: Option<&CudaSlice<f32>>,
15865 eps: f32,
15866 kc: &mut CudaSlice<u8>,
15867 vc: &mut CudaSlice<u8>,
15868 t: usize,
15869 k_tok_bytes: usize,
15870 v_tok_bytes: usize,
15871 g: bool,
15872 ) -> Result<(), Box<dyn std::error::Error>> {
15873 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
15874 let rows = rq + rk + rk;
15875 let theta_scale = base.powf(-2.0 / head_dim as f32);
15876 let (nc, rqi, rki, nhq, nhk) = (
15877 head_dim as i32,
15878 rq as i32,
15879 rk as i32,
15880 nh_q as i32,
15881 nh_k as i32,
15882 );
15883 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
15884 let ti = t as i32;
15885 if Self::pdl_on() && Self::pdl_wb_on() {
15886 use cudarc::driver::{DevicePtr, DevicePtrMut};
15887 let s = &self.gpu.stream();
15888 let (p0, _a0) = q0.device_ptr(s);
15889 let (p1, _a1) = k0.device_ptr(s);
15890 let (p2, _a2) = v0.device_ptr(s);
15891 let (pwq, _a3) = wq.device_ptr(s);
15892 let (pwk, _a4) = wk.device_ptr(s);
15893 let (pwv, _a5) = wv.device_ptr(s);
15894 let (pq, _a6) = q.device_ptr_mut(s);
15895 let (pk, _a7) = k.device_ptr_mut(s);
15896 let (pv, _a8) = v.device_ptr_mut(s);
15897 let (pp, _a9) = pos.device_ptr(s);
15898 let pff: u64 = match ff {
15899 Some(t) => {
15900 let (p, _gg) = t.device_ptr(s);
15901 p
15902 }
15903 None => 0,
15904 };
15905 let (pkc, _a10) = kc.device_ptr_mut(s);
15906 let (pvc, _a11) = vc.device_ptr_mut(s);
15907 let mut ps = [
15908 &p0 as *const _ as *mut std::ffi::c_void,
15909 &p1 as *const _ as *mut _,
15910 &p2 as *const _ as *mut _,
15911 &pwq as *const _ as *mut _,
15912 &pwk as *const _ as *mut _,
15913 &pwv as *const _ as *mut _,
15914 &pq as *const _ as *mut _,
15915 &pk as *const _ as *mut _,
15916 &pv as *const _ as *mut _,
15917 &nc as *const _ as *mut _,
15918 &rqi as *const _ as *mut _,
15919 &rki as *const _ as *mut _,
15920 &pp as *const _ as *mut _,
15921 &nhq as *const _ as *mut _,
15922 &nhk as *const _ as *mut _,
15923 &theta_scale as *const _ as *mut _,
15924 &freq_scale as *const _ as *mut _,
15925 &pff as *const _ as *mut _,
15926 &eps as *const _ as *mut _,
15927 &pkc as *const _ as *mut _,
15928 &pvc as *const _ as *mut _,
15929 &ti as *const _ as *mut _,
15930 &ktb as *const _ as *mut _,
15931 &vtb as *const _ as *mut _,
15932 ];
15933 unsafe {
15934 self.launch_pdl_flash(
15935 g,
15936 "rms_norm_qkv_rope_append_f32",
15937 (rows as u32, 1, 1),
15938 (rms_block(), 1, 1),
15939 0,
15940 &mut ps,
15941 )?;
15942 }
15943 return Ok(());
15944 }
15945 let f = if g {
15946 self.func_g("rms_norm_qkv_rope_append_f32")
15947 } else {
15948 self.func("rms_norm_qkv_rope_append_f32")
15949 };
15950 let cfg = LaunchConfig {
15951 grid_dim: (rows as u32, 1, 1),
15952 block_dim: (rms_block(), 1, 1),
15953 shared_mem_bytes: 0,
15954 };
15955 let __s_b = self.gpu.stream();
15956 let mut b = __s_b.launch_builder(&f);
15957 let null: u64 = 0;
15958 b.arg(q0)
15959 .arg(k0)
15960 .arg(v0)
15961 .arg(wq)
15962 .arg(wk)
15963 .arg(wv)
15964 .arg(&mut *q)
15965 .arg(&mut *k)
15966 .arg(&mut *v)
15967 .arg(&nc)
15968 .arg(&rqi)
15969 .arg(&rki)
15970 .arg(pos)
15971 .arg(&nhq)
15972 .arg(&nhk)
15973 .arg(&theta_scale)
15974 .arg(&freq_scale);
15975 match ff {
15976 Some(t) => {
15977 b.arg(t);
15978 }
15979 None => {
15980 b.arg(&null);
15981 }
15982 }
15983 b.arg(&eps)
15984 .arg(&mut *kc)
15985 .arg(&mut *vc)
15986 .arg(&ti)
15987 .arg(&ktb)
15988 .arg(&vtb);
15989 unsafe {
15990 b.launch(cfg)?;
15991 }
15992 Ok(())
15993 }
15994
15995 pub fn add_q8_1(
15996 &self,
15997 a: &CudaSlice<f32>,
15998 b: &CudaSlice<f32>,
15999 res: &mut CudaSlice<f32>,
16000 ncols: usize,
16001 nrows: usize,
16002 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16003 debug_assert!(ncols.is_multiple_of(128));
16004 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
16005 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
16006 let f = self.func("add_q8_1_f32");
16007 let cfg = LaunchConfig {
16008 grid_dim: (nrows as u32, 1, 1),
16009 block_dim: (rms_block(), 1, 1),
16010 shared_mem_bytes: 0,
16011 };
16012 let nc = ncols as i32;
16013 let __s_b2 = self.gpu.stream();
16014 let mut b2 = __s_b2.launch_builder(&f);
16015 b2.arg(a)
16016 .arg(b)
16017 .arg(&mut *res)
16018 .arg(&mut out_q)
16019 .arg(&mut out_d)
16020 .arg(&nc);
16021 unsafe {
16022 b2.launch(cfg)?;
16023 }
16024 Ok((out_q, out_d))
16025 }
16026
16027 #[allow(clippy::too_many_arguments)] pub fn rms_pre_add_q8_1(
16032 &self,
16033 a: &CudaSlice<f32>,
16034 wa: &CudaSlice<f32>,
16035 b: &CudaSlice<f32>,
16036 res: &mut CudaSlice<f32>,
16037 ncols: usize,
16038 nrows: usize,
16039 eps: f32,
16040 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16041 debug_assert!(ncols.is_multiple_of(128));
16042 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
16043 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
16044 let f = self.func("rms_pre_add_q8_1_f32");
16045 let cfg = LaunchConfig {
16046 grid_dim: (nrows as u32, 1, 1),
16047 block_dim: (rms_block(), 1, 1),
16048 shared_mem_bytes: 0,
16049 };
16050 let (nc, ep) = (ncols as i32, eps);
16051 let __s_b2 = self.gpu.stream();
16052 let mut b2 = __s_b2.launch_builder(&f);
16053 b2.arg(a)
16054 .arg(wa)
16055 .arg(b)
16056 .arg(&mut *res)
16057 .arg(&mut out_q)
16058 .arg(&mut out_d)
16059 .arg(&nc)
16060 .arg(&ep);
16061 unsafe {
16062 b2.launch(cfg)?;
16063 }
16064 Ok((out_q, out_d))
16065 }
16066
16067 pub fn l2_v2_on(ncols: usize) -> bool {
16071 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
16072 }
16073
16074 pub fn l2_norm_pp(
16075 &self,
16076 x: &CudaSlice<f32>,
16077 dst: &mut CudaSlice<f32>,
16078 dst16: Option<&mut CudaSlice<u8>>,
16079 ncols: usize,
16080 nrows: usize,
16081 eps: f32,
16082 ) -> Result<(), Box<dyn std::error::Error>> {
16083 if Self::l2_v2_on(ncols) {
16084 let f = self.func("l2_norm_pp_v2_f32");
16085 let rows_per_block = 8u32; let cfg = LaunchConfig {
16087 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
16088 block_dim: (256, 1, 1),
16089 shared_mem_bytes: 0,
16090 };
16091 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
16092 let d16: u64 = match dst16 {
16094 Some(d) => self.addr_u8(d),
16095 None => 0,
16096 };
16097 let __s_b = self.gpu.stream();
16098 let mut b = __s_b.launch_builder(&f);
16099 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
16100 unsafe {
16101 b.launch(cfg)?;
16102 }
16103 return Ok(());
16104 }
16105 self.l2_norm(x, dst, ncols, nrows, eps)
16106 }
16107
16108 pub fn l2_norm(
16109 &self,
16110 x: &CudaSlice<f32>,
16111 dst: &mut CudaSlice<f32>,
16112 ncols: usize,
16113 nrows: usize,
16114 eps: f32,
16115 ) -> Result<(), Box<dyn std::error::Error>> {
16116 let f = self.func("l2_norm_f32");
16117 let cfg = LaunchConfig {
16118 grid_dim: (nrows as u32, 1, 1),
16119 block_dim: (256, 1, 1),
16120 shared_mem_bytes: 0,
16121 };
16122 let (nc, e) = (ncols as i32, eps);
16123 let __s_b = self.gpu.stream();
16124 let mut b = __s_b.launch_builder(&f);
16125 b.arg(x).arg(dst).arg(&nc).arg(&e);
16126 unsafe {
16127 b.launch(cfg)?;
16128 }
16129 Ok(())
16130 }
16131
16132 pub fn l2_norm_decode(
16138 &self,
16139 x: &CudaSlice<f32>,
16140 dst: &mut CudaSlice<f32>,
16141 ncols: usize,
16142 nrows: usize,
16143 eps: f32,
16144 ) -> Result<(), Box<dyn std::error::Error>> {
16145 let f = self.func("l2_norm_f32");
16146 let cfg = LaunchConfig {
16147 grid_dim: (nrows as u32, 1, 1),
16148 block_dim: (32, 1, 1),
16149 shared_mem_bytes: 0,
16150 };
16151 let (nc, e) = (ncols as i32, eps);
16152 let __s_b = self.gpu.stream();
16153 let mut b = __s_b.launch_builder(&f);
16154 b.arg(x).arg(dst).arg(&nc).arg(&e);
16155 unsafe {
16156 b.launch(cfg)?;
16157 }
16158 Ok(())
16159 }
16160
16161 #[allow(clippy::too_many_arguments)] pub fn rope_neox(
16164 &self,
16165 x: &mut CudaSlice<f32>,
16166 pos: &CudaSlice<i32>,
16167 head_dim: usize,
16168 n_dims: usize,
16169 n_heads: usize,
16170 n_tokens: usize,
16171 freq_base: f32,
16172 freq_scale: f32,
16173 ) -> Result<(), Box<dyn std::error::Error>> {
16174 let f = self.func("rope_neox_f32");
16175 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16176 let grid = (n_heads * n_tokens) as u32;
16177 let cfg = LaunchConfig {
16178 grid_dim: (grid, 1, 1),
16179 block_dim: ((head_dim / 2) as u32, 1, 1),
16180 shared_mem_bytes: 0,
16181 };
16182 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
16183 let __s_b = self.gpu.stream();
16184 let mut b = __s_b.launch_builder(&f);
16185 b.arg(x)
16186 .arg(pos)
16187 .arg(&hd)
16188 .arg(&nd)
16189 .arg(&nh)
16190 .arg(&theta_scale)
16191 .arg(&freq_scale);
16192 unsafe {
16193 b.launch(cfg)?;
16194 }
16195 Ok(())
16196 }
16197
16198 #[allow(clippy::too_many_arguments)] pub fn rope_neox_ff(
16201 &self,
16202 x: &mut CudaSlice<f32>,
16203 pos: &CudaSlice<i32>,
16204 head_dim: usize,
16205 n_dims: usize,
16206 n_heads: usize,
16207 n_tokens: usize,
16208 freq_base: f32,
16209 freq_scale: f32,
16210 ff: &CudaSlice<f32>,
16211 ) -> Result<(), Box<dyn std::error::Error>> {
16212 let f = self.func("rope_neox_ff_f32");
16213 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16214 let grid = (n_heads * n_tokens) as u32;
16215 let cfg = LaunchConfig {
16216 grid_dim: (grid, 1, 1),
16217 block_dim: ((head_dim / 2) as u32, 1, 1),
16218 shared_mem_bytes: 0,
16219 };
16220 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
16221 let __s_b = self.gpu.stream();
16222 let mut b = __s_b.launch_builder(&f);
16223 b.arg(x)
16224 .arg(pos)
16225 .arg(&hd)
16226 .arg(&nd)
16227 .arg(&nh)
16228 .arg(&theta_scale)
16229 .arg(&freq_scale)
16230 .arg(ff);
16231 unsafe {
16232 b.launch(cfg)?;
16233 }
16234 Ok(())
16235 }
16236
16237 #[allow(clippy::too_many_arguments)]
16241 pub fn rope_neox_ffm(
16242 &self,
16243 x: &mut CudaSlice<f32>,
16244 pos: &CudaSlice<i32>,
16245 head_dim: usize,
16246 n_dims: usize,
16247 n_heads: usize,
16248 n_tokens: usize,
16249 freq_base: f32,
16250 freq_scale: f32,
16251 ff: &CudaSlice<f32>,
16252 mscale: f32,
16253 ) -> Result<(), Box<dyn std::error::Error>> {
16254 let f = self.func("rope_neox_ffm_f32");
16255 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16256 let grid = (n_heads * n_tokens) as u32;
16257 let cfg = LaunchConfig {
16258 grid_dim: (grid, 1, 1),
16259 block_dim: ((head_dim / 2) as u32, 1, 1),
16260 shared_mem_bytes: 0,
16261 };
16262 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
16263 let __s_b = self.gpu.stream();
16264 let mut b = __s_b.launch_builder(&f);
16265 b.arg(x)
16266 .arg(pos)
16267 .arg(&hd)
16268 .arg(&nd)
16269 .arg(&nh)
16270 .arg(&theta_scale)
16271 .arg(&freq_scale)
16272 .arg(ff)
16273 .arg(&mscale);
16274 unsafe {
16275 b.launch(cfg)?;
16276 }
16277 Ok(())
16278 }
16279
16280 #[allow(clippy::too_many_arguments)]
16282 pub fn rope_neox2(
16283 &self,
16284 q: &mut CudaSlice<f32>,
16285 k: &mut CudaSlice<f32>,
16286 pos: &CudaSlice<i32>,
16287 head_dim: usize,
16288 n_dims: usize,
16289 nh_q: usize,
16290 nh_k: usize,
16291 n_tokens: usize,
16292 freq_base: f32,
16293 freq_scale: f32,
16294 ff: Option<&CudaSlice<f32>>,
16295 ) -> Result<(), Box<dyn std::error::Error>> {
16296 let f = self.func("rope_neox2_f32");
16297 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16298 let grid = ((nh_q + nh_k) * n_tokens) as u32;
16299 let cfg = LaunchConfig {
16300 grid_dim: (grid, 1, 1),
16301 block_dim: ((head_dim / 2) as u32, 1, 1),
16302 shared_mem_bytes: 0,
16303 };
16304 let (hd, nd, nq, nk, nt) = (
16305 head_dim as i32,
16306 n_dims as i32,
16307 nh_q as i32,
16308 nh_k as i32,
16309 n_tokens as i32,
16310 );
16311 let __s_b = self.gpu.stream();
16312 let mut b = __s_b.launch_builder(&f);
16313 b.arg(q)
16314 .arg(k)
16315 .arg(pos)
16316 .arg(&hd)
16317 .arg(&nd)
16318 .arg(&nq)
16319 .arg(&nk)
16320 .arg(&nt)
16321 .arg(&theta_scale)
16322 .arg(&freq_scale);
16323 match ff {
16324 Some(ffv) => {
16325 b.arg(ffv);
16326 unsafe {
16327 b.launch(cfg)?;
16328 }
16329 }
16330 None => {
16331 let null: u64 = 0;
16332 b.arg(&null);
16333 unsafe {
16334 b.launch(cfg)?;
16335 }
16336 }
16337 }
16338 Ok(())
16339 }
16340
16341 pub fn gelu_tanh_mul(
16343 &self,
16344 gate: &CudaSlice<f32>,
16345 up: &CudaSlice<f32>,
16346 dst: &mut CudaSlice<f32>,
16347 n: usize,
16348 ) -> Result<(), Box<dyn std::error::Error>> {
16349 let f = self.func("gelu_tanh_mul_f32");
16350 let cfg = LaunchConfig::for_num_elems(n as u32);
16351 let ni = n as i32;
16352 let __s_b = self.gpu.stream();
16353 let mut b = __s_b.launch_builder(&f);
16354 b.arg(gate).arg(up).arg(dst).arg(&ni);
16355 unsafe {
16356 b.launch(cfg)?;
16357 }
16358 Ok(())
16359 }
16360
16361 pub fn silu_mul(
16362 &self,
16363 gate: &CudaSlice<f32>,
16364 up: &CudaSlice<f32>,
16365 dst: &mut CudaSlice<f32>,
16366 n: usize,
16367 ) -> Result<(), Box<dyn std::error::Error>> {
16368 let f = self.func("silu_mul_f32");
16369 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16371 let ni = n as i32;
16372 let __s_b = self.gpu.stream();
16373 let mut b = __s_b.launch_builder(&f);
16374 b.arg(gate).arg(up).arg(dst).arg(&ni);
16375 unsafe {
16376 b.launch(cfg)?;
16377 }
16378 Ok(())
16379 }
16380
16381 pub fn silu_mul_host_expf(
16383 &self,
16384 gate: &CudaSlice<f32>,
16385 up: &CudaSlice<f32>,
16386 dst: &mut CudaSlice<f32>,
16387 n: usize,
16388 ) -> Result<(), Box<dyn std::error::Error>> {
16389 let f = self.func("silu_mul_host_expf_f32");
16390 let cfg = LaunchConfig::for_num_elems(n as u32);
16391 let ni = n as i32;
16392 let __s_b = self.gpu.stream();
16393 let mut b = __s_b.launch_builder(&f);
16394 b.arg(gate).arg(up).arg(dst).arg(&ni);
16395 unsafe {
16396 b.launch(cfg)?;
16397 }
16398 Ok(())
16399 }
16400
16401 pub fn silu_clamped_mul_host_expf(
16403 &self,
16404 gate: &CudaSlice<f32>,
16405 up: &CudaSlice<f32>,
16406 limit: f32,
16407 dst: &mut CudaSlice<f32>,
16408 n: usize,
16409 ) -> Result<(), Box<dyn std::error::Error>> {
16410 if !limit.is_finite() || limit <= 0.0 {
16411 return Err(
16412 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
16413 );
16414 }
16415 let f = self.func("silu_clamped_mul_host_expf_f32");
16416 let cfg = LaunchConfig::for_num_elems(n as u32);
16417 let ni = n as i32;
16418 let __s_b = self.gpu.stream();
16419 let mut b = __s_b.launch_builder(&f);
16420 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
16421 unsafe {
16422 b.launch(cfg)?;
16423 }
16424 Ok(())
16425 }
16426
16427 pub fn silu_mul_f16out(
16430 &self,
16431 gate: &CudaSlice<f32>,
16432 up: &CudaSlice<f32>,
16433 dst: &mut CudaSlice<f32>,
16434 dst16: &mut CudaSlice<u8>,
16435 n: usize,
16436 ) -> Result<(), Box<dyn std::error::Error>> {
16437 let f = self.func("silu_mul_f16out_f32");
16438 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16439 let ni = n as i32;
16440 let __s_b = self.gpu.stream();
16441 let mut b = __s_b.launch_builder(&f);
16442 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
16443 unsafe {
16444 b.launch(cfg)?;
16445 }
16446 Ok(())
16447 }
16448
16449 pub fn silu_mul_scaled(
16456 &self,
16457 gate: &CudaSlice<f32>,
16458 up: &CudaSlice<f32>,
16459 gs: f32,
16460 us: f32,
16461 dst: &mut CudaSlice<f32>,
16462 n: usize,
16463 ) -> Result<(), Box<dyn std::error::Error>> {
16464 let f = self.func("silu_mul_scaled_f32");
16465 let cfg = LaunchConfig::for_num_elems(n as u32);
16466 let ni = n as i32;
16467 let (gsf, usf) = (gs, us);
16468 let __s_b = self.gpu.stream();
16469 let mut b = __s_b.launch_builder(&f);
16470 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
16471 unsafe {
16472 b.launch(cfg)?;
16473 }
16474 Ok(())
16475 }
16476
16477 #[allow(clippy::too_many_arguments)]
16481 pub fn swigluoai_mul_scaled(
16482 &self,
16483 gate: &CudaSlice<f32>,
16484 up: &CudaSlice<f32>,
16485 gs: f32,
16486 us: f32,
16487 alpha: f32,
16488 limit: f32,
16489 dst: &mut CudaSlice<f32>,
16490 n: usize,
16491 ) -> Result<(), Box<dyn std::error::Error>> {
16492 let f = self.func("swigluoai_mul_scaled_f32");
16493 let cfg = LaunchConfig::for_num_elems(n as u32);
16494 let ni = n as i32;
16495 let __s_b = self.gpu.stream();
16496 let mut b = __s_b.launch_builder(&f);
16497 b.arg(gate)
16498 .arg(up)
16499 .arg(&gs)
16500 .arg(&us)
16501 .arg(&alpha)
16502 .arg(&limit)
16503 .arg(dst)
16504 .arg(&ni);
16505 unsafe {
16506 b.launch(cfg)?;
16507 }
16508 Ok(())
16509 }
16510
16511 pub fn silu_mul_scaled_q8_1(
16519 &self,
16520 gate: &CudaSlice<f32>,
16521 up: &CudaSlice<f32>,
16522 gs: f32,
16523 us: f32,
16524 n: usize,
16525 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
16526 let f = self.func("silu_mul_scaled_q8_1");
16527 let nblk = n / 32;
16528 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);
16532 let (gsf, usf, ni) = (gs, us, n as i32);
16533 let __s_b = self.gpu.stream();
16534 let mut b = __s_b.launch_builder(&f);
16535 b.arg(gate)
16536 .arg(up)
16537 .arg(&gsf)
16538 .arg(&usf)
16539 .arg(&mut aq)
16540 .arg(&mut ad)
16541 .arg(&ni);
16542 unsafe {
16543 b.launch(cfg)?;
16544 }
16545 Ok((aq, ad))
16546 }
16547
16548 pub fn add(
16549 &self,
16550 a: &CudaSlice<f32>,
16551 b_in: &CudaSlice<f32>,
16552 dst: &mut CudaSlice<f32>,
16553 n: usize,
16554 ) -> Result<(), Box<dyn std::error::Error>> {
16555 let f = self.func("add_f32");
16556 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
16558 let ni = n as i32;
16559 let __s_bld = self.gpu.stream();
16560 let mut bld = __s_bld.launch_builder(&f);
16561 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
16562 unsafe {
16563 bld.launch(cfg)?;
16564 }
16565 Ok(())
16566 }
16567
16568 pub fn mul(
16569 &self,
16570 a: &CudaSlice<f32>,
16571 b_in: &CudaSlice<f32>,
16572 dst: &mut CudaSlice<f32>,
16573 n: usize,
16574 ) -> Result<(), Box<dyn std::error::Error>> {
16575 let f = self.func("mul_f32");
16576 let cfg = LaunchConfig::for_num_elems(n as u32);
16577 let ni = n as i32;
16578 let __s_bld = self.gpu.stream();
16579 let mut bld = __s_bld.launch_builder(&f);
16580 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
16581 unsafe {
16582 bld.launch(cfg)?;
16583 }
16584 Ok(())
16585 }
16586
16587 pub fn matmul(
16590 &self,
16591 w: &crate::model::GpuTensor,
16592 x: &CudaSlice<f32>,
16593 m: usize,
16594 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16595 use crate::model::GpuTensor;
16596 let in_f = w.in_features();
16597 let out_f = w.out_features();
16598 #[allow(non_snake_case)]
16606 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
16609 usize::MAX
16610 } else {
16611 16usize
16612 };
16613
16614 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
16639 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
16640 return Ok(y);
16641 }
16642 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
16649 return Ok(y);
16650 }
16651 if let Some(y) = self.try_f16_gemm(w, x, m)? {
16654 return Ok(y);
16655 }
16656 }
16657 if let GpuTensor::Quant { qtype, .. } = w
16672 && *qtype == QT_F8_E4M3_BLK
16673 {
16674 if m >= GEMM_M_THRESHOLD
16675 && let Some(y) = self.try_e4m3_blk_prefill(w, x, m)?
16676 {
16677 return Ok(y);
16678 }
16679 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16680 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
16681 return Ok(y);
16682 }
16683 }
16684 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
16685 return self.qmatvec_mmq(w, x, m);
16686 }
16687 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
16688 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16689 return self.qmatvec_gemm(w, &aq, &ad, m);
16690 }
16691 if m >= GEMM_M_THRESHOLD
16694 && let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)?
16695 {
16696 return Ok(y);
16697 }
16698 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
16702 if m == 1
16707 && fast
16708 && let GpuTensor::Quant {
16709 bytes,
16710 qtype,
16711 row_bytes,
16712 rp,
16713 rp4,
16714 scale,
16715 ..
16716 } = w
16717 && self.mmvq_supports(*qtype)
16718 {
16719 let (bytes, rp) = match rp4 {
16723 Some(m4) => (m4, true),
16724 None => (bytes, *rp),
16725 };
16726 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16727 return self.qmatvec_mmvq(
16728 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
16729 );
16730 }
16731 if (2..=16).contains(&m)
16747 && fast
16748 && std::env::var("MEMRA_NO_BATCHED").is_err()
16749 && (m <= 4 || Self::b8_enabled())
16750 {
16751 let m_ok = m <= 8
16761 || matches!(w, GpuTensor::Quant { qtype, .. }
16762 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
16763 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
16764 if m_ok
16765 && let GpuTensor::Quant {
16766 bytes,
16767 qtype,
16768 row_bytes,
16769 rp,
16770 rp4,
16771 ..
16772 } = w
16773 && self.batched_supports(*qtype)
16774 && self.mmvq_supports(*qtype)
16775 {
16776 let (bytes, rp) = match rp4 {
16777 Some(m4) => (m4, true),
16778 None => (bytes, *rp),
16779 };
16780 let mcols = Self::batched_mcols(m);
16781 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16782 let mut y = self.qmatvec_mmvq_batched(
16783 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
16784 )?;
16785 if let GpuTensor::Quant { scale, .. } = w
16786 && *scale != 1.0
16787 {
16788 self.scale_inplace(&mut y, *scale, m * out_f)?;
16789 }
16790 return Ok(y);
16791 }
16792 }
16793 if fast
16799 && let GpuTensor::Quant {
16800 bytes,
16801 qtype,
16802 row_bytes,
16803 scale,
16804 ..
16805 } = w
16806 && *qtype == QT_F8_E4M3
16807 {
16808 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16809 return self.qmatvec_mmvq(
16810 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
16811 );
16812 }
16813 let mut y = match w {
16814 GpuTensor::Quant {
16815 bytes,
16816 qtype,
16817 row_bytes,
16818 ..
16819 } if fast && *qtype == QT_Q8_0 => {
16820 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16821 }
16822 GpuTensor::Quant {
16823 bytes,
16824 qtype,
16825 row_bytes,
16826 ..
16827 } if fast && *qtype == QT_Q4_K => {
16828 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16829 }
16830 GpuTensor::Quant {
16831 bytes,
16832 qtype,
16833 row_bytes,
16834 ..
16835 } if fast && *qtype == QT_Q6_K => {
16836 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16837 }
16838 GpuTensor::Quant {
16839 bytes,
16840 qtype,
16841 row_bytes,
16842 ..
16843 } if fast && *qtype == QT_Q5_K => {
16844 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16845 }
16846 GpuTensor::Quant {
16847 bytes,
16848 qtype,
16849 row_bytes,
16850 ..
16851 } if fast && *qtype == QT_Q3_K => {
16852 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16853 }
16854 GpuTensor::Quant {
16855 bytes,
16856 qtype,
16857 row_bytes,
16858 rp,
16859 ..
16860 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
16861 if *rp {
16862 "qmatvec_nvfp4_dp4a_rp"
16863 } else {
16864 "qmatvec_nvfp4_dp4a"
16865 },
16866 &bytes.slice(0..bytes.len()),
16867 x,
16868 m,
16869 in_f,
16870 out_f,
16871 *row_bytes,
16872 )?,
16873 GpuTensor::Quant {
16877 bytes,
16878 qtype,
16879 row_bytes,
16880 ..
16881 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
16882 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
16883 }
16884 GpuTensor::Quant {
16889 bytes,
16890 qtype,
16891 row_bytes,
16892 rp,
16893 ..
16894 } =>
16895 {
16898 self.qmatvec(
16899 bytes,
16900 x,
16901 m,
16902 in_f,
16903 out_f,
16904 if *rp && *qtype == QT_NVFP4 {
16905 QT_NVFP4_RP
16906 } else {
16907 *qtype
16908 },
16909 *row_bytes,
16910 )?
16911 }
16912 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
16913 GpuTensor::FloatBf16 { data, .. } => {
16916 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f.is_multiple_of(8) {
16923 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16924 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
16925 y
16926 } else {
16927 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
16928 }
16929 }
16930 };
16931 if let GpuTensor::Quant { scale, .. } = w
16933 && *scale != 1.0
16934 {
16935 self.scale_inplace(&mut y, *scale, m * out_f)?;
16936 }
16937 Ok(y)
16938 }
16939
16940 pub fn stage_a_raw_needed() -> bool {
16950 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16951 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
16952 }
16953
16954 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
16957 use crate::model::GpuTensor;
16958 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
16959 return false;
16960 }
16961 match w {
16962 GpuTensor::Quant { qtype, .. } => {
16969 matches!(
16970 *qtype,
16971 QT_Q8_0
16972 | QT_Q4_K
16973 | QT_Q6_K
16974 | QT_Q5_K
16975 | QT_Q3_K
16976 | QT_NVFP4
16977 | QT_F8_E4M3
16978 | QT_F8_E4M3_BLK
16979 | QT_Q4_0
16980 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
16981 }
16982 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16983 }
16984 }
16985
16986 pub fn matmul_pre(
16991 &self,
16992 w: &crate::model::GpuTensor,
16993 aq: &CudaSlice<i8>,
16994 ad: &CudaSlice<f32>,
16995 x_fallback: &CudaSlice<f32>,
16996 m: usize,
16997 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16998 use crate::model::GpuTensor;
16999 let x_raw_ok = x_fallback.len() >= m * w.in_features();
17005 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
17008 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
17009 return Ok(y);
17010 }
17011 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
17014 return Ok(y);
17015 }
17016 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
17018 return Ok(y);
17019 }
17020 }
17021 if m >= 16
17027 && x_raw_ok
17028 && !self.verify_exact_on()
17029 && let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)?
17030 {
17031 return Ok(y);
17032 }
17033 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
17034 return Ok(y);
17035 }
17036 if m >= 16
17041 && w.out_features() >= 128
17042 && self.mmq_supports(w)
17043 && !self.verify_exact_on()
17044 && x_raw_ok
17045 {
17046 return self.qmatvec_mmq(w, x_fallback, m);
17047 }
17048 if m >= 16
17051 && x_raw_ok
17052 && !self.verify_exact_on()
17053 && let Some(y) =
17054 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
17055 {
17056 return Ok(y);
17057 }
17058 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
17061 return self.qmatvec_gemm(w, aq, ad, m);
17062 }
17063 if !self.uses_q8_1_fast(w) {
17082 if !x_raw_ok {
17083 return Err(format!(
17084 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
17085 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
17086 dropped the f32, so there is nothing to fall back to — pass the real f32 \
17087 activation (see Engine::rms_norm_decode, which is bit-identical to \
17088 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
17089 x_fallback.len(),
17090 m,
17091 w.in_features(),
17092 m * w.in_features()
17093 )
17094 .into());
17095 }
17096 return self.matmul(w, x_fallback, m);
17097 }
17098 let in_f = w.in_features();
17099 let out_f = w.out_features();
17100 let (bytes, qtype, row_bytes, scale, rp) = match w {
17101 GpuTensor::Quant {
17102 bytes,
17103 qtype,
17104 row_bytes,
17105 scale,
17106 rp,
17107 ..
17108 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17109 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
17110 };
17111 let (mbytes, mrp) = match w {
17114 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
17115 _ => (bytes, rp),
17116 };
17117 if m == 1 && self.mmvq_supports(qtype) {
17121 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
17122 }
17123 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
17136 && std::env::var("MEMRA_NO_BATCHED").is_err()
17137 && (m <= 4 || Self::b8_enabled())
17138 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
17142 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
17143 {
17144 let mcols = Self::batched_mcols(m);
17145 return self.qmatvec_mmvq_batched(
17146 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
17147 );
17148 }
17149 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
17155 let (b2, r2) = if qtype == QT_Q4_0 {
17156 (mbytes, mrp)
17157 } else {
17158 (bytes, rp)
17159 };
17160 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
17161 }
17162 let name = match qtype {
17163 QT_Q8_0 => "qmatvec_q8_0_dp4a",
17164 QT_Q4_K => "qmatvec_q4_K_dp4a",
17165 QT_Q6_K => "qmatvec_q6_K_dp4a",
17166 QT_Q5_K => "qmatvec_q5_K_dp4a",
17167 QT_Q3_K => "qmatvec_q3_K_dp4a",
17168 QT_NVFP4 => {
17169 if rp {
17170 "qmatvec_nvfp4_dp4a_rp"
17171 } else {
17172 "qmatvec_nvfp4_dp4a"
17173 }
17174 }
17175 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
17176 _ => unreachable!(),
17177 };
17178 let f = self.func(name);
17179 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
17181 grid_dim: (out_f as u32, m as u32, 1),
17182 block_dim: (128, 1, 1),
17183 shared_mem_bytes: 0,
17184 };
17185 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17186 let __s_b = self.gpu.stream();
17187 let mut b = __s_b.launch_builder(&f);
17188 b.arg(bytes)
17189 .arg(aq)
17190 .arg(ad)
17191 .arg(&mut y)
17192 .arg(&inf)
17193 .arg(&outf)
17194 .arg(&mi)
17195 .arg(&rb);
17196 unsafe {
17197 b.launch(cfg)?;
17198 }
17199 if scale != 1.0 {
17200 self.scale_inplace(&mut y, scale, m * out_f)?;
17201 }
17202 Ok(y)
17203 }
17204
17205 pub fn matmul_decode_exact(
17213 &self,
17214 w: &crate::model::GpuTensor,
17215 x: &CudaSlice<f32>,
17216 m: usize,
17217 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17218 use crate::model::GpuTensor;
17219 if let GpuTensor::Float { data, .. } = w {
17227 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
17228 }
17229 if let GpuTensor::FloatBf16 { data, .. } = w {
17232 let (in_f, out_f) = (w.in_features(), w.out_features());
17233 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
17236 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17237 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
17238 return Ok(y);
17239 }
17240 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
17241 }
17242 if !self.uses_q8_1_fast(w) {
17243 return self.matmul(w, x, m);
17244 }
17245 let in_f = w.in_features();
17246 let out_f = w.out_features();
17247 let (bytes, qtype, row_bytes, scale, rp) = match w {
17248 GpuTensor::Quant {
17249 bytes,
17250 qtype,
17251 row_bytes,
17252 scale,
17253 rp,
17254 ..
17255 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17256 _ => return self.matmul(w, x, m),
17257 };
17258 let (bytes, rp) = match w {
17261 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
17262 _ => (bytes, rp),
17263 };
17264 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17265 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
17269 return Ok(y);
17270 }
17271 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
17280 && std::env::var("MEMRA_NO_BATCHED").is_err()
17281 && (m <= 4 || Self::b8_enabled())
17282 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
17285 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
17286 {
17287 let mcols = Self::batched_mcols(m);
17288 return self.qmatvec_mmvq_batched(
17289 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
17290 );
17291 }
17292 if self.mmvq_supports(qtype) {
17293 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
17296 }
17297 self.matmul_pre(w, &aq, &ad, x, m)
17300 }
17301
17302 pub fn matmul_decode_exact_pre(
17312 &self,
17313 w: &crate::model::GpuTensor,
17314 aq: &CudaSlice<i8>,
17315 ad: &CudaSlice<f32>,
17316 m: usize,
17317 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17318 use crate::model::GpuTensor;
17319 debug_assert!(
17320 self.uses_q8_1_fast(w),
17321 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
17322 );
17323 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
17325 return Ok(y);
17326 }
17327 let in_f = w.in_features();
17328 let out_f = w.out_features();
17329 let (bytes, qtype, row_bytes, scale, rp) = match w {
17330 GpuTensor::Quant {
17331 bytes,
17332 qtype,
17333 row_bytes,
17334 scale,
17335 rp,
17336 ..
17337 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17338 _ => {
17339 return Err(
17340 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
17341 );
17342 }
17343 };
17344 let (bytes, rp) = match w {
17346 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
17347 _ => (bytes, rp),
17348 };
17349 if (2..=16).contains(&m)
17351 && self.batched_supports(qtype)
17352 && self.mmvq_supports(qtype)
17353 && std::env::var("MEMRA_NO_BATCHED").is_err()
17354 && (m <= 4 || Self::b8_enabled())
17355 && (m <= 8
17356 || qtype == QT_Q4_0
17357 || qtype == QT_Q6_K
17358 || qtype == QT_F8_E4M3
17359 || qtype == QT_NVFP4
17360 || qtype == QT_Q4_K
17361 || qtype == QT_Q5_K
17362 || qtype == QT_Q8_0)
17363 {
17364 let mcols = Self::batched_mcols(m);
17365 return self.qmatvec_mmvq_batched(
17366 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
17367 );
17368 }
17369 if self.mmvq_supports(qtype) {
17370 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
17371 }
17372 let x0 = self.zeros(0)?;
17375 self.matmul_pre(w, aq, ad, &x0, m)
17376 }
17377
17378 #[allow(clippy::type_complexity)] pub fn matmul_decode_exact_dual_pre(
17388 &self,
17389 w0: &crate::model::GpuTensor,
17390 w1: &crate::model::GpuTensor,
17391 aq: &CudaSlice<i8>,
17392 ad: &CudaSlice<f32>,
17393 m: usize,
17394 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
17395 {
17396 use crate::model::GpuTensor;
17397 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17398 let on = *ON.get_or_init(|| {
17399 std::env::var("MEMRA_SPEC_DUAL_T")
17400 .map(|v| v != "0")
17401 .unwrap_or(true)
17402 });
17403 if !on
17404 || !(2..=7).contains(&m)
17405 || std::env::var("MEMRA_NO_BATCHED").is_ok()
17406 || !self.uses_q8_1_fast(w0)
17407 || !self.uses_q8_1_fast(w1)
17408 {
17409 return Ok(None);
17410 }
17411 if !self.mmvq_supports(QT_NVFP4) {
17416 return Ok(None);
17417 }
17418 let (in_f, out_f) = (w0.in_features(), w0.out_features());
17419 if w1.in_features() != in_f || w1.out_features() != out_f {
17420 return Ok(None);
17421 }
17422 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
17423 (
17424 GpuTensor::Quant {
17425 bytes: b0,
17426 qtype: q0,
17427 row_bytes: rb0,
17428 scale: s0,
17429 rp: rp0,
17430 rp4: None,
17431 ..
17432 },
17433 GpuTensor::Quant {
17434 bytes: b1,
17435 qtype: q1,
17436 row_bytes: rb1,
17437 scale: s1,
17438 rp: rp1,
17439 rp4: None,
17440 ..
17441 },
17442 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
17443 (b0, b1, *rb0, *s0, *s1, *rp0)
17444 }
17445 _ => return Ok(None),
17446 };
17447 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
17450 {
17451 return Ok(None);
17452 }
17453 let (y0, y1) =
17454 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
17455 Ok(Some(((y0, s0), (y1, s1))))
17456 }
17457
17458 pub fn matmul_decode_exact_group4_pre(
17470 &self,
17471 ws: [&crate::model::GpuTensor; 4],
17472 aq: &CudaSlice<i8>,
17473 ad: &CudaSlice<f32>,
17474 m: usize,
17475 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17476 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17477 let on = *ON.get_or_init(|| {
17478 std::env::var("MEMRA_TK_GDN_GROUP")
17479 .map(|v| v != "0")
17480 .unwrap_or(true)
17481 });
17482 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
17483 }
17484
17485 pub fn matmul_decode_exact_group3_pre(
17490 &self,
17491 ws: [&crate::model::GpuTensor; 3],
17492 aq: &CudaSlice<i8>,
17493 ad: &CudaSlice<f32>,
17494 m: usize,
17495 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17496 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17497 let on = *ON.get_or_init(|| {
17498 std::env::var("MEMRA_TK_FA_GROUP")
17499 .map(|v| v != "0")
17500 .unwrap_or(true)
17501 });
17502 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
17503 }
17504
17505 #[allow(clippy::manual_div_ceil)] fn matmul_decode_exact_group_pre(
17510 &self,
17511 ws: &[&crate::model::GpuTensor],
17512 aq: &CudaSlice<i8>,
17513 ad: &CudaSlice<f32>,
17514 m: usize,
17515 on: bool,
17516 tag: &'static str,
17517 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17518 use crate::model::GpuTensor;
17519 if !on
17520 || !(2..=16).contains(&m)
17521 || std::env::var("MEMRA_NO_BATCHED").is_ok()
17522 || (m > 4 && !Self::b8_enabled())
17523 || !self.mmvq_supports(QT_NVFP4)
17524 || !self.batched_supports(QT_NVFP4)
17525 {
17526 return Ok(None);
17527 }
17528 let in_f = ws[0].in_features();
17529 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
17530 for w in ws {
17531 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
17532 return Ok(None);
17533 }
17534 match w {
17535 GpuTensor::Quant {
17536 bytes,
17537 qtype,
17538 scale,
17539 rp: true,
17540 rp4: None,
17541 ..
17542 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
17543 parts.push((bytes, w.out_features(), *scale));
17544 }
17545 _ => return Ok(None),
17546 }
17547 }
17548 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17550 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
17551 let mcols = if (5..=7).contains(&m) && b567 {
17552 m
17553 } else {
17554 Self::batched_mcols(m)
17555 };
17556 let kname: &'static str = match mcols {
17557 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
17558 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
17559 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
17560 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
17561 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
17562 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
17563 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
17564 _ => return Ok(None),
17565 };
17566 if std::env::var("MEMRA_DEBUG").is_ok() {
17569 use std::sync::Mutex;
17570 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
17571 let mut seen = SEEN.lock().unwrap();
17572 if !seen.contains(&tag) {
17573 seen.push(tag);
17574 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
17575 }
17576 }
17577 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
17580 let three = parts.len() == 3;
17581 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
17582 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
17583 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
17584 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
17587 let cfg = LaunchConfig {
17588 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
17589 block_dim: (32, ROWS_PER_BLOCK, 1),
17590 shared_mem_bytes: 0,
17591 };
17592 let (inf, mi) = (in_f as i32, m as i32);
17593 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
17594 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
17595 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
17596 let s3 = if three { 1.0f32 } else { parts[3].2 };
17597 let w3 = if three { parts[0].0 } else { parts[3].0 };
17598 let f = self.func(kname);
17599 let __s_b = self.gpu.stream();
17600 let mut b = __s_b.launch_builder(&f);
17601 b.arg(parts[0].0)
17602 .arg(parts[1].0)
17603 .arg(parts[2].0)
17604 .arg(w3)
17605 .arg(aq)
17606 .arg(ad)
17607 .arg(&mut y0)
17608 .arg(&mut y1)
17609 .arg(&mut y2)
17610 .arg(&mut y3)
17611 .arg(&inf)
17612 .arg(&n0)
17613 .arg(&n1)
17614 .arg(&n2)
17615 .arg(&n3)
17616 .arg(&mi)
17617 .arg(&s0)
17618 .arg(&s1)
17619 .arg(&s2)
17620 .arg(&s3);
17621 unsafe {
17622 b.launch(cfg)?;
17623 }
17624 Ok(Some(if three {
17625 vec![y0, y1, y2]
17626 } else {
17627 vec![y0, y1, y2, y3]
17628 }))
17629 }
17630
17631 #[allow(clippy::type_complexity)] pub fn matmul_decode_exact_dual(
17648 &self,
17649 w0: &crate::model::GpuTensor,
17650 w1: &crate::model::GpuTensor,
17651 x: &CudaSlice<f32>,
17652 m: usize,
17653 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
17654 use crate::model::GpuTensor;
17655 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17656 let on = *ON.get_or_init(|| {
17657 std::env::var("MEMRA_SPEC_DUAL_T")
17658 .map(|v| v != "0")
17659 .unwrap_or(true)
17660 });
17661 if !on
17662 || !(2..=4).contains(&m)
17663 || std::env::var("MEMRA_NO_BATCHED").is_ok()
17664 || !self.uses_q8_1_fast(w0)
17665 || !self.uses_q8_1_fast(w1)
17666 {
17667 return Ok(None);
17668 }
17669 if !self.mmvq_supports(QT_NVFP4) {
17674 return Ok(None);
17675 }
17676 let (in_f, out_f) = (w0.in_features(), w0.out_features());
17677 if w1.in_features() != in_f || w1.out_features() != out_f {
17678 return Ok(None);
17679 }
17680 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
17681 (
17682 GpuTensor::Quant {
17683 bytes: b0,
17684 qtype: q0,
17685 row_bytes: rb0,
17686 scale: s0,
17687 rp: rp0,
17688 rp4: None,
17689 ..
17690 },
17691 GpuTensor::Quant {
17692 bytes: b1,
17693 qtype: q1,
17694 row_bytes: rb1,
17695 scale: s1,
17696 rp: rp1,
17697 rp4: None,
17698 ..
17699 },
17700 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
17701 (b0, b1, *rb0, *s0, *s1, *rp0)
17702 }
17703 _ => return Ok(None),
17704 };
17705 if std::env::var("MEMRA_DEBUG").is_ok() {
17708 static ONCE: std::sync::Once = std::sync::Once::new();
17709 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
17710 }
17711 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17712 let (y0, y1) =
17713 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
17714 let mut y0 = y0;
17715 let mut y1 = y1;
17716 if s0 != 1.0 {
17717 self.scale_inplace(&mut y0, s0, m * out_f)?;
17718 }
17719 if s1 != 1.0 {
17720 self.scale_inplace(&mut y1, s1, m * out_f)?;
17721 }
17722 Ok(Some((y0, y1)))
17723 }
17724
17725 #[allow(clippy::too_many_arguments)]
17730 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_batched_dual_raw(
17732 &self,
17733 b0: &CudaSlice<u8>,
17734 b1: &CudaSlice<u8>,
17735 aq: &CudaSlice<i8>,
17736 ad: &CudaSlice<f32>,
17737 m: usize,
17738 in_f: usize,
17739 out_f: usize,
17740 row_bytes: usize,
17741 rp: bool,
17742 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17743 const ROWS_PER_BLOCK: u32 = 4;
17744 let mcols = Self::batched_mcols(m);
17745 let tiny_rp1 = rp
17748 && mcols == 4
17749 && out_f <= 128
17750 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
17751 let (name, rows_per_block) = if tiny_rp1 {
17752 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
17753 } else {
17754 match (mcols, rp, m) {
17755 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
17756 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
17757 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
17758 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
17759 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
17760 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
17761 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
17762 _ => {
17763 return Err(
17764 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
17765 );
17766 }
17767 }
17768 };
17769 let f = self.func(name);
17770 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
17771 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
17772 let cfg = LaunchConfig {
17773 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
17774 block_dim: (32, ROWS_PER_BLOCK, 1),
17775 shared_mem_bytes: 0,
17776 };
17777 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17778 let __s_b = self.gpu.stream();
17779 let mut b = __s_b.launch_builder(&f);
17780 b.arg(b0)
17781 .arg(b1)
17782 .arg(aq)
17783 .arg(ad)
17784 .arg(&mut y0)
17785 .arg(&mut y1)
17786 .arg(&inf)
17787 .arg(&outf)
17788 .arg(&mi)
17789 .arg(&rb);
17790 unsafe {
17791 b.launch(cfg)?;
17792 }
17793 Ok((y0, y1))
17794 }
17795
17796 #[allow(clippy::type_complexity)] #[allow(clippy::manual_div_ceil)] pub fn matmul_pre_dual_noscale(
17810 &self,
17811 w0: &crate::model::GpuTensor,
17812 w1: &crate::model::GpuTensor,
17813 aq: &CudaSlice<i8>,
17814 ad: &CudaSlice<f32>,
17815 m: usize,
17816 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
17817 {
17818 use crate::model::GpuTensor;
17819 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
17820 return Ok(None);
17821 }
17822 if !self.mmvq_supports(QT_NVFP4) {
17832 return Ok(None);
17833 }
17834 let (in_f, out_f) = (w0.in_features(), w0.out_features());
17835 if w1.in_features() != in_f || w1.out_features() != out_f {
17836 return Ok(None);
17837 }
17838 let no_mirror =
17851 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
17852 if self.q8_ffn_fuse2_on()
17853 && no_mirror(w0)
17854 && no_mirror(w1)
17855 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
17856 {
17857 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
17858 return Ok(Some(((y0, 1.0), (y1, 1.0))));
17859 }
17860 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
17870 let (y0, y1) =
17871 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
17872 return Ok(Some(((y0, p0.3), (y1, p1.3))));
17873 }
17874 let (b0, q0, rb0, s0, rp0) = match w0 {
17875 GpuTensor::Quant {
17876 bytes,
17877 qtype,
17878 row_bytes,
17879 scale,
17880 rp,
17881 ..
17882 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17883 _ => return Ok(None),
17884 };
17885 let (b1, q1, rb1, s1, rp1) = match w1 {
17886 GpuTensor::Quant {
17887 bytes,
17888 qtype,
17889 row_bytes,
17890 scale,
17891 rp,
17892 ..
17893 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17894 _ => return Ok(None),
17895 };
17896 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
17897 return Ok(None);
17898 }
17899 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
17901 let rows_per_block = ROWS_PER_BLOCK * RPW;
17902 let f = self.func(if rp0 {
17903 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
17904 } else {
17905 "qmatvec_nvfp4_mmvq_dual_mr2"
17906 });
17907 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
17908 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
17909 let cfg = LaunchConfig {
17910 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
17911 block_dim: (32, ROWS_PER_BLOCK, 1),
17912 shared_mem_bytes: 0,
17913 };
17914 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
17915 let one = 1.0f32;
17918 let __s_b = self.gpu.stream();
17919 let mut b = __s_b.launch_builder(&f);
17920 b.arg(b0)
17921 .arg(b1)
17922 .arg(aq)
17923 .arg(ad)
17924 .arg(&mut y0)
17925 .arg(&mut y1)
17926 .arg(&inf)
17927 .arg(&outf)
17928 .arg(&mi)
17929 .arg(&rb)
17930 .arg(&one)
17931 .arg(&one);
17932 unsafe {
17933 b.launch(cfg)?;
17934 }
17935 Ok(Some(((y0, s0), (y1, s1))))
17936 }
17937
17938 #[allow(clippy::too_many_arguments)]
17946 #[allow(clippy::type_complexity)] pub fn matmul_nvfp4_fused3(
17948 &self,
17949 w0: &crate::model::GpuTensor,
17950 w1: &crate::model::GpuTensor,
17951 w2: &crate::model::GpuTensor,
17952 aq: &CudaSlice<i8>,
17953 ad: &CudaSlice<f32>,
17954 m: usize,
17955 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
17956 {
17957 use crate::model::GpuTensor;
17958 if !self.mmvq_supports(QT_NVFP4)
17965 || !self.uses_q8_1_fast(w0)
17966 || !self.uses_q8_1_fast(w1)
17967 || !self.uses_q8_1_fast(w2)
17968 {
17969 return Ok(None);
17970 }
17971 if (9..=16).contains(&m) {
17974 return Ok(
17975 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
17976 Some(mut ys) => {
17977 let y2 = ys.pop().unwrap();
17978 let y1 = ys.pop().unwrap();
17979 let y0 = ys.pop().unwrap();
17980 Some((y0, y1, y2))
17981 }
17982 None => None,
17983 },
17984 );
17985 }
17986 if !(1..=8).contains(&m) {
17987 return Ok(None);
17988 }
17989 if m > 1 {
17990 let in_f = w0.in_features();
17991 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
17992 || !self.batched_supports(QT_NVFP4)
17993 || std::env::var("MEMRA_NO_BATCHED").is_ok()
17994 || (m > 4 && !Self::b8_enabled())
17995 || !in_f.is_multiple_of(512)
17996 || in_f / 64 > 272
17997 {
17998 return Ok(None);
17999 }
18000 }
18001 let unpack = |w: &crate::model::GpuTensor| match w {
18002 GpuTensor::Quant {
18003 bytes,
18004 qtype,
18005 scale,
18006 rp,
18007 ..
18008 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18009 _ => None,
18010 };
18011 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
18012 return Ok(None);
18013 };
18014 let in_f = w0.in_features();
18015 if w1.in_features() != in_f || w2.in_features() != in_f {
18016 return Ok(None);
18017 }
18018 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
18019 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
18021 let rows_pb = ROWS_PER_BLOCK * RPW;
18022 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18023 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
18024 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
18025 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
18026 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
18027 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
18030 if m > 1 {
18031 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
18033 return Ok(None);
18034 }
18035 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
18036 let cfg = LaunchConfig {
18037 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
18038 block_dim: (32, ROWS_PER_BLOCK, 1),
18039 shared_mem_bytes: 0,
18040 };
18041 let __s_b = self.gpu.stream();
18042 let mut b = __s_b.launch_builder(&f);
18043 b.arg(b0)
18044 .arg(b1)
18045 .arg(b2)
18046 .arg(aq)
18047 .arg(ad)
18048 .arg(&mut y0)
18049 .arg(&mut y1)
18050 .arg(&mut y2)
18051 .arg(&inf)
18052 .arg(&oi0)
18053 .arg(&oi1)
18054 .arg(&oi2)
18055 .arg(&mi);
18056 unsafe {
18057 b.launch(cfg)?;
18058 }
18059 return Ok(Some((y0, y1, y2)));
18060 }
18061 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
18062 let cfg = LaunchConfig {
18063 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
18064 block_dim: (32, ROWS_PER_BLOCK, 1),
18065 shared_mem_bytes: 0,
18066 };
18067 let __s_b = self.gpu.stream();
18068 let mut b = __s_b.launch_builder(&f);
18069 b.arg(b0)
18070 .arg(b1)
18071 .arg(b2)
18072 .arg(aq)
18073 .arg(ad)
18074 .arg(&mut y0)
18075 .arg(&mut y1)
18076 .arg(&mut y2)
18077 .arg(&inf)
18078 .arg(&oi0)
18079 .arg(&oi1)
18080 .arg(&oi2)
18081 .arg(&mi)
18082 .arg(&p0.1)
18083 .arg(&p1.1)
18084 .arg(&p2.1);
18085 unsafe {
18086 b.launch(cfg)?;
18087 }
18088 Ok(Some((y0, y1, y2)))
18089 }
18090
18091 #[allow(clippy::type_complexity)] pub fn matmul_nvfp4_fused2(
18101 &self,
18102 w0: &crate::model::GpuTensor,
18103 w1: &crate::model::GpuTensor,
18104 aq: &CudaSlice<i8>,
18105 ad: &CudaSlice<f32>,
18106 m: usize,
18107 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18108 use crate::model::GpuTensor;
18109 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18110 let off =
18111 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
18112 if off
18115 || m != 1
18116 || !self.mmvq_supports(QT_NVFP4)
18117 || !self.uses_q8_1_fast(w0)
18118 || !self.uses_q8_1_fast(w1)
18119 {
18120 return Ok(None);
18121 }
18122 let unpack = |w: &crate::model::GpuTensor| match w {
18123 GpuTensor::Quant {
18124 bytes,
18125 qtype,
18126 scale,
18127 rp,
18128 ..
18129 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18130 _ => None,
18131 };
18132 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
18133 return Ok(None);
18134 };
18135 let in_f = w0.in_features();
18136 if w1.in_features() != in_f {
18137 return Ok(None);
18138 }
18139 let (o0, o1) = (w0.out_features(), w1.out_features());
18140 const ROWS_PER_BLOCK: u32 = 4; let mut rpw: u32 = 2;
18142 let mut kname = "qmatvec_nvfp4_mmvq_fused2_rp";
18143 if b200_matvec_arm_on() {
18150 let waves_at_rpw2 =
18151 (o0 as u32).div_ceil(ROWS_PER_BLOCK * 2) + (o1 as u32).div_ceil(ROWS_PER_BLOCK * 2);
18152 if waves_at_rpw2 < 2 * self.sm_count() as u32 {
18153 rpw = 1;
18154 kname = "qmatvec_nvfp4_mmvq_fused2_rp_g2";
18155 }
18156 }
18157 let rows_pb = ROWS_PER_BLOCK * rpw;
18158 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18159 let f = self.func(kname);
18160 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
18161 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
18162 let cfg = LaunchConfig {
18163 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
18164 block_dim: (32, ROWS_PER_BLOCK, 1),
18165 shared_mem_bytes: 0,
18166 };
18167 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
18168 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
18171 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
18174 {
18175 use cudarc::driver::{DevicePtr, DevicePtrMut};
18176 let s = &self.gpu.stream();
18177 let (pw0, _g0) = b0.device_ptr(s);
18178 let (pw1, _g1) = b1.device_ptr(s);
18179 let (paq, _g2) = aq.device_ptr(s);
18180 let (pad, _g3) = ad.device_ptr(s);
18181 let (py0, _g4) = y0.device_ptr_mut(s);
18182 let (py1, _g5) = y1.device_ptr_mut(s);
18183 let (s0, s1) = (p0.1, p1.1);
18184 let mut ps = [
18185 &pw0 as *const _ as *mut std::ffi::c_void,
18186 &pw1 as *const _ as *mut _,
18187 &paq as *const _ as *mut _,
18188 &pad as *const _ as *mut _,
18189 &py0 as *const _ as *mut _,
18190 &py1 as *const _ as *mut _,
18191 &inf as *const _ as *mut _,
18192 &oi0 as *const _ as *mut _,
18193 &oi1 as *const _ as *mut _,
18194 &mi as *const _ as *mut _,
18195 &s0 as *const _ as *mut _,
18196 &s1 as *const _ as *mut _,
18197 ];
18198 unsafe {
18199 self.launch_pdl(kname, cfg.grid_dim, cfg.block_dim, &mut ps)?;
18200 }
18201 }
18202 return Ok(Some((y0, y1)));
18203 }
18204 let __s_b = self.gpu.stream();
18205 let mut b = __s_b.launch_builder(&f);
18206 b.arg(b0)
18207 .arg(b1)
18208 .arg(aq)
18209 .arg(ad)
18210 .arg(&mut y0)
18211 .arg(&mut y1)
18212 .arg(&inf)
18213 .arg(&oi0)
18214 .arg(&oi1)
18215 .arg(&mi)
18216 .arg(&p0.1)
18217 .arg(&p1.1);
18218 unsafe {
18219 b.launch(cfg)?;
18220 }
18221 Ok(Some((y0, y1)))
18222 }
18223
18224 pub fn matmul_nvfp4_fused2_into(
18229 &self,
18230 w0: &crate::model::GpuTensor,
18231 w1: &crate::model::GpuTensor,
18232 aq: &CudaSlice<i8>,
18233 ad: &CudaSlice<f32>,
18234 y0: &mut CudaSlice<f32>,
18235 y1: &mut CudaSlice<f32>,
18236 ) -> Result<bool, Box<dyn std::error::Error>> {
18237 use crate::model::GpuTensor;
18238 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18239 let off =
18240 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
18241 if off
18242 || !self.mmvq_supports(QT_NVFP4)
18243 || !self.uses_q8_1_fast(w0)
18244 || !self.uses_q8_1_fast(w1)
18245 {
18246 return Ok(false);
18247 }
18248 let unpack = |w: &crate::model::GpuTensor| match w {
18249 GpuTensor::Quant {
18250 bytes,
18251 qtype,
18252 scale,
18253 rp,
18254 ..
18255 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18256 _ => None,
18257 };
18258 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
18259 return Ok(false);
18260 };
18261 let in_f = w0.in_features();
18262 if w1.in_features() != in_f {
18263 return Ok(false);
18264 }
18265 let (o0, o1) = (w0.out_features(), w1.out_features());
18266 if y0.len() < o0 || y1.len() < o1 {
18267 return Ok(false);
18268 }
18269 const ROWS_PER_BLOCK: u32 = 4; let mut rpw: u32 = 2;
18271 let mut kname = "qmatvec_nvfp4_mmvq_fused2_rp";
18272 if b200_matvec_arm_on() {
18276 let waves_at_rpw2 =
18277 (o0 as u32).div_ceil(ROWS_PER_BLOCK * 2) + (o1 as u32).div_ceil(ROWS_PER_BLOCK * 2);
18278 if waves_at_rpw2 < 2 * self.sm_count() as u32 {
18279 rpw = 1;
18280 kname = "qmatvec_nvfp4_mmvq_fused2_rp_g2";
18281 }
18282 }
18283 let rows_pb = ROWS_PER_BLOCK * rpw;
18284 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18285 let f = self.func(kname);
18286 let cfg = LaunchConfig {
18287 grid_dim: (nb(o0) + nb(o1), 1, 1),
18288 block_dim: (32, ROWS_PER_BLOCK, 1),
18289 shared_mem_bytes: 0,
18290 };
18291 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
18292 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
18295 let __s_b = self.gpu.stream();
18296 let mut b = __s_b.launch_builder(&f);
18297 b.arg(b0)
18298 .arg(b1)
18299 .arg(aq)
18300 .arg(ad)
18301 .arg(&mut *y0)
18302 .arg(&mut *y1)
18303 .arg(&inf)
18304 .arg(&oi0)
18305 .arg(&oi1)
18306 .arg(&mi)
18307 .arg(&p0.1)
18308 .arg(&p1.1);
18309 unsafe {
18310 b.launch(cfg)?;
18311 }
18312 Ok(true)
18313 }
18314
18315 #[allow(clippy::type_complexity)]
18320 #[allow(clippy::too_many_arguments)] pub fn matmul_nvfp4_fused4(
18322 &self,
18323 w0: &crate::model::GpuTensor,
18324 w1: &crate::model::GpuTensor,
18325 w2: &crate::model::GpuTensor,
18326 w3: &crate::model::GpuTensor,
18327 aq: &CudaSlice<i8>,
18328 ad: &CudaSlice<f32>,
18329 m: usize,
18330 ) -> Result<
18331 Option<(
18332 CudaSlice<f32>,
18333 CudaSlice<f32>,
18334 CudaSlice<f32>,
18335 CudaSlice<f32>,
18336 )>,
18337 Box<dyn std::error::Error>,
18338 > {
18339 use crate::model::GpuTensor;
18340 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
18347 || !self.mmvq_supports(QT_NVFP4)
18348 || !self.uses_q8_1_fast(w0)
18349 || !self.uses_q8_1_fast(w1)
18350 || !self.uses_q8_1_fast(w2)
18351 || !self.uses_q8_1_fast(w3)
18352 {
18353 return Ok(None);
18354 }
18355 if (9..=16).contains(&m) {
18360 return Ok(
18361 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
18362 Some(mut ys) => {
18363 let y3 = ys.pop().unwrap();
18364 let y2 = ys.pop().unwrap();
18365 let y1 = ys.pop().unwrap();
18366 let y0 = ys.pop().unwrap();
18367 Some((y0, y1, y2, y3))
18368 }
18369 None => None,
18370 },
18371 );
18372 }
18373 if !(1..=8).contains(&m) {
18374 return Ok(None);
18375 }
18376 if m > 1 {
18377 let in_f = w0.in_features();
18380 if !self.batched_supports(QT_NVFP4)
18381 || std::env::var("MEMRA_NO_BATCHED").is_ok()
18382 || (m > 4 && !Self::b8_enabled())
18383 || !in_f.is_multiple_of(512)
18384 || in_f / 64 > 272
18385 {
18386 return Ok(None);
18387 }
18388 }
18389 let unpack = |w: &crate::model::GpuTensor| match w {
18390 GpuTensor::Quant {
18391 bytes,
18392 qtype,
18393 scale,
18394 rp,
18395 ..
18396 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
18397 _ => None,
18398 };
18399 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
18400 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
18401 else {
18402 return Ok(None);
18403 };
18404 let in_f = w0.in_features();
18405 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
18406 return Ok(None);
18407 }
18408 let (o0, o1, o2, o3) = (
18409 w0.out_features(),
18410 w1.out_features(),
18411 w2.out_features(),
18412 w3.out_features(),
18413 );
18414 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
18416 let rows_pb = ROWS_PER_BLOCK * RPW;
18417 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
18418 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
18419 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
18420 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
18421 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
18422 let (inf, oi0, oi1, oi2, oi3, mi) = (
18423 in_f as i32,
18424 o0 as i32,
18425 o1 as i32,
18426 o2 as i32,
18427 o3 as i32,
18428 m as i32,
18429 );
18430 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
18433 if m > 1 {
18434 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
18437 return Ok(None);
18438 }
18439 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
18440 let cfg = LaunchConfig {
18441 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
18442 block_dim: (32, ROWS_PER_BLOCK, 1),
18443 shared_mem_bytes: 0,
18444 };
18445 let __s_b = self.gpu.stream();
18446 let mut b = __s_b.launch_builder(&f);
18447 b.arg(b0)
18448 .arg(b1)
18449 .arg(b2)
18450 .arg(b3)
18451 .arg(aq)
18452 .arg(ad)
18453 .arg(&mut y0)
18454 .arg(&mut y1)
18455 .arg(&mut y2)
18456 .arg(&mut y3)
18457 .arg(&inf)
18458 .arg(&oi0)
18459 .arg(&oi1)
18460 .arg(&oi2)
18461 .arg(&oi3)
18462 .arg(&mi);
18463 unsafe {
18464 b.launch(cfg)?;
18465 }
18466 return Ok(Some((y0, y1, y2, y3)));
18467 }
18468 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
18469 let cfg = LaunchConfig {
18470 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
18471 block_dim: (32, ROWS_PER_BLOCK, 1),
18472 shared_mem_bytes: 0,
18473 };
18474 let __s_b = self.gpu.stream();
18475 let mut b = __s_b.launch_builder(&f);
18476 b.arg(b0)
18477 .arg(b1)
18478 .arg(b2)
18479 .arg(b3)
18480 .arg(aq)
18481 .arg(ad)
18482 .arg(&mut y0)
18483 .arg(&mut y1)
18484 .arg(&mut y2)
18485 .arg(&mut y3)
18486 .arg(&inf)
18487 .arg(&oi0)
18488 .arg(&oi1)
18489 .arg(&oi2)
18490 .arg(&oi3)
18491 .arg(&mi)
18492 .arg(&p0.1)
18493 .arg(&p1.1)
18494 .arg(&p2.1)
18495 .arg(&p3.1);
18496 unsafe {
18497 b.launch(cfg)?;
18498 }
18499 Ok(Some((y0, y1, y2, y3)))
18500 }
18501
18502 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused2(
18511 &self,
18512 w0: &crate::model::GpuTensor,
18513 w1: &crate::model::GpuTensor,
18514 aq: &CudaSlice<i8>,
18515 ad: &CudaSlice<f32>,
18516 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18517 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
18523 return Ok(Some(self.e4m3_fused2_core(
18524 p0.0,
18525 p1.0,
18526 aq,
18527 ad,
18528 w0.in_features(),
18529 p0.1,
18530 p1.1,
18531 p0.2,
18532 p0.3,
18533 p1.3,
18534 )?));
18535 }
18536 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
18537 return Ok(None);
18538 };
18539 Ok(Some(self.q8_fused2_core(
18540 p0.0,
18541 p1.0,
18542 aq,
18543 ad,
18544 w0.in_features(),
18545 p0.1,
18546 p1.1,
18547 p0.2,
18548 )?))
18549 }
18550
18551 #[allow(clippy::too_many_arguments)]
18552 fn q8_fused2_core(
18553 &self,
18554 b0: &CudaSlice<u8>,
18555 b1: &CudaSlice<u8>,
18556 aq: &CudaSlice<i8>,
18557 ad: &CudaSlice<f32>,
18558 in_f: usize,
18559 out0: usize,
18560 out1: usize,
18561 row_bytes: usize,
18562 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18563 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
18565 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
18566 let f = self.func("qmatvec_q8_0_mmvq_fused2");
18567 let mut y0 = self.alloc_uninit::<f32>(out0)?;
18568 let mut y1 = self.alloc_uninit::<f32>(out1)?;
18569 let cfg = LaunchConfig {
18570 grid_dim: (nb0 + nb1, 1, 1),
18571 block_dim: (32, ROWS_PER_BLOCK, 1),
18572 shared_mem_bytes: 0,
18573 };
18574 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
18575 let __s_b = self.gpu.stream();
18576 let mut b = __s_b.launch_builder(&f);
18577 b.arg(b0)
18578 .arg(b1)
18579 .arg(aq)
18580 .arg(ad)
18581 .arg(&mut y0)
18582 .arg(&mut y1)
18583 .arg(&inf)
18584 .arg(&o0)
18585 .arg(&o1)
18586 .arg(&rbl);
18587 unsafe {
18588 b.launch(cfg)?;
18589 }
18590 Ok((y0, y1))
18591 }
18592
18593 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused2_x(
18600 &self,
18601 w0: &crate::model::GpuTensor,
18602 w1: &crate::model::GpuTensor,
18603 x: &CudaSlice<f32>,
18604 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18605 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
18606 return Ok(None);
18607 }
18608 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
18609 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
18610 return Ok(Some(self.e4m3_fused2_core(
18611 p0.0,
18612 p1.0,
18613 &aq,
18614 &ad,
18615 w0.in_features(),
18616 p0.1,
18617 p1.1,
18618 p0.2,
18619 p0.3,
18620 p1.3,
18621 )?));
18622 }
18623 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
18624 return Ok(None);
18625 };
18626 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
18627 Ok(Some(self.q8_fused2_core(
18628 p0.0,
18629 p1.0,
18630 &aq,
18631 &ad,
18632 w0.in_features(),
18633 p0.1,
18634 p1.1,
18635 p0.2,
18636 )?))
18637 }
18638
18639 #[allow(clippy::too_many_arguments)]
18642 pub fn qmatvec_q8_fused2_raw(
18643 &self,
18644 b0: &CudaSlice<u8>,
18645 b1: &CudaSlice<u8>,
18646 x: &CudaSlice<f32>,
18647 in_f: usize,
18648 out0: usize,
18649 out1: usize,
18650 row_bytes: usize,
18651 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18652 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
18653 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
18654 }
18655
18656 #[allow(clippy::type_complexity)] pub fn matmul_q4_fused3(
18663 &self,
18664 w0: &crate::model::GpuTensor,
18665 w1: &crate::model::GpuTensor,
18666 w2: &crate::model::GpuTensor,
18667 aq: &CudaSlice<i8>,
18668 ad: &CudaSlice<f32>,
18669 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
18670 {
18671 use crate::model::GpuTensor;
18672 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18673 match w {
18674 GpuTensor::Quant {
18675 qtype, row_bytes, ..
18676 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18677 _ => None,
18678 }
18679 };
18680 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
18681 return Ok(None);
18682 };
18683 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
18684 return Ok(None);
18685 }
18686 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18690 match w {
18691 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18692 Some(m) => (m, true),
18693 None => (bytes, *rp),
18694 },
18695 _ => unreachable!(),
18696 }
18697 }
18698 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
18699 if rp0 != rp1 || rp1 != rp2 {
18700 return Ok(None);
18701 }
18702 let rp = rp0;
18703 let rpb: u32 = 4;
18704 let mr1 = rp && Self::q40_mr1_on();
18708 let nb = |o: usize| {
18709 if mr1 {
18710 (o as u32).div_ceil(rpb)
18711 } else {
18712 (o as u32).div_ceil(2).div_ceil(rpb)
18713 }
18714 };
18715 let grid = nb(o0) + nb(o1) + nb(o2);
18716 let mut y0 = self.alloc_uninit::<f32>(o0)?;
18717 let mut y1 = self.alloc_uninit::<f32>(o1)?;
18718 let mut y2 = self.alloc_uninit::<f32>(o2)?;
18719 let f = self.func(if mr1 {
18720 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
18721 } else if rp {
18722 "qmatvec_q4_0_mmvq_fused3_rp"
18723 } else {
18724 "qmatvec_q4_0_mmvq_fused3"
18725 });
18726 let cfg = LaunchConfig {
18727 grid_dim: (grid, 1, 1),
18728 block_dim: (32, rpb, 1),
18729 shared_mem_bytes: 0,
18730 };
18731 let inf = w0.in_features() as i32;
18732 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
18733 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
18734 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18737 {
18738 use cudarc::driver::{DevicePtr, DevicePtrMut};
18739 let s = &self.gpu.stream();
18740 let (p0, _g0) = b0.device_ptr(s);
18741 let (p1, _g1) = b1.device_ptr(s);
18742 let (p2, _g2) = b2.device_ptr(s);
18743 let (paq, _g3) = aq.device_ptr(s);
18744 let (pad, _g4) = ad.device_ptr(s);
18745 let (py0, _g5) = y0.device_ptr_mut(s);
18746 let (py1, _g6) = y1.device_ptr_mut(s);
18747 let (py2, _g7) = y2.device_ptr_mut(s);
18748 let mut ps = [
18749 &p0 as *const _ as *mut std::ffi::c_void,
18750 &p1 as *const _ as *mut _,
18751 &p2 as *const _ as *mut _,
18752 &paq as *const _ as *mut _,
18753 &pad as *const _ as *mut _,
18754 &py0 as *const _ as *mut _,
18755 &py1 as *const _ as *mut _,
18756 &py2 as *const _ as *mut _,
18757 &inf as *const _ as *mut _,
18758 &oo0 as *const _ as *mut _,
18759 &oo1 as *const _ as *mut _,
18760 &oo2 as *const _ as *mut _,
18761 &r0 as *const _ as *mut _,
18762 &r1 as *const _ as *mut _,
18763 &r2 as *const _ as *mut _,
18764 ];
18765 unsafe {
18766 self.launch_pdl(
18767 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
18768 (grid, 1, 1),
18769 (32, rpb, 1),
18770 &mut ps,
18771 )?;
18772 }
18773 }
18774 return Ok(Some((y0, y1, y2)));
18775 }
18776 let __s_b = self.gpu.stream();
18777 let mut b = __s_b.launch_builder(&f);
18778 b.arg(b0)
18779 .arg(b1)
18780 .arg(b2)
18781 .arg(aq)
18782 .arg(ad)
18783 .arg(&mut y0)
18784 .arg(&mut y1)
18785 .arg(&mut y2)
18786 .arg(&inf)
18787 .arg(&oo0)
18788 .arg(&oo1)
18789 .arg(&oo2)
18790 .arg(&r0)
18791 .arg(&r1)
18792 .arg(&r2);
18793 unsafe {
18794 b.launch(cfg)?;
18795 }
18796 Ok(Some((y0, y1, y2)))
18797 }
18798
18799 #[allow(clippy::too_many_arguments)]
18802 pub fn matmul_q4_fused3_into(
18803 &self,
18804 w0: &crate::model::GpuTensor,
18805 w1: &crate::model::GpuTensor,
18806 w2: &crate::model::GpuTensor,
18807 aq: &CudaSlice<i8>,
18808 ad: &CudaSlice<f32>,
18809 y0: &mut CudaSlice<f32>,
18810 y1: &mut CudaSlice<f32>,
18811 y2: &mut CudaSlice<f32>,
18812 ) -> Result<bool, Box<dyn std::error::Error>> {
18813 use crate::model::GpuTensor;
18814 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18815 match w {
18816 GpuTensor::Quant {
18817 qtype, row_bytes, ..
18818 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18819 _ => None,
18820 }
18821 };
18822 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
18823 return Ok(false);
18824 };
18825 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
18826 return Ok(false);
18827 }
18828 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18829 match w {
18830 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18831 Some(m) => (m, true),
18832 None => (bytes, *rp),
18833 },
18834 _ => unreachable!(),
18835 }
18836 }
18837 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
18838 if rp0 != rp1 || rp1 != rp2 {
18839 return Ok(false);
18840 }
18841 let rp = rp0;
18842 let rpb: u32 = 4;
18843 let mr1 = rp && Self::q40_mr1_on();
18844 let nb = |o: usize| {
18845 if mr1 {
18846 (o as u32).div_ceil(rpb)
18847 } else {
18848 (o as u32).div_ceil(2).div_ceil(rpb)
18849 }
18850 };
18851 let grid = nb(o0) + nb(o1) + nb(o2);
18852 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
18853 let f = self.func(if mr1 {
18854 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
18855 } else if rp {
18856 "qmatvec_q4_0_mmvq_fused3_rp"
18857 } else {
18858 "qmatvec_q4_0_mmvq_fused3"
18859 });
18860 let cfg = LaunchConfig {
18861 grid_dim: (grid, 1, 1),
18862 block_dim: (32, rpb, 1),
18863 shared_mem_bytes: 0,
18864 };
18865 let inf = w0.in_features() as i32;
18866 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
18867 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
18868 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18870 use cudarc::driver::{DevicePtr, DevicePtrMut};
18871 let s = &self.gpu.stream();
18872 let (p0, _g0) = b0.device_ptr(s);
18873 let (p1, _g1) = b1.device_ptr(s);
18874 let (p2, _g2) = b2.device_ptr(s);
18875 let (paq, _g3) = aq.device_ptr(s);
18876 let (pad, _g4) = ad.device_ptr(s);
18877 let (py0, _g5) = y0.device_ptr_mut(s);
18878 let (py1, _g6) = y1.device_ptr_mut(s);
18879 let (py2, _g7) = y2.device_ptr_mut(s);
18880 let mut ps = [
18881 &p0 as *const _ as *mut std::ffi::c_void,
18882 &p1 as *const _ as *mut _,
18883 &p2 as *const _ as *mut _,
18884 &paq as *const _ as *mut _,
18885 &pad as *const _ as *mut _,
18886 &py0 as *const _ as *mut _,
18887 &py1 as *const _ as *mut _,
18888 &py2 as *const _ as *mut _,
18889 &inf as *const _ as *mut _,
18890 &oo0 as *const _ as *mut _,
18891 &oo1 as *const _ as *mut _,
18892 &oo2 as *const _ as *mut _,
18893 &r0 as *const _ as *mut _,
18894 &r1 as *const _ as *mut _,
18895 &r2 as *const _ as *mut _,
18896 ];
18897 unsafe {
18898 self.launch_pdl(
18899 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
18900 (grid, 1, 1),
18901 (32, rpb, 1),
18902 &mut ps,
18903 )?;
18904 }
18905 return Ok(true);
18906 }
18907 let __s_b = self.gpu.stream();
18908 let mut b = __s_b.launch_builder(&f);
18909 b.arg(b0)
18910 .arg(b1)
18911 .arg(b2)
18912 .arg(aq)
18913 .arg(ad)
18914 .arg(&mut *y0)
18915 .arg(&mut *y1)
18916 .arg(&mut *y2)
18917 .arg(&inf)
18918 .arg(&oo0)
18919 .arg(&oo1)
18920 .arg(&oo2)
18921 .arg(&r0)
18922 .arg(&r1)
18923 .arg(&r2);
18924 unsafe {
18925 b.launch(cfg)?;
18926 }
18927 Ok(true)
18928 }
18929
18930 #[allow(clippy::type_complexity)] pub fn matmul_q4_fused2(
18933 &self,
18934 w0: &crate::model::GpuTensor,
18935 w1: &crate::model::GpuTensor,
18936 aq: &CudaSlice<i8>,
18937 ad: &CudaSlice<f32>,
18938 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
18939 use crate::model::GpuTensor;
18940 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
18941 match w {
18942 GpuTensor::Quant {
18943 qtype, row_bytes, ..
18944 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
18945 _ => None,
18946 }
18947 };
18948 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
18949 return Ok(None);
18950 };
18951 if w0.in_features() != w1.in_features() {
18952 return Ok(None);
18953 }
18954 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
18956 match w {
18957 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
18958 Some(m) => (m, true),
18959 None => (bytes, *rp),
18960 },
18961 _ => unreachable!(),
18962 }
18963 }
18964 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
18965 if rp0 != rp1 {
18966 return Ok(None);
18967 }
18968 let rp = rp0;
18969 let rpb: u32 = 4;
18970 let mr1 = rp && Self::q40_mr1_on();
18972 let nb = |o: usize| {
18973 if mr1 {
18974 (o as u32).div_ceil(rpb)
18975 } else {
18976 (o as u32).div_ceil(2).div_ceil(rpb)
18977 }
18978 };
18979 let grid = nb(o0) + nb(o1);
18980 let mut y0 = self.alloc_uninit::<f32>(o0)?;
18981 let mut y1 = self.alloc_uninit::<f32>(o1)?;
18982 let f = self.func(if mr1 {
18983 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
18984 } else if rp {
18985 "qmatvec_q4_0_mmvq_fused2_rp"
18986 } else {
18987 "qmatvec_q4_0_mmvq_fused2"
18988 });
18989 let cfg = LaunchConfig {
18990 grid_dim: (grid, 1, 1),
18991 block_dim: (32, rpb, 1),
18992 shared_mem_bytes: 0,
18993 };
18994 let inf = w0.in_features() as i32;
18995 let (oo0, oo1) = (o0 as i32, o1 as i32);
18996 let (r0, r1) = (rb0 as i64, rb1 as i64);
18997 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
18999 {
19000 use cudarc::driver::{DevicePtr, DevicePtrMut};
19001 let s = &self.gpu.stream();
19002 let (p0, _g0) = b0.device_ptr(s);
19003 let (p1, _g1) = b1.device_ptr(s);
19004 let (paq, _g2) = aq.device_ptr(s);
19005 let (pad, _g3) = ad.device_ptr(s);
19006 let (py0, _g4) = y0.device_ptr_mut(s);
19007 let (py1, _g5) = y1.device_ptr_mut(s);
19008 let mut ps = [
19009 &p0 as *const _ as *mut std::ffi::c_void,
19010 &p1 as *const _ as *mut _,
19011 &paq as *const _ as *mut _,
19012 &pad as *const _ as *mut _,
19013 &py0 as *const _ as *mut _,
19014 &py1 as *const _ as *mut _,
19015 &inf as *const _ as *mut _,
19016 &oo0 as *const _ as *mut _,
19017 &oo1 as *const _ as *mut _,
19018 &r0 as *const _ as *mut _,
19019 &r1 as *const _ as *mut _,
19020 ];
19021 unsafe {
19022 self.launch_pdl(
19023 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
19024 (grid, 1, 1),
19025 (32, rpb, 1),
19026 &mut ps,
19027 )?;
19028 }
19029 }
19030 return Ok(Some((y0, y1)));
19031 }
19032 let __s_b = self.gpu.stream();
19033 let mut b = __s_b.launch_builder(&f);
19034 b.arg(b0)
19035 .arg(b1)
19036 .arg(aq)
19037 .arg(ad)
19038 .arg(&mut y0)
19039 .arg(&mut y1)
19040 .arg(&inf)
19041 .arg(&oo0)
19042 .arg(&oo1)
19043 .arg(&r0)
19044 .arg(&r1);
19045 unsafe {
19046 b.launch(cfg)?;
19047 }
19048 Ok(Some((y0, y1)))
19049 }
19050
19051 pub fn matmul_q4_fused2_into(
19053 &self,
19054 w0: &crate::model::GpuTensor,
19055 w1: &crate::model::GpuTensor,
19056 aq: &CudaSlice<i8>,
19057 ad: &CudaSlice<f32>,
19058 y0: &mut CudaSlice<f32>,
19059 y1: &mut CudaSlice<f32>,
19060 ) -> Result<bool, Box<dyn std::error::Error>> {
19061 use crate::model::GpuTensor;
19062 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
19063 match w {
19064 GpuTensor::Quant {
19065 qtype, row_bytes, ..
19066 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
19067 _ => None,
19068 }
19069 };
19070 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
19071 return Ok(false);
19072 };
19073 if w0.in_features() != w1.in_features() {
19074 return Ok(false);
19075 }
19076 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
19077 match w {
19078 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
19079 Some(m) => (m, true),
19080 None => (bytes, *rp),
19081 },
19082 _ => unreachable!(),
19083 }
19084 }
19085 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
19086 if rp0 != rp1 {
19087 return Ok(false);
19088 }
19089 let rp = rp0;
19090 let rpb: u32 = 4;
19091 let mr1 = rp && Self::q40_mr1_on();
19092 let nb = |o: usize| {
19093 if mr1 {
19094 (o as u32).div_ceil(rpb)
19095 } else {
19096 (o as u32).div_ceil(2).div_ceil(rpb)
19097 }
19098 };
19099 let grid = nb(o0) + nb(o1);
19100 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
19101 let f = self.func(if mr1 {
19102 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
19103 } else if rp {
19104 "qmatvec_q4_0_mmvq_fused2_rp"
19105 } else {
19106 "qmatvec_q4_0_mmvq_fused2"
19107 });
19108 let cfg = LaunchConfig {
19109 grid_dim: (grid, 1, 1),
19110 block_dim: (32, rpb, 1),
19111 shared_mem_bytes: 0,
19112 };
19113 let inf = w0.in_features() as i32;
19114 let (oo0, oo1) = (o0 as i32, o1 as i32);
19115 let (r0, r1) = (rb0 as i64, rb1 as i64);
19116 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
19118 use cudarc::driver::{DevicePtr, DevicePtrMut};
19119 let s = &self.gpu.stream();
19120 let (p0, _g0) = b0.device_ptr(s);
19121 let (p1, _g1) = b1.device_ptr(s);
19122 let (paq, _g2) = aq.device_ptr(s);
19123 let (pad, _g3) = ad.device_ptr(s);
19124 let (py0, _g4) = y0.device_ptr_mut(s);
19125 let (py1, _g5) = y1.device_ptr_mut(s);
19126 let mut ps = [
19127 &p0 as *const _ as *mut std::ffi::c_void,
19128 &p1 as *const _ as *mut _,
19129 &paq as *const _ as *mut _,
19130 &pad as *const _ as *mut _,
19131 &py0 as *const _ as *mut _,
19132 &py1 as *const _ as *mut _,
19133 &inf as *const _ as *mut _,
19134 &oo0 as *const _ as *mut _,
19135 &oo1 as *const _ as *mut _,
19136 &r0 as *const _ as *mut _,
19137 &r1 as *const _ as *mut _,
19138 ];
19139 unsafe {
19140 self.launch_pdl(
19141 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
19142 (grid, 1, 1),
19143 (32, rpb, 1),
19144 &mut ps,
19145 )?;
19146 }
19147 return Ok(true);
19148 }
19149 let __s_b = self.gpu.stream();
19150 let mut b = __s_b.launch_builder(&f);
19151 b.arg(b0)
19152 .arg(b1)
19153 .arg(aq)
19154 .arg(ad)
19155 .arg(&mut *y0)
19156 .arg(&mut *y1)
19157 .arg(&inf)
19158 .arg(&oo0)
19159 .arg(&oo1)
19160 .arg(&r0)
19161 .arg(&r1);
19162 unsafe {
19163 b.launch(cfg)?;
19164 }
19165 Ok(true)
19166 }
19167
19168 #[allow(clippy::type_complexity)] pub fn matmul_q4_fused2_batched(
19174 &self,
19175 w0: &crate::model::GpuTensor,
19176 w1: &crate::model::GpuTensor,
19177 aq: &CudaSlice<i8>,
19178 ad: &CudaSlice<f32>,
19179 m: usize,
19180 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
19181 use crate::model::GpuTensor;
19182 if !(2..=8).contains(&m) {
19183 return Ok(None);
19184 }
19185 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
19186 match w {
19187 GpuTensor::Quant {
19188 qtype, row_bytes, ..
19189 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
19190 _ => None,
19191 }
19192 };
19193 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
19194 return Ok(None);
19195 };
19196 if w0.in_features() != w1.in_features() {
19197 return Ok(None);
19198 }
19199 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
19200 match w {
19201 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
19202 Some(mr) => (mr, true),
19203 None => (bytes, *rp),
19204 },
19205 _ => unreachable!(),
19206 }
19207 }
19208 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
19209 if !rp0 || !rp1 {
19210 return Ok(None);
19211 }
19212 let mcols = Self::batched_mcols(m);
19213 let rpb: u32 = 4;
19214 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
19215 let grid = nb(o0) + nb(o1);
19216 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
19217 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
19218 let f = self.func(match mcols {
19219 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
19220 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
19221 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
19222 });
19223 let cfg = LaunchConfig {
19224 grid_dim: (grid, 1, 1),
19225 block_dim: (32, rpb, 1),
19226 shared_mem_bytes: 0,
19227 };
19228 let inf = w0.in_features() as i32;
19229 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
19230 let rb = rb0 as i64;
19231 let __s_b = self.gpu.stream();
19232 let mut b = __s_b.launch_builder(&f);
19233 b.arg(b0)
19234 .arg(b1)
19235 .arg(aq)
19236 .arg(ad)
19237 .arg(&mut y0)
19238 .arg(&mut y1)
19239 .arg(&inf)
19240 .arg(&oo0)
19241 .arg(&oo1)
19242 .arg(&mi)
19243 .arg(&rb);
19244 unsafe {
19245 b.launch(cfg)?;
19246 }
19247 Ok(Some((y0, y1)))
19248 }
19249
19250 #[allow(clippy::too_many_arguments)]
19253 #[allow(clippy::type_complexity)] pub fn matmul_q4_fused3_batched(
19255 &self,
19256 w0: &crate::model::GpuTensor,
19257 w1: &crate::model::GpuTensor,
19258 w2: &crate::model::GpuTensor,
19259 aq: &CudaSlice<i8>,
19260 ad: &CudaSlice<f32>,
19261 m: usize,
19262 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
19263 {
19264 use crate::model::GpuTensor;
19265 if !(2..=8).contains(&m) {
19266 return Ok(None);
19267 }
19268 let q4 = |w: &GpuTensor| -> Option<usize> {
19269 match w {
19270 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
19271 _ => None,
19272 }
19273 };
19274 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
19275 return Ok(None);
19276 };
19277 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
19278 return Ok(None);
19279 }
19280 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
19281 match w {
19282 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
19283 Some(mr) => (mr, true),
19284 None => (bytes, *rp),
19285 },
19286 _ => unreachable!(),
19287 }
19288 }
19289 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
19290 if !rp0 || !rp1 || !rp2 {
19291 return Ok(None);
19292 }
19293 let mcols = Self::batched_mcols(m);
19294 let rpb: u32 = 4;
19295 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
19296 let grid = nb(o0) + nb(o1) + nb(o2);
19297 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
19298 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
19299 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
19300 let f = self.func(match mcols {
19301 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
19302 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
19303 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
19304 });
19305 let cfg = LaunchConfig {
19306 grid_dim: (grid, 1, 1),
19307 block_dim: (32, rpb, 1),
19308 shared_mem_bytes: 0,
19309 };
19310 let inf = w0.in_features() as i32;
19311 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
19312 let rb = 0i64;
19313 let __s_b = self.gpu.stream();
19314 let mut b = __s_b.launch_builder(&f);
19315 b.arg(b0)
19316 .arg(b1)
19317 .arg(b2)
19318 .arg(aq)
19319 .arg(ad)
19320 .arg(&mut y0)
19321 .arg(&mut y1)
19322 .arg(&mut y2)
19323 .arg(&inf)
19324 .arg(&oo0)
19325 .arg(&oo1)
19326 .arg(&oo2)
19327 .arg(&mi)
19328 .arg(&rb);
19329 unsafe {
19330 b.launch(cfg)?;
19331 }
19332 Ok(Some((y0, y1, y2)))
19333 }
19334
19335 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused3(
19337 &self,
19338 w0: &crate::model::GpuTensor,
19339 w1: &crate::model::GpuTensor,
19340 w2: &crate::model::GpuTensor,
19341 aq: &CudaSlice<i8>,
19342 ad: &CudaSlice<f32>,
19343 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
19344 {
19345 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
19348 return Ok(Some(self.e4m3_fused3_core(
19349 p0.0,
19350 p1.0,
19351 p2.0,
19352 aq,
19353 ad,
19354 w0.in_features(),
19355 p0.1,
19356 p1.1,
19357 p2.1,
19358 p0.2,
19359 p0.3,
19360 p1.3,
19361 p2.3,
19362 )?));
19363 }
19364 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
19365 return Ok(None);
19366 };
19367 Ok(Some(self.q8_fused3_core(
19368 p0.0,
19369 p1.0,
19370 p2.0,
19371 aq,
19372 ad,
19373 w0.in_features(),
19374 p0.1,
19375 p1.1,
19376 p2.1,
19377 p0.2,
19378 )?))
19379 }
19380
19381 #[allow(clippy::too_many_arguments)]
19382 #[allow(clippy::type_complexity)] fn q8_fused3_core(
19384 &self,
19385 b0: &CudaSlice<u8>,
19386 b1: &CudaSlice<u8>,
19387 b2: &CudaSlice<u8>,
19388 aq: &CudaSlice<i8>,
19389 ad: &CudaSlice<f32>,
19390 in_f: usize,
19391 out0: usize,
19392 out1: usize,
19393 out2: usize,
19394 row_bytes: usize,
19395 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19396 const ROWS_PER_BLOCK: u32 = 4;
19397 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19398 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19399 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19400 let f = self.func("qmatvec_q8_0_mmvq_fused3");
19401 let mut y0 = self.alloc_uninit::<f32>(out0)?;
19402 let mut y1 = self.alloc_uninit::<f32>(out1)?;
19403 let mut y2 = self.alloc_uninit::<f32>(out2)?;
19404 let cfg = LaunchConfig {
19405 grid_dim: (nb0 + nb1 + nb2, 1, 1),
19406 block_dim: (32, ROWS_PER_BLOCK, 1),
19407 shared_mem_bytes: 0,
19408 };
19409 let (inf, o0, o1, o2, rbl) = (
19410 in_f as i32,
19411 out0 as i32,
19412 out1 as i32,
19413 out2 as i32,
19414 row_bytes as i64,
19415 );
19416 let __s_b = self.gpu.stream();
19417 let mut b = __s_b.launch_builder(&f);
19418 b.arg(b0)
19419 .arg(b1)
19420 .arg(b2)
19421 .arg(aq)
19422 .arg(ad)
19423 .arg(&mut y0)
19424 .arg(&mut y1)
19425 .arg(&mut y2)
19426 .arg(&inf)
19427 .arg(&o0)
19428 .arg(&o1)
19429 .arg(&o2)
19430 .arg(&rbl);
19431 unsafe {
19432 b.launch(cfg)?;
19433 }
19434 Ok((y0, y1, y2))
19435 }
19436
19437 #[allow(clippy::too_many_arguments)]
19439 #[allow(clippy::type_complexity)] pub fn qmatvec_q8_fused3_raw(
19441 &self,
19442 b0: &CudaSlice<u8>,
19443 b1: &CudaSlice<u8>,
19444 b2: &CudaSlice<u8>,
19445 x: &CudaSlice<f32>,
19446 in_f: usize,
19447 out0: usize,
19448 out1: usize,
19449 out2: usize,
19450 row_bytes: usize,
19451 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19452 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
19453 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
19454 }
19455
19456 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused2_t(
19468 &self,
19469 w0: &crate::model::GpuTensor,
19470 w1: &crate::model::GpuTensor,
19471 aq: &CudaSlice<i8>,
19472 ad: &CudaSlice<f32>,
19473 m: usize,
19474 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
19475 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
19479 return Ok(None);
19480 }
19481 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
19484 if m > 4 && !Self::b8_enabled() {
19485 return Ok(None);
19486 }
19487 return Ok(Some(self.e4m3_fused2_t_core(
19488 p0.0,
19489 p1.0,
19490 aq,
19491 ad,
19492 m,
19493 w0.in_features(),
19494 p0.1,
19495 p1.1,
19496 p0.2,
19497 p0.3,
19498 p1.3,
19499 )?));
19500 }
19501 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
19502 return Ok(None);
19503 };
19504 Ok(Some(self.q8_fused2_t_core(
19505 p0.0,
19506 p1.0,
19507 aq,
19508 ad,
19509 m,
19510 w0.in_features(),
19511 p0.1,
19512 p1.1,
19513 p0.2,
19514 )?))
19515 }
19516
19517 #[allow(clippy::too_many_arguments)]
19518 fn q8_fused2_t_core(
19519 &self,
19520 b0: &CudaSlice<u8>,
19521 b1: &CudaSlice<u8>,
19522 aq: &CudaSlice<i8>,
19523 ad: &CudaSlice<f32>,
19524 m: usize,
19525 in_f: usize,
19526 out0: usize,
19527 out1: usize,
19528 row_bytes: usize,
19529 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19530 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19532 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19533 let f = self.func(match Self::batched_mcols(m) {
19534 2 => "qmatvec_q8_0_mmvq_fused2_b2",
19535 4 => "qmatvec_q8_0_mmvq_fused2_b4",
19536 _ => "qmatvec_q8_0_mmvq_fused2_b8",
19538 });
19539 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
19540 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
19541 let cfg = LaunchConfig {
19542 grid_dim: (nb0 + nb1, 1, 1),
19543 block_dim: (32, ROWS_PER_BLOCK, 1),
19544 shared_mem_bytes: 0,
19545 };
19546 let (inf, o0, o1, mi, rbl) = (
19547 in_f as i32,
19548 out0 as i32,
19549 out1 as i32,
19550 m as i32,
19551 row_bytes as i64,
19552 );
19553 let __s_b = self.gpu.stream();
19554 let mut b = __s_b.launch_builder(&f);
19555 b.arg(b0)
19556 .arg(b1)
19557 .arg(aq)
19558 .arg(ad)
19559 .arg(&mut y0)
19560 .arg(&mut y1)
19561 .arg(&inf)
19562 .arg(&o0)
19563 .arg(&o1)
19564 .arg(&mi)
19565 .arg(&rbl);
19566 unsafe {
19567 b.launch(cfg)?;
19568 }
19569 Ok((y0, y1))
19570 }
19571
19572 #[allow(clippy::too_many_arguments)]
19575 pub fn qmatvec_q8_fused2_t_raw(
19576 &self,
19577 b0: &CudaSlice<u8>,
19578 b1: &CudaSlice<u8>,
19579 x: &CudaSlice<f32>,
19580 m: usize,
19581 in_f: usize,
19582 out0: usize,
19583 out1: usize,
19584 row_bytes: usize,
19585 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19586 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19587 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
19588 }
19589
19590 #[allow(clippy::too_many_arguments)]
19593 #[allow(clippy::type_complexity)] pub fn matmul_q8_fused3_t(
19595 &self,
19596 w0: &crate::model::GpuTensor,
19597 w1: &crate::model::GpuTensor,
19598 w2: &crate::model::GpuTensor,
19599 aq: &CudaSlice<i8>,
19600 ad: &CudaSlice<f32>,
19601 m: usize,
19602 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
19603 {
19604 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
19605 return Ok(None);
19606 }
19607 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
19608 return Ok(Some(self.e4m3_fused3_t_core(
19609 p0.0,
19610 p1.0,
19611 p2.0,
19612 aq,
19613 ad,
19614 m,
19615 w0.in_features(),
19616 p0.1,
19617 p1.1,
19618 p2.1,
19619 p0.2,
19620 p0.3,
19621 p1.3,
19622 p2.3,
19623 )?));
19624 }
19625 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
19626 return Ok(None);
19627 };
19628 Ok(Some(self.q8_fused3_t_core(
19629 p0.0,
19630 p1.0,
19631 p2.0,
19632 aq,
19633 ad,
19634 m,
19635 w0.in_features(),
19636 p0.1,
19637 p1.1,
19638 p2.1,
19639 p0.2,
19640 )?))
19641 }
19642
19643 #[allow(clippy::too_many_arguments)]
19644 #[allow(clippy::type_complexity)] fn q8_fused3_t_core(
19646 &self,
19647 b0: &CudaSlice<u8>,
19648 b1: &CudaSlice<u8>,
19649 b2: &CudaSlice<u8>,
19650 aq: &CudaSlice<i8>,
19651 ad: &CudaSlice<f32>,
19652 m: usize,
19653 in_f: usize,
19654 out0: usize,
19655 out1: usize,
19656 out2: usize,
19657 row_bytes: usize,
19658 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19659 const ROWS_PER_BLOCK: u32 = 4;
19660 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19661 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19662 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19663 let f = self.func(if Self::batched_mcols(m) == 2 {
19664 "qmatvec_q8_0_mmvq_fused3_b2"
19665 } else {
19666 "qmatvec_q8_0_mmvq_fused3_b4"
19667 });
19668 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
19669 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
19670 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
19671 let cfg = LaunchConfig {
19672 grid_dim: (nb0 + nb1 + nb2, 1, 1),
19673 block_dim: (32, ROWS_PER_BLOCK, 1),
19674 shared_mem_bytes: 0,
19675 };
19676 let (inf, o0, o1, o2, mi, rbl) = (
19677 in_f as i32,
19678 out0 as i32,
19679 out1 as i32,
19680 out2 as i32,
19681 m as i32,
19682 row_bytes as i64,
19683 );
19684 let __s_b = self.gpu.stream();
19685 let mut b = __s_b.launch_builder(&f);
19686 b.arg(b0)
19687 .arg(b1)
19688 .arg(b2)
19689 .arg(aq)
19690 .arg(ad)
19691 .arg(&mut y0)
19692 .arg(&mut y1)
19693 .arg(&mut y2)
19694 .arg(&inf)
19695 .arg(&o0)
19696 .arg(&o1)
19697 .arg(&o2)
19698 .arg(&mi)
19699 .arg(&rbl);
19700 unsafe {
19701 b.launch(cfg)?;
19702 }
19703 Ok((y0, y1, y2))
19704 }
19705
19706 #[allow(clippy::too_many_arguments)]
19708 #[allow(clippy::type_complexity)] pub fn qmatvec_q8_fused3_t_raw(
19710 &self,
19711 b0: &CudaSlice<u8>,
19712 b1: &CudaSlice<u8>,
19713 b2: &CudaSlice<u8>,
19714 x: &CudaSlice<f32>,
19715 m: usize,
19716 in_f: usize,
19717 out0: usize,
19718 out1: usize,
19719 out2: usize,
19720 row_bytes: usize,
19721 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19722 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
19723 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
19724 }
19725
19726 pub fn q8_ffn_fuse2_on(&self) -> bool {
19730 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19731 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
19732 }
19733
19734 #[allow(clippy::type_complexity)]
19740 fn q8_fused_params<'w, const N: usize>(
19741 &self,
19742 ws: &[&'w crate::model::GpuTensor; N],
19743 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
19744 use crate::model::GpuTensor;
19745 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
19746 return None;
19747 }
19748 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
19749 return None;
19750 }
19751 let in_f = ws[0].in_features();
19752 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
19753 for (i, w) in ws.iter().enumerate() {
19754 match w {
19755 GpuTensor::Quant {
19756 bytes,
19757 qtype,
19758 row_bytes,
19759 scale,
19760 ..
19761 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
19762 out[i] = Some((bytes, w.out_features(), *row_bytes))
19763 }
19764 _ => return None,
19765 }
19766 }
19767 Some(out.map(|o| o.unwrap()))
19768 }
19769
19770 pub fn e4m3_dual_on(&self) -> bool {
19773 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19774 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
19775 }
19776
19777 #[allow(clippy::type_complexity)]
19789 fn e4m3_fused_params<'w, const N: usize>(
19790 &self,
19791 ws: &[&'w crate::model::GpuTensor; N],
19792 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
19793 use crate::model::GpuTensor;
19794 if !self.e4m3_dual_on() {
19795 return None;
19796 }
19797 let in_f = ws[0].in_features();
19798 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
19799 for (i, w) in ws.iter().enumerate() {
19800 match w {
19801 GpuTensor::Quant {
19802 bytes,
19803 qtype,
19804 row_bytes,
19805 scale,
19806 rp,
19807 rp4,
19808 ..
19809 } if *qtype == QT_F8_E4M3
19810 && w.in_features() == in_f
19811 && *row_bytes == in_f
19812 && !*rp
19813 && rp4.is_none() =>
19814 {
19815 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
19816 }
19817 _ => return None,
19818 }
19819 }
19820 Some(out.map(|o| o.unwrap()))
19821 }
19822
19823 #[allow(clippy::too_many_arguments)]
19827 fn e4m3_fused2_core(
19828 &self,
19829 b0: &CudaSlice<u8>,
19830 b1: &CudaSlice<u8>,
19831 aq: &CudaSlice<i8>,
19832 ad: &CudaSlice<f32>,
19833 in_f: usize,
19834 out0: usize,
19835 out1: usize,
19836 row_bytes: usize,
19837 ws0: f32,
19838 ws1: f32,
19839 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19840 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19842 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19843 let f = self.func("qmatvec_e4m3_mmvq_fused2");
19844 let mut y0 = self.alloc_uninit::<f32>(out0)?;
19845 let mut y1 = self.alloc_uninit::<f32>(out1)?;
19846 let cfg = LaunchConfig {
19847 grid_dim: (nb0 + nb1, 1, 1),
19848 block_dim: (32, ROWS_PER_BLOCK, 1),
19849 shared_mem_bytes: 0,
19850 };
19851 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
19852 let __s_b = self.gpu.stream();
19853 let mut b = __s_b.launch_builder(&f);
19854 b.arg(b0)
19855 .arg(b1)
19856 .arg(aq)
19857 .arg(ad)
19858 .arg(&mut y0)
19859 .arg(&mut y1)
19860 .arg(&inf)
19861 .arg(&o0)
19862 .arg(&o1)
19863 .arg(&rbl)
19864 .arg(&ws0)
19865 .arg(&ws1);
19866 unsafe {
19867 b.launch(cfg)?;
19868 }
19869 Ok((y0, y1))
19870 }
19871
19872 #[allow(clippy::too_many_arguments)]
19874 #[allow(clippy::type_complexity)] fn e4m3_fused3_core(
19876 &self,
19877 b0: &CudaSlice<u8>,
19878 b1: &CudaSlice<u8>,
19879 b2: &CudaSlice<u8>,
19880 aq: &CudaSlice<i8>,
19881 ad: &CudaSlice<f32>,
19882 in_f: usize,
19883 out0: usize,
19884 out1: usize,
19885 out2: usize,
19886 row_bytes: usize,
19887 ws0: f32,
19888 ws1: f32,
19889 ws2: f32,
19890 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19891 const ROWS_PER_BLOCK: u32 = 4;
19892 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
19893 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
19894 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
19895 let f = self.func("qmatvec_e4m3_mmvq_fused3");
19896 let mut y0 = self.alloc_uninit::<f32>(out0)?;
19897 let mut y1 = self.alloc_uninit::<f32>(out1)?;
19898 let mut y2 = self.alloc_uninit::<f32>(out2)?;
19899 let cfg = LaunchConfig {
19900 grid_dim: (nb0 + nb1 + nb2, 1, 1),
19901 block_dim: (32, ROWS_PER_BLOCK, 1),
19902 shared_mem_bytes: 0,
19903 };
19904 let (inf, o0, o1, o2, rbl) = (
19905 in_f as i32,
19906 out0 as i32,
19907 out1 as i32,
19908 out2 as i32,
19909 row_bytes as i64,
19910 );
19911 let __s_b = self.gpu.stream();
19912 let mut b = __s_b.launch_builder(&f);
19913 b.arg(b0)
19914 .arg(b1)
19915 .arg(b2)
19916 .arg(aq)
19917 .arg(ad)
19918 .arg(&mut y0)
19919 .arg(&mut y1)
19920 .arg(&mut y2)
19921 .arg(&inf)
19922 .arg(&o0)
19923 .arg(&o1)
19924 .arg(&o2)
19925 .arg(&rbl)
19926 .arg(&ws0)
19927 .arg(&ws1)
19928 .arg(&ws2);
19929 unsafe {
19930 b.launch(cfg)?;
19931 }
19932 Ok((y0, y1, y2))
19933 }
19934
19935 #[allow(clippy::too_many_arguments)]
19942 pub fn e4m3_fused6_into(
19943 &self,
19944 w: [&CudaSlice<u8>; 6],
19945 aq: &CudaSlice<i8>,
19946 ad: &CudaSlice<f32>,
19947 in_f: usize,
19948 dims: [usize; 6],
19949 row_bytes: usize,
19950 ws: [f32; 6],
19951 outs: &mut [CudaSlice<f32>; 6],
19952 ) -> Result<(), Box<dyn std::error::Error>> {
19953 self.e4m3_fused6_into_arm(
19954 w,
19955 aq,
19956 ad,
19957 in_f,
19958 dims,
19959 row_bytes,
19960 ws,
19961 outs,
19962 e4m3_row_ilp_level(),
19963 )
19964 }
19965
19966 #[allow(clippy::too_many_arguments)]
19971 pub fn e4m3_fused6_into_arm(
19972 &self,
19973 w: [&CudaSlice<u8>; 6],
19974 aq: &CudaSlice<i8>,
19975 ad: &CudaSlice<f32>,
19976 in_f: usize,
19977 dims: [usize; 6],
19978 row_bytes: usize,
19979 ws: [f32; 6],
19980 outs: &mut [CudaSlice<f32>; 6],
19981 arm: u32,
19982 ) -> Result<(), Box<dyn std::error::Error>> {
19983 const ROWS_PER_BLOCK: u32 = 4;
19984 let blocks: u32 = dims
19985 .iter()
19986 .map(|&o| (o as u32).div_ceil(ROWS_PER_BLOCK))
19987 .sum();
19988 let f = self.func(match arm {
19991 0 => "qmatvec_e4m3_mmvq_fused6",
19992 1 => "qmatvec_e4m3_mmvq_fused6_ilp",
19993 2 => "qmatvec_e4m3_mmvq_fused6_ilp2",
19994 3 => "qmatvec_e4m3_mmvq_fused6_ilp8",
19995 _ => "qmatvec_e4m3_mmvq_fused6_ilp16",
19996 });
19997 let cfg = LaunchConfig {
19998 grid_dim: (blocks, 1, 1),
19999 block_dim: (32, ROWS_PER_BLOCK, 1),
20000 shared_mem_bytes: 0,
20001 };
20002 let inf = in_f as i32;
20003 let o: [i32; 6] = std::array::from_fn(|i| dims[i] as i32);
20004 let rbl = row_bytes as i64;
20005 let (o0, o1) = outs.split_at_mut(1);
20006 let (o1, o2) = o1.split_at_mut(1);
20007 let (o2, o3) = o2.split_at_mut(1);
20008 let (o3, o4) = o3.split_at_mut(1);
20009 let (o4, o5) = o4.split_at_mut(1);
20010 let __s_b = self.gpu.stream();
20011 let mut b = __s_b.launch_builder(&f);
20012 b.arg(w[0])
20013 .arg(w[1])
20014 .arg(w[2])
20015 .arg(w[3])
20016 .arg(w[4])
20017 .arg(w[5])
20018 .arg(aq)
20019 .arg(ad)
20020 .arg(&mut o0[0])
20021 .arg(&mut o1[0])
20022 .arg(&mut o2[0])
20023 .arg(&mut o3[0])
20024 .arg(&mut o4[0])
20025 .arg(&mut o5[0])
20026 .arg(&inf)
20027 .arg(&o[0])
20028 .arg(&o[1])
20029 .arg(&o[2])
20030 .arg(&o[3])
20031 .arg(&o[4])
20032 .arg(&o[5])
20033 .arg(&rbl)
20034 .arg(&ws[0])
20035 .arg(&ws[1])
20036 .arg(&ws[2])
20037 .arg(&ws[3])
20038 .arg(&ws[4])
20039 .arg(&ws[5]);
20040 unsafe {
20041 b.launch(cfg)?;
20042 }
20043 Ok(())
20044 }
20045
20046 #[allow(clippy::too_many_arguments)]
20050 fn e4m3_fused2_t_core(
20051 &self,
20052 b0: &CudaSlice<u8>,
20053 b1: &CudaSlice<u8>,
20054 aq: &CudaSlice<i8>,
20055 ad: &CudaSlice<f32>,
20056 m: usize,
20057 in_f: usize,
20058 out0: usize,
20059 out1: usize,
20060 row_bytes: usize,
20061 ws0: f32,
20062 ws1: f32,
20063 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20064 const ROWS_PER_BLOCK: u32 = 4;
20065 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
20066 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
20067 let f = self.func(match Self::batched_mcols(m) {
20068 2 => "qmatvec_e4m3_mmvq_fused2_b2",
20069 4 => "qmatvec_e4m3_mmvq_fused2_b4",
20070 _ => "qmatvec_e4m3_mmvq_fused2_b8",
20071 });
20072 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
20073 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
20074 let cfg = LaunchConfig {
20075 grid_dim: (nb0 + nb1, 1, 1),
20076 block_dim: (32, ROWS_PER_BLOCK, 1),
20077 shared_mem_bytes: 0,
20078 };
20079 let (inf, o0, o1, mi, rbl) = (
20080 in_f as i32,
20081 out0 as i32,
20082 out1 as i32,
20083 m as i32,
20084 row_bytes as i64,
20085 );
20086 let __s_b = self.gpu.stream();
20087 let mut b = __s_b.launch_builder(&f);
20088 b.arg(b0)
20089 .arg(b1)
20090 .arg(aq)
20091 .arg(ad)
20092 .arg(&mut y0)
20093 .arg(&mut y1)
20094 .arg(&inf)
20095 .arg(&o0)
20096 .arg(&o1)
20097 .arg(&mi)
20098 .arg(&rbl);
20099 unsafe {
20100 b.launch(cfg)?;
20101 }
20102 if ws0 != 1.0 {
20103 self.scale_inplace(&mut y0, ws0, m * out0)?;
20104 }
20105 if ws1 != 1.0 {
20106 self.scale_inplace(&mut y1, ws1, m * out1)?;
20107 }
20108 Ok((y0, y1))
20109 }
20110
20111 #[allow(clippy::too_many_arguments)]
20113 #[allow(clippy::type_complexity)] fn e4m3_fused3_t_core(
20115 &self,
20116 b0: &CudaSlice<u8>,
20117 b1: &CudaSlice<u8>,
20118 b2: &CudaSlice<u8>,
20119 aq: &CudaSlice<i8>,
20120 ad: &CudaSlice<f32>,
20121 m: usize,
20122 in_f: usize,
20123 out0: usize,
20124 out1: usize,
20125 out2: usize,
20126 row_bytes: usize,
20127 ws0: f32,
20128 ws1: f32,
20129 ws2: f32,
20130 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20131 const ROWS_PER_BLOCK: u32 = 4;
20132 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
20133 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
20134 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
20135 let f = self.func(if Self::batched_mcols(m) == 2 {
20136 "qmatvec_e4m3_mmvq_fused3_b2"
20137 } else {
20138 "qmatvec_e4m3_mmvq_fused3_b4"
20139 });
20140 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
20141 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
20142 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
20143 let cfg = LaunchConfig {
20144 grid_dim: (nb0 + nb1 + nb2, 1, 1),
20145 block_dim: (32, ROWS_PER_BLOCK, 1),
20146 shared_mem_bytes: 0,
20147 };
20148 let (inf, o0, o1, o2, mi, rbl) = (
20149 in_f as i32,
20150 out0 as i32,
20151 out1 as i32,
20152 out2 as i32,
20153 m as i32,
20154 row_bytes as i64,
20155 );
20156 let __s_b = self.gpu.stream();
20157 let mut b = __s_b.launch_builder(&f);
20158 b.arg(b0)
20159 .arg(b1)
20160 .arg(b2)
20161 .arg(aq)
20162 .arg(ad)
20163 .arg(&mut y0)
20164 .arg(&mut y1)
20165 .arg(&mut y2)
20166 .arg(&inf)
20167 .arg(&o0)
20168 .arg(&o1)
20169 .arg(&o2)
20170 .arg(&mi)
20171 .arg(&rbl);
20172 unsafe {
20173 b.launch(cfg)?;
20174 }
20175 if ws0 != 1.0 {
20176 self.scale_inplace(&mut y0, ws0, m * out0)?;
20177 }
20178 if ws1 != 1.0 {
20179 self.scale_inplace(&mut y1, ws1, m * out1)?;
20180 }
20181 if ws2 != 1.0 {
20182 self.scale_inplace(&mut y2, ws2, m * out2)?;
20183 }
20184 Ok((y0, y1, y2))
20185 }
20186
20187 #[allow(clippy::too_many_arguments)] pub fn qmatvec_e4m3_blk_mmvq(
20198 &self,
20199 bytes: &CudaSlice<u8>,
20200 aq: &CudaSlice<i8>,
20201 ad: &CudaSlice<f32>,
20202 scales: &CudaSlice<f32>,
20203 m: usize,
20204 in_f: usize,
20205 out_f: usize,
20206 row_bytes: usize,
20207 scale_cols: usize,
20208 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20209 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
20211 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
20212 )?;
20213 Ok(y)
20214 }
20215
20216 #[allow(clippy::too_many_arguments)]
20218 pub fn qmatvec_e4m3_blk_mmvq_into(
20219 &self,
20220 bytes: &CudaSlice<u8>,
20221 aq: &CudaSlice<i8>,
20222 ad: &CudaSlice<f32>,
20223 scales: &CudaSlice<f32>,
20224 m: usize,
20225 in_f: usize,
20226 out_f: usize,
20227 row_bytes: usize,
20228 scale_cols: usize,
20229 y: &mut CudaSlice<f32>,
20230 ) -> Result<(), Box<dyn std::error::Error>> {
20231 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
20233 let cfg = LaunchConfig {
20234 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
20235 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
20238 let (inf, outf, mi, rb, sc) = (
20239 in_f as i32,
20240 out_f as i32,
20241 m as i32,
20242 row_bytes as i64,
20243 scale_cols as i32,
20244 );
20245 let __s_b = self.gpu.stream();
20246 let mut b = __s_b.launch_builder(&f);
20247 b.arg(bytes)
20248 .arg(aq)
20249 .arg(ad)
20250 .arg(scales)
20251 .arg(&mut *y)
20252 .arg(&inf)
20253 .arg(&outf)
20254 .arg(&mi)
20255 .arg(&rb)
20256 .arg(&sc);
20257 unsafe {
20258 b.launch(cfg)?;
20259 }
20260 Ok(())
20261 }
20262
20263 #[allow(clippy::too_many_arguments)]
20269 pub fn qmatvec_e4m3_blk_mmvq_batched(
20270 &self,
20271 bytes: &CudaSlice<u8>,
20272 aq: &CudaSlice<i8>,
20273 ad: &CudaSlice<f32>,
20274 scales: &CudaSlice<f32>,
20275 m: usize,
20276 in_f: usize,
20277 out_f: usize,
20278 row_bytes: usize,
20279 scale_cols: usize,
20280 mcols: usize,
20281 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20282 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
20284 let name = match mcols {
20285 2 => "qmatvec_e4m3_blk_mmvq_b2",
20286 4 => "qmatvec_e4m3_blk_mmvq_b4",
20287 8 => "qmatvec_e4m3_blk_mmvq_b8",
20288 16 => "qmatvec_e4m3_blk_mmvq_b16",
20289 _ => {
20290 return Err(
20291 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
20292 );
20293 }
20294 };
20295 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20296 let f = self.func(name);
20297 let cfg = LaunchConfig {
20298 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
20299 block_dim: (32, ROWS_PER_BLOCK, 1),
20300 shared_mem_bytes: 0,
20301 };
20302 let (inf, outf, mi, rb, sc) = (
20303 in_f as i32,
20304 out_f as i32,
20305 m as i32,
20306 row_bytes as i64,
20307 scale_cols as i32,
20308 );
20309 let __s_b = self.gpu.stream();
20310 let mut b = __s_b.launch_builder(&f);
20311 b.arg(bytes)
20312 .arg(aq)
20313 .arg(ad)
20314 .arg(scales)
20315 .arg(&mut y)
20316 .arg(&inf)
20317 .arg(&outf)
20318 .arg(&mi)
20319 .arg(&rb)
20320 .arg(&sc);
20321 unsafe {
20322 b.launch(cfg)?;
20323 }
20324 Ok(y)
20325 }
20326
20327 #[allow(clippy::too_many_arguments)]
20330 pub fn qmatvec_e4m3_blk_batched_raw(
20331 &self,
20332 bytes: &CudaSlice<u8>,
20333 x: &CudaSlice<f32>,
20334 scales: &CudaSlice<f32>,
20335 m: usize,
20336 in_f: usize,
20337 out_f: usize,
20338 row_bytes: usize,
20339 scale_cols: usize,
20340 mcols: usize,
20341 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20342 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20343 self.qmatvec_e4m3_blk_mmvq_batched(
20344 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
20345 )
20346 }
20347
20348 #[allow(clippy::too_many_arguments)]
20351 pub fn qmatvec_e4m3_blk_mmvq_raw(
20352 &self,
20353 bytes: &CudaSlice<u8>,
20354 x: &CudaSlice<f32>,
20355 scales: &CudaSlice<f32>,
20356 m: usize,
20357 in_f: usize,
20358 out_f: usize,
20359 row_bytes: usize,
20360 scale_cols: usize,
20361 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20362 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20363 self.qmatvec_e4m3_blk_mmvq(
20364 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
20365 )
20366 }
20367
20368 #[allow(clippy::too_many_arguments)]
20371 pub fn qmatvec_e4m3_fused2_raw(
20372 &self,
20373 b0: &CudaSlice<u8>,
20374 b1: &CudaSlice<u8>,
20375 x: &CudaSlice<f32>,
20376 in_f: usize,
20377 out0: usize,
20378 out1: usize,
20379 row_bytes: usize,
20380 ws0: f32,
20381 ws1: f32,
20382 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20383 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20384 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
20385 }
20386
20387 #[allow(clippy::too_many_arguments)]
20388 #[allow(clippy::type_complexity)] pub fn qmatvec_e4m3_fused3_raw(
20390 &self,
20391 b0: &CudaSlice<u8>,
20392 b1: &CudaSlice<u8>,
20393 b2: &CudaSlice<u8>,
20394 x: &CudaSlice<f32>,
20395 in_f: usize,
20396 out0: usize,
20397 out1: usize,
20398 out2: usize,
20399 row_bytes: usize,
20400 ws0: f32,
20401 ws1: f32,
20402 ws2: f32,
20403 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20404 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20405 self.e4m3_fused3_core(
20406 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
20407 )
20408 }
20409
20410 #[allow(clippy::too_many_arguments)]
20414 pub fn qmatvec_e4m3_fused6_raw(
20415 &self,
20416 w: [&CudaSlice<u8>; 6],
20417 x: &CudaSlice<f32>,
20418 in_f: usize,
20419 dims: [usize; 6],
20420 row_bytes: usize,
20421 ws: [f32; 6],
20422 arm: u32,
20423 ) -> Result<[CudaSlice<f32>; 6], Box<dyn std::error::Error>> {
20424 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
20425 let mut outs = [
20426 self.alloc_uninit::<f32>(dims[0])?,
20427 self.alloc_uninit::<f32>(dims[1])?,
20428 self.alloc_uninit::<f32>(dims[2])?,
20429 self.alloc_uninit::<f32>(dims[3])?,
20430 self.alloc_uninit::<f32>(dims[4])?,
20431 self.alloc_uninit::<f32>(dims[5])?,
20432 ];
20433 self.e4m3_fused6_into_arm(w, &aq, &ad, in_f, dims, row_bytes, ws, &mut outs, arm)?;
20434 Ok(outs)
20435 }
20436
20437 #[allow(clippy::too_many_arguments)]
20438 pub fn qmatvec_e4m3_fused2_t_raw(
20439 &self,
20440 b0: &CudaSlice<u8>,
20441 b1: &CudaSlice<u8>,
20442 x: &CudaSlice<f32>,
20443 m: usize,
20444 in_f: usize,
20445 out0: usize,
20446 out1: usize,
20447 row_bytes: usize,
20448 ws0: f32,
20449 ws1: f32,
20450 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20451 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20452 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
20453 }
20454
20455 #[allow(clippy::too_many_arguments)]
20456 #[allow(clippy::type_complexity)] pub fn qmatvec_e4m3_fused3_t_raw(
20458 &self,
20459 b0: &CudaSlice<u8>,
20460 b1: &CudaSlice<u8>,
20461 b2: &CudaSlice<u8>,
20462 x: &CudaSlice<f32>,
20463 m: usize,
20464 in_f: usize,
20465 out0: usize,
20466 out1: usize,
20467 out2: usize,
20468 row_bytes: usize,
20469 ws0: f32,
20470 ws1: f32,
20471 ws2: f32,
20472 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20473 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
20474 self.e4m3_fused3_t_core(
20475 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
20476 )
20477 }
20478
20479 fn try_e4m3_blk_pre(
20490 &self,
20491 w: &crate::model::GpuTensor,
20492 aq: &CudaSlice<i8>,
20493 ad: &CudaSlice<f32>,
20494 m: usize,
20495 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20496 use crate::model::GpuTensor;
20497 if let GpuTensor::Quant {
20498 bytes,
20499 qtype,
20500 row_bytes,
20501 blk: Some(g),
20502 ..
20503 } = w
20504 && *qtype == QT_F8_E4M3_BLK
20505 {
20506 if (2..=16).contains(&m)
20512 && std::env::var("MEMRA_NO_BATCHED").is_err()
20513 && (m <= 4 || Self::b8_enabled())
20514 {
20515 let mcols = Self::batched_mcols(m);
20516 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
20517 bytes,
20518 aq,
20519 ad,
20520 &g.scales,
20521 m,
20522 w.in_features(),
20523 w.out_features(),
20524 *row_bytes,
20525 g.cols,
20526 mcols,
20527 )?));
20528 }
20529 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
20530 bytes,
20531 aq,
20532 ad,
20533 &g.scales,
20534 m,
20535 w.in_features(),
20536 w.out_features(),
20537 *row_bytes,
20538 g.cols,
20539 )?));
20540 }
20541 Ok(None)
20542 }
20543
20544 fn try_e4m3_blk_prefill(
20591 &self,
20592 w: &crate::model::GpuTensor,
20593 x: &CudaSlice<f32>,
20594 m: usize,
20595 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
20596 use crate::model::GpuTensor;
20597 let GpuTensor::Quant {
20598 bytes,
20599 qtype,
20600 blk: Some(g),
20601 ..
20602 } = w
20603 else {
20604 return Ok(None);
20605 };
20606 if *qtype != QT_F8_E4M3_BLK {
20607 return Ok(None);
20608 }
20609 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
20614 return Ok(Some(y));
20615 }
20616 let (in_f, out_f) = (w.in_features(), w.out_features());
20617 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
20618 let tmp = GpuTensor::Quant {
20619 bytes: slab,
20620 qtype: QT_Q8_0,
20621 row_bytes: in_f / 32 * 34,
20622 ne: vec![in_f as u64, out_f as u64],
20623 scale: 1.0,
20624 rp: false,
20625 #[cfg(memra_cutlass)]
20626 cutlass: None,
20627 fp8: None,
20628 blk: None,
20629 f16: None,
20630 rp4: None,
20631 };
20632 Ok(Some(self.matmul(&tmp, x, m)?))
20634 }
20635
20636 #[allow(clippy::type_complexity)] pub fn matmul_pre_noscale(
20638 &self,
20639 w: &crate::model::GpuTensor,
20640 aq: &CudaSlice<i8>,
20641 ad: &CudaSlice<f32>,
20642 m: usize,
20643 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
20644 use crate::model::GpuTensor;
20645 if m == 1
20649 && let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)?
20650 {
20651 return Ok(Some((y, 1.0)));
20652 }
20653 if m != 1 || !self.uses_q8_1_fast(w) {
20655 return Ok(None);
20656 }
20657 let in_f = w.in_features();
20658 let out_f = w.out_features();
20659 let (bytes, qtype, row_bytes, scale, rp) = match w {
20660 GpuTensor::Quant {
20661 bytes,
20662 qtype,
20663 row_bytes,
20664 scale,
20665 rp,
20666 ..
20667 } => (bytes, *qtype, *row_bytes, *scale, *rp),
20668 _ => return Ok(None),
20669 };
20670 if self.mmvq_supports(qtype) {
20672 let (mbytes, mrp) = match w {
20674 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
20675 _ => (bytes, rp),
20676 };
20677 let y = self.qmatvec_mmvq(
20678 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
20679 )?;
20680 return Ok(Some((y, scale)));
20681 }
20682 let name = match qtype {
20684 QT_Q8_0 => "qmatvec_q8_0_dp4a",
20685 QT_Q4_K => "qmatvec_q4_K_dp4a",
20686 QT_Q6_K => "qmatvec_q6_K_dp4a",
20687 QT_Q5_K => "qmatvec_q5_K_dp4a",
20688 QT_Q3_K => "qmatvec_q3_K_dp4a",
20689 QT_NVFP4 => {
20690 if rp {
20691 "qmatvec_nvfp4_dp4a_rp"
20692 } else {
20693 "qmatvec_nvfp4_dp4a"
20694 }
20695 }
20696 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
20697 _ => return Ok(None),
20698 };
20699 let f = self.func(name);
20700 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
20701 let cfg = LaunchConfig {
20702 grid_dim: (out_f as u32, m as u32, 1),
20703 block_dim: (128, 1, 1),
20704 shared_mem_bytes: 0,
20705 };
20706 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20707 let __s_b = self.gpu.stream();
20708 let mut b = __s_b.launch_builder(&f);
20709 b.arg(bytes)
20710 .arg(aq)
20711 .arg(ad)
20712 .arg(&mut y)
20713 .arg(&inf)
20714 .arg(&outf)
20715 .arg(&mi)
20716 .arg(&rb);
20717 unsafe {
20718 b.launch(cfg)?;
20719 }
20720 Ok(Some((y, scale)))
20721 }
20722
20723 pub fn mmvq_supports(&self, qtype: i32) -> bool {
20726 if qtype == QT_F8_E4M3 {
20731 return true;
20732 }
20733 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
20734 return false;
20735 }
20736 matches!(
20737 qtype,
20738 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
20739 )
20740 }
20741
20742 #[allow(clippy::too_many_arguments)] pub fn qmatvec_mmvq(
20748 &self,
20749 bytes: &CudaSlice<u8>,
20750 aq: &CudaSlice<i8>,
20751 ad: &CudaSlice<f32>,
20752 m: usize,
20753 in_f: usize,
20754 out_f: usize,
20755 qtype: i32,
20756 row_bytes: usize,
20757 scale: f32,
20758 rp: bool,
20759 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20760 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
20762 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
20763 )?;
20764 Ok(y)
20765 }
20766
20767 #[allow(clippy::too_many_arguments)]
20769 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_mmvq_into(
20771 &self,
20772 bytes: &CudaSlice<u8>,
20773 aq: &CudaSlice<i8>,
20774 ad: &CudaSlice<f32>,
20775 m: usize,
20776 in_f: usize,
20777 out_f: usize,
20778 qtype: i32,
20779 row_bytes: usize,
20780 scale: f32,
20781 rp: bool,
20782 y: &mut CudaSlice<f32>,
20783 ) -> Result<(), Box<dyn std::error::Error>> {
20784 debug_assert!(y.len() >= m * out_f);
20785 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
20791 && rp
20792 && m == 1
20793 && out_f >= 64
20794 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
20795 && {
20796 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20797 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
20798 }
20799 {
20800 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
20801 let cfg = LaunchConfig {
20802 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
20803 block_dim: (32, 2, 1),
20804 shared_mem_bytes: 0,
20805 };
20806 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
20807 let __s_b = self.gpu.stream();
20808 let mut b = __s_b.launch_builder(&f);
20809 b.arg(bytes)
20810 .arg(aq)
20811 .arg(ad)
20812 .arg(&mut *y)
20813 .arg(&inf)
20814 .arg(&outf)
20815 .arg(&mi)
20816 .arg(&rb);
20817 unsafe {
20818 b.launch(cfg)?;
20819 }
20820 if scale != 1.0 {
20821 self.scale_inplace(y, scale, out_f)?;
20822 }
20823 return Ok(());
20824 }
20825 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
20834 2
20835 } else {
20836 1
20837 };
20838 if m == 1 && qtype == QT_Q4_0 {
20843 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
20844 mr = *Q40MR.get_or_init(|| {
20847 std::env::var("MEMRA_Q40_MR")
20848 .ok()
20849 .and_then(|v| v.parse().ok())
20850 .unwrap_or(1)
20851 });
20852 }
20853 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
20864 let q5_force = q5_mode.as_deref() == Some("2");
20865 let q5_il = qtype == QT_Q5_K
20868 && m == 1
20869 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
20870 if q5_il && !q5_force && out_f > 65536 {
20871 mr = 1;
20872 }
20873 if qtype == QT_Q4_0 && rp && mr != 1 {
20876 mr = 2;
20877 }
20878 if qtype == QT_Q8_0 && rp {
20882 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
20883 mr = *Q80MR.get_or_init(|| {
20884 std::env::var("MEMRA_Q80_MR")
20885 .ok()
20886 .and_then(|v| v.parse().ok())
20887 .unwrap_or(1)
20888 });
20889 }
20890 if qtype == QT_NVFP4
20902 && mr == 2
20903 && rp
20904 && m == 1
20905 && b200_matvec_arm_on()
20906 && (out_f as u32).div_ceil(ROWS_PER_BLOCK * 2)
20907 < b200_mr1_fill() * self.sm_count() as u32
20908 {
20909 mr = 1;
20910 }
20911 let nv_ilp = qtype == QT_NVFP4 && rp && nvfp4_row_ilp_on();
20914 if nv_ilp
20915 && NVFP4_ROW_ILP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0
20916 {
20917 eprintln!(
20918 "[nvfp4-row-ilp] engaged: split-plane NVFP4 trunk matvec with four groups' loads \
20919 per lane ahead of the dp4a chains (MEMRA_NVFP4_ROW_ILP=1, mr={mr})"
20920 );
20921 }
20922 let name = match (qtype, mr, rp) {
20923 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
20924 (QT_NVFP4, 2, true) if nv_ilp => "qmatvec_nvfp4_mmvq_mr2_rp_ilp",
20925 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
20926 (QT_NVFP4, _, true) if nv_ilp => "qmatvec_nvfp4_mmvq_rp_ilp",
20927 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
20928 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
20929 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
20930 (QT_Q5_K, 2, _) => {
20931 if q5_il {
20932 "qmatvec_q5_K_mmvq_mr2_il"
20933 } else {
20934 "qmatvec_q5_K_mmvq_mr2"
20935 }
20936 }
20937 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
20938 (QT_Q8_0, _, true)
20943 if in_f.is_multiple_of(1024) && {
20944 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20945 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
20946 } =>
20947 {
20948 "qmatvec_q8_0_mmvq_rpca"
20949 }
20950 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
20951 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
20952 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
20956 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
20957 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
20958 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
20959 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
20960 (QT_Q5_K, _, _) => {
20961 if q5_il {
20962 "qmatvec_q5_K_mmvq_il"
20963 } else {
20964 "qmatvec_q5_K_mmvq"
20965 }
20966 }
20967 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
20968 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
20969 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
20970 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
20971 };
20972 let f = self.func(name);
20973 let rows_per_block = ROWS_PER_BLOCK * mr;
20975 let cfg = LaunchConfig {
20976 grid_dim: (
20977 (out_f as u32 + rows_per_block - 1) / rows_per_block,
20978 m as u32,
20979 1,
20980 ),
20981 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
20984 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
20985 let __s_b = self.gpu.stream();
20986 let mut b = __s_b.launch_builder(&f);
20987 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
20992 if Self::pdl_on()
20998 && Self::pdl_mmvq_on()
20999 && Self::pdl_nvfp4q8_on()
21000 && matches!(
21001 name,
21002 "qmatvec_nvfp4_mmvq_mr2_rp"
21003 | "qmatvec_nvfp4_mmvq_mr2_rp_ilp"
21004 | "qmatvec_nvfp4_mmvq_rp"
21005 | "qmatvec_nvfp4_mmvq_rp_ilp"
21006 )
21007 {
21008 use cudarc::driver::{DevicePtr, DevicePtrMut};
21009 let s = &self.gpu.stream();
21010 let (pw, _g0) = bytes.device_ptr(s);
21011 let (paq, _g1) = aq.device_ptr(s);
21012 let (pad, _g2) = ad.device_ptr(s);
21013 let (py, _g3) = y.device_ptr_mut(s);
21014 let mut ps = [
21015 &pw as *const _ as *mut std::ffi::c_void,
21016 &paq as *const _ as *mut _,
21017 &pad as *const _ as *mut _,
21018 &py as *const _ as *mut _,
21019 &inf as *const _ as *mut _,
21020 &outf as *const _ as *mut _,
21021 &mi as *const _ as *mut _,
21022 &rb as *const _ as *mut _,
21023 &scale as *const _ as *mut _,
21024 ];
21025 unsafe {
21026 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
21027 }
21028 return Ok(());
21029 }
21030 b.arg(bytes)
21031 .arg(aq)
21032 .arg(ad)
21033 .arg(&mut *y)
21034 .arg(&inf)
21035 .arg(&outf)
21036 .arg(&mi)
21037 .arg(&rb)
21038 .arg(&scale);
21039 unsafe {
21040 b.launch(cfg)?;
21041 }
21042 } else if Self::pdl_on()
21043 && Self::pdl_mmvq_on()
21044 && (matches!(
21045 name,
21046 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
21047 ) || (Self::pdl_nvfp4q8_on()
21048 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
21049 {
21050 {
21054 use cudarc::driver::{DevicePtr, DevicePtrMut};
21055 let s = &self.gpu.stream();
21056 let (pw, _g0) = bytes.device_ptr(s);
21057 let (paq, _g1) = aq.device_ptr(s);
21058 let (pad, _g2) = ad.device_ptr(s);
21059 let (py, _g3) = y.device_ptr_mut(s);
21060 let mut ps = [
21061 &pw as *const _ as *mut std::ffi::c_void,
21062 &paq as *const _ as *mut _,
21063 &pad as *const _ as *mut _,
21064 &py as *const _ as *mut _,
21065 &inf as *const _ as *mut _,
21066 &outf as *const _ as *mut _,
21067 &mi as *const _ as *mut _,
21068 &rb as *const _ as *mut _,
21069 ];
21070 unsafe {
21071 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
21072 }
21073 }
21074 if scale != 1.0 {
21075 self.scale_inplace(y, scale, m * out_f)?;
21076 }
21077 } else {
21078 b.arg(bytes)
21079 .arg(aq)
21080 .arg(ad)
21081 .arg(&mut *y)
21082 .arg(&inf)
21083 .arg(&outf)
21084 .arg(&mi)
21085 .arg(&rb);
21086 unsafe {
21087 b.launch(cfg)?;
21088 }
21089 if scale != 1.0 {
21090 self.scale_inplace(y, scale, m * out_f)?;
21091 }
21092 }
21093 Ok(())
21094 }
21095
21096 #[allow(clippy::too_many_arguments)] pub fn qmatvec_mmvq_raw(
21101 &self,
21102 bytes: &CudaSlice<u8>,
21103 x: &CudaSlice<f32>,
21104 m: usize,
21105 in_f: usize,
21106 out_f: usize,
21107 qtype: i32,
21108 row_bytes: usize,
21109 rp: bool,
21110 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21111 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
21112 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
21113 }
21114
21115 #[allow(clippy::too_many_arguments)]
21124 pub fn qmatvec_nvfp4_rp_arm_raw(
21125 &self,
21126 bytes: &CudaSlice<u8>,
21127 aq: &CudaSlice<i8>,
21128 ad: &CudaSlice<f32>,
21129 m: usize,
21130 in_f: usize,
21131 out_f: usize,
21132 row_bytes: usize,
21133 yscale: f32,
21134 use_arm: bool,
21135 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21136 const ROWS_PER_BLOCK: u32 = 4; let mr: u32 = if use_arm { 1 } else { 2 };
21138 let name = if use_arm {
21139 "qmatvec_nvfp4_mmvq_rp"
21140 } else {
21141 "qmatvec_nvfp4_mmvq_mr2_rp"
21142 };
21143 let f = self.func(name);
21144 let rows_per_block = ROWS_PER_BLOCK * mr;
21145 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
21146 let cfg = LaunchConfig {
21147 grid_dim: ((out_f as u32).div_ceil(rows_per_block), m as u32, 1),
21148 block_dim: (32, ROWS_PER_BLOCK, 1),
21149 shared_mem_bytes: 0,
21150 };
21151 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
21152 let __s_b = self.gpu.stream();
21153 let mut b = __s_b.launch_builder(&f);
21154 b.arg(bytes)
21155 .arg(aq)
21156 .arg(ad)
21157 .arg(&mut y)
21158 .arg(&inf)
21159 .arg(&outf)
21160 .arg(&mi)
21161 .arg(&rb)
21162 .arg(&yscale);
21163 unsafe {
21164 b.launch(cfg)?;
21165 }
21166 Ok(y)
21167 }
21168
21169 #[allow(clippy::too_many_arguments)]
21172 pub fn qmatvec_nvfp4_rp_arm_raw_ilp(
21173 &self,
21174 bytes: &CudaSlice<u8>,
21175 aq: &CudaSlice<i8>,
21176 ad: &CudaSlice<f32>,
21177 m: usize,
21178 in_f: usize,
21179 out_f: usize,
21180 row_bytes: usize,
21181 yscale: f32,
21182 use_arm: bool,
21183 ilp: bool,
21184 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21185 const ROWS_PER_BLOCK: u32 = 4; let mr: u32 = if use_arm { 1 } else { 2 };
21187 let name = match (use_arm, ilp) {
21188 (true, false) => "qmatvec_nvfp4_mmvq_rp",
21189 (true, true) => "qmatvec_nvfp4_mmvq_rp_ilp",
21190 (false, false) => "qmatvec_nvfp4_mmvq_mr2_rp",
21191 (false, true) => "qmatvec_nvfp4_mmvq_mr2_rp_ilp",
21192 };
21193 let f = self.func(name);
21194 let rows_per_block = ROWS_PER_BLOCK * mr;
21195 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
21196 let cfg = LaunchConfig {
21197 grid_dim: ((out_f as u32).div_ceil(rows_per_block), m as u32, 1),
21198 block_dim: (32, ROWS_PER_BLOCK, 1),
21199 shared_mem_bytes: 0,
21200 };
21201 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
21202 let __s_b = self.gpu.stream();
21203 let mut b = __s_b.launch_builder(&f);
21204 b.arg(bytes)
21205 .arg(aq)
21206 .arg(ad)
21207 .arg(&mut y)
21208 .arg(&inf)
21209 .arg(&outf)
21210 .arg(&mi)
21211 .arg(&rb)
21212 .arg(&yscale);
21213 unsafe {
21214 b.launch(cfg)?;
21215 }
21216 Ok(y)
21217 }
21218
21219 #[allow(clippy::too_many_arguments)]
21224 pub fn qmatvec_nvfp4_fused2_rp_arm_raw(
21225 &self,
21226 w0: &CudaSlice<u8>,
21227 w1: &CudaSlice<u8>,
21228 aq: &CudaSlice<i8>,
21229 ad: &CudaSlice<f32>,
21230 in_f: usize,
21231 out0: usize,
21232 out1: usize,
21233 s0: f32,
21234 s1: f32,
21235 use_arm: bool,
21236 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21237 const ROWS_PER_BLOCK: u32 = 4; let rpw: u32 = if use_arm { 1 } else { 2 };
21239 let name = if use_arm {
21240 "qmatvec_nvfp4_mmvq_fused2_rp_g2"
21241 } else {
21242 "qmatvec_nvfp4_mmvq_fused2_rp"
21243 };
21244 let rows_pb = ROWS_PER_BLOCK * rpw;
21245 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
21246 let f = self.func(name);
21247 let mut y0 = self.alloc_uninit::<f32>(out0)?;
21248 let mut y1 = self.alloc_uninit::<f32>(out1)?;
21249 let cfg = LaunchConfig {
21250 grid_dim: (nb(out0) + nb(out1), 1, 1),
21251 block_dim: (32, ROWS_PER_BLOCK, 1),
21252 shared_mem_bytes: 0,
21253 };
21254 let (inf, oi0, oi1, mi) = (in_f as i32, out0 as i32, out1 as i32, 1i32);
21255 let __s_b = self.gpu.stream();
21256 let mut b = __s_b.launch_builder(&f);
21257 b.arg(w0)
21258 .arg(w1)
21259 .arg(aq)
21260 .arg(ad)
21261 .arg(&mut y0)
21262 .arg(&mut y1)
21263 .arg(&inf)
21264 .arg(&oi0)
21265 .arg(&oi1)
21266 .arg(&mi)
21267 .arg(&s0)
21268 .arg(&s1);
21269 unsafe {
21270 b.launch(cfg)?;
21271 }
21272 Ok((y0, y1))
21273 }
21274
21275 #[allow(clippy::too_many_arguments)]
21280 pub fn matvec_bf16_f32acc_x4_rows_arm_raw(
21281 &self,
21282 w: &CudaSlice<u8>,
21283 x: &CudaSlice<f32>,
21284 y: &mut CudaSlice<f32>,
21285 in_f: usize,
21286 out_f: usize,
21287 t: usize,
21288 use_arm: bool,
21289 ) -> Result<(), Box<dyn std::error::Error>> {
21290 let name = if use_arm {
21291 "matvec_bf16_f32acc_x4_rows_pf"
21292 } else {
21293 "matvec_bf16_f32acc_x4_rows"
21294 };
21295 let f = self.func(name);
21296 let cfg = LaunchConfig {
21297 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
21298 block_dim: (mmv_block(), 1, 1),
21299 shared_mem_bytes: 0,
21300 };
21301 let (ini, outi) = (in_f as i32, out_f as i32);
21302 let __s_b = self.gpu.stream();
21303 let mut b = __s_b.launch_builder(&f);
21304 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21305 unsafe {
21306 b.launch(cfg)?;
21307 }
21308 Ok(())
21309 }
21310
21311 pub fn gemv_v2_ksplit(&self, in_f: usize, out_f: usize, t: usize) -> usize {
21317 let blocks = out_f.div_ceil(GEMV_V2_ROWS) * t.max(1);
21318 let want = 2 * self.sm_count().max(1) as usize;
21319 if blocks >= want || blocks == 0 {
21320 return 1;
21321 }
21322 let per_step = mmv_block() as usize * 8;
21323 let max_ks = (in_f / per_step).max(1);
21324 want.div_ceil(blocks).clamp(1, max_ks)
21325 }
21326
21327 #[allow(clippy::too_many_arguments)]
21332 pub fn matvec_bf16_v2_raw(
21333 &self,
21334 w: &CudaSlice<u8>,
21335 x: &CudaSlice<f32>,
21336 y: &mut CudaSlice<f32>,
21337 in_f: usize,
21338 out_f: usize,
21339 t: usize,
21340 ksplit: usize,
21341 ) -> Result<(), Box<dyn std::error::Error>> {
21342 if t == 0 || !in_f.is_multiple_of(8) || x.len() < t * in_f || y.len() < t * out_f {
21343 return Err("matvec_bf16_v2 geometry".into());
21344 }
21345 let nb = mmv_block();
21346 let smem = (GEMV_V2_ROWS as u32) * nb * 4;
21347 let rows = out_f.div_ceil(GEMV_V2_ROWS) as u32;
21348 let (ini, outi) = (in_f as i32, out_f as i32);
21349 if ksplit <= 1 {
21350 let f = self.func("matvec_bf16_v2");
21351 let cfg = LaunchConfig {
21352 grid_dim: (rows, t as u32, 1),
21353 block_dim: (nb, 1, 1),
21354 shared_mem_bytes: smem,
21355 };
21356 let __s_b = self.gpu.stream();
21357 let mut b = __s_b.launch_builder(&f);
21358 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21359 unsafe {
21360 b.launch(cfg)?;
21361 }
21362 return Ok(());
21363 }
21364 let chunk = in_f.div_ceil(8).div_ceil(ksplit) * 8;
21369 let eff = in_f.div_ceil(chunk).max(1);
21370 let mut part = self.alloc_uninit::<f32>(eff * t * out_f)?;
21371 let ks = eff as i32;
21372 {
21373 let f = self.func("matvec_bf16_v2_sk");
21374 let cfg = LaunchConfig {
21375 grid_dim: (rows, t as u32, eff as u32),
21376 block_dim: (nb, 1, 1),
21377 shared_mem_bytes: smem,
21378 };
21379 let __s_b = self.gpu.stream();
21380 let mut b = __s_b.launch_builder(&f);
21381 b.arg(w).arg(x).arg(&mut part).arg(&ini).arg(&outi).arg(&ks);
21382 unsafe {
21383 b.launch(cfg)?;
21384 }
21385 }
21386 let f = self.func("matvec_bf16_v2_sk_combine");
21387 let n = (out_f * t) as u32;
21388 let cfg = LaunchConfig {
21389 grid_dim: (n.div_ceil(256), 1, 1),
21390 block_dim: (256, 1, 1),
21391 shared_mem_bytes: 0,
21392 };
21393 let ti = t as i32;
21394 let __s_b = self.gpu.stream();
21395 let mut b = __s_b.launch_builder(&f);
21396 b.arg(&part).arg(&mut *y).arg(&outi).arg(&ti).arg(&ks);
21397 unsafe {
21398 b.launch(cfg)?;
21399 }
21400 Ok(())
21401 }
21402
21403 pub fn gemv_v3_fits(&self) -> bool {
21407 gemv_v3_fits()
21408 }
21409
21410 pub fn matvec_bf16_v3_raw(
21418 &self,
21419 w: &CudaSlice<u8>,
21420 x: &CudaSlice<f32>,
21421 y: &mut CudaSlice<f32>,
21422 in_f: usize,
21423 out_f: usize,
21424 t: usize,
21425 ) -> Result<(), Box<dyn std::error::Error>> {
21426 if t == 0 || !in_f.is_multiple_of(8) || x.len() < t * in_f || y.len() < t * out_f {
21427 return Err("matvec_bf16_v3 geometry".into());
21428 }
21429 let nb = mmv_block();
21430 let smem = gemv_v3_smem_bytes(nb as usize);
21431 if smem > 48 * 1024 {
21432 return Err("matvec_bf16_v3 smem over the 48 KB default cap".into());
21433 }
21434 let f = self.func("matvec_bf16_v3");
21435 let cfg = LaunchConfig {
21436 grid_dim: (out_f.div_ceil(GEMV_V2_ROWS) as u32, t as u32, 1),
21437 block_dim: (nb, 1, 1),
21438 shared_mem_bytes: smem as u32,
21439 };
21440 let (ini, outi) = (in_f as i32, out_f as i32);
21441 let __s_b = self.gpu.stream();
21442 let mut b = __s_b.launch_builder(&f);
21443 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
21444 unsafe {
21445 b.launch(cfg)?;
21446 }
21447 Ok(())
21448 }
21449
21450 #[allow(clippy::too_many_arguments)]
21457 pub fn qmatvec_q8_0_rp_v2_raw(
21458 &self,
21459 w: &CudaSlice<u8>,
21460 aq: &CudaSlice<i8>,
21461 ad: &CudaSlice<f32>,
21462 y: &mut CudaSlice<f32>,
21463 in_f: usize,
21464 out_f: usize,
21465 t: usize,
21466 ) -> Result<(), Box<dyn std::error::Error>> {
21467 self.qmatvec_q8_0_rp_v2_raw_arm(w, aq, ad, y, in_f, out_f, t, q8_row_ilp_on())
21468 }
21469
21470 #[allow(clippy::too_many_arguments)]
21473 pub fn qmatvec_q8_0_rp_v2_raw_arm(
21474 &self,
21475 w: &CudaSlice<u8>,
21476 aq: &CudaSlice<i8>,
21477 ad: &CudaSlice<f32>,
21478 y: &mut CudaSlice<f32>,
21479 in_f: usize,
21480 out_f: usize,
21481 t: usize,
21482 ilp: bool,
21483 ) -> Result<(), Box<dyn std::error::Error>> {
21484 if t == 0 || !in_f.is_multiple_of(32) || y.len() < t * out_f {
21485 return Err("qmatvec_q8_0_rp_v2 geometry".into());
21486 }
21487 let smem = Self::q8_v2_smem_bytes(in_f);
21488 if smem > 48 * 1024 {
21489 return Err("qmatvec_q8_0_rp_v2 smem over the 48 KB default cap".into());
21490 }
21491 if ilp {
21492 q8_row_ilp_note("qmatvec_q8_0_mmvq_rp_v2");
21493 }
21494 let f = self.func(if ilp {
21495 "qmatvec_q8_0_mmvq_rp_v2_ilp"
21496 } else {
21497 "qmatvec_q8_0_mmvq_rp_v2"
21498 });
21499 let cfg = LaunchConfig {
21500 grid_dim: ((out_f as u32).div_ceil(Q8_V2_ROWS), t as u32, 1),
21501 block_dim: (32, Q8_V2_ROWS, 1),
21502 shared_mem_bytes: smem as u32,
21503 };
21504 let (ini, outi, mi, rb) = (in_f as i32, out_f as i32, t as i32, 0i64);
21505 let __s_b = self.gpu.stream();
21506 let mut b = __s_b.launch_builder(&f);
21507 b.arg(w)
21508 .arg(aq)
21509 .arg(ad)
21510 .arg(&mut *y)
21511 .arg(&ini)
21512 .arg(&outi)
21513 .arg(&mi)
21514 .arg(&rb);
21515 unsafe {
21516 b.launch(cfg)?;
21517 }
21518 Ok(())
21519 }
21520
21521 #[allow(clippy::too_many_arguments)]
21527 pub fn qmatvec_q8_0_rows_tw_v2_raw(
21528 &self,
21529 w: &CudaSlice<u8>,
21530 aq: &CudaSlice<i8>,
21531 ad: &CudaSlice<f32>,
21532 y: &mut CudaSlice<f32>,
21533 in_f: usize,
21534 out_f: usize,
21535 t: usize,
21536 ) -> Result<(), Box<dyn std::error::Error>> {
21537 if t == 0 || t > 8 || !in_f.is_multiple_of(32) || y.len() < t * out_f {
21538 return Err("qmatvec_q8_0_rows_tw_v2 geometry".into());
21539 }
21540 let f = self.func("qmatvec_q8_0_rows_tw_v2");
21541 let cfg = LaunchConfig {
21542 grid_dim: ((out_f as u32).div_ceil(Q8_V2_ROWS), 1, 1),
21543 block_dim: (32, Q8_V2_ROWS, 1),
21544 shared_mem_bytes: 0,
21545 };
21546 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
21547 let __s_b = self.gpu.stream();
21548 let mut b = __s_b.launch_builder(&f);
21549 b.arg(w)
21550 .arg(aq)
21551 .arg(ad)
21552 .arg(&mut *y)
21553 .arg(&ini)
21554 .arg(&outi)
21555 .arg(&ti);
21556 unsafe {
21557 b.launch(cfg)?;
21558 }
21559 Ok(())
21560 }
21561
21562 #[allow(clippy::too_many_arguments)]
21567 pub fn qmatvec_q8_0_rows_tw_arm_raw(
21568 &self,
21569 w: &CudaSlice<u8>,
21570 aq: &CudaSlice<i8>,
21571 ad: &CudaSlice<f32>,
21572 y: &mut CudaSlice<f32>,
21573 in_f: usize,
21574 out_f: usize,
21575 t: usize,
21576 use_v2: bool,
21577 ) -> Result<(), Box<dyn std::error::Error>> {
21578 if use_v2 {
21579 return self.qmatvec_q8_0_rows_tw_v2_raw(w, aq, ad, y, in_f, out_f, t);
21580 }
21581 let f = self.func("qmatvec_q8_0_rows_tw");
21582 let cfg = LaunchConfig {
21583 grid_dim: ((out_f as u32).div_ceil(4), 1, 1),
21584 block_dim: (32, 4, 1),
21585 shared_mem_bytes: 0,
21586 };
21587 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
21588 let __s_b = self.gpu.stream();
21589 let mut b = __s_b.launch_builder(&f);
21590 b.arg(w)
21591 .arg(aq)
21592 .arg(ad)
21593 .arg(&mut *y)
21594 .arg(&ini)
21595 .arg(&outi)
21596 .arg(&ti);
21597 unsafe {
21598 b.launch(cfg)?;
21599 }
21600 Ok(())
21601 }
21602
21603 #[allow(clippy::too_many_arguments)]
21618 #[allow(clippy::map_entry)] pub fn kda_proj_fused6_q8rp_raw(
21620 &self,
21621 bq: &CudaSlice<u8>,
21622 bk: &CudaSlice<u8>,
21623 bv: &CudaSlice<u8>,
21624 wfa: &CudaSlice<f32>,
21625 wga: &CudaSlice<f32>,
21626 wb: &CudaSlice<f32>,
21627 x: &CudaSlice<f32>,
21628 outs: &mut [CudaSlice<f32>; 6],
21629 in_f: usize,
21630 dims: [usize; 6],
21631 t: usize,
21632 ) -> Result<(), Box<dyn std::error::Error>> {
21633 self.kda_proj_fused6_q8rp_raw_arm(
21634 bq,
21635 bk,
21636 bv,
21637 wfa,
21638 wga,
21639 wb,
21640 x,
21641 outs,
21642 in_f,
21643 dims,
21644 t,
21645 q8_row_ilp_on(),
21646 )
21647 }
21648
21649 #[allow(clippy::too_many_arguments)]
21652 pub fn kda_proj_fused6_q8rp_raw_arm(
21653 &self,
21654 bq: &CudaSlice<u8>,
21655 bk: &CudaSlice<u8>,
21656 bv: &CudaSlice<u8>,
21657 wfa: &CudaSlice<f32>,
21658 wga: &CudaSlice<f32>,
21659 wb: &CudaSlice<f32>,
21660 x: &CudaSlice<f32>,
21661 outs: &mut [CudaSlice<f32>; 6],
21662 in_f: usize,
21663 dims: [usize; 6],
21664 t: usize,
21665 ilp: bool,
21666 ) -> Result<(), Box<dyn std::error::Error>> {
21667 self.kda_proj_fused6_q8rp_raw_pre(
21668 bq, bk, bv, wfa, wga, wb, x, outs, in_f, dims, t, ilp, None,
21669 )
21670 }
21671
21672 #[allow(clippy::too_many_arguments)]
21679 pub fn kda_proj_fused6_q8rp_raw_pre(
21680 &self,
21681 bq: &CudaSlice<u8>,
21682 bk: &CudaSlice<u8>,
21683 bv: &CudaSlice<u8>,
21684 wfa: &CudaSlice<f32>,
21685 wga: &CudaSlice<f32>,
21686 wb: &CudaSlice<f32>,
21687 x: &CudaSlice<f32>,
21688 outs: &mut [CudaSlice<f32>; 6],
21689 in_f: usize,
21690 dims: [usize; 6],
21691 t: usize,
21692 ilp: bool,
21693 pre_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
21694 ) -> Result<(), Box<dyn std::error::Error>> {
21695 use cudarc::driver::DevicePtr;
21696 if t == 0 || t > 32 || !in_f.is_multiple_of(32) || x.len() < t * in_f {
21697 return Err("kda_proj_fused6_q8rp geometry".into());
21698 }
21699 let smem = Self::q8_v2_smem_bytes(in_f);
21700 if smem > 48 * 1024 {
21701 return Err("kda_proj_fused6_q8rp smem over the 48 KB default cap".into());
21702 }
21703 for (i, (o, want)) in outs.iter().zip(dims).enumerate() {
21704 if o.len() < t * want {
21705 return Err(format!("kda_proj_fused6_q8rp: output {i} too small").into());
21706 }
21707 }
21708 let nblk = in_f / 32;
21709 let keys: Vec<(u64, u32, u32)> = {
21711 let s = self.gpu.stream();
21712 [(bq, dims[0]), (bk, dims[1]), (bv, dims[2])]
21713 .iter()
21714 .map(|(d, out)| {
21715 let (p, _g) = d.device_ptr(&s);
21716 (p, in_f as u32, *out as u32)
21717 })
21718 .collect()
21719 };
21720 {
21721 let mut mirrors = self
21722 .w8_mirrors
21723 .lock()
21724 .map_err(|_| "w8 mirror map is poisoned")?;
21725 for ((d, out), key) in [(bq, dims[0]), (bk, dims[1]), (bv, dims[2])]
21726 .iter()
21727 .zip(&keys)
21728 {
21729 if !mirrors.contains_key(key) {
21730 let mut interleaved = self.alloc_u8_uninit(out * Self::q8_0_row_bytes(in_f))?;
21731 self.encode_q8_0_from_bf16(d, &mut interleaved, in_f, *out)?;
21732 let planar = self.build_q8_rp4_raw(&interleaved, in_f, *out)?;
21733 mirrors.insert(*key, planar);
21734 }
21735 }
21736 }
21737 let akey = in_f * 64 + t.min(32);
21740 if pre_q8.is_none() {
21741 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21742 if let std::collections::hash_map::Entry::Vacant(slot) = act.entry(akey) {
21743 let aq = self.alloc_i8_uninit(32 * in_f)?;
21744 let ad = self.alloc_uninit::<f32>(32 * nblk)?;
21745 slot.insert((aq, ad));
21746 }
21747 let (aq, ad) = act.get_mut(&akey).expect("just inserted");
21748 self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
21749 }
21750 let mirrors = self
21751 .w8_mirrors
21752 .lock()
21753 .map_err(|_| "w8 mirror map is poisoned")?;
21754 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
21755 let (aq, ad): (&CudaSlice<i8>, &CudaSlice<f32>) = match pre_q8 {
21756 Some((aq, ad)) => {
21757 if aq.len() < t * in_f || ad.len() < t * nblk {
21758 return Err("kda_proj_fused6_q8rp: pre-quantized activation too small".into());
21759 }
21760 (aq, ad)
21761 }
21762 None => {
21763 let (aq, ad) = act.get(&akey).expect("built above");
21764 (aq, ad)
21765 }
21766 };
21767 let m0 = mirrors.get(&keys[0]).expect("built above");
21768 let m1 = mirrors.get(&keys[1]).expect("built above");
21769 let m2 = mirrors.get(&keys[2]).expect("built above");
21770 let blocks: usize = dims.iter().map(|d| d.div_ceil(Q8_V2_ROWS as usize)).sum();
21771 if ilp {
21772 q8_row_ilp_note("qmatvec_kda6_q8f32_rp_v2");
21773 }
21774 let f = self.func(if ilp {
21775 "qmatvec_kda6_q8f32_rp_v2_ilp"
21776 } else {
21777 "qmatvec_kda6_q8f32_rp_v2"
21778 });
21779 let cfg = LaunchConfig {
21780 grid_dim: (blocks as u32, t as u32, 1),
21781 block_dim: (32, Q8_V2_ROWS, 1),
21782 shared_mem_bytes: smem as u32,
21783 };
21784 let inf = in_f as i32;
21785 let d = dims.map(|v| v as i32);
21786 let mi = t as i32;
21787 let [o0, o1, o2, o3, o4, o5] = outs;
21788 let stream = self.gpu.stream();
21789 let mut b = stream.launch_builder(&f);
21790 b.arg(m0)
21791 .arg(m1)
21792 .arg(m2)
21793 .arg(wfa)
21794 .arg(wga)
21795 .arg(wb)
21796 .arg(aq)
21797 .arg(ad)
21798 .arg(x)
21799 .arg(&mut *o0)
21800 .arg(&mut *o1)
21801 .arg(&mut *o2)
21802 .arg(&mut *o3)
21803 .arg(&mut *o4)
21804 .arg(&mut *o5)
21805 .arg(&inf)
21806 .arg(&d[0])
21807 .arg(&d[1])
21808 .arg(&d[2])
21809 .arg(&d[3])
21810 .arg(&d[4])
21811 .arg(&d[5])
21812 .arg(&mi);
21813 unsafe {
21814 b.launch(cfg)?;
21815 }
21816 Ok(())
21817 }
21818
21819 #[allow(clippy::too_many_arguments)]
21824 pub fn moe_gate_up_preclamp8_q8_v2(
21825 &self,
21826 gp: WPtr8,
21827 up: WPtr8,
21828 aq: &CudaSlice<i8>,
21829 ad: &CudaSlice<f32>,
21830 gs: F32x8,
21831 us: F32x8,
21832 limit: f32,
21833 in_f: usize,
21834 n_ff: usize,
21835 n_used: usize,
21836 qt_g: i32,
21837 qt_u: i32,
21838 rb_g: usize,
21839 rb_u: usize,
21840 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21841 debug_assert!(
21842 limit > 1e-6,
21843 "moe_gate_up_preclamp8_q8_v2 needs a live limit; use moe_gate_up_silu8_q8"
21844 );
21845 const ROWS: u32 = 8;
21846 let f = self.func("moe_gate_up_preclamp8_q8_v2");
21847 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
21848 let cfg = LaunchConfig {
21849 grid_dim: ((n_ff as u32).div_ceil(ROWS), n_used as u32, 1),
21850 block_dim: (32, ROWS, 1),
21851 shared_mem_bytes: 0,
21852 };
21853 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
21854 let __s_b = self.gpu.stream();
21855 let mut b = __s_b.launch_builder(&f);
21856 b.arg(&gp)
21857 .arg(&up)
21858 .arg(aq)
21859 .arg(ad)
21860 .arg(&gs)
21861 .arg(&us)
21862 .arg(&limit)
21863 .arg(&mut act)
21864 .arg(&inf)
21865 .arg(&nff)
21866 .arg(&qt_g)
21867 .arg(&qt_u)
21868 .arg(&rbg)
21869 .arg(&rbu);
21870 unsafe {
21871 b.launch(cfg)?;
21872 }
21873 Ok(act)
21874 }
21875
21876 #[allow(clippy::too_many_arguments)]
21882 pub fn moe_down8_fma_q8_v2(
21883 &self,
21884 dp: WPtr8,
21885 w: F32x8,
21886 aq2: &CudaSlice<i8>,
21887 ad2: &CudaSlice<f32>,
21888 dst: &mut cudarc::driver::CudaViewMut<f32>,
21889 in_f: usize,
21890 out_f: usize,
21891 n_used: usize,
21892 qt: i32,
21893 rb: usize,
21894 ) -> Result<(), Box<dyn std::error::Error>> {
21895 let f = self.func("moe_down8_fma_q8_v2");
21896 let cfg = LaunchConfig {
21897 grid_dim: (out_f as u32, 1, 1),
21898 block_dim: (32, 8, 1),
21899 shared_mem_bytes: 0,
21900 };
21901 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
21902 let __s_b = self.gpu.stream();
21903 let mut b = __s_b.launch_builder(&f);
21904 b.arg(&dp)
21905 .arg(&w)
21906 .arg(aq2)
21907 .arg(ad2)
21908 .arg(dst)
21909 .arg(&inf)
21910 .arg(&outf)
21911 .arg(&nu)
21912 .arg(&qt)
21913 .arg(&rbi);
21914 unsafe {
21915 b.launch(cfg)?;
21916 }
21917 Ok(())
21918 }
21919
21920 pub fn batched_supports(&self, qtype: i32) -> bool {
21924 matches!(
21925 qtype,
21926 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
21927 )
21928 }
21929
21930 pub fn iq_fast_enabled() -> bool {
21938 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21939 *ON.get_or_init(|| {
21940 std::env::var("MEMRA_IQ_FAST")
21941 .map(|v| v != "0")
21942 .unwrap_or(true)
21943 })
21944 }
21945
21946 pub fn b8_enabled() -> bool {
21949 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21950 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
21951 }
21952
21953 pub fn batched_mcols(m: usize) -> usize {
21955 if m == 2 {
21956 2
21957 } else if m <= 4 {
21958 4
21959 } else if m <= 8 {
21960 8
21961 } else {
21962 16
21963 }
21964 }
21965
21966 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
21971 Some(match (qtype, mcols) {
21972 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
21973 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
21974 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
21975 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
21981 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
21982 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
21983 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
21984 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
21987 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
21988 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
21989 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
21990 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
21993 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
21994 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
21995 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
21996 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
21997 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
21998 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
21999 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
22000 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
22004 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
22005 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
22006 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
22007 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
22011 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
22012 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
22013 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
22014 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
22015 _ => return None,
22016 })
22017 }
22018
22019 pub fn sm_count(&self) -> i32 {
22054 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
22055 *SMS.get_or_init(|| {
22056 use cudarc::driver::sys::CUdevice_attribute_enum as A;
22057 self.gpu
22058 .ctx
22059 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
22060 .unwrap_or(82)
22061 })
22062 }
22063
22064 #[allow(clippy::too_many_arguments)]
22065 #[allow(clippy::if_same_then_else)] pub fn batched_variant(
22068 &self,
22069 _m: usize,
22070 in_f: usize,
22071 out_f: usize,
22072 qtype: i32,
22073 row_bytes: usize,
22074 mcols: usize,
22075 rp: bool,
22076 ) -> &'static str {
22077 if qtype == QT_Q8_0 {
22082 return if rp { "rp" } else { "base" };
22083 }
22084 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
22085 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
22086 Ok("base") => "base",
22087 Ok("pf") => "pf",
22088 Ok("r2") => "r2",
22089 Ok("r2w8") => "r2w8",
22090 Ok("pfr2") => "pfr2",
22091 Ok("ca") => "ca",
22092 Ok("car2") => "car2",
22093 Ok("rp") => "rp",
22096 Ok("rpr2") => "rpr2",
22097 Ok("rpr2w8") => "rpr2w8",
22098 Ok("rpca") => "rpca",
22101 Ok("rpcar2") => "rpcar2",
22102 Ok("rpsc") => "rpsc",
22109 Ok("rpms") => "rpms",
22110 Ok("rpmsc") => "rpmsc",
22111 Ok("rpks") => "rpks",
22112 Ok("rpksc") => "rpksc",
22113 _ => "auto",
22114 });
22115 let ca_ok = qtype == QT_NVFP4 && row_bytes.is_multiple_of(16) && in_f.is_multiple_of(1024);
22119 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22124 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
22125 let sc_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(256) && (in_f / 64 <= 272);
22126 let ks_ok = ks_on && qtype == QT_NVFP4 && in_f.is_multiple_of(512) && (in_f / 64 <= 272);
22127 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
22128 let sms = *SMS.get_or_init(|| {
22129 use cudarc::driver::sys::CUdevice_attribute_enum as A;
22130 self.gpu
22131 .ctx
22132 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
22133 .unwrap_or(82)
22134 });
22135 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
22155 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
22158 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
22159 Ok("base") => "base",
22160 Ok("r2") => "r2",
22161 Ok("r2w8") => "r2w8",
22162 _ => "auto",
22163 });
22164 let variant: &'static str = if qtype == QT_Q4_0 {
22165 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
22169 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
22170 Ok("base") => "base",
22176 Ok("r2") => "r2",
22177 Ok("ms") => "ms",
22178 Ok("sm") => "sm",
22179 Ok("la") => "la",
22180 _ => "auto",
22181 });
22182 let v = if q40 != "auto" {
22183 q40
22184 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
22185 "r2"
22186 } else {
22187 "base"
22188 };
22189 if rp {
22194 match v {
22195 "ms" => "r2ms_rp",
22196 "sm" => "r2sm_rp",
22197 "la" => "r2la_rp",
22198 "r2" => "r2_rp",
22199 _ => "rp",
22200 }
22201 } else if matches!(v, "ms" | "sm" | "la") {
22202 "r2"
22203 } else {
22204 v
22205 }
22206 } else if qtype != QT_NVFP4 && !kq_r2 {
22207 "base"
22208 } else if kq_r2 && rp {
22209 "rp"
22213 } else if kq_r2 {
22214 if kq_bv != "auto" {
22217 if kq_bv == "r2w8" && mcols != 4 {
22218 "r2"
22219 } else {
22220 kq_bv
22221 }
22222 } else if bv != "auto" {
22223 match bv {
22224 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
22225 "r2w8" | "rpr2w8" => {
22226 if mcols != 4 {
22227 "r2"
22228 } else {
22229 "r2w8"
22230 }
22231 }
22232 _ => "base", }
22234 } else {
22235 #[allow(clippy::manual_div_ceil)]
22236 let blocks = (out_f + 7) / 8;
22238 let waves = blocks as f64 / (7 * sms as usize) as f64;
22239 let filled = blocks >= 4 * sms as usize;
22240 let use_r2 = if qtype == QT_Q4_K {
22241 filled
22242 } else {
22243 waves >= 2.0
22244 };
22245 if use_r2 { "r2" } else { "base" }
22246 }
22247 } else if bv != "auto" {
22248 let v = if bv == "r2w8" && mcols == 2 {
22253 "r2"
22254 } else if bv == "ca" && (!ca_ok || mcols == 8) {
22255 "pf"
22256 } else if bv == "car2" && (!ca_ok || mcols == 8) {
22257 "r2"
22258 } else if bv == "pfr2" && mcols == 8 {
22259 "r2"
22260 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
22261 "rpr2"
22262 }
22263 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
22265 if mcols == 8 { "rpr2w8" } else { "rpr2" }
22266 } else if bv == "rpcar2" && mcols == 2 {
22267 "rpca"
22268 }
22269 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
22272 "rpr2"
22273 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
22274 "rpr2"
22275 } else {
22276 bv
22277 };
22278 if rp {
22279 match v {
22280 "base" | "pf" | "ca" | "rp" => "rp",
22281 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
22282 "r2w8" | "rpr2w8" => {
22283 if mcols == 2 {
22284 "rpr2"
22285 } else {
22286 "rpr2w8"
22287 }
22288 }
22289 other => other, }
22291 } else {
22292 v
22293 }
22294 } else if mcols == 8 {
22295 if rp {
22306 if sc_ok { "rpsc" } else { "rpr2w8" }
22307 } else {
22308 "r2w8"
22309 }
22310 } else if mcols >= 4 {
22311 #[allow(clippy::manual_div_ceil)]
22315 let blocks = (out_f + 7) / 8;
22317 let r7 = 7 * sms as usize;
22318 let r8 = 8 * sms as usize;
22319 let waves = blocks as f64 / r7 as f64;
22320 let filled = blocks >= 4 * sms as usize;
22321 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
22325 if rp { "rpr2w8" } else { "r2w8" }
22329 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
22330 if rp { "rpr2" } else { "r2" }
22333 } else {
22334 if rp { "rp" } else { "pf" }
22338 }
22339 } else if in_f >= 6144 {
22340 if rp { "rpr2" } else { "r2" }
22344 } else if rp {
22345 #[allow(clippy::manual_div_ceil)]
22350 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
22352 if sc_ok && (0.9..=1.1).contains(&waves) {
22353 "rpsc"
22354 } else {
22355 "rp"
22356 }
22357 } else {
22358 "base"
22359 };
22360 variant
22361 }
22362
22363 #[allow(clippy::too_many_arguments)]
22364 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_mmvq_batched(
22367 &self,
22368 bytes: &CudaSlice<u8>,
22369 aq: &CudaSlice<i8>,
22370 ad: &CudaSlice<f32>,
22371 m: usize,
22372 in_f: usize,
22373 out_f: usize,
22374 qtype: i32,
22375 row_bytes: usize,
22376 mcols: usize,
22377 scale: f32,
22378 rp: bool,
22379 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22380 const ROWS_PER_BLOCK: u32 = 4;
22381 let forced: Option<&'static str> = {
22386 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
22387 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
22388 .as_deref()
22389 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
22390 };
22391 let variant = match forced {
22392 Some(v) if !rp || v.contains("rp") => v,
22393 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
22394 };
22395 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
22396 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
22397 })?;
22398 let variant = if mcols == 16 {
22402 if rp { "rp" } else { "base" }
22403 } else {
22404 variant
22405 };
22406 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22413 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
22414 if b567
22415 && qtype == QT_NVFP4
22416 && rp
22417 && mcols == 8
22418 && (5..=7).contains(&m)
22419 && matches!(variant, "rpsc" | "rpr2w8")
22420 {
22421 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
22422 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
22424 let cfg = LaunchConfig {
22425 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
22426 block_dim: (32, ROWS_PER_BLOCK, 1),
22427 shared_mem_bytes: 0,
22428 };
22429 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22430 let __s_b = self.gpu.stream();
22431 let mut b = __s_b.launch_builder(&f);
22432 b.arg(bytes)
22433 .arg(aq)
22434 .arg(ad)
22435 .arg(&mut y)
22436 .arg(&inf)
22437 .arg(&outf)
22438 .arg(&mi)
22439 .arg(&rb);
22440 unsafe {
22441 b.launch(cfg)?;
22442 }
22443 if scale != 1.0 {
22444 self.scale_inplace(&mut y, scale, m * out_f)?;
22445 }
22446 return Ok(y);
22447 }
22448 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
22449 "base" => (base_name.into(), ROWS_PER_BLOCK),
22450 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
22451 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
22452 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
22453 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
22457 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
22458 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
22459 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
22460 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
22461 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
22462 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
22463 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
22465 debug_assert!(
22466 !rp || name.contains("_rp"),
22467 "rp weight dispatched to a GGUF-layout kernel"
22468 );
22469 let f = self.func(&name);
22470 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
22471 let smem = if name.contains("_r2sm_rp") {
22473 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
22474 } else {
22475 0
22476 };
22477 let cfg = LaunchConfig {
22478 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
22479 block_dim: (32, ROWS_PER_BLOCK, 1),
22480 shared_mem_bytes: smem,
22481 };
22482 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22483 let __s_b = self.gpu.stream();
22484 let mut b = __s_b.launch_builder(&f);
22485 b.arg(bytes)
22486 .arg(aq)
22487 .arg(ad)
22488 .arg(&mut y)
22489 .arg(&inf)
22490 .arg(&outf)
22491 .arg(&mi)
22492 .arg(&rb);
22493 unsafe {
22494 b.launch(cfg)?;
22495 }
22496 if scale != 1.0 {
22497 self.scale_inplace(&mut y, scale, m * out_f)?;
22498 }
22499 Ok(y)
22500 }
22501
22502 #[allow(clippy::too_many_arguments)] pub fn qmatvec_batched_raw(
22507 &self,
22508 bytes: &CudaSlice<u8>,
22509 x: &CudaSlice<f32>,
22510 m: usize,
22511 in_f: usize,
22512 out_f: usize,
22513 qtype: i32,
22514 row_bytes: usize,
22515 mcols: usize,
22516 rp: bool,
22517 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22518 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
22519 self.qmatvec_mmvq_batched(
22520 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
22521 )
22522 }
22523
22524 #[allow(clippy::too_many_arguments)] pub fn qmatvec_nvfp4_batched_raw(
22527 &self,
22528 bytes: &CudaSlice<u8>,
22529 x: &CudaSlice<f32>,
22530 m: usize,
22531 in_f: usize,
22532 out_f: usize,
22533 row_bytes: usize,
22534 mcols: usize,
22535 rp: bool,
22536 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22537 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
22538 }
22539
22540 fn try_fp4_gemm(
22544 &self,
22545 w: &crate::model::GpuTensor,
22546 x: &CudaSlice<f32>,
22547 m: usize,
22548 in_f: usize,
22549 out_f: usize,
22550 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22551 use crate::model::GpuTensor;
22552 if cfg!(memra_portable_cuda) {
22553 return Ok(None);
22554 }
22555 if std::env::var("MEMRA_FP4").is_ok() {
22564 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
22565 assert!(
22566 konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a"),
22567 "MEMRA_FP4 forces the native mxf4 block-scale GEMM (qmatvec_gemm_nvfp4_fp4), \
22568 which only the sm_120a fatbin contains — this is an sm_{} build. Unset \
22569 MEMRA_FP4; the W4A8 int8 path is the correct default for NVFP4 weights.",
22570 env!("MEMRA_BUILT_CUDA_ARCH")
22571 );
22572 }
22573 if std::env::var("MEMRA_FP4").is_err() {
22574 return Ok(None);
22575 }
22576 #[cfg(memra_cutlass)]
22585 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
22586 if let GpuTensor::Quant {
22587 bytes,
22588 qtype,
22589 scale,
22590 row_bytes,
22591 cutlass,
22592 ..
22593 } = w
22594 {
22595 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
22596 if let Some(cw) = cutlass {
22597 let y = self.cutlass_fp4_gemm(
22599 &cw.b_packed,
22600 &cw.sfb_swizzled,
22601 x,
22602 *scale,
22603 m,
22604 out_f,
22605 in_f,
22606 )?;
22607 return Ok(Some(y));
22608 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
22609 let (b_packed, sfb_sw) =
22614 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
22615 let y =
22616 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
22617 return Ok(Some(y));
22618 }
22619 }
22620 }
22621 }
22622 if let GpuTensor::Quant {
22623 bytes,
22624 qtype,
22625 row_bytes,
22626 scale,
22627 rp,
22628 ..
22629 } = w
22630 {
22631 if *qtype == QT_NVFP4 && in_f.is_multiple_of(64) && !*rp {
22634 let y =
22635 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
22636 return Ok(Some(y));
22637 }
22638 }
22639 Ok(None)
22640 }
22641
22642 #[allow(clippy::too_many_arguments)] pub fn rms_norm_f16out(
22647 &self,
22648 x: &CudaSlice<f32>,
22649 w: &CudaSlice<f32>,
22650 dst: &mut CudaSlice<f32>,
22651 dst16: &mut CudaSlice<u8>,
22652 ncols: usize,
22653 nrows: usize,
22654 eps: f32,
22655 ) -> Result<(), Box<dyn std::error::Error>> {
22656 let f = self.func("rms_norm_f16out_f32");
22657 let cfg = LaunchConfig {
22658 grid_dim: (nrows as u32, 1, 1),
22659 block_dim: (rms_block(), 1, 1),
22660 shared_mem_bytes: 0,
22661 };
22662 let (nc, e) = (ncols as i32, eps);
22663 let __s_b = self.gpu.stream();
22664 let mut b = __s_b.launch_builder(&f);
22665 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
22666 unsafe {
22667 b.launch(cfg)?;
22668 }
22669 Ok(())
22670 }
22671
22672 #[allow(clippy::too_many_arguments)]
22675 pub fn add_rms_norm_f16out(
22676 &self,
22677 a: &CudaSlice<f32>,
22678 b: &CudaSlice<f32>,
22679 w: &CudaSlice<f32>,
22680 res: &mut CudaSlice<f32>,
22681 dst: &mut CudaSlice<f32>,
22682 dst16: &mut CudaSlice<u8>,
22683 ncols: usize,
22684 nrows: usize,
22685 eps: f32,
22686 ) -> Result<(), Box<dyn std::error::Error>> {
22687 let f = self.func("add_rms_norm_f16out_f32");
22688 let cfg = LaunchConfig {
22689 grid_dim: (nrows as u32, 1, 1),
22690 block_dim: (rms_block(), 1, 1),
22691 shared_mem_bytes: 0,
22692 };
22693 let (nc, e) = (ncols as i32, eps);
22694 let __s_lb = self.gpu.stream();
22695 let mut lb = __s_lb.launch_builder(&f);
22696 lb.arg(a)
22697 .arg(b)
22698 .arg(w)
22699 .arg(res)
22700 .arg(dst)
22701 .arg(dst16)
22702 .arg(&nc)
22703 .arg(&e);
22704 unsafe {
22705 lb.launch(cfg)?;
22706 }
22707 Ok(())
22708 }
22709
22710 pub fn matmul_group_xh(
22713 &self,
22714 ws: &[&crate::model::GpuTensor],
22715 x: &CudaSlice<f32>,
22716 xh: &CudaSlice<u8>,
22717 m: usize,
22718 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22719 let mut out = Vec::with_capacity(ws.len());
22720 let in_f = ws[0].in_features();
22721 for w in ws {
22722 if w.in_features() == in_f
22723 && m >= 16
22724 && !self.verify_exact_on()
22725 && let Some(y) = self.try_f16_gemm_pre(w, xh, m)?
22726 {
22727 out.push(y);
22728 continue;
22729 }
22730 out.push(self.matmul(w, x, m)?);
22731 }
22732 Ok(out)
22733 }
22734
22735 pub fn gdn_pad_mask(
22738 &self,
22739 beta: &mut CudaSlice<f32>,
22740 g_log: &mut CudaSlice<f32>,
22741 len_d: &CudaSlice<i32>,
22742 h: usize,
22743 t: usize,
22744 ) -> Result<(), Box<dyn std::error::Error>> {
22745 let f = self.func("gdn_pad_mask_f32");
22746 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
22747 let (hi, ti) = (h as i32, t as i32);
22748 let __s_b = self.gpu.stream();
22749 let mut b = __s_b.launch_builder(&f);
22750 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
22751 unsafe {
22752 b.launch(cfg)?;
22753 }
22754 Ok(())
22755 }
22756
22757 pub fn row_gather_dev(
22760 &self,
22761 src: &CudaSlice<f32>,
22762 dst: &mut CudaSlice<f32>,
22763 len_d: &CudaSlice<i32>,
22764 ncols: usize,
22765 ) -> Result<(), Box<dyn std::error::Error>> {
22766 let f = self.func("row_gather_dev_f32");
22767 let cfg = LaunchConfig::for_num_elems(ncols as u32);
22768 let nc = ncols as i32;
22769 let __s_b = self.gpu.stream();
22770 let mut b = __s_b.launch_builder(&f);
22771 b.arg(src).arg(dst).arg(len_d).arg(&nc);
22772 unsafe {
22773 b.launch(cfg)?;
22774 }
22775 Ok(())
22776 }
22777
22778 pub fn matmul_group(
22785 &self,
22786 ws: &[&crate::model::GpuTensor],
22787 x: &CudaSlice<f32>,
22788 m: usize,
22789 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22790 use crate::model::GpuTensor;
22791 let mut out = Vec::with_capacity(ws.len());
22792 let any_mirror = ws
22793 .iter()
22794 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
22795 if m >= 16 && any_mirror && !self.verify_exact_on() {
22796 let in_f = ws[0].in_features();
22797 let xh = self.f16_act(x, m * in_f, in_f)?;
22798 for w in ws {
22799 if w.in_features() == in_f
22800 && let Some(y) = self.try_f16_gemm_pre(w, &xh, m)?
22801 {
22802 out.push(y);
22803 continue;
22804 }
22805 out.push(self.matmul(w, x, m)?);
22806 }
22807 return Ok(out);
22808 }
22809 for w in ws {
22810 out.push(self.matmul(w, x, m)?);
22811 }
22812 Ok(out)
22813 }
22814
22815 pub fn matmul_group_multi(
22822 &self,
22823 ws: &[&crate::model::GpuTensor],
22824 xs: &[&CudaSlice<f32>],
22825 ms: &[usize],
22826 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
22827 assert_eq!(xs.len(), ms.len());
22828 let in_f = ws[0].in_features();
22829 let total: usize = ms.iter().sum();
22830 let mut xcat = self.uninit(total * in_f)?;
22831 let mut off = 0usize;
22832 for (x, &m) in xs.iter().zip(ms) {
22833 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
22834 off += m;
22835 }
22836 let ys = self.matmul_group(ws, &xcat, total)?;
22837 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
22838 for (w, y) in ws.iter().zip(ys) {
22839 let out_f = w.out_features();
22840 let mut off = 0usize;
22841 for (s, &m) in ms.iter().enumerate() {
22842 let mut ys_s = self.uninit(m * out_f)?;
22843 let src = y.slice(off * out_f..(off + m) * out_f);
22844 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
22845 out[s].push(ys_s);
22846 off += m;
22847 }
22848 }
22849 Ok(out)
22850 }
22851
22852 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
22862 use crate::model::GpuTensor;
22863 if !legacy_quant_gemm_allowed(
22864 cfg!(memra_portable_cuda),
22865 cfg!(memra_hopper_mma),
22866 std::env::var_os("MEMRA_NO_GEMM").is_some(),
22867 ) {
22868 return false;
22869 }
22870 match w {
22871 GpuTensor::Quant { qtype, .. } => {
22872 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
22873 || (*qtype == QT_NVFP4 && w.in_features().is_multiple_of(64))
22874 }
22875 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
22876 }
22877 }
22878
22879 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_gemm(
22887 &self,
22888 w: &crate::model::GpuTensor,
22889 aq: &CudaSlice<i8>,
22890 ad: &CudaSlice<f32>,
22891 m: usize,
22892 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22893 use crate::model::GpuTensor;
22894 let in_f = w.in_features();
22895 let out_f = w.out_features();
22896 let (bytes, qtype, row_bytes, scale, rp) = match w {
22897 GpuTensor::Quant {
22898 bytes,
22899 qtype,
22900 row_bytes,
22901 scale,
22902 rp,
22903 ..
22904 } => (bytes, *qtype, *row_bytes, *scale, *rp),
22905 _ => unreachable!("gemm_supports guaranteed Quant"),
22906 };
22907 if cfg!(memra_hopper_mma)
22913 && qtype == QT_Q8_0
22914 && out_f.is_multiple_of(64)
22915 && wgmma_gemm_enabled()
22916 && let GpuTensor::Quant { rp4: Some(m4), .. } = w
22917 {
22918 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
22919 if scale != 1.0 {
22920 self.scale_inplace(&mut y, scale, m * out_f)?;
22921 }
22922 return Ok(y);
22923 }
22924 let name = match qtype {
22925 QT_Q8_0 => "qmatvec_gemm_q8_0",
22926 QT_Q4_K => "qmatvec_gemm_q4_K",
22927 QT_Q4_0 => {
22928 if rp {
22929 "qmatvec_gemm_q4_0_rp"
22930 } else {
22931 "qmatvec_gemm_q4_0"
22932 }
22933 }
22934 QT_Q5_K => "qmatvec_gemm_q5_K",
22935 QT_Q6_K => "qmatvec_gemm_q6_K",
22936 QT_NVFP4 => {
22937 if rp {
22938 "qmatvec_gemm_nvfp4_rp"
22939 } else {
22940 "qmatvec_gemm_nvfp4"
22941 }
22942 }
22943 _ => unreachable!(),
22944 };
22945 let f = self.func(name);
22946 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);
22951 let k1_tile = if is_k1 {
22953 k1_launch_override().unwrap_or((128, 128, 8))
22954 } else {
22955 (128, 128, 8)
22956 };
22957 let (bm, bn): (u32, u32) = if is_k1 {
22958 (k1_tile.0, k1_tile.1)
22959 } else {
22960 (64, 256)
22961 };
22962 let warps: u32 = if is_k1 {
22963 k1_tile.2
22964 } else {
22965 match qtype {
22966 QT_NVFP4 => 8,
22967 _ => 4,
22968 }
22969 };
22970 let cfg = LaunchConfig {
22971 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
22972 block_dim: (32, warps, 1),
22973 shared_mem_bytes: 0,
22974 };
22975 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
22976 let __s_b = self.gpu.stream();
22977 let mut b = __s_b.launch_builder(&f);
22978 b.arg(bytes)
22979 .arg(aq)
22980 .arg(ad)
22981 .arg(&mut y)
22982 .arg(&inf)
22983 .arg(&outf)
22984 .arg(&mi)
22985 .arg(&rb);
22986 unsafe {
22987 b.launch(cfg)?;
22988 }
22989 if scale != 1.0 {
22990 self.scale_inplace(&mut y, scale, m * out_f)?;
22991 }
22992 Ok(y)
22993 }
22994
22995 #[allow(clippy::too_many_arguments)]
23000 #[allow(clippy::manual_div_ceil)] pub fn qmatvec_gemm_raw(
23003 &self,
23004 bytes: &CudaSlice<u8>,
23005 x: &CudaSlice<f32>,
23006 m: usize,
23007 in_f: usize,
23008 out_f: usize,
23009 qtype: i32,
23010 row_bytes: usize,
23011 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23012 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
23013 let name = match qtype {
23014 QT_Q8_0 => "qmatvec_gemm_q8_0",
23015 QT_Q4_K => "qmatvec_gemm_q4_K",
23016 QT_Q4_0 => "qmatvec_gemm_q4_0",
23017 QT_Q5_K => "qmatvec_gemm_q5_K",
23018 QT_Q6_K => "qmatvec_gemm_q6_K",
23019 QT_NVFP4 => "qmatvec_gemm_nvfp4",
23020 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
23021 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
23022 };
23023 let f = self.func(name);
23024 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);
23028 let k1_tile = if is_k1 {
23030 k1_launch_override().unwrap_or((128, 128, 8))
23031 } else {
23032 (128, 128, 8)
23033 };
23034 let (bm, bn): (u32, u32) = if is_k1 {
23035 (k1_tile.0, k1_tile.1)
23036 } else {
23037 (64, 256)
23038 };
23039 let warps: u32 = if is_k1 {
23040 k1_tile.2
23041 } else {
23042 match qtype {
23043 QT_NVFP4 | QT_NVFP4_RP => 8,
23044 _ => 4,
23045 }
23046 };
23047 let cfg = LaunchConfig {
23048 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
23049 block_dim: (32, warps, 1),
23050 shared_mem_bytes: 0,
23051 };
23052 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
23053 let __s_b = self.gpu.stream();
23054 let mut b = __s_b.launch_builder(&f);
23055 b.arg(bytes)
23056 .arg(&aq)
23057 .arg(&ad)
23058 .arg(&mut y)
23059 .arg(&inf)
23060 .arg(&outf)
23061 .arg(&mi)
23062 .arg(&rb);
23063 unsafe {
23064 b.launch(cfg)?;
23065 }
23066 Ok(y)
23067 }
23068
23069 pub fn qmatvec_gemm_q8_0_wgmma_raw(
23076 &self,
23077 rp4: &CudaSlice<u8>,
23078 aq: &CudaSlice<i8>,
23079 ad: &CudaSlice<f32>,
23080 m: usize,
23081 in_f: usize,
23082 out_f: usize,
23083 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23084 assert!(
23085 out_f.is_multiple_of(64) && in_f.is_multiple_of(32),
23086 "wgmma GEMM needs out_f%64==0, in_f%32==0"
23087 );
23088 let f = self.func("qmatvec_gemm_q8_0_wgmma");
23089 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
23091 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
23092 block_dim: (128, 1, 1),
23093 shared_mem_bytes: 0,
23094 };
23095 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
23096 let __s_b = self.gpu.stream();
23097 let mut b = __s_b.launch_builder(&f);
23098 b.arg(rp4)
23099 .arg(aq)
23100 .arg(ad)
23101 .arg(&mut y)
23102 .arg(&inf)
23103 .arg(&outf)
23104 .arg(&mi);
23105 unsafe {
23106 b.launch(cfg)?;
23107 }
23108 Ok(y)
23109 }
23110
23111 pub fn scale_inplace(
23113 &self,
23114 y: &mut CudaSlice<f32>,
23115 s: f32,
23116 n: usize,
23117 ) -> Result<(), Box<dyn std::error::Error>> {
23118 let f = self.func("scale_f32");
23119 let cfg = LaunchConfig::for_num_elems(n as u32);
23120 let (sf, ni) = (s, n as i32);
23121 let __s_b = self.gpu.stream();
23122 let mut b = __s_b.launch_builder(&f);
23123 b.arg(y).arg(&sf).arg(&ni);
23124 unsafe {
23125 b.launch(cfg)?;
23126 }
23127 Ok(())
23128 }
23129
23130 pub fn bf16_to_f32(
23135 &self,
23136 data: &cudarc::driver::CudaView<'_, u8>,
23137 n: usize,
23138 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23139 let mut out = self.alloc_uninit::<f32>(n)?;
23140 let f = self.func("bf16_to_f32");
23141 let cfg = LaunchConfig::for_num_elems(n as u32);
23142 let ni = n as i32;
23143 let __s_b = self.gpu.stream();
23144 let mut b = __s_b.launch_builder(&f);
23145 b.arg(data).arg(&mut out).arg(&ni);
23146 unsafe {
23147 b.launch(cfg)?;
23148 }
23149 Ok(out)
23150 }
23151
23152 #[allow(clippy::too_many_arguments)] fn linear_bf16_chunked(
23160 &self,
23161 x: &CudaSlice<f32>,
23162 data: &CudaSlice<u8>,
23163 m: usize,
23164 in_f: usize,
23165 out_f: usize,
23166 exact: bool,
23167 canonical_chunk_rows: Option<usize>,
23168 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23169 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23173 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23174 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23175 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
23176 let started = timing.then(std::time::Instant::now);
23177 let result =
23178 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
23179 if let Some(started) = started {
23180 use std::sync::atomic::Ordering;
23181 self.stream().synchronize()?;
23182 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
23183 + started.elapsed().as_nanos() as u64;
23184 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
23185 + (in_f * out_f * 2) as u64;
23186 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
23187 if calls.is_multiple_of(1024) {
23188 eprintln!(
23189 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
23190 weight_gb={:.2}",
23191 ns as f64 / 1.0e6,
23192 ns as f64 / calls as f64 / 1.0e3,
23193 wb as f64 / 1.0e9,
23194 );
23195 }
23196 }
23197 result
23198 }
23199
23200 pub(crate) fn bf16_mmv_on() -> bool {
23205 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23206 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
23207 }
23208
23209 fn matvec_bf16(
23212 &self,
23213 data: &CudaSlice<u8>,
23214 x: &CudaSlice<f32>,
23215 in_f: usize,
23216 out_f: usize,
23217 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23218 if data.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) {
23219 return Err(format!(
23220 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
23221 data.len(),
23222 x.len()
23223 )
23224 .into());
23225 }
23226 let mut y = self.alloc_uninit::<f32>(out_f)?;
23227 let f = self.func("matvec_bf16_f32acc");
23228 let cfg = LaunchConfig {
23229 grid_dim: (out_f as u32, 1, 1),
23230 block_dim: (mmv_block(), 1, 1),
23231 shared_mem_bytes: 0,
23232 };
23233 let ini = in_f as i32;
23234 let __s_bld = self.gpu.stream();
23235 let mut bld = __s_bld.launch_builder(&f);
23236 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
23237 unsafe {
23238 bld.launch(cfg)?;
23239 }
23240 Ok(y)
23241 }
23242
23243 #[allow(clippy::too_many_arguments)]
23247 #[allow(clippy::too_many_arguments)]
23252 #[allow(clippy::too_many_arguments)]
23257 pub fn qk_norm_rope_append_inc_dcw_rows(
23258 &self,
23259 q_raw_t: &CudaSlice<f32>,
23260 k_raw_t: &CudaSlice<f32>,
23261 v_raw_t: &CudaSlice<f32>,
23262 qw: &CudaSlice<f32>,
23263 kw: &CudaSlice<f32>,
23264 q_out_t: &mut CudaSlice<f32>,
23265 k_out_t: &mut CudaSlice<f32>,
23266 tab: &CudaSlice<u64>,
23267 pos_t: &CudaSlice<i32>,
23268 same_session: bool,
23269 t: usize,
23270 kv_dim_k: usize,
23271 kv_dim_v: usize,
23272 k_tok_bytes: usize,
23273 v_tok_bytes: usize,
23274 head_dim: usize,
23275 n_dims: usize,
23276 nh_q: usize,
23277 nh_k: usize,
23278 eps: f32,
23279 freq_base: f32,
23280 freq_scale: f32,
23281 ff: Option<&CudaSlice<f32>>,
23282 ) -> Result<(), Box<dyn std::error::Error>> {
23283 if head_dim != 128
23284 || kv_dim_v != kv_dim_k
23285 || kv_dim_k != nh_k * head_dim
23286 || t == 0
23287 || t > 32
23288 || tab.len() < t * 6
23289 || pos_t.len() < t
23290 || q_raw_t.len() < t * nh_q * head_dim
23291 || k_raw_t.len() < t * nh_k * head_dim
23292 || v_raw_t.len() < t * kv_dim_v
23293 || q_out_t.len() < t * nh_q * head_dim
23294 || k_out_t.len() < t * nh_k * head_dim
23295 {
23296 return Err(format!(
23297 "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
23298 nh_q={nh_q} nh_k={nh_k}"
23299 )
23300 .into());
23301 }
23302 let f = self.func("qk_norm_rope_append_inc_dcw_rows");
23303 let same_t: i32 = if same_session { t as i32 } else { 0 };
23304 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
23305 let cfg = LaunchConfig {
23306 grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
23307 block_dim: (128, 1, 1),
23308 shared_mem_bytes: 0,
23309 };
23310 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
23311 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23312 let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
23313 let null: u64 = 0;
23314 let __s_b = self.gpu.stream();
23315 let mut b = __s_b.launch_builder(&f);
23316 b.arg(q_raw_t)
23317 .arg(k_raw_t)
23318 .arg(v_raw_t)
23319 .arg(qw)
23320 .arg(kw)
23321 .arg(q_out_t)
23322 .arg(k_out_t)
23323 .arg(tab)
23324 .arg(pos_t)
23325 .arg(&same_t)
23326 .arg(&kvk)
23327 .arg(&kvv)
23328 .arg(&ktb)
23329 .arg(&vtb)
23330 .arg(&hd)
23331 .arg(&nd)
23332 .arg(&nq)
23333 .arg(&nk)
23334 .arg(&eps)
23335 .arg(&theta_scale)
23336 .arg(&freq_scale);
23337 match ff {
23338 Some(freqs) => {
23339 b.arg(freqs);
23340 }
23341 None => {
23342 b.arg(&null);
23343 }
23344 }
23345 unsafe {
23346 b.launch(cfg)?;
23347 }
23348 Ok(())
23349 }
23350
23351 #[allow(clippy::too_many_arguments)] pub fn qk_norm_rope_append_inc_dcw(
23353 &self,
23354 q_raw: &CudaSlice<f32>,
23355 k_raw: &CudaSlice<f32>,
23356 v_raw: &CudaSlice<f32>,
23357 qw: &CudaSlice<f32>,
23358 kw: &CudaSlice<f32>,
23359 q_out: &mut CudaSlice<f32>,
23360 k_out: &mut CudaSlice<f32>,
23361 pos: &CudaSlice<i32>,
23362 k_plane: &mut CudaSlice<u8>,
23363 v_plane: &mut CudaSlice<u8>,
23364 len_dev: &CudaSlice<i32>,
23367 base_dev: Option<&CudaSlice<i32>>,
23368 done_ctr: &mut CudaSlice<u32>,
23369 kv_dim_k: usize,
23370 kv_dim_v: usize,
23371 k_tok_bytes: usize,
23372 v_tok_bytes: usize,
23373 head_dim: usize,
23374 n_dims: usize,
23375 nh_q: usize,
23376 nh_k: usize,
23377 eps: f32,
23378 freq_base: f32,
23379 freq_scale: f32,
23380 ff: Option<&CudaSlice<f32>>,
23381 ) -> Result<(), Box<dyn std::error::Error>> {
23382 if head_dim != 128
23383 || kv_dim_v != kv_dim_k
23384 || kv_dim_k != nh_k * head_dim
23385 || q_raw.len() < nh_q * head_dim
23386 || k_raw.len() < nh_k * head_dim
23387 || v_raw.len() < kv_dim_v
23388 || q_out.len() < nh_q * head_dim
23389 || k_out.len() < nh_k * head_dim
23390 || pos.is_empty()
23391 || done_ctr.is_empty()
23392 {
23393 return Err(format!(
23394 "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}"
23395 )
23396 .into());
23397 }
23398 let f = self.func("qk_norm_rope_append_inc_dcw");
23399 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
23400 let cfg = LaunchConfig {
23401 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
23402 block_dim: (128, 1, 1),
23403 shared_mem_bytes: 0,
23404 };
23405 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
23406 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23407 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
23408 let null: u64 = 0;
23409 let __s_b = self.gpu.stream();
23410 let mut b = __s_b.launch_builder(&f);
23411 b.arg(q_raw)
23412 .arg(k_raw)
23413 .arg(v_raw)
23414 .arg(qw)
23415 .arg(kw)
23416 .arg(q_out)
23417 .arg(k_out)
23418 .arg(pos)
23419 .arg(&mut *k_plane)
23420 .arg(&mut *v_plane)
23421 .arg(len_dev);
23422 match base_dev {
23423 Some(base) => {
23424 b.arg(base);
23425 }
23426 None => {
23427 b.arg(&null);
23428 }
23429 }
23430 b.arg(&mut *done_ctr)
23431 .arg(&kvk)
23432 .arg(&kvv)
23433 .arg(&ktb)
23434 .arg(&vtb)
23435 .arg(&hd)
23436 .arg(&nd)
23437 .arg(&nq)
23438 .arg(&eps)
23439 .arg(&theta_scale)
23440 .arg(&freq_scale);
23441 match ff {
23442 Some(freqs) => {
23443 b.arg(freqs);
23444 }
23445 None => {
23446 b.arg(&null);
23447 }
23448 }
23449 unsafe {
23450 b.launch(cfg)?;
23451 }
23452 Ok(())
23453 }
23454
23455 #[allow(clippy::too_many_arguments)] pub fn qk_norm_rope_into(
23457 &self,
23458 q_raw: &CudaSlice<f32>,
23459 k_raw: &CudaSlice<f32>,
23460 qw: &CudaSlice<f32>,
23461 kw: &CudaSlice<f32>,
23462 q_out: &mut CudaSlice<f32>,
23463 k_out: &mut CudaSlice<f32>,
23464 pos: &CudaSlice<i32>,
23465 head_dim: usize,
23466 n_dims: usize,
23467 nh_q: usize,
23468 nh_k: usize,
23469 eps: f32,
23470 freq_base: f32,
23471 freq_scale: f32,
23472 ff: Option<&CudaSlice<f32>>,
23473 ) -> Result<(), Box<dyn std::error::Error>> {
23474 if head_dim > 512
23475 || q_raw.len() < nh_q * head_dim
23476 || k_raw.len() < nh_k * head_dim
23477 || q_out.len() < nh_q * head_dim
23478 || k_out.len() < nh_k * head_dim
23479 || qw.len() < head_dim
23480 || kw.len() < head_dim
23481 || pos.is_empty()
23482 {
23483 return Err(format!(
23484 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
23485 )
23486 .into());
23487 }
23488 let f = self.func("qk_norm_rope_f32");
23489 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
23490 let cfg = LaunchConfig {
23491 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
23492 block_dim: (128, 1, 1),
23493 shared_mem_bytes: 0,
23494 };
23495 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
23496 let __s_b = self.gpu.stream();
23497 let mut b = __s_b.launch_builder(&f);
23498 b.arg(q_raw)
23499 .arg(k_raw)
23500 .arg(qw)
23501 .arg(kw)
23502 .arg(q_out)
23503 .arg(k_out)
23504 .arg(pos)
23505 .arg(&hd)
23506 .arg(&nd)
23507 .arg(&nq)
23508 .arg(&eps)
23509 .arg(&theta_scale)
23510 .arg(&freq_scale);
23511 match ff {
23512 Some(ffv) => {
23513 b.arg(ffv);
23514 unsafe {
23515 b.launch(cfg)?;
23516 }
23517 }
23518 None => {
23519 let null: u64 = 0;
23520 b.arg(&null);
23521 unsafe {
23522 b.launch(cfg)?;
23523 }
23524 }
23525 }
23526 Ok(())
23527 }
23528
23529 #[allow(clippy::too_many_arguments)]
23532 pub fn matvec_f32_b4_into(
23533 &self,
23534 w: [&CudaSlice<f32>; 4],
23535 x: &CudaSlice<f32>,
23536 y: &mut CudaSlice<f32>,
23537 block_cols: usize,
23538 out_f: usize,
23539 ) -> Result<(), Box<dyn std::error::Error>> {
23540 if !block_cols.is_multiple_of(4)
23541 || x.len() < 4 * block_cols
23542 || y.len() < out_f
23543 || w.iter().any(|w| w.len() != out_f * block_cols)
23544 {
23545 return Err(format!(
23546 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
23547 x.len()
23548 )
23549 .into());
23550 }
23551 let f = self.func("matvec_f32_b4");
23552 let cfg = LaunchConfig {
23553 grid_dim: (out_f as u32, 1, 1),
23554 block_dim: (128, 1, 1),
23555 shared_mem_bytes: 0,
23556 };
23557 let (bc, of) = (block_cols as i32, out_f as i32);
23558 let __s_b = self.gpu.stream();
23559 let mut b = __s_b.launch_builder(&f);
23560 b.arg(w[0])
23561 .arg(w[1])
23562 .arg(w[2])
23563 .arg(w[3])
23564 .arg(x)
23565 .arg(y)
23566 .arg(&bc)
23567 .arg(&of);
23568 unsafe {
23569 b.launch(cfg)?;
23570 }
23571 Ok(())
23572 }
23573
23574 pub fn axpy_rows_seq_into(
23577 &self,
23578 x: &CudaSlice<f32>,
23579 w: &CudaSlice<f32>,
23580 y: &mut CudaSlice<f32>,
23581 width: usize,
23582 n_rows: usize,
23583 ) -> Result<(), Box<dyn std::error::Error>> {
23584 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
23585 return Err(format!(
23586 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
23587 x.len(),
23588 w.len(),
23589 y.len()
23590 )
23591 .into());
23592 }
23593 let f = self.func("axpy_rows_seq_f32");
23594 let cfg = LaunchConfig::for_num_elems(width as u32);
23595 let (wi, nr) = (width as i32, n_rows as i32);
23596 let __s_b = self.gpu.stream();
23597 let mut b = __s_b.launch_builder(&f);
23598 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
23599 unsafe {
23600 b.launch(cfg)?;
23601 }
23602 Ok(())
23603 }
23604
23605 pub fn axpy_rows_seq_tokens_into(
23608 &self,
23609 x: &CudaSlice<f32>,
23610 w: &CudaSlice<f32>,
23611 y: &mut CudaSlice<f32>,
23612 width: usize,
23613 slots: usize,
23614 tokens: usize,
23615 ) -> Result<(), Box<dyn std::error::Error>> {
23616 let rows = slots
23617 .checked_mul(tokens)
23618 .ok_or("axpy_rows_seq_tokens row count overflow")?;
23619 if x.len() < rows * width || w.len() < rows || y.len() < tokens * width {
23620 return Err(format!(
23621 "axpy_rows_seq_tokens geometry x={} w={} y={} width={width} \
23622 slots={slots} tokens={tokens}",
23623 x.len(),
23624 w.len(),
23625 y.len()
23626 )
23627 .into());
23628 }
23629 let f = self.func("axpy_rows_seq_tokens_f32");
23630 let block = 256u32;
23631 let cfg = LaunchConfig {
23632 grid_dim: ((width as u32).div_ceil(block), tokens as u32, 1),
23633 block_dim: (block, 1, 1),
23634 shared_mem_bytes: 0,
23635 };
23636 let (wi, sl, tk) = (width as i32, slots as i32, tokens as i32);
23637 let __s_b = self.gpu.stream();
23638 let mut b = __s_b.launch_builder(&f);
23639 b.arg(x).arg(w).arg(y).arg(&wi).arg(&sl).arg(&tk);
23640 unsafe {
23641 b.launch(cfg)?;
23642 }
23643 Ok(())
23644 }
23645
23646 #[allow(clippy::too_many_arguments)]
23650 pub fn axpy_rows_seq_md_off_into(
23651 &self,
23652 x: &CudaSlice<f32>,
23653 w_route: &CudaSlice<f32>,
23654 md: &CudaSlice<f32>,
23655 sel: &CudaSlice<i32>,
23656 y: &mut CudaSlice<f32>,
23657 width: usize,
23658 n_rows: usize,
23659 row0: usize,
23660 ) -> Result<(), Box<dyn std::error::Error>> {
23661 if x.len() < (row0 + n_rows) * width
23662 || w_route.len() < row0 + n_rows
23663 || sel.len() < row0 + n_rows
23664 || y.len() < width
23665 {
23666 return Err(format!(
23667 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
23668 rows={n_rows} row0={row0}",
23669 x.len(),
23670 w_route.len(),
23671 sel.len(),
23672 y.len()
23673 )
23674 .into());
23675 }
23676 let f = self.func("axpy_rows_seq_md_off_f32");
23677 let cfg = LaunchConfig::for_num_elems(width as u32);
23678 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
23679 let __s_b = self.gpu.stream();
23680 let mut b = __s_b.launch_builder(&f);
23681 b.arg(x)
23682 .arg(w_route)
23683 .arg(md)
23684 .arg(sel)
23685 .arg(y)
23686 .arg(&wi)
23687 .arg(&nr)
23688 .arg(&r0);
23689 unsafe {
23690 b.launch(cfg)?;
23691 }
23692 Ok(())
23693 }
23694
23695 #[allow(clippy::too_many_arguments)]
23698 pub fn axpy_rows_seq_md_into(
23699 &self,
23700 x: &CudaSlice<f32>,
23701 w_route: &CudaSlice<f32>,
23702 md: &CudaSlice<f32>,
23703 sel: &CudaSlice<i32>,
23704 y: &mut CudaSlice<f32>,
23705 width: usize,
23706 n_rows: usize,
23707 ) -> Result<(), Box<dyn std::error::Error>> {
23708 if x.len() < n_rows * width
23709 || w_route.len() < n_rows
23710 || sel.len() < n_rows
23711 || y.len() < width
23712 {
23713 return Err(format!(
23714 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
23715 x.len(),
23716 w_route.len(),
23717 sel.len(),
23718 y.len()
23719 )
23720 .into());
23721 }
23722 let f = self.func("axpy_rows_seq_md_f32");
23723 let cfg = LaunchConfig::for_num_elems(width as u32);
23724 let (wi, nr) = (width as i32, n_rows as i32);
23725 let __s_b = self.gpu.stream();
23726 let mut b = __s_b.launch_builder(&f);
23727 b.arg(x)
23728 .arg(w_route)
23729 .arg(md)
23730 .arg(sel)
23731 .arg(y)
23732 .arg(&wi)
23733 .arg(&nr);
23734 unsafe {
23735 b.launch(cfg)?;
23736 }
23737 Ok(())
23738 }
23739
23740 #[allow(clippy::too_many_arguments)]
23742 #[allow(clippy::too_many_arguments)]
23746 pub fn matvec_bf16_qkvg_tcol_into(
23747 &self,
23748 wq: &CudaSlice<u8>,
23749 wk: &CudaSlice<u8>,
23750 wv: &CudaSlice<u8>,
23751 wg: &CudaSlice<u8>,
23752 x_t: &CudaSlice<f32>,
23753 yq: &mut CudaSlice<f32>,
23754 yk: &mut CudaSlice<f32>,
23755 yv: &mut CudaSlice<f32>,
23756 yg: &mut CudaSlice<f32>,
23757 in_f: usize,
23758 out_q: usize,
23759 out_kv: usize,
23760 out_g: usize,
23761 t: usize,
23762 ) -> Result<(), Box<dyn std::error::Error>> {
23763 if t == 0
23764 || t > 8
23765 || !in_f.is_multiple_of(8)
23766 || x_t.len() < t * in_f
23767 || yq.len() < t * out_q
23768 || yk.len() < t * out_kv
23769 || yv.len() < t * out_kv
23770 || (out_g > 0 && yg.len() < t * out_g)
23771 {
23772 return Err("matvec_bf16_qkvg_tcol geometry".into());
23773 }
23774 let grid = out_q + 2 * out_kv + out_g;
23775 let cfg = LaunchConfig {
23776 grid_dim: (grid as u32, 1, 1),
23777 block_dim: (mmv_block(), 1, 1),
23778 shared_mem_bytes: 0,
23779 };
23780 let (ini, oq, okv, og, ti) = (
23781 in_f as i32,
23782 out_q as i32,
23783 out_kv as i32,
23784 out_g as i32,
23785 t as i32,
23786 );
23787 let __s_b = self.gpu.stream();
23788 let f = self.func("matvec_bf16_qkvg_tcol");
23794 let mut b = __s_b.launch_builder(&f);
23795 b.arg(wq)
23796 .arg(wk)
23797 .arg(wv)
23798 .arg(wg)
23799 .arg(x_t)
23800 .arg(yq)
23801 .arg(yk)
23802 .arg(yv)
23803 .arg(yg)
23804 .arg(&ini)
23805 .arg(&oq)
23806 .arg(&okv)
23807 .arg(&og)
23808 .arg(&ti);
23809 unsafe {
23810 b.launch(cfg)?;
23811 }
23812 Ok(())
23813 }
23814
23815 #[allow(clippy::too_many_arguments)] pub fn matvec_bf16_qkvg_into(
23817 &self,
23818 wq: &CudaSlice<u8>,
23819 wk: &CudaSlice<u8>,
23820 wv: &CudaSlice<u8>,
23821 wg: &CudaSlice<u8>,
23822 x: &CudaSlice<f32>,
23823 yq: &mut CudaSlice<f32>,
23824 yk: &mut CudaSlice<f32>,
23825 yv: &mut CudaSlice<f32>,
23826 yg: &mut CudaSlice<f32>,
23827 in_f: usize,
23828 out_q: usize,
23829 out_kv: usize,
23830 out_g: usize,
23831 ) -> Result<(), Box<dyn std::error::Error>> {
23832 if !in_f.is_multiple_of(8)
23833 || wq.len() != out_q * in_f * 2
23834 || wk.len() != out_kv * in_f * 2
23835 || wv.len() != out_kv * in_f * 2
23836 || wg.len() < out_g * in_f * 2
23837 || x.len() < in_f
23838 || yq.len() < out_q
23839 || yk.len() < out_kv
23840 || yv.len() < out_kv
23841 || (out_g > 0 && yg.len() < out_g)
23842 {
23843 return Err(format!(
23844 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
23845 )
23846 .into());
23847 }
23848 let f = self.func("matvec_bf16_qkvg");
23849 let cfg = LaunchConfig {
23850 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
23851 block_dim: (mmv_block(), 1, 1),
23852 shared_mem_bytes: 0,
23853 };
23854 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
23855 let __s_b = self.gpu.stream();
23856 let mut b = __s_b.launch_builder(&f);
23857 b.arg(wq)
23858 .arg(wk)
23859 .arg(wv)
23860 .arg(wg)
23861 .arg(x)
23862 .arg(yq)
23863 .arg(yk)
23864 .arg(yv)
23865 .arg(yg)
23866 .arg(&inf)
23867 .arg(&oq)
23868 .arg(&okv)
23869 .arg(&og);
23870 unsafe {
23871 b.launch(cfg)?;
23872 }
23873 Ok(())
23874 }
23875
23876 pub fn matvec_bf16_b4_into(
23878 &self,
23879 w: [&CudaSlice<u8>; 4],
23880 x: &CudaSlice<f32>,
23881 y: &mut CudaSlice<f32>,
23882 block_cols: usize,
23883 out_f: usize,
23884 ) -> Result<(), Box<dyn std::error::Error>> {
23885 if !block_cols.is_multiple_of(8)
23886 || x.len() < 4 * block_cols
23887 || y.len() < out_f
23888 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
23889 {
23890 return Err(format!(
23891 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
23892 x.len()
23893 )
23894 .into());
23895 }
23896 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23899 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
23900 let f = self.func(if x2 {
23901 "matvec_bf16_b4_x2"
23902 } else {
23903 "matvec_bf16_b4"
23904 });
23905 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
23906 let cfg = LaunchConfig {
23907 grid_dim: (grid as u32, 1, 1),
23908 block_dim: (mmv_block(), 1, 1),
23909 shared_mem_bytes: 0,
23910 };
23911 let (bc, of) = (block_cols as i32, out_f as i32);
23912 let __s_b = self.gpu.stream();
23913 let mut b = __s_b.launch_builder(&f);
23914 b.arg(w[0])
23915 .arg(w[1])
23916 .arg(w[2])
23917 .arg(w[3])
23918 .arg(x)
23919 .arg(y)
23920 .arg(&bc)
23921 .arg(&of);
23922 unsafe {
23923 b.launch(cfg)?;
23924 }
23925 Ok(())
23926 }
23927
23928 pub fn matvec_bf16_b4_tcol_into(
23934 &self,
23935 w: [&CudaSlice<u8>; 4],
23936 x_t: &CudaSlice<f32>,
23937 y_t: &mut CudaSlice<f32>,
23938 block_cols: usize,
23939 out_f: usize,
23940 t: usize,
23941 ) -> Result<(), Box<dyn std::error::Error>> {
23942 if !block_cols.is_multiple_of(8)
23943 || t == 0
23944 || t > 8
23945 || x_t.len() < t * 4 * block_cols
23946 || y_t.len() < t * out_f
23947 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
23948 {
23949 return Err(format!(
23950 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
23951 x_t.len()
23952 )
23953 .into());
23954 }
23955 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
23956 return Err(
23957 "b4 tcol verify is qualified against the plain b4 kernel only \
23958 (MEMRA_B4_X2=1 is a different t=1 program)"
23959 .into(),
23960 );
23961 }
23962 let cfg = LaunchConfig {
23966 grid_dim: (out_f as u32, 1, 1),
23967 block_dim: (mmv_block(), 1, 1),
23968 shared_mem_bytes: 0,
23969 };
23970 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
23971 let __s_b = self.gpu.stream();
23972 let f = self.func("matvec_bf16_b4_tcol");
23973 let mut b = __s_b.launch_builder(&f);
23974 b.arg(w[0])
23975 .arg(w[1])
23976 .arg(w[2])
23977 .arg(w[3])
23978 .arg(x_t)
23979 .arg(y_t)
23980 .arg(&bc)
23981 .arg(&of)
23982 .arg(&ti);
23983 unsafe {
23984 b.launch(cfg)?;
23985 }
23986 Ok(())
23987 }
23988
23989 pub fn q8_0_row_bytes(in_f: usize) -> usize {
23992 in_f / 32 * 34
23993 }
23994
23995 pub fn q8_v2_smem_bytes(in_f: usize) -> usize {
23999 in_f + (in_f / 32) * 4
24000 }
24001
24002 pub fn encode_q8_0_from_bf16(
24006 &self,
24007 w_bf16: &CudaSlice<u8>,
24008 out: &mut CudaSlice<u8>,
24009 in_f: usize,
24010 out_f: usize,
24011 ) -> Result<(), Box<dyn std::error::Error>> {
24012 if !in_f.is_multiple_of(32)
24013 || w_bf16.len() < in_f * out_f * 2
24014 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
24015 {
24016 return Err(format!(
24017 "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
24018 w_bf16.len(),
24019 out.len()
24020 )
24021 .into());
24022 }
24023 let f = self.func("encode_q8_0_rows_from_bf16");
24024 const PAIRS_PER_BLOCK: u32 = 4;
24027 let pairs = (out_f * (in_f / 32)) as u64;
24028 let cfg = LaunchConfig {
24029 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
24030 block_dim: (32, PAIRS_PER_BLOCK, 1),
24031 shared_mem_bytes: 0,
24032 };
24033 let (ini, outi) = (in_f as i32, out_f as i32);
24034 let __s_b = self.gpu.stream();
24035 let mut b = __s_b.launch_builder(&f);
24036 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
24037 unsafe {
24038 b.launch(cfg)?;
24039 }
24040 Ok(())
24041 }
24042
24043 pub fn encode_q8_0_from_bf16_view(
24047 &self,
24048 w_bf16: &cudarc::driver::CudaView<'_, u8>,
24049 out: &mut CudaSlice<u8>,
24050 in_f: usize,
24051 out_f: usize,
24052 ) -> Result<(), Box<dyn std::error::Error>> {
24053 if !in_f.is_multiple_of(32)
24054 || w_bf16.len() < in_f * out_f * 2
24055 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
24056 {
24057 return Err(format!(
24058 "encode_q8_0_from_bf16_view geometry in={in_f} out={out_f} src={} dst={}",
24059 w_bf16.len(),
24060 out.len()
24061 )
24062 .into());
24063 }
24064 let f = self.func("encode_q8_0_rows_from_bf16");
24065 const PAIRS_PER_BLOCK: u32 = 4;
24066 let pairs = (out_f * (in_f / 32)) as u64;
24067 let cfg = LaunchConfig {
24068 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
24069 block_dim: (32, PAIRS_PER_BLOCK, 1),
24070 shared_mem_bytes: 0,
24071 };
24072 let (ini, outi) = (in_f as i32, out_f as i32);
24073 let __s_b = self.gpu.stream();
24074 let mut b = __s_b.launch_builder(&f);
24075 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
24076 unsafe {
24077 b.launch(cfg)?;
24078 }
24079 Ok(())
24080 }
24081
24082 #[allow(clippy::too_many_arguments)]
24087 pub fn qmatvec_q8_0_qkv_rp_into(
24088 &self,
24089 wq: &CudaSlice<u8>,
24090 wk: &CudaSlice<u8>,
24091 wv: &CudaSlice<u8>,
24092 aq: &CudaSlice<i8>,
24093 ad: &CudaSlice<f32>,
24094 yq: &mut CudaSlice<f32>,
24095 yk: &mut CudaSlice<f32>,
24096 yv: &mut CudaSlice<f32>,
24097 in_f: usize,
24098 out_q: usize,
24099 out_kv: usize,
24100 ) -> Result<(), Box<dyn std::error::Error>> {
24101 const ROWS_PER_BLOCK: u32 = 4; let rows = out_q + 2 * out_kv;
24103 let nblk = in_f / 32;
24104 if !in_f.is_multiple_of(32)
24105 || aq.len() < in_f
24106 || ad.len() < nblk
24107 || yq.len() < out_q
24108 || yk.len() < out_kv
24109 || yv.len() < out_kv
24110 || wq.len() < out_q * nblk * 34
24111 || wk.len() < out_kv * nblk * 34
24112 || wv.len() < out_kv * nblk * 34
24113 {
24114 return Err(
24115 format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
24116 );
24117 }
24118 let f = self.func("qmatvec_q8_0_qkv_rp");
24119 let cfg = LaunchConfig {
24120 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24121 block_dim: (32, ROWS_PER_BLOCK, 1),
24122 shared_mem_bytes: 0,
24123 };
24124 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
24125 let __s_b = self.gpu.stream();
24126 let mut b = __s_b.launch_builder(&f);
24127 b.arg(wq)
24128 .arg(wk)
24129 .arg(wv)
24130 .arg(aq)
24131 .arg(ad)
24132 .arg(yq)
24133 .arg(yk)
24134 .arg(yv)
24135 .arg(&ini)
24136 .arg(&oq)
24137 .arg(&okv);
24138 unsafe {
24139 b.launch(cfg)?;
24140 }
24141 Ok(())
24142 }
24143
24144 #[allow(clippy::too_many_arguments)]
24148 pub fn qmatvec_q8_0_b4_rp_into(
24149 &self,
24150 w: [&CudaSlice<u8>; 4],
24151 aq: &CudaSlice<i8>,
24152 ad: &CudaSlice<f32>,
24153 y: &mut CudaSlice<f32>,
24154 block_cols: usize,
24155 out_f: usize,
24156 ) -> Result<(), Box<dyn std::error::Error>> {
24157 const ROWS_PER_BLOCK: u32 = 4; let nblk = block_cols / 32;
24159 if !block_cols.is_multiple_of(32)
24160 || aq.len() < 4 * block_cols
24161 || ad.len() < 4 * nblk
24162 || y.len() < out_f
24163 || w.iter().any(|p| p.len() < out_f * nblk * 34)
24164 {
24165 return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
24166 }
24167 let f = self.func("qmatvec_q8_0_b4_rp");
24168 let cfg = LaunchConfig {
24169 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24170 block_dim: (32, ROWS_PER_BLOCK, 1),
24171 shared_mem_bytes: 0,
24172 };
24173 let (bc, of) = (block_cols as i32, out_f as i32);
24174 let __s_b = self.gpu.stream();
24175 let mut b = __s_b.launch_builder(&f);
24176 b.arg(w[0])
24177 .arg(w[1])
24178 .arg(w[2])
24179 .arg(w[3])
24180 .arg(aq)
24181 .arg(ad)
24182 .arg(y)
24183 .arg(&bc)
24184 .arg(&of);
24185 unsafe {
24186 b.launch(cfg)?;
24187 }
24188 Ok(())
24189 }
24190
24191 #[allow(clippy::map_entry)] fn matvec_bf16_via_q8_mirror_t(
24195 &self,
24196 data: &CudaSlice<u8>,
24197 x: &CudaSlice<f32>,
24198 y: &mut CudaSlice<f32>,
24199 in_f: usize,
24200 out_f: usize,
24201 t: usize,
24202 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
24203 use cudarc::driver::DevicePtr;
24204 let key = {
24205 let s = self.gpu.stream();
24206 let (p, _g) = data.device_ptr(&s);
24207 (p, in_f as u32, out_f as u32)
24208 };
24209 {
24210 let mut mirrors = self
24211 .w8_mirrors
24212 .lock()
24213 .map_err(|_| "w8 mirror map is poisoned")?;
24214 if !mirrors.contains_key(&key) {
24215 if crate::glm5_graph_capture_open() {
24222 return Err(format!(
24223 "MEMRA_GLM5_W8: the q8_0 mirror for a {in_f}x{out_f} weight would be built inside an \
24224 open CUDA graph capture (its quantize kernels recorded, never executed; memra#131). \
24225 Warm the walk before capturing."
24226 )
24227 .into());
24228 }
24229 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
24230 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
24231 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
24232 mirrors.insert(key, planar);
24233 }
24234 }
24235 let nblk = in_f / 32;
24236 let akey = in_f * 64 + t.min(32);
24238 {
24239 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24240 if let std::collections::hash_map::Entry::Vacant(slot) = act.entry(akey) {
24241 let aq = self.alloc_i8_uninit(32 * in_f)?;
24242 let ad = self.alloc_uninit::<f32>(32 * nblk)?;
24243 slot.insert((aq, ad));
24244 }
24245 let (aq, ad) = act.get_mut(&akey).expect("just inserted");
24246 self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
24247 }
24248 let mirrors = self
24249 .w8_mirrors
24250 .lock()
24251 .map_err(|_| "w8 mirror map is poisoned")?;
24252 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24253 let mirror = mirrors.get(&key).expect("built above");
24254 let (aq, ad) = act.get(&akey).expect("built above");
24255 const ROWS_PER_BLOCK: u32 = 4;
24256 let (ini, of) = (in_f as i32, out_f as i32);
24257 if b200_gemv_v2_on() && q8t_wonce_on() && t <= 8 {
24264 if GEMV_V2_Q8_ROWS_TW_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0
24265 {
24266 eprintln!(
24267 "[b200-gemv-v2] engaged arm=q8_rows_tw_v2 t={t} in_f={in_f} out_f={out_f} \
24268 (W8 verify walk; MEMRA_B200_GEMV_V2=1)"
24269 );
24270 }
24271 self.qmatvec_q8_0_rows_tw_v2_raw(mirror, aq, ad, y, in_f, out_f, t)?;
24272 return Ok(Some(()));
24273 }
24274 if q8t_wonce_on() && t <= 32 {
24275 let f = self.func(if t <= 8 {
24276 "qmatvec_q8_0_rows_tw"
24277 } else {
24278 "qmatvec_q8_0_rows_tw32"
24279 });
24280 let cfg = LaunchConfig {
24281 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24282 block_dim: (32, ROWS_PER_BLOCK, 1),
24283 shared_mem_bytes: 0,
24284 };
24285 let ti = t as i32;
24286 let __s_b = self.gpu.stream();
24287 let mut b = __s_b.launch_builder(&f);
24288 b.arg(mirror)
24289 .arg(aq)
24290 .arg(ad)
24291 .arg(&mut *y)
24292 .arg(&ini)
24293 .arg(&of)
24294 .arg(&ti);
24295 unsafe {
24296 b.launch(cfg)?;
24297 }
24298 return Ok(Some(()));
24299 }
24300 let f = self.func("qmatvec_q8_0_rows_t");
24301 let cfg = LaunchConfig {
24302 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
24303 block_dim: (32, ROWS_PER_BLOCK, 1),
24304 shared_mem_bytes: 0,
24305 };
24306 let __s_b = self.gpu.stream();
24307 let mut b = __s_b.launch_builder(&f);
24308 b.arg(mirror)
24309 .arg(aq)
24310 .arg(ad)
24311 .arg(&mut *y)
24312 .arg(&ini)
24313 .arg(&of);
24314 unsafe {
24315 b.launch(cfg)?;
24316 }
24317 Ok(Some(()))
24318 }
24319
24320 #[allow(clippy::map_entry)] fn matvec_bf16_via_q8_mirror(
24324 &self,
24325 data: &CudaSlice<u8>,
24326 x: &CudaSlice<f32>,
24327 y: &mut CudaSlice<f32>,
24328 in_f: usize,
24329 out_f: usize,
24330 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
24331 use cudarc::driver::DevicePtr;
24332 let key = {
24333 let s = self.gpu.stream();
24334 let (p, _g) = data.device_ptr(&s);
24335 (p, in_f as u32, out_f as u32)
24336 };
24337 {
24338 let mut mirrors = self
24339 .w8_mirrors
24340 .lock()
24341 .map_err(|_| "w8 mirror map is poisoned")?;
24342 if !mirrors.contains_key(&key) {
24343 if crate::glm5_graph_capture_open() {
24350 return Err(format!(
24351 "MEMRA_GLM5_W8: the q8_0 mirror for a {in_f}x{out_f} weight would be built inside an \
24352 open CUDA graph capture (its quantize kernels recorded, never executed; memra#131). \
24353 Warm the walk before capturing."
24354 )
24355 .into());
24356 }
24357 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
24358 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
24359 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
24360 mirrors.insert(key, planar);
24361 if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
24367 eprintln!(
24368 "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
24369 mirrors.len()
24370 );
24371 }
24372 }
24373 }
24374 let nblk = in_f / 32;
24375 {
24376 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24377 if !act.contains_key(&in_f) {
24378 let aq = self.alloc_uninit::<i8>(in_f)?;
24379 let ad = self.alloc_uninit::<f32>(nblk)?;
24380 act.insert(in_f, (aq, ad));
24381 }
24382 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
24383 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
24384 }
24385 let mirrors = self
24386 .w8_mirrors
24387 .lock()
24388 .map_err(|_| "w8 mirror map is poisoned")?;
24389 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24390 let mirror = mirrors.get(&key).expect("built above");
24391 let (aq, ad) = act.get(&in_f).expect("built above");
24392 if b200_gemv_v2_on() && in_f.is_multiple_of(32) && Self::q8_v2_smem_bytes(in_f) <= 48 * 1024
24402 {
24403 if GEMV_V2_Q8_RP_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
24404 eprintln!(
24405 "[b200-gemv-v2] engaged arm=q8_rp_v2 t=1 in_f={in_f} out_f={out_f} \
24406 (W8 posture; MEMRA_B200_GEMV_V2=1)"
24407 );
24408 }
24409 self.qmatvec_q8_0_rp_v2_raw(mirror, aq, ad, y, in_f, out_f, 1)?;
24410 return Ok(Some(()));
24411 }
24412 self.qmatvec_mmvq_into(
24413 mirror,
24414 aq,
24415 ad,
24416 1,
24417 in_f,
24418 out_f,
24419 QT_Q8_0,
24420 Self::q8_0_row_bytes(in_f),
24421 1.0,
24422 true,
24423 y,
24424 )?;
24425 Ok(Some(()))
24426 }
24427
24428 #[allow(clippy::too_many_arguments)]
24433 pub fn qmatvec_q8_0_qkv_rp_t_into(
24434 &self,
24435 wq: &CudaSlice<u8>,
24436 wk: &CudaSlice<u8>,
24437 wv: &CudaSlice<u8>,
24438 aq: &CudaSlice<i8>,
24439 ad: &CudaSlice<f32>,
24440 yq: &mut CudaSlice<f32>,
24441 yk: &mut CudaSlice<f32>,
24442 yv: &mut CudaSlice<f32>,
24443 in_f: usize,
24444 out_q: usize,
24445 out_kv: usize,
24446 t: usize,
24447 ) -> Result<(), Box<dyn std::error::Error>> {
24448 const ROWS_PER_BLOCK: u32 = 4;
24449 let rows = out_q + 2 * out_kv;
24450 let nblk = in_f / 32;
24451 if !in_f.is_multiple_of(32)
24452 || t == 0
24453 || aq.len() < t * in_f
24454 || ad.len() < t * nblk
24455 || yq.len() < t * out_q
24456 || yk.len() < t * out_kv
24457 || yv.len() < t * out_kv
24458 {
24459 return Err(format!("q8_0 qkv rp_t geometry in={in_f} t={t}").into());
24460 }
24461 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
24462 if q8t_wonce_on() && t <= 32 {
24466 let f = self.func(if t <= 8 {
24467 "qmatvec_q8_0_qkv_rp_tw"
24468 } else {
24469 "qmatvec_q8_0_qkv_rp_tw32"
24470 });
24471 let cfg = LaunchConfig {
24472 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24473 block_dim: (32, ROWS_PER_BLOCK, 1),
24474 shared_mem_bytes: 0,
24475 };
24476 let ti = t as i32;
24477 let __s_b = self.gpu.stream();
24478 let mut b = __s_b.launch_builder(&f);
24479 b.arg(wq)
24480 .arg(wk)
24481 .arg(wv)
24482 .arg(aq)
24483 .arg(ad)
24484 .arg(yq)
24485 .arg(yk)
24486 .arg(yv)
24487 .arg(&ini)
24488 .arg(&oq)
24489 .arg(&okv)
24490 .arg(&ti);
24491 unsafe {
24492 b.launch(cfg)?;
24493 }
24494 return Ok(());
24495 }
24496 let f = self.func("qmatvec_q8_0_qkv_rp_t");
24497 let cfg = LaunchConfig {
24498 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
24499 block_dim: (32, ROWS_PER_BLOCK, 1),
24500 shared_mem_bytes: 0,
24501 };
24502 let __s_b = self.gpu.stream();
24503 let mut b = __s_b.launch_builder(&f);
24504 b.arg(wq)
24505 .arg(wk)
24506 .arg(wv)
24507 .arg(aq)
24508 .arg(ad)
24509 .arg(yq)
24510 .arg(yk)
24511 .arg(yv)
24512 .arg(&ini)
24513 .arg(&oq)
24514 .arg(&okv);
24515 unsafe {
24516 b.launch(cfg)?;
24517 }
24518 Ok(())
24519 }
24520
24521 #[allow(clippy::too_many_arguments)]
24524 pub fn qmatvec_q8_0_b4_rp_t_into(
24525 &self,
24526 w: [&CudaSlice<u8>; 4],
24527 aq: &CudaSlice<i8>,
24528 ad: &CudaSlice<f32>,
24529 y: &mut CudaSlice<f32>,
24530 block_cols: usize,
24531 out_f: usize,
24532 t: usize,
24533 ) -> Result<(), Box<dyn std::error::Error>> {
24534 const ROWS_PER_BLOCK: u32 = 4;
24535 let nblk = block_cols / 32;
24536 if !block_cols.is_multiple_of(32)
24537 || t == 0
24538 || aq.len() < t * 4 * block_cols
24539 || ad.len() < t * 4 * nblk
24540 || y.len() < t * out_f
24541 {
24542 return Err(format!("q8_0 b4 rp_t geometry cols={block_cols} t={t}").into());
24543 }
24544 let (bc, of) = (block_cols as i32, out_f as i32);
24545 if q8t_wonce_on() && t <= 32 {
24548 let f = self.func(if t <= 8 {
24549 "qmatvec_q8_0_b4_rp_tw"
24550 } else {
24551 "qmatvec_q8_0_b4_rp_tw32"
24552 });
24553 let cfg = LaunchConfig {
24554 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
24555 block_dim: (32, ROWS_PER_BLOCK, 1),
24556 shared_mem_bytes: 0,
24557 };
24558 let ti = t as i32;
24559 let __s_b = self.gpu.stream();
24560 let mut b = __s_b.launch_builder(&f);
24561 b.arg(w[0])
24562 .arg(w[1])
24563 .arg(w[2])
24564 .arg(w[3])
24565 .arg(aq)
24566 .arg(ad)
24567 .arg(y)
24568 .arg(&bc)
24569 .arg(&of)
24570 .arg(&ti);
24571 unsafe {
24572 b.launch(cfg)?;
24573 }
24574 return Ok(());
24575 }
24576 let f = self.func("qmatvec_q8_0_b4_rp_t");
24577 let cfg = LaunchConfig {
24578 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
24579 block_dim: (32, ROWS_PER_BLOCK, 1),
24580 shared_mem_bytes: 0,
24581 };
24582 let __s_b = self.gpu.stream();
24583 let mut b = __s_b.launch_builder(&f);
24584 b.arg(w[0])
24585 .arg(w[1])
24586 .arg(w[2])
24587 .arg(w[3])
24588 .arg(aq)
24589 .arg(ad)
24590 .arg(y)
24591 .arg(&bc)
24592 .arg(&of);
24593 unsafe {
24594 b.launch(cfg)?;
24595 }
24596 Ok(())
24597 }
24598
24599 #[allow(clippy::map_entry)] fn matvec_bf16_view_via_q8_mirror(
24610 &self,
24611 data: &cudarc::driver::CudaView<'_, u8>,
24612 x: &CudaSlice<f32>,
24613 y: &mut CudaSlice<f32>,
24614 in_f: usize,
24615 out_f: usize,
24616 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
24617 use cudarc::driver::DevicePtr;
24618 let key = {
24619 let s = self.gpu.stream();
24620 let (p, _g) = data.device_ptr(&s);
24621 (p, in_f as u32, out_f as u32)
24622 };
24623 {
24624 let mut mirrors = self
24625 .w8_mirrors
24626 .lock()
24627 .map_err(|_| "w8 mirror map is poisoned")?;
24628 if !mirrors.contains_key(&key) {
24629 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
24630 self.encode_q8_0_from_bf16_view(data, &mut interleaved, in_f, out_f)?;
24631 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
24632 mirrors.insert(key, planar);
24633 eprintln!("[w8-view] mirror built in_f={in_f} out_f={out_f}");
24637 }
24638 }
24639 let nblk = in_f / 32;
24640 {
24641 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24642 if !act.contains_key(&in_f) {
24643 let aq = self.alloc_uninit::<i8>(in_f)?;
24644 let ad = self.alloc_uninit::<f32>(nblk)?;
24645 act.insert(in_f, (aq, ad));
24646 }
24647 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
24648 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
24649 }
24650 let mirrors = self
24651 .w8_mirrors
24652 .lock()
24653 .map_err(|_| "w8 mirror map is poisoned")?;
24654 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
24655 let mirror = mirrors.get(&key).expect("built above");
24656 let (aq, ad) = act.get(&in_f).expect("built above");
24657 self.qmatvec_mmvq_into(
24658 mirror,
24659 aq,
24660 ad,
24661 1,
24662 in_f,
24663 out_f,
24664 QT_Q8_0,
24665 Self::q8_0_row_bytes(in_f),
24666 1.0,
24667 true,
24668 y,
24669 )?;
24670 Ok(Some(()))
24671 }
24672
24673 pub fn matvec_bf16_into(
24674 &self,
24675 data: &CudaSlice<u8>,
24676 x: &CudaSlice<f32>,
24677 y: &mut CudaSlice<f32>,
24678 in_f: usize,
24679 out_f: usize,
24680 ) -> Result<(), Box<dyn std::error::Error>> {
24681 if data.len() != in_f * out_f * 2
24682 || x.len() < in_f
24683 || !in_f.is_multiple_of(8)
24684 || y.len() < out_f
24685 {
24686 return Err(format!(
24687 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24688 data.len(),
24689 x.len(),
24690 y.len()
24691 )
24692 .into());
24693 }
24694 if step_tp_w8_on()
24701 && w8_hybrid_on()
24702 && in_f.is_multiple_of(32)
24703 && out_f >= 64
24704 && let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)?
24705 {
24706 return Ok(());
24707 }
24708 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24712 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
24713 && in_f <= 2048;
24714 if x4 {
24715 let f = self.func("matvec_bf16_f32acc_x4");
24716 let cfg = LaunchConfig {
24717 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
24718 block_dim: (mmv_block(), 1, 1),
24719 shared_mem_bytes: 0,
24720 };
24721 let (ini, outi) = (in_f as i32, out_f as i32);
24722 let __s_b = self.gpu.stream();
24723 let mut b = __s_b.launch_builder(&f);
24724 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
24725 unsafe {
24726 b.launch(cfg)?;
24727 }
24728 return Ok(());
24729 }
24730 let f = self.func("matvec_bf16_f32acc");
24731 let cfg = LaunchConfig {
24732 grid_dim: (out_f as u32, 1, 1),
24733 block_dim: (mmv_block(), 1, 1),
24734 shared_mem_bytes: 0,
24735 };
24736 let ini = in_f as i32;
24737 let __s_b = self.gpu.stream();
24738 let mut b = __s_b.launch_builder(&f);
24739 b.arg(data).arg(x).arg(y).arg(&ini);
24740 unsafe {
24741 b.launch(cfg)?;
24742 }
24743 Ok(())
24744 }
24745
24746 pub fn matvec_bf16_views_into(
24751 &self,
24752 data: &CudaSlice<u8>,
24753 x: &cudarc::driver::CudaView<'_, f32>,
24754 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
24755 in_f: usize,
24756 out_f: usize,
24757 ) -> Result<(), Box<dyn std::error::Error>> {
24758 if data.len() != in_f * out_f * 2
24759 || x.len() < in_f
24760 || !in_f.is_multiple_of(8)
24761 || y.len() < out_f
24762 {
24763 return Err(format!(
24764 "matvec_bf16_views_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24765 data.len(),
24766 x.len(),
24767 y.len()
24768 )
24769 .into());
24770 }
24771 let f = self.func("matvec_bf16_f32acc");
24772 let cfg = LaunchConfig {
24773 grid_dim: (out_f as u32, 1, 1),
24774 block_dim: (mmv_block(), 1, 1),
24775 shared_mem_bytes: 0,
24776 };
24777 let ini = in_f as i32;
24778 let __s_b = self.gpu.stream();
24779 let mut b = __s_b.launch_builder(&f);
24780 b.arg(data).arg(x).arg(y).arg(&ini);
24781 unsafe {
24782 b.launch(cfg)?;
24783 }
24784 Ok(())
24785 }
24786
24787 pub fn matvec_bf16_view_into(
24790 &self,
24791 data: &cudarc::driver::CudaView<'_, u8>,
24792 x: &CudaSlice<f32>,
24793 y: &mut CudaSlice<f32>,
24794 in_f: usize,
24795 out_f: usize,
24796 ) -> Result<(), Box<dyn std::error::Error>> {
24797 if data.len() != in_f * out_f * 2
24798 || x.len() < in_f
24799 || !in_f.is_multiple_of(8)
24800 || y.len() < out_f
24801 {
24802 return Err(format!(
24803 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
24804 data.len(),
24805 x.len(),
24806 y.len()
24807 )
24808 .into());
24809 }
24810 if w8_view_on()
24811 && step_tp_w8_on()
24812 && w8_hybrid_on()
24813 && in_f.is_multiple_of(32)
24814 && out_f >= 64
24815 && let Some(()) = self.matvec_bf16_view_via_q8_mirror(data, x, y, in_f, out_f)?
24816 {
24817 return Ok(());
24818 }
24819 let f = self.func("matvec_bf16_f32acc");
24820 let cfg = LaunchConfig {
24821 grid_dim: (out_f as u32, 1, 1),
24822 block_dim: (mmv_block(), 1, 1),
24823 shared_mem_bytes: 0,
24824 };
24825 let ini = in_f as i32;
24826 let __s_b = self.gpu.stream();
24827 let mut b = __s_b.launch_builder(&f);
24828 b.arg(data).arg(x).arg(y).arg(&ini);
24829 unsafe {
24830 b.launch(cfg)?;
24831 }
24832 Ok(())
24833 }
24834
24835 pub fn matvec_bf16_raw_out(
24838 &self,
24839 w: &CudaSlice<u8>,
24840 x: &CudaSlice<f32>,
24841 y_raw: u64,
24842 in_f: usize,
24843 out_f: usize,
24844 ) -> Result<(), Box<dyn std::error::Error>> {
24845 if w.len() != in_f * out_f * 2 || x.len() < in_f || !in_f.is_multiple_of(8) || y_raw == 0 {
24846 return Err("matvec_bf16_raw_out geometry".into());
24847 }
24848 let f = self.func("matvec_bf16_f32acc");
24849 let cfg = LaunchConfig {
24850 grid_dim: (out_f as u32, 1, 1),
24851 block_dim: (mmv_block(), 1, 1),
24852 shared_mem_bytes: 0,
24853 };
24854 let ini = in_f as i32;
24855 let __s_b = self.gpu.stream();
24856 let mut b = __s_b.launch_builder(&f);
24857 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
24858 unsafe {
24859 b.launch(cfg)?;
24860 }
24861 Ok(())
24862 }
24863
24864 pub fn add3_raw(
24868 &self,
24869 a: &CudaSlice<f32>,
24870 b: &CudaSlice<f32>,
24871 sh_raw: u64,
24872 scale_raw: u64,
24873 dst: &mut CudaSlice<f32>,
24874 n: usize,
24875 ) -> Result<(), Box<dyn std::error::Error>> {
24876 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
24877 return Err("add3_raw geometry".into());
24878 }
24879 let f = self.func("add3_f32");
24880 let cfg = LaunchConfig {
24881 grid_dim: ((n as u32).div_ceil(256), 1, 1),
24882 block_dim: (256, 1, 1),
24883 shared_mem_bytes: 0,
24884 };
24885 let ni = n as i32;
24886 let __s_b = self.gpu.stream();
24887 let mut bld = __s_b.launch_builder(&f);
24888 bld.arg(a)
24889 .arg(b)
24890 .arg(&sh_raw)
24891 .arg(&scale_raw)
24892 .arg(dst)
24893 .arg(&ni);
24894 unsafe {
24895 bld.launch(cfg)?;
24896 }
24897 Ok(())
24898 }
24899
24900 pub fn matvec_bf16_down_addscale_into(
24903 &self,
24904 w: &CudaSlice<u8>,
24905 x: &CudaSlice<f32>,
24906 scale: &CudaSlice<f32>,
24907 dst: &mut CudaSlice<f32>,
24908 in_f: usize,
24909 out_f: usize,
24910 ) -> Result<(), Box<dyn std::error::Error>> {
24911 if w.len() != in_f * out_f * 2
24912 || x.len() < in_f
24913 || !in_f.is_multiple_of(8)
24914 || dst.len() < out_f
24915 || scale.is_empty()
24916 {
24917 return Err("matvec_bf16_down_addscale geometry".into());
24918 }
24919 let f = self.func("matvec_bf16_down_addscale");
24920 let cfg = LaunchConfig {
24921 grid_dim: (out_f as u32, 1, 1),
24922 block_dim: (mmv_block(), 1, 1),
24923 shared_mem_bytes: 0,
24924 };
24925 let ini = in_f as i32;
24926 let __s_b = self.gpu.stream();
24927 let mut b = __s_b.launch_builder(&f);
24928 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
24929 unsafe {
24930 b.launch(cfg)?;
24931 }
24932 Ok(())
24933 }
24934
24935 #[allow(clippy::too_many_arguments)]
24939 pub fn matvec_bf16_dual_silu_rows_into(
24940 &self,
24941 wg: &CudaSlice<u8>,
24942 wu: &CudaSlice<u8>,
24943 x: &CudaSlice<f32>,
24944 act: &mut CudaSlice<f32>,
24945 in_f: usize,
24946 out_f: usize,
24947 limit: Option<f32>,
24948 t: usize,
24949 ) -> Result<(), Box<dyn std::error::Error>> {
24950 if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
24951 return Err("matvec_bf16_dual_silu_rows geometry".into());
24952 }
24953 let f = self.func("matvec_bf16_dual_silu_rows");
24954 let cfg = LaunchConfig {
24955 grid_dim: (out_f as u32, t as u32, 1),
24956 block_dim: (mmv_block(), 1, 1),
24957 shared_mem_bytes: 0,
24958 };
24959 let (ini, outi) = (in_f as i32, out_f as i32);
24960 let lim = limit.unwrap_or(0.0);
24961 let __s_b = self.gpu.stream();
24962 let mut b = __s_b.launch_builder(&f);
24963 b.arg(wg)
24964 .arg(wu)
24965 .arg(x)
24966 .arg(&mut *act)
24967 .arg(&ini)
24968 .arg(&outi)
24969 .arg(&lim);
24970 unsafe {
24971 b.launch(cfg)?;
24972 }
24973 Ok(())
24974 }
24975
24976 pub fn matvec_bf16_rows_into(
24978 &self,
24979 w: &CudaSlice<u8>,
24980 x: &CudaSlice<f32>,
24981 y: &mut CudaSlice<f32>,
24982 in_f: usize,
24983 out_f: usize,
24984 t: usize,
24985 ) -> Result<(), Box<dyn std::error::Error>> {
24986 if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || !in_f.is_multiple_of(8)
24987 {
24988 return Err("matvec_bf16_rows geometry".into());
24989 }
24990 if (2..=32).contains(&t)
24995 && step_tp_w8_on()
24996 && w8_hybrid_on()
24997 && in_f.is_multiple_of(32)
24998 && out_f >= 64
24999 && let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)?
25000 {
25001 return Ok(());
25002 }
25003 if t == 1
25009 && step_tp_w8_on()
25010 && w8_hybrid_on()
25011 && in_f.is_multiple_of(32)
25012 && out_f >= 64
25013 && let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)?
25014 {
25015 return Ok(());
25016 }
25017 if (2..=32).contains(&t)
25023 && glm5_w8_on()
25024 && in_f.is_multiple_of(32)
25025 && out_f >= 64
25026 && let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)?
25027 {
25028 if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25029 eprintln!(
25030 "[glm5-w8] engaged t={t} in_f={in_f} out_f={out_f} (q8_0 mirror, \
25031 MEMRA_GLM5_W8=1)"
25032 );
25033 }
25034 return Ok(());
25035 }
25036 if t == 1
25040 && glm5_w8_on()
25041 && in_f.is_multiple_of(32)
25042 && out_f >= 64
25043 && let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)?
25044 {
25045 if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25046 eprintln!(
25047 "[glm5-w8] engaged t=1 in_f={in_f} out_f={out_f} (q8_0 mirror, \
25048 MEMRA_GLM5_W8=1)"
25049 );
25050 }
25051 return Ok(());
25052 }
25053 if (2..=16).contains(&t) && bf16_tcols_wide_on() {
25061 if BF16_TCOLS_WIDE_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25062 eprintln!(
25063 "[bf16-tcols-wide] engaged: t={t} in_f={in_f} out_f={out_f} rides the \
25064 weight-once tcols class (MEMRA_BF16_TCOLS_WIDE=1)"
25065 );
25066 }
25067 if t <= 8 {
25068 return self.matvec_bf16_tcols_into(w, x, y, in_f, out_f, t);
25069 }
25070 return self.matvec_bf16_tcols16_into(w, x, y, in_f, out_f, t);
25071 }
25072 if b200_gemv_v2_on() {
25081 let ksplit = self.gemv_v2_ksplit(in_f, out_f, t);
25082 let v3 = b200_gemv_v2_level() >= 2 && ksplit == 1 && gemv_v3_fits();
25086 if GEMV_V2_BF16_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25087 let arm = if v3 { "v3 (cp.async staged)" } else { "v2" };
25088 eprintln!(
25089 "[b200-gemv-v2] engaged arm={arm} t={t} in_f={in_f} out_f={out_f} \
25090 ksplit={ksplit} (MEMRA_B200_GEMV_V2={})",
25091 b200_gemv_v2_level()
25092 );
25093 }
25094 if v3 {
25095 return self.matvec_bf16_v3_raw(w, x, y, in_f, out_f, t);
25096 }
25097 return self.matvec_bf16_v2_raw(w, x, y, in_f, out_f, t, ksplit);
25098 }
25099 if b200_bf16_gemv_lt_on() && self.bf16_gemv_lt_into(w, x, y, in_f, out_f, t)? {
25107 return Ok(());
25108 }
25109 let kname = if b200_matvec_arm_on() {
25116 "matvec_bf16_f32acc_x4_rows_pf"
25117 } else {
25118 "matvec_bf16_f32acc_x4_rows"
25119 };
25120 let f = self.func(kname);
25121 let cfg = LaunchConfig {
25122 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
25123 block_dim: (mmv_block(), 1, 1),
25124 shared_mem_bytes: 0,
25125 };
25126 let (ini, outi) = (in_f as i32, out_f as i32);
25127 let __s_b = self.gpu.stream();
25128 let mut b = __s_b.launch_builder(&f);
25129 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
25130 unsafe {
25131 b.launch(cfg)?;
25132 }
25133 Ok(())
25134 }
25135
25136 #[allow(clippy::too_many_arguments)]
25142 pub fn matvec_bf16_col_range_into(
25143 &self,
25144 w: &CudaSlice<u8>,
25145 x: &CudaSlice<f32>,
25146 y: &mut CudaSlice<f32>,
25147 in_f: usize,
25148 out_f: usize,
25149 k_start: usize,
25150 k_len: usize,
25151 ) -> Result<(), Box<dyn std::error::Error>> {
25152 let weight_bytes = in_f
25153 .checked_mul(out_f)
25154 .and_then(|elements| elements.checked_mul(2))
25155 .ok_or("BF16 column-range matvec geometry")?;
25156 let k_end = k_start
25157 .checked_add(k_len)
25158 .ok_or("BF16 column-range matvec geometry")?;
25159 if w.len() < weight_bytes
25160 || x.len() < k_len
25161 || y.len() < out_f
25162 || in_f > i32::MAX as usize
25163 || out_f > i32::MAX as usize
25164 || k_start > i32::MAX as usize
25165 || k_len > i32::MAX as usize
25166 || in_f == 0
25167 || out_f == 0
25168 || k_len == 0
25169 || !in_f.is_multiple_of(8)
25170 || !k_start.is_multiple_of(8)
25171 || !k_len.is_multiple_of(8)
25172 || k_end > in_f
25173 {
25174 return Err("BF16 column-range matvec geometry".into());
25175 }
25176 let function = self.func("matvec_bf16_f32acc_x4_range");
25177 let config = LaunchConfig {
25178 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
25179 block_dim: (mmv_block(), 1, 1),
25180 shared_mem_bytes: 0,
25181 };
25182 let (in_f, out_f, k_start, k_len) =
25183 (in_f as i32, out_f as i32, k_start as i32, k_len as i32);
25184 let stream = self.gpu.stream();
25185 let mut launch = stream.launch_builder(&function);
25186 launch
25187 .arg(w)
25188 .arg(x)
25189 .arg(y)
25190 .arg(&in_f)
25191 .arg(&out_f)
25192 .arg(&k_start)
25193 .arg(&k_len);
25194 unsafe {
25195 launch.launch(config)?;
25196 }
25197 Ok(())
25198 }
25199
25200 pub fn matvec_bf16_tcols_into(
25208 &self,
25209 w: &CudaSlice<u8>,
25210 x: &CudaSlice<f32>,
25211 y: &mut CudaSlice<f32>,
25212 in_f: usize,
25213 out_f: usize,
25214 t: usize,
25215 ) -> Result<(), Box<dyn std::error::Error>> {
25216 if x.len() < t * in_f
25217 || y.len() < t * out_f
25218 || !(2..=8).contains(&t)
25219 || !in_f.is_multiple_of(8)
25220 {
25221 return Err("matvec_bf16_tcols geometry".into());
25222 }
25223 let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
25234 if rf
25235 && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
25236 == 0
25237 {
25238 eprintln!(
25239 "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
25240 shared across the t token columns + intra-warp shuffles at the identical \
25241 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
25242 );
25243 }
25244 let x1 = bf16_tcols_x1_on();
25245 let (fname, grid_x) = match (x1, rf) {
25246 (true, true) => ("matvec_bf16_f32acc_x1_tcols_rf", out_f as u32),
25247 (true, false) => ("matvec_bf16_f32acc_x1_tcols", out_f as u32),
25248 (false, true) => ("matvec_bf16_f32acc_x4_tcols_rf", out_f.div_ceil(4) as u32),
25249 (false, false) => ("matvec_bf16_f32acc_x4_tcols", out_f.div_ceil(4) as u32),
25250 };
25251 if x1 && BF16_TCOLS_X1_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25252 eprintln!(
25253 "[bf16-tcols-x1] engaged: one-row-per-block tcols grid \
25254 (MEMRA_BF16_TCOLS_X1=1)"
25255 );
25256 }
25257 let f = self.func(fname);
25258 let cfg = LaunchConfig {
25259 grid_dim: (grid_x, 1, 1),
25260 block_dim: (mmv_block(), 1, 1),
25261 shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
25262 };
25263 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
25264 let __s_b = self.gpu.stream();
25265 let mut b = __s_b.launch_builder(&f);
25266 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
25267 unsafe {
25268 b.launch(cfg)?;
25269 }
25270 Ok(())
25271 }
25272
25273 pub fn matvec_bf16_tcols16_into(
25279 &self,
25280 w: &CudaSlice<u8>,
25281 x: &CudaSlice<f32>,
25282 y: &mut CudaSlice<f32>,
25283 in_f: usize,
25284 out_f: usize,
25285 t: usize,
25286 ) -> Result<(), Box<dyn std::error::Error>> {
25287 if x.len() < t * in_f
25288 || y.len() < t * out_f
25289 || !(9..=16).contains(&t)
25290 || !in_f.is_multiple_of(8)
25291 {
25292 return Err("matvec_bf16_tcols16 geometry".into());
25293 }
25294 let rf = bf16_tcols_red_fused_on() && mmv_block().is_power_of_two();
25298 if rf
25299 && BF16_TCOLS_RED_FUSED_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
25300 == 0
25301 {
25302 eprintln!(
25303 "[bf16-tcols-red-fused] engaged: fused-t reduce tail, one barrier sequence \
25304 shared across the t token columns + intra-warp shuffles at the identical \
25305 pairing (MEMRA_BF16_TCOLS_RED_FUSED=1)"
25306 );
25307 }
25308 let f = self.func(if rf {
25309 "matvec_bf16_f32acc_x4_tcols16_rf"
25310 } else {
25311 "matvec_bf16_f32acc_x4_tcols16"
25312 });
25313 let cfg = LaunchConfig {
25314 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
25315 block_dim: (mmv_block(), 1, 1),
25316 shared_mem_bytes: if rf { (t as u32) * mmv_block() * 4 } else { 0 },
25317 };
25318 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
25319 let __s_b = self.gpu.stream();
25320 let mut b = __s_b.launch_builder(&f);
25321 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
25322 unsafe {
25323 b.launch(cfg)?;
25324 }
25325 Ok(())
25326 }
25327
25328 #[allow(clippy::too_many_arguments)]
25334 pub fn matvec_bf16_tcols_gate_kernel_into(
25335 &self,
25336 kernel: &str,
25337 w: &CudaSlice<u8>,
25338 x: &CudaSlice<f32>,
25339 y: &mut CudaSlice<f32>,
25340 in_f: usize,
25341 out_f: usize,
25342 t: usize,
25343 ) -> Result<(), Box<dyn std::error::Error>> {
25344 let (grid_x, t_max) = match kernel {
25345 "matvec_bf16_f32acc_x1_tcols_rf" | "matvec_bf16_f32acc_x1_tcols_rf_redshift" => {
25346 (out_f as u32, 8usize)
25347 }
25348 "matvec_bf16_f32acc_x4_tcols_rf" => (out_f.div_ceil(4) as u32, 8usize),
25349 "matvec_bf16_f32acc_x4_tcols16_rf" => (out_f.div_ceil(4) as u32, 16usize),
25350 _ => return Err("matvec_bf16_tcols_gate_kernel_into: unknown kernel".into()),
25351 };
25352 if x.len() < t * in_f
25353 || y.len() < t * out_f
25354 || !(1..=t_max).contains(&t)
25355 || !in_f.is_multiple_of(8)
25356 || !mmv_block().is_power_of_two()
25357 {
25358 return Err("matvec_bf16_tcols_gate_kernel geometry".into());
25359 }
25360 let f = self.func(kernel);
25361 let cfg = LaunchConfig {
25362 grid_dim: (grid_x, 1, 1),
25363 block_dim: (mmv_block(), 1, 1),
25364 shared_mem_bytes: (t as u32) * mmv_block() * 4,
25365 };
25366 let (ini, outi, ti) = (in_f as i32, out_f as i32, t as i32);
25367 let __s_b = self.gpu.stream();
25368 let mut b = __s_b.launch_builder(&f);
25369 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi).arg(&ti);
25370 unsafe {
25371 b.launch(cfg)?;
25372 }
25373 Ok(())
25374 }
25375
25376 pub fn matmul_rows_exact(
25384 &self,
25385 w: &crate::model::GpuTensor,
25386 x: &CudaSlice<f32>,
25387 m: usize,
25388 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25389 use crate::model::GpuTensor;
25390 if let GpuTensor::FloatBf16 { data, .. } = w
25395 && (2..=32).contains(&m)
25396 && Self::bf16_mmv_on()
25397 && w.in_features().is_multiple_of(32)
25398 && w.out_features() >= 64
25399 && glm5_w8_on()
25400 {
25401 let (in_f, out_f) = (w.in_features(), w.out_features());
25402 let mut y = self.vws_uninit(m * out_f)?;
25403 if let Some(()) = self.matvec_bf16_via_q8_mirror_t(data, x, &mut y, in_f, out_f, m)? {
25404 if GLM5_W8_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
25405 eprintln!(
25406 "[glm5-w8] engaged rows-exact m={m} in_f={in_f} out_f={out_f} \
25407 (q8_0 mirror, MEMRA_GLM5_W8=1)"
25408 );
25409 }
25410 return Ok(y);
25411 }
25412 }
25413 if let GpuTensor::FloatBf16 { data, .. } = w
25414 && (2..=8).contains(&m)
25415 && Self::bf16_mmv_on()
25416 && w.in_features().is_multiple_of(8)
25417 && !(step_tp_w8_on() && w8_hybrid_on())
25418 && !glm5_w8_on()
25419 {
25420 let (in_f, out_f) = (w.in_features(), w.out_features());
25421 let mut y = self.vws_uninit(m * out_f)?;
25424 self.matvec_bf16_tcols_into(data, x, &mut y, in_f, out_f, m)?;
25425 return Ok(y);
25426 }
25427 self.matmul_decode_exact(w, x, m)
25428 }
25429
25430 #[allow(clippy::too_many_arguments)] pub fn matvec_bf16_dual_silu_into(
25432 &self,
25433 wg: &CudaSlice<u8>,
25434 wu: &CudaSlice<u8>,
25435 x: &CudaSlice<f32>,
25436 act: &mut CudaSlice<f32>,
25437 in_f: usize,
25438 out_f: usize,
25439 limit: Option<f32>,
25440 ) -> Result<(), Box<dyn std::error::Error>> {
25441 if wg.len() != in_f * out_f * 2
25442 || wu.len() != in_f * out_f * 2
25443 || x.len() < in_f
25444 || !in_f.is_multiple_of(8)
25445 || act.len() < out_f
25446 {
25447 return Err("matvec_bf16_dual_silu geometry".into());
25448 }
25449 let f = self.func("matvec_bf16_dual_silu");
25450 let cfg = LaunchConfig {
25451 grid_dim: (out_f as u32, 1, 1),
25452 block_dim: (mmv_block(), 1, 1),
25453 shared_mem_bytes: 0,
25454 };
25455 let (ini, outi) = (in_f as i32, out_f as i32);
25456 let lim = limit.unwrap_or(0.0);
25457 let __s_b = self.gpu.stream();
25458 let mut b = __s_b.launch_builder(&f);
25459 b.arg(wg)
25460 .arg(wu)
25461 .arg(x)
25462 .arg(act)
25463 .arg(&ini)
25464 .arg(&outi)
25465 .arg(&lim);
25466 unsafe {
25467 b.launch(cfg)?;
25468 }
25469 Ok(())
25470 }
25471
25472 #[allow(clippy::too_many_arguments)]
25475 pub fn matvec_bf16_dual_view_into(
25476 &self,
25477 wg: &cudarc::driver::CudaView<'_, u8>,
25478 wu: &cudarc::driver::CudaView<'_, u8>,
25479 x: &CudaSlice<f32>,
25480 yg: &mut CudaSlice<f32>,
25481 yu: &mut CudaSlice<f32>,
25482 in_f: usize,
25483 out_f: usize,
25484 ) -> Result<(), Box<dyn std::error::Error>> {
25485 if wg.len() != in_f * out_f * 2
25486 || wu.len() != in_f * out_f * 2
25487 || x.len() < in_f
25488 || !in_f.is_multiple_of(8)
25489 || yg.len() < out_f
25490 || yu.len() < out_f
25491 {
25492 return Err(format!(
25493 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
25494 wg.len(),
25495 wu.len(),
25496 x.len()
25497 )
25498 .into());
25499 }
25500 let f = self.func("matvec_bf16_dual");
25501 let cfg = LaunchConfig {
25502 grid_dim: ((2 * out_f) as u32, 1, 1),
25503 block_dim: (mmv_block(), 1, 1),
25504 shared_mem_bytes: 0,
25505 };
25506 let (ini, outi) = (in_f as i32, out_f as i32);
25507 let __s_b = self.gpu.stream();
25508 let mut b = __s_b.launch_builder(&f);
25509 b.arg(wg)
25510 .arg(wu)
25511 .arg(x)
25512 .arg(yg)
25513 .arg(yu)
25514 .arg(&ini)
25515 .arg(&outi);
25516 unsafe {
25517 b.launch(cfg)?;
25518 }
25519 Ok(())
25520 }
25521
25522 #[allow(clippy::too_many_arguments)]
25524 pub fn matvec_bf16_dual_into(
25525 &self,
25526 wg: &CudaSlice<u8>,
25527 wu: &CudaSlice<u8>,
25528 x: &CudaSlice<f32>,
25529 yg: &mut CudaSlice<f32>,
25530 yu: &mut CudaSlice<f32>,
25531 in_f: usize,
25532 out_f: usize,
25533 ) -> Result<(), Box<dyn std::error::Error>> {
25534 if wg.len() != in_f * out_f * 2
25535 || wu.len() != in_f * out_f * 2
25536 || x.len() < in_f
25537 || !in_f.is_multiple_of(8)
25538 || yg.len() < out_f
25539 || yu.len() < out_f
25540 {
25541 return Err(format!(
25542 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
25543 wg.len(),
25544 wu.len(),
25545 x.len()
25546 )
25547 .into());
25548 }
25549 let f = self.func("matvec_bf16_dual");
25550 let cfg = LaunchConfig {
25551 grid_dim: ((2 * out_f) as u32, 1, 1),
25552 block_dim: (mmv_block(), 1, 1),
25553 shared_mem_bytes: 0,
25554 };
25555 let (ini, outi) = (in_f as i32, out_f as i32);
25556 let __s_b = self.gpu.stream();
25557 let mut b = __s_b.launch_builder(&f);
25558 b.arg(wg)
25559 .arg(wu)
25560 .arg(x)
25561 .arg(yg)
25562 .arg(yu)
25563 .arg(&ini)
25564 .arg(&outi);
25565 unsafe {
25566 b.launch(cfg)?;
25567 }
25568 Ok(())
25569 }
25570
25571 #[allow(dead_code)] pub(crate) fn matvec_bf16_dual(
25575 &self,
25576 wg: &CudaSlice<u8>,
25577 wu: &CudaSlice<u8>,
25578 x: &CudaSlice<f32>,
25579 in_f: usize,
25580 out_f: usize,
25581 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25582 if wg.len() != in_f * out_f * 2
25583 || wu.len() != in_f * out_f * 2
25584 || x.len() < in_f
25585 || !in_f.is_multiple_of(8)
25586 {
25587 return Err(format!(
25588 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
25589 wg.len(),
25590 wu.len(),
25591 x.len()
25592 )
25593 .into());
25594 }
25595 let mut yg = self.alloc_uninit::<f32>(out_f)?;
25596 let mut yu = self.alloc_uninit::<f32>(out_f)?;
25597 let f = self.func("matvec_bf16_dual");
25598 let cfg = LaunchConfig {
25599 grid_dim: ((2 * out_f) as u32, 1, 1),
25600 block_dim: (mmv_block(), 1, 1),
25601 shared_mem_bytes: 0,
25602 };
25603 let (ini, outi) = (in_f as i32, out_f as i32);
25604 let __s_b = self.gpu.stream();
25605 let mut b = __s_b.launch_builder(&f);
25606 b.arg(wg)
25607 .arg(wu)
25608 .arg(x)
25609 .arg(&mut yg)
25610 .arg(&mut yu)
25611 .arg(&ini)
25612 .arg(&outi);
25613 unsafe {
25614 b.launch(cfg)?;
25615 }
25616 Ok((yg, yu))
25617 }
25618
25619 #[allow(clippy::too_many_arguments)]
25620 #[allow(clippy::manual_is_multiple_of)] fn linear_bf16_chunked_inner(
25622 &self,
25623 x: &CudaSlice<f32>,
25624 data: &CudaSlice<u8>,
25625 m: usize,
25626 in_f: usize,
25627 out_f: usize,
25628 exact: bool,
25629 canonical_chunk_rows: Option<usize>,
25630 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25631 const CHUNK_BYTES: usize = 256 << 20;
25632 if m == 1
25635 && !exact
25636 && canonical_chunk_rows.is_none()
25637 && in_f.is_multiple_of(8)
25638 && Self::bf16_mmv_on()
25639 {
25640 return self.matvec_bf16(data, x, in_f, out_f);
25641 }
25642 if m >= 16
25647 && !exact
25648 && canonical_chunk_rows.is_none()
25649 && data.len() == in_f * out_f * 2
25650 && crate::f16_ffi::pp_bf16_enabled()
25651 {
25652 if let Some(y) = self.bf16_tc_gemm(data, x, m, in_f, out_f)? {
25655 return Ok(y);
25656 }
25657 }
25658 let row_bytes = in_f
25659 .checked_mul(std::mem::size_of::<f32>())
25660 .ok_or("BF16 chunk row byte count overflow")?;
25661 if row_bytes == 0 || out_f == 0 {
25662 return Err("BF16 chunk dimensions must be nonzero".into());
25663 }
25664 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
25665 let chunk_rows = match canonical_chunk_rows {
25666 Some(0) => {
25667 return Err("canonical BF16 chunk rows must be nonzero".into());
25668 }
25669 Some(rows) if rows > max_chunk_rows => {
25670 return Err(format!(
25671 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
25672 )
25673 .into());
25674 }
25675 Some(rows) if out_f % rows != 0 => {
25676 return Err(format!(
25677 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
25678 )
25679 .into());
25680 }
25681 Some(rows) => rows,
25682 None => max_chunk_rows,
25683 };
25684 if chunk_rows >= out_f {
25685 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
25686 return if exact {
25687 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
25688 } else {
25689 self.linear(x, &wf32, m, in_f, out_f)
25690 };
25691 }
25692 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
25693 let mut r0 = 0usize;
25694 while r0 < out_f {
25695 let rows = chunk_rows.min(out_f - r0);
25696 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
25697 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
25698 let yc = if exact {
25699 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
25700 } else {
25701 self.linear(x, &wf32, m, in_f, rows)?
25702 };
25703 for mi in 0..m {
25705 let src = yc.slice(mi * rows..(mi + 1) * rows);
25706 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
25707 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
25708 }
25709 r0 += rows;
25710 }
25711 Ok(y)
25712 }
25713
25714 pub fn linear_bf16_resident(
25718 &self,
25719 x: &CudaSlice<f32>,
25720 data: &CudaSlice<u8>,
25721 m: usize,
25722 in_f: usize,
25723 out_f: usize,
25724 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25725 if data.len() != in_f * out_f * 2 {
25726 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
25727 }
25728 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
25729 }
25730
25731 pub fn linear_bf16_resident_canonical_rows(
25737 &self,
25738 x: &CudaSlice<f32>,
25739 data: &CudaSlice<u8>,
25740 m: usize,
25741 in_f: usize,
25742 out_f: usize,
25743 canonical_chunk_rows: usize,
25744 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25745 if data.len() != in_f * out_f * 2 {
25746 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
25747 }
25748 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
25749 }
25750
25751 pub fn linear_f32_resident_canonical_rows(
25756 &self,
25757 x: &CudaSlice<f32>,
25758 data: &CudaSlice<f32>,
25759 m: usize,
25760 in_f: usize,
25761 out_f: usize,
25762 canonical_chunk_rows: usize,
25763 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25764 self.linear_f32_resident_canonical_rows_inner(
25765 x,
25766 data,
25767 m,
25768 in_f,
25769 out_f,
25770 canonical_chunk_rows,
25771 false,
25772 )
25773 }
25774
25775 pub fn linear_f32_resident_canonical_rows_strided(
25781 &self,
25782 x: &CudaSlice<f32>,
25783 data: &CudaSlice<f32>,
25784 m: usize,
25785 in_f: usize,
25786 out_f: usize,
25787 canonical_chunk_rows: usize,
25788 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25789 self.linear_f32_resident_canonical_rows_inner(
25790 x,
25791 data,
25792 m,
25793 in_f,
25794 out_f,
25795 canonical_chunk_rows,
25796 true,
25797 )
25798 }
25799
25800 #[allow(clippy::too_many_arguments)]
25801 #[allow(clippy::manual_is_multiple_of)] fn linear_f32_resident_canonical_rows_inner(
25804 &self,
25805 x: &CudaSlice<f32>,
25806 data: &CudaSlice<f32>,
25807 m: usize,
25808 in_f: usize,
25809 out_f: usize,
25810 canonical_chunk_rows: usize,
25811 strided_output: bool,
25812 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25813 if data.len() != in_f * out_f {
25814 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
25815 }
25816 if canonical_chunk_rows == 0
25817 || canonical_chunk_rows > out_f
25818 || out_f % canonical_chunk_rows != 0
25819 {
25820 return Err(format!(
25821 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
25822 )
25823 .into());
25824 }
25825 if canonical_chunk_rows == out_f {
25826 return self.linear(x, data, m, in_f, out_f);
25827 }
25828
25829 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
25830 let input = x.slice(0..x.len());
25831 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
25832 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
25833 if m == 1 {
25834 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
25835 self.linear_device_into(
25836 &input,
25837 &weights,
25838 &mut destination,
25839 1,
25840 in_f,
25841 canonical_chunk_rows,
25842 )?;
25843 continue;
25844 }
25845 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
25846 if strided_output {
25847 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
25848 } else {
25849 for token in 0..m {
25850 let source = chunk
25851 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
25852 let mut destination =
25853 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
25854 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
25855 }
25856 }
25857 }
25858 Ok(y)
25859 }
25860
25861 #[allow(clippy::manual_is_multiple_of)] pub fn linear_f32_resident_canonical_rows_t1_into(
25867 &self,
25868 x: &CudaSlice<f32>,
25869 data: &CudaSlice<f32>,
25870 y: &mut CudaSlice<f32>,
25871 in_f: usize,
25872 out_f: usize,
25873 canonical_chunk_rows: usize,
25874 ) -> Result<(), Box<dyn std::error::Error>> {
25875 if data.len() != in_f * out_f {
25876 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
25877 }
25878 if y.len() != out_f || x.len() != in_f {
25879 return Err(format!(
25880 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
25881 x.len(),
25882 y.len()
25883 )
25884 .into());
25885 }
25886 if canonical_chunk_rows == 0
25887 || canonical_chunk_rows > out_f
25888 || out_f % canonical_chunk_rows != 0
25889 {
25890 return Err(format!(
25891 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
25892 )
25893 .into());
25894 }
25895 let input = x.slice(0..x.len());
25896 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
25897 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
25898 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
25899 self.linear_device_into(
25900 &input,
25901 &weights,
25902 &mut destination,
25903 1,
25904 in_f,
25905 canonical_chunk_rows,
25906 )?;
25907 }
25908 Ok(())
25909 }
25910
25911 pub fn linear_t1_into(
25914 &self,
25915 x: &cudarc::driver::CudaView<'_, f32>,
25916 w: &cudarc::driver::CudaView<'_, f32>,
25917 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
25918 in_f: usize,
25919 out_f: usize,
25920 ) -> Result<(), Box<dyn std::error::Error>> {
25921 self.linear_device_into(x, w, y, 1, in_f, out_f)
25922 }
25923
25924 pub fn linear_decode_exact(
25931 &self,
25932 x: &CudaSlice<f32>,
25933 w: &CudaSlice<f32>,
25934 m_tokens: usize,
25935 in_f: usize,
25936 out_f: usize,
25937 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25938 if m_tokens == 1 {
25939 return self.linear(x, w, 1, in_f, out_f);
25940 }
25941 let xv = self.view(x, m_tokens * in_f);
25942 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
25943 for t in 0..m_tokens {
25944 let row = xv.slice(t * in_f..(t + 1) * in_f);
25945 let mut xr = self.alloc_uninit::<f32>(in_f)?;
25946 self.copy_view_into(&mut xr, 0, &row, in_f)?;
25947 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
25948 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
25949 }
25950 Ok(y)
25951 }
25952
25953 pub fn linear(
25954 &self,
25955 x: &CudaSlice<f32>,
25956 w: &CudaSlice<f32>,
25957 m_tokens: usize,
25958 in_f: usize,
25959 out_f: usize,
25960 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25961 self.linear_device(x, w, m_tokens, in_f, out_f)
25962 }
25963
25964 fn linear_device<I>(
25965 &self,
25966 x: &I,
25967 w: &I,
25968 m_tokens: usize,
25969 in_f: usize,
25970 out_f: usize,
25971 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
25972 where
25973 I: cudarc::driver::DevicePtr<f32>,
25974 {
25975 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)?;
25977 Ok(c)
25978 }
25979
25980 fn linear_device_into<I, O>(
25981 &self,
25982 x: &I,
25983 w: &I,
25984 c: &mut O,
25985 m_tokens: usize,
25986 in_f: usize,
25987 out_f: usize,
25988 ) -> Result<(), Box<dyn std::error::Error>>
25989 where
25990 I: cudarc::driver::DevicePtr<f32>,
25991 O: cudarc::driver::DevicePtrMut<f32>,
25992 {
25993 use cudarc::cublaslt::{Matmul, MatmulConfig};
25994 let cfg = MatmulConfig {
25995 transa: true,
25996 transb: false,
25997 transc: false,
25998 m: out_f as u64,
25999 n: m_tokens as u64,
26000 k: in_f as u64,
26001 alpha: 1.0,
26002 lda: in_f as i64,
26003 ldb: in_f as i64,
26004 beta: 0.0,
26005 ldc: out_f as i64,
26006 stride_a: None,
26007 stride_b: None,
26008 stride_c: None,
26009 stride_bias: None,
26010 batch_size: None,
26011 };
26012 let blas = self.gpu.blas();
26013 unsafe {
26014 blas.matmul(cfg, w, x, c, None, None)?;
26015 }
26016 Ok(())
26017 }
26018
26019 #[allow(clippy::too_many_arguments)] pub fn sdpa_naive(
26029 &self,
26030 q: &CudaSlice<f32>,
26031 k: &CudaSlice<f32>,
26032 v: &CudaSlice<f32>,
26033 o: &mut CudaSlice<f32>,
26034 head_dim: usize,
26035 n_head: usize,
26036 n_head_kv: usize,
26037 t: usize,
26038 t_kv: usize,
26039 scale: f32,
26040 causal: bool,
26041 ) -> Result<(), Box<dyn std::error::Error>> {
26042 if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
26043 return self.sdpa_naive_gmem(
26044 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26045 );
26046 }
26047 let f = self.func("sdpa_naive_f32");
26048 let cfg = LaunchConfig {
26049 grid_dim: (n_head as u32, t as u32, 1),
26050 block_dim: (128, 1, 1),
26051 shared_mem_bytes: (t_kv * 4) as u32,
26052 };
26053 let (hd, nh, nhkv, ti, tkvi, cz) = (
26054 head_dim as i32,
26055 n_head as i32,
26056 n_head_kv as i32,
26057 t as i32,
26058 t_kv as i32,
26059 causal as i32,
26060 );
26061 let __s_b = self.gpu.stream();
26062 let mut b = __s_b.launch_builder(&f);
26063 b.arg(q)
26064 .arg(k)
26065 .arg(v)
26066 .arg(o)
26067 .arg(&hd)
26068 .arg(&nh)
26069 .arg(&nhkv)
26070 .arg(&ti)
26071 .arg(&tkvi)
26072 .arg(&scale)
26073 .arg(&cz);
26074 unsafe {
26075 b.launch(cfg)?;
26076 }
26077 Ok(())
26078 }
26079
26080 #[allow(clippy::too_many_arguments)]
26089 pub fn sdpa_naive_gmem(
26090 &self,
26091 q: &CudaSlice<f32>,
26092 k: &CudaSlice<f32>,
26093 v: &CudaSlice<f32>,
26094 o: &mut CudaSlice<f32>,
26095 head_dim: usize,
26096 n_head: usize,
26097 n_head_kv: usize,
26098 t: usize,
26099 t_kv: usize,
26100 scale: f32,
26101 causal: bool,
26102 ) -> Result<(), Box<dyn std::error::Error>> {
26103 let ws_len = n_head
26104 .checked_mul(t)
26105 .and_then(|x| x.checked_mul(t_kv))
26106 .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
26107 let ws_bytes = ws_len
26108 .checked_mul(std::mem::size_of::<f32>())
26109 .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
26110 if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
26111 return Err(format!(
26112 "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
26113 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
26114 needs a tiled/flash kernel, not the naive oracle"
26115 )
26116 .into());
26117 }
26118 let mut scores = self.uninit(ws_len)?;
26119 let f = self.func("sdpa_naive_gmem_f32");
26120 let cfg = LaunchConfig {
26121 grid_dim: (n_head as u32, t as u32, 1),
26122 block_dim: (128, 1, 1),
26123 shared_mem_bytes: 0,
26124 };
26125 let (hd, nh, nhkv, ti, tkvi, cz) = (
26126 head_dim as i32,
26127 n_head as i32,
26128 n_head_kv as i32,
26129 t as i32,
26130 t_kv as i32,
26131 causal as i32,
26132 );
26133 let __s_b = self.gpu.stream();
26134 let mut b = __s_b.launch_builder(&f);
26135 b.arg(q)
26136 .arg(k)
26137 .arg(v)
26138 .arg(o)
26139 .arg(&mut scores)
26140 .arg(&hd)
26141 .arg(&nh)
26142 .arg(&nhkv)
26143 .arg(&ti)
26144 .arg(&tkvi)
26145 .arg(&scale)
26146 .arg(&cz);
26147 unsafe {
26148 b.launch(cfg)?;
26149 }
26150 Ok(())
26151 }
26152
26153 #[allow(clippy::too_many_arguments)]
26158 pub fn sdpa_naive_island(
26159 &self,
26160 q: &CudaSlice<f32>,
26161 k: &CudaSlice<f32>,
26162 v: &CudaSlice<f32>,
26163 o: &mut CudaSlice<f32>,
26164 span_id: &CudaSlice<i32>,
26165 head_dim: usize,
26166 n_head: usize,
26167 n_head_kv: usize,
26168 t: usize,
26169 t_kv: usize,
26170 scale: f32,
26171 window: usize,
26172 ) -> Result<(), Box<dyn std::error::Error>> {
26173 let f = self.func("sdpa_naive_island_f32");
26174 let cfg = LaunchConfig {
26175 grid_dim: (n_head as u32, t as u32, 1),
26176 block_dim: (128, 1, 1),
26177 shared_mem_bytes: (t_kv * 4) as u32,
26178 };
26179 let (hd, nh, nhkv, ti, tkvi, wi) = (
26180 head_dim as i32,
26181 n_head as i32,
26182 n_head_kv as i32,
26183 t as i32,
26184 t_kv as i32,
26185 window as i32,
26186 );
26187 let __s_b = self.gpu.stream();
26188 let mut b = __s_b.launch_builder(&f);
26189 b.arg(q)
26190 .arg(k)
26191 .arg(v)
26192 .arg(o)
26193 .arg(span_id)
26194 .arg(&hd)
26195 .arg(&nh)
26196 .arg(&nhkv)
26197 .arg(&ti)
26198 .arg(&tkvi)
26199 .arg(&scale)
26200 .arg(&wi);
26201 unsafe {
26202 b.launch(cfg)?;
26203 }
26204 Ok(())
26205 }
26206
26207 #[allow(clippy::too_many_arguments)]
26209 pub fn sdpa_naive_w(
26210 &self,
26211 q: &CudaSlice<f32>,
26212 k: &CudaSlice<f32>,
26213 v: &CudaSlice<f32>,
26214 o: &mut CudaSlice<f32>,
26215 head_dim: usize,
26216 n_head: usize,
26217 n_head_kv: usize,
26218 t: usize,
26219 t_kv: usize,
26220 scale: f32,
26221 causal: bool,
26222 window: usize,
26223 ) -> Result<(), Box<dyn std::error::Error>> {
26224 let f = self.func("sdpa_naive_w_f32");
26225 let cfg = LaunchConfig {
26226 grid_dim: (n_head as u32, t as u32, 1),
26227 block_dim: (128, 1, 1),
26228 shared_mem_bytes: (t_kv * 4) as u32,
26229 };
26230 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26231 head_dim as i32,
26232 n_head as i32,
26233 n_head_kv as i32,
26234 t as i32,
26235 t_kv as i32,
26236 causal as i32,
26237 window as i32,
26238 );
26239 let __s_b = self.gpu.stream();
26240 let mut b = __s_b.launch_builder(&f);
26241 b.arg(q)
26242 .arg(k)
26243 .arg(v)
26244 .arg(o)
26245 .arg(&hd)
26246 .arg(&nh)
26247 .arg(&nhkv)
26248 .arg(&ti)
26249 .arg(&tkvi)
26250 .arg(&scale)
26251 .arg(&cz)
26252 .arg(&wi);
26253 unsafe {
26254 b.launch(cfg)?;
26255 }
26256 Ok(())
26257 }
26258
26259 #[allow(clippy::too_many_arguments)]
26275 pub fn sdpa_naive_w_lo(
26276 &self,
26277 q: &CudaSlice<f32>,
26278 k: &CudaSlice<f32>,
26279 v: &CudaSlice<f32>,
26280 o: &mut CudaSlice<f32>,
26281 head_dim: usize,
26282 n_head: usize,
26283 n_head_kv: usize,
26284 t: usize,
26285 t_kv: usize,
26286 scale: f32,
26287 causal: bool,
26288 window: usize,
26289 kv_floor: usize,
26290 ) -> Result<(), Box<dyn std::error::Error>> {
26291 if kv_floor > t_kv - t {
26292 return Err(format!(
26293 "sdpa_naive_w_lo: kv_floor {kv_floor} would clip the query rows themselves \
26294 (t_kv {t_kv} - t {t})"
26295 )
26296 .into());
26297 }
26298 let kv_lo = if window > 0 {
26299 (t_kv - t + 1).saturating_sub(window)
26300 } else {
26301 0
26302 }
26303 .max(kv_floor);
26304 let smem = (t_kv - kv_lo) * 4;
26305 if smem > 48 * 1024 {
26306 return Err(format!(
26307 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
26308 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
26309 a window this wide needs the multi-pass long-ctx kernel"
26310 )
26311 .into());
26312 }
26313 let f = self.func("sdpa_naive_w_lo_f32");
26314 let cfg = LaunchConfig {
26315 grid_dim: (n_head as u32, t as u32, 1),
26316 block_dim: (128, 1, 1),
26317 shared_mem_bytes: smem as u32,
26318 };
26319 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
26320 head_dim as i32,
26321 n_head as i32,
26322 n_head_kv as i32,
26323 t as i32,
26324 t_kv as i32,
26325 causal as i32,
26326 window as i32,
26327 kv_lo as i32,
26328 );
26329 let __s_b = self.gpu.stream();
26330 let mut b = __s_b.launch_builder(&f);
26331 b.arg(q)
26332 .arg(k)
26333 .arg(v)
26334 .arg(o)
26335 .arg(&hd)
26336 .arg(&nh)
26337 .arg(&nhkv)
26338 .arg(&ti)
26339 .arg(&tkvi)
26340 .arg(&scale)
26341 .arg(&cz)
26342 .arg(&wi)
26343 .arg(&lo);
26344 unsafe {
26345 b.launch(cfg)?;
26346 }
26347 Ok(())
26348 }
26349
26350 #[allow(clippy::too_many_arguments)] pub fn sdpa_naive_view(
26353 &self,
26354 q: &CudaSlice<f32>,
26355 k: &cudarc::driver::CudaView<f32>,
26356 v: &cudarc::driver::CudaView<f32>,
26357 o: &mut CudaSlice<f32>,
26358 head_dim: usize,
26359 n_head: usize,
26360 n_head_kv: usize,
26361 t: usize,
26362 t_kv: usize,
26363 scale: f32,
26364 causal: bool,
26365 ) -> Result<(), Box<dyn std::error::Error>> {
26366 let f = self.func("sdpa_naive_f32");
26367 let cfg = LaunchConfig {
26368 grid_dim: (n_head as u32, t as u32, 1),
26369 block_dim: (128, 1, 1),
26370 shared_mem_bytes: (t_kv * 4) as u32,
26371 };
26372 let (hd, nh, nhkv, ti, tkvi, cz) = (
26373 head_dim as i32,
26374 n_head as i32,
26375 n_head_kv as i32,
26376 t as i32,
26377 t_kv as i32,
26378 causal as i32,
26379 );
26380 let __s_b = self.gpu.stream();
26381 let mut b = __s_b.launch_builder(&f);
26382 b.arg(q)
26383 .arg(k)
26384 .arg(v)
26385 .arg(o)
26386 .arg(&hd)
26387 .arg(&nh)
26388 .arg(&nhkv)
26389 .arg(&ti)
26390 .arg(&tkvi)
26391 .arg(&scale)
26392 .arg(&cz);
26393 unsafe {
26394 b.launch(cfg)?;
26395 }
26396 Ok(())
26397 }
26398
26399 #[allow(clippy::too_many_arguments)]
26407 pub fn fa_dequant_kv_view_f32(
26408 &self,
26409 k: &cudarc::driver::CudaView<u8>,
26410 v: &cudarc::driver::CudaView<u8>,
26411 kf: &mut CudaSlice<f32>,
26412 vf: &mut CudaSlice<f32>,
26413 kv_dim_k: usize,
26414 kv_dim_v: usize,
26415 t_kv: usize,
26416 k_tok_bytes: usize,
26417 v_tok_bytes: usize,
26418 g: bool,
26419 ) -> Result<(), Box<dyn std::error::Error>> {
26420 let f = if g {
26421 self.func_g("fa_dequant_kv_ws_f32")
26422 } else {
26423 self.func("fa_dequant_kv_ws_f32")
26424 };
26425 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
26426 #[allow(clippy::manual_div_ceil)]
26427 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26429 let cfg = LaunchConfig {
26430 grid_dim: (nblk.max(1), 1, 1),
26431 block_dim: (256, 1, 1),
26432 shared_mem_bytes: 0,
26433 };
26434 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
26435 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
26436 let __s_b = self.gpu.stream();
26437 let mut b = __s_b.launch_builder(&f);
26438 b.arg(k)
26439 .arg(v)
26440 .arg(&mut *kf)
26441 .arg(&mut *vf)
26442 .arg(&kdk)
26443 .arg(&kdv)
26444 .arg(&tkvi)
26445 .arg(&ktb)
26446 .arg(&vtb);
26447 unsafe {
26448 b.launch(cfg)?;
26449 }
26450 Ok(())
26451 }
26452
26453 #[allow(clippy::too_many_arguments)]
26454 pub fn sdpa_naive_quantized_view(
26455 &self,
26456 q: &CudaSlice<f32>,
26457 k: &cudarc::driver::CudaView<u8>,
26458 v: &cudarc::driver::CudaView<u8>,
26459 o: &mut CudaSlice<f32>,
26460 head_dim: usize,
26461 n_head: usize,
26462 n_head_kv: usize,
26463 t: usize,
26464 t_kv: usize,
26465 scale: f32,
26466 causal: bool,
26467 k_tok_bytes: usize,
26468 v_tok_bytes: usize,
26469 ) -> Result<(), Box<dyn std::error::Error>> {
26470 let kv_dim = n_head_kv * head_dim;
26471 let mut kf = self.uninit(t_kv * kv_dim)?;
26472 let mut vf = self.uninit(t_kv * kv_dim)?;
26473 let f = self.func("fa_dequant_kv_ws_f32");
26474 let total = (2 * t_kv * kv_dim) as u64;
26475 #[allow(clippy::manual_div_ceil)]
26476 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26478 let cfg = LaunchConfig {
26479 grid_dim: (nblk.max(1), 1, 1),
26480 block_dim: (256, 1, 1),
26481 shared_mem_bytes: 0,
26482 };
26483 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
26484 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
26485 let __s_b = self.gpu.stream();
26486 let mut b = __s_b.launch_builder(&f);
26487 b.arg(k)
26488 .arg(v)
26489 .arg(&mut kf)
26490 .arg(&mut vf)
26491 .arg(&kv_dim_i)
26492 .arg(&kv_dim_i)
26493 .arg(&t_kv_i)
26494 .arg(&k_tok_bytes_i)
26495 .arg(&v_tok_bytes_i);
26496 unsafe { b.launch(cfg)? };
26497 self.sdpa_naive(
26498 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26499 )
26500 }
26501
26502 #[allow(clippy::too_many_arguments)]
26514 pub fn sdpa_naive_w_quantized_view(
26515 &self,
26516 q: &CudaSlice<f32>,
26517 k: &cudarc::driver::CudaView<u8>,
26518 v: &cudarc::driver::CudaView<u8>,
26519 o: &mut CudaSlice<f32>,
26520 head_dim: usize,
26521 n_head: usize,
26522 n_head_kv: usize,
26523 t: usize,
26524 t_kv: usize,
26525 scale: f32,
26526 causal: bool,
26527 window: usize,
26528 k_tok_bytes: usize,
26529 v_tok_bytes: usize,
26530 ) -> Result<(), Box<dyn std::error::Error>> {
26531 let kv_dim = n_head_kv * head_dim;
26532 let mut kf = self.uninit(t_kv * kv_dim)?;
26533 let mut vf = self.uninit(t_kv * kv_dim)?;
26534 let f = self.func("fa_dequant_kv_ws_f32");
26535 let total = (2 * t_kv * kv_dim) as u64;
26536 #[allow(clippy::manual_div_ceil)]
26537 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
26539 let cfg = LaunchConfig {
26540 grid_dim: (nblk.max(1), 1, 1),
26541 block_dim: (256, 1, 1),
26542 shared_mem_bytes: 0,
26543 };
26544 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
26545 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
26546 let __s_b = self.gpu.stream();
26547 let mut b = __s_b.launch_builder(&f);
26548 b.arg(k)
26549 .arg(v)
26550 .arg(&mut kf)
26551 .arg(&mut vf)
26552 .arg(&kv_dim_i)
26553 .arg(&kv_dim_i)
26554 .arg(&t_kv_i)
26555 .arg(&k_tok_bytes_i)
26556 .arg(&v_tok_bytes_i);
26557 unsafe { b.launch(cfg)? };
26558 self.sdpa_naive_w(
26559 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
26560 )
26561 }
26562
26563 #[allow(clippy::too_many_arguments)]
26567 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill(
26570 &self,
26571 q: &CudaSlice<f32>,
26572 k: &CudaSlice<f32>,
26573 v: &CudaSlice<f32>,
26574 o: &mut CudaSlice<f32>,
26575 head_dim: usize,
26576 n_head: usize,
26577 n_head_kv: usize,
26578 t: usize,
26579 t_kv: usize,
26580 scale: f32,
26581 causal: bool,
26582 ) -> Result<(), Box<dyn std::error::Error>> {
26583 if portable_mma_gated() {
26584 return self.sdpa_naive(
26585 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
26586 );
26587 }
26588 let fa3_on = head_dim == 256
26596 && causal
26597 && t == t_kv
26598 && match std::env::var("MEMRA_FA3").as_deref() {
26599 Ok("0") => false,
26600 Ok("1") => {
26604 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
26605 true
26606 }
26607 _ => cfg!(memra_hopper_mma),
26608 };
26609 if fa3_on {
26610 let n = t * n_head * head_dim;
26611 let nkv = t * n_head_kv * head_dim;
26612 let mut q16 = self.alloc_u8_uninit(n * 2)?;
26613 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
26614 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
26615 self.f32_to_bf16_into(q, &mut q16, n)?;
26616 self.f32_to_bf16_into(k, &mut k16, nkv)?;
26617 self.f32_to_bf16_into(v, &mut v16, nkv)?;
26618 let rc = {
26619 use cudarc::driver::{DevicePtr, DevicePtrMut};
26620 let stream = self.gpu.stream();
26621 let (qp, _g1) = q16.device_ptr(&stream);
26622 let (kp, _g2) = k16.device_ptr(&stream);
26623 let (vp, _g3) = v16.device_ptr(&stream);
26624 let (op, _g4) = o.device_ptr_mut(&stream);
26625 unsafe {
26626 memra_fa3_prefill(
26627 qp as *const core::ffi::c_void,
26628 kp as *const core::ffi::c_void,
26629 vp as *const core::ffi::c_void,
26630 op as *mut f32,
26631 t as i32,
26632 n_head as i32,
26633 n_head_kv as i32,
26634 head_dim as i32,
26635 scale,
26636 stream.cu_stream() as *mut core::ffi::c_void,
26637 )
26638 }
26639 };
26640 if rc != 0 {
26641 return Err(format!("memra_fa3_prefill rc={rc}").into());
26642 }
26643 return Ok(());
26644 }
26645 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26650 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
26651 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
26652 const BLOCK_Q: usize = 64;
26653 const BKX: usize = 32;
26654 let f = self.func("fa_prefill_bf16_p1");
26655 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
26656 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
26657 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26658 f.set_attribute(
26659 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26660 shmem as i32,
26661 )?;
26662 let cfg = LaunchConfig {
26663 grid_dim: (
26664 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
26665 n_head as u32,
26666 1,
26667 ),
26668 block_dim: (32, 4, 1),
26669 shared_mem_bytes: shmem,
26670 };
26671 let (hd, nh, nhkv, ti, tkvi, cz) = (
26672 head_dim as i32,
26673 n_head as i32,
26674 n_head_kv as i32,
26675 t as i32,
26676 t_kv as i32,
26677 causal as i32,
26678 );
26679 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
26680 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
26681 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
26682 let __s_b = self.gpu.stream();
26683 let mut b = __s_b.launch_builder(&f);
26684 b.arg(&qb)
26685 .arg(&kb)
26686 .arg(&vb)
26687 .arg(o)
26688 .arg(&hd)
26689 .arg(&nh)
26690 .arg(&nhkv)
26691 .arg(&ti)
26692 .arg(&tkvi)
26693 .arg(&scale)
26694 .arg(&cz);
26695 unsafe {
26696 b.launch(cfg)?;
26697 }
26698 return Ok(());
26699 }
26700 const BK: usize = 32;
26706 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
26709 let (block_q, warps, w2_sfx): (usize, u32, &str) =
26710 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
26711 let hd_sfx = fa_hd_suffix(head_dim)?;
26715 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
26716 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
26721 let (kb16, vb16) = if bf16kv {
26722 let n = t_kv * n_head_kv * head_dim;
26723 let mut kb = self.alloc_u8_uninit(n * 2)?;
26724 let mut vb = self.alloc_u8_uninit(n * 2)?;
26725 let fcv = self.func("f32_to_bf16_bulk");
26726 let ni = n as i64;
26727 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26728 let __s_b = self.gpu.stream();
26729 let mut b = __s_b.launch_builder(&fcv);
26730 b.arg(k).arg(&mut kb).arg(&ni);
26731 unsafe {
26732 b.launch(cfgc)?;
26733 }
26734 let __s_b = self.gpu.stream();
26735 let mut b = __s_b.launch_builder(&fcv);
26736 b.arg(v).arg(&mut vb).arg(&ni);
26737 unsafe {
26738 b.launch(cfgc)?;
26739 }
26740 (Some(kb), Some(vb))
26741 } else {
26742 (None, None)
26743 };
26744 let f = self.func(&if bf16kv {
26745 format!("fa_prefill_bf16kv_pp{hd_sfx}")
26746 } else {
26747 format!(
26748 "fa_prefill_f32{}{}{hd_sfx}",
26749 if floor { "" } else { "_pp" },
26750 if floor { "" } else { w2_sfx }
26751 )
26752 });
26753 let kv_stages = if bf16kv { 2 } else { 1 };
26756 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
26757 + 4 * (block_q * BK + 2 * block_q)) as u32;
26758 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26759 f.set_attribute(
26760 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26761 shmem as i32,
26762 )?;
26763 let cfg = LaunchConfig {
26764 grid_dim: (
26765 (t as u32 + block_q as u32 - 1) / block_q as u32,
26766 n_head as u32,
26767 1,
26768 ),
26769 block_dim: (32, warps, 1),
26770 shared_mem_bytes: shmem,
26771 };
26772 let (hd, nh, nhkv, ti, tkvi, cz) = (
26773 head_dim as i32,
26774 n_head as i32,
26775 n_head_kv as i32,
26776 t as i32,
26777 t_kv as i32,
26778 causal as i32,
26779 );
26780 let __s_b = self.gpu.stream();
26781 let mut b = __s_b.launch_builder(&f);
26782 b.arg(q);
26783 match (&kb16, &vb16) {
26784 (Some(kb), Some(vb)) => {
26785 b.arg(kb).arg(vb);
26786 }
26787 _ => {
26788 b.arg(k).arg(v);
26789 }
26790 }
26791 b.arg(o)
26792 .arg(&hd)
26793 .arg(&nh)
26794 .arg(&nhkv)
26795 .arg(&ti)
26796 .arg(&tkvi)
26797 .arg(&scale)
26798 .arg(&cz);
26799 unsafe {
26800 b.launch(cfg)?;
26801 }
26802 Ok(())
26803 }
26804
26805 #[allow(clippy::too_many_arguments)]
26809 pub fn fa_prefill_w(
26810 &self,
26811 q: &CudaSlice<f32>,
26812 k: &CudaSlice<f32>,
26813 v: &CudaSlice<f32>,
26814 o: &mut CudaSlice<f32>,
26815 head_dim: usize,
26816 n_head: usize,
26817 n_head_kv: usize,
26818 t: usize,
26819 t_kv: usize,
26820 scale: f32,
26821 causal: bool,
26822 window: usize,
26823 ) -> Result<(), Box<dyn std::error::Error>> {
26824 if portable_mma_gated() {
26827 return self.sdpa_naive_w(
26828 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
26829 );
26830 }
26831 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26835 let faw_f32 =
26836 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
26837 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
26838 self.fa_prefill_w_arm(
26839 q,
26840 k,
26841 v,
26842 o,
26843 head_dim,
26844 n_head,
26845 n_head_kv,
26846 t,
26847 t_kv,
26848 scale,
26849 causal,
26850 window,
26851 floor || faw_f32,
26852 floor,
26853 )
26854 }
26855
26856 #[allow(clippy::too_many_arguments)]
26859 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_w_pre(
26861 &self,
26862 qb: &CudaSlice<u8>,
26863 kb: &CudaSlice<u8>,
26864 vb: &CudaSlice<u8>,
26865 o: &mut CudaSlice<f32>,
26866 head_dim: usize,
26867 n_head: usize,
26868 n_head_kv: usize,
26869 t: usize,
26870 t_kv: usize,
26871 scale: f32,
26872 causal: bool,
26873 window: usize,
26874 v_f16: bool,
26875 ) -> Result<(), Box<dyn std::error::Error>> {
26876 const BLOCK_Q: usize = 64;
26877 const BK: usize = 32;
26878 debug_assert_eq!(head_dim, 256);
26879 let hp = fa_f16pv_on()
26880 && faw_hp_on()
26881 && n_head.is_multiple_of(2)
26882 && (n_head / n_head_kv).is_multiple_of(2);
26883 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
26884 if hp {
26885 const BLOCK_QH: usize = 32;
26886 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
26889 let vh: &CudaSlice<u8> = if v_f16 {
26890 vb
26891 } else {
26892 let n = t_kv * n_head_kv * head_dim;
26893 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
26894 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
26895 }
26896 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
26897 vguard.as_ref().unwrap()
26898 };
26899 let f = self.func("fa_prefill_w_bf16_p1h2");
26900 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
26901 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26902 f.set_attribute(
26903 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26904 shmem as i32,
26905 )?;
26906 let cfg = LaunchConfig {
26907 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
26908 block_dim: (32, 4, 1),
26909 shared_mem_bytes: shmem,
26910 };
26911 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26912 head_dim as i32,
26913 n_head as i32,
26914 n_head_kv as i32,
26915 t as i32,
26916 t_kv as i32,
26917 causal as i32,
26918 window as i32,
26919 );
26920 let __s_b = self.gpu.stream();
26921 let mut b = __s_b.launch_builder(&f);
26922 b.arg(qb)
26923 .arg(kb)
26924 .arg(vh)
26925 .arg(o)
26926 .arg(&hd)
26927 .arg(&nh)
26928 .arg(&nhkv)
26929 .arg(&ti)
26930 .arg(&tkvi)
26931 .arg(&scale)
26932 .arg(&cz)
26933 .arg(&wi);
26934 unsafe {
26935 b.launch(cfg)?;
26936 }
26937 return Ok(());
26938 }
26939 let f = self.func("fa_prefill_w_bf16_p1");
26940 let shmem =
26941 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
26942 use cudarc::driver::sys::CUfunction_attribute_enum as A;
26943 f.set_attribute(
26944 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26945 shmem as i32,
26946 )?;
26947 let cfg = LaunchConfig {
26948 grid_dim: (
26949 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
26950 n_head as u32,
26951 1,
26952 ),
26953 block_dim: (32, 4, 1),
26954 shared_mem_bytes: shmem,
26955 };
26956 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
26957 head_dim as i32,
26958 n_head as i32,
26959 n_head_kv as i32,
26960 t as i32,
26961 t_kv as i32,
26962 causal as i32,
26963 window as i32,
26964 );
26965 let __s_b = self.gpu.stream();
26966 let mut b = __s_b.launch_builder(&f);
26967 b.arg(qb)
26968 .arg(kb)
26969 .arg(vb)
26970 .arg(o)
26971 .arg(&hd)
26972 .arg(&nh)
26973 .arg(&nhkv)
26974 .arg(&ti)
26975 .arg(&tkvi)
26976 .arg(&scale)
26977 .arg(&cz)
26978 .arg(&wi);
26979 unsafe {
26980 b.launch(cfg)?;
26981 }
26982 Ok(())
26983 }
26984
26985 #[allow(clippy::too_many_arguments)]
26987 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_w_arm(
26989 &self,
26990 q: &CudaSlice<f32>,
26991 k: &CudaSlice<f32>,
26992 v: &CudaSlice<f32>,
26993 o: &mut CudaSlice<f32>,
26994 head_dim: usize,
26995 n_head: usize,
26996 n_head_kv: usize,
26997 t: usize,
26998 t_kv: usize,
26999 scale: f32,
27000 causal: bool,
27001 window: usize,
27002 f32_stage: bool,
27003 floor: bool,
27004 ) -> Result<(), Box<dyn std::error::Error>> {
27005 const BLOCK_Q: usize = 64;
27006 const BK: usize = 32;
27007 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
27008 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27012 let p1 = !floor
27013 && !f32_stage
27014 && *P1_ON.get_or_init(|| {
27015 std::env::var("MEMRA_FAW_P1")
27016 .map(|v| v != "0")
27017 .unwrap_or(true)
27018 });
27019 let hp = p1
27020 && fa_f16pv_on()
27021 && faw_hp_on()
27022 && n_head.is_multiple_of(2)
27023 && (n_head / n_head_kv).is_multiple_of(2);
27024 if hp {
27025 const BLOCK_QH: usize = 32;
27026 let f = self.func("fa_prefill_w_bf16_p1h2");
27027 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
27028 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27029 f.set_attribute(
27030 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27031 shmem as i32,
27032 )?;
27033 let cfg = LaunchConfig {
27034 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
27035 block_dim: (32, 4, 1),
27036 shared_mem_bytes: shmem,
27037 };
27038 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
27039 head_dim as i32,
27040 n_head as i32,
27041 n_head_kv as i32,
27042 t as i32,
27043 t_kv as i32,
27044 causal as i32,
27045 window as i32,
27046 );
27047 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27048 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27049 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
27050 let __s_b = self.gpu.stream();
27051 let mut b = __s_b.launch_builder(&f);
27052 b.arg(&qb)
27053 .arg(&kb)
27054 .arg(&vh)
27055 .arg(o)
27056 .arg(&hd)
27057 .arg(&nh)
27058 .arg(&nhkv)
27059 .arg(&ti)
27060 .arg(&tkvi)
27061 .arg(&scale)
27062 .arg(&cz)
27063 .arg(&wi);
27064 unsafe {
27065 b.launch(cfg)?;
27066 }
27067 return Ok(());
27068 }
27069 if p1 {
27070 let f = self.func("fa_prefill_w_bf16_p1");
27071 let shmem =
27072 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
27073 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27074 f.set_attribute(
27075 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27076 shmem as i32,
27077 )?;
27078 let cfg = LaunchConfig {
27079 grid_dim: (
27080 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27081 n_head as u32,
27082 1,
27083 ),
27084 block_dim: (32, 4, 1),
27085 shared_mem_bytes: shmem,
27086 };
27087 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
27088 head_dim as i32,
27089 n_head as i32,
27090 n_head_kv as i32,
27091 t as i32,
27092 t_kv as i32,
27093 causal as i32,
27094 window as i32,
27095 );
27096 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27097 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27098 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27099 let __s_b = self.gpu.stream();
27100 let mut b = __s_b.launch_builder(&f);
27101 b.arg(&qb)
27102 .arg(&kb)
27103 .arg(&vb)
27104 .arg(o)
27105 .arg(&hd)
27106 .arg(&nh)
27107 .arg(&nhkv)
27108 .arg(&ti)
27109 .arg(&tkvi)
27110 .arg(&scale)
27111 .arg(&cz)
27112 .arg(&wi);
27113 unsafe {
27114 b.launch(cfg)?;
27115 }
27116 return Ok(());
27117 }
27118 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27121 let g4 = !floor
27122 && !f32_stage
27123 && n_head_kv == 1
27124 && n_head.is_multiple_of(4)
27125 && *G4_ON.get_or_init(|| {
27126 std::env::var("MEMRA_FAW_G4")
27127 .map(|v| v != "0")
27128 .unwrap_or(true)
27129 });
27130 if g4 {
27131 const SP_M: usize = 16;
27132 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27135 let o2 = *O2_ON.get_or_init(|| {
27136 std::env::var("MEMRA_FAW_O2")
27137 .map(|v| v != "0")
27138 .unwrap_or(true)
27139 });
27140 let f = self.func(if o2 {
27141 "fa_prefill_w_bf16_g4o2"
27142 } else {
27143 "fa_prefill_w_bf16_g4"
27144 });
27145 let shmem = if o2 {
27146 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
27147 } else {
27148 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
27149 as u32
27150 };
27151 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27152 f.set_attribute(
27153 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27154 shmem as i32,
27155 )?;
27156 let cfg = LaunchConfig {
27157 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
27158 block_dim: (32, 4, 1),
27159 shared_mem_bytes: shmem,
27160 };
27161 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
27162 head_dim as i32,
27163 n_head as i32,
27164 n_head_kv as i32,
27165 t as i32,
27166 t_kv as i32,
27167 causal as i32,
27168 window as i32,
27169 );
27170 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27171 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27172 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27173 let __s_b = self.gpu.stream();
27174 let mut b = __s_b.launch_builder(&f);
27175 b.arg(&qb)
27176 .arg(&kb)
27177 .arg(&vb)
27178 .arg(o)
27179 .arg(&hd)
27180 .arg(&nh)
27181 .arg(&nhkv)
27182 .arg(&ti)
27183 .arg(&tkvi)
27184 .arg(&scale)
27185 .arg(&cz)
27186 .arg(&wi);
27187 unsafe {
27188 b.launch(cfg)?;
27189 }
27190 return Ok(());
27191 }
27192 let f = self.func(if floor {
27193 "fa_prefill_w_f32"
27194 } else if f32_stage {
27195 "fa_prefill_w_f32_pp"
27196 } else {
27197 "fa_prefill_w_bf16_pp"
27198 });
27199 let shmem =
27200 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
27201 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27202 f.set_attribute(
27203 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27204 shmem as i32,
27205 )?;
27206 let cfg = LaunchConfig {
27207 grid_dim: (
27208 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27209 n_head as u32,
27210 1,
27211 ),
27212 block_dim: (32, 4, 1),
27213 shared_mem_bytes: shmem,
27214 };
27215 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
27216 head_dim as i32,
27217 n_head as i32,
27218 n_head_kv as i32,
27219 t as i32,
27220 t_kv as i32,
27221 causal as i32,
27222 window as i32,
27223 );
27224 if f32_stage {
27225 let __s_b = self.gpu.stream();
27226 let mut b = __s_b.launch_builder(&f);
27227 b.arg(q)
27228 .arg(k)
27229 .arg(v)
27230 .arg(o)
27231 .arg(&hd)
27232 .arg(&nh)
27233 .arg(&nhkv)
27234 .arg(&ti)
27235 .arg(&tkvi)
27236 .arg(&scale)
27237 .arg(&cz)
27238 .arg(&wi);
27239 unsafe {
27240 b.launch(cfg)?;
27241 }
27242 } else {
27243 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27244 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27245 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27246 let __s_b = self.gpu.stream();
27247 let mut b = __s_b.launch_builder(&f);
27248 b.arg(&qb)
27249 .arg(&kb)
27250 .arg(&vb)
27251 .arg(o)
27252 .arg(&hd)
27253 .arg(&nh)
27254 .arg(&nhkv)
27255 .arg(&ti)
27256 .arg(&tkvi)
27257 .arg(&scale)
27258 .arg(&cz)
27259 .arg(&wi);
27260 unsafe {
27261 b.launch(cfg)?;
27262 }
27263 }
27264 Ok(())
27265 }
27266
27267 #[allow(clippy::too_many_arguments)]
27271 pub fn fa_prefill_hd512(
27272 &self,
27273 q: &CudaSlice<f32>,
27274 k: &CudaSlice<f32>,
27275 v: &CudaSlice<f32>,
27276 o: &mut CudaSlice<f32>,
27277 head_dim: usize,
27278 n_head: usize,
27279 n_head_kv: usize,
27280 t: usize,
27281 t_kv: usize,
27282 scale: f32,
27283 causal: bool,
27284 ) -> Result<(), Box<dyn std::error::Error>> {
27285 if portable_mma_gated() {
27287 return self.sdpa_naive(
27288 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
27289 );
27290 }
27291 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27297 let f32_stage =
27298 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
27299 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27303 let sp = !f32_stage
27304 && *SP_ON.get_or_init(|| {
27305 std::env::var("MEMRA_FA512_SP")
27306 .map(|v| v != "0")
27307 .unwrap_or(true)
27308 });
27309 self.fa_prefill_hd512_arm(
27310 q,
27311 k,
27312 v,
27313 o,
27314 head_dim,
27315 n_head,
27316 n_head_kv,
27317 t,
27318 t_kv,
27319 scale,
27320 causal,
27321 f32_stage,
27322 sp,
27323 sp && fa_f16pv_on(),
27324 )
27325 }
27326
27327 #[allow(clippy::too_many_arguments)]
27329 pub fn fa_prefill_hd512_pre(
27330 &self,
27331 qb: &CudaSlice<u8>,
27332 kb: &CudaSlice<u8>,
27333 vb: &CudaSlice<u8>,
27334 o: &mut CudaSlice<f32>,
27335 head_dim: usize,
27336 n_head: usize,
27337 n_head_kv: usize,
27338 t: usize,
27339 t_kv: usize,
27340 scale: f32,
27341 causal: bool,
27342 v_f16: bool,
27343 ) -> Result<(), Box<dyn std::error::Error>> {
27344 debug_assert_eq!(head_dim, 512);
27345 const SP_M: usize = 16;
27346 const BKS: usize = 32;
27347 let f16pv = fa_f16pv_on();
27351 let nw = if f16pv { fa512_wide_warps() } else { 2 };
27352 let hp = f16pv
27353 && fa512_hp_on()
27354 && n_head.is_multiple_of(2)
27355 && (n_head / n_head_kv).is_multiple_of(2);
27356 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
27357 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
27358 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
27359 let n = t_kv * n_head_kv * head_dim;
27361 let need = n * 2;
27362 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
27363 *vguard = Some(self.alloc_uninit::<u8>(need)?);
27364 }
27365 let dst = vguard.as_mut().unwrap();
27366 self.bf16_to_f16_into(vb, n, dst)?;
27367 vguard.as_ref().unwrap()
27368 } else {
27369 vb
27370 };
27371 let f = self.func(if hp {
27372 "fa_prefill_bf16_hd512_sp16h2"
27373 } else {
27374 match (f16pv, nw) {
27375 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
27376 (true, _) => "fa_prefill_bf16_hd512_sp16",
27377 _ => "fa_prefill_bf16_hd512_sp",
27378 }
27379 });
27380 let (nwarp, npart) = if hp {
27381 (4usize, 4usize)
27382 } else if nw > 2 {
27383 (nw, nw)
27384 } else {
27385 (2, 1)
27386 };
27387 let shmem = if hp {
27389 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
27390 as u32
27391 } else {
27392 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
27393 + 4 * (npart * SP_M * BKS + SP_M)) as u32
27394 };
27395 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27396 f.set_attribute(
27397 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27398 shmem as i32,
27399 )?;
27400 let grid_y = if hp {
27401 (n_head / 2) as u32
27402 } else {
27403 n_head as u32
27404 };
27405 let cfg = LaunchConfig {
27406 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
27407 block_dim: (32, nwarp as u32, 1),
27408 shared_mem_bytes: shmem,
27409 };
27410 let (hd, nh, nhkv, ti, tkvi, cz) = (
27411 head_dim as i32,
27412 n_head as i32,
27413 n_head_kv as i32,
27414 t as i32,
27415 t_kv as i32,
27416 causal as i32,
27417 );
27418 let __s_b = self.gpu.stream();
27419 let mut b = __s_b.launch_builder(&f);
27420 b.arg(qb)
27421 .arg(kb)
27422 .arg(vref)
27423 .arg(o)
27424 .arg(&hd)
27425 .arg(&nh)
27426 .arg(&nhkv)
27427 .arg(&ti)
27428 .arg(&tkvi)
27429 .arg(&scale)
27430 .arg(&cz);
27431 unsafe {
27432 b.launch(cfg)?;
27433 }
27434 Ok(())
27435 }
27436
27437 #[allow(clippy::too_many_arguments)]
27444 pub fn mla_attn_gathered_tc(
27445 &self,
27446 q_lat_bf: &CudaSlice<u8>, cache_bf: &CudaSlice<u8>, idx: &CudaSlice<i32>, o_lat: &mut CudaSlice<f32>, n_head: usize,
27451 kv_rank: usize,
27452 t_q: usize,
27453 width: usize,
27454 scale: f32,
27455 ) -> Result<(), Box<dyn std::error::Error>> {
27456 if kv_rank != 512 {
27457 return Err(format!(
27458 "mla_attn_gathered_tc is stamped at kv_rank 512 (the glm5_next latent width); \
27459 got {kv_rank} — the caller's door must fall back to the f32 gathered kernel"
27460 )
27461 .into());
27462 }
27463 if t_q == 0 || n_head == 0 {
27464 return Ok(());
27465 }
27466 const SP_M: usize = 16;
27467 const BKS: usize = 32;
27468 const HD: usize = 512;
27469 let f = self.func("fa_mla_gathered_bf16");
27470 let shmem =
27472 (2 * (SP_M * HD + BKS * HD + SP_M * BKS) + 4 * (SP_M * BKS + SP_M) + 4 * BKS) as u32;
27473 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27474 f.set_attribute(
27475 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27476 shmem as i32,
27477 )?;
27478 let cfg = LaunchConfig {
27479 grid_dim: (t_q as u32, (n_head as u32).div_ceil(SP_M as u32), 1),
27480 block_dim: (32, 2, 1),
27481 shared_mem_bytes: shmem,
27482 };
27483 let (nh, tq, w) = (n_head as i32, t_q as i32, width as i32);
27484 let __s_b = self.gpu.stream();
27485 let mut b = __s_b.launch_builder(&f);
27486 b.arg(q_lat_bf)
27487 .arg(cache_bf)
27488 .arg(idx)
27489 .arg(o_lat)
27490 .arg(&nh)
27491 .arg(&tq)
27492 .arg(&w)
27493 .arg(&scale);
27494 unsafe {
27495 b.launch(cfg)?;
27496 }
27497 Ok(())
27498 }
27499
27500 #[allow(clippy::too_many_arguments)]
27503 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_hd512_arm(
27505 &self,
27506 q: &CudaSlice<f32>,
27507 k: &CudaSlice<f32>,
27508 v: &CudaSlice<f32>,
27509 o: &mut CudaSlice<f32>,
27510 head_dim: usize,
27511 n_head: usize,
27512 n_head_kv: usize,
27513 t: usize,
27514 t_kv: usize,
27515 scale: f32,
27516 causal: bool,
27517 f32_stage: bool,
27518 sp: bool,
27519 f16pv: bool,
27520 ) -> Result<(), Box<dyn std::error::Error>> {
27521 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
27522 if sp && !f32_stage {
27523 const SP_M: usize = 16;
27527 const BKS: usize = 32;
27528 let nw = if f16pv { fa512_wide_warps() } else { 2 };
27529 let hp = f16pv
27530 && fa512_hp_on()
27531 && n_head.is_multiple_of(2)
27532 && (n_head / n_head_kv).is_multiple_of(2);
27533 let f = self.func(if hp {
27534 "fa_prefill_bf16_hd512_sp16h2"
27535 } else {
27536 match (f16pv, nw) {
27537 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
27538 (true, _) => "fa_prefill_bf16_hd512_sp16",
27539 _ => "fa_prefill_bf16_hd512_sp",
27540 }
27541 });
27542 let (nwarp, npart) = if hp {
27543 (4usize, 4usize)
27544 } else if nw > 2 {
27545 (nw, nw)
27546 } else {
27547 (2, 1)
27548 };
27549 let shmem = if hp {
27550 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
27551 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
27552 } else {
27553 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
27554 + 4 * (npart * SP_M * BKS + SP_M)) as u32
27555 };
27556 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27557 f.set_attribute(
27558 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27559 shmem as i32,
27560 )?;
27561 let grid_y = if hp {
27562 (n_head / 2) as u32
27563 } else {
27564 n_head as u32
27565 };
27566 let cfg = LaunchConfig {
27567 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
27568 block_dim: (32, nwarp as u32, 1),
27569 shared_mem_bytes: shmem,
27570 };
27571 let (hd, nh, nhkv, ti, tkvi, cz) = (
27572 head_dim as i32,
27573 n_head as i32,
27574 n_head_kv as i32,
27575 t as i32,
27576 t_kv as i32,
27577 causal as i32,
27578 );
27579 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27580 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27581 let vb = if f16pv {
27582 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
27583 } else {
27584 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
27585 };
27586 let __s_b = self.gpu.stream();
27587 let mut b = __s_b.launch_builder(&f);
27588 b.arg(&qb)
27589 .arg(&kb)
27590 .arg(&vb)
27591 .arg(o)
27592 .arg(&hd)
27593 .arg(&nh)
27594 .arg(&nhkv)
27595 .arg(&ti)
27596 .arg(&tkvi)
27597 .arg(&scale)
27598 .arg(&cz);
27599 unsafe {
27600 b.launch(cfg)?;
27601 }
27602 return Ok(());
27603 }
27604 const BLOCK_Q: usize = 32;
27605 const BK: usize = 32;
27606 const HALF: usize = 256;
27607 let f = self.func(if f32_stage {
27608 "fa_prefill_f32_hd512"
27609 } else {
27610 "fa_prefill_bf16_hd512"
27611 });
27612 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
27614 + 4 * BLOCK_Q) as u32;
27615 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27616 f.set_attribute(
27617 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27618 shmem as i32,
27619 )?;
27620 let cfg = LaunchConfig {
27621 grid_dim: (
27622 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
27623 n_head as u32,
27624 2,
27625 ),
27626 block_dim: (32, 2, 1),
27627 shared_mem_bytes: shmem,
27628 };
27629 let (hd, nh, nhkv, ti, tkvi, cz) = (
27630 head_dim as i32,
27631 n_head as i32,
27632 n_head_kv as i32,
27633 t as i32,
27634 t_kv as i32,
27635 causal as i32,
27636 );
27637 if f32_stage {
27638 let __s_b = self.gpu.stream();
27639 let mut b = __s_b.launch_builder(&f);
27640 b.arg(q)
27641 .arg(k)
27642 .arg(v)
27643 .arg(o)
27644 .arg(&hd)
27645 .arg(&nh)
27646 .arg(&nhkv)
27647 .arg(&ti)
27648 .arg(&tkvi)
27649 .arg(&scale)
27650 .arg(&cz);
27651 unsafe {
27652 b.launch(cfg)?;
27653 }
27654 } else {
27655 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
27656 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
27657 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
27658 let __s_b = self.gpu.stream();
27659 let mut b = __s_b.launch_builder(&f);
27660 b.arg(&qb)
27661 .arg(&kb)
27662 .arg(&vb)
27663 .arg(o)
27664 .arg(&hd)
27665 .arg(&nh)
27666 .arg(&nhkv)
27667 .arg(&ti)
27668 .arg(&tkvi)
27669 .arg(&scale)
27670 .arg(&cz);
27671 unsafe {
27672 b.launch(cfg)?;
27673 }
27674 }
27675 Ok(())
27676 }
27677
27678 #[allow(clippy::too_many_arguments)]
27682 pub fn rope_neox2_bf16e(
27683 &self,
27684 q: &mut CudaSlice<f32>,
27685 k: &mut CudaSlice<f32>,
27686 qb: &mut CudaSlice<u8>,
27687 kb: &mut CudaSlice<u8>,
27688 pos: &CudaSlice<i32>,
27689 head_dim: usize,
27690 n_dims: usize,
27691 nh_q: usize,
27692 nh_k: usize,
27693 n_tokens: usize,
27694 base: f32,
27695 freq_scale: f32,
27696 ff: Option<&CudaSlice<f32>>,
27697 ) -> Result<(), Box<dyn std::error::Error>> {
27698 let f = self.func("rope_neox2_bf16e_f32");
27699 let rows = ((nh_q + nh_k) * n_tokens) as u32;
27700 let cfg = LaunchConfig {
27701 grid_dim: (rows, 1, 1),
27702 block_dim: ((head_dim / 2) as u32, 1, 1),
27703 shared_mem_bytes: 0,
27704 };
27705 let theta_scale = base.powf(-2.0 / n_dims as f32);
27706 let (hd, nd, nhq, nhk, nt) = (
27707 head_dim as i32,
27708 n_dims as i32,
27709 nh_q as i32,
27710 nh_k as i32,
27711 n_tokens as i32,
27712 );
27713 let __s_b = self.gpu.stream();
27714 let mut b = __s_b.launch_builder(&f);
27715 match ff {
27716 Some(t) => {
27717 b.arg(&mut *q)
27718 .arg(&mut *k)
27719 .arg(&mut *qb)
27720 .arg(&mut *kb)
27721 .arg(pos)
27722 .arg(&hd)
27723 .arg(&nd)
27724 .arg(&nhq)
27725 .arg(&nhk)
27726 .arg(&nt)
27727 .arg(&theta_scale)
27728 .arg(&freq_scale)
27729 .arg(t);
27730 unsafe {
27731 b.launch(cfg)?;
27732 }
27733 }
27734 None => {
27735 let null: u64 = 0;
27736 b.arg(&mut *q)
27737 .arg(&mut *k)
27738 .arg(&mut *qb)
27739 .arg(&mut *kb)
27740 .arg(pos)
27741 .arg(&hd)
27742 .arg(&nd)
27743 .arg(&nhq)
27744 .arg(&nhk)
27745 .arg(&nt)
27746 .arg(&theta_scale)
27747 .arg(&freq_scale)
27748 .arg(&null);
27749 unsafe {
27750 b.launch(cfg)?;
27751 }
27752 }
27753 }
27754 Ok(())
27755 }
27756
27757 pub fn f32_to_bf16(
27760 &self,
27761 x: &CudaSlice<f32>,
27762 n: usize,
27763 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27764 assert!(
27765 n.is_multiple_of(4),
27766 "f32_to_bf16 requires n % 4 == 0, got {n}"
27767 );
27768 let mut y = self.alloc_uninit::<u8>(n * 2)?;
27769 let f = self.func("f32_to_bf16_flat");
27770 let n_i = n as i64;
27771 let cfg = LaunchConfig {
27772 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
27773 block_dim: (256, 1, 1),
27774 shared_mem_bytes: 0,
27775 };
27776 let __s_b = self.gpu.stream();
27777 let mut b = __s_b.launch_builder(&f);
27778 b.arg(x).arg(&mut y).arg(&n_i);
27779 unsafe {
27780 b.launch(cfg)?;
27781 }
27782 Ok(y)
27783 }
27784
27785 pub fn f32_to_f16(
27786 &self,
27787 x: &CudaSlice<f32>,
27788 n: usize,
27789 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27790 assert!(
27791 n.is_multiple_of(4),
27792 "f32_to_f16 requires n % 4 == 0, got {n}"
27793 );
27794 let mut y = self.alloc_uninit::<u8>(n * 2)?;
27795 let f = self.func("f32_to_f16_flat");
27796 let n_i = n as i64;
27797 let cfg = LaunchConfig {
27798 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
27799 block_dim: (256, 1, 1),
27800 shared_mem_bytes: 0,
27801 };
27802 let __s_b = self.gpu.stream();
27803 let mut b = __s_b.launch_builder(&f);
27804 b.arg(x).arg(&mut y).arg(&n_i);
27805 unsafe {
27806 b.launch(cfg)?;
27807 }
27808 Ok(y)
27809 }
27810
27811 pub fn bf16_to_f16(
27813 &self,
27814 xb: &CudaSlice<u8>,
27815 n: usize,
27816 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27817 let mut y = self.alloc_uninit::<u8>(n * 2)?;
27818 self.bf16_to_f16_into(xb, n, &mut y)?;
27819 Ok(y)
27820 }
27821
27822 pub fn bf16_to_f16_into(
27824 &self,
27825 xb: &CudaSlice<u8>,
27826 n: usize,
27827 y: &mut CudaSlice<u8>,
27828 ) -> Result<(), Box<dyn std::error::Error>> {
27829 assert!(
27830 n.is_multiple_of(2),
27831 "bf16_to_f16 requires n % 2 == 0, got {n}"
27832 );
27833 assert!(y.len() >= n * 2);
27834 let f = self.func("bf16_to_f16_flat");
27835 let n2 = (n / 2) as i64;
27836 let cfg = LaunchConfig {
27837 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
27838 block_dim: (256, 1, 1),
27839 shared_mem_bytes: 0,
27840 };
27841 let __s_b = self.gpu.stream();
27842 let mut b = __s_b.launch_builder(&f);
27843 b.arg(xb).arg(y).arg(&n2);
27844 unsafe {
27845 b.launch(cfg)?;
27846 }
27847 Ok(())
27848 }
27849
27850 #[allow(clippy::too_many_arguments)]
27855 pub fn fa_prefill_vl8(
27856 &self,
27857 seqs: &[FaSeqVl],
27858 head_dim: usize,
27859 n_head: usize,
27860 n_head_kv: usize,
27861 scale: f32,
27862 ) -> Result<(), Box<dyn std::error::Error>> {
27863 const BK: usize = 32;
27864 let b = seqs.len();
27865 assert!((1..=8).contains(&b));
27866 let mut packed = [FaSeqVl::default(); 8];
27867 packed[..b].copy_from_slice(seqs);
27868 let v = FaVl8(packed);
27869 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
27870 let ept = (n_head_kv * head_dim) as i32;
27871 {
27872 let f = self.func("fa_mirror_vl");
27873 let max_n = (max_t as i64) * ept as i64;
27874 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
27875 for which in 0..2i32 {
27876 let cfg = LaunchConfig {
27877 grid_dim: (blocks, 1, b as u32),
27878 block_dim: (256, 1, 1),
27879 shared_mem_bytes: 0,
27880 };
27881 let __s_lb = self.gpu.stream();
27882 let mut lb = __s_lb.launch_builder(&f);
27883 lb.arg(&v).arg(&ept).arg(&which);
27884 unsafe {
27885 lb.launch(cfg)?;
27886 }
27887 }
27888 }
27889 let hd_sfx = fa_hd_suffix(head_dim)?;
27890 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
27891 let block_q = 64usize;
27892 let kv_stages = 2usize;
27893 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
27894 + 4 * (block_q * BK + 2 * block_q)) as u32;
27895 use cudarc::driver::sys::CUfunction_attribute_enum as A;
27896 f.set_attribute(
27897 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
27898 shmem as i32,
27899 )?;
27900 let cfg = LaunchConfig {
27901 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
27902 block_dim: (32, 4, 1),
27903 shared_mem_bytes: shmem,
27904 };
27905 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27906 let __s_lb = self.gpu.stream();
27907 let mut lb = __s_lb.launch_builder(&f);
27908 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
27909 unsafe {
27910 lb.launch(cfg)?;
27911 }
27912 Ok(())
27913 }
27914
27915 #[allow(clippy::too_many_arguments)]
27919 pub fn attn_pre_vl8(
27920 &self,
27921 seqs: &[AttnPreVl],
27922 wq: &CudaSlice<f32>,
27923 wk: &CudaSlice<f32>,
27924 head_dim: usize,
27925 rope_dims: usize,
27926 n_head: usize,
27927 n_head_kv: usize,
27928 eps: f32,
27929 freq_base: f32,
27930 freq_scale: f32,
27931 kv_dim_k: usize,
27932 kv_dim_v: usize,
27933 k_tok_bytes: usize,
27934 v_tok_bytes: usize,
27935 ) -> Result<(), Box<dyn std::error::Error>> {
27936 let b = seqs.len();
27937 assert!((1..=8).contains(&b));
27938 let mut packed = [AttnPreVl::default(); 8];
27939 packed[..b].copy_from_slice(seqs);
27940 let v = AttnPreVl8(packed);
27941 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
27942 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
27943 {
27944 let f = self.func("q_gate_split_vl");
27945 let n = max_t * (n_head * head_dim) as u32;
27946 let cfg = LaunchConfig {
27947 grid_dim: (n.div_ceil(256), 1, b as u32),
27948 block_dim: (256, 1, 1),
27949 shared_mem_bytes: 0,
27950 };
27951 let __s_lb = self.gpu.stream();
27952 let mut lb = __s_lb.launch_builder(&f);
27953 lb.arg(&v).arg(&hd).arg(&nh);
27954 unsafe {
27955 lb.launch(cfg)?;
27956 }
27957 }
27958 {
27959 let f = self.func("attn_rms_vl");
27960 let cfg = LaunchConfig {
27961 grid_dim: (max_t * n_head as u32, 2, b as u32),
27962 block_dim: (rms_block(), 1, 1),
27963 shared_mem_bytes: 0,
27964 };
27965 let __s_lb = self.gpu.stream();
27966 let mut lb = __s_lb.launch_builder(&f);
27967 lb.arg(&v)
27968 .arg(wq)
27969 .arg(wk)
27970 .arg(&hd)
27971 .arg(&nh)
27972 .arg(&nhkv)
27973 .arg(&eps);
27974 unsafe {
27975 lb.launch(cfg)?;
27976 }
27977 }
27978 {
27979 let f = self.func("attn_rope_vl");
27980 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
27981 let nd = rope_dims as i32;
27982 let cfg = LaunchConfig {
27983 grid_dim: (max_t * n_head as u32, 2, b as u32),
27984 block_dim: ((head_dim / 2) as u32, 1, 1),
27985 shared_mem_bytes: 0,
27986 };
27987 let __s_lb = self.gpu.stream();
27988 let mut lb = __s_lb.launch_builder(&f);
27989 lb.arg(&v)
27990 .arg(&hd)
27991 .arg(&nd)
27992 .arg(&nh)
27993 .arg(&nhkv)
27994 .arg(&theta_scale)
27995 .arg(&freq_scale);
27996 unsafe {
27997 lb.launch(cfg)?;
27998 }
27999 }
28000 {
28001 let f = self.func("append_kv_vl");
28002 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
28003 let cfg = LaunchConfig {
28004 grid_dim: (nblk, max_t, b as u32),
28005 block_dim: (32, 1, 1),
28006 shared_mem_bytes: 0,
28007 };
28008 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
28009 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28010 let __s_lb = self.gpu.stream();
28011 let mut lb = __s_lb.launch_builder(&f);
28012 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
28013 unsafe {
28014 lb.launch(cfg)?;
28015 }
28016 }
28017 Ok(())
28018 }
28019
28020 #[allow(clippy::too_many_arguments)]
28025 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_view(
28028 &self,
28029 q: &CudaSlice<f32>,
28030 k: &cudarc::driver::CudaView<u8>,
28031 v: &cudarc::driver::CudaView<u8>,
28032 o: &mut CudaSlice<f32>,
28033 head_dim: usize,
28034 n_head: usize,
28035 n_head_kv: usize,
28036 t: usize,
28037 t_kv: usize,
28038 scale: f32,
28039 causal: bool,
28040 k_tok_bytes: usize,
28041 v_tok_bytes: usize,
28042 g: bool,
28043 ) -> Result<(), Box<dyn std::error::Error>> {
28044 if portable_mma_gated() {
28045 return self.sdpa_naive_quantized_view(
28046 q,
28047 k,
28048 v,
28049 o,
28050 head_dim,
28051 n_head,
28052 n_head_kv,
28053 t,
28054 t_kv,
28055 scale,
28056 causal,
28057 k_tok_bytes,
28058 v_tok_bytes,
28059 );
28060 }
28061 const BLOCK_Q: usize = 64;
28062 const BK: usize = 32;
28063 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
28066 let f = if g {
28067 self.func_g(&name)
28068 } else {
28069 self.func(&name)
28070 };
28071 let shmem =
28072 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
28073 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28074 f.set_attribute(
28075 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28076 shmem as i32,
28077 )?;
28078 let cfg = LaunchConfig {
28079 grid_dim: (
28080 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
28081 n_head as u32,
28082 1,
28083 ),
28084 block_dim: (32, 4, 1),
28085 shared_mem_bytes: shmem,
28086 };
28087 let (hd, nh, nhkv, ti, tkvi, cz) = (
28088 head_dim as i32,
28089 n_head as i32,
28090 n_head_kv as i32,
28091 t as i32,
28092 t_kv as i32,
28093 causal as i32,
28094 );
28095 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28096 let __s_b = self.gpu.stream();
28097 let mut b = __s_b.launch_builder(&f);
28098 b.arg(q)
28099 .arg(k)
28100 .arg(v)
28101 .arg(o)
28102 .arg(&hd)
28103 .arg(&nh)
28104 .arg(&nhkv)
28105 .arg(&ti)
28106 .arg(&tkvi)
28107 .arg(&scale)
28108 .arg(&cz)
28109 .arg(&ktb)
28110 .arg(&vtb);
28111 unsafe {
28112 b.launch(cfg)?;
28113 }
28114 Ok(())
28115 }
28116
28117 #[allow(clippy::too_many_arguments)]
28127 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_view_ws(
28129 &self,
28130 q: &CudaSlice<f32>,
28131 k: &cudarc::driver::CudaView<u8>,
28132 v: &cudarc::driver::CudaView<u8>,
28133 o: &mut CudaSlice<f32>,
28134 head_dim: usize,
28135 n_head: usize,
28136 n_head_kv: usize,
28137 t: usize,
28138 t_kv: usize,
28139 scale: f32,
28140 causal: bool,
28141 k_tok_bytes: usize,
28142 v_tok_bytes: usize,
28143 g: bool,
28144 ) -> Result<(), Box<dyn std::error::Error>> {
28145 if portable_mma_gated() {
28146 return self.sdpa_naive_quantized_view(
28147 q,
28148 k,
28149 v,
28150 o,
28151 head_dim,
28152 n_head,
28153 n_head_kv,
28154 t,
28155 t_kv,
28156 scale,
28157 causal,
28158 k_tok_bytes,
28159 v_tok_bytes,
28160 );
28161 }
28162 const BLOCK_Q: usize = 64;
28163 const BK: usize = 32;
28164 let kv_dim_k = n_head_kv * head_dim;
28165 let kv_dim_v = n_head_kv * head_dim;
28166 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
28168 let mut guard = self.prime_deqw_ws.lock().unwrap();
28170 let need_grow = match guard.as_ref() {
28171 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
28172 None => true,
28173 };
28174 if need_grow {
28175 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
28176 let (ck, cv) = guard
28177 .as_ref()
28178 .map(|(a, b)| (a.len(), b.len()))
28179 .unwrap_or((0, 0));
28180 *guard = Some((
28181 self.alloc_u8(grow(ck, k_ws_bytes))?,
28182 self.alloc_u8(grow(cv, v_ws_bytes))?,
28183 ));
28184 }
28185 let (kw, vw) = guard.as_mut().unwrap();
28186 {
28188 let f = if g {
28190 self.func_g("fa_dequant_kv_ws_bf16")
28191 } else {
28192 self.func("fa_dequant_kv_ws_bf16")
28193 };
28194 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
28195 #[allow(clippy::manual_div_ceil)]
28196 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
28198 let cfg = LaunchConfig {
28199 grid_dim: (nblk.max(1), 1, 1),
28200 block_dim: (256, 1, 1),
28201 shared_mem_bytes: 0,
28202 };
28203 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
28204 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28205 let __s_b = self.gpu.stream();
28206 let mut b = __s_b.launch_builder(&f);
28207 b.arg(k)
28208 .arg(v)
28209 .arg(&mut *kw)
28210 .arg(&mut *vw)
28211 .arg(&kdk)
28212 .arg(&kdv)
28213 .arg(&tkvi)
28214 .arg(&ktb)
28215 .arg(&vtb);
28216 unsafe {
28217 b.launch(cfg)?;
28218 }
28219 }
28220 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
28228 .map(|v| v != "0")
28229 .unwrap_or(true);
28230 {
28231 let hd_sfx = fa_hd_suffix(head_dim)?;
28232 let f = self.func(&format!(
28233 "fa_prefill_qw{}{hd_sfx}",
28234 if db { "_db" } else { "" }
28235 ));
28236 let shmem = if db {
28237 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
28239 } else {
28240 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
28241 };
28242 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28243 f.set_attribute(
28244 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28245 shmem as i32,
28246 )?;
28247 let cfg = LaunchConfig {
28248 grid_dim: (
28249 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
28250 n_head as u32,
28251 1,
28252 ),
28253 block_dim: (32, 4, 1),
28254 shared_mem_bytes: shmem,
28255 };
28256 let (hd, nh, nhkv, ti, tkvi, cz) = (
28257 head_dim as i32,
28258 n_head as i32,
28259 n_head_kv as i32,
28260 t as i32,
28261 t_kv as i32,
28262 causal as i32,
28263 );
28264 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
28265 let __s_b = self.gpu.stream();
28266 let mut b = __s_b.launch_builder(&f);
28267 b.arg(q)
28268 .arg(&*kw)
28269 .arg(&*vw)
28270 .arg(o)
28271 .arg(&hd)
28272 .arg(&nh)
28273 .arg(&nhkv)
28274 .arg(&ti)
28275 .arg(&tkvi)
28276 .arg(&scale)
28277 .arg(&cz)
28278 .arg(&kdk)
28279 .arg(&kdv);
28280 unsafe {
28281 b.launch(cfg)?;
28282 }
28283 }
28284 Ok(())
28285 }
28286
28287 #[allow(clippy::too_many_arguments)]
28303 #[allow(clippy::manual_div_ceil)] pub fn fa_prefill_view_ws_w_hd128(
28305 &self,
28306 q: &CudaSlice<f32>,
28307 k: &cudarc::driver::CudaView<u8>,
28308 v: &cudarc::driver::CudaView<u8>,
28309 o: &mut CudaSlice<f32>,
28310 head_dim: usize,
28311 n_head: usize,
28312 n_head_kv: usize,
28313 t: usize,
28314 t_kv: usize,
28315 scale: f32,
28316 causal: bool,
28317 window: usize,
28318 k_tok_bytes: usize,
28319 v_tok_bytes: usize,
28320 ) -> Result<(), Box<dyn std::error::Error>> {
28321 assert_eq!(
28322 head_dim, 128,
28323 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
28324 );
28325 if portable_mma_gated() {
28326 return self.sdpa_naive_w_quantized_view(
28327 q,
28328 k,
28329 v,
28330 o,
28331 head_dim,
28332 n_head,
28333 n_head_kv,
28334 t,
28335 t_kv,
28336 scale,
28337 causal,
28338 window,
28339 k_tok_bytes,
28340 v_tok_bytes,
28341 );
28342 }
28343 const BLOCK_Q: usize = 64;
28344 const BK: usize = 32;
28345 let kv_dim_k = n_head_kv * head_dim;
28346 let kv_dim_v = n_head_kv * head_dim;
28347 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
28349 let mut guard = self.prime_deqw_ws.lock().unwrap();
28350 let need_grow = match guard.as_ref() {
28351 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
28352 None => true,
28353 };
28354 if need_grow {
28355 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
28356 let (ck, cv) = guard
28357 .as_ref()
28358 .map(|(a, b)| (a.len(), b.len()))
28359 .unwrap_or((0, 0));
28360 *guard = Some((
28361 self.alloc_u8(grow(ck, k_ws_bytes))?,
28362 self.alloc_u8(grow(cv, v_ws_bytes))?,
28363 ));
28364 }
28365 let (kw, vw) = guard.as_mut().unwrap();
28366 {
28369 let f = self.func("fa_dequant_kv_ws_bf16");
28370 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
28371 #[allow(clippy::manual_div_ceil)]
28372 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
28374 let cfg = LaunchConfig {
28375 grid_dim: (nblk.max(1), 1, 1),
28376 block_dim: (256, 1, 1),
28377 shared_mem_bytes: 0,
28378 };
28379 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
28380 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28381 let __s_b = self.gpu.stream();
28382 let mut b = __s_b.launch_builder(&f);
28383 b.arg(k)
28384 .arg(v)
28385 .arg(&mut *kw)
28386 .arg(&mut *vw)
28387 .arg(&kdk)
28388 .arg(&kdv)
28389 .arg(&tkvi)
28390 .arg(&ktb)
28391 .arg(&vtb);
28392 unsafe {
28393 b.launch(cfg)?;
28394 }
28395 }
28396 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
28398 .map(|v| v != "0")
28399 .unwrap_or(true);
28400 {
28401 let f = self.func(if db {
28402 "fa_prefill_qw_db_w_hd128"
28403 } else {
28404 "fa_prefill_qw_w_hd128"
28405 });
28406 let shmem = if db {
28407 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
28408 } else {
28409 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
28410 };
28411 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28412 f.set_attribute(
28413 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28414 shmem as i32,
28415 )?;
28416 let cfg = LaunchConfig {
28417 grid_dim: (
28418 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
28419 n_head as u32,
28420 1,
28421 ),
28422 block_dim: (32, 4, 1),
28423 shared_mem_bytes: shmem,
28424 };
28425 let (hd, nh, nhkv, ti, tkvi, cz) = (
28426 head_dim as i32,
28427 n_head as i32,
28428 n_head_kv as i32,
28429 t as i32,
28430 t_kv as i32,
28431 causal as i32,
28432 );
28433 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
28434 let __s_b = self.gpu.stream();
28435 let mut b = __s_b.launch_builder(&f);
28436 b.arg(q)
28437 .arg(&*kw)
28438 .arg(&*vw)
28439 .arg(o)
28440 .arg(&hd)
28441 .arg(&nh)
28442 .arg(&nhkv)
28443 .arg(&ti)
28444 .arg(&tkvi)
28445 .arg(&scale)
28446 .arg(&cz)
28447 .arg(&kdk)
28448 .arg(&kdv)
28449 .arg(&wnd);
28450 unsafe {
28451 b.launch(cfg)?;
28452 }
28453 }
28454 Ok(())
28455 }
28456
28457 #[allow(clippy::too_many_arguments)] pub fn fa_decode(
28462 &self,
28463 q: &CudaSlice<f32>,
28464 k: &cudarc::driver::CudaView<u8>,
28465 v: &cudarc::driver::CudaView<u8>,
28466 o: &mut CudaSlice<f32>,
28467 head_dim: usize,
28468 n_head: usize,
28469 n_head_kv: usize,
28470 t_kv: usize,
28471 scale: f32,
28472 k_tok_bytes: usize,
28473 v_tok_bytes: usize,
28474 ) -> Result<(), Box<dyn std::error::Error>> {
28475 self.fa_decode_kvmod(
28476 q,
28477 k,
28478 v,
28479 o,
28480 head_dim,
28481 n_head,
28482 n_head_kv,
28483 t_kv,
28484 scale,
28485 k_tok_bytes,
28486 v_tok_bytes,
28487 false,
28488 )
28489 }
28490
28491 #[allow(clippy::too_many_arguments)]
28495 #[allow(clippy::too_many_arguments)]
28499 #[allow(clippy::too_many_arguments)]
28500 fn fa_decode_scalar_unified(
28501 &self,
28502 q: &cudarc::driver::CudaView<f32>,
28503 k: &cudarc::driver::CudaView<u8>,
28504 v: &cudarc::driver::CudaView<u8>,
28505 o: &mut cudarc::driver::CudaViewMut<f32>,
28506 head_dim: usize,
28507 n_head: usize,
28508 n_head_kv: usize,
28509 t_kv_host: usize,
28510 t_kv_dev: Option<&CudaSlice<i32>>,
28511 scale: f32,
28512 n_splits: usize,
28513 split_keys: usize,
28514 k_tok_bytes: usize,
28515 v_tok_bytes: usize,
28516 g: bool,
28517 part_o: &mut CudaSlice<f32>,
28518 part_m: &mut CudaSlice<f32>,
28519 part_l: &mut CudaSlice<f32>,
28520 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
28521 ) -> Result<(), Box<dyn std::error::Error>> {
28522 let f = if g {
28523 self.func_g("fa_decode_f32")
28524 } else {
28525 self.fa_func("fa_decode_f32", head_dim)
28526 };
28527 let cfg = LaunchConfig {
28528 grid_dim: (n_head as u32, n_splits as u32, 1),
28529 block_dim: (head_dim as u32, 1, 1),
28530 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
28531 };
28532 let (hd, nh, nhkv, nsp) = (
28533 head_dim as i32,
28534 n_head as i32,
28535 n_head_kv as i32,
28536 n_splits as i32,
28537 );
28538 let (ktb, vtb, tkvi, ski) = (
28539 k_tok_bytes as i64,
28540 v_tok_bytes as i64,
28541 t_kv_host as i32,
28542 split_keys as i32,
28543 );
28544 let __s_b = self.gpu.stream();
28545 let mut b = __s_b.launch_builder(&f);
28546 match t_kv_dev {
28547 Some(d) => {
28548 b.arg(q)
28549 .arg(k)
28550 .arg(v)
28551 .arg(&mut *part_o)
28552 .arg(&mut *part_m)
28553 .arg(&mut *part_l)
28554 .arg(&hd)
28555 .arg(&nh)
28556 .arg(&nhkv)
28557 .arg(&tkvi)
28558 .arg(d)
28559 .arg(&scale)
28560 .arg(&nsp)
28561 .arg(&ski)
28562 .arg(&ktb)
28563 .arg(&vtb);
28564 unsafe {
28565 b.launch(cfg)?;
28566 }
28567 }
28568 None => {
28569 let null: u64 = 0;
28570 b.arg(q)
28571 .arg(k)
28572 .arg(v)
28573 .arg(&mut *part_o)
28574 .arg(&mut *part_m)
28575 .arg(&mut *part_l)
28576 .arg(&hd)
28577 .arg(&nh)
28578 .arg(&nhkv)
28579 .arg(&tkvi)
28580 .arg(&null)
28581 .arg(&scale)
28582 .arg(&nsp)
28583 .arg(&ski)
28584 .arg(&ktb)
28585 .arg(&vtb);
28586 unsafe {
28587 b.launch(cfg)?;
28588 }
28589 }
28590 }
28591 let cfg2 = LaunchConfig {
28592 grid_dim: (n_head as u32, 1, 1),
28593 block_dim: (head_dim as u32, 1, 1),
28594 shared_mem_bytes: 0,
28595 };
28596 if let Some((oq, od)) = q8_out {
28597 let fc = if g {
28599 self.func_g("fa_decode_combine_q8_1")
28600 } else {
28601 self.fa_func("fa_decode_combine_q8_1", head_dim)
28602 };
28603 let __s_b2 = self.gpu.stream();
28604 let mut b2 = __s_b2.launch_builder(&fc);
28605 b2.arg(&*part_o)
28606 .arg(&*part_m)
28607 .arg(&*part_l)
28608 .arg(oq)
28609 .arg(od)
28610 .arg(&hd)
28611 .arg(&nh)
28612 .arg(&nsp);
28613 unsafe {
28614 b2.launch(cfg2)?;
28615 }
28616 return Ok(());
28617 }
28618 let fc = if g {
28619 self.func_g("fa_decode_combine_f32")
28620 } else {
28621 self.fa_func("fa_decode_combine_f32", head_dim)
28622 };
28623 let __s_b2 = self.gpu.stream();
28624 let mut b2 = __s_b2.launch_builder(&fc);
28625 b2.arg(&*part_o)
28626 .arg(&*part_m)
28627 .arg(&*part_l)
28628 .arg(o)
28629 .arg(&hd)
28630 .arg(&nh)
28631 .arg(&nsp);
28632 unsafe {
28633 b2.launch(cfg2)?;
28634 }
28635 Ok(())
28636 }
28637
28638 #[allow(clippy::too_many_arguments)] pub fn fa_decode_kvmod(
28640 &self,
28641 q: &CudaSlice<f32>,
28642 k: &cudarc::driver::CudaView<u8>,
28643 v: &cudarc::driver::CudaView<u8>,
28644 o: &mut CudaSlice<f32>,
28645 head_dim: usize,
28646 n_head: usize,
28647 n_head_kv: usize,
28648 t_kv: usize,
28649 scale: f32,
28650 k_tok_bytes: usize,
28651 v_tok_bytes: usize,
28652 g: bool,
28653 ) -> Result<(), Box<dyn std::error::Error>> {
28654 let q_view = q.as_view();
28655 let mut o_view = o.as_view_mut();
28656 self.fa_decode_kvmod_view(
28657 &q_view,
28658 k,
28659 v,
28660 &mut o_view,
28661 head_dim,
28662 n_head,
28663 n_head_kv,
28664 t_kv,
28665 scale,
28666 k_tok_bytes,
28667 v_tok_bytes,
28668 g,
28669 )
28670 }
28671
28672 #[allow(clippy::too_many_arguments)]
28677 #[allow(clippy::manual_div_ceil)] pub fn fa_decode_kvmod_view(
28679 &self,
28680 q: &cudarc::driver::CudaView<f32>,
28681 k: &cudarc::driver::CudaView<u8>,
28682 v: &cudarc::driver::CudaView<u8>,
28683 o: &mut cudarc::driver::CudaViewMut<f32>,
28684 head_dim: usize,
28685 n_head: usize,
28686 n_head_kv: usize,
28687 t_kv: usize,
28688 scale: f32,
28689 k_tok_bytes: usize,
28690 v_tok_bytes: usize,
28691 g: bool,
28692 ) -> Result<(), Box<dyn std::error::Error>> {
28693 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
28714 if g && head_dim == 256 && !fa_v4_at(t_kv) {
28718 fa_vec = false;
28719 }
28720 let sp = fa_split_keys(t_kv, n_head_kv);
28721 let n_splits = if fa_vec {
28722 ((t_kv + sp - 1) / sp).max(1)
28723 } else {
28724 ((t_kv + 255) / 256).max(1)
28725 };
28726 let o_len = n_head * n_splits * head_dim;
28727 let ml_len = n_head * n_splits;
28728 let mut part_guard = self.fa_part_pool.lock().unwrap();
28729 if part_guard
28730 .as_ref()
28731 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
28732 .unwrap_or(true)
28733 {
28734 let old = part_guard.take();
28745 let (co, cm) = old
28746 .as_ref()
28747 .map(|pp| (pp.0.len(), pp.1.len()))
28748 .unwrap_or((0, 0));
28749 if let Some(old) = old {
28750 self.fa_part_retired.lock().unwrap().push(old);
28751 }
28752 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
28753 eprintln!(
28754 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
28755 co, o_len, cm, ml_len
28756 );
28757 }
28758 *part_guard =
28759 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
28760 }
28761 let pg = part_guard.as_mut().unwrap();
28762 self.gpu
28763 .stream()
28764 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
28765 self.gpu
28766 .stream()
28767 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
28768 self.gpu
28769 .stream()
28770 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
28771 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
28772 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
28773 let (hd, nh, nhkv, tkvi, nsp) = (
28774 head_dim as i32,
28775 n_head as i32,
28776 n_head_kv as i32,
28777 t_kv as i32,
28778 n_splits as i32,
28779 );
28780 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
28781 let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
28785 let fa512_min = fa512_min_tkv();
28790 let deep = fa_vec
28793 && head_dim == 256
28794 && fa_v4_at(t_kv)
28795 && !g
28796 && fa_deep_at(t_kv)
28797 && !matches!(fa_v4_mode(), "noB3" | "stage");
28798 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
28799 let gqa = (n_head / n_head_kv).max(1) as u32;
28802 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
28803 (
28804 fv,
28805 LaunchConfig {
28806 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28807 block_dim: (32, gqa, 1),
28808 shared_mem_bytes: 0,
28809 },
28810 )
28811 } else if fa_vec && head_dim <= 256 {
28812 let gqa = (n_head / n_head_kv).max(1) as u32;
28813 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
28824 let smem_tkv = *SMEM_TKV.get_or_init(|| {
28825 std::env::var("MEMRA_FA_SMEM_TKV")
28826 .ok()
28827 .and_then(|v| v.parse().ok())
28828 .unwrap_or_else(|| {
28829 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
28830 })
28831 });
28832 if fa_v4_at(t_kv) && head_dim == 256 {
28833 let v4name = match fa_v4_mode() {
28837 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
28840 _ => "fa_decode_vec_q_v4",
28841 };
28842 let fv = if g {
28843 self.func_g(v4name)
28844 } else {
28845 self.func(v4name)
28846 };
28847 let shmem = (if deep { 12160 } else { 11520 }
28850 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
28851 use cudarc::driver::sys::CUfunction_attribute_enum as A;
28852 fv.set_attribute(
28853 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28854 shmem as i32,
28855 )?;
28856 (
28857 fv,
28858 LaunchConfig {
28859 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28860 block_dim: (32, gqa, 1),
28861 shared_mem_bytes: shmem,
28862 },
28863 )
28864 } else if fa_v3_active(head_dim) {
28865 let fv = if g {
28868 self.func_g("fa_decode_vec_q_v3")
28869 } else {
28870 self.func("fa_decode_vec_q_v3")
28871 };
28872 let shmem = (32 * head_dim * 2) as u32; (
28874 fv,
28875 LaunchConfig {
28876 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28877 block_dim: (32, gqa, 1),
28878 shared_mem_bytes: shmem,
28879 },
28880 )
28881 } else if fa_v2_on() {
28882 let fv = if g {
28886 self.func_g("fa_decode_vec_q_v2")
28887 } else {
28888 self.func("fa_decode_vec_q_v2")
28889 };
28890 let shmem = (2 * 32 * head_dim * 2) as u32; (
28892 fv,
28893 LaunchConfig {
28894 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28895 block_dim: (32, gqa, 1),
28896 shared_mem_bytes: shmem,
28897 },
28898 )
28899 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
28900 {
28901 let fv = if g {
28905 self.func_g("fa_decode_vec_q_smem")
28906 } else {
28907 self.func("fa_decode_vec_q_smem")
28908 };
28909 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
28911 fv.set_attribute(
28912 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
28913 shmem as i32,
28914 )?;
28915 (
28916 fv,
28917 LaunchConfig {
28918 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28919 block_dim: (32, gqa, 1),
28920 shared_mem_bytes: shmem,
28921 },
28922 )
28923 } else {
28924 let fv = if g {
28927 self.func_g("fa_decode_vec_q")
28928 } else {
28929 self.func("fa_decode_vec_q")
28930 };
28931 (
28932 fv,
28933 LaunchConfig {
28934 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
28935 block_dim: (32, gqa, 1),
28936 shared_mem_bytes: 0,
28937 },
28938 )
28939 }
28940 } else {
28941 return self.fa_decode_scalar_unified(
28944 q,
28945 k,
28946 v,
28947 o,
28948 head_dim,
28949 n_head,
28950 n_head_kv,
28951 t_kv,
28952 None,
28953 scale,
28954 n_splits,
28955 if fa_vec { sp } else { 256 },
28956 k_tok_bytes,
28957 v_tok_bytes,
28958 g,
28959 part_o,
28960 part_m,
28961 part_l,
28962 None,
28963 );
28964 };
28965 let __s_b = self.gpu.stream();
28966 let mut b = __s_b.launch_builder(&f);
28967 b.arg(q)
28968 .arg(k)
28969 .arg(v)
28970 .arg(&mut *part_o)
28971 .arg(&mut *part_m)
28972 .arg(&mut *part_l)
28973 .arg(&hd)
28974 .arg(&nh)
28975 .arg(&nhkv)
28976 .arg(&tkvi)
28977 .arg(&scale)
28978 .arg(&nsp)
28979 .arg(&ktb)
28980 .arg(&vtb);
28981 unsafe {
28982 b.launch(cfg)?;
28983 }
28984 let (fc, cfg2) = (
28987 if g {
28988 self.func_g("fa_decode_combine_f32")
28989 } else {
28990 self.fa_func("fa_decode_combine_f32", head_dim)
28991 },
28992 LaunchConfig {
28993 grid_dim: (n_head as u32, 1, 1),
28994 block_dim: (head_dim as u32, 1, 1),
28995 shared_mem_bytes: 0,
28996 },
28997 );
28998 let __s_b2 = self.gpu.stream();
28999 let mut b2 = __s_b2.launch_builder(&fc);
29000 b2.arg(&*part_o)
29001 .arg(&*part_m)
29002 .arg(&*part_l)
29003 .arg(o)
29004 .arg(&hd)
29005 .arg(&nh)
29006 .arg(&nsp);
29007 unsafe {
29008 b2.launch(cfg2)?;
29009 }
29010 Ok(())
29011 }
29012
29013 #[allow(clippy::too_many_arguments)]
29024 pub fn fa_decode_batch_seqs_v4(
29025 &self,
29026 q: &CudaSlice<f32>,
29027 kv_ptrs: &cudarc::driver::CudaView<u64>,
29028 pos_seq: &CudaSlice<i32>,
29029 o: &mut CudaSlice<f32>,
29030 head_dim: usize,
29031 n_head: usize,
29032 n_head_kv: usize,
29033 b_n: usize,
29034 t_kv_max: usize,
29035 scale: f32,
29036 split_keys: usize,
29037 k_tok_bytes: usize,
29038 v_tok_bytes: usize,
29039 ) -> Result<(), Box<dyn std::error::Error>> {
29040 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
29041 #[allow(clippy::manual_div_ceil)]
29042 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
29044 let o_len = b_n * n_head * n_splits_max * head_dim;
29045 let ml_len = b_n * n_head * n_splits_max;
29046 let mut part_guard = self.fa_part_pool.lock().unwrap();
29047 if part_guard
29048 .as_ref()
29049 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29050 .unwrap_or(true)
29051 {
29052 let old = part_guard.take();
29063 let (co, cm) = old
29064 .as_ref()
29065 .map(|pp| (pp.0.len(), pp.1.len()))
29066 .unwrap_or((0, 0));
29067 if let Some(old) = old {
29068 self.fa_part_retired.lock().unwrap().push(old);
29069 }
29070 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29071 eprintln!(
29072 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29073 co, o_len, cm, ml_len
29074 );
29075 }
29076 *part_guard =
29077 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29078 }
29079 let pg = part_guard.as_mut().unwrap();
29080 self.gpu
29081 .stream()
29082 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29083 self.gpu
29084 .stream()
29085 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29086 self.gpu
29087 .stream()
29088 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29089 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29090 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29091 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
29092 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29093 let gqa = (n_head / n_head_kv).max(1) as u32;
29094 let f = self.func("fa_decode_vec_q_seqs_v4");
29095 let shmem = (11520 + 32 * head_dim * 2) as u32;
29097 use cudarc::driver::sys::CUfunction_attribute_enum as A;
29098 f.set_attribute(
29099 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29100 shmem as i32,
29101 )?;
29102 let cfg = LaunchConfig {
29103 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
29104 block_dim: (32, gqa, 1),
29105 shared_mem_bytes: shmem,
29106 };
29107 {
29108 let __s_b = self.gpu.stream();
29109 let mut b = __s_b.launch_builder(&f);
29110 b.arg(q)
29111 .arg(kv_ptrs)
29112 .arg(pos_seq)
29113 .arg(&mut *part_o)
29114 .arg(&mut *part_m)
29115 .arg(&mut *part_l)
29116 .arg(&hd)
29117 .arg(&nh)
29118 .arg(&nhkv)
29119 .arg(&scale)
29120 .arg(&nspm)
29121 .arg(&spk)
29122 .arg(&ktb)
29123 .arg(&vtb);
29124 unsafe {
29125 b.launch(cfg)?;
29126 }
29127 }
29128 let fc = self.func("fa_decode_combine_seqs");
29129 let cfg2 = LaunchConfig {
29130 grid_dim: (n_head as u32, b_n as u32, 1),
29131 block_dim: (head_dim as u32, 1, 1),
29132 shared_mem_bytes: 0,
29133 };
29134 let __s_b2 = self.gpu.stream();
29135 let mut b2 = __s_b2.launch_builder(&fc);
29136 b2.arg(&*part_o)
29137 .arg(&*part_m)
29138 .arg(&*part_l)
29139 .arg(o)
29140 .arg(&hd)
29141 .arg(&nh)
29142 .arg(pos_seq)
29143 .arg(&nspm)
29144 .arg(&spk);
29145 unsafe {
29146 b2.launch(cfg2)?;
29147 }
29148 Ok(())
29149 }
29150
29151 #[allow(clippy::too_many_arguments)]
29158 pub fn append_kv_quantized_seqs(
29159 &self,
29160 k_rows: &CudaSlice<f32>,
29161 v_rows: &CudaSlice<f32>,
29162 kv_ptrs: &cudarc::driver::CudaView<u64>,
29163 pos_seq: &CudaSlice<i32>,
29164 b_n: usize,
29165 kv_dim_k: usize,
29166 kv_dim_v: usize,
29167 k_tok_bytes: usize,
29168 v_tok_bytes: usize,
29169 ) -> Result<(), Box<dyn std::error::Error>> {
29170 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
29171 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
29172 let cfg = LaunchConfig {
29173 grid_dim: (nblk, b_n as u32, 1),
29174 block_dim: (32, 1, 1),
29175 shared_mem_bytes: 0,
29176 };
29177 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
29178 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29179 let __s_b = self.gpu.stream();
29180 let mut b = __s_b.launch_builder(&f);
29181 b.arg(k_rows)
29182 .arg(v_rows)
29183 .arg(kv_ptrs)
29184 .arg(pos_seq)
29185 .arg(&kdk)
29186 .arg(&kdv)
29187 .arg(&ktb)
29188 .arg(&vtb);
29189 unsafe {
29190 b.launch(cfg)?;
29191 }
29192 Ok(())
29193 }
29194
29195 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
29201 std::env::var("MEMRA_NO_FA_VEC").is_err()
29202 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
29203 && base_len + 1 >= fa_vec_min_tkv()
29204 && head_dim <= 256
29205 && head_dim.is_multiple_of(32)
29206 }
29207
29208 #[allow(clippy::too_many_arguments)]
29217 pub fn fa_decode_rows(
29218 &self,
29219 q: &CudaSlice<f32>,
29220 k: &cudarc::driver::CudaView<u8>,
29221 v: &cudarc::driver::CudaView<u8>,
29222 o: &mut CudaSlice<f32>,
29223 head_dim: usize,
29224 n_head: usize,
29225 n_head_kv: usize,
29226 base_len: usize,
29227 t: usize,
29228 scale: f32,
29229 k_tok_bytes: usize,
29230 v_tok_bytes: usize,
29231 base_dev: Option<(&CudaSlice<i32>, i32)>,
29235 kv_shared: bool,
29238 g: bool,
29242 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
29245 ) -> Result<(), Box<dyn std::error::Error>> {
29246 debug_assert!(
29247 base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim.is_multiple_of(32)
29248 );
29249 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
29256 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29257 let v = *SP512.get_or_init(|| {
29260 std::env::var("MEMRA_FA_SP512")
29261 .ok()
29262 .and_then(|x| x.parse().ok())
29263 .unwrap_or(0)
29264 });
29265 sp = if v >= 8 {
29266 v
29267 } else {
29268 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
29269 };
29270 }
29271 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29272 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29273 let gqa = (n_head / n_head_kv).max(1) as u32;
29274 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
29285 groups.push((0, t, sp));
29286 } else {
29287 let mut r0 = 0usize;
29288 while r0 < t {
29289 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
29290 let mut r1 = r0 + 1;
29291 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
29292 r1 += 1;
29293 }
29294 groups.push((r0, r1 - r0, sp_g));
29295 r0 = r1;
29296 }
29297 }
29298 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29302 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
29303 std::env::var("MEMRA_FA_SMEM_TKV")
29304 .ok()
29305 .and_then(|v| v.parse().ok())
29306 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
29307 });
29308 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
29309 let v3 = fa_v3_active(head_dim);
29310 let smem_rows =
29311 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
29312 let _ = kv_shared;
29317 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
29320 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
29334 let tb512 = head_dim == 512
29336 && sp <= 32
29337 && n_head / n_head_kv.max(1) <= 16
29338 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
29339 let fname = if tb512 {
29340 "fa_decode_vec_q_rows_v4_512_tb"
29341 } else if i2 {
29342 "fa_decode_vec_q_rows_dpl16_i2"
29343 } else if head_dim == 512 {
29344 "fa_decode_vec_q_rows_dpl16"
29345 }
29346 else if v4 {
29348 "fa_decode_vec_q_rows_v4"
29349 } else if v3 {
29350 "fa_decode_vec_q_rows_v3"
29351 } else if fa_v2_on() {
29352 "fa_decode_vec_q_rows_v2"
29353 } else if smem_rows {
29354 "fa_decode_vec_q_rows_smem"
29355 } else {
29356 "fa_decode_vec_q_rows"
29357 };
29358 let f = if head_dim == 512 {
29359 self.fa_func(fname, head_dim)
29360 } else if g {
29361 self.func_g(if smem_rows {
29369 "fa_decode_vec_q_rows"
29370 } else {
29371 fname
29372 })
29373 } else {
29374 self.func(fname)
29375 };
29376 let shmem = if tb512 {
29377 let gk = Self::gkv_on();
29379 let sh =
29380 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
29381 use cudarc::driver::sys::CUfunction_attribute_enum as A;
29382 f.set_attribute(
29383 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29384 sh as i32,
29385 )?;
29386 sh
29387 } else if v4 || v3 || smem_rows || fa_v2_on() {
29388 let sh = (if v4 {
29390 11520 + 32 * head_dim * if g { 1 } else { 2 }
29391 } else if v3 {
29392 32 * head_dim * 2
29393 } else {
29394 2 * 32 * head_dim * 2
29395 }) as u32;
29396 use cudarc::driver::sys::CUfunction_attribute_enum as A;
29397 f.set_attribute(
29398 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29399 sh as i32,
29400 )?;
29401 sh
29402 } else {
29403 0
29404 };
29405 for &(r0, t_g, sp_g) in &groups {
29409 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
29410 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
29411 let base_i = (base_len + r0) as i32;
29412 let o_len = t_g * n_head * n_splits_g * head_dim;
29413 let ml_len = t_g * n_head * n_splits_g;
29414 let mut part_guard = self.fa_part_pool.lock().unwrap();
29415 if part_guard
29416 .as_ref()
29417 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29418 .unwrap_or(true)
29419 {
29420 let old = part_guard.take();
29431 let (co, cm) = old
29432 .as_ref()
29433 .map(|pp| (pp.0.len(), pp.1.len()))
29434 .unwrap_or((0, 0));
29435 if let Some(old) = old {
29436 self.fa_part_retired.lock().unwrap().push(old);
29437 }
29438 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29439 eprintln!(
29440 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29441 co, o_len, cm, ml_len
29442 );
29443 }
29444 *part_guard =
29445 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29446 }
29447 let pg = part_guard.as_mut().unwrap();
29448 self.gpu
29449 .stream()
29450 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29451 self.gpu
29452 .stream()
29453 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29454 self.gpu
29455 .stream()
29456 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29457 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29458 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
29459 let qv = self.view(q, t * n_head * head_dim);
29460 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
29461 let cfg = LaunchConfig {
29462 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
29463 block_dim: (32, gqa, 1),
29464 shared_mem_bytes: shmem,
29465 };
29466 {
29467 let __s_b = self.gpu.stream();
29468 let mut b = __s_b.launch_builder(&f);
29469 if tb512 {
29470 let (bd, plus) =
29472 base_dev.expect("hd512 rows twin requires a device base counter");
29473 let plus_g = plus + r0 as i32;
29474 let nr = t_g as i32;
29475 if Self::pdl_on() && Self::pdl_wb_on() {
29476 use cudarc::driver::{DevicePtr, DevicePtrMut};
29478 let s = &self.gpu.stream();
29479 let (pq, _b0) = q_g.device_ptr(s);
29480 let (pk, _b1) = k.device_ptr(s);
29481 let (pv, _b2) = v.device_ptr(s);
29482 let (po, _b3) = part_o.device_ptr_mut(s);
29483 let (pm, _b4) = part_m.device_ptr_mut(s);
29484 let (pl, _b5) = part_l.device_ptr_mut(s);
29485 let (pb, _b6) = bd.device_ptr(s);
29486 let mut ps = [
29487 &pq as *const _ as *mut std::ffi::c_void,
29488 &pk as *const _ as *mut _,
29489 &pv as *const _ as *mut _,
29490 &po as *const _ as *mut _,
29491 &pm as *const _ as *mut _,
29492 &pl as *const _ as *mut _,
29493 &hd as *const _ as *mut _,
29494 &nh as *const _ as *mut _,
29495 &nhkv as *const _ as *mut _,
29496 &pb as *const _ as *mut _,
29497 &plus_g as *const _ as *mut _,
29498 &scale as *const _ as *mut _,
29499 &nspm as *const _ as *mut _,
29500 &spk as *const _ as *mut _,
29501 &ktb as *const _ as *mut _,
29502 &vtb as *const _ as *mut _,
29503 &nr as *const _ as *mut _,
29504 ];
29505 unsafe {
29506 self.launch_pdl_flash(
29507 Self::gkv_on(),
29508 "fa_decode_vec_q_rows_v4_512_tb",
29509 (n_head_kv as u32, n_splits_g as u32, 1),
29510 (32, gqa, 1),
29511 shmem,
29512 &mut ps,
29513 )?;
29514 }
29515 } else {
29516 let cfg_tb = LaunchConfig {
29517 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
29518 block_dim: (32, gqa, 1),
29519 shared_mem_bytes: shmem,
29520 };
29521 b.arg(&q_g)
29522 .arg(k)
29523 .arg(v)
29524 .arg(&mut *part_o)
29525 .arg(&mut *part_m)
29526 .arg(&mut *part_l)
29527 .arg(&hd)
29528 .arg(&nh)
29529 .arg(&nhkv)
29530 .arg(bd)
29531 .arg(&plus_g)
29532 .arg(&scale)
29533 .arg(&nspm)
29534 .arg(&spk)
29535 .arg(&ktb)
29536 .arg(&vtb)
29537 .arg(&nr);
29538 unsafe {
29539 b.launch(cfg_tb)?;
29540 }
29541 }
29542 } else if head_dim == 512 {
29543 let (bd, plus) =
29544 base_dev.expect("hd512 rows twin requires a device base counter");
29545 let plus_g = plus + r0 as i32;
29546 b.arg(&q_g)
29547 .arg(k)
29548 .arg(v)
29549 .arg(&mut *part_o)
29550 .arg(&mut *part_m)
29551 .arg(&mut *part_l)
29552 .arg(&hd)
29553 .arg(&nh)
29554 .arg(&nhkv)
29555 .arg(bd)
29556 .arg(&plus_g)
29557 .arg(&scale)
29558 .arg(&nspm)
29559 .arg(&spk)
29560 .arg(&ktb)
29561 .arg(&vtb);
29562 unsafe {
29563 b.launch(cfg)?;
29564 }
29565 } else {
29566 b.arg(&q_g)
29567 .arg(k)
29568 .arg(v)
29569 .arg(&mut *part_o)
29570 .arg(&mut *part_m)
29571 .arg(&mut *part_l)
29572 .arg(&hd)
29573 .arg(&nh)
29574 .arg(&nhkv)
29575 .arg(&base_i)
29576 .arg(&scale)
29577 .arg(&nspm)
29578 .arg(&spk)
29579 .arg(&ktb)
29580 .arg(&vtb);
29581 unsafe {
29582 b.launch(cfg)?;
29583 }
29584 }
29585 }
29586 let cfg2 = LaunchConfig {
29587 grid_dim: (n_head as u32, t_g as u32, 1),
29588 block_dim: (head_dim as u32, 1, 1),
29589 shared_mem_bytes: 0,
29590 };
29591 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
29592 if head_dim == 512 {
29593 let (bd, plus) = base_dev.unwrap();
29596 let plus_g = plus + r0 as i32;
29597 if let Some((oq, od)) = q8_out.as_mut() {
29598 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
29600 if Self::pdl_on() && Self::pdl_wb_on() {
29601 use cudarc::driver::{DevicePtr, DevicePtrMut};
29603 let s = &self.gpu.stream();
29604 let (po, _g0) = part_o.device_ptr(s);
29605 let (pm, _g1) = part_m.device_ptr(s);
29606 let (pl, _g2) = part_l.device_ptr(s);
29607 let (pq, _g3) = oq.device_ptr_mut(s);
29608 let (pd, _g4) = od.device_ptr_mut(s);
29609 let (pb, _g5) = bd.device_ptr(s);
29610 let mut ps = [
29611 &po as *const _ as *mut std::ffi::c_void,
29612 &pm as *const _ as *mut _,
29613 &pl as *const _ as *mut _,
29614 &pq as *const _ as *mut _,
29615 &pd as *const _ as *mut _,
29616 &hd as *const _ as *mut _,
29617 &nh as *const _ as *mut _,
29618 &pb as *const _ as *mut _,
29619 &plus_g as *const _ as *mut _,
29620 &nspm as *const _ as *mut _,
29621 &spk as *const _ as *mut _,
29622 ];
29623 unsafe {
29624 self.launch_pdl_flash(
29625 Self::gkv_on(),
29626 "fa_decode_combine_rows_dc_q8_1",
29627 cfg2.grid_dim,
29628 cfg2.block_dim,
29629 0,
29630 &mut ps,
29631 )?;
29632 }
29633 continue;
29634 }
29635 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
29636 let __s_b2 = self.gpu.stream();
29637 let mut b2 = __s_b2.launch_builder(&fc);
29638 b2.arg(&*part_o)
29639 .arg(&*part_m)
29640 .arg(&*part_l)
29641 .arg(&mut **oq)
29642 .arg(&mut **od)
29643 .arg(&hd)
29644 .arg(&nh)
29645 .arg(bd)
29646 .arg(&plus_g)
29647 .arg(&nspm)
29648 .arg(&spk);
29649 unsafe {
29650 b2.launch(cfg2)?;
29651 }
29652 continue;
29653 }
29654 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
29655 let __s_b2 = self.gpu.stream();
29656 let mut b2 = __s_b2.launch_builder(&fc);
29657 b2.arg(&*part_o)
29658 .arg(&*part_m)
29659 .arg(&*part_l)
29660 .arg(&mut o_g)
29661 .arg(&hd)
29662 .arg(&nh)
29663 .arg(bd)
29664 .arg(&plus_g)
29665 .arg(&nspm)
29666 .arg(&spk);
29667 unsafe {
29668 b2.launch(cfg2)?;
29669 }
29670 } else {
29671 assert!(
29674 q8_out.is_none(),
29675 "rows q8 emit requires the hd512 dc combine"
29676 );
29677 let fc = self.func("fa_decode_combine_rows");
29678 let __s_b2 = self.gpu.stream();
29679 let mut b2 = __s_b2.launch_builder(&fc);
29680 b2.arg(&*part_o)
29681 .arg(&*part_m)
29682 .arg(&*part_l)
29683 .arg(&mut o_g)
29684 .arg(&hd)
29685 .arg(&nh)
29686 .arg(&base_i)
29687 .arg(&nspm)
29688 .arg(&spk);
29689 unsafe {
29690 b2.launch(cfg2)?;
29691 }
29692 }
29693 }
29694 Ok(())
29695 }
29696
29697 #[allow(clippy::too_many_arguments)]
29701 pub fn fa_decode_rows_w(
29702 &self,
29703 q: &CudaSlice<f32>,
29704 k: &cudarc::driver::CudaView<u8>,
29705 v: &cudarc::driver::CudaView<u8>,
29706 o: &mut CudaSlice<f32>,
29707 head_dim: usize,
29708 n_head: usize,
29709 n_head_kv: usize,
29710 base_dev: &CudaSlice<i32>,
29711 base_plus: i32,
29712 t: usize,
29713 scale: f32,
29714 window: usize,
29715 k_tok_bytes: usize,
29716 v_tok_bytes: usize,
29717 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
29718 ) -> Result<(), Box<dyn std::error::Error>> {
29719 debug_assert!(head_dim == 256);
29724 let sp = {
29732 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29733 let v = *SPW.get_or_init(|| {
29734 std::env::var("MEMRA_FA_SPW")
29735 .ok()
29736 .and_then(|x| x.parse().ok())
29737 .unwrap_or(0)
29738 });
29739 if v >= 8 {
29740 v
29741 } else {
29742 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
29743 }
29744 };
29745 #[allow(clippy::manual_div_ceil)]
29746 let n_splits_max = (window + sp - 1) / sp;
29748 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
29749 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
29750 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
29751 let gqa = (n_head / n_head_kv).max(1) as u32;
29752 let o_len = t * n_head * n_splits_max * head_dim;
29753 let ml_len = t * n_head * n_splits_max;
29754 let mut part_guard = self.fa_part_pool.lock().unwrap();
29755 if part_guard
29756 .as_ref()
29757 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
29758 .unwrap_or(true)
29759 {
29760 let old = part_guard.take();
29771 let (co, cm) = old
29772 .as_ref()
29773 .map(|pp| (pp.0.len(), pp.1.len()))
29774 .unwrap_or((0, 0));
29775 if let Some(old) = old {
29776 self.fa_part_retired.lock().unwrap().push(old);
29777 }
29778 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
29779 eprintln!(
29780 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
29781 co, o_len, cm, ml_len
29782 );
29783 }
29784 *part_guard =
29785 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
29786 }
29787 let pg = part_guard.as_mut().unwrap();
29788 self.gpu
29789 .stream()
29790 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
29791 self.gpu
29792 .stream()
29793 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
29794 self.gpu
29795 .stream()
29796 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
29797 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
29798 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
29804 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
29805 std::env::var("MEMRA_FA_SMEM_TKV")
29806 .ok()
29807 .and_then(|v| v.parse().ok())
29808 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
29809 });
29810 use cudarc::driver::sys::CUfunction_attribute_enum as A;
29816 let wg = Self::wkv_on();
29821 let sp2 =
29824 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
29825 if sp2 {
29826 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
29827 if Self::pdl_on() && Self::pdl_wb_on() {
29828 use cudarc::driver::{DevicePtr, DevicePtrMut};
29830 let s = &self.gpu.stream();
29831 let (pq, _b0) = q.device_ptr(s);
29832 let (pk, _b1) = k.device_ptr(s);
29833 let (pv, _b2) = v.device_ptr(s);
29834 let (po, _b3) = part_o.device_ptr_mut(s);
29835 let (pm, _b4) = part_m.device_ptr_mut(s);
29836 let (pl, _b5) = part_l.device_ptr_mut(s);
29837 let (pb, _b6) = base_dev.device_ptr(s);
29838 let mut ps = [
29839 &pq as *const _ as *mut std::ffi::c_void,
29840 &pk as *const _ as *mut _,
29841 &pv as *const _ as *mut _,
29842 &po as *const _ as *mut _,
29843 &pm as *const _ as *mut _,
29844 &pl as *const _ as *mut _,
29845 &hd as *const _ as *mut _,
29846 &nh as *const _ as *mut _,
29847 &nhkv as *const _ as *mut _,
29848 &pb as *const _ as *mut _,
29849 &base_plus as *const _ as *mut _,
29850 &scale as *const _ as *mut _,
29851 &nspm as *const _ as *mut _,
29852 &spk as *const _ as *mut _,
29853 &ktb as *const _ as *mut _,
29854 &vtb as *const _ as *mut _,
29855 &wini as *const _ as *mut _,
29856 ];
29857 unsafe {
29858 self.launch_pdl_flash(
29859 wg,
29860 "fa_decode_vec_q_rows_v4_w_sp",
29861 (n_head_kv as u32, n_splits_max as u32, t as u32),
29862 (32, gqa + 1, 1),
29863 sh,
29864 &mut ps,
29865 )?;
29866 }
29867 } else {
29868 let f = if wg {
29869 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
29870 } else {
29871 self.func("fa_decode_vec_q_rows_v4_w_sp")
29872 };
29873 f.set_attribute(
29874 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29875 sh as i32,
29876 )?;
29877 let cfg = LaunchConfig {
29878 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
29879 block_dim: (32, gqa + 1, 1),
29880 shared_mem_bytes: sh,
29881 };
29882 let __s_b = self.gpu.stream();
29883 let mut b = __s_b.launch_builder(&f);
29884 b.arg(q)
29885 .arg(k)
29886 .arg(v)
29887 .arg(&mut *part_o)
29888 .arg(&mut *part_m)
29889 .arg(&mut *part_l)
29890 .arg(&hd)
29891 .arg(&nh)
29892 .arg(&nhkv)
29893 .arg(base_dev)
29894 .arg(&base_plus)
29895 .arg(&scale)
29896 .arg(&nspm)
29897 .arg(&spk)
29898 .arg(&ktb)
29899 .arg(&vtb)
29900 .arg(&wini);
29901 unsafe {
29902 b.launch(cfg)?;
29903 }
29904 }
29905 } else {
29906 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
29907 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
29909 use cudarc::driver::{DevicePtr, DevicePtrMut};
29910 let s = &self.gpu.stream();
29911 let (pq, _b0) = q.device_ptr(s);
29912 let (pk, _b1) = k.device_ptr(s);
29913 let (pv, _b2) = v.device_ptr(s);
29914 let (po, _b3) = part_o.device_ptr_mut(s);
29915 let (pm, _b4) = part_m.device_ptr_mut(s);
29916 let (pl, _b5) = part_l.device_ptr_mut(s);
29917 let (pb, _b6) = base_dev.device_ptr(s);
29918 let mut ps = [
29919 &pq as *const _ as *mut std::ffi::c_void,
29920 &pk as *const _ as *mut _,
29921 &pv as *const _ as *mut _,
29922 &po as *const _ as *mut _,
29923 &pm as *const _ as *mut _,
29924 &pl as *const _ as *mut _,
29925 &hd as *const _ as *mut _,
29926 &nh as *const _ as *mut _,
29927 &nhkv as *const _ as *mut _,
29928 &pb as *const _ as *mut _,
29929 &base_plus as *const _ as *mut _,
29930 &scale as *const _ as *mut _,
29931 &nspm as *const _ as *mut _,
29932 &spk as *const _ as *mut _,
29933 &ktb as *const _ as *mut _,
29934 &vtb as *const _ as *mut _,
29935 &wini as *const _ as *mut _,
29936 ];
29937 unsafe {
29938 self.launch_pdl_flash(
29939 wg,
29940 "fa_decode_vec_q_rows_v4_w",
29941 (n_head_kv as u32, n_splits_max as u32, t as u32),
29942 (32, gqa, 1),
29943 sh,
29944 &mut ps,
29945 )?;
29946 }
29947 } else {
29948 let pick = |name: &str| {
29949 if wg {
29950 self.func_g(name)
29951 } else {
29952 self.func(name)
29953 }
29954 };
29955 let (f, sh) = if fa_v4_at(window) {
29956 let f = pick("fa_decode_vec_q_rows_v4_w");
29957 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
29958 } else if smem_tkv > 0 && window >= smem_tkv {
29959 (
29962 pick("fa_decode_vec_q_rows_smem_w"),
29963 (2 * 32 * head_dim * 2) as u32,
29964 )
29965 } else {
29966 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
29967 };
29968 f.set_attribute(
29969 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
29970 sh as i32,
29971 )?;
29972 let cfg = LaunchConfig {
29973 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
29974 block_dim: (32, gqa, 1),
29975 shared_mem_bytes: sh,
29976 };
29977 let __s_b = self.gpu.stream();
29978 let mut b = __s_b.launch_builder(&f);
29979 b.arg(q)
29980 .arg(k)
29981 .arg(v)
29982 .arg(&mut *part_o)
29983 .arg(&mut *part_m)
29984 .arg(&mut *part_l)
29985 .arg(&hd)
29986 .arg(&nh)
29987 .arg(&nhkv)
29988 .arg(base_dev)
29989 .arg(&base_plus)
29990 .arg(&scale)
29991 .arg(&nspm)
29992 .arg(&spk)
29993 .arg(&ktb)
29994 .arg(&vtb)
29995 .arg(&wini);
29996 unsafe {
29997 b.launch(cfg)?;
29998 }
29999 }
30000 }
30001 let cfg2 = LaunchConfig {
30002 grid_dim: (n_head as u32, t as u32, 1),
30003 block_dim: (head_dim as u32, 1, 1),
30004 shared_mem_bytes: 0,
30005 };
30006 if let Some((oq, od)) = q8_out {
30007 if Self::pdl_on() && Self::pdl_wb_on() {
30010 use cudarc::driver::{DevicePtr, DevicePtrMut};
30012 let s = &self.gpu.stream();
30013 let (po, _g0) = part_o.device_ptr(s);
30014 let (pm, _g1) = part_m.device_ptr(s);
30015 let (pl, _g2) = part_l.device_ptr(s);
30016 let (pq, _g3) = oq.device_ptr_mut(s);
30017 let (pd, _g4) = od.device_ptr_mut(s);
30018 let mut ps = [
30019 &po as *const _ as *mut std::ffi::c_void,
30020 &pm as *const _ as *mut _,
30021 &pl as *const _ as *mut _,
30022 &pq as *const _ as *mut _,
30023 &pd as *const _ as *mut _,
30024 &hd as *const _ as *mut _,
30025 &nh as *const _ as *mut _,
30026 &nspm as *const _ as *mut _,
30027 &spk as *const _ as *mut _,
30028 &wini as *const _ as *mut _,
30029 ];
30030 unsafe {
30031 self.launch_pdl_flash(
30032 wg,
30033 "fa_decode_combine_rows_w_q8_1",
30034 cfg2.grid_dim,
30035 cfg2.block_dim,
30036 0,
30037 &mut ps,
30038 )?;
30039 }
30040 return Ok(());
30041 }
30042 let fc = if wg {
30043 self.func_g("fa_decode_combine_rows_w_q8_1")
30044 } else {
30045 self.func("fa_decode_combine_rows_w_q8_1")
30046 };
30047 let __s_b2 = self.gpu.stream();
30048 let mut b2 = __s_b2.launch_builder(&fc);
30049 b2.arg(&*part_o)
30050 .arg(&*part_m)
30051 .arg(&*part_l)
30052 .arg(oq)
30053 .arg(od)
30054 .arg(&hd)
30055 .arg(&nh)
30056 .arg(&nspm)
30057 .arg(&spk)
30058 .arg(&wini);
30059 unsafe {
30060 b2.launch(cfg2)?;
30061 }
30062 return Ok(());
30063 }
30064 let fc = if wg {
30065 self.func_g("fa_decode_combine_rows_w")
30066 } else {
30067 self.func("fa_decode_combine_rows_w")
30068 };
30069 let __s_b2 = self.gpu.stream();
30070 let mut b2 = __s_b2.launch_builder(&fc);
30071 b2.arg(&*part_o)
30072 .arg(&*part_m)
30073 .arg(&*part_l)
30074 .arg(o)
30075 .arg(&hd)
30076 .arg(&nh)
30077 .arg(&nspm)
30078 .arg(&spk)
30079 .arg(&wini);
30080 unsafe {
30081 b2.launch(cfg2)?;
30082 }
30083 Ok(())
30084 }
30085
30086 #[allow(clippy::too_many_arguments)]
30092 pub fn fa_decode_rows_dc(
30093 &self,
30094 q: &CudaSlice<f32>,
30095 k: &cudarc::driver::CudaView<u8>,
30096 v: &cudarc::driver::CudaView<u8>,
30097 o: &mut CudaSlice<f32>,
30098 head_dim: usize,
30099 n_head: usize,
30100 n_head_kv: usize,
30101 base_dev: &CudaSlice<i32>,
30102 t_kv_upper: usize,
30103 t: usize,
30104 scale: f32,
30105 k_tok_bytes: usize,
30106 v_tok_bytes: usize,
30107 base_plus: i32,
30108 g: bool,
30109 ) -> Result<(), Box<dyn std::error::Error>> {
30110 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
30111 assert!(
30112 v4 || fa_v3_active(head_dim),
30113 "stream fa rows requires the v3 or v4 lane"
30114 );
30115 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
30116 if v4 {
30117 let sp = fa_split_keys(t_kv_upper, n_head_kv);
30118 #[allow(clippy::manual_div_ceil)]
30119 let n_splits_max = (t_kv_upper + sp - 1) / sp;
30121 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
30122 let (nspm, spk) = (n_splits_max as i32, sp as i32);
30123 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30124 let gqa = (n_head / n_head_kv).max(1) as u32;
30125 let o_len = t * n_head * n_splits_max * head_dim;
30126 let ml_len = t * n_head * n_splits_max;
30127 let mut part_guard = self.fa_part_pool.lock().unwrap();
30128 if part_guard
30129 .as_ref()
30130 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30131 .unwrap_or(true)
30132 {
30133 let old = part_guard.take();
30144 let (co, cm) = old
30145 .as_ref()
30146 .map(|pp| (pp.0.len(), pp.1.len()))
30147 .unwrap_or((0, 0));
30148 if let Some(old) = old {
30149 self.fa_part_retired.lock().unwrap().push(old);
30150 }
30151 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
30152 eprintln!(
30153 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
30154 co, o_len, cm, ml_len
30155 );
30156 }
30157 *part_guard =
30158 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30159 }
30160 let pg = part_guard.as_mut().unwrap();
30161 self.gpu
30162 .stream()
30163 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
30164 self.gpu
30165 .stream()
30166 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
30167 self.gpu
30168 .stream()
30169 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
30170 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30171 let f = if g {
30172 self.func_g("fa_decode_vec_q_rows_v4_dc")
30173 } else {
30174 self.func("fa_decode_vec_q_rows_v4_dc")
30175 };
30176 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
30177 use cudarc::driver::sys::CUfunction_attribute_enum as A;
30178 f.set_attribute(
30179 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30180 sh as i32,
30181 )?;
30182 let cfg = LaunchConfig {
30183 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
30184 block_dim: (32, gqa, 1),
30185 shared_mem_bytes: sh,
30186 };
30187 let __s_b = self.gpu.stream();
30188 let mut b = __s_b.launch_builder(&f);
30189 b.arg(q)
30190 .arg(k)
30191 .arg(v)
30192 .arg(&mut *part_o)
30193 .arg(&mut *part_m)
30194 .arg(&mut *part_l)
30195 .arg(&hd)
30196 .arg(&nh)
30197 .arg(&nhkv)
30198 .arg(base_dev)
30199 .arg(&base_plus)
30200 .arg(&scale)
30201 .arg(&nspm)
30202 .arg(&spk)
30203 .arg(&ktb)
30204 .arg(&vtb);
30205 unsafe {
30206 b.launch(cfg)?;
30207 }
30208 let fc = self.func("fa_decode_combine_rows_dc");
30209 let cfg2 = LaunchConfig {
30210 grid_dim: (n_head as u32, t as u32, 1),
30211 block_dim: (head_dim as u32, 1, 1),
30212 shared_mem_bytes: 0,
30213 };
30214 let __s_b2 = self.gpu.stream();
30215 let mut b2 = __s_b2.launch_builder(&fc);
30216 b2.arg(&*part_o)
30217 .arg(&*part_m)
30218 .arg(&*part_l)
30219 .arg(o)
30220 .arg(&hd)
30221 .arg(&nh)
30222 .arg(base_dev)
30223 .arg(&base_plus)
30224 .arg(&nspm)
30225 .arg(&spk);
30226 unsafe {
30227 b2.launch(cfg2)?;
30228 }
30229 return Ok(());
30230 }
30231 let sp = fa_split_keys(t_kv_upper, n_head_kv);
30232 #[allow(clippy::manual_div_ceil)]
30233 let n_splits_max = (t_kv_upper + sp - 1) / sp;
30235 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
30236 let (nspm, spk) = (n_splits_max as i32, sp as i32);
30237 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30238 let gqa = (n_head / n_head_kv).max(1) as u32;
30239 let o_len = t * n_head * n_splits_max * head_dim;
30240 let ml_len = t * n_head * n_splits_max;
30241 let mut part_guard = self.fa_part_pool.lock().unwrap();
30242 if part_guard
30243 .as_ref()
30244 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30245 .unwrap_or(true)
30246 {
30247 let old = part_guard.take();
30258 let (co, cm) = old
30259 .as_ref()
30260 .map(|pp| (pp.0.len(), pp.1.len()))
30261 .unwrap_or((0, 0));
30262 if let Some(old) = old {
30263 self.fa_part_retired.lock().unwrap().push(old);
30264 }
30265 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
30266 eprintln!(
30267 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
30268 co, o_len, cm, ml_len
30269 );
30270 }
30271 *part_guard =
30272 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30273 }
30274 let pg = part_guard.as_mut().unwrap();
30275 self.gpu
30276 .stream()
30277 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
30278 self.gpu
30279 .stream()
30280 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
30281 self.gpu
30282 .stream()
30283 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
30284 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30285 let f = self.func("fa_decode_vec_q_rows_v3_dc");
30286 let sh = (32 * head_dim * 2) as u32;
30287 use cudarc::driver::sys::CUfunction_attribute_enum as A;
30288 f.set_attribute(
30289 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30290 sh as i32,
30291 )?;
30292 let cfg = LaunchConfig {
30293 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
30294 block_dim: (32, gqa, 1),
30295 shared_mem_bytes: sh,
30296 };
30297 let __s_b = self.gpu.stream();
30298 let mut b = __s_b.launch_builder(&f);
30299 b.arg(q)
30300 .arg(k)
30301 .arg(v)
30302 .arg(&mut *part_o)
30303 .arg(&mut *part_m)
30304 .arg(&mut *part_l)
30305 .arg(&hd)
30306 .arg(&nh)
30307 .arg(&nhkv)
30308 .arg(base_dev)
30309 .arg(&scale)
30310 .arg(&nspm)
30311 .arg(&spk)
30312 .arg(&ktb)
30313 .arg(&vtb);
30314 unsafe {
30315 b.launch(cfg)?;
30316 }
30317 let fc = self.func("fa_decode_combine_rows_dc");
30318 let cfg2 = LaunchConfig {
30319 grid_dim: (n_head as u32, t as u32, 1),
30320 block_dim: (head_dim as u32, 1, 1),
30321 shared_mem_bytes: 0,
30322 };
30323 let plus0 = 0i32;
30324 let __s_b2 = self.gpu.stream();
30325 let mut b2 = __s_b2.launch_builder(&fc);
30326 b2.arg(&*part_o)
30327 .arg(&*part_m)
30328 .arg(&*part_l)
30329 .arg(o)
30330 .arg(&hd)
30331 .arg(&nh)
30332 .arg(base_dev)
30333 .arg(&plus0)
30334 .arg(&nspm)
30335 .arg(&spk);
30336 unsafe {
30337 b2.launch(cfg2)?;
30338 }
30339 Ok(())
30340 }
30341
30342 #[allow(clippy::too_many_arguments)] pub fn fa_decode_dc(
30354 &self,
30355 q: &CudaSlice<f32>,
30356 k: &cudarc::driver::CudaView<u8>,
30357 v: &cudarc::driver::CudaView<u8>,
30358 o: &mut CudaSlice<f32>,
30359 head_dim: usize,
30360 n_head: usize,
30361 n_head_kv: usize,
30362 t_kv_dev: &CudaSlice<i32>,
30363 bucket_max: usize,
30364 scale: f32,
30365 k_tok_bytes: usize,
30366 v_tok_bytes: usize,
30367 g: bool,
30368 ) -> Result<(), Box<dyn std::error::Error>> {
30369 self.fa_decode_dc_q8(
30370 q,
30371 k,
30372 v,
30373 o,
30374 head_dim,
30375 n_head,
30376 n_head_kv,
30377 t_kv_dev,
30378 bucket_max,
30379 scale,
30380 k_tok_bytes,
30381 v_tok_bytes,
30382 g,
30383 None,
30384 )
30385 }
30386
30387 #[allow(clippy::too_many_arguments)]
30390 #[allow(clippy::manual_div_ceil)] pub fn fa_decode_dc_q8(
30392 &self,
30393 q: &CudaSlice<f32>,
30394 k: &cudarc::driver::CudaView<u8>,
30395 v: &cudarc::driver::CudaView<u8>,
30396 o: &mut CudaSlice<f32>,
30397 head_dim: usize,
30398 n_head: usize,
30399 n_head_kv: usize,
30400 t_kv_dev: &CudaSlice<i32>,
30401 bucket_max: usize,
30402 scale: f32,
30403 k_tok_bytes: usize,
30404 v_tok_bytes: usize,
30405 g: bool,
30406 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
30407 ) -> Result<(), Box<dyn std::error::Error>> {
30408 let mut fa_vec =
30416 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
30417 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
30418 fa_vec = false;
30419 } let sp = fa_split_keys(bucket_max, n_head_kv);
30421 let n_splits = if fa_vec {
30422 ((bucket_max + sp - 1) / sp).max(1)
30423 } else {
30424 ((bucket_max + 255) / 256).max(1)
30425 };
30426 let o_len = n_head * n_splits * head_dim;
30427 let ml_len = n_head * n_splits;
30428 let mut part_guard = self.fa_part_pool.lock().unwrap();
30429 if part_guard
30430 .as_ref()
30431 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30432 .unwrap_or(true)
30433 {
30434 let old = part_guard.take();
30445 let (co, cm) = old
30446 .as_ref()
30447 .map(|pp| (pp.0.len(), pp.1.len()))
30448 .unwrap_or((0, 0));
30449 if let Some(old) = old {
30450 self.fa_part_retired.lock().unwrap().push(old);
30451 }
30452 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
30453 eprintln!(
30454 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
30455 co, o_len, cm, ml_len
30456 );
30457 }
30458 *part_guard =
30459 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30460 }
30461 let pg = part_guard.as_mut().unwrap();
30462 self.gpu
30463 .stream()
30464 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
30465 self.gpu
30466 .stream()
30467 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
30468 self.gpu
30469 .stream()
30470 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
30471 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30472 let (hd, nh, nhkv, nsp) = (
30473 head_dim as i32,
30474 n_head as i32,
30475 n_head_kv as i32,
30476 n_splits as i32,
30477 );
30478 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30479 let fa_vec = fa_vec && head_dim <= 512 && head_dim.is_multiple_of(32);
30480 let deep = fa_vec
30483 && head_dim == 256
30484 && fa_v4_at(bucket_max)
30485 && !g
30486 && fa_deep_at(bucket_max)
30487 && !matches!(fa_v4_mode(), "noB3" | "stage");
30488 let (f, cfg) = if fa_vec
30489 && head_dim == 512
30490 && bucket_max >= {
30491 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
30492 *FA512_MIN_DC.get_or_init(|| {
30493 std::env::var("MEMRA_FA512_MIN")
30494 .ok()
30495 .and_then(|v| v.parse().ok())
30496 .unwrap_or(512)
30497 })
30498 } {
30499 let gqa = (n_head / n_head_kv).max(1) as u32;
30501 (
30502 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
30503 LaunchConfig {
30504 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30505 block_dim: (32, gqa, 1),
30506 shared_mem_bytes: 0,
30507 },
30508 )
30509 } else if fa_vec && head_dim == 512 {
30510 let q_view = q.as_view();
30513 let mut o_view = o.as_view_mut();
30514 return self.fa_decode_scalar_unified(
30515 &q_view,
30516 k,
30517 v,
30518 &mut o_view,
30519 head_dim,
30520 n_head,
30521 n_head_kv,
30522 0,
30523 Some(t_kv_dev),
30524 scale,
30525 n_splits,
30526 sp,
30527 k_tok_bytes,
30528 v_tok_bytes,
30529 g,
30530 &mut *part_o,
30531 &mut *part_m,
30532 &mut *part_l,
30533 q8_out,
30534 );
30535 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
30536 let gqa = (n_head / n_head_kv).max(1) as u32;
30539 let fv = if g {
30540 self.func_g("fa_decode_vec_q_v4_dc")
30541 } else if deep {
30542 self.func("fa_decode_vec_q_v4_deep_dc")
30543 } else {
30544 self.func("fa_decode_vec_q_v4_dc")
30545 };
30546 let shmem =
30547 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
30548 use cudarc::driver::sys::CUfunction_attribute_enum as A;
30549 fv.set_attribute(
30550 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
30551 shmem as i32,
30552 )?;
30553 (
30554 fv,
30555 LaunchConfig {
30556 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30557 block_dim: (32, gqa, 1),
30558 shared_mem_bytes: shmem,
30559 },
30560 )
30561 } else if fa_vec && fa_v3_active(head_dim) {
30562 let gqa = (n_head / n_head_kv).max(1) as u32;
30565 let fv = if g {
30566 self.func_g("fa_decode_vec_q_v3_dc")
30567 } else {
30568 self.func("fa_decode_vec_q_v3_dc")
30569 };
30570 let shmem = (32 * head_dim * 2) as u32; (
30572 fv,
30573 LaunchConfig {
30574 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30575 block_dim: (32, gqa, 1),
30576 shared_mem_bytes: shmem,
30577 },
30578 )
30579 } else if fa_vec && fa_v2_on() {
30580 let gqa = (n_head / n_head_kv).max(1) as u32;
30584 let fv = if g {
30585 self.func_g("fa_decode_vec_q_v2_dc")
30586 } else {
30587 self.func("fa_decode_vec_q_v2_dc")
30588 };
30589 let shmem = (2 * 32 * head_dim * 2) as u32; (
30591 fv,
30592 LaunchConfig {
30593 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30594 block_dim: (32, gqa, 1),
30595 shared_mem_bytes: shmem,
30596 },
30597 )
30598 } else if fa_vec {
30599 let gqa = (n_head / n_head_kv).max(1) as u32;
30600 let fv = if g {
30602 self.func_g("fa_decode_vec_q_dc")
30603 } else {
30604 self.func("fa_decode_vec_q_dc")
30605 };
30606 (
30607 fv,
30608 LaunchConfig {
30609 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30610 block_dim: (32, gqa, 1),
30611 shared_mem_bytes: 0,
30612 },
30613 )
30614 } else {
30615 let q_view = q.as_view();
30616 let mut o_view = o.as_view_mut();
30617 return self.fa_decode_scalar_unified(
30618 &q_view,
30619 k,
30620 v,
30621 &mut o_view,
30622 head_dim,
30623 n_head,
30624 n_head_kv,
30625 0,
30626 Some(t_kv_dev),
30627 scale,
30628 n_splits,
30629 if fa_vec { sp } else { 256 },
30630 k_tok_bytes,
30631 v_tok_bytes,
30632 g,
30633 &mut *part_o,
30634 &mut *part_m,
30635 &mut *part_l,
30636 q8_out,
30637 );
30638 };
30639 let ski = sp as i32; let __s_b = self.gpu.stream();
30641 let mut b = __s_b.launch_builder(&f);
30642 b.arg(q)
30643 .arg(k)
30644 .arg(v)
30645 .arg(&mut *part_o)
30646 .arg(&mut *part_m)
30647 .arg(&mut *part_l)
30648 .arg(&hd)
30649 .arg(&nh)
30650 .arg(&nhkv)
30651 .arg(t_kv_dev)
30652 .arg(&scale)
30653 .arg(&nsp)
30654 .arg(&ski)
30655 .arg(&ktb)
30656 .arg(&vtb);
30657 unsafe {
30658 b.launch(cfg)?;
30659 }
30660 let cfg2 = LaunchConfig {
30661 grid_dim: (n_head as u32, 1, 1),
30662 block_dim: (head_dim as u32, 1, 1),
30663 shared_mem_bytes: 0,
30664 };
30665 if let Some((oq, od)) = q8_out {
30666 let fc = if g {
30667 self.func_g("fa_decode_combine_q8_1")
30668 } else {
30669 self.fa_func("fa_decode_combine_q8_1", head_dim)
30670 };
30671 let __s_b2 = self.gpu.stream();
30672 let mut b2 = __s_b2.launch_builder(&fc);
30673 b2.arg(&*part_o)
30674 .arg(&*part_m)
30675 .arg(&*part_l)
30676 .arg(oq)
30677 .arg(od)
30678 .arg(&hd)
30679 .arg(&nh)
30680 .arg(&nsp);
30681 unsafe {
30682 b2.launch(cfg2)?;
30683 }
30684 return Ok(());
30685 }
30686 let fc = if g {
30687 self.func_g("fa_decode_combine_f32")
30688 } else {
30689 self.fa_func("fa_decode_combine_f32", head_dim)
30690 };
30691 let __s_b2 = self.gpu.stream();
30692 let mut b2 = __s_b2.launch_builder(&fc);
30693 b2.arg(&*part_o)
30694 .arg(&*part_m)
30695 .arg(&*part_l)
30696 .arg(o)
30697 .arg(&hd)
30698 .arg(&nh)
30699 .arg(&nsp);
30700 unsafe {
30701 b2.launch(cfg2)?;
30702 }
30703 Ok(())
30704 }
30705
30706 #[allow(clippy::too_many_arguments)]
30710 pub fn append_kv_quantized_dcw(
30711 &self,
30712 k_row: &CudaSlice<f32>,
30713 v_row: &CudaSlice<f32>,
30714 kc: &mut CudaSlice<u8>,
30715 vc: &mut CudaSlice<u8>,
30716 len_dev: &CudaSlice<i32>,
30717 base_dev: Option<&CudaSlice<i32>>,
30718 kv_dim_k: usize,
30719 kv_dim_v: usize,
30720 k_tok_bytes: usize,
30721 v_tok_bytes: usize,
30722 ) -> Result<(), Box<dyn std::error::Error>> {
30723 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
30724 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
30725 let cfg = LaunchConfig {
30726 grid_dim: (nblk, 1, 1),
30727 block_dim: (32, 1, 1),
30728 shared_mem_bytes: 0,
30729 };
30730 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
30731 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30732 let null: u64 = 0;
30733 let __s_b = self.gpu.stream();
30734 let mut b = __s_b.launch_builder(&f);
30735 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
30736 match base_dev {
30737 Some(base) => {
30738 b.arg(base);
30739 }
30740 None => {
30741 b.arg(&null);
30742 }
30743 }
30744 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
30745 unsafe {
30746 b.launch(cfg)?;
30747 }
30748 Ok(())
30749 }
30750
30751 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
30753 let f = self.func("inc_i32");
30754 let cfg = LaunchConfig {
30755 grid_dim: (1, 1, 1),
30756 block_dim: (1, 1, 1),
30757 shared_mem_bytes: 0,
30758 };
30759 let __s_b = self.gpu.stream();
30760 let mut b = __s_b.launch_builder(&f);
30761 b.arg(counter);
30762 unsafe {
30763 b.launch(cfg)?;
30764 }
30765 Ok(())
30766 }
30767
30768 #[allow(clippy::too_many_arguments)]
30777 #[allow(clippy::type_complexity)] fn fa_part_alloc(
30802 &self,
30803 o_len: usize,
30804 ml_len: usize,
30805 co: usize,
30806 cm: usize,
30807 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
30808 static GROWS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
30809 let n = GROWS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
30810 if n < 64 {
30811 eprintln!(
30812 "[fa-pool] grow #{n} dev={} o_len {co} -> {o_len} ml_len {cm} -> {ml_len} (retired kept, zero={})",
30813 self.ctx().ordinal(),
30814 fa_part_zero_on()
30815 );
30816 }
30817 let mut po = self.alloc_uninit::<f32>(o_len)?;
30818 let mut pm = self.alloc_uninit::<f32>(ml_len)?;
30819 let mut pl = self.alloc_uninit::<f32>(ml_len)?;
30820 if fa_part_zero_on() {
30821 self.gpu.stream().memset_zeros(&mut po)?;
30822 self.gpu.stream().memset_zeros(&mut pm)?;
30823 self.gpu.stream().memset_zeros(&mut pl)?;
30824 }
30825 Ok((po, pm, pl))
30826 }
30827
30828 fn fa_part_pool_grow(
30829 &self,
30830 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
30831 o_len: usize,
30832 ml_len: usize,
30833 ) -> Result<(), Box<dyn std::error::Error>> {
30834 if part_guard
30835 .as_ref()
30836 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
30837 .unwrap_or(true)
30838 {
30839 let old = part_guard.take();
30840 let (co, cm) = old
30841 .as_ref()
30842 .map(|pp| (pp.0.len(), pp.1.len()))
30843 .unwrap_or((0, 0));
30844 if let Some(old) = old {
30845 self.fa_part_retired.lock().unwrap().push(old);
30846 }
30847 *part_guard =
30860 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
30861 }
30862 Ok(())
30863 }
30864
30865 pub fn fa_dcw_pool_ensure(
30868 &self,
30869 head_dim: usize,
30870 n_head: usize,
30871 n_head_kv: usize,
30872 bucket_max: usize,
30873 ) -> Result<(), Box<dyn std::error::Error>> {
30874 let sp = fa_split_keys(bucket_max, n_head_kv);
30875 #[allow(clippy::manual_div_ceil)]
30876 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
30878 let o_len = n_head * n_splits * head_dim;
30879 let ml_len = n_head * n_splits;
30880 let mut part_guard = self.fa_part_pool.lock().unwrap();
30881 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
30882 }
30883
30884 #[allow(clippy::too_many_arguments)]
30892 pub fn fa_decode_dcw2(
30893 &self,
30894 q2: &CudaSlice<f32>,
30895 k_ring: &cudarc::driver::CudaView<u8>,
30896 v_ring: &cudarc::driver::CudaView<u8>,
30897 o2: &mut CudaSlice<f32>,
30898 head_dim: usize,
30899 n_head: usize,
30900 n_head_kv: usize,
30901 len_dev: &CudaSlice<i32>,
30902 base_dev: Option<&CudaSlice<i32>>,
30903 window: usize,
30904 bucket_max: usize,
30905 scale: f32,
30906 k_tok_bytes: usize,
30907 v_tok_bytes: usize,
30908 gate2: &CudaSlice<f32>,
30909 ) -> Result<(), Box<dyn std::error::Error>> {
30910 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
30911 if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) || !fa_v3_on() {
30912 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
30913 }
30914 let sp = fa_split_keys(bucket_max, n_head_kv);
30915 #[allow(clippy::manual_div_ceil)]
30916 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
30918 let o_len = 2 * n_head * n_splits * head_dim;
30920 let ml_len = 2 * n_head * n_splits;
30921 let mut part_guard = self.fa_part_pool.lock().unwrap();
30922 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
30923 let pg = part_guard.as_mut().unwrap();
30924 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
30925 let (hd, nh, nhkv, nsp) = (
30926 head_dim as i32,
30927 n_head as i32,
30928 n_head_kv as i32,
30929 n_splits as i32,
30930 );
30931 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
30932 let (ski, win) = (sp as i32, window as i32);
30933 let gqa = (n_head / n_head_kv).max(1) as u32;
30934 let smem = (32 * head_dim * 2) as u32;
30935 let f = self.func("fa_decode_vec_q_v3_dcw2");
30936 let cfg = LaunchConfig {
30937 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
30938 block_dim: (32, gqa, 1),
30939 shared_mem_bytes: smem,
30940 };
30941 let null: u64 = 0;
30942 {
30943 let __s_b = self.gpu.stream();
30944 let mut b = __s_b.launch_builder(&f);
30945 b.arg(q2)
30946 .arg(k_ring)
30947 .arg(v_ring)
30948 .arg(&mut *part_o)
30949 .arg(&mut *part_m)
30950 .arg(&mut *part_l)
30951 .arg(&hd)
30952 .arg(&nh)
30953 .arg(&nhkv)
30954 .arg(len_dev);
30955 match base_dev {
30956 Some(base) => {
30957 b.arg(base);
30958 }
30959 None => {
30960 b.arg(&null);
30961 }
30962 }
30963 b.arg(&win)
30964 .arg(&scale)
30965 .arg(&nsp)
30966 .arg(&ski)
30967 .arg(&ktb)
30968 .arg(&vtb);
30969 unsafe {
30970 b.launch(cfg)?;
30971 }
30972 }
30973 let fc = {
30977 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
30978 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
30979 self.func("fa_decode_combine_gate_f32_s")
30980 } else {
30981 self.func("fa_decode_combine_gate_f32")
30982 }
30983 };
30984 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
30985 let nh2 = (2 * n_head) as i32;
30986 let cfg2 = LaunchConfig {
30987 grid_dim: ((2 * n_head) as u32, 1, 1),
30988 block_dim: (head_dim as u32, 1, 1),
30989 shared_mem_bytes: if combine_shared {
30990 (2 * n_splits * 4) as u32
30991 } else {
30992 0
30993 },
30994 };
30995 let __s_b2 = self.gpu.stream();
30996 let mut b2 = __s_b2.launch_builder(&fc);
30997 b2.arg(&*part_o)
30998 .arg(&*part_m)
30999 .arg(&*part_l)
31000 .arg(gate2)
31001 .arg(o2)
31002 .arg(&hd)
31003 .arg(&nh2)
31004 .arg(&nsp);
31005 unsafe {
31006 b2.launch(cfg2)?;
31007 }
31008 Ok(())
31009 }
31010
31011 #[allow(clippy::too_many_arguments)]
31020 pub fn fa_decode_dcw_rows(
31021 &self,
31022 q_rows: &CudaSlice<f32>,
31023 tab: &CudaSlice<u64>,
31024 o_rows: &mut CudaSlice<f32>,
31025 t: usize,
31026 head_dim: usize,
31027 n_head: usize,
31028 n_head_kv: usize,
31029 window: usize,
31030 max_ns: usize,
31031 scale: f32,
31032 k_tok_bytes: usize,
31033 v_tok_bytes: usize,
31034 gate_rows: &CudaSlice<f32>,
31035 ) -> Result<(), Box<dyn std::error::Error>> {
31036 if std::env::var("MEMRA_NO_FA_VEC").is_ok()
31037 || head_dim > 256
31038 || !head_dim.is_multiple_of(32)
31039 || !fa_v3_on()
31040 {
31041 return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
31042 }
31043 if fa_sm_count() < 128
31044 || std::env::var("MEMRA_FA_SPLIT").is_ok()
31045 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
31046 || std::env::var("MEMRA_FA_SP16").is_ok()
31047 {
31048 return Err(
31049 "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
31050 (or a <128-SM rig) keep the per-row path"
31051 .into(),
31052 );
31053 }
31054 if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
31055 return Err("fa_decode_dcw_rows geometry".into());
31056 }
31057 let o_len = t * n_head * max_ns * head_dim;
31058 let ml_len = t * n_head * max_ns;
31059 let mut part_guard = self.fa_part_pool.lock().unwrap();
31060 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
31061 let pg = part_guard.as_mut().unwrap();
31062 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
31063 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
31064 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
31065 let (win, mns) = (window as i32, max_ns as i32);
31066 let gqa = (n_head / n_head_kv).max(1) as u32;
31067 let smem = (32 * head_dim * 2) as u32;
31068 let f = self.func("fa_decode_vec_q_v3_dcw_rows");
31069 let cfg = LaunchConfig {
31070 grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
31071 block_dim: (32, gqa, 1),
31072 shared_mem_bytes: smem,
31073 };
31074 {
31075 let __s_b = self.gpu.stream();
31076 let mut b = __s_b.launch_builder(&f);
31077 b.arg(q_rows)
31078 .arg(tab)
31079 .arg(&mut *part_o)
31080 .arg(&mut *part_m)
31081 .arg(&mut *part_l)
31082 .arg(&hd)
31083 .arg(&nh)
31084 .arg(&nhkv)
31085 .arg(&win)
31086 .arg(&scale)
31087 .arg(&mns)
31088 .arg(&ktb)
31089 .arg(&vtb);
31090 unsafe {
31091 b.launch(cfg)?;
31092 }
31093 }
31094 let fc = {
31098 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31099 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
31100 self.func("fa_decode_combine_gate_f32_s")
31101 } else {
31102 self.func("fa_decode_combine_gate_f32")
31103 }
31104 };
31105 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
31106 let nht = (t * n_head) as i32;
31107 let cfg2 = LaunchConfig {
31108 grid_dim: ((t * n_head) as u32, 1, 1),
31109 block_dim: (head_dim as u32, 1, 1),
31110 shared_mem_bytes: if combine_shared {
31111 (2 * max_ns * 4) as u32
31112 } else {
31113 0
31114 },
31115 };
31116 let __s_b2 = self.gpu.stream();
31117 let mut b2 = __s_b2.launch_builder(&fc);
31118 b2.arg(&*part_o)
31119 .arg(&*part_m)
31120 .arg(&*part_l)
31121 .arg(gate_rows)
31122 .arg(o_rows)
31123 .arg(&hd)
31124 .arg(&nht)
31125 .arg(&mns);
31126 unsafe {
31127 b2.launch(cfg2)?;
31128 }
31129 Ok(())
31130 }
31131
31132 #[allow(clippy::too_many_arguments)] pub fn fa_decode_dcw(
31134 &self,
31135 q: &CudaSlice<f32>,
31136 k_ring: &cudarc::driver::CudaView<u8>,
31137 v_ring: &cudarc::driver::CudaView<u8>,
31138 o: &mut CudaSlice<f32>,
31139 head_dim: usize,
31140 n_head: usize,
31141 n_head_kv: usize,
31142 len_dev: &CudaSlice<i32>,
31143 base_dev: Option<&CudaSlice<i32>>,
31144 window: usize,
31145 bucket_max: usize,
31146 scale: f32,
31147 k_tok_bytes: usize,
31148 v_tok_bytes: usize,
31149 fused_gate: Option<&CudaSlice<f32>>,
31153 ) -> Result<(), Box<dyn std::error::Error>> {
31154 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
31155 if !fa_vec || head_dim > 256 || !head_dim.is_multiple_of(32) || !fa_v3_on() {
31156 return Err("fa_decode_dcw supports the default v3-vec class only (bucket >= vec floor, head_dim <= 256, MEMRA_FA_V3 on); keep eager outside it"
31157 .into());
31158 }
31159 let sp = fa_split_keys(bucket_max, n_head_kv);
31160 #[allow(clippy::manual_div_ceil)]
31161 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
31163 let o_len = n_head * n_splits * head_dim;
31164 let ml_len = n_head * n_splits;
31165 let mut part_guard = self.fa_part_pool.lock().unwrap();
31166 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
31167 let pg = part_guard.as_mut().unwrap();
31168 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31173 let memset_on = *MEMSET_ON
31178 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
31179 || crate::tp::token_graph_building();
31180 if memset_on {
31181 self.gpu
31182 .stream()
31183 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
31184 self.gpu
31185 .stream()
31186 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
31187 self.gpu
31188 .stream()
31189 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
31190 }
31191 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
31192 let (hd, nh, nhkv, nsp) = (
31193 head_dim as i32,
31194 n_head as i32,
31195 n_head_kv as i32,
31196 n_splits as i32,
31197 );
31198 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
31199 let (ski, win) = (sp as i32, window as i32);
31200 let gqa = (n_head / n_head_kv).max(1) as u32;
31201 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31205 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
31206 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
31207 Ok("2") => 2,
31208 Ok("1") => 1,
31209 _ => 0,
31210 });
31211 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31216 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
31217 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
31218 std::sync::Mutex::new(None);
31219 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31223 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
31224 && (n_head / n_head_kv).is_multiple_of(2)
31225 && (n_head / n_head_kv) >= 2;
31226 let f = if fprof {
31227 self.func("fa_decode_vec_q_v3_dcw_prof")
31228 } else if hs2 {
31229 self.func("fa_decode_vec_q_v3_dcw_hs2")
31230 } else if hoist == 2 {
31231 self.func("fa_decode_vec_q_v3_dcw_hc")
31233 } else if hoist == 1 {
31234 self.func("fa_decode_vec_q_v3_dcw_h")
31236 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
31237 self.func("fa_decode_vec_q_v3_dcw_u8")
31238 } else {
31239 self.func("fa_decode_vec_q_v3_dcw")
31240 };
31241 let cfg = LaunchConfig {
31242 grid_dim: if hs2 {
31243 ((2 * n_head_kv) as u32, n_splits as u32, 1)
31244 } else {
31245 (n_head_kv as u32, n_splits as u32, 1)
31246 },
31247 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
31248 shared_mem_bytes: smem,
31249 };
31250 let null: u64 = 0;
31251 let __s_b = self.gpu.stream();
31252 let mut b = __s_b.launch_builder(&f);
31253 b.arg(q)
31254 .arg(k_ring)
31255 .arg(v_ring)
31256 .arg(&mut *part_o)
31257 .arg(&mut *part_m)
31258 .arg(&mut *part_l)
31259 .arg(&hd)
31260 .arg(&nh)
31261 .arg(&nhkv)
31262 .arg(len_dev);
31263 match base_dev {
31264 Some(base) => {
31265 b.arg(base);
31266 }
31267 None => {
31268 b.arg(&null);
31269 }
31270 }
31271 b.arg(&win)
31272 .arg(&scale)
31273 .arg(&nsp)
31274 .arg(&ski)
31275 .arg(&ktb)
31276 .arg(&vtb);
31277 if fprof {
31278 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
31279 if guard
31280 .as_ref()
31281 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
31282 {
31283 *guard = Some((self.ctx().ordinal(), self.htod_u64(&[0u64; 8])?));
31284 }
31285 let (_, buf) = guard.as_mut().expect("armed above");
31286 b.arg(&*buf);
31287 unsafe {
31288 b.launch(cfg)?;
31289 }
31290 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
31291 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
31292 if n.is_multiple_of(430) {
31293 self.stream().synchronize()?;
31294 let h = self.dtoh_u64(buf)?;
31295 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
31296 let tot: u64 = h[..6].iter().sum();
31297 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
31298 for (i, name) in phases.iter().enumerate() {
31299 let pct = if tot > 0 {
31300 h[i] as f64 / tot as f64 * 100.0
31301 } else {
31302 0.0
31303 };
31304 line.push_str(&format!(" {name}={pct:.1}%"));
31305 }
31306 if h[6] > 0 {
31307 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
31308 }
31309 eprintln!("{line}");
31310 }
31311 } else {
31312 unsafe {
31313 b.launch(cfg)?;
31314 }
31315 }
31316 let mut combine_shared = false;
31317 let fc = if fused_gate.is_some() {
31318 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31321 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
31322 combine_shared = true;
31323 self.func("fa_decode_combine_gate_f32_s")
31324 } else {
31325 self.func("fa_decode_combine_gate_f32")
31326 }
31327 } else {
31328 self.fa_func("fa_decode_combine_f32", head_dim)
31329 };
31330 let cfg2 = LaunchConfig {
31331 grid_dim: (n_head as u32, 1, 1),
31332 block_dim: (head_dim as u32, 1, 1),
31333 shared_mem_bytes: if combine_shared {
31334 (2 * n_splits * 4) as u32
31335 } else {
31336 0
31337 },
31338 };
31339 let __s_b2 = self.gpu.stream();
31340 let mut b2 = __s_b2.launch_builder(&fc);
31341 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
31342 if let Some(gate_row) = fused_gate {
31343 b2.arg(gate_row);
31344 }
31345 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
31346 unsafe {
31347 b2.launch(cfg2)?;
31348 }
31349 Ok(())
31350 }
31351
31352 #[allow(clippy::manual_div_ceil)] pub fn fa_geom_eager(
31359 &self,
31360 t_kv: usize,
31361 head_dim: usize,
31362 n_head_kv: usize,
31363 g: bool,
31364 ) -> (bool, usize) {
31365 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
31369 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
31375 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim.is_multiple_of(32));
31376 if g && head_dim == 256 && !fa_v4_at(t_kv) {
31382 fa_vec = false;
31383 }
31384 let sp = fa_split_keys(t_kv, n_head_kv);
31385 let n_splits = if fa_vec {
31386 ((t_kv + sp - 1) / sp).max(1)
31387 } else {
31388 ((t_kv + 255) / 256).max(1)
31389 };
31390 (fa_vec, n_splits)
31391 }
31392
31393 pub fn fa_bucket_key(
31399 &self,
31400 t_kv: usize,
31401 head_dim: usize,
31402 n_head_kv: usize,
31403 g: bool,
31404 ) -> (bool, usize) {
31405 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
31406 }
31407
31408 #[allow(clippy::type_complexity)] pub fn capture_graph_retained<F>(
31421 &self,
31422 step: F,
31423 ) -> Result<
31424 (
31425 cudarc::driver::CudaGraph,
31426 Vec<Box<dyn std::any::Any + Send>>,
31427 ),
31428 Box<dyn std::error::Error>,
31429 >
31430 where
31431 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
31432 {
31433 use cudarc::driver::sys::CUgraphInstantiate_flags;
31434 self.capture_graph_retained_flags(
31435 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
31436 step,
31437 )
31438 }
31439
31440 #[allow(clippy::type_complexity)] pub fn capture_graph_retained_flags<F>(
31446 &self,
31447 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
31448 mut step: F,
31449 ) -> Result<
31450 (
31451 cudarc::driver::CudaGraph,
31452 Vec<Box<dyn std::any::Any + Send>>,
31453 ),
31454 Box<dyn std::error::Error>,
31455 >
31456 where
31457 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
31458 {
31459 use cudarc::driver::sys::CUstreamCaptureMode;
31460 self.capture_keep.lock().unwrap().clear();
31468 let was_tracking = self.gpu.ctx.is_event_tracking();
31469 if was_tracking {
31470 unsafe {
31471 self.gpu.ctx.disable_event_tracking();
31472 }
31473 }
31474 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
31475 self.capture_keep_on
31476 .store(true, std::sync::atomic::Ordering::Relaxed);
31477 let w = (|| {
31478 step(self)?;
31479 step(self)
31480 })();
31481 self.capture_keep_on
31482 .store(false, std::sync::atomic::Ordering::Relaxed);
31483 w?;
31484 self.gpu.stream().synchronize()?;
31485 self.gpu
31486 .stream()
31487 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
31488 let r = step(self);
31489 let g = self.gpu.stream().end_capture(flags);
31490 r?;
31491 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
31492 graph.upload()?;
31493 Ok(graph)
31494 };
31495 let result = run();
31496 self.capture_keep_on
31497 .store(false, std::sync::atomic::Ordering::Relaxed);
31498 if was_tracking {
31499 unsafe {
31500 self.gpu.ctx.enable_event_tracking();
31501 }
31502 }
31503 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
31504 Ok((result?, keeper))
31505 }
31506
31507 #[allow(clippy::type_complexity)] pub fn capture_graph_retained_nowarm<F>(
31514 &self,
31515 mut step: F,
31516 ) -> Result<
31517 (
31518 cudarc::driver::CudaGraph,
31519 Vec<Box<dyn std::any::Any + Send>>,
31520 ),
31521 Box<dyn std::error::Error>,
31522 >
31523 where
31524 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
31525 {
31526 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
31527 let was_tracking = self.gpu.ctx.is_event_tracking();
31528 if was_tracking {
31529 unsafe {
31530 self.gpu.ctx.disable_event_tracking();
31531 }
31532 }
31533 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
31534 self.gpu.stream().synchronize()?;
31535 self.gpu
31536 .stream()
31537 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
31538 let r = step(self);
31539 let g = self.gpu.stream().end_capture(
31540 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
31541 );
31542 r?;
31543 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
31544 graph.upload()?;
31545 Ok(graph)
31546 };
31547 let result = run();
31548 if was_tracking {
31549 unsafe {
31550 self.gpu.ctx.enable_event_tracking();
31551 }
31552 }
31553 Ok((result?, Vec::new()))
31554 }
31555
31556 pub fn capture_graph<F>(
31557 &self,
31558 mut step: F,
31559 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
31560 where
31561 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
31562 {
31563 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
31564 let was_tracking = self.gpu.ctx.is_event_tracking();
31572 if was_tracking {
31573 unsafe {
31574 self.gpu.ctx.disable_event_tracking();
31575 }
31576 }
31577 let iflag = {
31584 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
31585 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
31586 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
31589 Ok("priority") => {
31590 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
31591 }
31592 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
31593 })
31594 };
31595 let ct = {
31602 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31603 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
31604 };
31605 let warmups = {
31628 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
31629 *W.get_or_init(|| {
31630 std::env::var("MEMRA_GRAPH_WARMUPS")
31631 .ok()
31632 .and_then(|v| v.parse().ok())
31633 .filter(|n| *n >= 1)
31634 .unwrap_or(1)
31635 })
31636 };
31637 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
31638 let t_w = std::time::Instant::now();
31639 for _ in 0..warmups {
31641 step(self)?;
31642 }
31643 self.gpu.stream().synchronize()?;
31644 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
31645 let t_c = std::time::Instant::now();
31647 self.gpu
31648 .stream()
31649 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
31650 let r = step(self);
31653 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
31654 let t_i = std::time::Instant::now();
31655 let g = self.gpu.stream().end_capture(iflag);
31656 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
31657 r?;
31658 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
31659 let t_u = std::time::Instant::now();
31660 graph.upload()?;
31661 if ct {
31662 println!(
31663 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
31664 instantiate {ms_inst:.2} ms upload {:.2} ms",
31665 t_u.elapsed().as_secs_f64() * 1e3
31666 );
31667 }
31668 Ok(graph)
31669 };
31670 let result = run();
31671 if was_tracking {
31672 unsafe {
31673 self.gpu.ctx.enable_event_tracking();
31674 }
31675 }
31676 result
31677 }
31678
31679 #[allow(clippy::too_many_arguments)] pub fn gdn_scan_s128_view(
31682 &self,
31683 q: &CudaSlice<f32>,
31684 k: &CudaSlice<f32>,
31685 v: &CudaSlice<f32>,
31686 g: &CudaSlice<f32>,
31687 beta: &CudaSlice<f32>,
31688 state_in: &cudarc::driver::CudaView<f32>,
31689 state_out: &mut cudarc::driver::CudaViewMut<f32>,
31690 o: &mut CudaSlice<f32>,
31691 n_head: usize,
31692 t: usize,
31693 scale: f32,
31694 ) -> Result<(), Box<dyn std::error::Error>> {
31695 let f = self.func("gdn_scan_s128");
31696 const S_V: u32 = 128;
31697 const WARP: u32 = 32;
31698 const COLS: u32 = 4;
31699 let cfg = LaunchConfig {
31700 grid_dim: (n_head as u32, 1, S_V / COLS),
31701 block_dim: (WARP, COLS, 1),
31702 shared_mem_bytes: 0,
31703 };
31704 let (h, ti) = (n_head as i32, t as i32);
31705 let __s_b = self.gpu.stream();
31706 let mut b = __s_b.launch_builder(&f);
31707 b.arg(q)
31708 .arg(k)
31709 .arg(v)
31710 .arg(g)
31711 .arg(beta)
31712 .arg(state_in)
31713 .arg(state_out)
31714 .arg(o)
31715 .arg(&h)
31716 .arg(&ti)
31717 .arg(&scale);
31718 unsafe {
31719 b.launch(cfg)?;
31720 }
31721 Ok(())
31722 }
31723
31724 #[allow(clippy::too_many_arguments)]
31726 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_view(
31729 &self,
31730 x: &cudarc::driver::CudaView<f32>,
31731 w: &CudaSlice<f32>,
31732 y: &mut CudaSlice<f32>,
31733 conv_dim: usize,
31734 t: usize,
31735 d_conv: usize,
31736 silu: bool,
31737 ) -> Result<(), Box<dyn std::error::Error>> {
31738 let f = self.func("ssm_conv1d_silu_f32");
31739 let cfg = LaunchConfig {
31741 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
31742 block_dim: (256, 1, 1),
31743 shared_mem_bytes: 0,
31744 };
31745 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
31746 let __s_b = self.gpu.stream();
31747 let mut b = __s_b.launch_builder(&f);
31748 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
31749 unsafe {
31750 b.launch(cfg)?;
31751 }
31752 Ok(())
31753 }
31754
31755 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_tm(
31763 &self,
31764 qkv_tm: &CudaSlice<f32>,
31765 w: &CudaSlice<f32>,
31766 y: &mut CudaSlice<f32>,
31767 conv_dim: usize,
31768 t: usize,
31769 d_conv: usize,
31770 ) -> Result<(), Box<dyn std::error::Error>> {
31771 let f = self.func("ssm_conv1d_tm_f32");
31772 let cfg = LaunchConfig {
31773 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31774 block_dim: (256, 1, 1),
31775 shared_mem_bytes: 0,
31776 };
31777 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31778 let __s_b = self.gpu.stream();
31779 let mut b = __s_b.launch_builder(&f);
31780 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
31781 unsafe {
31782 b.launch(cfg)?;
31783 }
31784 Ok(())
31785 }
31786
31787 #[allow(clippy::too_many_arguments)] pub fn ssm_conv1d_tm_state(
31796 &self,
31797 qkv_tm: &CudaSlice<f32>,
31798 conv_state: &mut CudaSlice<f32>,
31799 w: &CudaSlice<f32>,
31800 y: &mut CudaSlice<f32>,
31801 conv_dim: usize,
31802 t: usize,
31803 d_conv: usize,
31804 ) -> Result<(), Box<dyn std::error::Error>> {
31805 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
31806 }
31807
31808 #[allow(clippy::too_many_arguments)]
31811 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_tm_state_pad(
31813 &self,
31814 qkv_tm: &CudaSlice<f32>,
31815 conv_state: &mut CudaSlice<f32>,
31816 w: &CudaSlice<f32>,
31817 y: &mut CudaSlice<f32>,
31818 conv_dim: usize,
31819 t: usize,
31820 d_conv: usize,
31821 pad_len: Option<&CudaSlice<i32>>,
31822 ) -> Result<(), Box<dyn std::error::Error>> {
31823 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
31824 let ring_old = if t < d_conv - 1 {
31828 Some(self.clone_dtod(conv_state)?)
31829 } else {
31830 None
31831 };
31832 {
31833 let f = self.func("ssm_conv1d_tm_state_f32");
31834 let cfg = LaunchConfig {
31835 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31836 block_dim: (256, 1, 1),
31837 shared_mem_bytes: 0,
31838 };
31839 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31840 let __s_b = self.gpu.stream();
31841 let mut b = __s_b.launch_builder(&f);
31842 b.arg(qkv_tm)
31843 .arg(&*conv_state)
31844 .arg(w)
31845 .arg(y)
31846 .arg(&cd)
31847 .arg(&ti)
31848 .arg(&dc);
31849 unsafe {
31850 b.launch(cfg)?;
31851 }
31852 }
31853 match (ring_old, pad_len) {
31854 (None, Some(len_d)) => {
31855 let f = self.func("ssm_conv_ring_update_dev_f32");
31856 let n = conv_dim * (d_conv - 1);
31857 let cfg = LaunchConfig::for_num_elems(n as u32);
31858 let (cd, dc) = (conv_dim as i32, d_conv as i32);
31859 let __s_b = self.gpu.stream();
31860 let mut b = __s_b.launch_builder(&f);
31861 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
31862 unsafe {
31863 b.launch(cfg)?;
31864 }
31865 }
31866 (None, None) => {
31867 let f = self.func("ssm_conv_ring_update_f32");
31868 let n = conv_dim * (d_conv - 1);
31869 let cfg = LaunchConfig::for_num_elems(n as u32);
31870 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31871 let __s_b = self.gpu.stream();
31872 let mut b = __s_b.launch_builder(&f);
31873 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
31874 unsafe {
31875 b.launch(cfg)?;
31876 }
31877 }
31878 (Some(old), _) => {
31879 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
31880 }
31881 }
31882 Ok(())
31883 }
31884
31885 #[allow(clippy::too_many_arguments)]
31887 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_tm_state_pad_v(
31890 &self,
31891 qkv_tm: &cudarc::driver::CudaView<f32>,
31892 conv_state: &mut CudaSlice<f32>,
31893 w: &CudaSlice<f32>,
31894 y: &mut CudaSlice<f32>,
31895 conv_dim: usize,
31896 t: usize,
31897 d_conv: usize,
31898 pad_len: Option<&CudaSlice<i32>>,
31899 ) -> Result<(), Box<dyn std::error::Error>> {
31900 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
31901 let ring_old = if t < d_conv - 1 {
31905 Some(self.clone_dtod(conv_state)?)
31906 } else {
31907 None
31908 };
31909 {
31910 let f = self.func("ssm_conv1d_tm_state_f32");
31911 let cfg = LaunchConfig {
31912 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
31913 block_dim: (256, 1, 1),
31914 shared_mem_bytes: 0,
31915 };
31916 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31917 let __s_b = self.gpu.stream();
31918 let mut b = __s_b.launch_builder(&f);
31919 b.arg(qkv_tm)
31920 .arg(&*conv_state)
31921 .arg(w)
31922 .arg(y)
31923 .arg(&cd)
31924 .arg(&ti)
31925 .arg(&dc);
31926 unsafe {
31927 b.launch(cfg)?;
31928 }
31929 }
31930 match (ring_old, pad_len) {
31931 (None, Some(len_d)) => {
31932 let f = self.func("ssm_conv_ring_update_dev_f32");
31933 let n = conv_dim * (d_conv - 1);
31934 let cfg = LaunchConfig::for_num_elems(n as u32);
31935 let (cd, dc) = (conv_dim as i32, d_conv as i32);
31936 let __s_b = self.gpu.stream();
31937 let mut b = __s_b.launch_builder(&f);
31938 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
31939 unsafe {
31940 b.launch(cfg)?;
31941 }
31942 }
31943 (None, None) => {
31944 let f = self.func("ssm_conv_ring_update_f32");
31945 let n = conv_dim * (d_conv - 1);
31946 let cfg = LaunchConfig::for_num_elems(n as u32);
31947 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
31948 let __s_b = self.gpu.stream();
31949 let mut b = __s_b.launch_builder(&f);
31950 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
31951 unsafe {
31952 b.launch(cfg)?;
31953 }
31954 }
31955 (Some(_), _) => unreachable!(
31956 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
31957 ),
31958 }
31959 Ok(())
31960 }
31961
31962 pub fn ssm_conv_ring_rebuild(
31967 &self,
31968 qkv_tm: &CudaSlice<f32>,
31969 ring_old: &CudaSlice<f32>,
31970 conv_state: &mut CudaSlice<f32>,
31971 conv_dim: usize,
31972 tc: usize,
31973 d_conv: usize,
31974 ) -> Result<(), Box<dyn std::error::Error>> {
31975 let f = self.func("ssm_conv_ring_rebuild_f32");
31976 let n = conv_dim * (d_conv - 1);
31977 let cfg = LaunchConfig::for_num_elems(n as u32);
31978 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
31979 let __s_b = self.gpu.stream();
31980 let mut b = __s_b.launch_builder(&f);
31981 b.arg(qkv_tm)
31982 .arg(ring_old)
31983 .arg(conv_state)
31984 .arg(&cd)
31985 .arg(&ti)
31986 .arg(&dc);
31987 unsafe {
31988 b.launch(cfg)?;
31989 }
31990 Ok(())
31991 }
31992
31993 #[allow(clippy::too_many_arguments)]
31998 pub fn gdn_prep_decode(
31999 &self,
32000 conv_out: &CudaSlice<f32>,
32001 beta_raw: &CudaSlice<f32>,
32002 alpha: &CudaSlice<f32>,
32003 dt_bias: &CudaSlice<f32>,
32004 a: &CudaSlice<f32>,
32005 q_l2: &mut CudaSlice<f32>,
32006 k_l2: &mut CudaSlice<f32>,
32007 v_g: &mut CudaSlice<f32>,
32008 beta: &mut CudaSlice<f32>,
32009 g_log: &mut CudaSlice<f32>,
32010 d_state: usize,
32011 num_v: usize,
32012 num_k: usize,
32013 key_dim: usize,
32014 eps: f32,
32015 ) -> Result<(), Box<dyn std::error::Error>> {
32016 let f = self.func("gdn_prep_decode_f32");
32017 let cfg = LaunchConfig {
32018 grid_dim: (num_v as u32, 1, 1),
32019 block_dim: (32, 4, 1),
32020 shared_mem_bytes: 0,
32021 };
32022 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
32023 let __s_b = self.gpu.stream();
32024 let mut b = __s_b.launch_builder(&f);
32025 b.arg(conv_out)
32026 .arg(beta_raw)
32027 .arg(alpha)
32028 .arg(dt_bias)
32029 .arg(a)
32030 .arg(q_l2)
32031 .arg(k_l2)
32032 .arg(v_g)
32033 .arg(beta)
32034 .arg(g_log)
32035 .arg(&ds)
32036 .arg(&nv)
32037 .arg(&nk)
32038 .arg(&kd)
32039 .arg(&eps);
32040 unsafe {
32041 b.launch(cfg)?;
32042 }
32043 Ok(())
32044 }
32045
32046 #[allow(clippy::too_many_arguments)]
32050 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_gdn(
32052 &self,
32053 qkv_tm: &CudaSlice<f32>,
32054 w: &CudaSlice<f32>,
32055 q_g: &mut CudaSlice<f32>,
32056 k_g: &mut CudaSlice<f32>,
32057 v_g: &mut CudaSlice<f32>,
32058 conv_dim: usize,
32059 t: usize,
32060 d_conv: usize,
32061 d_state: usize,
32062 num_v: usize,
32063 num_k: usize,
32064 key_dim: usize,
32065 ) -> Result<(), Box<dyn std::error::Error>> {
32066 let f = self.func("ssm_conv1d_gdn_f32");
32067 let cfg = LaunchConfig {
32068 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
32069 block_dim: (256, 1, 1),
32070 shared_mem_bytes: 0,
32071 };
32072 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
32073 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
32074 let __s_b = self.gpu.stream();
32075 let mut b = __s_b.launch_builder(&f);
32076 b.arg(qkv_tm)
32077 .arg(w)
32078 .arg(q_g)
32079 .arg(k_g)
32080 .arg(v_g)
32081 .arg(&cd)
32082 .arg(&ti)
32083 .arg(&dc)
32084 .arg(&ds)
32085 .arg(&nv)
32086 .arg(&nk)
32087 .arg(&kd);
32088 unsafe {
32089 b.launch(cfg)?;
32090 }
32091 Ok(())
32092 }
32093
32094 #[allow(clippy::too_many_arguments)]
32095 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d(
32098 &self,
32099 x: &CudaSlice<f32>,
32100 w: &CudaSlice<f32>,
32101 y: &mut CudaSlice<f32>,
32102 conv_dim: usize,
32103 t: usize,
32104 d_conv: usize,
32105 silu: bool,
32106 ) -> Result<(), Box<dyn std::error::Error>> {
32107 let f = self.func("ssm_conv1d_silu_f32");
32108 let cfg = LaunchConfig {
32109 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
32110 block_dim: (256, 1, 1),
32111 shared_mem_bytes: 0,
32112 };
32113 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
32114 let __s_b = self.gpu.stream();
32115 let mut b = __s_b.launch_builder(&f);
32116 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
32117 unsafe {
32118 b.launch(cfg)?;
32119 }
32120 Ok(())
32121 }
32122
32123 #[allow(clippy::too_many_arguments)] pub fn gdn_scan_s128(
32127 &self,
32128 q: &CudaSlice<f32>,
32129 k: &CudaSlice<f32>,
32130 v: &CudaSlice<f32>,
32131 g: &CudaSlice<f32>,
32132 beta: &CudaSlice<f32>,
32133 state_in: &CudaSlice<f32>,
32134 state_out: &mut CudaSlice<f32>,
32135 o: &mut CudaSlice<f32>,
32136 n_head: usize,
32137 t: usize,
32138 scale: f32,
32139 ) -> Result<(), Box<dyn std::error::Error>> {
32140 let f = self.func("gdn_scan_s128");
32141 const S_V: u32 = 128;
32142 const WARP: u32 = 32;
32143 const COLS_PER_BLOCK: u32 = 4;
32144 let cfg = LaunchConfig {
32145 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
32146 block_dim: (WARP, COLS_PER_BLOCK, 1),
32147 shared_mem_bytes: 0,
32148 };
32149 let (h, ti) = (n_head as i32, t as i32);
32150 let __s_b = self.gpu.stream();
32151 let mut b = __s_b.launch_builder(&f);
32152 b.arg(q)
32153 .arg(k)
32154 .arg(v)
32155 .arg(g)
32156 .arg(beta)
32157 .arg(state_in)
32158 .arg(state_out)
32159 .arg(o)
32160 .arg(&h)
32161 .arg(&ti)
32162 .arg(&scale);
32163 unsafe {
32164 b.launch(cfg)?;
32165 }
32166 Ok(())
32167 }
32168
32169 #[allow(clippy::too_many_arguments)]
32174 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_fused_decode_b(
32176 &self,
32177 qkv_cols: &CudaSlice<f32>,
32178 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
32179 w: &CudaSlice<f32>,
32180 conv_outs: &mut CudaSlice<f32>,
32181 conv_dim: usize,
32182 d_conv: usize,
32183 b_n: usize,
32184 ) -> Result<(), Box<dyn std::error::Error>> {
32185 let f = self.func("ssm_conv1d_fused_decode_b_f32");
32186 let cfg = LaunchConfig {
32187 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
32188 block_dim: (256, 1, 1),
32189 shared_mem_bytes: 0,
32190 };
32191 let (cd, dc) = (conv_dim as i32, d_conv as i32);
32192 let __s_b = self.gpu.stream();
32193 let mut b = __s_b.launch_builder(&f);
32194 b.arg(qkv_cols)
32195 .arg(conv_state_ptrs)
32196 .arg(w)
32197 .arg(conv_outs)
32198 .arg(&cd)
32199 .arg(&dc);
32200 unsafe {
32201 b.launch(cfg)?;
32202 }
32203 Ok(())
32204 }
32205
32206 #[allow(clippy::too_many_arguments)]
32207 pub fn gdn_prep_decode_b(
32208 &self,
32209 conv_outs: &CudaSlice<f32>,
32210 beta_raws: &CudaSlice<f32>,
32211 alphas: &CudaSlice<f32>,
32212 dt_bias: &CudaSlice<f32>,
32213 a: &CudaSlice<f32>,
32214 q_l2: &mut CudaSlice<f32>,
32215 k_l2: &mut CudaSlice<f32>,
32216 v_g: &mut CudaSlice<f32>,
32217 beta: &mut CudaSlice<f32>,
32218 g_log: &mut CudaSlice<f32>,
32219 d_state: usize,
32220 num_v: usize,
32221 num_k: usize,
32222 key_dim: usize,
32223 eps: f32,
32224 conv_dim: usize,
32225 b_n: usize,
32226 ) -> Result<(), Box<dyn std::error::Error>> {
32227 let f = self.func("gdn_prep_decode_b_f32");
32228 let cfg = LaunchConfig {
32229 grid_dim: (num_v as u32, 1, b_n as u32),
32230 block_dim: (32, 4, 1),
32231 shared_mem_bytes: 0,
32232 };
32233 let (ds, nv, nk, kd, cd) = (
32234 d_state as i32,
32235 num_v as i32,
32236 num_k as i32,
32237 key_dim as i32,
32238 conv_dim as i32,
32239 );
32240 let __s_b = self.gpu.stream();
32241 let mut b = __s_b.launch_builder(&f);
32242 b.arg(conv_outs)
32243 .arg(beta_raws)
32244 .arg(alphas)
32245 .arg(dt_bias)
32246 .arg(a)
32247 .arg(q_l2)
32248 .arg(k_l2)
32249 .arg(v_g)
32250 .arg(beta)
32251 .arg(g_log)
32252 .arg(&ds)
32253 .arg(&nv)
32254 .arg(&nk)
32255 .arg(&kd)
32256 .arg(&eps)
32257 .arg(&cd);
32258 unsafe {
32259 b.launch(cfg)?;
32260 }
32261 Ok(())
32262 }
32263
32264 #[allow(clippy::too_many_arguments)]
32265 pub fn gdn_scan_s128_batched(
32266 &self,
32267 q: &CudaSlice<f32>,
32268 k: &CudaSlice<f32>,
32269 v: &CudaSlice<f32>,
32270 g: &CudaSlice<f32>,
32271 beta: &CudaSlice<f32>,
32272 state_in_ptrs: &cudarc::driver::CudaView<u64>,
32273 state_out_ptrs: &cudarc::driver::CudaView<u64>,
32274 o: &mut CudaSlice<f32>,
32275 n_head: usize,
32276 b_n: usize,
32277 scale: f32,
32278 ) -> Result<(), Box<dyn std::error::Error>> {
32279 let f = self.func("gdn_scan_s128_b");
32280 const S_V: u32 = 128;
32281 const WARP: u32 = 32;
32282 const COLS_PER_BLOCK: u32 = 4;
32283 let cfg = LaunchConfig {
32284 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
32285 block_dim: (WARP, COLS_PER_BLOCK, 1),
32286 shared_mem_bytes: 0,
32287 };
32288 let h = n_head as i32;
32289 let __s_b = self.gpu.stream();
32290 let mut b = __s_b.launch_builder(&f);
32291 b.arg(q)
32292 .arg(k)
32293 .arg(v)
32294 .arg(g)
32295 .arg(beta)
32296 .arg(state_in_ptrs)
32297 .arg(state_out_ptrs)
32298 .arg(o)
32299 .arg(&h)
32300 .arg(&scale);
32301 unsafe {
32302 b.launch(cfg)?;
32303 }
32304 Ok(())
32305 }
32306
32307 #[allow(clippy::too_many_arguments)]
32313 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_fused_decode_b_view(
32315 &self,
32316 qkv_cols: &cudarc::driver::CudaView<f32>,
32317 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
32318 w: &CudaSlice<f32>,
32319 conv_outs: &mut CudaSlice<f32>,
32320 conv_dim: usize,
32321 d_conv: usize,
32322 b_n: usize,
32323 ) -> Result<(), Box<dyn std::error::Error>> {
32324 let f = self.func("ssm_conv1d_fused_decode_b_f32");
32325 let cfg = LaunchConfig {
32326 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
32327 block_dim: (256, 1, 1),
32328 shared_mem_bytes: 0,
32329 };
32330 let (cd, dc) = (conv_dim as i32, d_conv as i32);
32331 let __s_b = self.gpu.stream();
32332 let mut b = __s_b.launch_builder(&f);
32333 b.arg(qkv_cols)
32334 .arg(conv_state_ptrs)
32335 .arg(w)
32336 .arg(conv_outs)
32337 .arg(&cd)
32338 .arg(&dc);
32339 unsafe {
32340 b.launch(cfg)?;
32341 }
32342 Ok(())
32343 }
32344
32345 #[allow(clippy::too_many_arguments)]
32346 pub fn gdn_prep_decode_b_view(
32347 &self,
32348 conv_outs: &CudaSlice<f32>,
32349 beta_raws: &cudarc::driver::CudaView<f32>,
32350 alphas: &cudarc::driver::CudaView<f32>,
32351 dt_bias: &CudaSlice<f32>,
32352 a: &CudaSlice<f32>,
32353 q_l2: &mut CudaSlice<f32>,
32354 k_l2: &mut CudaSlice<f32>,
32355 v_g: &mut CudaSlice<f32>,
32356 beta: &mut CudaSlice<f32>,
32357 g_log: &mut CudaSlice<f32>,
32358 d_state: usize,
32359 num_v: usize,
32360 num_k: usize,
32361 key_dim: usize,
32362 eps: f32,
32363 conv_dim: usize,
32364 b_n: usize,
32365 ) -> Result<(), Box<dyn std::error::Error>> {
32366 let f = self.func("gdn_prep_decode_b_f32");
32367 let cfg = LaunchConfig {
32368 grid_dim: (num_v as u32, 1, b_n as u32),
32369 block_dim: (32, 4, 1),
32370 shared_mem_bytes: 0,
32371 };
32372 let (ds, nv, nk, kd, cd) = (
32373 d_state as i32,
32374 num_v as i32,
32375 num_k as i32,
32376 key_dim as i32,
32377 conv_dim as i32,
32378 );
32379 let __s_b = self.gpu.stream();
32380 let mut b = __s_b.launch_builder(&f);
32381 b.arg(conv_outs)
32382 .arg(beta_raws)
32383 .arg(alphas)
32384 .arg(dt_bias)
32385 .arg(a)
32386 .arg(q_l2)
32387 .arg(k_l2)
32388 .arg(v_g)
32389 .arg(beta)
32390 .arg(g_log)
32391 .arg(&ds)
32392 .arg(&nv)
32393 .arg(&nk)
32394 .arg(&kd)
32395 .arg(&eps)
32396 .arg(&cd);
32397 unsafe {
32398 b.launch(cfg)?;
32399 }
32400 Ok(())
32401 }
32402
32403 #[allow(clippy::too_many_arguments)]
32404 pub fn gdn_scan_s128_batched_view(
32405 &self,
32406 q: &CudaSlice<f32>,
32407 k: &CudaSlice<f32>,
32408 v: &CudaSlice<f32>,
32409 g: &CudaSlice<f32>,
32410 beta: &CudaSlice<f32>,
32411 state_in_ptrs: &cudarc::driver::CudaView<u64>,
32412 state_out_ptrs: &cudarc::driver::CudaView<u64>,
32413 o: &mut cudarc::driver::CudaViewMut<f32>,
32414 n_head: usize,
32415 b_n: usize,
32416 scale: f32,
32417 ) -> Result<(), Box<dyn std::error::Error>> {
32418 let f = self.func("gdn_scan_s128_b");
32419 const S_V: u32 = 128;
32420 const WARP: u32 = 32;
32421 const COLS_PER_BLOCK: u32 = 4;
32422 let cfg = LaunchConfig {
32423 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
32424 block_dim: (WARP, COLS_PER_BLOCK, 1),
32425 shared_mem_bytes: 0,
32426 };
32427 let h = n_head as i32;
32428 let __s_b = self.gpu.stream();
32429 let mut b = __s_b.launch_builder(&f);
32430 b.arg(q)
32431 .arg(k)
32432 .arg(v)
32433 .arg(g)
32434 .arg(beta)
32435 .arg(state_in_ptrs)
32436 .arg(state_out_ptrs)
32437 .arg(o)
32438 .arg(&h)
32439 .arg(&scale);
32440 unsafe {
32441 b.launch(cfg)?;
32442 }
32443 Ok(())
32444 }
32445
32446 pub fn gdn_chunked_enabled() -> bool {
32455 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
32456 *E.get_or_init(|| {
32457 std::env::var("MEMRA_GDN_CHUNKED")
32458 .map(|v| v != "0")
32459 .unwrap_or(true)
32460 })
32461 }
32462
32463 pub fn gdn_chunk_size() -> usize {
32468 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
32469 *C.get_or_init(|| {
32470 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
32471 .ok()
32472 .and_then(|v| v.parse().ok())
32473 .unwrap_or(32);
32474 c.clamp(32, 128) / 32 * 32
32475 })
32476 }
32477
32478 #[allow(clippy::too_many_arguments)]
32483 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
32486 #[allow(clippy::too_many_arguments)]
32487 pub fn gdn_chunk_k123(
32488 &self,
32489 q: &CudaSlice<f32>,
32490 k: &CudaSlice<f32>,
32491 v: &CudaSlice<f32>,
32492 g: &CudaSlice<f32>,
32493 beta: &CudaSlice<f32>,
32494 wb16: Option<&mut CudaSlice<u8>>,
32495 n_head: usize,
32496 t: usize,
32497 c: usize,
32498 hk: usize,
32499 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
32500 ) -> Result<
32501 (
32502 CudaSlice<f32>,
32503 CudaSlice<f32>,
32504 CudaSlice<f32>,
32505 CudaSlice<f32>,
32506 ),
32507 Box<dyn std::error::Error>,
32508 > {
32509 const D: usize = 128;
32510 let h = n_head;
32511 #[allow(clippy::manual_div_ceil)]
32512 let nc = (t + c - 1) / c;
32514 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
32515 let mut gcum = self.uninit(t * h)?;
32516 let mut a = self.uninit(nc * h * c * c)?;
32517 let mut p = self.uninit(nc * h * c * c)?;
32518 let mut u = self.uninit(nc * h * c * D)?;
32519 let mut w = self.uninit(nc * h * c * D)?;
32520 {
32521 let f = self.func("gdn_chunk_cumgate_f32");
32523 let cfg = LaunchConfig {
32524 grid_dim: (nc as u32, h as u32, 1),
32525 block_dim: (32, 1, 1),
32526 shared_mem_bytes: 0,
32527 };
32528 let __s_b = self.gpu.stream();
32529 let mut b = __s_b.launch_builder(&f);
32530 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
32531 unsafe {
32532 b.launch(cfg)?;
32533 }
32534 }
32535 if let Some((qb, kb, pb)) = k2w {
32536 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
32539 let f = self.func("gdn_k2_wgmma");
32540 let cfg = LaunchConfig {
32541 grid_dim: (nc as u32, h as u32, 1),
32542 block_dim: (128, 1, 1),
32543 shared_mem_bytes: 0,
32544 };
32545 let hki = hk as i32;
32546 let __s_b = self.gpu.stream();
32547 let mut b = __s_b.launch_builder(&f);
32548 b.arg(qb)
32549 .arg(kb)
32550 .arg(&gcum)
32551 .arg(beta)
32552 .arg(&mut a)
32553 .arg(&mut *pb)
32554 .arg(&hi)
32555 .arg(&ti)
32556 .arg(&ci)
32557 .arg(&hki);
32558 unsafe {
32559 b.launch(cfg)?;
32560 }
32561 } else if c <= 64 && !portable_mma_gated() {
32562 let f = self.func("gdn_chunk_attn_f32");
32564 f.set_attribute(
32565 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
32566 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
32567 )?;
32568 #[allow(clippy::manual_div_ceil)]
32569 let jt = ((c + 31) / 32) as u32;
32571 let cfg = LaunchConfig {
32572 grid_dim: (nc as u32, h as u32, jt),
32573 block_dim: (256, 1, 1),
32574 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
32575 };
32576 let hki = hk as i32;
32577 let __s_b = self.gpu.stream();
32578 let mut b = __s_b.launch_builder(&f);
32579 b.arg(q)
32580 .arg(k)
32581 .arg(&gcum)
32582 .arg(beta)
32583 .arg(&mut a)
32584 .arg(&mut p)
32585 .arg(&hi)
32586 .arg(&ti)
32587 .arg(&ci)
32588 .arg(&hki);
32589 unsafe {
32590 b.launch(cfg)?;
32591 }
32592 } else {
32593 assert!(
32595 hk == h,
32596 "generic K2 is broadcast-only (de-broadcast rides C==32)"
32597 );
32598 let f = self.func("gdn_chunk_attn_g_f32");
32599 let cfg = LaunchConfig {
32600 grid_dim: (nc as u32, h as u32, 1),
32601 block_dim: (32, 8, 1),
32602 shared_mem_bytes: 0,
32603 };
32604 let __s_b = self.gpu.stream();
32605 let mut b = __s_b.launch_builder(&f);
32606 b.arg(q)
32607 .arg(k)
32608 .arg(&gcum)
32609 .arg(beta)
32610 .arg(&mut a)
32611 .arg(&mut p)
32612 .arg(&hi)
32613 .arg(&ti)
32614 .arg(&ci);
32615 unsafe {
32616 b.launch(cfg)?;
32617 }
32618 }
32619 {
32620 let cfg = LaunchConfig {
32622 grid_dim: (nc as u32, h as u32, 1),
32623 block_dim: (256, 1, 1),
32624 shared_mem_bytes: 0,
32625 };
32626 match c {
32627 32 | 64 => {
32628 let f = self.func(if c == 32 {
32629 "gdn_chunk_solve32_f32"
32630 } else {
32631 "gdn_chunk_solve64_f32"
32632 });
32633 let wb: u64 = match wb16 {
32635 Some(d) => self.addr_u8(d),
32636 None => 0,
32637 };
32638 let hki = hk as i32;
32639 let __s_b = self.gpu.stream();
32640 let mut b = __s_b.launch_builder(&f);
32641 b.arg(v)
32642 .arg(k)
32643 .arg(&a)
32644 .arg(&gcum)
32645 .arg(&mut u)
32646 .arg(&mut w)
32647 .arg(&wb)
32648 .arg(&hi)
32649 .arg(&ti)
32650 .arg(&hki);
32651 unsafe {
32652 b.launch(cfg)?;
32653 }
32654 }
32655 _ => {
32656 assert!(hk == h, "generic K3 is broadcast-only");
32657 let f = self.func("gdn_chunk_solve_f32");
32658 let __s_b = self.gpu.stream();
32659 let mut b = __s_b.launch_builder(&f);
32660 b.arg(v)
32661 .arg(k)
32662 .arg(&a)
32663 .arg(&gcum)
32664 .arg(&mut u)
32665 .arg(&mut w)
32666 .arg(&hi)
32667 .arg(&ti)
32668 .arg(&ci);
32669 unsafe {
32670 b.launch(cfg)?;
32671 }
32672 }
32673 }
32674 }
32675 Ok((gcum, p, u, w))
32676 }
32677
32678 pub fn gdn_db_on() -> bool {
32682 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
32683 }
32684
32685 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
32694 !portable_mma_gated()
32695 && c == 32
32696 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
32697 Ok("1") => true,
32698 Ok("0") => false,
32699 _ => gdn_mma_default_on(),
32700 }
32701 }
32702
32703 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
32710 cfg!(memra_hopper_mma)
32711 && self.gdn_mma_enabled(c)
32712 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
32713 }
32714
32715 #[allow(clippy::too_many_arguments)]
32720 #[allow(clippy::manual_div_ceil)] pub fn ssm_conv1d_gdn_state_pad(
32722 &self,
32723 qkv_tm: &cudarc::driver::CudaView<f32>,
32724 conv_state: &mut CudaSlice<f32>,
32725 w: &CudaSlice<f32>,
32726 q_g: &mut CudaSlice<f32>,
32727 k_g: &mut CudaSlice<f32>,
32728 v_g: &mut CudaSlice<f32>,
32729 conv_dim: usize,
32730 t: usize,
32731 d_conv: usize,
32732 d_state: usize,
32733 num_v: usize,
32734 num_k: usize,
32735 key_dim: usize,
32736 hk: usize,
32737 pad_len: Option<&CudaSlice<i32>>,
32738 ) -> Result<(), Box<dyn std::error::Error>> {
32739 assert!(
32740 t >= d_conv - 1,
32741 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
32742 );
32743 {
32744 let f = self.func("ssm_conv1d_gdn_state_f32");
32745 let cfg = LaunchConfig {
32746 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
32747 block_dim: (256, 1, 1),
32748 shared_mem_bytes: 0,
32749 };
32750 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
32751 let (ds, nv, nk, kd, hki) = (
32752 d_state as i32,
32753 num_v as i32,
32754 num_k as i32,
32755 key_dim as i32,
32756 hk as i32,
32757 );
32758 let __s_b = self.gpu.stream();
32759 let mut b = __s_b.launch_builder(&f);
32760 b.arg(qkv_tm)
32761 .arg(&*conv_state)
32762 .arg(w)
32763 .arg(q_g)
32764 .arg(k_g)
32765 .arg(v_g)
32766 .arg(&cd)
32767 .arg(&ti)
32768 .arg(&dc)
32769 .arg(&ds)
32770 .arg(&nv)
32771 .arg(&nk)
32772 .arg(&kd)
32773 .arg(&hki);
32774 unsafe {
32775 b.launch(cfg)?;
32776 }
32777 }
32778 match pad_len {
32779 Some(len_d) => {
32780 let f = self.func("ssm_conv_ring_update_dev_f32");
32781 let n = conv_dim * (d_conv - 1);
32782 let cfg = LaunchConfig::for_num_elems(n as u32);
32783 let (cd, dc) = (conv_dim as i32, d_conv as i32);
32784 let __s_b = self.gpu.stream();
32785 let mut b = __s_b.launch_builder(&f);
32786 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
32787 unsafe {
32788 b.launch(cfg)?;
32789 }
32790 }
32791 None => {
32792 let f = self.func("ssm_conv_ring_update_f32");
32793 let n = conv_dim * (d_conv - 1);
32794 let cfg = LaunchConfig::for_num_elems(n as u32);
32795 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
32796 let __s_b = self.gpu.stream();
32797 let mut b = __s_b.launch_builder(&f);
32798 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
32799 unsafe {
32800 b.launch(cfg)?;
32801 }
32802 }
32803 }
32804 Ok(())
32805 }
32806
32807 pub fn gdn_chunk_alloc(
32811 &self,
32812 n_head: usize,
32813 t: usize,
32814 c: usize,
32815 hk: usize,
32816 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
32817 const D: usize = 128;
32818 assert!(
32819 c == 32,
32820 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
32821 );
32822 let h = n_head;
32823 #[allow(clippy::manual_div_ceil)]
32824 let nc = (t + c - 1) / c;
32826 Ok(GdnChunkBufs {
32827 gcum: self.uninit(t * h)?,
32828 a: self.uninit(nc * h * c * c)?,
32829 p: self.uninit(nc * h * c * c)?,
32830 u: self.uninit(nc * h * c * D)?,
32831 w: self.uninit(nc * h * c * D)?,
32832 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
32833 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
32834 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
32835 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
32836 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
32837 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
32838 o: self.uninit(D * h * t)?,
32839 t,
32840 nc,
32841 })
32842 }
32843
32844 pub fn f32_to_bf16_v(
32846 &self,
32847 x: &cudarc::driver::CudaView<f32>,
32848 dst: &mut CudaSlice<u8>,
32849 n: usize,
32850 ) -> Result<(), Box<dyn std::error::Error>> {
32851 let f = self.func("f32_to_bf16_bulk");
32852 let ni = n as i64;
32853 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
32854 let __s_b = self.gpu.stream();
32855 let mut b = __s_b.launch_builder(&f);
32856 b.arg(x).arg(dst).arg(&ni);
32857 unsafe {
32858 b.launch(cfg)?;
32859 }
32860 Ok(())
32861 }
32862
32863 pub fn f32_to_bf16_into(
32865 &self,
32866 x: &CudaSlice<f32>,
32867 dst: &mut CudaSlice<u8>,
32868 n: usize,
32869 ) -> Result<(), Box<dyn std::error::Error>> {
32870 let f = self.func("f32_to_bf16_bulk");
32871 let ni = n as i64;
32872 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
32873 let __s_b = self.gpu.stream();
32874 let mut b = __s_b.launch_builder(&f);
32875 b.arg(x).arg(dst).arg(&ni);
32876 unsafe {
32877 b.launch(cfg)?;
32878 }
32879 Ok(())
32880 }
32881
32882 pub fn gdn_chunk_k123_vl8(
32885 &self,
32886 seqs: &[GdnSeqVl],
32887 n_head: usize,
32888 hk: usize,
32889 wq: Option<&GdnWVl8>,
32890 ) -> Result<(), Box<dyn std::error::Error>> {
32891 let b = seqs.len();
32892 assert!((1..=8).contains(&b), "gdn_chunk_k123_vl8: 1..=8 sequences");
32893 let mut packed = [GdnSeqVl::default(); 8];
32894 packed[..b].copy_from_slice(seqs);
32895 let v = GdnVl8(packed);
32896 let (hi, ci) = (n_head as i32, 32i32);
32897 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
32898 {
32899 let f = self.func("gdn_chunk_cumgate_vl");
32900 let cfg = LaunchConfig {
32901 grid_dim: (max_nc, n_head as u32, b as u32),
32902 block_dim: (32, 1, 1),
32903 shared_mem_bytes: 0,
32904 };
32905 let __s_lb = self.gpu.stream();
32906 let mut lb = __s_lb.launch_builder(&f);
32907 lb.arg(&v).arg(&hi).arg(&ci);
32908 unsafe {
32909 lb.launch(cfg)?;
32910 }
32911 }
32912 let hki = hk as i32;
32913 if let Some(w) = wq {
32914 let f = self.func("gdn_k2_wgmma_vl");
32916 let cfg = LaunchConfig {
32917 grid_dim: (max_nc, n_head as u32, b as u32),
32918 block_dim: (128, 1, 1),
32919 shared_mem_bytes: 0,
32920 };
32921 let __s_lb = self.gpu.stream();
32922 let mut lb = __s_lb.launch_builder(&f);
32923 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
32924 unsafe {
32925 lb.launch(cfg)?;
32926 }
32927 } else {
32928 let f = self.func("gdn_chunk_attn_vl");
32929 f.set_attribute(
32930 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
32931 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
32932 )?;
32933 let cfg = LaunchConfig {
32934 grid_dim: (max_nc, n_head as u32, b as u32),
32935 block_dim: (256, 1, 1),
32936 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
32937 };
32938 let __s_lb = self.gpu.stream();
32939 let mut lb = __s_lb.launch_builder(&f);
32940 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
32941 unsafe {
32942 lb.launch(cfg)?;
32943 }
32944 }
32945 {
32946 let f = self.func("gdn_chunk_solve32_vl");
32947 let cfg = LaunchConfig {
32948 grid_dim: (max_nc, n_head as u32, b as u32),
32949 block_dim: (256, 1, 1),
32950 shared_mem_bytes: 0,
32951 };
32952 let __s_lb = self.gpu.stream();
32953 let mut lb = __s_lb.launch_builder(&f);
32954 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
32955 unsafe {
32956 lb.launch(cfg)?;
32957 }
32958 }
32959 Ok(())
32960 }
32961
32962 #[allow(clippy::too_many_arguments)]
32966 pub fn gdn_prep_vl8(
32967 &self,
32968 seqs: &[GdnPrepVl],
32969 conv_w: &CudaSlice<f32>,
32970 dt_bias: &CudaSlice<f32>,
32971 a: &CudaSlice<f32>,
32972 conv_dim: usize,
32973 d_conv: usize,
32974 d_state: usize,
32975 num_v: usize,
32976 num_k: usize,
32977 key_dim: usize,
32978 hk: usize,
32979 eps: f32,
32980 ) -> Result<(), Box<dyn std::error::Error>> {
32981 let b = seqs.len();
32982 assert!((1..=8).contains(&b));
32983 let mut packed = [GdnPrepVl::default(); 8];
32984 packed[..b].copy_from_slice(seqs);
32985 let v = GdnPrepVl8(packed);
32986 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
32987 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
32988 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
32989 assert!(
32990 conv_fuse || hk == num_v,
32991 "de-broadcast requires the fused conv"
32992 );
32993 if conv_fuse {
32994 let f = self.func("ssm_conv1d_gdn_state_vl");
32995 let cfg = LaunchConfig {
32996 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
32997 block_dim: (256, 1, 1),
32998 shared_mem_bytes: 0,
32999 };
33000 let (dsi, nvi, nki, kdi, hki) = (
33001 d_state as i32,
33002 num_v as i32,
33003 num_k as i32,
33004 key_dim as i32,
33005 hk as i32,
33006 );
33007 let __s_lb = self.gpu.stream();
33008 let mut lb = __s_lb.launch_builder(&f);
33009 lb.arg(&v)
33010 .arg(conv_w)
33011 .arg(&cdi)
33012 .arg(&dci)
33013 .arg(&dsi)
33014 .arg(&nvi)
33015 .arg(&nki)
33016 .arg(&kdi)
33017 .arg(&hki);
33018 unsafe {
33019 lb.launch(cfg)?;
33020 }
33021 } else {
33022 let f = self.func("ssm_conv1d_tm_state_vl");
33023 let cfg = LaunchConfig {
33024 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
33025 block_dim: (256, 1, 1),
33026 shared_mem_bytes: 0,
33027 };
33028 let __s_lb = self.gpu.stream();
33029 let mut lb = __s_lb.launch_builder(&f);
33030 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
33031 unsafe {
33032 lb.launch(cfg)?;
33033 }
33034 }
33035 {
33036 let f = self.func("ssm_conv_ring_update_vl");
33037 let n = (conv_dim * (d_conv - 1)) as u32;
33038 let cfg = LaunchConfig {
33039 grid_dim: (n.div_ceil(256), 1, b as u32),
33040 block_dim: (256, 1, 1),
33041 shared_mem_bytes: 0,
33042 };
33043 let __s_lb = self.gpu.stream();
33044 let mut lb = __s_lb.launch_builder(&f);
33045 lb.arg(&v).arg(&cdi).arg(&dci);
33046 unsafe {
33047 lb.launch(cfg)?;
33048 }
33049 }
33050 if !conv_fuse {
33051 let f = self.func("qkv_to_gdn_repack_vl");
33052 let n = max_t * (num_v * d_state) as u32;
33053 let cfg = LaunchConfig {
33054 grid_dim: (n.div_ceil(256), 1, b as u32),
33055 block_dim: (256, 1, 1),
33056 shared_mem_bytes: 0,
33057 };
33058 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
33059 let __s_lb = self.gpu.stream();
33060 let mut lb = __s_lb.launch_builder(&f);
33061 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
33062 unsafe {
33063 lb.launch(cfg)?;
33064 }
33065 }
33066 if Self::l2_v2_on(d_state) {
33067 let f = self.func("gdn_l2_v2_vl");
33068 let cfg = LaunchConfig {
33069 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
33070 block_dim: (256, 1, 1),
33071 shared_mem_bytes: 0,
33072 };
33073 let (dsi, nvi) = (d_state as i32, hk as i32);
33074 let __s_lb = self.gpu.stream();
33075 let mut lb = __s_lb.launch_builder(&f);
33076 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
33077 unsafe {
33078 lb.launch(cfg)?;
33079 }
33080 } else {
33081 let f = self.func("gdn_l2_vl");
33082 let cfg = LaunchConfig {
33083 grid_dim: (max_t * hk as u32, 2, b as u32),
33084 block_dim: (256, 1, 1),
33085 shared_mem_bytes: 0,
33086 };
33087 let (dsi, nvi) = (d_state as i32, hk as i32);
33088 let __s_lb = self.gpu.stream();
33089 let mut lb = __s_lb.launch_builder(&f);
33090 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
33091 unsafe {
33092 lb.launch(cfg)?;
33093 }
33094 }
33095 {
33096 let f = self.func("gdn_gate_prep_vl");
33097 let n = max_t * num_v as u32;
33098 let cfg = LaunchConfig {
33099 grid_dim: (n.div_ceil(256), 1, b as u32),
33100 block_dim: (256, 1, 1),
33101 shared_mem_bytes: 0,
33102 };
33103 let nvi = num_v as i32;
33104 let __s_lb = self.gpu.stream();
33105 let mut lb = __s_lb.launch_builder(&f);
33106 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
33107 unsafe {
33108 lb.launch(cfg)?;
33109 }
33110 }
33111 Ok(())
33112 }
33113
33114 pub fn gdn_mirror_vl8(
33116 &self,
33117 seqs: &[GdnSeqVl],
33118 n_head: usize,
33119 which: i32,
33120 hk: usize,
33121 ) -> Result<(), Box<dyn std::error::Error>> {
33122 let b = seqs.len();
33123 assert!((1..=8).contains(&b));
33124 let mut packed = [GdnSeqVl::default(); 8];
33125 packed[..b].copy_from_slice(seqs);
33126 let v = GdnVl8(packed);
33127 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
33128 let max_n = seqs
33129 .iter()
33130 .map(|s| {
33131 if which == 0 {
33132 s.t as i64 * ept as i64
33133 } else {
33134 s.nc as i64 * ept as i64 * 32
33135 }
33136 })
33137 .max()
33138 .unwrap();
33139 let f = self.func("gdn_mirror_vl");
33140 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
33141 let cfg = LaunchConfig {
33142 grid_dim: (blocks, 1, b as u32),
33143 block_dim: (256, 1, 1),
33144 shared_mem_bytes: 0,
33145 };
33146 let __s_lb = self.gpu.stream();
33147 let mut lb = __s_lb.launch_builder(&f);
33148 lb.arg(&v).arg(&ept).arg(&which);
33149 unsafe {
33150 lb.launch(cfg)?;
33151 }
33152 Ok(())
33153 }
33154
33155 pub fn gdn_tail_vl8(
33157 &self,
33158 seqs: &[GdnPrepVl],
33159 norm_w: &CudaSlice<f32>,
33160 d_state: usize,
33161 num_v: usize,
33162 eps: f32,
33163 ) -> Result<(), Box<dyn std::error::Error>> {
33164 let b = seqs.len();
33165 assert!((1..=8).contains(&b));
33166 let mut packed = [GdnPrepVl::default(); 8];
33167 packed[..b].copy_from_slice(seqs);
33168 let v = GdnPrepVl8(packed);
33169 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
33170 let f = self.func("gated_rmsnorm_f16out_vl");
33171 let cfg = LaunchConfig {
33173 grid_dim: (max_t * num_v as u32, 1, b as u32),
33174 block_dim: (128, 1, 1),
33175 shared_mem_bytes: 0,
33176 };
33177 let (dsi, nvi) = (d_state as i32, num_v as i32);
33178 let __s_lb = self.gpu.stream();
33179 let mut lb = __s_lb.launch_builder(&f);
33180 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
33181 unsafe {
33182 lb.launch(cfg)?;
33183 }
33184 Ok(())
33185 }
33186
33187 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
33190 use cudarc::driver::DevicePtr;
33191 let s = self.gpu.stream();
33192 let (p, _g) = x.device_ptr(&s);
33193 p
33194 }
33195 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
33196 use cudarc::driver::DevicePtrMut;
33197 let s = self.gpu.stream();
33198 let (p, _g) = x.device_ptr_mut(&s);
33199 p
33200 }
33201 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
33202 use cudarc::driver::DevicePtr;
33203 let s = self.gpu.stream();
33204 let (p, _g) = x.device_ptr(&s);
33205 p
33206 }
33207 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
33208 use cudarc::driver::DevicePtr;
33209 let s = self.gpu.stream();
33210 let (p, _g) = x.device_ptr(&s);
33211 p
33212 }
33213
33214 pub fn gdn_chunk_vl8(
33218 &self,
33219 seqs: &[GdnSeqVl],
33220 n_head: usize,
33221 scale: f32,
33222 hk: usize,
33223 wq: Option<&GdnWVl8>,
33224 ) -> Result<(), Box<dyn std::error::Error>> {
33225 const NSPLIT: u32 = 4;
33226 let b = seqs.len();
33227 assert!((1..=8).contains(&b), "gdn_chunk_vl8: 1..=8 sequences");
33228 let mut packed = [GdnSeqVl::default(); 8];
33229 packed[..b].copy_from_slice(seqs);
33230 let v = GdnVl8(packed);
33231 let (hi, ci) = (n_head as i32, 32i32);
33232 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
33233 let hki = hk as i32;
33234 if let Some(w) = wq {
33235 let f = self.func("gdn_k45_wgmma_vl");
33237 let cfg = LaunchConfig {
33238 grid_dim: (n_head as u32, NSPLIT, b as u32),
33239 block_dim: (256, 1, 1),
33240 shared_mem_bytes: 0,
33241 };
33242 let __s_lb = self.gpu.stream();
33243 let mut lb = __s_lb.launch_builder(&f);
33244 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
33245 unsafe {
33246 lb.launch(cfg)?;
33247 }
33248 let _ = max_nc;
33249 return Ok(());
33250 }
33251 {
33252 let f = self.func("gdn_chunk_state_mma_vl");
33253 let cfg = LaunchConfig {
33254 grid_dim: (n_head as u32, NSPLIT, b as u32),
33255 block_dim: (256, 1, 1),
33256 shared_mem_bytes: 0,
33257 };
33258 let __s_lb = self.gpu.stream();
33259 let mut lb = __s_lb.launch_builder(&f);
33260 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
33261 unsafe {
33262 lb.launch(cfg)?;
33263 }
33264 }
33265 {
33266 let f = self.func("gdn_chunk_output_mma_vl");
33267 let cfg = LaunchConfig {
33268 grid_dim: (max_nc, n_head as u32, b as u32),
33269 block_dim: (256, 1, 1),
33270 shared_mem_bytes: 0,
33271 };
33272 let __s_lb = self.gpu.stream();
33273 let mut lb = __s_lb.launch_builder(&f);
33274 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
33275 unsafe {
33276 lb.launch(cfg)?;
33277 }
33278 }
33279 Ok(())
33280 }
33281 #[allow(clippy::too_many_arguments)] pub fn gdn_scan_chunked(
33283 &self,
33284 q: &CudaSlice<f32>,
33285 k: &CudaSlice<f32>,
33286 v: &CudaSlice<f32>,
33287 g: &CudaSlice<f32>,
33288 beta: &CudaSlice<f32>,
33289 kb16_pre: Option<&CudaSlice<u8>>,
33290 qb16_pre: Option<&CudaSlice<u8>>,
33291 state_in: &CudaSlice<f32>,
33292 state_out: &mut CudaSlice<f32>,
33293 o: &mut CudaSlice<f32>,
33294 n_head: usize,
33295 t: usize,
33296 scale: f32,
33297 c: usize,
33298 hk: usize,
33299 ) -> Result<(), Box<dyn std::error::Error>> {
33300 const D: usize = 128;
33301 const NSPLIT: u32 = 4;
33302 assert!(
33303 (1..=128).contains(&c),
33304 "gdn_scan_chunked: C must be in 1..=128"
33305 );
33306 let h = n_head;
33307 #[allow(clippy::manual_div_ceil)]
33308 let nc = (t + c - 1) / c;
33310 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
33311 let gdn_mma_pre = !portable_mma_gated()
33316 && c == 32
33317 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
33318 Ok("1") => true,
33319 Ok("0") => false,
33320 _ => gdn_mma_default_on(),
33321 };
33322 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
33323 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
33324 } else {
33325 None
33326 };
33327 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
33332 && gdn_mma_pre
33333 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
33334 let nk = t * hk * D;
33335 let mut kb16_local: Option<CudaSlice<u8>> = None;
33336 if gdn_mma_pre && kb16_pre.is_none() {
33337 let mut kb = self.alloc_u8_uninit(nk * 2)?;
33338 let f = self.func("f32_to_bf16_bulk");
33339 let n2 = nk as i64;
33340 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
33341 let __s_b = self.gpu.stream();
33342 let mut b = __s_b.launch_builder(&f);
33343 b.arg(k).arg(&mut kb).arg(&n2);
33344 unsafe {
33345 b.launch(cfg2)?;
33346 }
33347 kb16_local = Some(kb);
33348 }
33349 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
33350 if let Some(kb) = kb16_pre {
33351 assert!(kb.len() >= nk * 2, "kb16_pre too small");
33352 }
33353 let mut qb16: Option<CudaSlice<u8>> = None;
33354 let mut pb16: Option<CudaSlice<u8>> = None;
33355 if gdn_wgmma_pre {
33356 if qb16_pre.is_none() {
33359 let mut qb = self.alloc_u8_uninit(nk * 2)?;
33360 let f = self.func("f32_to_bf16_bulk");
33361 let n2 = nk as i64;
33362 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
33363 let __s_b = self.gpu.stream();
33364 let mut b = __s_b.launch_builder(&f);
33365 b.arg(q).arg(&mut qb).arg(&n2);
33366 unsafe {
33367 b.launch(cfg2)?;
33368 }
33369 qb16 = Some(qb);
33370 } else if let Some(qb) = qb16_pre {
33371 assert!(qb.len() >= nk * 2, "qb16_pre too small");
33372 }
33373 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
33374 }
33375 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
33376 let k2w = if gdn_wgmma_pre {
33377 Some((
33378 *qb16_ref0.as_ref().unwrap(),
33379 *kb16_ref0.as_ref().unwrap(),
33380 pb16.as_mut().unwrap(),
33381 ))
33382 } else {
33383 None
33384 };
33385 let (gcum, p, u, w) =
33386 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
33387 let _ = &w;
33388 let mut y = self.uninit(nc * h * c * D)?;
33389 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
33403 && c == 32
33404 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
33405 Ok("1") => true,
33406 Ok("0") => false,
33407 _ => gdn_mma_default_on(),
33408 };
33409 if gdn_mma {
33410 let wb16 = wb16_pre
33411 .take()
33412 .expect("mma path pre-allocates wb16 (K3 store fold)");
33413 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
33414 if gdn_wgmma_pre {
33426 let qb16 = qb16_ref0.unwrap();
33428 let pb16 = pb16.as_ref().unwrap();
33429 {
33430 let f = self.func("gdn_k45_wgmma");
33431 let cfg = LaunchConfig {
33432 grid_dim: (h as u32, 4, 1),
33433 block_dim: (256, 1, 1),
33434 shared_mem_bytes: 0,
33435 };
33436 let hki = hk as i32;
33437 let __s_b = self.gpu.stream();
33438 let mut b = __s_b.launch_builder(&f);
33439 b.arg(kb16_ref)
33440 .arg(&gcum)
33441 .arg(beta)
33442 .arg(&u)
33443 .arg(&wb16)
33444 .arg(qb16)
33445 .arg(pb16)
33446 .arg(o)
33447 .arg(&scale)
33448 .arg(state_in)
33449 .arg(&mut *state_out)
33450 .arg(&hi)
33451 .arg(&ti)
33452 .arg(&ci)
33453 .arg(&hki);
33454 unsafe {
33455 b.launch(cfg)?;
33456 }
33457 }
33458 return Ok(());
33459 }
33460 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
33464 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
33465 {
33466 let f = self.func("gdn_chunk_state_mma");
33467 let cfg = LaunchConfig {
33468 grid_dim: (h as u32, NSPLIT, 1),
33469 block_dim: (256, 1, 1),
33470 shared_mem_bytes: 0,
33471 };
33472 let hki = hk as i32;
33473 let __s_b = self.gpu.stream();
33474 let mut b = __s_b.launch_builder(&f);
33475 b.arg(kb16_ref)
33476 .arg(&gcum)
33477 .arg(beta)
33478 .arg(&u)
33479 .arg(&wb16)
33480 .arg(&mut y16)
33481 .arg(&mut ssnap16)
33482 .arg(state_in)
33483 .arg(&mut *state_out)
33484 .arg(&hi)
33485 .arg(&ti)
33486 .arg(&ci)
33487 .arg(&hki);
33488 unsafe {
33489 b.launch(cfg)?;
33490 }
33491 }
33492 {
33493 let f = self.func("gdn_chunk_output_mma");
33495 #[allow(clippy::manual_div_ceil)]
33496 let jt = ((c + 31) / 32) as u32;
33498 let cfg = LaunchConfig {
33499 grid_dim: (nc as u32, h as u32, jt),
33500 block_dim: (256, 1, 1),
33501 shared_mem_bytes: 0,
33502 };
33503 let hki = hk as i32;
33504 let __s_b = self.gpu.stream();
33505 let mut b = __s_b.launch_builder(&f);
33506 b.arg(q)
33507 .arg(&gcum)
33508 .arg(&p)
33509 .arg(&y16)
33510 .arg(&ssnap16)
33511 .arg(o)
33512 .arg(&hi)
33513 .arg(&ti)
33514 .arg(&ci)
33515 .arg(&scale)
33516 .arg(&hki);
33517 unsafe {
33518 b.launch(cfg)?;
33519 }
33520 }
33521 return Ok(());
33522 }
33523 {
33524 let f = self.func("gdn_chunk_state_f32");
33526 let cfg = LaunchConfig {
33527 grid_dim: (h as u32, NSPLIT, 1),
33528 block_dim: (256, 1, 1),
33529 shared_mem_bytes: 0,
33530 };
33531 let __s_b = self.gpu.stream();
33532 let mut b = __s_b.launch_builder(&f);
33533 b.arg(k)
33534 .arg(&gcum)
33535 .arg(beta)
33536 .arg(&u)
33537 .arg(&w)
33538 .arg(&mut y)
33539 .arg(&mut ssnap)
33540 .arg(state_in)
33541 .arg(&mut *state_out)
33542 .arg(&hi)
33543 .arg(&ti)
33544 .arg(&ci);
33545 unsafe {
33546 b.launch(cfg)?;
33547 }
33548 }
33549 {
33550 let f = self.func("gdn_chunk_output_f32");
33552 #[allow(clippy::manual_div_ceil)]
33553 let jt = ((c + 31) / 32) as u32;
33555 let cfg = LaunchConfig {
33556 grid_dim: (nc as u32, h as u32, jt),
33557 block_dim: (256, 1, 1),
33558 shared_mem_bytes: 0,
33559 };
33560 let __s_b = self.gpu.stream();
33561 let mut b = __s_b.launch_builder(&f);
33562 b.arg(q)
33563 .arg(&gcum)
33564 .arg(&p)
33565 .arg(&y)
33566 .arg(&ssnap)
33567 .arg(o)
33568 .arg(&hi)
33569 .arg(&ti)
33570 .arg(&ci)
33571 .arg(&scale);
33572 unsafe {
33573 b.launch(cfg)?;
33574 }
33575 }
33576 Ok(())
33577 }
33578
33579 #[allow(clippy::too_many_arguments)]
33588 #[allow(clippy::too_many_arguments)]
33589 pub fn gdn_scan_prefill(
33590 &self,
33591 q: &CudaSlice<f32>,
33592 k: &CudaSlice<f32>,
33593 v: &CudaSlice<f32>,
33594 g: &CudaSlice<f32>,
33595 beta: &CudaSlice<f32>,
33596 kb16_pre: Option<&CudaSlice<u8>>,
33597 qb16_pre: Option<&CudaSlice<u8>>,
33598 state_in: &CudaSlice<f32>,
33599 state_out: &mut CudaSlice<f32>,
33600 o: &mut CudaSlice<f32>,
33601 n_head: usize,
33602 t: usize,
33603 scale: f32,
33604 hk: usize,
33605 ) -> Result<(), Box<dyn std::error::Error>> {
33606 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
33607 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
33608 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
33609 }
33610 if Self::gdn_chunked_enabled() && t >= 16 {
33611 self.gdn_scan_chunked(
33612 q,
33613 k,
33614 v,
33615 g,
33616 beta,
33617 kb16_pre,
33618 qb16_pre,
33619 state_in,
33620 state_out,
33621 o,
33622 n_head,
33623 t,
33624 scale,
33625 Self::gdn_chunk_size(),
33626 hk,
33627 )
33628 } else {
33629 assert!(
33630 hk == n_head,
33631 "s128 scan is broadcast-only (prep guarantees by predicate)"
33632 );
33633 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
33634 }
33635 }
33636
33637 #[allow(clippy::too_many_arguments)]
33639 fn gdn_scan_diff(
33640 &self,
33641 q: &CudaSlice<f32>,
33642 k: &CudaSlice<f32>,
33643 v: &CudaSlice<f32>,
33644 g: &CudaSlice<f32>,
33645 beta: &CudaSlice<f32>,
33646 state_in: &CudaSlice<f32>,
33647 state_out: &mut CudaSlice<f32>,
33648 o: &mut CudaSlice<f32>,
33649 n_head: usize,
33650 t: usize,
33651 scale: f32,
33652 ) -> Result<(), Box<dyn std::error::Error>> {
33653 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
33654 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
33655 let mut o_c = self.uninit(o.len())?;
33656 let mut st_c = self.uninit(state_out.len())?;
33657 self.gdn_scan_chunked(
33658 q,
33659 k,
33660 v,
33661 g,
33662 beta,
33663 None,
33664 None,
33665 state_in,
33666 &mut st_c,
33667 &mut o_c,
33668 n_head,
33669 t,
33670 scale,
33671 Self::gdn_chunk_size(),
33672 n_head,
33673 )?;
33674 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
33675 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
33676 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
33677 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
33678 let mut max_abs = 0f32;
33679 let mut max_rel = 0f32;
33680 let mut sum_rel = 0f64;
33681 for (x, y) in a.iter().zip(b) {
33682 let ad = (x - y).abs();
33683 let rel = ad / x.abs().max(y.abs()).max(1e-3);
33684 if ad > max_abs {
33685 max_abs = ad;
33686 }
33687 if rel > max_rel {
33688 max_rel = rel;
33689 }
33690 sum_rel += rel as f64;
33691 }
33692 (max_abs, max_rel, sum_rel / a.len() as f64)
33693 };
33694 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
33695 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
33696 println!(
33697 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
33698 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
33699 Self::gdn_chunk_size()
33700 );
33701 Ok(())
33702 }
33703
33704 pub fn gdn_glog(
33706 &self,
33707 alpha: &CudaSlice<f32>,
33708 dt_bias: &CudaSlice<f32>,
33709 a: &CudaSlice<f32>,
33710 g_log: &mut CudaSlice<f32>,
33711 n_head: usize,
33712 t: usize,
33713 ) -> Result<(), Box<dyn std::error::Error>> {
33714 let f = self.func("gdn_glog_f32");
33715 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
33716 let (h, ti) = (n_head as i32, t as i32);
33717 let __s_b = self.gpu.stream();
33718 let mut b = __s_b.launch_builder(&f);
33719 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
33720 unsafe {
33721 b.launch(cfg)?;
33722 }
33723 Ok(())
33724 }
33725
33726 pub fn sigmoid_v(
33729 &self,
33730 x: &cudarc::driver::CudaView<f32>,
33731 y: &mut CudaSlice<f32>,
33732 n: usize,
33733 ) -> Result<(), Box<dyn std::error::Error>> {
33734 let f = self.func("sigmoid_f32");
33735 let cfg = LaunchConfig::for_num_elems(n as u32);
33736 let ni = n as i32;
33737 let __s_b = self.gpu.stream();
33738 let mut b = __s_b.launch_builder(&f);
33739 b.arg(x).arg(y).arg(&ni);
33740 unsafe {
33741 b.launch(cfg)?;
33742 }
33743 Ok(())
33744 }
33745
33746 pub fn gdn_glog_v(
33747 &self,
33748 alpha: &cudarc::driver::CudaView<f32>,
33749 dt_bias: &CudaSlice<f32>,
33750 a: &CudaSlice<f32>,
33751 g_log: &mut CudaSlice<f32>,
33752 n_head: usize,
33753 t: usize,
33754 ) -> Result<(), Box<dyn std::error::Error>> {
33755 let f = self.func("gdn_glog_f32");
33756 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
33757 let (h, ti) = (n_head as i32, t as i32);
33758 let __s_b = self.gpu.stream();
33759 let mut b = __s_b.launch_builder(&f);
33760 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
33761 unsafe {
33762 b.launch(cfg)?;
33763 }
33764 Ok(())
33765 }
33766
33767 pub fn sigmoid(
33768 &self,
33769 x: &CudaSlice<f32>,
33770 y: &mut CudaSlice<f32>,
33771 n: usize,
33772 ) -> Result<(), Box<dyn std::error::Error>> {
33773 let f = self.func("sigmoid_f32");
33774 let cfg = LaunchConfig::for_num_elems(n as u32);
33775 let ni = n as i32;
33776 let __s_b = self.gpu.stream();
33777 let mut b = __s_b.launch_builder(&f);
33778 b.arg(x).arg(y).arg(&ni);
33779 unsafe {
33780 b.launch(cfg)?;
33781 }
33782 Ok(())
33783 }
33784
33785 pub fn sig_mul_f16out(
33788 &self,
33789 a: &CudaSlice<f32>,
33790 g: &CudaSlice<f32>,
33791 dst: &mut CudaSlice<f32>,
33792 dst16: &mut CudaSlice<u8>,
33793 n: usize,
33794 ) -> Result<(), Box<dyn std::error::Error>> {
33795 let f = self.func("sig_mul_f16out_f32");
33796 let cfg = LaunchConfig::for_num_elems(n as u32);
33797 let ni = n as i32;
33798 let __s_b = self.gpu.stream();
33799 let mut b = __s_b.launch_builder(&f);
33800 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
33801 unsafe {
33802 b.launch(cfg)?;
33803 }
33804 Ok(())
33805 }
33806
33807 #[allow(clippy::too_many_arguments)]
33816 pub fn attn_head_gate(
33817 &self,
33818 a: &CudaSlice<f32>,
33819 g: &CudaSlice<f32>,
33820 dst: &mut CudaSlice<f32>,
33821 dst16: Option<&mut CudaSlice<u8>>,
33822 head_dim: usize,
33823 n_head: usize,
33824 t: usize,
33825 ) -> Result<(), Box<dyn std::error::Error>> {
33826 let f = self.func("attn_head_gate_f32");
33827 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
33828 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
33829 let d16: u64 = match dst16 {
33831 Some(d) => self.addr_u8(d),
33832 None => 0,
33833 };
33834 let __s_b = self.gpu.stream();
33835 let mut b = __s_b.launch_builder(&f);
33836 b.arg(a)
33837 .arg(g)
33838 .arg(dst)
33839 .arg(&d16)
33840 .arg(&hd)
33841 .arg(&nh)
33842 .arg(&ti);
33843 unsafe {
33844 b.launch(cfg)?;
33845 }
33846 Ok(())
33847 }
33848
33849 #[allow(clippy::too_many_arguments)]
33858 pub fn swiglu_clamped_mul_scaled(
33859 &self,
33860 gate: &CudaSlice<f32>,
33861 up: &CudaSlice<f32>,
33862 gs: f32,
33863 us: f32,
33864 limit: f32,
33865 dst: &mut CudaSlice<f32>,
33866 n: usize,
33867 ) -> Result<(), Box<dyn std::error::Error>> {
33868 debug_assert!(
33869 limit > 1e-6,
33870 "swiglu_clamped needs a live limit; use silu_mul_scaled"
33871 );
33872 let f = self.func("swiglu_clamped_mul_scaled_f32");
33873 let cfg = LaunchConfig::for_num_elems(n as u32);
33874 let ni = n as i32;
33875 let __s_b = self.gpu.stream();
33876 let mut b = __s_b.launch_builder(&f);
33877 b.arg(gate)
33878 .arg(up)
33879 .arg(&gs)
33880 .arg(&us)
33881 .arg(&limit)
33882 .arg(dst)
33883 .arg(&ni);
33884 unsafe {
33885 b.launch(cfg)?;
33886 }
33887 Ok(())
33888 }
33889
33890 #[allow(clippy::too_many_arguments)]
33899 pub fn swiglu_preclamped_mul_scaled(
33900 &self,
33901 gate: &CudaSlice<f32>,
33902 up: &CudaSlice<f32>,
33903 gs: f32,
33904 us: f32,
33905 limit: f32,
33906 dst: &mut CudaSlice<f32>,
33907 n: usize,
33908 ) -> Result<(), Box<dyn std::error::Error>> {
33909 debug_assert!(
33910 limit > 1e-6,
33911 "swiglu_preclamped needs a live limit; use silu_mul_scaled"
33912 );
33913 let f = self.func("swiglu_preclamped_mul_scaled_f32");
33914 let cfg = LaunchConfig::for_num_elems(n as u32);
33915 let ni = n as i32;
33916 let __s_b = self.gpu.stream();
33917 let mut b = __s_b.launch_builder(&f);
33918 b.arg(gate)
33919 .arg(up)
33920 .arg(&gs)
33921 .arg(&us)
33922 .arg(&limit)
33923 .arg(dst)
33924 .arg(&ni);
33925 unsafe {
33926 b.launch(cfg)?;
33927 }
33928 Ok(())
33929 }
33930
33931 #[allow(clippy::too_many_arguments)] pub fn gated_rmsnorm(
33934 &self,
33935 o: &CudaSlice<f32>,
33936 w: &CudaSlice<f32>,
33937 z: &CudaSlice<f32>,
33938 dst: &mut CudaSlice<f32>,
33939 ncols: usize,
33940 nrows: usize,
33941 eps: f32,
33942 ) -> Result<(), Box<dyn std::error::Error>> {
33943 let f = self.func("gated_rmsnorm_f32");
33944 let cfg = LaunchConfig {
33945 grid_dim: (nrows as u32, 1, 1),
33946 block_dim: (128, 1, 1),
33947 shared_mem_bytes: 0,
33948 };
33949 let (nc, e) = (ncols as i32, eps);
33950 let __s_b = self.gpu.stream();
33951 let mut b = __s_b.launch_builder(&f);
33952 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
33953 unsafe {
33954 b.launch(cfg)?;
33955 }
33956 Ok(())
33957 }
33958
33959 #[allow(clippy::too_many_arguments)] pub fn gated_rmsnorm_f16out(
33963 &self,
33964 o: &CudaSlice<f32>,
33965 w: &CudaSlice<f32>,
33966 z: &CudaSlice<f32>,
33967 dst: &mut CudaSlice<f32>,
33968 dst16: &mut CudaSlice<u8>,
33969 ncols: usize,
33970 nrows: usize,
33971 eps: f32,
33972 ) -> Result<(), Box<dyn std::error::Error>> {
33973 let f = self.func("gated_rmsnorm_f16out_f32");
33974 let cfg = LaunchConfig {
33976 grid_dim: (nrows as u32, 1, 1),
33977 block_dim: (128, 1, 1),
33978 shared_mem_bytes: 0,
33979 };
33980 let (nc, e) = (ncols as i32, eps);
33981 let __s_b = self.gpu.stream();
33982 let mut b = __s_b.launch_builder(&f);
33983 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
33984 unsafe {
33985 b.launch(cfg)?;
33986 }
33987 Ok(())
33988 }
33989
33990 #[allow(clippy::too_many_arguments)]
33994 pub fn add_rms_norm_zq8(
33995 &self,
33996 a: &CudaSlice<f32>,
33997 b_in: &CudaSlice<f32>,
33998 w: &CudaSlice<f32>,
33999 res: &mut CudaSlice<f32>,
34000 z: &mut CudaSlice<f32>,
34001 ncols: usize,
34002 nrows: usize,
34003 eps: f32,
34004 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
34005 assert!(ncols.is_multiple_of(32));
34006 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
34007 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
34008 let f = self.func("add_rms_norm_zq8");
34009 let cfg = LaunchConfig {
34010 grid_dim: (nrows as u32, 1, 1),
34011 block_dim: (1024, 1, 1),
34012 shared_mem_bytes: 0,
34013 };
34014 let (nc, ep) = (ncols as i32, eps);
34015 let __s_b = self.gpu.stream();
34016 let mut b = __s_b.launch_builder(&f);
34017 b.arg(a)
34018 .arg(b_in)
34019 .arg(w)
34020 .arg(res)
34021 .arg(z)
34022 .arg(&mut q)
34023 .arg(&mut d)
34024 .arg(&nc)
34025 .arg(&ep);
34026 unsafe {
34027 b.launch(cfg)?;
34028 }
34029 Ok((q, d))
34030 }
34031
34032 #[allow(clippy::too_many_arguments)] pub fn gated_rmsnorm_zv(
34038 &self,
34039 o: &CudaSlice<f32>,
34040 w: &CudaSlice<f32>,
34041 z: &cudarc::driver::CudaView<f32>,
34042 dst: &mut CudaSlice<f32>,
34043 ncols: usize,
34044 nrows: usize,
34045 eps: f32,
34046 ) -> Result<(), Box<dyn std::error::Error>> {
34047 let f = self.func("gated_rmsnorm_f32");
34048 let cfg = LaunchConfig {
34049 grid_dim: (nrows as u32, 1, 1),
34050 block_dim: (128, 1, 1),
34051 shared_mem_bytes: 0,
34052 };
34053 let (nc, e) = (ncols as i32, eps);
34054 let __s_b = self.gpu.stream();
34055 let mut b = __s_b.launch_builder(&f);
34056 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
34057 unsafe {
34058 b.launch(cfg)?;
34059 }
34060 Ok(())
34061 }
34062
34063 #[allow(clippy::too_many_arguments)] pub fn gated_rmsnorm_f16out_zv(
34065 &self,
34066 o: &CudaSlice<f32>,
34067 w: &CudaSlice<f32>,
34068 z: &cudarc::driver::CudaView<f32>,
34069 dst: &mut CudaSlice<f32>,
34070 dst16: &mut CudaSlice<u8>,
34071 ncols: usize,
34072 nrows: usize,
34073 eps: f32,
34074 ) -> Result<(), Box<dyn std::error::Error>> {
34075 let f = self.func("gated_rmsnorm_f16out_f32");
34076 let cfg = LaunchConfig {
34078 grid_dim: (nrows as u32, 1, 1),
34079 block_dim: (128, 1, 1),
34080 shared_mem_bytes: 0,
34081 };
34082 let (nc, e) = (ncols as i32, eps);
34083 let __s_b = self.gpu.stream();
34084 let mut b = __s_b.launch_builder(&f);
34085 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
34086 unsafe {
34087 b.launch(cfg)?;
34088 }
34089 Ok(())
34090 }
34091
34092 pub fn gated_rmsnorm_q8_1(
34093 &self,
34094 o: &CudaSlice<f32>,
34095 w: &CudaSlice<f32>,
34096 z: &CudaSlice<f32>,
34097 ncols: usize,
34098 nrows: usize,
34099 eps: f32,
34100 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
34101 assert!(ncols.is_multiple_of(32));
34102 let f = self.func("gated_rmsnorm_q8_1");
34103 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
34104 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
34105 let cfg = LaunchConfig {
34106 grid_dim: (nrows as u32, 1, 1),
34107 block_dim: (128, 1, 1),
34108 shared_mem_bytes: 0,
34109 };
34110 let (nc, ep) = (ncols as i32, eps);
34111 let __s_b = self.gpu.stream();
34112 let mut b = __s_b.launch_builder(&f);
34113 b.arg(o)
34114 .arg(w)
34115 .arg(z)
34116 .arg(&mut out_q)
34117 .arg(&mut out_d)
34118 .arg(&nc)
34119 .arg(&ep);
34120 unsafe {
34121 b.launch(cfg)?;
34122 }
34123 Ok((out_q, out_d))
34124 }
34125
34126 pub fn transpose(
34128 &self,
34129 inp: &CudaSlice<f32>,
34130 rows: usize,
34131 cols: usize,
34132 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34133 let f = self.func("transpose_f32");
34134 let mut out = self.zeros(rows * cols)?;
34135 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
34136 let (r, c) = (rows as i32, cols as i32);
34137 let __s_b = self.gpu.stream();
34138 let mut b = __s_b.launch_builder(&f);
34139 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
34140 unsafe {
34141 b.launch(cfg)?;
34142 }
34143 Ok(out)
34144 }
34145
34146 pub fn repeat_heads(
34148 &self,
34149 inp: &CudaSlice<f32>,
34150 out: &mut CudaSlice<f32>,
34151 head_dim: usize,
34152 n_in: usize,
34153 n_out: usize,
34154 t: usize,
34155 ) -> Result<(), Box<dyn std::error::Error>> {
34156 let f = self.func("repeat_heads_f32");
34157 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
34158 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
34159 let __s_b = self.gpu.stream();
34160 let mut b = __s_b.launch_builder(&f);
34161 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
34162 unsafe {
34163 b.launch(cfg)?;
34164 }
34165 Ok(())
34166 }
34167
34168 pub fn q_gate_split(
34175 &self,
34176 qf: &CudaSlice<f32>,
34177 q_out: &mut CudaSlice<f32>,
34178 gate_out: &mut CudaSlice<f32>,
34179 head_dim: usize,
34180 n_head: usize,
34181 t: usize,
34182 ) -> Result<(), Box<dyn std::error::Error>> {
34183 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
34184 let out_need = head_dim * n_head * t;
34185 if q_out.len() < out_need || gate_out.len() < out_need {
34186 return Err(format!(
34187 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
34188 q_out.len(),
34189 gate_out.len()
34190 )
34191 .into());
34192 }
34193 let f = self.func("q_gate_split_f32");
34194 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
34195 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
34196 let __s_b = self.gpu.stream();
34197 let mut b = __s_b.launch_builder(&f);
34198 b.arg(qf)
34199 .arg(q_out)
34200 .arg(gate_out)
34201 .arg(&hd)
34202 .arg(&nh)
34203 .arg(&ti);
34204 unsafe {
34205 b.launch(cfg)?;
34206 }
34207 Ok(())
34208 }
34209
34210 #[allow(clippy::too_many_arguments)] pub fn qkv_to_gdn_repack(
34215 &self,
34216 conv_out: &CudaSlice<f32>,
34217 q_g: &mut CudaSlice<f32>,
34218 k_g: &mut CudaSlice<f32>,
34219 v_g: &mut CudaSlice<f32>,
34220 d_state: usize,
34221 num_v: usize,
34222 num_k: usize,
34223 key_dim: usize,
34224 t: usize,
34225 ) -> Result<(), Box<dyn std::error::Error>> {
34226 let f = self.func("qkv_to_gdn_repack_f32");
34227 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
34228 let (ds, nv, nk, kd, ti) = (
34229 d_state as i32,
34230 num_v as i32,
34231 num_k as i32,
34232 key_dim as i32,
34233 t as i32,
34234 );
34235 let __s_b = self.gpu.stream();
34236 let mut b = __s_b.launch_builder(&f);
34237 b.arg(conv_out)
34238 .arg(q_g)
34239 .arg(k_g)
34240 .arg(v_g)
34241 .arg(&ds)
34242 .arg(&nv)
34243 .arg(&nk)
34244 .arg(&kd)
34245 .arg(&ti);
34246 unsafe {
34247 b.launch(cfg)?;
34248 }
34249 Ok(())
34250 }
34251
34252 pub fn conv_left_pad(
34255 &self,
34256 src: &CudaSlice<f32>,
34257 dst: &mut CudaSlice<f32>,
34258 conv_dim: usize,
34259 t: usize,
34260 pad: usize,
34261 ) -> Result<(), Box<dyn std::error::Error>> {
34262 let f = self.func("conv_left_pad_f32");
34263 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
34264 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
34265 let __s_b = self.gpu.stream();
34266 let mut b = __s_b.launch_builder(&f);
34267 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
34268 unsafe {
34269 b.launch(cfg)?;
34270 }
34271 Ok(())
34272 }
34273
34274 pub fn conv_assemble_and_roll(
34278 &self,
34279 qkv_col: &CudaSlice<f32>,
34280 conv_state: &mut CudaSlice<f32>,
34281 conv_in: &mut CudaSlice<f32>,
34282 conv_dim: usize,
34283 pad: usize,
34284 ) -> Result<(), Box<dyn std::error::Error>> {
34285 let f = self.func("conv_assemble_and_roll_f32");
34286 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
34287 let (cd, p) = (conv_dim as i32, pad as i32);
34288 let __s_b = self.gpu.stream();
34289 let mut b = __s_b.launch_builder(&f);
34290 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
34291 unsafe {
34292 b.launch(cfg)?;
34293 }
34294 Ok(())
34295 }
34296
34297 pub fn ssm_conv1d_fused_decode(
34303 &self,
34304 qkv_col: &CudaSlice<f32>,
34305 conv_state: &mut CudaSlice<f32>,
34306 w: &CudaSlice<f32>,
34307 conv_out: &mut CudaSlice<f32>,
34308 conv_dim: usize,
34309 d_conv: usize,
34310 ) -> Result<(), Box<dyn std::error::Error>> {
34311 let f = self.func("ssm_conv1d_fused_decode_f32");
34312 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
34313 let (cd, dc) = (conv_dim as i32, d_conv as i32);
34314 let __s_b = self.gpu.stream();
34315 let mut b = __s_b.launch_builder(&f);
34316 b.arg(qkv_col)
34317 .arg(conv_state)
34318 .arg(w)
34319 .arg(conv_out)
34320 .arg(&cd)
34321 .arg(&dc);
34322 unsafe {
34323 b.launch(cfg)?;
34324 }
34325 Ok(())
34326 }
34327
34328 pub fn slice_range(
34331 &self,
34332 src: &CudaSlice<f32>,
34333 start: usize,
34334 len: usize,
34335 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34336 let host = self.gpu.stream().clone_dtoh(src)?;
34337 self.gpu.stream().synchronize()?;
34338 self.htod(&host[start..start + len])
34339 }
34340}
34341
34342#[cfg(test)]
34343mod target_dispatch_tests {
34344 use super::legacy_quant_gemm_allowed;
34345
34346 #[test]
34347 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
34348 assert!(legacy_quant_gemm_allowed(false, false, false));
34350 assert!(!legacy_quant_gemm_allowed(false, false, true));
34351 assert!(!legacy_quant_gemm_allowed(true, false, false));
34353 assert!(!legacy_quant_gemm_allowed(true, false, true));
34354 assert!(legacy_quant_gemm_allowed(true, true, false));
34356 assert!(!legacy_quant_gemm_allowed(true, true, true));
34357 }
34358
34359 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
34360 #[test]
34361 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
34362 assert!(!legacy_quant_gemm_allowed(
34363 cfg!(memra_portable_cuda),
34364 cfg!(memra_hopper_mma),
34365 false
34366 ));
34367 }
34368
34369 #[cfg(memra_hopper_mma)]
34370 #[test]
34371 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
34372 assert!(legacy_quant_gemm_allowed(
34373 cfg!(memra_portable_cuda),
34374 cfg!(memra_hopper_mma),
34375 false
34376 ));
34377 assert!(super::portable_mma_gated() == false);
34378 }
34379}
34380
34381impl memra_kv::KvDev for Engine {
34384 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34385 Engine::zeros(self, n)
34386 }
34387 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34388 Engine::uninit(self, n)
34389 }
34390 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
34391 Engine::alloc_u8(self, n)
34392 }
34393 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
34394 Engine::htod_i32(self, v)
34395 }
34396 fn clone_dtod(
34397 &self,
34398 src: &CudaSlice<f32>,
34399 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
34400 Engine::clone_dtod(self, src)
34401 }
34402 fn copy_into(
34403 &self,
34404 dst: &mut CudaSlice<f32>,
34405 off: usize,
34406 src: &CudaSlice<f32>,
34407 len: usize,
34408 ) -> Result<(), Box<dyn std::error::Error>> {
34409 Engine::copy_into(self, dst, off, src, len)
34410 }
34411 fn copy_range_into(
34412 &self,
34413 dst: &mut CudaSlice<f32>,
34414 dst_off: usize,
34415 src: &CudaSlice<f32>,
34416 src_off: usize,
34417 len: usize,
34418 ) -> Result<(), Box<dyn std::error::Error>> {
34419 Engine::copy_range_into(self, dst, dst_off, src, src_off, len)
34420 }
34421 fn set_i32_one(
34422 &self,
34423 d: &mut CudaSlice<i32>,
34424 v: i32,
34425 ) -> Result<(), Box<dyn std::error::Error>> {
34426 Engine::set_i32_one(self, d, v)
34427 }
34428}
34429
34430#[cfg(test)]
34431mod fused_gate_bounds_tests {
34432 use super::*;
34433
34434 #[test]
34447 #[ignore = "requires a CUDA GPU"]
34448 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
34449 let e = Engine::new(0).unwrap();
34450 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
34451 let fused = 2 * head_dim * n_head * t;
34452 let out_n = head_dim * n_head * t;
34453
34454 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
34456 let mut q = e.uninit(out_n).unwrap();
34457 let mut gate = e.uninit(out_n).unwrap();
34458 let err = e
34459 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
34460 .expect_err("half-width wq must be refused, not read past")
34461 .to_string();
34462 assert!(err.contains("NO fused gate"), "{err}");
34463 assert!(err.contains(&format!("{fused}")), "{err}");
34464
34465 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
34468 let wide = e.htod(&host).unwrap();
34469 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
34470 .expect("full-width wq splits");
34471 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
34472 for tok in 0..t {
34473 for hh in 0..n_head {
34474 for d in 0..head_dim {
34475 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
34476 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
34477 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
34478 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
34479 }
34480 }
34481 }
34482
34483 let mut small = e.uninit(out_n - 1).unwrap();
34485 assert!(
34486 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
34487 .is_err()
34488 );
34489 }
34490}
34491
34492#[cfg(test)]
34496mod fused_rope_width_tests {
34497 use super::Engine;
34498
34499 #[test]
34502 fn full_width_is_accepted() {
34503 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
34504 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
34505 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
34506 }
34507
34508 #[test]
34521 fn gemma4_official_artifact_widths_pass() {
34522 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
34523 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
34524 }
34525
34526 #[test]
34529 fn partial_rotary_is_refused_with_the_geometry_named() {
34530 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
34532 .expect_err("partial rotary must refuse");
34533 let msg = err.to_string();
34534 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
34535 assert!(msg.contains("n_rot 64"), "{msg}");
34536 assert!(msg.contains("head_dim 256"), "{msg}");
34537 assert!(
34538 msg.contains("64..256"),
34539 "names the band it would corrupt: {msg}"
34540 );
34541 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
34543 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
34545 }
34546}